diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/TARGETS b/.Package.swift/kernels_torchao/dummy.swift similarity index 100% rename from examples/demo-apps/apple_ios/LLaMA/LLaMARunner/TARGETS rename to .Package.swift/kernels_torchao/dummy.swift diff --git a/examples/demo-apps/apple_ios/LLaMA/TARGETS b/.Package.swift/kernels_torchao_debug/dummy.swift similarity index 100% rename from examples/demo-apps/apple_ios/LLaMA/TARGETS rename to .Package.swift/kernels_torchao_debug/dummy.swift diff --git a/.ci/docker/ci_commit_pins/optimum-executorch.txt b/.ci/docker/ci_commit_pins/optimum-executorch.txt index 9c1dac7fa91..ef3282ba6cc 100644 --- a/.ci/docker/ci_commit_pins/optimum-executorch.txt +++ b/.ci/docker/ci_commit_pins/optimum-executorch.txt @@ -1 +1 @@ -36e3dd54effb3f6d13d792029609292fdd5502bb +40b02a2dc61bbf901a2df91719f47c98d65368ec diff --git a/.ci/docker/ci_commit_pins/pytorch.txt b/.ci/docker/ci_commit_pins/pytorch.txt index 6305196d2ad..8c9330d6f2c 100644 --- a/.ci/docker/ci_commit_pins/pytorch.txt +++ b/.ci/docker/ci_commit_pins/pytorch.txt @@ -1 +1 @@ -6fc0ad22f0a07b6f38d138861c56a765d5a9bb02 +4d4abec80f03cd8fdefe1d9cb3a60d3690cd777e diff --git a/.ci/scripts/build-qnn-sdk.sh b/.ci/scripts/build-qnn-sdk.sh index 971581eb053..7f34e8afb63 100755 --- a/.ci/scripts/build-qnn-sdk.sh +++ b/.ci/scripts/build-qnn-sdk.sh @@ -11,8 +11,10 @@ set -o xtrace build_qnn_backend() { echo "Start building qnn backend." - export ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT:-/opt/ndk} - export QNN_SDK_ROOT=${QNN_SDK_ROOT:-/tmp/qnn/2.28.0.241029} + # Source QNN configuration + source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh" + setup_android_ndk + install_qnn export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" parallelism=$(( $(nproc) - 1 )) diff --git a/.ci/scripts/setup-conda.sh b/.ci/scripts/setup-conda.sh index 5466cc0d60d..a725c90dd82 100755 --- a/.ci/scripts/setup-conda.sh +++ b/.ci/scripts/setup-conda.sh @@ -9,7 +9,7 @@ set -ex install_conda() { pushd .ci/docker || return - ${CONDA_INSTALL} -y --file conda-env-ci.txt + ${CONDA_INSTALL} -c conda-forge -y --file conda-env-ci.txt popd || return } diff --git a/.ci/scripts/setup-linux.sh b/.ci/scripts/setup-linux.sh index a090571ab49..feb8a128b17 100755 --- a/.ci/scripts/setup-linux.sh +++ b/.ci/scripts/setup-linux.sh @@ -11,6 +11,7 @@ set -exu source "$(dirname "${BASH_SOURCE[0]}")/utils.sh" read -r BUILD_TOOL BUILD_MODE EDITABLE < <(parse_args "$@") +echo "Build tool: $BUILD_TOOL, Mode: $BUILD_MODE" # As Linux job is running inside a Docker container, all of its dependencies # have already been installed, so we use PyTorch build from source here instead diff --git a/.ci/scripts/setup-qnn-deps.sh b/.ci/scripts/setup-qnn-deps.sh index d962f7b908a..0de9db5c030 100644 --- a/.ci/scripts/setup-qnn-deps.sh +++ b/.ci/scripts/setup-qnn-deps.sh @@ -10,4 +10,5 @@ set -ex source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh" setup_libcpp 12 -install_qnn +setup_android_ndk +install_qnn \ No newline at end of file diff --git a/.ci/scripts/setup-samsung-linux-deps.sh b/.ci/scripts/setup-samsung-linux-deps.sh new file mode 100644 index 00000000000..ed704b2bfbd --- /dev/null +++ b/.ci/scripts/setup-samsung-linux-deps.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 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. + +set -ex + + +download_ai_lite_core() { + API_BASE="https://soc-developer.semiconductor.samsung.com/api/v1/resource/ai-litecore/download" + API_KEY="kn10SoSY3hkC-9Qny5TqD2mnqVrlupv3krnjLeBt5cY" + + VERSION="0.5" + OS_NAME="Ubuntu 22.04" + OUT_FILE="/tmp/exynos-ai-litecore-v${VERSION}.tar.gz" + TARGET_PATH="/tmp/exynos_ai_lite_core" + + mkdir -p ${TARGET_PATH} + # Presigned issue URL + JSON_RESP=$(curl -sS -G \ + --location --fail --retry 3 \ + -H "apikey: ${API_KEY}" \ + --data-urlencode "version=${VERSION}" \ + --data-urlencode "os=${OS_NAME}" \ + "${API_BASE}") + + DOWNLOAD_URL=$(echo "$JSON_RESP" | sed -n 's/.*"data":[[:space:]]*"\([^"]*\)".*/\1/p') + + if [[ -z "$DOWNLOAD_URL" ]]; then + echo "Failed to extract download URL" + echo "$JSON_RESP" + exit 1 + fi + + # Download LiteCore + curl -sS -L --fail --retry 3 \ + --output "$OUT_FILE" \ + "$DOWNLOAD_URL" + + echo "Download done: $OUT_FILE" + + + tar -C "${TARGET_PATH}" --strip-components=1 -xzvf "${OUT_FILE}" + + export EXYNOS_AI_LITECORE_ROOT=${TARGET_PATH} + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:${EXYNOS_AI_LITECORE_ROOT}/lib/x86_64-linux +} + +install_enn_backend() { + NDK_INSTALLATION_DIR=/opt/ndk + rm -rf "${NDK_INSTALLATION_DIR}" && sudo mkdir -p "${NDK_INSTALLATION_DIR}" + ANDROID_NDK_VERSION=r27b + + # build Exynos backend + export ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT:-/opt/ndk} + bash backends/samsung/build.sh --build all + # set env variable + export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." && pwd)" + export PYTHONPATH=${PYTHONPATH:-}:${EXECUTORCH_ROOT}/.. +} + +AI_LITE_CORE_VERSION=0.5.0 + +download_ai_lite_core ${AI_LITE_CORE_VERSION} +install_enn_backend diff --git a/.ci/scripts/setup-vulkan-linux-deps.sh b/.ci/scripts/setup-vulkan-linux-deps.sh index c0b2596f20e..cd99ff0d6ff 100755 --- a/.ci/scripts/setup-vulkan-linux-deps.sh +++ b/.ci/scripts/setup-vulkan-linux-deps.sh @@ -23,6 +23,7 @@ install_swiftshader() { export VK_ICD_FILENAMES="${_swiftshader_dir}/swiftshader/build/Linux/vk_swiftshader_icd.json" export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/" + export ETVK_USING_SWIFTSHADER=1 } install_vulkan_sdk() { @@ -42,7 +43,7 @@ install_vulkan_sdk() { export PATH="${PATH}:${_vulkan_sdk_dir}/${VULKAN_SDK_VERSION}/x86_64/bin/" } -VULKAN_SDK_VERSION="1.3.296.0" +VULKAN_SDK_VERSION="1.4.321.1" install_swiftshader install_vulkan_sdk "${VULKAN_SDK_VERSION}" diff --git a/.ci/scripts/setup-windows.ps1 b/.ci/scripts/setup-windows.ps1 new file mode 100644 index 00000000000..329e81b3cf0 --- /dev/null +++ b/.ci/scripts/setup-windows.ps1 @@ -0,0 +1,24 @@ +param ( + [string]$editable = "false" +) + +conda create --yes --quiet -n et python=3.12 +conda activate et + +# Activate the VS environment - this is required for Dynamo to work, as it uses MSVC. +# There are a bunch of environment variables that it requires. +# See https://learn.microsoft.com/en-us/cpp/build/building-on-the-command-line. +& "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\Launch-VsDevShell.ps1" -Arch amd64 + +# Install test dependencies +pip install -r .ci/docker/requirements-ci.txt + +if ($editable -eq 'true') { + install_executorch.bat --editable +} else { + install_executorch.bat +} +if ($LASTEXITCODE -ne 0) { + Write-Host "Installation was unsuccessful. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} diff --git a/.ci/scripts/test_backend_linux.sh b/.ci/scripts/test_backend_linux.sh new file mode 100755 index 00000000000..243602fea21 --- /dev/null +++ b/.ci/scripts/test_backend_linux.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env 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. +set -eux + +SUITE=$1 +FLOW=$2 +ARTIFACT_DIR=$3 + +REPORT_FILE="$ARTIFACT_DIR/test-report-$FLOW-$SUITE.csv" + +echo "Running backend test job for suite $SUITE, flow $FLOW." +echo "Saving job artifacts to $ARTIFACT_DIR." + +# The generic Linux job chooses to use base env, not the one setup by the image +eval "$(conda shell.bash hook)" +CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") +conda activate "${CONDA_ENV}" + +export PYTHON_EXECUTABLE=python + +# CMake options to use, in addition to the defaults. +EXTRA_BUILD_ARGS="" + +if [[ "$FLOW" == *qnn* ]]; then + # Setup QNN sdk and deps - note that this is a bit hacky due to the nature of the + # Qualcomm build. TODO (gjcomer) Clean this up once the QNN pybinding integration is + # cleaned up. + PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool cmake + PYTHON_EXECUTABLE=python source .ci/scripts/build-qnn-sdk.sh + QNN_X86_LIB_DIR=`realpath build-x86/lib/` + export LD_LIBRARY_PATH"=$QNN_X86_LIB_DIR:$QNN_SDK_ROOT/lib/x86_64-linux-clang/:${LD_LIBRARY_PATH:-}" + + # TODO Get SDK root from install scripts + EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_QNN=ON -DQNN_SDK_ROOT=$QNN_SDK_ROOT" +fi + +if [[ "$FLOW" == *vulkan* ]]; then + # Setup swiftshader and Vulkan SDK which are required to build the Vulkan delegate + source .ci/scripts/setup-vulkan-linux-deps.sh + + EXTRA_BUILD_ARGS+=" -DEXECUTORCH_BUILD_VULKAN=ON" +fi + +# We need the runner to test the built library. +PYTHON_EXECUTABLE=python CMAKE_ARGS="$EXTRA_BUILD_ARGS" .ci/scripts/setup-linux.sh --build-tool cmake --build-mode Release --editable true + +EXIT_CODE=0 +python -m executorch.backends.test.suite.runner $SUITE --flow $FLOW --report "$REPORT_FILE" || EXIT_CODE=$? + +# Generate markdown summary. +python -m executorch.backends.test.suite.generate_markdown_summary "$REPORT_FILE" > ${GITHUB_STEP_SUMMARY:-"step_summary.md"} --exit-code $EXIT_CODE diff --git a/.ci/scripts/test_backend_macos.sh b/.ci/scripts/test_backend_macos.sh new file mode 100755 index 00000000000..c31fd504b03 --- /dev/null +++ b/.ci/scripts/test_backend_macos.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env 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. +set -eux + +SUITE=$1 +FLOW=$2 +ARTIFACT_DIR=$3 + +REPORT_FILE="$ARTIFACT_DIR/test-report-$FLOW-$SUITE.csv" + +echo "Running backend test job for suite $SUITE, flow $FLOW." +echo "Saving job artifacts to $ARTIFACT_DIR." + +${CONDA_RUN} --no-capture-output pip install awscli==1.37.21 + +bash .ci/scripts/setup-conda.sh +eval "$(conda shell.bash hook)" + +PYTHON_EXECUTABLE=python +${CONDA_RUN} --no-capture-output .ci/scripts/setup-macos.sh --build-tool cmake --build-mode Release + +EXIT_CODE=0 +${CONDA_RUN} --no-capture-output python -m executorch.backends.test.suite.runner $SUITE --flow $FLOW --report "$REPORT_FILE" || EXIT_CODE=$? + +# Generate markdown summary. +${CONDA_RUN} --no-capture-output python -m executorch.backends.test.suite.generate_markdown_summary "$REPORT_FILE" > ${GITHUB_STEP_SUMMARY:-"step_summary.md"} --exit-code $EXIT_CODE diff --git a/.ci/scripts/test_huggingface_optimum_model.py b/.ci/scripts/test_huggingface_optimum_model.py index 8a0b244c549..05b25299522 100644 --- a/.ci/scripts/test_huggingface_optimum_model.py +++ b/.ci/scripts/test_huggingface_optimum_model.py @@ -1,7 +1,11 @@ import argparse +import gc +import logging +import math import subprocess import tempfile from pathlib import Path +from typing import List import torch from datasets import load_dataset @@ -15,6 +19,7 @@ ) from transformers import ( AutoConfig, + AutoModelForCausalLM, AutoModelForImageClassification, AutoProcessor, AutoTokenizer, @@ -37,6 +42,56 @@ def cli_export(command, model_dir): print(f"Export failed with error: {e}") +def check_causal_lm_output_quality( + model_id: str, generated_tokens: List[int], max_perplexity_threshold: float = 100.0 +): + """ + Evaluates the quality of text generated by a causal language model by calculating its perplexity. + + Args: + model_id: HuggingFace model identifier (e.g., "google/gemma2-2b") + generated_tokens: The tokens generated by the exported model to evaluate + max_perplexity_threshold: Maximum acceptable perplexity (lower is better) + + Returns: + tuple: (is_quality_ok, reason) with boolean result and explanation + """ + logging.info(f"Starting perplexity check with model '{model_id}' ...") + # Load model + model = AutoModelForCausalLM.from_pretrained( + model_id, + low_cpu_mem_usage=True, + use_cache=False, + torch_dtype=torch.bfloat16, + ) + + with torch.no_grad(): + outputs = model(input_ids=generated_tokens, labels=generated_tokens) + + # Get the loss (negative log-likelihood) + loss = outputs.loss.item() + + # Calculate perplexity (exp of the average negative log-likelihood) + perplexity = math.exp(loss) + + is_quality_ok = perplexity <= max_perplexity_threshold + if is_quality_ok: + logging.info( + f"✓ Perplexity check passed: {perplexity:.2f} <= {max_perplexity_threshold}" + ) + else: + logging.warning( + f"✗ Perplexity check failed: {perplexity:.2f} > {max_perplexity_threshold}" + ) + + # Clean up immediately + del model + del outputs + gc.collect() + + return is_quality_ok + + def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only=False): command = [ "optimum-cli", @@ -51,7 +106,19 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only "--output_dir", model_dir, ] - if "coreml" in recipe: + if "xnnpack" in recipe: + command += [ + "--use_custom_sdpa", + "--use_custom_kv_cache", + ] + if quantize: + command += [ + "--qlinear", + "8da4w", + "--qembedding", + "8w", + ] + elif "coreml" in recipe: command += [ "--disable_dynamic_shapes", ] @@ -63,7 +130,9 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only "8w", ] else: - assert not quantize, "Quantization is not supported for non-CoreML recipes yet" + assert ( + not quantize + ), "Quantization is only supported for XnnPack and CoreML recipes at the moment." if not run_only: cli_export(command, model_dir) @@ -77,6 +146,14 @@ def test_text_generation(model_id, model_dir, recipe, *, quantize=True, run_only max_seq_len=64, ) print(f"\nGenerated text:\n\t{generated_text}") + generated_tokens = tokenizer(generated_text, return_tensors="pt").input_ids + + # Free memory before loading eager for quality check + del model + del tokenizer + gc.collect() + + assert check_causal_lm_output_quality(model_id, generated_tokens) is True def test_fill_mask(model_id, model_dir, recipe, *, quantize=True, run_only=False): @@ -262,7 +339,7 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False): assert torch.allclose( eager_output.logits, et_output, atol=1e-02, rtol=1e-02 - ), "CoreML output does not match eager" + ), "Model output does not match eager" if __name__ == "__main__": @@ -270,35 +347,57 @@ def test_vit(model_id, model_dir, recipe, *, quantize=False, run_only=False): parser.add_argument("--model", type=str, required=True) parser.add_argument("--recipe", type=str, required=True) parser.add_argument("--quantize", action="store_true", help="Enable quantization") + parser.add_argument( + "--model_dir", + type=str, + required=False, + help="When provided, write the pte file to this directory. Otherwise, a temporary directory is created for the test.", + ) args = parser.parse_args() - model_to_model_id_and_test_function = { - "smollm": ("HuggingFaceTB/SmolLM2-135M", test_text_generation), # works - "qwen3": ("Qwen/Qwen3-0.6B", test_text_generation), # works - "olmo": ("allenai/OLMo-1B-hf", test_text_generation), # works - "gemma3": ("unsloth/gemma-3-1b-it", test_text_generation), # does not export - "phi4": ( + _text_generation_mapping = { + "llama3.2-1b": ("NousResearch/Llama-3.2-1B", test_text_generation), + "qwen3-0.6b": ("Qwen/Qwen3-0.6B", test_text_generation), + "qwen3-1.7b": ("Qwen/Qwen3-1.7B", test_text_generation), + "gemma3-1b": ( + "unsloth/gemma-3-1b-it", + test_text_generation, + ), # does not export for CoreML + "phi4-mini": ( "microsoft/Phi-4-mini-instruct", test_text_generation, - ), # fails to lower - "llama3": ("NousResearch/Llama-3.2-1B", test_text_generation), # works - "bert": ("google-bert/bert-base-uncased", test_fill_mask), # works - "roberta": ("FacebookAI/xlmcl-roberta-base", test_fill_mask), # works - "distilbert": ("distilbert/distilbert-base-uncased", test_fill_mask), # works - "whisper": ("openai/whisper-tiny", test_whisper), # works + ), # fails to lower for CoreML + "smollm2-135m": ("HuggingFaceTB/SmolLM2-135M", test_text_generation), + "smollm3-3b": ("HuggingFaceTB/SmolLM3-3B", test_text_generation), + "olmo-1b": ("allenai/OLMo-1B-hf", test_text_generation), + } + + _mask_fill_mapping = { + "bert": ("google-bert/bert-base-uncased", test_fill_mask), + "roberta": ("FacebookAI/xlmcl-roberta-base", test_fill_mask), + "distilbert": ("distilbert/distilbert-base-uncased", test_fill_mask), + } + + _misc_model_mapping = { + "whisper": ("openai/whisper-tiny", test_whisper), "t5": ("google-t5/t5-small", test_t5), # CoreML runime failure - "vit": ("google/vit-base-patch16-224", test_vit), # works + "vit": ("google/vit-base-patch16-224", test_vit), } + + model_to_model_id_and_test_function = ( + _text_generation_mapping | _mask_fill_mapping | _misc_model_mapping + ) + if args.model not in model_to_model_id_and_test_function: raise ValueError( f"Unknown model name: {args.model}. Available models: {model_to_model_id_and_test_function.keys()}" ) + model_id, test_fn = model_to_model_id_and_test_function[args.model] with tempfile.TemporaryDirectory() as tmp_dir: - model_id, test_fn = model_to_model_id_and_test_function[args.model] test_fn( model_id=model_id, - model_dir=tmp_dir, + model_dir=tmp_dir if args.model_dir is None else args.model_dir, recipe=args.recipe, quantize=args.quantize, ) diff --git a/.ci/scripts/test_ios_ci.sh b/.ci/scripts/test_ios_ci.sh index 6908d61483c..a89c2cc5809 100755 --- a/.ci/scripts/test_ios_ci.sh +++ b/.ci/scripts/test_ios_ci.sh @@ -36,7 +36,7 @@ say() { say "Cloning the Demo App" -git clone --depth 1 https://github.com/pytorch-labs/executorch-examples.git +git clone --depth 1 https://github.com/meta-pytorch/executorch-examples.git say "Installing CoreML Backend Requirements" diff --git a/.ci/scripts/test_llama.sh b/.ci/scripts/test_llama.sh index 20fd7939a0d..84278e290f6 100644 --- a/.ci/scripts/test_llama.sh +++ b/.ci/scripts/test_llama.sh @@ -119,8 +119,12 @@ echo "COREML option ${COREML}" if [[ "${MODE}" =~ .*qnn.* ]]; then QNN=ON + + # Download QNN_SDK. If already downloaded, export environment path + source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh" + install_qnn + export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" - export QNN_SDK_ROOT=/tmp/qnn/2.28.0.241029 export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang" export PYTHONPATH=".." cp schema/program.fbs exir/_serialize/program.fbs @@ -150,6 +154,7 @@ cmake_install_executorch_libraries() { echo "Installing libexecutorch.a, libextension_module.so, libportable_ops_lib.a" rm -rf cmake-out retry cmake --preset llm \ + -DEXECUTORCH_BUILD_TESTS=ON \ -DBUILD_TESTING=OFF \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ @@ -166,6 +171,7 @@ cmake_build_llama_runner() { popd dir="examples/models/llama" retry cmake \ + -DEXECUTORCH_BUILD_TESTS=ON \ -DBUILD_TESTING=OFF \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ diff --git a/.ci/scripts/test_llama_lora.sh b/.ci/scripts/test_llama_lora.sh index 5c87cb8da72..6337bbf76a2 100644 --- a/.ci/scripts/test_llama_lora.sh +++ b/.ci/scripts/test_llama_lora.sh @@ -48,8 +48,17 @@ DOWNLOADED_PATH=$( --model_id "${HF_MODEL_REPO}" \ --files "adapter_config.json" "adapter_model.pt" "consolidated.00.pth" "params.json" "tokenizer.model" ) -EXPORTED_MODEL_NAME="llama_3_2_1B_lora.pte" -# Export model. +# Build llama runner. +cmake_install_executorch_libraries +cmake_build_llama_runner + +# Constants. +RUNTIME_ARGS="--tokenizer_path=${DOWNLOADED_PATH}/tokenizer.model --temperature=0 --seq_len=20 --warmup=1" +PROMPT="What happens if you eat watermelon seeds?" +EXPECTED_PREFIX="What happens if you eat watermelon seeds? Watermelon seeds are a good source of vitamin C," + +# Export LoRA PTE file. +MODEL_NAME="llama_3_2_1B_lora" $PYTHON_EXECUTABLE -m extension.llm.export.export_llm \ base.checkpoint="${DOWNLOADED_PATH}/consolidated.00.pth" \ base.params="${DOWNLOADED_PATH}/params.json" \ @@ -61,36 +70,64 @@ $PYTHON_EXECUTABLE -m extension.llm.export.export_llm \ model.dtype_override="fp32" \ backend.xnnpack.enabled=true \ backend.xnnpack.extended_ops=true \ - export.output_name="${EXPORTED_MODEL_NAME}" - -# Build llama runner. -cmake_install_executorch_libraries -cmake_build_llama_runner + export.output_name="${MODEL_NAME}.pte" -PROMPT="What happens if you eat watermelon seeds?" # Run llama runner -RUNTIME_ARGS="--model_path=${EXPORTED_MODEL_NAME} --tokenizer_path=${DOWNLOADED_PATH}/tokenizer.model --temperature=0 --seq_len=20 --warmup=1" - NOW=$(date +"%H:%M:%S") echo "Starting to run llama runner at ${NOW}" # shellcheck source=/dev/null -cmake-out/examples/models/llama/llama_main --prompt="${PROMPT}" ${RUNTIME_ARGS} > result.txt +cmake-out/examples/models/llama/llama_main --model_path=${MODEL_NAME}.pte --prompt="${PROMPT}" ${RUNTIME_ARGS} > result.txt NOW=$(date +"%H:%M:%S") echo "Finished at ${NOW}" RESULT=$(cat result.txt) -EXPECTED_PREFIX="What happens if you eat watermelon seeds? Watermelon seeds are a good source of vitamin C," - if [[ "${RESULT}" == "${EXPECTED_PREFIX}"* ]]; then echo "Expected result prefix: ${EXPECTED_PREFIX}" echo "Actual result: ${RESULT}" + # Do not clean up files if test passes, as they're re-used in the next test. echo "Success" - cleanup_files else echo "Expected result prefix: ${EXPECTED_PREFIX}" echo "Actual result: ${RESULT}" echo "Failure; results not the same" + cleanup_files + exit 1 +fi +# Export LoRA PTE, PTD file. +MODEL_SEPARATE="${MODEL_NAME}_separate" +$PYTHON_EXECUTABLE -m extension.llm.export.export_llm \ + base.checkpoint="${DOWNLOADED_PATH}/consolidated.00.pth" \ + base.params="${DOWNLOADED_PATH}/params.json" \ + base.adapter_checkpoint="${DOWNLOADED_PATH}/adapter_model.pt" \ + base.adapter_config="${DOWNLOADED_PATH}/adapter_config.json" \ + base.tokenizer_path="${DOWNLOADED_PATH}/tokenizer.model" \ + model.use_kv_cache=true \ + model.use_sdpa_with_kv_cache=true \ + model.dtype_override="fp32" \ + backend.xnnpack.enabled=true \ + backend.xnnpack.extended_ops=true \ + export.output_name="${MODEL_SEPARATE}.pte" \ + export.foundation_weights_file="${MODEL_SEPARATE}.ptd" + +# Run llama runner. +NOW=$(date +"%H:%M:%S") +echo "Starting to run llama runner at ${NOW}" +# shellcheck source=/dev/null +cmake-out/examples/models/llama/llama_main --model_path=${MODEL_SEPARATE}.pte --data_path=${MODEL_SEPARATE}.ptd --prompt="${PROMPT}" ${RUNTIME_ARGS} > result2.txt +NOW=$(date +"%H:%M:%S") +echo "Finished at ${NOW}" + +RESULT2=$(cat result2.txt) +if [[ "${RESULT2}" == "${EXPECTED_PREFIX}"* ]]; then + echo "Expected result prefix: ${EXPECTED_PREFIX}" + echo "Actual result: ${RESULT2}" + echo "Success" + cleanup_files +else + echo "Expected result prefix: ${EXPECTED_PREFIX}" + echo "Actual result: ${RESULT2}" + echo "Failure; results not the same" cleanup_files exit 1 fi diff --git a/.ci/scripts/test_llama_torchao_lowbit.sh b/.ci/scripts/test_llama_torchao_lowbit.sh index ae8f74a5df5..5f472fad63b 100644 --- a/.ci/scripts/test_llama_torchao_lowbit.sh +++ b/.ci/scripts/test_llama_torchao_lowbit.sh @@ -29,27 +29,22 @@ cmake -DPYTHON_EXECUTABLE=python \ -DEXECUTORCH_ENABLE_LOGGING=1 \ -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ - -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ - -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ - -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ -DEXECUTORCH_BUILD_XNNPACK=OFF \ -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_TORCHAO=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out . -cmake --build cmake-out -j16 --target install --config Release +cmake --build cmake-out -j16 --config Release --target install # Install llama runner with torchao cmake -DPYTHON_EXECUTABLE=python \ - -DBUILD_TESTING=OFF \ -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ - -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ - -DEXECUTORCH_BUILD_XNNPACK=OFF \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ - -DEXECUTORCH_BUILD_TORCHAO=ON \ -Bcmake-out/examples/models/llama \ examples/models/llama cmake --build cmake-out/examples/models/llama -j16 --config Release diff --git a/.ci/scripts/test_model.ps1 b/.ci/scripts/test_model.ps1 new file mode 100644 index 00000000000..19d70c591ac --- /dev/null +++ b/.ci/scripts/test_model.ps1 @@ -0,0 +1,89 @@ +param ( + [string]$modelName, + [string]$backend, + [string]$buildDir = "cmake-out", + [bool]$strict = $false +) + +Set-PSDebug -Trace 1 +$ErrorActionPreference = 'Stop' +$PSNativeCommandUseErrorActionPreference = $true + +function ExportModel-Portable { + param ( + [string]$model_name, + [bool]$strict + ) + + $exportParams = "--model_name", "$modelName" + if ($strict) { + $exportParams += "--strict" + } + python -m examples.portable.scripts.export @exportParams | Write-Host + if ($LASTEXITCODE -ne 0) { + Write-Host "Model export failed. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE + } + + "$modelName.pte" +} + +function ExportModel-Xnnpack { + param ( + [string]$model_name, + [bool]$quantize + ) + + if ($quantize) { + python -m examples.xnnpack.aot_compiler --model_name="${MODEL_NAME}" --delegate --quantize | Write-Host + $modelFile = "$($modelName)_xnnpack_q8.pte" + } else { + python -m examples.xnnpack.aot_compiler --model_name="${MODEL_NAME}" --delegate | Write-Host + $modelFile = "$($modelName)_xnnpack_fp32.pte" + } + if ($LASTEXITCODE -ne 0) { + Write-Host "Model export failed. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE + } + + $modelFile +} + +# Build the runner +if (Test-Path -Path $buildDir) { + Remove-Item -Path $buildDir -Recurse -Force +} +New-Item -Path $buildDir -ItemType Directory +Push-Location $buildDir +cmake .. --preset windows +cmake --build . -t executor_runner -j16 --config Release +if ($LASTEXITCODE -ne 0) { + Write-Host "Runner build failed. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} +$executorBinaryPath = Join-Path -Path $buildDir -ChildPath "Release\executor_runner.exe" +Pop-Location + +# Export the model +switch ($backend) { + "portable" { + $model_path = ExportModel-Portable -model_name $modelName -strict $strict + } + "xnnpack-f32" { + $model_path = ExportModel-Xnnpack -model_name $modelName -quantize $false + } + "xnnpack-q8" { + $model_path = ExportModel-Xnnpack -model_name $modelName -quantize $true + } + default { + Write-Host "Unknown backend $backend." + exit 1 + } +} + +# Run the runner +& "$executorBinaryPath" --model_path="$model_path" +if ($LASTEXITCODE -ne 0) { + Write-Host "Model execution failed. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} diff --git a/.ci/scripts/test_model.sh b/.ci/scripts/test_model.sh index 5d9f694b0b6..74eb75c6ddd 100755 --- a/.ci/scripts/test_model.sh +++ b/.ci/scripts/test_model.sh @@ -97,7 +97,7 @@ test_model() { bash examples/models/llava/install_requirements.sh STRICT="--no-strict" fi - if [[ "${MODEL_NAME}" == "qwen2_5" ]]; then + if [[ "${MODEL_NAME}" == "qwen2_5_1_5b" ]]; then # Install requirements for export_llama bash examples/models/llama/install_requirements.sh # Test export_llm script: python3 -m extension.llm.export.export_llm. @@ -199,6 +199,8 @@ test_model_with_qnn() { EXPORT_SCRIPT=albert elif [[ "${MODEL_NAME}" == "bert" ]]; then EXPORT_SCRIPT=bert + elif [[ "${MODEL_NAME}" == "conv_former" ]]; then + EXPORT_SCRIPT=conv_former elif [[ "${MODEL_NAME}" == "cvt" ]]; then EXPORT_SCRIPT=cvt elif [[ "${MODEL_NAME}" == "distilbert" ]]; then @@ -238,7 +240,7 @@ test_model_with_qnn() { "cvt"|"dit"|"focalnet"|"mobilevit_v2"|"pvt"|"swin") SCRIPT_FOLDER=oss_scripts ;; - "albert"|"bert"|"distilbert"|"roberta"|"efficientnet"|"mobilevit_v1") + "albert"|"bert"|"conv_former"|"distilbert"|"roberta"|"efficientnet"|"mobilevit_v1") pip install evaluate SCRIPT_FOLDER=oss_scripts # 16bit models will encounter op validation fail on some operations, diff --git a/.ci/scripts/test_qnn_static_llama.sh b/.ci/scripts/test_qnn_static_llama.sh index a5f194ba0b9..7898d03b3b9 100644 --- a/.ci/scripts/test_qnn_static_llama.sh +++ b/.ci/scripts/test_qnn_static_llama.sh @@ -9,8 +9,11 @@ set -euxo pipefail source "$(dirname "${BASH_SOURCE[0]}")/utils.sh" +# Download QNN_SDK. If already downloaded, export environment path +source "$(dirname "${BASH_SOURCE[0]}")/../../backends/qualcomm/scripts/install_qnn_sdk.sh" +install_qnn + export EXECUTORCH_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" -export QNN_SDK_ROOT=/tmp/qnn/2.28.0.241029 export LD_LIBRARY_PATH="${QNN_SDK_ROOT}/lib/x86_64-linux-clang" export PYTHONPATH=".." cp schema/program.fbs exir/_serialize/program.fbs @@ -33,12 +36,12 @@ echo "Creating tokenizer.bin" $PYTHON_EXECUTABLE -m pytorch_tokenizers.tools.llama2c.convert -t tokenizer.model -o tokenizer.bin set +e -# Compile only as weight sharing is not applicable on x86 -$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-android/ --executorch_root . --artifact_dir . --llama_artifacts . --compile_only +# Compile only as weight sharing is not applicable on x86. +$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-android/ --executorch_root . --artifact_dir ./stories_110m_pte_size --llama_artifacts . --compile_only exit_code1=$? # Checks accuracy with weight sharing disabled since x86 does not support weight sharing. -$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-x86/ --executorch_root . --artifact_dir . --llama_artifacts . --enable_x86_64 +$PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_110m --model SM8650 --build_folder build-x86/ --executorch_root . --artifact_dir ./stories_110m_accuracy --llama_artifacts . --enable_x86_64 exit_code2=$? # Check BC diff --git a/.ci/scripts/test_torchao_huggingface_checkpoints.sh b/.ci/scripts/test_torchao_huggingface_checkpoints.sh new file mode 100644 index 00000000000..c0910b47826 --- /dev/null +++ b/.ci/scripts/test_torchao_huggingface_checkpoints.sh @@ -0,0 +1,139 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ------------------------- +# Args / flags +# ------------------------- +TEST_WITH_RUNNER=0 +MODEL_NAME="" + +# Parse args +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [--test_with_runner]" + echo "Supported model_name values: qwen3_4b, phi_4_mini" + exit 1 +fi + +MODEL_NAME="$1" +shift + +while [[ $# -gt 0 ]]; do + case "$1" in + --test_with_runner) + TEST_WITH_RUNNER=1 + ;; + -h|--help) + echo "Usage: $0 [--test_with_runner]" + echo " model_name: qwen3_4b | phi_4_mini" + echo " --test_with_runner: build ET + run llama_main to sanity-check the export" + exit 0 + ;; + *) + echo "Unknown option: $1" + exit 1 + ;; + esac + shift +done + +if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then + PYTHON_EXECUTABLE=python3 +fi + +MODEL_OUT=model.pte + +case "$MODEL_NAME" in + qwen3_4b) + echo "Running Qwen3-4B export..." + HF_MODEL_DIR=$(hf download pytorch/Qwen3-4B-INT8-INT4) + EXPECTED_MODEL_SIZE_UPPER_BOUND=$((3 * 1024 * 1024 * 1024)) # 3GB + $PYTHON_EXECUTABLE -m executorch.examples.models.qwen3.convert_weights \ + $HF_MODEL_DIR \ + pytorch_model_converted.bin + + $PYTHON_EXECUTABLE -m executorch.examples.models.llama.export_llama \ + --model "qwen3_4b" \ + --checkpoint pytorch_model_converted.bin \ + --params examples/models/qwen3/config/4b_config.json \ + --output_name $MODEL_OUT \ + -kv \ + --use_sdpa_with_kv_cache \ + -X \ + --xnnpack-extended-ops \ + --max_context_length 1024 \ + --max_seq_length 1024 \ + --dtype fp32 \ + --metadata '{"get_bos_id":199999, "get_eos_ids":[200020,199999]}' + ;; + + phi_4_mini) + echo "Running Phi-4-mini export..." + HF_MODEL_DIR=$(hf download pytorch/Phi-4-mini-instruct-INT8-INT4) + EXPECTED_MODEL_SIZE_UPPER_BOUND=$((3 * 1024 * 1024 * 1024)) # 3GB + $PYTHON_EXECUTABLE -m executorch.examples.models.phi_4_mini.convert_weights \ + $HF_MODEL_DIR \ + pytorch_model_converted.bin + + $PYTHON_EXECUTABLE -m executorch.examples.models.llama.export_llama \ + --model "phi_4_mini" \ + --checkpoint pytorch_model_converted.bin \ + --params examples/models/phi_4_mini/config/config.json \ + --output_name $MODEL_OUT \ + -kv \ + --use_sdpa_with_kv_cache \ + -X \ + --xnnpack-extended-ops \ + --max_context_length 1024 \ + --max_seq_length 1024 \ + --dtype fp32 \ + --metadata '{"get_bos_id":199999, "get_eos_ids":[200020,199999]}' + ;; + + *) + echo "Error: unsupported model_name '$MODEL_NAME'" + echo "Supported values: qwen3_4b, phi_4_mini" + exit 1 + ;; +esac + +# Check file size +MODEL_SIZE=$(stat --printf="%s" $MODEL_OUT 2>/dev/null || stat -f%z $MODEL_OUT) +if [[ $MODEL_SIZE -gt $EXPECTED_MODEL_SIZE_UPPER_BOUND ]]; then + echo "Error: model size $MODEL_SIZE is greater than expected upper bound $EXPECTED_MODEL_SIZE_UPPER_BOUND" + exit 1 +fi + +# Install ET with CMake +if [[ "$TEST_WITH_RUNNER" -eq 1 ]]; then + echo "[runner] Building and testing llama_main ..." + cmake -DPYTHON_EXECUTABLE=python \ + -DCMAKE_INSTALL_PREFIX=cmake-out \ + -DEXECUTORCH_ENABLE_LOGGING=1 \ + -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_XNNPACK=ON \ + -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ + -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ + -Bcmake-out . + cmake --build cmake-out -j16 --config Release --target install + + + # Install llama runner + 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 + + # Run the model + ./cmake-out/examples/models/llama/llama_main --model_path=$MODEL_OUT --tokenizer_path="${HF_MODEL_DIR}/tokenizer.json" --prompt="Once upon a time," +fi + +# Clean up +rm -f pytorch_model_converted.bin "$MODEL_OUT" diff --git a/.ci/scripts/test_wheel_package_qnn.sh b/.ci/scripts/test_wheel_package_qnn.sh new file mode 100644 index 00000000000..39c52a4a396 --- /dev/null +++ b/.ci/scripts/test_wheel_package_qnn.sh @@ -0,0 +1,201 @@ +#!/bin/bash +# === CI Wheel Build & Test Script === + +# Exit immediately on error, print each command, and capture all output to build.log +set -e +set -x +exec > >(tee -i build.log) 2>&1 + +# Save repo root +REPO_ROOT=$(pwd) + +# ---------------------------- +# Dynamically create script_qnn_wheel_test.py +# ---------------------------- +cat > "/tmp/script_qnn_wheel_test.py" << 'EOF' +# pyre-ignore-all-errors +import argparse + +import torch +from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_soc_to_chipset_map, + to_edge_transform_and_lower_to_qnn, +) +from executorch.exir.backend.utils import format_delegated_graph +from executorch.examples.models.model_factory import EagerModelFactory +from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.extension.export_util.utils import save_pte_program +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e, prepare_qat_pt2e + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--output_folder", type=str, default="", help="The folder to store the exported program") + parser.add_argument("--soc", type=str, default="SM8650", help="Specify the SoC model.") + parser.add_argument("-q", "--quantization", choices=["ptq", "qat"], help="Run post-traininig quantization.") + args = parser.parse_args() + + class LinearModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(3, 3) + def forward(self, arg): + return self.linear(arg) + def get_example_inputs(self): + return (torch.randn(3, 3),) + + model = LinearModule() + example_inputs = model.get_example_inputs() + + if args.quantization: + quantizer = QnnQuantizer() + m = torch.export.export(model.eval(), example_inputs, strict=True).module() + if args.quantization == "qat": + m = prepare_qat_pt2e(m, quantizer) + m(*example_inputs) + elif args.quantization == "ptq": + m = prepare_pt2e(m, quantizer) + m(*example_inputs) + m = convert_pt2e(m) + else: + m = model + + use_fp16 = True if args.quantization is None else False + backend_options = generate_htp_compiler_spec(use_fp16=use_fp16) + compile_spec = generate_qnn_executorch_compiler_spec( + soc_model=get_soc_to_chipset_map()[args.soc], + backend_options=backend_options, + ) + delegated_program = to_edge_transform_and_lower_to_qnn(m, example_inputs, compile_spec) + output_graph = format_delegated_graph(delegated_program.exported_program().graph_module) + # Ensure QnnBackend is in the output graph + assert "QnnBackend" in output_graph + executorch_program = delegated_program.to_executorch( + config=ExecutorchBackendConfig(extract_delegate_segments=False) + ) + save_pte_program(executorch_program, "linear", args.output_folder) + +if __name__ == "__main__": + main() +EOF + +# ---------------------------- +# Wheel build and .so checks +# ---------------------------- +echo "=== Building Wheel Package ===" +source .ci/scripts/utils.sh +install_executorch +EXECUTORCH_BUILDING_WHEEL=1 python setup.py bdist_wheel +unset EXECUTORCH_BUILDING_WHEEL + +WHEEL_FILE=$(ls dist/*.whl | head -n 1) +echo "Found wheel: $WHEEL_FILE" + +PYTHON_VERSION=$1 +# ---------------------------- +# Check wheel does NOT contain qualcomm/sdk +# ---------------------------- +echo "Checking wheel does not contain qualcomm/sdk..." +SDK_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep "executorch/backends/qualcomm/sdk" || true) +if [ -n "$SDK_FILES" ]; then + echo "ERROR: Wheel package contains unexpected qualcomm/sdk files:" + echo "$SDK_FILES" + exit 1 +else + echo "OK: No qualcomm/sdk files found in wheel" +fi + +# ---------------------------- +# Check .so files in the wheel +# ---------------------------- +echo "Checking for .so files inside the wheel..." +WHEEL_SO_FILES=$(unzip -l "$WHEEL_FILE" | awk '{print $4}' | grep "executorch/backends/qualcomm/python" || true) +if [ -z "$WHEEL_SO_FILES" ]; then + echo "ERROR: No .so files found in wheel under executorch/backends/qualcomm/python" + exit 1 +else + echo "Wheel contains the following .so files:" + echo "$WHEEL_SO_FILES" +fi + +# ---------------------------- +# Helpers +# ---------------------------- +get_site_packages_dir () { + local PYBIN="$1" + "$PYBIN" - <<'PY' +import sysconfig, sys +print(sysconfig.get_paths().get("purelib") or sysconfig.get_paths().get("platlib")) +PY +} + +run_core_tests () { + local PYBIN="$1" # path to python + local PIPBIN="$2" # path to pip + local LABEL="$3" # label to print (conda/venv) + + echo "=== [$LABEL] Installing wheel & deps ===" + "$PIPBIN" install --upgrade pip + "$PIPBIN" install "$WHEEL_FILE" + "$PIPBIN" install torch=="2.9.0.dev20250906" --index-url "https://download.pytorch.org/whl/nightly/cpu" + "$PIPBIN" install --pre torchao --index-url "https://download.pytorch.org/whl/nightly/cpu" + + echo "=== [$LABEL] Import smoke tests ===" + "$PYBIN" -c "import executorch; print('executorch imported successfully')" + "$PYBIN" -c "import executorch.backends.qualcomm; print('executorch.backends.qualcomm imported successfully')" + + echo "=== [$LABEL] List installed executorch/backends/qualcomm/python ===" + local SITE_DIR + SITE_DIR="$(get_site_packages_dir "$PYBIN")" + local SO_DIR="$SITE_DIR/executorch/backends/qualcomm/python" + ls -l "$SO_DIR" || echo "Folder does not exist!" + + echo "=== [$LABEL] Run export script to generate linear.pte ===" + (cd "$REPO_ROOT" && "$PYBIN" "/tmp/script_qnn_wheel_test.py") + + if [ -f "$REPO_ROOT/linear.pte" ]; then + echo "[$LABEL] Model file linear.pte successfully created" + else + echo "ERROR: [$LABEL] Model file linear.pte was not created" + exit 1 + fi +} + +# ---------------------------- +# Conda environment setup & tests +# ---------------------------- +echo "=== Testing in Conda env ===" +TEMP_ENV_DIR=$(mktemp -d) +echo "Using temporary directory for conda: $TEMP_ENV_DIR" +conda create -y -p "$TEMP_ENV_DIR/env" python=$PYTHON_VERSION +# derive python/pip paths inside the conda env +CONDA_PY="$TEMP_ENV_DIR/env/bin/python" +CONDA_PIP="$TEMP_ENV_DIR/env/bin/pip" +# Some images require conda run; keep pip/python direct to simplify path math +run_core_tests "$CONDA_PY" "$CONDA_PIP" "conda" + +# Cleanup conda env +conda env remove -p "$TEMP_ENV_DIR/env" -y || true +rm -rf "$TEMP_ENV_DIR" + +# ---------------------------- +# Python venv setup & tests +# ---------------------------- +echo "=== Testing in Python venv ===" +TEMP_VENV_DIR=$(mktemp -d) +echo "Using temporary directory for venv: $TEMP_VENV_DIR" +python3 -m venv "$TEMP_VENV_DIR/venv" +VENV_PY="$TEMP_VENV_DIR/venv/bin/python" +VENV_PIP="$TEMP_VENV_DIR/venv/bin/pip" + +# Ensure venv has wheel/build basics if needed +"$VENV_PIP" install --upgrade pip + +run_core_tests "$VENV_PY" "$VENV_PIP" "venv" + +# Cleanup venv +rm -rf "$TEMP_VENV_DIR" + +echo "=== All tests completed! ===" diff --git a/.ci/scripts/unittest-buck2.sh b/.ci/scripts/unittest-buck2.sh index f56db8924be..f748be62ac1 100755 --- a/.ci/scripts/unittest-buck2.sh +++ b/.ci/scripts/unittest-buck2.sh @@ -11,9 +11,10 @@ set -eux # TODO: can't query //kernels/prim_ops because of non-buckified stuff in OSS. buck2 query "//backends/apple/... + //backends/example/... + \ //backends/mediatek/... + //backends/transforms/... + \ -//backends/xnnpack/... + //configurations/... + //kernels/aten/... + \ -//kernels/optimized/... + //kernels/portable/... + //kernels/quantized/... + \ -//kernels/test/... + //runtime/... + //schema/... + //test/... + //util/..." +//backends/xnnpack/... + //configurations/... + //extension/flat_tensor: + \ +//extension/llm/runner: + //kernels/aten/... + //kernels/optimized/... + \ +//kernels/portable/... + //kernels/quantized/... + //kernels/test/... + \ +//runtime/... + //schema/... + //test/... + //util/..." # TODO: optimized ops are unbuildable because they now use ATen; put # them back after we can use PyTorch in OSS buck. diff --git a/.ci/scripts/unittest-windows.ps1 b/.ci/scripts/unittest-windows.ps1 new file mode 100644 index 00000000000..6f1365bc3fc --- /dev/null +++ b/.ci/scripts/unittest-windows.ps1 @@ -0,0 +1,38 @@ +param ( + [string]$buildMode = "Release" +) + +Set-PSDebug -Trace 1 +$ErrorActionPreference = 'Stop' +$PSNativeCommandUseErrorActionPreference = $true + +# Run native unit tests (via ctest) +New-Item -Path "test-build" -ItemType Directory +cd "test-build" + +cmake .. --preset windows -B . -DEXECUTORCH_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=$buildMode +if ($LASTEXITCODE -ne 0) { + Write-Host "CMake configuration was unsuccessful. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} + +cmake --build . -j8 --config $buildMode --verbose +if ($LASTEXITCODE -ne 0) { + Write-Host "CMake build was unsuccessful. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} + +ctest -j8 . --build-config $buildMode --output-on-failure -E "method_test|tensor_parser_test" +if ($LASTEXITCODE -ne 0) { + Write-Host "CTest run was unsuccessful. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} + +cd .. + +# Run pytest +pytest -v -c pytest-windows.ini +if ($LASTEXITCODE -ne 0) { + Write-Host "Pytest invocation was unsuccessful. Exit code: $LASTEXITCODE." + exit $LASTEXITCODE +} \ No newline at end of file diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh index 6902cc3dec1..f6f6ece786b 100644 --- a/.ci/scripts/utils.sh +++ b/.ci/scripts/utils.sh @@ -131,8 +131,6 @@ build_executorch_runner_cmake() { else CXXFLAGS="" fi - # This command uses buck2 to gather source files and buck2 could crash flakily - # on MacOS CXXFLAGS="$CXXFLAGS" retry cmake -DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}" -DCMAKE_BUILD_TYPE="${1:-Release}" .. popd || return diff --git a/.ci/scripts/wheel/pre_build_script.sh b/.ci/scripts/wheel/pre_build_script.sh index 424529af864..f5d8cd362ee 100755 --- a/.ci/scripts/wheel/pre_build_script.sh +++ b/.ci/scripts/wheel/pre_build_script.sh @@ -9,9 +9,38 @@ set -euxo pipefail # This script is run before building ExecuTorch binaries +if [[ "$(uname -m)" == "aarch64" ]]; then + # On some Linux aarch64 systems, the "atomic" library is not found during linking. + # To work around this, replace "atomic" with the literal ${ATOMIC_LIB} so the + # build system uses the full path to the atomic library. + file="extension/llm/tokenizers/third-party/sentencepiece/src/CMakeLists.txt" + sed 's/list(APPEND SPM_LIBS "atomic")/list(APPEND SPM_LIBS ${ATOMIC_LIB})/' \ + "$file" > "${file}.tmp" && mv "${file}.tmp" "$file" + + grep -n 'list(APPEND SPM_LIBS ${ATOMIC_LIB})' "$file" && \ + echo "the file $file has been modified for atomic to use full path" +fi + +# Clone nested submodules for tokenizers - this is a workaround for recursive +# submodule clone failing due to path length limitations on Windows. Eventually, +# we should update the core job in test-infra to enable long paths before +# checkout to avoid needing to do this. +pushd extension/llm/tokenizers +git submodule update --init +popd + +# On Windows, enable symlinks and re-checkout the current revision to create +# the symlinked src/ directory. This is needed to build the wheel. +UNAME_S=$(uname -s) +if [[ $UNAME_S == *"MINGW"* || $UNAME_S == *"MSYS"* ]]; then + echo "Enabling symlinks on Windows" + git config core.symlinks true + git checkout -f HEAD +fi + # Manually install build requirements because `python setup.py bdist_wheel` does # not install them. TODO(dbort): Switch to using `python -m build --wheel`, # which does install them. Though we'd need to disable build isolation to be # able to see the installed torch package. -"${GITHUB_WORKSPACE}/${REPOSITORY}/install_requirements.sh" --example +"${GITHUB_WORKSPACE}/${REPOSITORY}/install_requirements.sh" --example diff --git a/.ci/scripts/wheel/test_linux_aarch64.py b/.ci/scripts/wheel/test_linux_aarch64.py new file mode 100644 index 00000000000..d514738c005 --- /dev/null +++ b/.ci/scripts/wheel/test_linux_aarch64.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python +# 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 test_base +from examples.models import Backend, Model + +if __name__ == "__main__": + # coremltools does not support linux aarch64 yet and install from the source fails on runtime + # https://github.com/apple/coremltools/issues/1254 + # https://github.com/apple/coremltools/issues/2195 + test_base.run_tests( + model_tests=[ + test_base.ModelTest( + model=Model.Mv3, + backend=Backend.XnnpackQuantizationDelegation, + ), + ] + ) diff --git a/.ci/scripts/wheel/test_windows.py b/.ci/scripts/wheel/test_windows.py new file mode 100644 index 00000000000..6ae953b5df5 --- /dev/null +++ b/.ci/scripts/wheel/test_windows.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +# 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 typing import List + +import torch +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL +from executorch.examples.models.model_factory import EagerModelFactory +from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS +from executorch.examples.xnnpack.quantization.utils import quantize as quantize_xnn +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) +from test_base import ModelTest + + +def test_model_xnnpack(model: Model, quantize: bool) -> None: + model_instance, example_inputs, _, _ = EagerModelFactory.create_model( + *MODEL_NAME_TO_MODEL[str(model)] + ) + + model_instance.eval() + ref_outputs = model_instance(*example_inputs) + + if quantize: + quant_type = MODEL_NAME_TO_OPTIONS[str(model)].quantization + model_instance = torch.export.export_for_training( + model_instance, example_inputs + ) + model_instance = quantize_xnn( + model_instance.module(), example_inputs, quant_type + ) + + lowered = to_edge_transform_and_lower( + torch.export.export(model_instance, example_inputs), + partitioner=[XnnpackPartitioner()], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + ), + ).to_executorch() + + loaded_model = _load_for_executorch_from_buffer(lowered.buffer) + et_outputs = loaded_model([*example_inputs]) + + if isinstance(ref_outputs, torch.Tensor): + ref_outputs = (ref_outputs,) + + assert len(ref_outputs) == len(et_outputs) + for i in range(len(ref_outputs)): + torch.testing.assert_close(ref_outputs[i], et_outputs[i], atol=1e-4, rtol=1e-5) + + +def run_tests(model_tests: List[ModelTest]) -> None: + for model_test in model_tests: + if model_test.backend == Backend.Xnnpack: + test_model_xnnpack(model_test.model, quantize=False) + else: + raise RuntimeError(f"Unsupported backend {model_test.backend}.") + + +if __name__ == "__main__": + run_tests( + model_tests=[ + ModelTest( + model=Model.Mv3, + backend=Backend.Xnnpack, + ), + ] + ) diff --git a/.ci/scripts/wheel/vc_env_helper.bat b/.ci/scripts/wheel/vc_env_helper.bat new file mode 100644 index 00000000000..d7fa2c1a596 --- /dev/null +++ b/.ci/scripts/wheel/vc_env_helper.bat @@ -0,0 +1,61 @@ +REM This is lightly modified from the torchvision Windows build logic. +REM See https://github.com/pytorch/vision/blob/main/packaging/windows/internal/vc_env_helper.bat + +@echo on + +set VC_VERSION_LOWER=17 +set VC_VERSION_UPPER=18 + +for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -legacy -products * -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do ( + if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" ( + set "VS15INSTALLDIR=%%i" + set "VS15VCVARSALL=%%i\VC\Auxiliary\Build\vcvarsall.bat" + goto vswhere + ) +) + +:vswhere +if "%VSDEVCMD_ARGS%" == "" ( + call "%VS15VCVARSALL%" x64 || exit /b 1 +) else ( + call "%VS15VCVARSALL%" x64 %VSDEVCMD_ARGS% || exit /b 1 +) + +@echo on + +if "%CU_VERSION%" == "xpu" call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" + +set DISTUTILS_USE_SDK=1 + +set args=%1 +shift +:start +if [%1] == [] goto done +set args=%args% %1 +shift +goto start + +:done +if "%args%" == "" ( + echo Usage: vc_env_helper.bat [command] [args] + echo e.g. vc_env_helper.bat cl /c test.cpp +) + +set work_dir=%CD% +if exist setup.py ( + echo "Creating symlink..." + REM Setup a symlink to shorten the path length. + REM Note that the ET directory has to be named "executorch". + cd %GITHUB_WORKSPACE% + if not exist et\ ( + mkdir et + ) + cd et + echo Work dir: %work_dir% + if not exist executorch\ ( + mklink /d executorch %work_dir% + ) + cd executorch +) + +%args% || exit /b 1 diff --git a/.ci/scripts/zephyr-utils.sh b/.ci/scripts/zephyr-utils.sh index 2b36c6b0427..28dca2c1dfb 100644 --- a/.ci/scripts/zephyr-utils.sh +++ b/.ci/scripts/zephyr-utils.sh @@ -6,9 +6,9 @@ # LICENSE file in the root directory of this source tree. download_arm_zephyr_sdk () { - wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.16.0/zephyr-sdk-0.16.0_linux-x86_64.tar.xz - tar -xf zephyr-sdk-0.16.0_linux-x86_64.tar.xz - rm -f zephyr-sdk-0.16.0_linux-x86_64.tar.xz + wget https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/zephyr-sdk-0.17.2_linux-x86_64.tar.xz + tar -xf zephyr-sdk-0.17.2_linux-x86_64.tar.xz + rm -f zephyr-sdk-0.17.2_linux-x86_64.tar.xz } setup_zephyr_et_module () { diff --git a/.github/workflows/_unittest.yml b/.github/workflows/_unittest.yml index 63f5e6693b7..587f2cf5e5a 100644 --- a/.github/workflows/_unittest.yml +++ b/.github/workflows/_unittest.yml @@ -19,6 +19,7 @@ on: required: false type: string description: Install ExecuTorch in editable mode or not. + default: 'false' python-version: required: false type: string @@ -52,3 +53,31 @@ jobs: # This is needed to get the prebuilt PyTorch wheel from S3 ${CONDA_RUN} --no-capture-output pip install awscli==1.37.21 .ci/scripts/unittest-macos.sh --build-tool "${{ inputs.build-tool }}" --build-mode "${{ inputs.build-mode }}" --editable "${{ inputs.editable }}" + + windows: + if: ${{ inputs.build-tool == 'cmake' }} + uses: pytorch/test-infra/.github/workflows/windows_job.yml@main + with: + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 120 + script: | + conda init powershell + + powershell -Command "& { + Set-PSDebug -Trace 1 + \$ErrorActionPreference = 'Stop' + \$PSNativeCommandUseErrorActionPreference = \$true + + .ci/scripts/setup-windows.ps1 -editable "${{ inputs.editable }}" + if (\$LASTEXITCODE -ne 0) { + Write-Host "Setup failed. Exit code: \$LASTEXITCODE." + exit \$LASTEXITCODE + } + + .ci/scripts/unittest-windows.ps1 -buildMode "${{ inputs.build-mode }}" + if (\$LASTEXITCODE -ne 0) { + Write-Host "Unit tests failed. Exit code: \$LASTEXITCODE." + exit \$LASTEXITCODE + } + }" diff --git a/.github/workflows/add-unanswered-to-project.yml b/.github/workflows/add-unanswered-to-project.yml new file mode 100644 index 00000000000..ba2bc6c8436 --- /dev/null +++ b/.github/workflows/add-unanswered-to-project.yml @@ -0,0 +1,93 @@ +name: Add Open External Contributor PRs and Issues to PyTorch Org Project 136 + +on: + workflow_dispatch: + 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) + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const projectId = "PVT_kwDOAUB9vs4A_PUL"; // PyTorch org project 136 + const owner = 'pytorch'; + const repo = 'executorch'; + + // List of authors to exclude + const excludedAuthors = new Set([ + "nil-is-all", "cbilgin", "KimishPatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", "shoumikhin", + "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka", + "Conarnar", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", "jackzhxng", + "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong", + "zhenyan-zhang-meta", "silverguo", "dbort", "jorgep31415", "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", "pytorchbot", + "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "Erik-Lundell", "zingo", "AdrianLundell", + "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", "mansnils", "martinlsm", "freddan80", + "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", "haowhsu-quic", "shewu-quic", + "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "cymbalrush", "DenisVieriu97", "billmguo", + "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "neuropilot-captain" + ]); + + async function addItem(contentId, type, number) { + try { + await github.graphql(` + mutation { + addProjectV2ItemById(input: {projectId: "${projectId}", contentId: "${contentId}"}) { + item { id } + } + } + `); + console.log(`Added ${type} #${number} to project`); + } catch (error) { + if (error.message && error.message.includes("A project item already exists for this content")) { + // Ignore if already exists + console.log(`${type} #${number} already in project`); + } else { + console.log(`Error adding ${type} #${number}: ${error.message}`); + } + } + } + + try { + // Add open issues (not PRs) and exclude by author + const issues = await github.paginate( + github.rest.issues.listForRepo, + { + owner, + repo, + state: 'open', + filter: 'all' + } + ); + for (const issue of issues) { + if (!issue.pull_request && !excludedAuthors.has(issue.user.login)) { + await addItem(issue.node_id, 'issue', issue.number); + } + } + + // Add open, non-draft PRs (regardless of review state), exclude by author + const prs = await github.paginate( + github.rest.pulls.list, + { + owner, + repo, + state: 'open', + draft: false, + } + ); + for (const pr of prs) { + if (!excludedAuthors.has(pr.user.login)) { + await addItem(pr.node_id, 'pr', pr.number); + } + } + } catch (error) { + core.setFailed(`Workflow failed: ${error.message}`); + } diff --git a/.github/workflows/android-perf.yml b/.github/workflows/android-perf.yml index 6f0e388fefd..33937531a01 100644 --- a/.github/workflows/android-perf.yml +++ b/.github/workflows/android-perf.yml @@ -72,7 +72,7 @@ jobs: # Separate default values from the workflow dispatch. To ensure defaults are accessible # during scheduled runs and to provide flexibility for different defaults between # on-demand and periodic benchmarking. - CRON_DEFAULT_MODELS: ${{ github.event_name == 'schedule' && 'mv3,mv2,ic4,ic3,resnet50,edsr,mobilebert,w2l,meta-llama/Llama-3.2-1B,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8,Qwen/Qwen3-0.6B,HuggingFaceTB/SmolLM2-135M,allenai/OLMo-1B-hf,google/gemma-3-1b-it' || 'Qwen/Qwen3-0.6B' }} + CRON_DEFAULT_MODELS: ${{ github.event_name == 'schedule' && 'mv3,mv2,ic4,ic3,resnet50,mobilebert,w2l,meta-llama/Llama-3.2-1B,meta-llama/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8,meta-llama/Llama-3.2-1B-Instruct-QLORA_INT4_EO8,Qwen/Qwen3-0.6B,HuggingFaceTB/SmolLM2-135M,allenai/OLMo-1B-hf,google/gemma-3-1b-it' || 'Qwen/Qwen3-0.6B' }} CRON_DEFAULT_DEVICES: samsung_galaxy_s22+public run: | set -eux @@ -292,7 +292,7 @@ jobs: export.output_name="${OUT_ET_MODEL_NAME}.pte" ls -lh "${OUT_ET_MODEL_NAME}.pte" elif [[ ${{ matrix.config }} == "llama3_qnn_htp" ]]; then - export QNN_SDK_ROOT=/tmp/qnn/2.28.0.241029 + export QNN_SDK_ROOT=/tmp/qnn/2.37.0.250724 export LD_LIBRARY_PATH=$QNN_SDK_ROOT/lib/x86_64-linux-clang/ export PYTHONPATH=$(pwd)/.. @@ -432,7 +432,7 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh mkdir -p aar-out - PYTHON_EXECUTABLE=python ANDROID_ABIS="arm64-v8a" BUILD_AAR_DIR=aar-out EXECUTORCH_BUILD_QNN=ON QNN_SDK_ROOT=/tmp/qnn/2.28.0.241029 EXECUTORCH_ANDROID_PROFILING=ON bash scripts/build_android_library.sh + PYTHON_EXECUTABLE=python ANDROID_ABIS="arm64-v8a" BUILD_AAR_DIR=aar-out EXECUTORCH_BUILD_QNN=ON QNN_SDK_ROOT=/tmp/qnn/2.37.0.250724 EXECUTORCH_ANDROID_PROFILING=ON bash scripts/build_android_library.sh mkdir -p extension/benchmark/android/benchmark/app/libs cp aar-out/executorch.aar extension/benchmark/android/benchmark/app/libs pushd extension/benchmark/android/benchmark diff --git a/.github/workflows/android-release-artifacts.yml b/.github/workflows/android-release-artifacts.yml index 9ef8d046b8b..f0b74342eb8 100644 --- a/.github/workflows/android-release-artifacts.yml +++ b/.github/workflows/android-release-artifacts.yml @@ -16,6 +16,7 @@ on: options: - "xnnpack" - "vulkan+xnnpack" + - "qnn" schedule: - cron: 0 10 * * * @@ -83,17 +84,29 @@ jobs: echo -n "$SECRET_EXECUTORCH_MAVEN_SIGNING_GPG_KEY_CONTENTS" | base64 -d > /tmp/secring.gpg + GRADLE_ARGS="" + # Update the version name in build.gradle in case of maven publish VERSION="${{ inputs.version }}" if [ ! -z "$VERSION" ]; then - sed -i "s/\(coordinates(\"org.pytorch\", \"executorch-android\", \"\)\([0-9]\+.[0-9]\+.[0-9]\+\)\(\")\)/\1$VERSION\3/" extension/android/executorch_android/build.gradle + GRADLE_ARGS+=" -DexecuTorchVersion=${VERSION}" fi FLAVOR="${{ inputs.flavor }}" - if [[ "$FLAVOR" == "vulkan+xnnpack" ]]; then + if [[ "$FLAVOR" == "vulkan+xnnpack" || -z "$FLAVOR" ]]; then export EXECUTORCH_BUILD_VULKAN=ON fi + if [[ "$FLAVOR" == "qnn" ]]; then + PYTHON_EXECUTABLE=python bash .ci/scripts/setup-qnn-deps.sh + PYTHON_EXECUTABLE=python bash .ci/scripts/build-qnn-sdk.sh + export EXECUTORCH_BUILD_QNN=ON + source backends/qualcomm/scripts/qnn_config.sh + export QNN_SDK_ROOT="/tmp/qnn/${QNN_VERSION}" + export ANDROID_ABIS=arm64-v8a + GRADLE_ARGS+=" -DqnnVersion=2.37.0" + fi + # Build AAR Package mkdir aar-out export BUILD_AAR_DIR=aar-out @@ -106,7 +119,7 @@ jobs: # Publish to maven staging UPLOAD_TO_MAVEN="${{ inputs.upload_to_maven }}" if [[ "$UPLOAD_TO_MAVEN" == "true" ]]; then - (cd extension/android; ANDROID_HOME="${ANDROID_SDK:-/opt/android/sdk}" ./gradlew :executorch_android:publishToMavenCentral) + (cd extension/android; ANDROID_HOME="${ANDROID_SDK:-/opt/android/sdk}" ./gradlew ${GRADLE_ARGS} :executorch_android:publishToMavenCentral) fi upload-release-aar: diff --git a/.github/workflows/apple-perf.yml b/.github/workflows/apple-perf.yml index f8c33ac11c5..56fc67d1617 100644 --- a/.github/workflows/apple-perf.yml +++ b/.github/workflows/apple-perf.yml @@ -230,7 +230,7 @@ jobs: model.use_sdpa_with_kv_cache=true \ backend.xnnpack.enabled=true \ backend.xnnpack.extended_ops=true \ - base.preq_mode="8da4w_output_8da8w" \ + base.preq_mode="preq_8da4w_out_8da8w" \ base.preq_group_size=32 \ export.max_seq_length=2048 \ export.max_context_length=2048 \ @@ -256,7 +256,7 @@ jobs: base.params="${DOWNLOADED_PATH}/params.json" \ quantization.use_qat=true \ base.use_lora=16 \ - base.preq_mode="8da4w_output_8da8w" \ + base.preq_mode="preq_8da4w_out_8da8w" \ base.preq_group_size=32 \ base.preq_embedding_quantize=\'8,0\' \ model.use_sdpa_with_kv_cache=true \ @@ -416,7 +416,7 @@ jobs: - set-parameters secrets: inherit with: - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} diff --git a/.github/workflows/apple.yml b/.github/workflows/apple.yml index 651956c4635..fb3c04d07fb 100644 --- a/.github/workflows/apple.yml +++ b/.github/workflows/apple.yml @@ -17,7 +17,6 @@ on: - scripts/build_apple_llm_demo.sh - scripts/create_frameworks.sh - .ci/scripts/test_ios_ci.sh - - examples/demo-apps/apple_ios/** - extension/apple/** - extension/benchmark/apple/** - extension/module/** @@ -49,7 +48,7 @@ jobs: uses: pytorch/test-infra/.github/workflows/macos_job.yml@main secrets: inherit with: - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -136,7 +135,7 @@ jobs: needs: set-version uses: pytorch/test-infra/.github/workflows/macos_job.yml@main with: - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -156,6 +155,7 @@ jobs: "kernels_llm" "kernels_optimized" "kernels_quantized" + "kernels_torchao" "threadpool" ) @@ -275,7 +275,7 @@ jobs: uses: pytorch/test-infra/.github/workflows/macos_job.yml@main secrets: inherit with: - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} diff --git a/.github/workflows/build-presets.yml b/.github/workflows/build-presets.yml index b0455140f62..66ab19eef3c 100644 --- a/.github/workflows/build-presets.yml +++ b/.github/workflows/build-presets.yml @@ -6,8 +6,6 @@ on: branches: - main - release/* - paths: - - .github/workflows/build-presets.yml workflow_dispatch: concurrency: @@ -24,7 +22,7 @@ jobs: with: job-name: build ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: 3.12 submodules: recursive timeout: 90 @@ -111,7 +109,7 @@ jobs: strategy: fail-fast: false matrix: - preset: [pybind] + preset: [pybind, windows] with: job-name: build ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -121,14 +119,24 @@ jobs: set -eux conda init powershell powershell -Command "& { - \$ErrorActionPreference = 'Stop' Set-PSDebug -Trace 1 + \$ErrorActionPreference = 'Stop' + \$PSNativeCommandUseErrorActionPreference = \$true conda create --yes --quiet -n et python=3.12 conda activate et - python install_requirements.py - cmake --preset ${{ matrix.preset }} + + cmake --preset ${{ matrix.preset }} -T ClangCL + if (\$LASTEXITCODE -ne 0) { + Write-Host "CMake configuration was unsuccessful. Exit code: \$LASTEXITCODE." + exit \$LASTEXITCODE + } + \$numCores = [System.Environment]::GetEnvironmentVariable('NUMBER_OF_PROCESSORS') - 1 cmake --build cmake-out -j \$numCores + if (\$LASTEXITCODE -ne 0) { + Write-Host "CMake build was unsuccessful. Exit code: \$LASTEXITCODE." + exit \$LASTEXITCODE + } }" diff --git a/.github/workflows/build-wheels-aarch64-linux.yml b/.github/workflows/build-wheels-aarch64-linux.yml new file mode 100644 index 00000000000..abc378f9061 --- /dev/null +++ b/.github/workflows/build-wheels-aarch64-linux.yml @@ -0,0 +1,66 @@ +# From https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows +name: Build Aarch64 Linux Wheels + +on: + pull_request: + paths: + - .ci/**/* + - .github/workflows/build-wheels-aarch64-linux.yml + - examples/**/* + - pyproject.toml + - setup.py + tags: + - ciflow/binaries/* + push: + branches: + - nightly + - release/* + tags: + # NOTE: Binary build pipelines should only get triggered on release candidate builds + # Release candidate tags look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + - ciflow/binaries/* + workflow_dispatch: + +jobs: + generate-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: linux-aarch64 + test-infra-repository: pytorch/test-infra + test-infra-ref: main + with-cuda: disabled + with-rocm: disabled + python-versions: '["3.10", "3.11", "3.12"]' + + build: + needs: generate-matrix + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/executorch + pre-script: .ci/scripts/wheel/pre_build_script.sh + post-script: .ci/scripts/wheel/post_build_script.sh + smoke-test-script: .ci/scripts/wheel/test_linux_aarch64.py + package-name: executorch + name: ${{ matrix.repository }} + uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + submodules: recursive + env-var-script: .ci/scripts/wheel/envvar_linux.sh + pre-script: ${{ matrix.pre-script }} + post-script: ${{ matrix.post-script }} + package-name: ${{ matrix.package-name }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} + architecture: aarch64 diff --git a/.github/workflows/build-wheels-macos.yml b/.github/workflows/build-wheels-macos.yml index 055543adae4..8db10c0335b 100644 --- a/.github/workflows/build-wheels-macos.yml +++ b/.github/workflows/build-wheels-macos.yml @@ -63,7 +63,7 @@ jobs: post-script: ${{ matrix.post-script }} package-name: ${{ matrix.package-name }} # Meta's macOS runners do not have Xcode, so use GitHub's runners. - runner-type: macos-latest-xlarge + runner-type: macos-14-xlarge setup-miniconda: true smoke-test-script: ${{ matrix.smoke-test-script }} trigger-event: ${{ github.event_name }} diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml new file mode 100644 index 00000000000..276edfb08d1 --- /dev/null +++ b/.github/workflows/build-wheels-windows.yml @@ -0,0 +1,59 @@ +name: Build Windows Wheels + +on: + pull_request: + push: + branches: + - nightly + - main + - release/* + tags: + # NOTE: Binary build pipelines should only get triggered on release candidate builds + # Release candidate tags look like: v1.11.0-rc1 + - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+ + workflow_dispatch: + +permissions: + id-token: write + contents: read + +jobs: + generate-matrix: + uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main + with: + package-type: wheel + os: windows + test-infra-repository: pytorch/test-infra + test-infra-ref: main + with-cuda: disabled + with-rocm: disabled + python-versions: '["3.10", "3.11", "3.12"]' + + build: + needs: generate-matrix + strategy: + fail-fast: false + matrix: + include: + - repository: pytorch/executorch + pre-script: .ci\\scripts\\wheel\\pre_build_script.sh + env-script: .ci\\scripts\\wheel\\vc_env_helper.bat + post-script: .ci\\scripts\\wheel\\post_build_script.sh + smoke-test-script: .ci/scripts/wheel/test_windows.py + package-name: executorch + name: ${{ matrix.repository }} + uses: pytorch/test-infra/.github/workflows/build_wheels_windows.yml@main + with: + repository: ${{ matrix.repository }} + ref: "" + test-infra-repository: pytorch/test-infra + test-infra-ref: main + build-matrix: ${{ needs.generate-matrix.outputs.matrix }} + pre-script: ${{ matrix.pre-script }} + env-script: ${{ matrix.env-script }} + post-script: ${{ matrix.post-script }} + package-name: ${{ matrix.package-name }} + smoke-test-script: ${{ matrix.smoke-test-script }} + trigger-event: ${{ github.event_name }} + wheel-build-params: "--verbose" + submodules: true diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 4658fdc0d26..c220b371c0a 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -36,3 +36,51 @@ jobs: uses: ./.github/workflows/_link_check.yml with: ref: ${{ github.sha }} + + backend-test-linux: + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + strategy: + fail-fast: false + matrix: + flow: [ + qnn, qnn_16a16w, qnn_16a8w, qnn_16a4w, qnn_16a4w_block, qnn_8a8w, + vulkan, vulkan_static_int8_per_channel, + xnnpack, xnnpack_dynamic_int8_per_channel, xnnpack_static_int8_per_channel, xnnpack_static_int8_per_tensor + ] + suite: [models, operators] + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: linux.4xlarge.memory + docker-image: ci-image:executorch-ubuntu-22.04-clang12 + submodules: recursive + timeout: 120 + upload-artifact: test-report-${{ matrix.flow }}-${{ matrix.suite }} + script: | + set -eux + + source .ci/scripts/test_backend_linux.sh "${{ matrix.suite }}" "${{ matrix.flow }}" "${RUNNER_ARTIFACT_DIR}" + + backend-test-macos: + uses: pytorch/test-infra/.github/workflows/macos_job.yml@main + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + matrix: + flow: [coreml, coreml_static_int8] + suite: [models, operators] + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: macos-m1-stable + python-version: 3.12 + submodules: recursive + timeout: 120 + upload-artifact: test-report-${{ matrix.flow }}-${{ matrix.suite }} + script: | + set -eux + + # This is needed to get the prebuilt PyTorch wheel from S3 + ${CONDA_RUN} --no-capture-output pip install awscli==1.37.21 + + source .ci/scripts/test_backend_macos.sh "${{ matrix.suite }}" "${{ matrix.flow }}" "${RUNNER_ARTIFACT_DIR}" diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml index 89e1692df97..01bff087124 100644 --- a/.github/workflows/periodic.yml +++ b/.github/workflows/periodic.yml @@ -11,6 +11,8 @@ on: branches: - release/* workflow_dispatch: + pull_request: + types: [opened, synchronize, reopened, labeled, unlabeled] 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' }}-${{ github.event.schedule }} @@ -32,10 +34,11 @@ jobs: python-version: '3.10' - name: Extract the list of models to test id: gather-models + env: + EFFECTIVE_EVENT: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'ciflow/periodic') && 'schedule' || github.event_name }} run: | set -eux - - PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --event "${GITHUB_EVENT_NAME}" + PYTHONPATH="${PWD}" python .ci/scripts/gather_test_models.py --event "${EFFECTIVE_EVENT}" test-models-linux: name: test-models-linux diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 47166721cf0..815e106ae1e 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -13,6 +13,33 @@ concurrency: cancel-in-progress: true jobs: + test-qnn-wheel-packages-linux: + name: test-qnn-wheel-packages-linux + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + strategy: + fail-fast: false + matrix: + python-version: [ "3.10", "3.11", "3.12" ] + with: + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 180 + 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}" + + # Create a clean env for each python version + conda create -y -n test_env_${{ matrix.python-version }} python=${{ matrix.python-version }} + conda activate test_env_${{ matrix.python-version }} + + PYTHON_EXECUTABLE=python bash .ci/scripts/test_wheel_package_qnn.sh "${{ matrix.python-version }}" + test-setup-linux-gcc: name: test-setup-linux-gcc uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main @@ -315,7 +342,7 @@ jobs: bash examples/models/moshi/mimi/install_requirements.sh # reinstall executorch - bash ./install_executorch.sh + bash ./install_executorch.sh --minimal # run python unittest python -m unittest examples.models.moshi.mimi.test_mimi @@ -406,7 +433,7 @@ jobs: output=$(ls -la cmake-out/test/size_test) arr=($output) size=${arr[4]} - threshold="51744" + threshold="51752" if [[ "$size" -le "$threshold" ]]; then echo "Success $size <= $threshold" else @@ -801,6 +828,8 @@ jobs: id-token: write contents: read strategy: + matrix: + enable-etdump: ['', '--enable-etdump'] fail-fast: false with: runner: linux.2xlarge @@ -820,7 +849,7 @@ jobs: source .ci/scripts/setup-emscripten.sh # Test selective build - bash scripts/build_wasm_tests.sh + bash scripts/build_wasm_tests.sh ${{ matrix.enable-etdump }} # Install Jest cd cmake-out-wasm/extension/wasm/test @@ -853,14 +882,133 @@ jobs: .ci/scripts/setup-linux.sh --build-tool "cmake" # Install test requirements - pip install -r backends/nxp/requirements-tests.txt + pip install -r backends/nxp/requirements-tests-pypi.txt + pip install -r backends/nxp/requirements-tests-eiq.txt # Run pytest PYTHON_EXECUTABLE=python bash backends/nxp/run_unittests.sh - # Run aot example: - PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh + # Run aot examples: + PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh cifar10 + PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh mobilenetv2 + + + test-samsung-models-linux: + name: test-samsung-models-linux + 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-clang12-android + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -ex + + # 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}" + + # Setup python + PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool "cmake" + + # Setup Samsung SDK (AI Lite Core) and install enn backend + source .ci/scripts/setup-samsung-linux-deps.sh + + # Test models serially + models="mv2 ic3 resnet18 resnet50 mv3 ic4 dl3 edsr vit w2l" + for model in $models; do + python -m executorch.examples.samsung.aot_compiler --model_name=$model -c E9955 + done + + # Test ops + python -m unittest discover -s backends/samsung/test/ops -p "test_*.py" + + + test-vulkan-models-linux: + name: test-vulkan-models-linux + 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-clang12 + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -eux + + # 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}" + + # Setup swiftshader and Vulkan SDK which are required to build the Vulkan delegate + source .ci/scripts/setup-vulkan-linux-deps.sh + + # Setup python + PYTHON_EXECUTABLE=python \ + CMAKE_ARGS="-DEXECUTORCH_BUILD_VULKAN=ON" \ + .ci/scripts/setup-linux.sh --build-tool "cmake" + + PYTHON_EXECUTABLE=python bash backends/vulkan/test/scripts/test_model.sh --build + + # Test models serially + models="mv2 mv3 edsr resnet18 resnet50 dl3" + for model in $models; do + python -m examples.vulkan.export --model_name=$model --test + done + + + test-vulkan-operators-linux: + name: test-vulkan-operators-linux + 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-clang12 + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + set -eux + + # 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}" + + # Setup swiftshader and Vulkan SDK which are required to build the Vulkan delegate + source .ci/scripts/setup-vulkan-linux-deps.sh + + # Setup python + PYTHON_EXECUTABLE=python \ + CMAKE_ARGS="-DEXECUTORCH_BUILD_VULKAN=ON" \ + .ci/scripts/setup-linux.sh --build-tool "cmake" + # Custom operator tests + PYTHON_EXECUTABLE=python bash backends/vulkan/test/custom_ops/build_and_run.sh add + ./cmake-out/backends/vulkan/test/custom_ops/q8csw_linear + ./cmake-out/backends/vulkan/test/custom_ops/q8csw_conv2d + ./cmake-out/backends/vulkan/test/custom_ops/q4gsw_linear + ./cmake-out/backends/vulkan/test/custom_ops/choose_qparams_per_row + + # "Classic" Operator tests + PYTHON_EXECUTABLE=python bash backends/vulkan/test/scripts/test_op.sh --build + # TODO(ssjia): figure out how to run custom op tests in CI. Currently, they are + # failing due to to the libstdc++.so.6 installed with conda not supporting + # GLIBCXX_3.4.30. These tests are still run in Meta internal CI. + # ./cmake-out/backends/vulkan/test/op_tests/vulkan_sdpa_test + + # Run e2e testing for selected operators. More operators will be tested via this + # route in the future. + python -m unittest backends/vulkan/test/test_vulkan_delegate.py -k "*pt2e*" + python -m unittest backends/vulkan/test/test_vulkan_delegate.py -k "*torchao*" nxp-build-test: name: nxp-build-test diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml new file mode 100644 index 00000000000..ae7cbe6857b --- /dev/null +++ b/.github/workflows/stale.yml @@ -0,0 +1,149 @@ +# The behavior is: +# - If a PR is not labeled stale, after 60 days inactivity label the PR as stale and comment about it. +# - If a PR is labeled stale, after 30 days inactivity close the PR. +# - `high priority` and `no-stale` PRs are exempt. + +name: Close stale pull requests + +on: + schedule: + # Run daily at 00:30 UTC. + - cron: '30 0 * * *' + workflow_dispatch: + +jobs: + stale: + if: ${{ github.repository == 'pytorch/executorch' }} + runs-on: linux.large + permissions: + contents: read + pull-requests: write + + steps: + - uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1 + with: + script: | + // Do some dumb retries on requests. + const retries = 7; + const baseBackoff = 100; + const sleep = timeout => new Promise(resolve => setTimeout(resolve, timeout)); + github.hook.wrap('request', async (request, options) => { + for (let attempt = 1; attempt <= retries; attempt++) { + try { + return await request(options); + } catch (err) { + if (attempt < retries) { + core.warning(`Request getting retried. Attempt: ${attempt}`); + await sleep(baseBackoff * Math.pow(2, attempt)); + continue; + } + throw err; + } + } + }); + + const MAX_API_REQUESTS = 100; + + // If a PRs not labeled stale, label them stale after no update for 60 days. + const STALE_LABEL_THRESHOLD_MS = 1000 * 60 * 60 * 24 * 60; + // For PRs already labeled stale, close after not update for 30 days. + const STALE_CLOSE_THRESHOLD_MS = 1000 * 60 * 60 * 24 * 30; + + const STALE_MESSAGE = + "Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as `Stale`.
" + + "Feel free to remove the `Stale` label if you feel this was a mistake.
" + + "If you are unable to remove the `Stale` label please contact a maintainer in order to do so.
" + + "If you want the bot to never mark this PR stale again, add the `no-stale` label.
" + + "`Stale` pull requests will automatically be closed after 30 days of inactivity.
"; + + let numAPIRequests = 0; + let numProcessed = 0; + + async function processPull(pull) { + core.info(`[${pull.number}] URL: ${pull.html_url}`); + numProcessed += 1; + const labels = pull.labels.map((label) => label.name); + + // Skip if certain labels are present. + if (labels.includes("no-stale") || labels.includes("high priority")) { + core.info(`[${pull.number}] Skipping because PR has an exempting label.`); + return false; + } + + // Check if the PR is stale, according to our configured thresholds. + let staleThresholdMillis; + if (labels.includes("Stale")) { + core.info(`[${pull.number}] PR is labeled stale, checking whether we should close it.`); + staleThresholdMillis = STALE_CLOSE_THRESHOLD_MS; + } else { + core.info(`[${pull.number}] Checking whether to label PR as stale.`); + staleThresholdMillis = STALE_LABEL_THRESHOLD_MS; + } + + const millisSinceLastUpdated = + new Date().getTime() - new Date(pull.updated_at).getTime(); + + if (millisSinceLastUpdated < staleThresholdMillis) { + core.info(`[${pull.number}] Skipping because PR was updated recently`); + return false; + } + + // At this point, we know we should do something. + // For PRs already labeled stale, close them. + if (labels.includes("Stale")) { + core.info(`[${pull.number}] Closing PR.`); + numAPIRequests += 1; + await github.rest.issues.update({ + owner: "pytorch", + repo: "executorch", + issue_number: pull.number, + state: "closed", + }); + } else { + // For PRs not labeled stale, label them stale. + core.info(`[${pull.number}] Labeling PR as stale.`); + + numAPIRequests += 1; + await github.rest.issues.createComment({ + owner: "pytorch", + repo: "executorch", + issue_number: pull.number, + body: STALE_MESSAGE, + }); + + numAPIRequests += 1; + await github.rest.issues.addLabels({ + owner: "pytorch", + repo: "executorch", + issue_number: pull.number, + labels: ["Stale"], + }); + } + } + + for await (const response of github.paginate.iterator( + github.rest.pulls.list, + { + owner: "pytorch", + repo: "executorch", + state: "open", + sort: "created", + direction: "asc", + per_page: 100, + } + )) { + numAPIRequests += 1; + const pulls = response.data; + // Awaiting in a loop is intentional here. We want to serialize execution so + // that log groups are printed correctl + for (const pull of pulls) { + if (numAPIRequests > MAX_API_REQUESTS) { + core.warning("Max API requests exceeded, exiting."); + process.exit(0); + } + await core.group(`Processing PR #${pull.number}`, async () => { + await processPull(pull); + }); + } + } + core.info(`Processed ${numProcessed} PRs total.`); diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index e7188652949..975a8ebbb30 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -55,48 +55,102 @@ jobs: # Build and test executorch PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" - test-models-arm-zephyr: - name: test-models-arm-zephyr - uses: pytorch/test-infra/.github/workflows/linux_job.yml@main - strategy: - matrix: - model: [add] - fail-fast: false - with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-zephyr-sdk - submodules: 'recursive' - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - timeout: 120 - script: | - MODEL_NAME=${{ matrix.model }} - CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") - conda activate "${CONDA_ENV}" - - source .ci/scripts/utils.sh - source .ci/scripts/zephyr-utils.sh - mkdir -p zephyr_scratch/ - cd zephyr_scratch - export ZEPHYR_PROJ_ROOT=$(realpath $(pwd)) - - download_arm_zephyr_sdk - ./zephyr-sdk-0.16.0/setup.sh -c -t arm-zephyr-eabi - - cd $ZEPHYR_PROJ_ROOT - setup_zephyr_et_module - - cd $ZEPHYR_PROJ_ROOT/modules/lib/executorch - install_executorch "--use-pt-pinned-commit" - .ci/scripts/setup-arm-baremetal-tools.sh --target-toolchain zephyr - source examples/arm/ethos-u-scratch/setup_path.sh - source $ZEPHYR_PROJ_ROOT/zephyr/zephyr-env.sh - cd $ZEPHYR_PROJ_ROOT/zephyr/samples/modules/executorch/arm/hello_world - west build -p always -b mps3/corstone300/fvp - FVP_Corstone_SSE-300_Ethos-U55 -a build/zephyr/zephyr.elf -C mps3_board.visualisation.disable-visualisation=1 -C mps3_board.telnetterminal0.start_telnet=0 -C mps3_board.uart0.out_file='sim.out' -C cpu0.CFGITCMSZ=15 -C cpu0.CFGDTCMSZ=15 --simlimit 120 - - grep -qF "Output[0][0]: (float) 2.000000" sim.out - exit_status=$? #store 0 if found (success), 1 if not (failure) - exit $exit_status +# test-models-arm-zephyr: +# name: test-models-arm-zephyr +# uses: pytorch/test-infra/.github/workflows/linux_job.yml@main +# strategy: +# matrix: +# model: [add, softmax, mv2] +# fail-fast: false +# with: +# runner: linux.2xlarge +# docker-image: ci-image:executorch-ubuntu-22.04-zephyr-sdk +# submodules: 'recursive' +# ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} +# timeout: 120 +# script: | +# MODEL_NAME=${{ matrix.model }} +# CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") +# conda activate "${CONDA_ENV}" +# if [[ ${{ matrix.model}} == "add" ]]; then +# SIM_LIMIT_SEC=60 +# elif [[ ${{ matrix.model}} == "softmax" ]]; then +# SIM_LIMIT_SEC=60 +# elif [[ ${{ matrix.model}} == "mv2" ]]; then +# SIM_LIMIT_SEC=5000 +# else +# echo "Failed unsupported model selection ${{ matrix.model }}" +# exit 1 +# fi +# +# source .ci/scripts/utils.sh +# source .ci/scripts/zephyr-utils.sh +# mkdir -p zephyr_scratch/ +# cd zephyr_scratch +# export ZEPHYR_PROJ_ROOT=$(realpath $(pwd)) +# export ARM_FVP_TUTORIALS_ROOT=$ZEPHYR_PROJ_ROOT/zephyr/samples/modules/executorch/arm-fvp-tutorials +# +# # TODO @Bujji: Should see if this can be moved into the docker image itself +# download_arm_zephyr_sdk +# ./zephyr-sdk-0.17.2/setup.sh -c -t arm-zephyr-eabi +# cd $ZEPHYR_PROJ_ROOT +# setup_zephyr_et_module +# +# # Run setup scripts for Arm FVP and Arm AOT Compilation +# cd $ZEPHYR_PROJ_ROOT/modules/lib/executorch +# install_executorch +# .ci/scripts/setup-arm-baremetal-tools.sh --target-toolchain zephyr +# source examples/arm/ethos-u-scratch/setup_path.sh +# source $ZEPHYR_PROJ_ROOT/zephyr/zephyr-env.sh +# +# # Get the model as PTE +# python -m examples.arm.aot_arm_compiler \ +# --model_name="${MODEL_NAME}" \ +# --output="${MODEL_NAME}.pte" +# +# # Generate the C-style header +# cd $ARM_FVP_TUTORIALS_ROOT +# python build_model.py \ +# --executorch-root $ZEPHYR_PROJ_ROOT/modules/lib/executorch \ +# --pte-file $ZEPHYR_PROJ_ROOT/modules/lib/executorch/${MODEL_NAME}.pte \ +# --output-path $ARM_FVP_TUTORIALS_ROOT/models/${MODEL_NAME}/src/ +# +# cd $ARM_FVP_TUTORIALS_ROOT/models/${MODEL_NAME}/ +# +# # Build the zephyr elf +# west build -p always -b mps3/corstone300/fvp -- \ +# -DET_PTE_FILE_PATH_FOR_SELECTIVE_BUILD=$ZEPHYR_PROJ_ROOT/modules/lib/executorch/${MODEL_NAME}.pte +# +# # Run the simulation +# FVP_Corstone_SSE-300_Ethos-U55 -a build/zephyr/zephyr.elf \ +# -C mps3_board.visualisation.disable-visualisation=1 \ +# -C mps3_board.telnetterminal0.start_telnet=0 \ +# -C mps3_board.uart0.out_file='sim.out' \ +# -C cpu0.CFGITCMSZ=15 \ +# -C cpu0.CFGDTCMSZ=15 \ +# --simlimit ${SIM_LIMIT_SEC} +# +# # Disable exit on error +# set +e +# # Report failure if any of the ouptut verification checks fail +# grep -qF "ERROR" sim.out +# exit_status=$? #store 0 if found (failure), 1 if not (success) +# if [[ "$exit_status" -eq "0" ]]; then +# cat sim.out +# set -e +# exit 1 +# fi +# +# # Report fail if simulation does not complete successfully +# grep -qF "SUCCESS: Program complete, exiting." sim.out +# exit_status=$? #store 0 if found (success), 1 if not (failure) +# if [[ "$exit_status" -eq "1" ]]; then +# cat sim.out +# set -e +# exit 1 +# fi +# # Re-enable exit on error +# set -e test-models-linux-aarch64: name: test-models-linux-aarch64 @@ -122,7 +176,7 @@ jobs: - model: phi_4_mini backend: portable runner: linux.arm64.m7g.4xlarge - - model: qwen2_5 + - model: qwen2_5_1_5b backend: portable runner: linux.arm64.2xlarge - model: llama3_2_vision_encoder @@ -234,6 +288,7 @@ jobs: - test_arm_baremetal: test_models_tosa - test_arm_baremetal: test_models_ethos-u55 - test_arm_baremetal: test_models_ethos-u85 + - test_arm_baremetal: test_smaller_stories_llama fail-fast: false with: runner: linux.2xlarge.memory @@ -380,7 +435,7 @@ jobs: name: test-coreml-delegate uses: pytorch/test-infra/.github/workflows/macos_job.yml@main with: - runner: macos-latest-xlarge + runner: macos-14-xlarge python-version: '3.11' submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -430,7 +485,7 @@ jobs: eval "$(conda shell.bash hook)" # Install requirements - ${CONDA_RUN} EXECUTORCH_BUILD_TORCHAO=1 python install_executorch.py + ${CONDA_RUN} EXECUTORCH_BUILD_KERNELS_TORCHAO=1 python install_executorch.py ${CONDA_RUN} sh examples/models/llama/install_requirements.sh # Run test @@ -530,6 +585,37 @@ jobs: # Test llama2 PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_llama.sh -model stories110M -build_tool cmake -dtype "${DTYPE}" -mode "${MODE}" + test-torchao-huggingface-checkpoints: + name: test-torchao-huggingface-checkpoints + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + strategy: + matrix: + model: [qwen3_4b, phi_4_mini] + include: + - model: qwen3_4b + test_with_runner: true + - model: phi_4_mini + test_with_runner: false + fail-fast: false + with: + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 900 + 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}" + + PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh --build-tool cmake + pip install -U "huggingface_hub[cli]" + + bash .ci/scripts/test_torchao_huggingface_checkpoints.sh ${{ matrix.model }} ${{ matrix.test_with_runner && '--test_with_runner' || '' }} + # # TODO(jackzhxng): Runner consistently runs out of memory before test finishes. Try to find a more powerful runner. # test-llava-runner-macos: # name: test-llava-runner-macos @@ -568,7 +654,7 @@ jobs: strategy: matrix: dtype: [fp32] - model: [dl3, mv3, mv2, ic4, ic3, vit, mb, w2l] + model: [dl3, mv3, mv2, ic4, ic3, vit, mb, w2l, conv_former] fail-fast: false with: runner: linux.2xlarge @@ -677,10 +763,10 @@ jobs: echo "::endgroup::" done - test-huggingface-transformers: + test-huggingface-transformers-xnnpack: # 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 + name: test-huggingface-transformers-xnnpack uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write @@ -688,12 +774,15 @@ jobs: secrets: inherit strategy: matrix: - hf_model_id: [ - google/gemma-3-1b-it, - Qwen/Qwen3-0.6B, - HuggingFaceTB/SmolLM2-135M, - meta-llama/Llama-3.2-1B, - allenai/OLMo-1B-hf, + config: [ + # XNNPack. + llama3.2-1b|xnnpack|--quantize, + qwen3-0.6b|xnnpack|--quantize, + qwen3-1.7b|xnnpack|--quantize, + gemma3-1b|xnnpack|--quantize, + phi4-mini|xnnpack|--quantize, + smollm2-135m|xnnpack|--quantize, + smollm3-3b|xnnpack|--quantize ] fail-fast: false with: @@ -705,6 +794,12 @@ jobs: timeout: 90 upload-artifact: profiling-artifacts-${{ strategy.job-index }} script: | + set -eux + IFS='|' read -r MODEL RECIPE QUANTIZE <<< "${{ matrix.config }}" + echo "Model: $MODEL" + echo "Recipe: $RECIPE" + echo "Quantize: $QUANTIZE" + echo "::group::Set up ExecuTorch" # 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]") @@ -742,79 +837,52 @@ jobs: pip list echo "::endgroup::" - echo "::group::Export to ExecuTorch" - # Pass matrix variable as environment variable - export MODEL_ID="${{ matrix.hf_model_id }}" - export OUTPUT_DIR="$(pwd)/${MODEL_ID}_custom_sdpa_kv_cache_8da4w" - pushd optimum-executorch - - ARGS=( - "--model" "${MODEL_ID}" - "--task" "text-generation" - "--recipe" "xnnpack" - "--use_custom_sdpa" - "--use_custom_kv_cache" - "--qlinear" "8da4w" - "--qembedding" "8w" - "--output_dir" "${OUTPUT_DIR}" - ) - - optimum-cli export executorch "${ARGS[@]}" - - ls -FlAGhp ${OUTPUT_DIR} - popd - echo "::endgroup::" - - echo "::group::Inference using python API" - pushd optimum-executorch - python -c " - import os - from optimum.executorch import ExecuTorchModelForCausalLM - from transformers import AutoTokenizer - - model_id = os.getenv('MODEL_ID') - pte_dir = os.getenv('OUTPUT_DIR') - print(f'Loading model {model_id} from {pte_dir}.') - model = ExecuTorchModelForCausalLM.from_pretrained(pte_dir) - generated_text = model.text_generation( - tokenizer=AutoTokenizer.from_pretrained(model_id), - prompt='Simply put, the theory of relativity states that', - max_seq_len=64 - ) - print(generated_text) - " - popd + echo "::group::Run tests" + export OUTPUT_DIR="$(pwd)/${MODEL}_${RECIPE}_${QUANTIZE}" + python .ci/scripts/test_huggingface_optimum_model.py --model ${MODEL} --recipe ${RECIPE} ${QUANTIZE} --model_dir ${OUTPUT_DIR} echo "::endgroup::" - echo "::group::Inference using executor_runner with ETDump" + echo "::group::Generate artifacts for performance profiling" ./cmake-out/executor_runner \ --model_path ${OUTPUT_DIR}/model.pte \ --etdump_path ${OUTPUT_DIR}/etdump.etdp - export TSV_PATH=artifacts-to-be-uploaded/${MODEL_ID}_op_prof.tsv + export TSV_PATH=artifacts-to-be-uploaded/${MODEL}_op_prof.tsv mkdir -p $(dirname "$TSV_PATH") python3 -m devtools.inspector.inspector_cli \ --etdump_path ${OUTPUT_DIR}/etdump.etdp \ --tsv_path ${TSV_PATH} - echo "::endgroup::" - test-huggingface-optimum-coreml: + test-huggingface-transformers-macos: # 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-optimum-coreml + name: test-huggingface-transformers-macos uses: pytorch/test-infra/.github/workflows/macos_job.yml@main permissions: id-token: write contents: read secrets: inherit + # Models below selected based on https://huggingface.co/models?pipeline_tag=text-generation&num_parameters=min:0,max:3B&sort=trending. strategy: matrix: config: [ - qwen3|coreml_fp32_gpu|--quantize, - smollm|coreml_fp32_gpu|--quantize, - llama3|coreml_fp32_gpu|--quantize, - olmo|coreml_fp32_gpu|--quantize, + # # XNNPack. (Skipping for now due to intermittent segmentation faults, see https://github.com/huggingface/optimum-executorch/issues/122.) + # llama3.2-1b|xnnpack|--quantize, + # qwen3-0.6b|xnnpack|--quantize, + # qwen3-1.7b|xnnpack|--quantize, + # gemma3-1b|xnnpack|--quantize, + # phi4-mini|xnnpack|--quantize, + # smollm2-135m|xnnpack|--quantize, + # smollm3-3b|xnnpack|--quantize, + # qwen3-1.7b|xnnpack|--quantize, + # CoreML. + llama3.2-1b|coreml_fp32_gpu|--quantize, + qwen3-0.6b|coreml_fp32_gpu|--quantize, + smollm2-135m|coreml_fp32_gpu|--quantize, + olmo-1b|coreml_fp32_gpu|--quantize, + bert|coreml_fp32_gpu|--quantize, + distilbert|coreml_fp32_gpu|--quantize ] fail-fast: false with: @@ -903,3 +971,66 @@ jobs: build-mode: Release build-tool: cmake docker-image: ci-image:executorch-ubuntu-22.04-clang12 + + test-mcu-models: + name: test-mcu-models + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + strategy: + matrix: + include: + - build-tool: cmake + fail-fast: false + 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: | + BUILD_TOOL=${{ matrix.build-tool }} + + # 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}" + + # Try to mirror these as closely as possible + source .ci/scripts/utils.sh + install_executorch "--use-pt-pinned-commit" + + .ci/scripts/setup-arm-baremetal-tools.sh + source examples/arm/ethos-u-scratch/setup_path.sh + + # Run selective Build + chmod +x examples/selective_build/test_selective_build.sh + examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" + + # Run MCU models + chmod +x examples/arm/run_mcu_models_fvp.sh + examples/arm/run_mcu_models_fvp.sh --target=cortex-m55 + + test-models-windows: + uses: pytorch/test-infra/.github/workflows/windows_job.yml@main + strategy: + fail-fast: false + matrix: + model: [linear, add, add_mul, ic3, ic4, mv2, mv3, resnet18, resnet50, vit, w2l, mobilebert, emformer_join, emformer_transcribe] + backend: [portable, xnnpack-f32, xnnpack-q8] + with: + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 60 + script: | + conda init powershell + + powershell -Command "& { + Set-PSDebug -Trace 1 + \$ErrorActionPreference = 'Stop' + \$PSNativeCommandUseErrorActionPreference = \$true + + .ci/scripts/setup-windows.ps1 + + .ci/scripts/test_model.ps1 -modelName ${{ matrix.model }} -backend ${{ matrix.backend }} + }" diff --git a/.gitignore b/.gitignore index 08d14e13582..b166f8c9512 100644 --- a/.gitignore +++ b/.gitignore @@ -16,14 +16,18 @@ cmake-android-out/ cmake-ios-out/ cmake-out* cmake-out-android/ +build-android/ +build-x86/ dist/ ethos-u-scratch/ executorch.egg-info pip-out/ +build-profiling/ # Any exported models and profiling outputs *.bin *.model +*.etdump tokenizer.json *.pte *.ptd @@ -58,6 +62,12 @@ xcuserdata/ /include/ /share/ /version.py +*.csv +*_etdump # Android *.aar + +# Windows +*.dll +*.pyd diff --git a/.gitmodules b/.gitmodules index 945ae5ed51e..1f202d4fdec 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,9 +1,6 @@ [submodule "backends/arm/third-party/ethos-u-core-driver"] path = backends/arm/third-party/ethos-u-core-driver url = https://git.gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-driver.git -[submodule "backends/arm/third-party/serialization_lib"] - path = backends/arm/third-party/serialization_lib - url = https://git.gitlab.arm.com/tosa/tosa-serialization.git [submodule "backends/vulkan/third-party/Vulkan-Headers"] path = backends/vulkan/third-party/Vulkan-Headers url = https://github.com/KhronosGroup/Vulkan-Headers @@ -30,7 +27,7 @@ url = https://github.com/google/pthreadpool.git [submodule "extension/llm/tokenizers"] path = extension/llm/tokenizers - url = https://github.com/pytorch-labs/tokenizers.git + url = https://github.com/meta-pytorch/tokenizers.git [submodule "kernels/optimized/third-party/eigen"] path = kernels/optimized/third-party/eigen url = https://gitlab.com/libeigen/eigen.git diff --git a/.lintrunner.toml b/.lintrunner.toml index 07227998c2c..0b6a6eb8908 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -73,7 +73,6 @@ exclude_patterns = [ '**/third-party/**', # NB: Objective-C is not supported 'examples/apple/**', - 'examples/demo-apps/apple_ios/**', 'examples/demo-apps/react-native/rnllama/ios/**', 'extension/apple/**', 'extension/llm/apple/**', @@ -136,6 +135,36 @@ init_command = [ '--requirement=requirements-lintrunner.txt', ] +[[linter]] +code = 'CMAKEFORMAT' +include_patterns = [ + "**/*.cmake", + "**/*.cmake.in", + "**/CMakeLists.txt", +] +exclude_patterns = [ + 'third-party/**', + '**/third-party/**', +] +command = [ + 'python', + '-m', + 'lintrunner_adapters', + 'run', + 'cmake_format_linter', + '--', + '@{{PATHSFILE}}', +] +init_command = [ + 'python', + '-m', + 'lintrunner_adapters', + 'run', + 'pip_init', + '--dry-run={{DRYRUN}}', + '--requirement=requirements-lintrunner.txt', +] + [[linter]] code = 'ETCAPITAL' include_patterns = [ diff --git a/CMakeLists.txt b/CMakeLists.txt index 9dc77596d37..fc427d517a9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,7 +50,10 @@ cmake_minimum_required(VERSION 3.29) project(executorch) +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}) + include(${PROJECT_SOURCE_DIR}/tools/cmake/common/preset.cmake) +include(${PROJECT_SOURCE_DIR}/tools/cmake/Codegen.cmake) include(${PROJECT_SOURCE_DIR}/tools/cmake/Utils.cmake) include(CMakeDependentOption) include(ExternalProject) @@ -76,11 +79,6 @@ if(NOT PYTHON_EXECUTABLE) endif() announce_configured_options(PYTHON_EXECUTABLE) -if(NOT BUCK2) - resolve_buck2() -endif() -announce_configured_options(BUCK2) - announce_configured_options(CMAKE_CXX_COMPILER_ID) announce_configured_options(CMAKE_TOOLCHAIN_FILE) announce_configured_options(BUILD_TESTING) @@ -123,8 +121,6 @@ set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON) # Instead please use `find_package(executorch REQUIRED)` in the example # directory and add a new executable in the example `CMakeLists.txt`. -set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}) - if(NOT EXECUTORCH_ENABLE_LOGGING) # Avoid pulling in the logging strings, which can be large. Note that this # will set the compiler flag for all targets in this directory, and for all @@ -147,9 +143,13 @@ endif() # -ffunction-sections -fdata-sections: breaks function and data into sections so # they can be properly gc'd. -s: strip symbol. -set(CMAKE_CXX_FLAGS_RELEASE - "-ffunction-sections -fdata-sections ${CMAKE_CXX_FLAGS_RELEASE}" -) +if(WIN32) + set(CMAKE_CXX_FLAGS_RELEASE "/Gy /Gw ${CMAKE_CXX_FLAGS_RELEASE}") +else() + set(CMAKE_CXX_FLAGS_RELEASE + "-ffunction-sections -fdata-sections ${CMAKE_CXX_FLAGS_RELEASE}" + ) +endif() if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -s") endif() @@ -278,25 +278,6 @@ if(EXECUTORCH_BUILD_PTHREADPOOL) ) endif() -if(EXECUTORCH_BUILD_KERNELS_TORCHAO) - set(TORCHAO_BUILD_ATEN_OPS OFF) - set(TORCHAO_BUILD_EXECUTORCH_OPS ON) - set(TORCHAO_BUILD_CPU_AARCH64 ON) - set(TORCHAO_ENABLE_ARM_NEON_DOT ON) - - list(APPEND TORCHAO_INCLUDE_DIRS - ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include - ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include - ${EXECUTORCH_ROOT}/third-party/ao - ) - - set(EXECUTORCH_INCLUDE_DIRS ${TORCHAO_INCLUDE_DIRS}) - - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/ao/torchao/experimental) - executorch_target_link_options_shared_lib(torchao_ops_executorch) - list(APPEND _executorch_kernels torchao_ops_executorch) -endif() - if(EXECUTORCH_BUILD_TESTS) set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) include(CTest) @@ -327,23 +308,15 @@ set(_common_include_directories ) # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # - -if(NOT EXECUTORCH_SRCS_FILE) - # A file wasn't provided. Run a script to extract the source lists from the - # buck2 build system and write them to a file we can include. - # - # NOTE: This will only happen once during cmake setup, so it will not re-run - # if the buck2 targets change. - message(STATUS "executorch: Generating source lists") - set(EXECUTORCH_SRCS_FILE "${CMAKE_CURRENT_BINARY_DIR}/executorch_srcs.cmake") - extract_sources(${EXECUTORCH_SRCS_FILE}) +if(EXECUTORCH_SRCS_FILE) + message( + WARNING + "EXECUTORCH_SRCS_FILE is no longer necessary and will not affect the build." + ) endif() - -# This file defines the `___srcs` variables used below. -message(STATUS "executorch: Using sources file ${EXECUTORCH_SRCS_FILE}") -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() # Detect if an iOS toolchain is set. if(CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") @@ -440,6 +413,12 @@ if(MAX_KERNEL_NUM) ) endif() +# Build devtools first if needed - some backends depend on protobuf from +# devtools +if(EXECUTORCH_BUILD_DEVTOOLS) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools) +endif() + if(EXECUTORCH_BUILD_PYBIND AND APPLE) # shared version add_library(executorch_core_shared SHARED ${_executorch_core__srcs}) @@ -596,6 +575,11 @@ if(EXECUTORCH_BUILD_QNN) list(APPEND _executorch_backends qnn_executorch_backend) endif() +if(EXECUTORCH_BUILD_ENN) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/samsung) + list(APPEND _executorch_backends enn_backend) +endif() + if(EXECUTORCH_BUILD_XNNPACK) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/xnnpack) list(APPEND _executorch_backends xnnpack_backend) @@ -605,13 +589,8 @@ if(EXECUTORCH_BUILD_CORTEX_M) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m) endif() -if(EXECUTORCH_BUILD_DEVTOOLS) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools) -endif() - if(EXECUTORCH_BUILD_EXTENSION_APPLE) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/apple) - list(APPEND _executorch_extensions apple_extension) endif() if(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER) @@ -678,7 +657,6 @@ endif() if(EXECUTORCH_BUILD_EXTENSION_LLM_APPLE) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/apple) - list(APPEND _executorch_extensions extension_llm_apple) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) @@ -701,6 +679,63 @@ if(EXECUTORCH_BUILD_PTHREADPOOL AND EXECUTORCH_BUILD_CPUINFO) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/threadpool) endif() +if(EXECUTORCH_BUILD_KERNELS_TORCHAO) + if(NOT TARGET cpuinfo) + message( + FATAL_ERROR + "EXECUTORCH_BUILD_KERNELS_TORCHAO requires EXECUTORCH_BUILD_CPUINFO be set ON" + ) + endif() + if(NOT TARGET pthreadpool) + message( + FATAL_ERROR + "EXECUTORCH_BUILD_KERNELS_TORCHAO requires EXECUTORCH_BUILD_PTHREADPOOL be set ON" + ) + endif() + + # Configure TorchAO kernels + set(TORCHAO_BUILD_ATEN_OPS OFF) + set(TORCHAO_BUILD_EXECUTORCH_OPS ON) + set(TORCHAO_BUILD_CPU_AARCH64 ON) + set(TORCHAO_ENABLE_ARM_NEON_DOT ON) + set(TORCHAO_BUILD_KLEIDIAI ON) + + # TorchAO kernels look for EXECUTORCH_INCLUDE_DIRS + if(DEFINED EXECUTORCH_INCLUDE_DIRS) + message(FATAL_ERROR "EXECUTORCH_INCLUDE_DIRS is already defined") + endif() + set(EXECUTORCH_INCLUDE_DIRS + ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/pthreadpool/include + ${EXECUTORCH_ROOT}/backends/xnnpack/third-party/cpuinfo/include + ) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/ao/torchao/csrc/cpu) + unset(EXECUTORCH_INCLUDE_DIRS) + + executorch_target_link_options_shared_lib(torchao_ops_executorch) + list(APPEND _executorch_kernels torchao_ops_executorch) + + install( + TARGETS torchao_ops_executorch torchao_kernels_aarch64 + EXPORT ExecuTorchTargets + DESTINATION lib + INCLUDES + DESTINATION ${_common_include_directories} + ) + # If using KleidiAI and XNNPACK has not installed it already, install it + if(TORCHAO_BUILD_KLEIDIAI AND NOT (EXECUTORCH_BUILD_XNNPACK + AND EXECUTORCH_XNNPACK_ENABLE_KLEIDI) + ) + install( + TARGETS kleidiai + EXPORT ExecuTorchTargets + DESTINATION lib + INCLUDES + DESTINATION ${_common_include_directories} + ) + endif() + +endif() + if(EXECUTORCH_BUILD_PYBIND) # Add codegen tools subdirectory for selective_build pybind module @@ -714,6 +749,30 @@ if(EXECUTORCH_BUILD_PYBIND) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools) endif() + # Create bundled_module target only for pybindings when bundled_program exists + # This target has hard dependencies on devtools generated headers + if(TARGET bundled_program) + add_library( + bundled_module STATIC + ${CMAKE_CURRENT_SOURCE_DIR}/extension/module/bundled_module.cpp + ) + + # Ensure bundled_module waits for bundled_program's generated headers + add_dependencies(bundled_module bundled_program) + + target_link_libraries(bundled_module PRIVATE extension_data_loader) + target_link_libraries( + bundled_module PUBLIC extension_module_static bundled_program + ) + + target_include_directories( + bundled_module PUBLIC ${_common_include_directories} + ) + target_compile_options( + bundled_module PUBLIC -Wno-deprecated-declarations -fPIC + ) + endif() + # find pytorch lib, to allow pybind to take at::Tensor as input/output find_package_torch() find_library( @@ -731,6 +790,16 @@ if(EXECUTORCH_BUILD_PYBIND) torch ) + if(EXECUTORCH_BUILD_EXTENSION_MODULE) + # Always use static linking for pybindings to avoid runtime symbol + # resolution issues + list(APPEND _dep_libs extension_module_static) + # Add bundled_module if available + if(TARGET bundled_module) + list(APPEND _dep_libs bundled_module) + endif() + endif() + if(EXECUTORCH_BUILD_TESTS) list(APPEND _dep_libs test_backend_compiler_lib) endif() @@ -753,12 +822,24 @@ if(EXECUTORCH_BUILD_PYBIND) list(APPEND _dep_libs openvino_backend) endif() + if(EXECUTORCH_BUILD_QNN) + list(APPEND _dep_libs qnn_executorch_backend) + endif() + + if(EXECUTORCH_BUILD_ENN) + list(APPEND _dep_libs enn_backend) + endif() + if(EXECUTORCH_BUILD_XNNPACK) # need to explicitly specify XNNPACK and xnnpack-microkernels-prod here # otherwise uses XNNPACK and microkernel-prod symbols from libtorch_cpu list(APPEND _dep_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod) endif() + if(EXECUTORCH_BUILD_VULKAN) + list(APPEND _dep_libs vulkan_backend) + endif() + # compile options for pybind set(_pybind_compile_options -Wno-deprecated-declarations -fPIC -frtti -fexceptions @@ -799,6 +880,10 @@ if(EXECUTORCH_BUILD_WASM) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/wasm) endif() +if(EXECUTORCH_BUILD_TOKENIZERS_WASM) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/wasm/tokenizers) +endif() + if(EXECUTORCH_BUILD_EXTENSION_TRAINING) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/training) list(APPEND _executorch_extensions extension_training) @@ -845,6 +930,19 @@ if(NOT EXECUTORCH_SELECT_OPS_YAML STREQUAL "" OR NOT EXECUTORCH_SELECT_OPS_LIST STREQUAL "" OR NOT EXECUTORCH_SELECT_OPS_MODEL STREQUAL "" ) + # Use optimized kernels when available. + if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED) + set(_selected_kernel_functions_yaml + ${CMAKE_CURRENT_BINARY_DIR}/configurations/merged.yaml + ) + set(_selected_kernel_libs "optimized_kernels" "optimized_portable_kernels") + else() + set(_selected_kernel_functions_yaml + ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml + ) + set(_selected_kernel_libs "portable_kernels") + endif() + gen_selected_ops( LIB_NAME "executorch_selected_kernels" @@ -864,7 +962,7 @@ if(NOT EXECUTORCH_SELECT_OPS_YAML STREQUAL "" LIB_NAME "executorch_selected_kernels" FUNCTIONS_YAML - ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml + ${_selected_kernel_functions_yaml} CUSTOM_OPS_YAML "" DTYPE_SELECTIVE_BUILD @@ -875,13 +973,19 @@ if(NOT EXECUTORCH_SELECT_OPS_YAML STREQUAL "" LIB_NAME "executorch_selected_kernels" KERNEL_LIBS - "portable_kernels" + ${_selected_kernel_libs} DEPS executorch_core DTYPE_SELECTIVE_BUILD "${EXECUTORCH_ENABLE_DTYPE_SELECTIVE_BUILD}" ) list(APPEND _executorch_kernels executorch_selected_kernels) + + install( + TARGETS executorch_selected_kernels + EXPORT ExecuTorchTargets + DESTINATION lib + ) else() # No selective build - link the full library. if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED) @@ -892,6 +996,10 @@ else() endif() target_link_libraries(executorch_kernels INTERFACE ${_executorch_kernels}) +install(TARGETS executorch_backends executorch_extensions executorch_kernels + EXPORT ExecuTorchTargets +) + if(EXECUTORCH_BUILD_EXECUTOR_RUNNER) # Baseline libraries that executor_runner will link against. set(_executor_runner_libs executorch extension_evalue_util diff --git a/CMakePresets.json b/CMakePresets.json index 9a3e9290d43..bcf3bbc8d83 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -6,6 +6,36 @@ "hidden": true, "binaryDir": "${sourceDir}/cmake-out" }, + { + "name": "android-arm64-v8a", + "displayName": "Build executorch core and JNI bindings on android arm64-v8a", + "inherits": ["common"], + "binaryDir": "${sourceDir}/cmake-out-android-arm64-v8a", + "cacheVariables": { + "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/android.cmake", + "ANDROID_ABI": "arm64-v8a" + }, + "condition": { + "type": "inList", + "string": "${hostSystemName}", + "list": ["Darwin", "Linux", "Windows"] + } + }, + { + "name": "android-x86_64", + "displayName": "Build executorch core and JNI bindings on android x86_64", + "inherits": ["common"], + "binaryDir": "${sourceDir}/cmake-out-android-x86_64", + "cacheVariables": { + "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/android.cmake", + "ANDROID_ABI": "x86_64" + }, + "condition": { + "type": "inList", + "string": "${hostSystemName}", + "list": ["Darwin", "Linux", "Windows"] + } + }, { "name": "macos", "displayName": "Build ExecuTorch for macOS", @@ -120,6 +150,21 @@ ] } }, + { + "name": "windows", + "displayName": "Build ExecuTorch for Windows", + "inherits": ["common"], + "cacheVariables": { + "CMAKE_SYSTEM_NAME": "Windows", + "EXECUTORCH_BUILD_PRESET_FILE": "${sourceDir}/tools/cmake/preset/windows.cmake" + }, + "toolset": "ClangCL", + "condition": { + "lhs": "${hostSystemName}", + "type": "equals", + "rhs": "Windows" + } + }, { "name": "zephyr", "displayName": "Build ExecuTorch for Zephyr RTOS", diff --git a/LICENSE b/LICENSE index fdd029058a6..f20b198d808 100644 --- a/LICENSE +++ b/LICENSE @@ -8,6 +8,7 @@ Copyright (c) Qualcomm Innovation Center, Inc. Copyright (c) 2023 Apple Inc. Copyright (c) 2024 MediaTek Inc. Copyright 2023 NXP +Copyright (c) 2025 Samsung Electronics Co. LTD Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: diff --git a/Package.swift b/Package.swift index ba61d162527..3186284f5f6 100644 --- a/Package.swift +++ b/Package.swift @@ -84,6 +84,11 @@ let products = deliverables([ ], ], "kernels_quantized": [:], + "kernels_torchao": [ + "targets": [ + "threadpool", + ], + ], ]) let targets = deliverables([ diff --git a/README-wheel.md b/README-wheel.md index 12906bfd382..a59af8ea05f 100644 --- a/README-wheel.md +++ b/README-wheel.md @@ -25,6 +25,6 @@ tutorials and documentation. Here are some starting points: * [Exporting to ExecuTorch](https://pytorch.org/executorch/main/tutorials/export-to-executorch-tutorial) * Learn the fundamentals of exporting a PyTorch `nn.Module` to ExecuTorch, and optimizing its performance using quantization and hardware delegation. -* Running LLaMA on [iOS](docs/source/llm/llama-demo-ios.md) and [Android](docs/source/llm/llama-demo-android.md) devices. +* Running etLLM on [iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) and [Android](docs/source/llm/llama-demo-android.md) devices. * Build and run LLaMA in a demo mobile app, and learn how to integrate models with your own apps. diff --git a/README.md b/README.md index 4ddc84ee253..17327990a1d 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ To get started you can: - Visit the [Step by Step Tutorial](https://pytorch.org/executorch/stable/getting-started.html) to get things running locally and deploy a model to a device - Use this [Colab Notebook](https://colab.research.google.com/drive/1qpxrXC3YdJQzly3mRg-4ayYiOjC6rue3?usp=sharing) to start playing around right away -- Jump straight into LLM use cases by following specific instructions for popular open-source models such as [Llama](examples/models/llama/README.md), [Qwen 3](examples/models/qwen3/README.md), [Phi-4-mini](examples/models/phi_4_mini/README.md), and [Llava](examples/models/llava/README.md) +- Jump straight into LLM use cases by following specific instructions for popular open-source models such as [Llama](examples/models/llama/README.md), [Qwen 3](examples/models/qwen3/README.md), [Phi-4-mini](examples/models/phi_4_mini/README.md), [Llava](examples/models/llava/README.md), [Voxtral](examples/models/voxtral/README.md), and [LFM2](examples/models/lfm2/README.md). ## Feedback and Engagement diff --git a/backends/apple/coreml/TARGETS b/backends/apple/coreml/TARGETS index 188d2b63b53..444e886b4e6 100644 --- a/backends/apple/coreml/TARGETS +++ b/backends/apple/coreml/TARGETS @@ -60,6 +60,45 @@ runtime.python_library( ], ) +runtime.python_library( + name = "coreml_recipes", + srcs = [ + "recipes/__init__.py", + "recipes/coreml_recipe_provider.py" + ], + visibility = [ + "@EXECUTORCH_CLIENTS", + "//executorch/export/...", + ], + deps = [ + "fbsource//third-party/pypi/coremltools:coremltools", + ":coreml_recipe_types", + ":backend", + ":partitioner", + ":quantizer", + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/exir/backend:compile_spec_schema", + "//executorch/exir/backend:partitioner", + "//executorch/exir/backend:utils", + "//executorch/export:lib", + ], +) + +runtime.python_library( + name = "coreml_recipe_types", + srcs = [ + "recipes/coreml_recipe_types.py", + ], + visibility = [ + "@EXECUTORCH_CLIENTS", + "//executorch/export/...", + ], + deps = [ + "//executorch/export:recipe", + ], +) + runtime.cxx_python_extension( name = "executorchcoreml", srcs = [ @@ -100,10 +139,13 @@ runtime.python_test( "test/*.py", ]), deps = [ + "fbsource//third-party/pypi/coremltools:coremltools", "fbsource//third-party/pypi/pytest:pytest", ":partitioner", ":quantizer", + ":coreml_recipes", "//caffe2:torch", "//pytorch/vision:torchvision", + "fbsource//third-party/pypi/scikit-learn:scikit-learn", ], ) diff --git a/backends/apple/coreml/compiler/coreml_preprocess.py b/backends/apple/coreml/compiler/coreml_preprocess.py index edf7aa97241..d1614f30451 100644 --- a/backends/apple/coreml/compiler/coreml_preprocess.py +++ b/backends/apple/coreml/compiler/coreml_preprocess.py @@ -17,6 +17,10 @@ import coremltools as ct import coremltools.optimize as cto from executorch.backends.apple.coreml import executorchcoreml +from executorch.backends.apple.coreml.compiler.enumerated_shape_utils import ( + _get_ct_inputs, + _SymbolicShapeToEnumeratedShapeMap, +) from executorch.backends.apple.coreml.logging import get_coreml_log_level from executorch.exir.backend.backend_details import ( BackendDetails, @@ -37,6 +41,7 @@ class COMPILE_SPEC_KEYS(Enum): MIN_DEPLOYMENT_TARGET = "min_deployment_target" MODEL_COMPUTE_PRECISION = "model_compute_precision" OP_LINEAR_QUANTIZER_CONFIG = "op_linear_quantizer_config" + ENUMERATED_SHAPES = "enumerated_shapes" class MODEL_PATHS(Enum): @@ -143,7 +148,7 @@ def generate_minimum_deployment_target_compile_spec( @staticmethod def min_deployment_target_from_compile_specs( compile_specs: List[CompileSpec], - ) -> ct.target: + ) -> Optional[ct.target]: """ Returns the minimum deployment target by parsing the list of compile specs. """ @@ -214,6 +219,54 @@ def op_linear_quantizer_config_from_compile_specs( return None + @staticmethod + def generate_enumerated_shapes_compile_spec( + ep: ExportedProgram, + enumerated_shapes: Dict[str, List[List[int]]], + ) -> CompileSpec: + """ + Returns the compile spec representing the model enumerated shapes + enumerated_shapes is a dictionary for each input to its enumerated shapes, e.g., + + enumerated_shapes = { + {"x": [[1, 1, 24], [8, 9, 24]] + {"y": [[1, 6], [30, 6]], + ] + + means the model can handle x can be shape [1, 1, 24] or [8, 9, 24] and y can be shape [1, 6] or [30, 6]. + + Only multiple inputs can have enumerated shapes if using iOS18 or later. + In this case, each input must have the same number of enumerated shapes, and these shapes are tied together + by their order in the list. For example, the model above can handle x with shape [1, 1, 24] and y with shape [1, 6], + or x with shape [8, 9, 24] and y with shape [30, 6], but not x with shape [1, 1, 24] and y with shape [30, 6]. + + Passing incorrect shapes at runtime will result in an error. + """ + emap = _SymbolicShapeToEnumeratedShapeMap.from_exported_program( + ep, + enumerated_shapes, + ) + str_representation = emap.to_json() + byte_representation = str_representation.encode("utf-8") + return CompileSpec( + COMPILE_SPEC_KEYS.ENUMERATED_SHAPES.value, + byte_representation, + ) + + @staticmethod + def enumerated_shapes_from_compile_specs( + compile_specs: List[CompileSpec], + ) -> cto.coreml.OpLinearQuantizerConfig: + """ + Returns the model's post conversion quantization by parsing the list of compile specs. + """ + for compile_spec in compile_specs: + if compile_spec.key == COMPILE_SPEC_KEYS.ENUMERATED_SHAPES.value: + emap_json = compile_spec.value.decode("utf-8") + emap = _SymbolicShapeToEnumeratedShapeMap.from_json(emap_json) + return emap + return None + @staticmethod def generate_compile_specs( compute_unit: ct.ComputeUnit = ct.ComputeUnit.ALL, @@ -446,6 +499,28 @@ def preprocess( op_linear_quantizer_config = ( CoreMLBackend.op_linear_quantizer_config_from_compile_specs(compile_specs) ) + enumerated_shapes = CoreMLBackend.enumerated_shapes_from_compile_specs( + compile_specs + ) + + # If using enumerated shapes, we need to pass the inputs to CoreML's convert() function + # explicitly + ct_inputs = None + if enumerated_shapes is not None: + ct_inputs = _get_ct_inputs(edge_program, enumerated_shapes) + + # Check there are not multiple enumerated inputs if iOS is below 18 + if (minimum_deployment_target is None) or ( + minimum_deployment_target < ct.target.iOS18 + ): + n_enumerated_inputs = 0 + for ct_in in ct_inputs: + if isinstance(ct_in.shape, ct.EnumeratedShapes): + n_enumerated_inputs += 1 + if n_enumerated_inputs > 1: + raise ValueError( + f"You're program has {n_enumerated_inputs}, but the minimum_deployment_target is set to {minimum_deployment_target}. Multiple enumerated inputs requires iOS18 or later." + ) # Load the model if MODEL_TYPE is 'COMPILED_MODEL'. This step is necessary because # get_compiled_model_path() requires a loaded model. @@ -459,6 +534,7 @@ def preprocess( compute_precision=model_compute_precision, minimum_deployment_target=minimum_deployment_target, compute_units=compute_units, + inputs=ct_inputs, ) if op_linear_quantizer_config is not None: diff --git a/backends/apple/coreml/compiler/enumerated_shape_utils.py b/backends/apple/coreml/compiler/enumerated_shape_utils.py new file mode 100644 index 00000000000..663830f702e --- /dev/null +++ b/backends/apple/coreml/compiler/enumerated_shape_utils.py @@ -0,0 +1,244 @@ +import json +from dataclasses import asdict, dataclass +from typing import Optional, Tuple + +import coremltools as ct +import torch +from coremltools.converters.mil.frontend.torch.utils import TORCH_DTYPE_TO_MIL_DTYPE + +_IGNORE_RANGE_CONSTRAINTS: bool = False + + +@dataclass(frozen=True, slots=True) +class _SymInt: + key_name: str + low: Optional[int] + high: Optional[int] + + @classmethod + def from_symint_and_range_constraints(cls, s: torch.SymInt, range_constraints=None): + # Canonicalize: "Sym(s0)" -> "s0", or leave "s0" as is + def _symkey(sym: torch.SymInt) -> str: + s = str(sym) + return s[4:-1] if s.startswith("Sym(") and s.endswith(")") else s + + # Convert symint to int. Infinity is converted to None + def _as_int_or_none(b): + if b is None: + return None + s = str(b) + if s in {"int_oo", "-int_oo", "oo", "-oo", "Infinity", "-Infinity"}: + return None + return int(s) + + # Get low/high from range_constraints if provided + low, high = None, None + if range_constraints is not None: + for k, v in range_constraints.items(): + if _symkey(k) == _symkey(s): + low = _as_int_or_none(getattr(v, "lower", getattr(v, "min", None))) + high = _as_int_or_none(getattr(v, "upper", getattr(v, "max", None))) + return _SymInt(_symkey(s), low, high) + + +@dataclass(frozen=True, slots=True) +class _SymbolicShape: + shape: Tuple[int | _SymInt] + + @classmethod + def from_shape_and_range_constraints(cls, shape, range_constraints=None): + out_shape = [] + for s in shape: + if isinstance(s, int): + assert s >= 0 + out_shape.append(s) + elif isinstance(s, torch.SymInt): + out_shape.append( + _SymInt.from_symint_and_range_constraints(s, range_constraints) + ) + else: + raise ValueError(f"Unexpected type found in shape: {type(s)}") + out_shape = tuple(out_shape) + return _SymbolicShape(out_shape) + + def is_static_shape(self): + for s in self.shape: + if isinstance(s, _SymInt): + return False + return True + + def __len__(self): + return len(self.shape) + + def __getitem__(self, key): + return self.shape[key] + + def to_dict(self): + return asdict(self) + + @classmethod + def from_dict(cls, d): + assert len(d) == 1 and "shape" in d + shape = [] + for s in d["shape"]: + if isinstance(s, int): + shape.append(s) + elif isinstance(s, dict): + assert len(s) == 3 and "key_name" in s and "low" in s and "high" in s + shape.append(_SymInt(**s)) + else: + raise ValueError(f"Unexpected type found in shape: {type(s)}") + shape = tuple(shape) + return _SymbolicShape(shape) + + +def _iterate_over_fake_user_inputs(ep): + user_inputs = ep.graph_signature.user_inputs + for node in ep.graph.nodes: + if node.op == "placeholder" and node.name in user_inputs: + yield (node.name, node.meta["val"]) + + +def _create_enumeration_map(ep, enumerated_shapes, *, ignore_range_constraints=False): + # Each input should have the same number of enumerations + assert len(enumerated_shapes) > 0, "No enumerated shapes provided" + num_enumerations = None + for name, eshapes in enumerated_shapes.items(): + if num_enumerations is None: + num_enumerations = len(eshapes) + else: + assert ( + len(eshapes) > 1 + ), f"Input {name} only has {len(eshapes)} enumerated shapes provided. You should not specify enumerated shapes for inputs with only 1 input." + assert ( + len(eshapes) == num_enumerations + ), f"Input {name} has {len(eshapes)} enumerated shape provided, but other inputs have {num_enumerations} enumerated shapes" + + symbolic_shape_to_enumerations = {} + for name, fake_input in _iterate_over_fake_user_inputs(ep): + shape = fake_input.shape + serialized_shape = _SymbolicShape.from_shape_and_range_constraints( + shape, ep.range_constraints if not ignore_range_constraints else None + ) + if serialized_shape.is_static_shape(): + continue + # Shape is dynamic + if name not in enumerated_shapes: + raise ValueError( + f"The input {name} has a symbolic shape, but you did not provide an enumeration for it" + ) + # Validate + for eshape in enumerated_shapes[name]: + assert len(serialized_shape) == len( + eshape + ), f"In {name}, the rank of the enumeration is {len(eshape)}, but the symbolic shape has rank {len(serialized_shape)}" + for i in range(len(eshape)): + assert isinstance( + eshape[i], int + ), f"Enumerated shapes must be ints, but got {type(eshape[i])}." + assert eshape[i] >= 1, "Each enumerated shape dimension must be >= 1" + if isinstance(serialized_shape[i], int): + assert ( + serialized_shape[i] == eshape[i] + ), f"In {name}, the shape enumeration {eshape} does not match {shape} on the non-symbolic value at index {i}" + else: + # Check eshape is within bound + if serialized_shape[i].low is not None: + # We add special case for when the low bound is 2. This is because Torch does not usually allow 1 as a lower bound + assert (eshape[i] >= serialized_shape[i].low) or ( + eshape[i] == 1 and serialized_shape[i].low == 2 + ), f"In {name}, the shape enumeration {eshape} violates the lower range-constraint on the symbolic shape {shape} at index {i}" + if serialized_shape[i].high is not None: + assert ( + eshape[i] <= serialized_shape[i].high + ), f"In {name}, the shape enumeration {eshape} violates the upper range-constraint on the symbolic shape {shape} at index {i}" + if serialized_shape in symbolic_shape_to_enumerations: + enumerations, names = symbolic_shape_to_enumerations[serialized_shape] + assert ( + enumerations == enumerated_shapes[name] + ), f"The symbolic shape {shape}, has multiple enumerations defined. A new enumeration is defined for input {name}, but the existing inputs {names} have a different one defined. If these inputs have different enumerations, they should be exported with different symbolic shapes." + names.append(name) + symbolic_shape_to_enumerations[serialized_shape] = (enumerations, names) + else: + symbolic_shape_to_enumerations[serialized_shape] = ( + enumerated_shapes[name], + [name], + ) + return symbolic_shape_to_enumerations + + +class _SymbolicShapeToEnumeratedShapeMap: + def __init__(self, emap): + self.emap = emap + + def to_json(self): + json_list = [] + for k in self.emap: + json_list.append((k.to_dict(), self.emap[k])) + return json.dumps(json_list) + + @classmethod + def from_json(cls, s): + emap = {} + json_list = json.loads(s) + for k, v in json_list: + k = _SymbolicShape.from_dict(k) + emap[k] = tuple(v) + return cls(emap) + + @classmethod + def from_exported_program( + cls, + ep, + enumerated_shapes, + *, + ignore_range_constraints=_IGNORE_RANGE_CONSTRAINTS, + ): + emap = _create_enumeration_map( + ep, enumerated_shapes, ignore_range_constraints=ignore_range_constraints + ) + return cls(emap) + + def __getitem__(self, key: _SymbolicShape): + return self.emap[key][0] + + def __contains__(self, key): + return key in self.emap + + def __repr__(self): + return f"_SymbolicShapeToEnumeratedShapeMap(emap={self.emap})" + + +def _get_ct_inputs(ep, emap: _SymbolicShapeToEnumeratedShapeMap): + ct_inputs = [] + for name, fake_input in _iterate_over_fake_user_inputs(ep): + + # CoreML can do funny conversions in ct.convert (e.g., int64 -> int32, int16 -> int32), so here + # we restrict users to use dtypes we know are supported + _ENUMERATED_SHAPE_INPUT_DTYPES = [torch.float16, torch.float32, torch.int32] + for dtype in _ENUMERATED_SHAPE_INPUT_DTYPES: + assert dtype in TORCH_DTYPE_TO_MIL_DTYPE + assert ( + fake_input.dtype in _ENUMERATED_SHAPE_INPUT_DTYPES + ), f"When using enumerated shapes, all inputs must have one of the following dtyeps {_ENUMERATED_SHAPE_INPUT_DTYPES}, but {name} has dtype {fake_input.dtype}" + + ct_dtype = TORCH_DTYPE_TO_MIL_DTYPE[fake_input.dtype] + shape = fake_input.shape + serializable_shape = _SymbolicShape.from_shape_and_range_constraints( + shape, ep.range_constraints if not _IGNORE_RANGE_CONSTRAINTS else None + ) + if serializable_shape.is_static_shape(): + ct_inputs.append( + ct.TensorType(name=name, shape=serializable_shape.shape, dtype=ct_dtype) + ) + continue + # Dynamic shape + assert ( + serializable_shape in emap + ), f"The shape of input {name} ({serializable_shape}) is not in the _SymbolicShapeToEnumeratedShapeMap={emap}" + enumerations = emap[serializable_shape] + ct_enumerated_shape = ct.EnumeratedShapes(shapes=enumerations) + ct_inputs.append( + ct.TensorType(name=name, shape=ct_enumerated_shape, dtype=ct_dtype) + ) + return ct_inputs diff --git a/backends/apple/coreml/compiler/torch_ops.py b/backends/apple/coreml/compiler/torch_ops.py index 11294a69a3d..53ac436fe38 100644 --- a/backends/apple/coreml/compiler/torch_ops.py +++ b/backends/apple/coreml/compiler/torch_ops.py @@ -8,12 +8,14 @@ # coremltools than is used by ExecuTorch. Each op registered here should have a link to a PR in coremltools that adds # the op to the coremltools library. +import numpy as np import torch as _torch from coremltools import _logger from coremltools.converters.mil.frontend import _utils from coremltools.converters.mil.frontend.torch.ops import ( _get_inputs, _get_kwinputs, + noop, NUM_TO_NUMPY_DTYPE, NUM_TO_TORCH_DTYPE, split, @@ -21,7 +23,6 @@ transpose, unbind, ) - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( register_torch_op, ) @@ -47,24 +48,70 @@ def split_copy(context, node): split(context, node) +def is_fbcode(): + return not hasattr(_torch.version, "git_version") + + +if not is_fbcode(): + from coremltools.converters.mil.frontend.torch.dim_order_ops import ( + _empty_dim_order, + _to_dim_order_copy, + ) + + # This is a temporary hack to register the alias "dim_order_ops._to_dim_order_copy", + # which was missed by coremltools + @register_torch_op(torch_alias=["dim_order_ops._to_dim_order_copy"], override=False) + def _to_dim_order_copy_TMP_EXECUTORCH_ALIAS_HACK(context, node): + _to_dim_order_copy(context, node) + + # This is a temporary hack to register the alias "dim_order_ops._empty_dim_order", + # which was missed by coremltools + @register_torch_op(torch_alias=["dim_order_ops._empty_dim_order"], override=False) + def _empty_dim_order_TMP_EXECUTORCH_ALIAS_HACK(context, node): + _empty_dim_order(context, node) + +else: + # TODO: remove this case when fbcode updates to coremltools 9.0 + @register_torch_op( + torch_alias=[ + "dim_order_ops::_to_dim_order_copy", + "dim_order_ops._to_dim_order_copy", + ], + override=False, + ) + def _to_dim_order_copy(context, node): + dim_order = _get_kwinputs(context, node, "dim_order", default=[None])[0] + node.kwinputs.pop("dim_order") + + # In CoreML, dim_order.val will be an ndarray, so we convert it to a list + dim_order = [int(d) for d in dim_order.val] + memory_format = get_memory_format(dim_order) + assert ( + memory_format == _torch.contiguous_format + ), "Only contiguous memory format is supported in CoreML" + to(context, node) + + @register_torch_op( torch_alias=[ - "dim_order_ops::_to_dim_order_copy", - "dim_order_ops._to_dim_order_copy", + "dim_order_ops::_clone_dim_order", + "dim_order_ops._clone_dim_order", ], override=False, ) -def _to_dim_order_copy(context, node): +def _clone_dim_order(context, node): dim_order = _get_kwinputs(context, node, "dim_order", default=[None])[0] node.kwinputs.pop("dim_order") - # In CoreML, dim_order.val will be an ndarray, so we convert it to a list + # In CoreML, dim_order.val will be a ndarray, so we convert it to a list to check memory format. dim_order = [int(d) for d in dim_order.val] memory_format = get_memory_format(dim_order) assert ( memory_format == _torch.contiguous_format ), "Only contiguous memory format is supported in CoreML" - to(context, node) + + # Since CoreML only supports contiguous format, no dim_order preservation is needed. Treat this as a no-op clone. + noop(context, node) # https://github.com/apple/coremltools/pull/2558 @@ -128,7 +175,57 @@ def dequantize_affine(context, node): int_data.astype(quantized_np_dtype), zero_point, scale, - axis=-1, + name=node.name, + ) + context.add(output, node.name) + + +@register_torch_op( + torch_alias=["quant::dequantize_codebook", "quant.dequantize_codebook"], + override=False, +) +def dequantize_codebook(context, node): + inputs = _get_inputs(context, node, expected=[4, 5]) + codes = inputs[0].val + codebook = inputs[1].val + nbits = inputs[2].val + + # information in block_size is redundant with codebook.shape + block_size = inputs[3].val # noqa: F841 + + assert len(codes.shape) == 2, "Only rank 2 inputs are supported" + + # Assert codebook is as expected. codebook.dim() = codes.dim() + 2 + assert len(codebook.shape) == 4, "Only rank 4 inputs are supported for codebook" + assert (codebook.shape[0] == 1) or ( + codebook.shape[1] == 1 + ), "Only grouped_channel granularity is supported" + if codebook.shape[0] == 1: + # LUT is per column group + n_luts = codebook.shape[1] + assert ( + codes.shape[1] % n_luts == 0 + ), "codes.shape[1] must be divisible by codebook.shape[1]" + else: + # LUT is per row group + n_luts = codebook.shape[0] + assert ( + codes.shape[0] % n_luts == 0 + ), "codes.shape[0] must be divisible by codebook.shape[0]" + + assert codebook.shape[2] == 2**nbits + assert codebook.shape[3] == 1, "Only scalar look up values are supported" + + if len(inputs) > 4: + output_dtype = inputs[4].val + out_np_dtype = NUM_TO_NUMPY_DTYPE[output_dtype] + _logger.warning( + f"Core ML ignores output_dtype {out_np_dtype} on torchao.dequantize_affine and instead uses the native precision." + ) + + output = _utils._construct_constexpr_lut_op( + codes.astype(np.int8), + codebook, name=node.name, ) context.add(output, node.name) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 93506e6d985..6b3a73599af 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -10,6 +10,9 @@ import torch from executorch.backends.apple.coreml.compiler import CoreMLBackend +from executorch.backends.apple.coreml.compiler.coreml_preprocess import ( + COMPILE_SPEC_KEYS, +) from executorch.backends.apple.coreml.logging import get_coreml_log_level from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -192,6 +195,13 @@ def __init__( if skip_ops_for_coreml_delegation is None: skip_ops_for_coreml_delegation = [] self.skip_ops_for_coreml_delegation = skip_ops_for_coreml_delegation + + for compile_spec in compile_specs or []: + if compile_spec.key == COMPILE_SPEC_KEYS.ENUMERATED_SHAPES.value: + assert ( + lower_full_graph + ), "lower_full_graph must be True in the CoreMLPartitioner when using an enumerated shape compile spec" + self.delegation_spec = DelegationSpec( backend_id=CoreMLBackend.__name__, compile_specs=compile_specs if compile_specs is not None else [], diff --git a/backends/apple/coreml/recipes/__init__.py b/backends/apple/coreml/recipes/__init__.py new file mode 100644 index 00000000000..8bcd1c254a8 --- /dev/null +++ b/backends/apple/coreml/recipes/__init__.py @@ -0,0 +1,17 @@ +# Copyright © 2025 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + + +from executorch.export import recipe_registry + +from .coreml_recipe_provider import CoreMLRecipeProvider +from .coreml_recipe_types import CoreMLRecipeType + +# Auto-register CoreML backend recipe provider +recipe_registry.register_backend_recipe_provider(CoreMLRecipeProvider()) + +__all__ = [ + "CoreMLRecipeProvider", + "CoreMLRecipeType", +] diff --git a/backends/apple/coreml/recipes/coreml_recipe_provider.py b/backends/apple/coreml/recipes/coreml_recipe_provider.py new file mode 100644 index 00000000000..ed29e57f26b --- /dev/null +++ b/backends/apple/coreml/recipes/coreml_recipe_provider.py @@ -0,0 +1,413 @@ +# Copyright © 2025 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + + +import logging +from typing import Any, Optional, Sequence + +import coremltools as ct +import torch + +from executorch.backends.apple.coreml.compiler import CoreMLBackend +from executorch.backends.apple.coreml.partition.coreml_partitioner import ( + CoreMLPartitioner, +) +from executorch.backends.apple.coreml.recipes.coreml_recipe_types import ( + COREML_BACKEND, + CoreMLRecipeType, +) + +from executorch.exir import EdgeCompileConfig +from executorch.export import ( + AOQuantizationConfig, + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) +from torchao.quantization.granularity import PerAxis, PerGroup +from torchao.quantization.quant_api import IntxWeightOnlyConfig + + +class CoreMLRecipeProvider(BackendRecipeProvider): + @property + def backend_name(self) -> str: + return COREML_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(CoreMLRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + """Create CoreML recipe with precision and compute unit combinations""" + + if recipe_type not in self.get_supported_recipes(): + return None + + if ct is None: + raise ImportError( + "coremltools is required for CoreML recipes. " + "Install it with: pip install coremltools" + ) + + # Validate kwargs + self._validate_recipe_kwargs(recipe_type, **kwargs) + + if recipe_type == CoreMLRecipeType.FP32: + return self._build_fp_recipe(recipe_type, ct.precision.FLOAT32, **kwargs) + elif recipe_type == CoreMLRecipeType.FP16: + return self._build_fp_recipe(recipe_type, ct.precision.FLOAT16, **kwargs) + elif recipe_type == CoreMLRecipeType.PT2E_INT8_STATIC: + return self._build_pt2e_quantized_recipe( + recipe_type, activation_dtype=torch.quint8, **kwargs + ) + elif recipe_type == CoreMLRecipeType.PT2E_INT8_WEIGHT_ONLY: + return self._build_pt2e_quantized_recipe( + recipe_type, activation_dtype=torch.float32, **kwargs + ) + elif recipe_type == CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL: + self._validate_and_set_deployment_target(kwargs, ct.target.iOS18, "torchao") + return self._build_torchao_quantized_recipe( + recipe_type, + weight_dtype=torch.int4, + is_per_channel=True, + **kwargs, + ) + elif recipe_type == CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP: + group_size = kwargs.pop("group_size", 32) + self._validate_and_set_deployment_target(kwargs, ct.target.iOS18, "torchao") + return self._build_torchao_quantized_recipe( + recipe_type, + weight_dtype=torch.int4, + is_per_channel=False, + group_size=group_size, + **kwargs, + ) + elif recipe_type == CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL: + self._validate_and_set_deployment_target(kwargs, ct.target.iOS16, "torchao") + return self._build_torchao_quantized_recipe( + recipe_type, weight_dtype=torch.int8, is_per_channel=True, **kwargs + ) + elif recipe_type == CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP: + group_size = kwargs.pop("group_size", 32) + # override minimum_deployment_target to ios18 for torchao (GH issue #13122) + self._validate_and_set_deployment_target(kwargs, ct.target.iOS18, "torchao") + return self._build_torchao_quantized_recipe( + recipe_type, + weight_dtype=torch.int8, + is_per_channel=False, + group_size=group_size, + **kwargs, + ) + elif recipe_type == CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY: + bits = kwargs.pop("bits") + block_size = kwargs.pop("block_size") + return self._build_codebook_quantized_recipe( + recipe_type, bits=bits, block_size=block_size, **kwargs + ) + + return None + + def _validate_recipe_kwargs(self, recipe_type: RecipeType, **kwargs: Any) -> None: + """Validate kwargs for each recipe type""" + expected_keys = self._get_expected_keys(recipe_type) + + unexpected = set(kwargs.keys()) - expected_keys + if unexpected: + logging.warning( + f"CoreML recipe '{recipe_type.value}' ignoring unexpected parameters: {list(unexpected)}. " + f"Expected parameters: {list(expected_keys)}" + ) + + self._validate_base_parameters(kwargs) + self._validate_group_size_parameter(recipe_type, kwargs) + self._validate_codebook_parameters(recipe_type, kwargs) + + def _get_expected_keys(self, recipe_type: RecipeType) -> set: + """Get expected parameter keys for a recipe type""" + common_keys = { + "minimum_deployment_target", + "compute_unit", + "skip_ops_for_coreml_delegation", + "lower_full_graph", + "take_over_constant_data", + } + + if recipe_type in [ + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, + ]: + return common_keys | {"group_size", "filter_fn"} + elif recipe_type in [ + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL, + ]: + return common_keys | {"filter_fn"} + elif recipe_type == CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY: + return common_keys | {"bits", "block_size", "filter_fn"} + else: + return common_keys + + def _validate_base_parameters(self, kwargs: Any) -> None: + """Validate minimum_deployment_target and compute_unit parameters""" + if "minimum_deployment_target" in kwargs: + minimum_deployment_target = kwargs["minimum_deployment_target"] + if not isinstance(minimum_deployment_target, ct.target): + raise ValueError( + f"Parameter 'minimum_deployment_target' must be an enum of type ct.target, got {type(minimum_deployment_target)}" + ) + + if "compute_unit" in kwargs: + compute_unit = kwargs["compute_unit"] + if not isinstance(compute_unit, ct.ComputeUnit): + raise ValueError( + f"Parameter 'compute_unit' must be an enum of type ct.ComputeUnit, got {type(compute_unit)}" + ) + + def _validate_group_size_parameter( + self, recipe_type: RecipeType, kwargs: Any + ) -> None: + """Validate group_size parameter for applicable recipe types""" + if ( + recipe_type + in [ + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, + ] + and "group_size" in kwargs + ): + group_size = kwargs["group_size"] + if not isinstance(group_size, int): + raise ValueError( + f"Parameter 'group_size' must be an integer, got {type(group_size).__name__}: {group_size}" + ) + if group_size <= 0: + raise ValueError( + f"Parameter 'group_size' must be positive, got: {group_size}" + ) + + def _validate_codebook_parameters( + self, recipe_type: RecipeType, kwargs: Any + ) -> None: + """Validate bits and block_size parameters for codebook recipe type""" + if recipe_type != CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY: + return + + # Both bits and block_size must be present + if not ("bits" in kwargs and "block_size" in kwargs): + raise ValueError( + "Parameters 'bits' and 'block_size' must be present for codebook recipes" + ) + + if "bits" in kwargs: + bits = kwargs["bits"] + if not isinstance(bits, int): + raise ValueError( + f"Parameter 'bits' must be an integer, got {type(bits).__name__}: {bits}" + ) + if not (1 <= bits <= 8): + raise ValueError( + f"Parameter 'bits' must be between 1 and 8, got: {bits}" + ) + + if "block_size" in kwargs: + block_size = kwargs["block_size"] + if not isinstance(block_size, list): + raise ValueError( + f"Parameter 'block_size' must be a list, got {type(block_size).__name__}: {block_size}" + ) + + def _validate_and_set_deployment_target( + self, kwargs: Any, min_target: ct.target, quantization_type: str + ) -> None: + """Validate or set minimum deployment target for quantization recipes""" + minimum_deployment_target = kwargs.get("minimum_deployment_target", None) + if minimum_deployment_target and minimum_deployment_target < min_target: + raise ValueError( + f"minimum_deployment_target must be {str(min_target)} or higher for {quantization_type} quantization" + ) + else: + # Default to the minimum target for this quantization type + kwargs["minimum_deployment_target"] = min_target + + def _build_fp_recipe( + self, + recipe_type: RecipeType, + precision: ct.precision, + **kwargs: Any, + ) -> ExportRecipe: + """Build FP32/FP16 recipe""" + lowering_recipe = self._get_coreml_lowering_recipe( + compute_precision=precision, + **kwargs, + ) + + return ExportRecipe( + name=recipe_type.value, + lowering_recipe=lowering_recipe, + ) + + def _build_pt2e_quantized_recipe( + self, + recipe_type: RecipeType, + activation_dtype: torch.dtype, + **kwargs: Any, + ) -> ExportRecipe: + """Build PT2E-based quantization recipe""" + from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer + + self._validate_and_set_deployment_target(kwargs, ct.target.iOS17, "pt2e") + + # Validate activation_dtype + assert activation_dtype in [ + torch.quint8, + torch.float32, + ], f"activation_dtype must be torch.quint8 or torch.float32, got {activation_dtype}" + + # Create quantization config + config = ct.optimize.torch.quantization.LinearQuantizerConfig( + global_config=ct.optimize.torch.quantization.ModuleLinearQuantizerConfig( + quantization_scheme="symmetric", + activation_dtype=activation_dtype, + weight_dtype=torch.qint8, + weight_per_channel=True, + ) + ) + + quantizer = CoreMLQuantizer(config) + quantization_recipe = QuantizationRecipe(quantizers=[quantizer]) + + lowering_recipe = self._get_coreml_lowering_recipe(**kwargs) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + ) + + def _build_torchao_quantized_recipe( + self, + recipe_type: RecipeType, + weight_dtype: torch.dtype, + is_per_channel: bool, + group_size: int = 32, + **kwargs: Any, + ) -> ExportRecipe: + """Build TorchAO-based quantization recipe""" + if is_per_channel: + weight_granularity = PerAxis(axis=0) + else: + weight_granularity = PerGroup(group_size=group_size) + + # Use user-provided filter_fn if provided + filter_fn = kwargs.get("filter_fn", None) + config = AOQuantizationConfig( + ao_base_config=IntxWeightOnlyConfig( + weight_dtype=weight_dtype, + granularity=weight_granularity, + ), + filter_fn=filter_fn, + ) + + quantization_recipe = QuantizationRecipe( + quantizers=None, + ao_quantization_configs=[config], + ) + + lowering_recipe = self._get_coreml_lowering_recipe(**kwargs) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + ) + + def _build_codebook_quantized_recipe( + self, + recipe_type: RecipeType, + bits: int, + block_size: list, + **kwargs: Any, + ) -> ExportRecipe: + """Build codebook/palettization quantization recipe""" + from torchao.prototype.quantization.codebook_coreml import ( + CodebookWeightOnlyConfig, + ) + + self._validate_and_set_deployment_target(kwargs, ct.target.iOS18, "codebook") + + # Get the appropriate dtype (torch.uint1 through torch.uint8) + dtype = getattr(torch, f"uint{bits}") + + # Use user-provided filter_fn or default to Linear/Embedding layers + filter_fn = kwargs.get( + "filter_fn", + lambda m, fqn: ( + isinstance(m, torch.nn.Embedding) or isinstance(m, torch.nn.Linear) + ), + ) + + config = AOQuantizationConfig( + ao_base_config=CodebookWeightOnlyConfig( + dtype=dtype, + block_size=block_size, + ), + filter_fn=filter_fn, + ) + + quantization_recipe = QuantizationRecipe( + quantizers=None, + ao_quantization_configs=[config], + ) + + lowering_recipe = self._get_coreml_lowering_recipe(**kwargs) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + ) + + def _get_coreml_lowering_recipe( + self, + compute_precision: ct.precision = ct.precision.FLOAT16, + **kwargs: Any, + ) -> LoweringRecipe: + """Get CoreML lowering recipe with optional precision""" + compile_specs = CoreMLBackend.generate_compile_specs( + compute_precision=compute_precision, + compute_unit=kwargs.get("compute_unit", ct.ComputeUnit.ALL), + minimum_deployment_target=kwargs.get("minimum_deployment_target", None), + ) + + minimum_deployment_target = kwargs.get("minimum_deployment_target", None) + take_over_mutable_buffer = True + if minimum_deployment_target and minimum_deployment_target < ct.target.iOS18: + take_over_mutable_buffer = False + + # Extract additional partitioner parameters + skip_ops_for_coreml_delegation = kwargs.get( + "skip_ops_for_coreml_delegation", None + ) + lower_full_graph = kwargs.get("lower_full_graph", False) + take_over_constant_data = kwargs.get("take_over_constant_data", True) + + partitioner = CoreMLPartitioner( + compile_specs=compile_specs, + take_over_mutable_buffer=take_over_mutable_buffer, + skip_ops_for_coreml_delegation=skip_ops_for_coreml_delegation, + lower_full_graph=lower_full_graph, + take_over_constant_data=take_over_constant_data, + ) + + edge_compile_config = EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=False, + ) + + return LoweringRecipe( + partitioners=[partitioner], edge_compile_config=edge_compile_config + ) diff --git a/backends/apple/coreml/recipes/coreml_recipe_types.py b/backends/apple/coreml/recipes/coreml_recipe_types.py new file mode 100644 index 00000000000..fc7292c3c58 --- /dev/null +++ b/backends/apple/coreml/recipes/coreml_recipe_types.py @@ -0,0 +1,53 @@ +# Copyright © 2025 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + + +from executorch.export import RecipeType + + +COREML_BACKEND: str = "coreml" + + +class CoreMLRecipeType(RecipeType): + """CoreML-specific generic recipe types""" + + ## All the recipes accept common kwargs + # 1. minimum_deployment_unit (default: None) + # 2. compute_unit (default: ct.ComputeUnit.ALL) + + # FP32 precision recipe, defaults to values published by the CoreML backend and partitioner + FP32 = "coreml_fp32" + + # FP16 precision recipe, defaults to values published by the CoreML backend and partitioner + FP16 = "coreml_fp16" + + ## PT2E-based quantization recipes + # INT8 Static Quantization (weights + activations), requires calibration dataset + PT2E_INT8_STATIC = "coreml_pt2e_int8_static" + # INT8 Weight-only Quantization (activations remain FP32) + PT2E_INT8_WEIGHT_ONLY = "coreml_pt2e_int8_weight_only" + + ## TorchAO-based quantization recipes + # All TorchAO recipes accept filter_fn kwarg to control which layers are quantized + # INT4 Weight-only Quantization, per-channel (axis=0) + # Additional kwargs: filter_fn (default: Embedding and linear layers) + TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL = "coreml_torchao_int4_weight_only_per_channel" + # INT4 Weight-only Quantization, per-group + # Additional kwargs: group_size (default: 32), filter_fn (default: Embedding and linear layers) + TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP = "coreml_torchao_int4_weight_only_per_group" + # INT8 Weight-only Quantization, per-channel (axis=0) + # Additional kwargs: filter_fn (default: Embedding and linear layers) + TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL = "coreml_torchao_int8_weight_only_per_channel" + # INT8 Weight-only Quantization, per-group + # Additional kwargs: group_size (default: 32), filter_fn (default: Embedding and linear layers) + TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP = "coreml_torchao_int8_weight_only_per_group" + + ## Codebook/Palettization Quantization + # Additional mandatory kwargs: bits (range: 1-8), block_size (list of ints), + # filter_fn (default: targets Linear and Embedding layers) + CODEBOOK_WEIGHT_ONLY = "coreml_codebook_weight_only" + + @classmethod + def get_backend_name(cls) -> str: + return COREML_BACKEND diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h index 11d957044e9..a9e06efa90d 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.h @@ -99,6 +99,17 @@ NS_ASSUME_NONNULL_BEGIN - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError* __autoreleasing*)error; +/// Executes a block with a unique temporary directory. +/// +/// A new temporary subdirectory URL is created inside the receiver’s designated +/// base directory. The directory is passed to the block, which can use it to +/// perform temporary file operations. After the block finishes executing, +/// the directory and its contents are removed. +/// +/// @param block A block to execute. The block receives a unique URL. +- (void)withTemporaryDirectory:(void (^)(NSURL* directoryURL))block; + + /// Purges the assets storage. The assets are moved to the trash directory and are asynchronously /// deleted. /// @@ -117,6 +128,12 @@ NS_ASSUME_NONNULL_BEGIN /// contents are deleted asynchronously. @property (copy, readonly, nonatomic) NSURL* trashDirectoryURL; + +/// The staging directory URL, used to hold assets that are being prepared or processed +/// before they are moved into their final location. The contents of this directory +/// are temporary and may be cleared when no longer needed. +@property (copy, readonly, nonatomic) NSURL* stagingDirectoryURL; + /// The file manager. @property (strong, readonly, nonatomic) NSFileManager* fileManager; diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm index 256026e1f09..53c3d1cdc69 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLAssetManager.mm @@ -254,6 +254,29 @@ BOOL is_asset_alive(NSMapTable *assets_in_use_map, return assets; } + +NSURL * _Nullable move_to_directory(NSURL *url, + NSURL *directoryURL, + NSFileManager *fileManager, + NSError * __autoreleasing *error) { + if (!url) { + ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: source URL is nil."); + return nil; + } + + if (!directoryURL) { + ETCoreMLLogErrorAndSetNSError(error, ETCoreMLErrorInternalError, "Move operation failed: destination URL is nil."); + return nil; + } + + NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + if (![fileManager moveItemAtURL:url toURL:dstURL error:error]) { + return nil; + } + + return dstURL; +} + } //namespace @interface ETCoreMLAssetManager () { @@ -299,12 +322,17 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr&)data if (!managedAssetsDirectoryURL) { return nil; } - + NSURL *managedTrashDirectoryURL = ::create_directory_if_needed(trashDirectoryURL, @"models", fileManager, error); if (!managedTrashDirectoryURL) { return nil; } - + + NSURL *managedStagingDirectoryURL = ::create_directory_if_needed(assetsDirectoryURL, @"staging", fileManager, error); + if (!managedStagingDirectoryURL) { + return nil; + } + // If directory is empty then purge the stores if (::is_directory_empty(managedAssetsDirectoryURL, fileManager, nil)) { assetsMetaStore.impl()->purge(ec); @@ -315,6 +343,7 @@ - (nullable instancetype)initWithDatabase:(const std::shared_ptr&)data _assetsStore = std::move(assetsStore); _assetsMetaStore = std::move(assetsMetaStore); _assetsDirectoryURL = managedAssetsDirectoryURL; + _stagingDirectoryURL = managedStagingDirectoryURL; _trashDirectoryURL = managedTrashDirectoryURL; _estimatedSizeInBytes = sizeInBytes.value(); _maxAssetsSizeInBytes = maxAssetsSizeInBytes; @@ -346,15 +375,15 @@ - (nullable instancetype)initWithDatabaseURL:(NSURL *)databaseURL error:error]; } -- (nullable NSURL *)moveURL:(NSURL *)url - toUniqueURLInDirectory:(NSURL *)directoryURL - error:(NSError * __autoreleasing *)error { - NSURL *dstURL = [directoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - if (![self.fileManager moveItemAtURL:url toURL:dstURL error:error]) { - return nil; +- (void)withTemporaryDirectory:(void (^)(NSURL *directoryURL))block { + NSURL *dstURL = [self.stagingDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; + block(dstURL); + if (![self.fileManager fileExistsAtPath:dstURL.path]) { + return; } - - return dstURL; + + move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil); + [self cleanupTrashDirectory]; } - (void)cleanupAssetIfNeeded:(ETCoreMLAsset *)asset { @@ -407,9 +436,8 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL return false; } - // If an asset exists move it - [self moveURL:dstURL toUniqueURLInDirectory:self.trashDirectoryURL error:nil]; - + // If a file already exists at `dstURL`, move it to the trash for removal. + move_to_directory(dstURL, self.trashDirectoryURL, self.fileManager, nil); // Move the asset to assets directory. if (![self.fileManager moveItemAtURL:srcURL toURL:dstURL error:error]) { return false; @@ -433,16 +461,25 @@ - (nullable ETCoreMLAsset *)_storeAssetAtURL:(NSURL *)srcURL } - (void)triggerCompaction { - if (self.estimatedSizeInBytes < self.maxAssetsSizeInBytes) { - return; + if (self.estimatedSizeInBytes >= self.maxAssetsSizeInBytes) { + __weak __typeof(self) weakSelf = self; + dispatch_async(self.syncQueue, ^{ + NSError *localError = nil; + if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) { + ETCoreMLLogError(localError, "Failed to compact asset store."); + } + }); } - + + // Always clean the trash directory to ensure a minimal footprint. + // The `trashQueue` is serialized, so only one cleanup will run at a time. + [self cleanupTrashDirectory]; +} + +- (void)cleanupTrashDirectory { __weak __typeof(self) weakSelf = self; - dispatch_async(self.syncQueue, ^{ - NSError *localError = nil; - if (![weakSelf _compact:self.maxAssetsSizeInBytes error:&localError]) { - ETCoreMLLogError(localError, "Failed to compact asset store."); - } + dispatch_async(self.trashQueue, ^{ + [weakSelf removeFilesInTrashDirectory]; }); } @@ -548,7 +585,7 @@ - (BOOL)_removeAssetWithIdentifier:(NSString *)identifier NSURL *assetURL = ::get_asset_url(assetValue); if ([self.fileManager fileExistsAtPath:assetURL.path] && - ![self moveURL:assetURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { + !move_to_directory(assetURL, self.trashDirectoryURL, self.fileManager, error)) { return false; } @@ -649,13 +686,7 @@ - (NSUInteger)_compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing identifier); } } - - // Trigger cleanup. - __weak __typeof(self) weakSelf = self; - dispatch_async(self.trashQueue, ^{ - [weakSelf removeFilesInTrashDirectory]; - }); - + return _estimatedSizeInBytes; } @@ -664,7 +695,10 @@ - (NSUInteger)compact:(NSUInteger)sizeInBytes error:(NSError * __autoreleasing * dispatch_sync(self.syncQueue, ^{ result = [self _compact:sizeInBytes error:error]; }); - + + // Always clean the trash directory to ensure a minimal footprint. + // The `trashQueue` is serialized, so only one cleanup will run at a time. + [self cleanupTrashDirectory]; return result; } @@ -708,7 +742,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { } // Move the the whole assets directory to the temp directory. - if (![self moveURL:self.assetsDirectoryURL toUniqueURLInDirectory:self.trashDirectoryURL error:error]) { + if (!move_to_directory(self.assetsDirectoryURL, self.trashDirectoryURL, self.fileManager, error)) { return false; } @@ -724,13 +758,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error { ::set_error_from_error_code(ec, error); // Trigger cleanup - if (status) { - __weak __typeof(self) weakSelf = self; - dispatch_async(self.trashQueue, ^{ - [weakSelf removeFilesInTrashDirectory]; - }); - } - + [self cleanupTrashDirectory]; return static_cast(status); } diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm index 05aa910d954..9e8ae04842e 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm @@ -62,21 +62,12 @@ + (nullable ETCoreMLModel *)loadModelWithContentsOfURL:(NSURL *)compiledModelURL if (model) { return model; } - - if (localError) { - ETCoreMLLogError(localError, - "Failed to load model from compiled asset with identifier = %@", - identifier); - } - - // If store failed then we will load the model from compiledURL. - auto backingAsset = Asset::make(compiledModelURL, identifier, assetManager.fileManager, error); - if (!backingAsset) { - return nil; + + if (error) { + *error = localError; } - - asset = [[ETCoreMLAsset alloc] initWithBackingAsset:backingAsset.value()]; - return ::get_model_from_asset(asset, configuration, metadata, error); + + return nil; } @end diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm index f4cfd2146ac..524ceaf7e28 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm @@ -345,6 +345,10 @@ void add_compute_unit(std::string& identifier, MLComputeUnits compute_units) { return [ETCoreMLModelDebugInfo modelDebugInfoFromData:file_data error:error]; } +NSString *raw_model_identifier(NSString *identifier) { + return [NSString stringWithFormat:@"raw_%@", identifier]; +} + #endif } //namespace @@ -408,7 +412,7 @@ - (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier { return modelAsset; } - NSError *localError = nil; + __block NSError *localError = nil; modelAsset = [self.assetManager assetWithIdentifier:identifier error:&localError]; if (localError) { ETCoreMLLogError(localError, @@ -420,8 +424,9 @@ - (nullable ETCoreMLAsset *)assetWithIdentifier:(NSString *)identifier { } - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier + modelURL:(nullable NSURL *)modelURL inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - assetManager:(ETCoreMLAssetManager *)assetManager + dstURL:(NSURL *)dstURL error:(NSError * __autoreleasing *)error { auto modelAssetType = get_model_asset_type(inMemoryFS); if (!modelAssetType) { @@ -430,78 +435,135 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier "AOT blob is missing model file."); return nil; } - - NSURL *dstURL = [self.assetManager.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - NSURL *modelURL = ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error); + + // If modelURL is not provided, write model files to the destination directory (dstURL) + // and obtain a URL pointing to them. Otherwise, use the provided modelURL. + modelURL = (modelURL == nil) ? ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error) : modelURL; + if (!modelURL) { + // Failed to generate or locate model files, return nil. + return nil; + } + + // Handle based on the type of the model asset. switch (modelAssetType.value()) { case ModelAssetType::CompiledModel: { - // Model is already compiled. + // The model is already compiled; no further action needed. + // Return the existing model URL. + ETCoreMLLogInfo("The model in the pte file is pre-compiled. Skipping compilation."); return modelURL; } - + case ModelAssetType::Model: { - // Compile the model. + // The model is not compiled yet. + // Compile the model at the specified URL with a maximum wait time of 5 minutes. + ETCoreMLLogInfo("The model in the pte file is not pre-compiled. Compiling with a 5 min timeout."); NSURL *compiledModelURL = [ETCoreMLModelCompiler compileModelAtURL:modelURL maxWaitTimeInSeconds:(5 * 60) error:error]; - + // Return the URL of the compiled model or nil if compilation fails. return compiledModelURL; } } } -#if ET_EVENT_TRACER_ENABLED -- (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata - inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS - configuration:(MLModelConfiguration *)configuration - error:(NSError * __autoreleasing *)error { +- (nullable ETCoreMLAsset *)compiledModelAssetWithMetadata:(const ModelMetadata&)metadata + modelURL:(nullable NSURL *)modelURL + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + error:(NSError * __autoreleasing *)error { NSString *identifier = @(metadata.identifier.c_str()); - // Otherwise try to retrieve the compiled asset. - ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; + __block ETCoreMLAsset *compiledModelAsset = [self assetWithIdentifier:identifier]; if (compiledModelAsset) { - ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); + ETCoreMLLogInfo("Cache Hit: Successfully retrieved compiled model with identifier=%@ from the models cache.", identifier); } else { - ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); + ETCoreMLLogInfo("Cache Miss: Compiled Model with identifier=%@ was not found in the models cache.", identifier); } - - // Create a unique directory for writing model files. - NSURL *dstURL = [self.assetManager.trashDirectoryURL URLByAppendingPathComponent:[NSUUID UUID].UUIDString]; - auto modelAssetType = get_model_asset_type(inMemoryFS); - ETCoreMLAsset *modelAsset = nil; - // Write the model files. - if (modelAssetType == ModelAssetType::Model) { - NSURL *modelURL = ::write_model_files(dstURL, self.fileManager, identifier, modelAssetType.value(), inMemoryFS, error); - if (modelURL) { - modelAsset = make_asset(modelURL, - identifier, - self.fileManager, - error); + + [self.assetManager withTemporaryDirectory:^(NSURL * _Nonnull directoryURL) { + if (compiledModelAsset) { + return; } - } - - if (!compiledModelAsset) { - // Compile the model. + + // The directory specified by `directoryURL` is unique and will be automatically cleaned up + // once the enclosing block completes. NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier + modelURL:modelURL inMemoryFS:inMemoryFS - assetManager:self.assetManager + dstURL:directoryURL error:error]; - compiledModelAsset = make_asset(compiledModelURL, - identifier, - self.fileManager, - error); - } - - if (!compiledModelAsset) { - return nil; + if (compiledModelURL) { + // Move the compiled model to the asset manager to transfer ownership. + ETCoreMLLogInfo("Storing compiled asset with identifier=%@ in the asset manager.", identifier); + compiledModelAsset = [self.assetManager storeAssetAtURL:compiledModelURL withIdentifier:identifier error:error]; + } + }]; + + return compiledModelAsset; +} + +#if ET_EVENT_TRACER_ENABLED +- (nullable ETCoreMLAsset *)modelAssetWithMetadata:(const ModelMetadata&)metadata + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + error:(NSError * __autoreleasing *)error { + NSString *identifier = @(metadata.identifier.c_str()); + NSString *rawIdentifier = raw_model_identifier(identifier); + __block ETCoreMLAsset *modelAsset = [self assetWithIdentifier:rawIdentifier]; + if (modelAsset) { + ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); + } else { + ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); } - + + [self.assetManager withTemporaryDirectory:^(NSURL * _Nonnull directoryURL) { + if (modelAsset) { + return; + } + + auto modelAssetType = get_model_asset_type(inMemoryFS); + if (modelAssetType != ModelAssetType::Model) { + return; + } + + // The directory specified by `directoryURL` is unique and will be automatically cleaned up + // once the enclosing block completes. + NSURL *modelURL = ::write_model_files(directoryURL, + self.fileManager, + identifier, + modelAssetType.value(), + inMemoryFS, + error); + if (modelURL) { + // Move the model to the asset manager to transfer ownership. + modelAsset = [self.assetManager storeAssetAtURL:modelURL withIdentifier:rawIdentifier error:error]; + } + }]; + + return modelAsset; +} + +- (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata + inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS + configuration:(MLModelConfiguration *)configuration + error:(NSError * __autoreleasing *)error { NSError *localError = nil; - ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, &localError); + ETCoreMLAsset *modelAsset = [self modelAssetWithMetadata:metadata inMemoryFS:inMemoryFS error:&localError]; if (localError) { - ETCoreMLLogError(localError, "Failed to parse debug info file"); + if (error) { + *error = localError; + } + + return nil; + } + + ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata + modelURL:modelAsset.contentURL + inMemoryFS:inMemoryFS + error:error]; + if (!compiledModelAsset) { + return nil; } - + ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, error); + // The analyzer requires both the raw (uncompiled) asset and the compiled model asset to perform analysis. return [[ETCoreMLModelAnalyzer alloc] initWithCompiledModelAsset:compiledModelAsset modelAsset:modelAsset modelDebugInfo:debug_info @@ -510,41 +572,33 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier assetManager:self.assetManager error:error]; } - #else - (nullable id)modelExecutorWithMetadata:(const ModelMetadata&)metadata inMemoryFS:(const inmemoryfs::InMemoryFileSystem*)inMemoryFS configuration:(MLModelConfiguration *)configuration error:(NSError * __autoreleasing *)error { - NSString *identifier = @(metadata.identifier.c_str()); - // Otherwise try to retrieve the compiled asset. - ETCoreMLAsset *asset = [self assetWithIdentifier:identifier]; - ETCoreMLModel *model = asset ? get_model_from_asset(asset, configuration, metadata, error) : nil; - if (model) { - ETCoreMLLogInfo("Cache Hit: Successfully retrieved model with identifier=%@ from the models cache.", identifier); - return [[ETCoreMLDefaultModelExecutor alloc] initWithModel:model]; + ETCoreMLAsset *compiledModelAsset = [self compiledModelAssetWithMetadata:metadata + modelURL:nil + inMemoryFS:inMemoryFS + error:error]; + if (!compiledModelAsset) { + return nil; } - - ETCoreMLLogInfo("Cache Miss: Model with identifier=%@ was not found in the models cache.", identifier); - // Compile the model. - NSURL *compiledModelURL = [self compiledModelURLWithIdentifier:identifier - inMemoryFS:inMemoryFS - assetManager:self.assetManager - error:error]; - if (!compiledModelURL) { + + ETCoreMLModel *model = [ETCoreMLModelLoader loadModelWithContentsOfURL:compiledModelAsset.contentURL + configuration:configuration + metadata:metadata + assetManager:self.assetManager + error:error]; + if (!model) { return nil; } - - model = [ETCoreMLModelLoader loadModelWithContentsOfURL:compiledModelURL - configuration:configuration - metadata:metadata - assetManager:self.assetManager - error:error]; - + return [[ETCoreMLDefaultModelExecutor alloc] initWithModel:model]; } #endif + - (nullable id)_modelExecutorWithAOTData:(NSData *)data configuration:(MLModelConfiguration *)configuration error:(NSError * __autoreleasing *)error { @@ -729,6 +783,7 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle args.count); return result; } + NSError *localError = nil; @autoreleasepool { NSArray *inputs = [args subarrayWithRange:NSMakeRange(0, model.orderedInputNames.count)]; @@ -748,11 +803,11 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle result = YES; } } - if (!result) { - if (error) { - *error = localError; - } + + if (localError && error) { + *error = localError; } + return result; } diff --git a/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm b/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm index 9a0b4facc89..04a95e8a5a3 100644 --- a/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm +++ b/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm @@ -46,6 +46,7 @@ using executorch::runtime::get_backend_class; using executorch::runtime::Result; using executorch::aten::SizesType; +using executorch::runtime::Span; using executorch::aten::Tensor; using executorch::runtime::kTensorDimensionLimit; @@ -88,17 +89,17 @@ ET_LOG(Error, "%s: DataType=%d is not supported", ETCoreMLStrings.delegateIdentifier.UTF8String, (int)tensor.scalar_type()); return std::nullopt; } - + std::vector strides(tensor.strides().begin(), tensor.strides().end()); std::vector shape(tensor.sizes().begin(), tensor.sizes().end()); - + // If tensor is rank 0, wrap in rank 1 // See https://github.com/apple/coremltools/blob/8.2/coremltools/converters/mil/frontend/torch/exir_utils.py#L73 if (shape.size() == 0) { shape.push_back(1); strides.push_back(1); } - + MultiArray::MemoryLayout layout(dataType.value(), std::move(shape), std::move(strides)); switch (argType) { case ArgType::Input: { @@ -197,7 +198,7 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) { Error CoreMLBackendDelegate::execute(BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const { + Span args) const { const auto& nArgs = impl_->get_num_arguments(handle); std::vector delegate_args; size_t nInputs = nArgs.first; @@ -281,9 +282,11 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) { } namespace { -auto cls = CoreMLBackendDelegate(); -Backend backend{ETCoreMLStrings.delegateIdentifier.UTF8String, &cls}; -static auto success_with_compiler = register_backend(backend); + #ifndef LAZY_LOAD_IOS_PYTORCH_INITIALIZER + auto cls = CoreMLBackendDelegate(); + Backend backend{ETCoreMLStrings.delegateIdentifier.UTF8String, &cls}; + static auto success_with_compiler = register_backend(backend); + #endif } } // namespace coreml diff --git a/backends/apple/coreml/runtime/delegate/executorch_operations.h b/backends/apple/coreml/runtime/delegate/executorch_operations.h new file mode 100644 index 00000000000..4853c7645be --- /dev/null +++ b/backends/apple/coreml/runtime/delegate/executorch_operations.h @@ -0,0 +1,5 @@ +#pragma once + +namespace executorch::core_ml_backend_delegate { +void register_backend_coreml(); +} // namespace executorch::core_ml_backend_delegate diff --git a/backends/apple/coreml/runtime/delegate/executorch_operations.mm b/backends/apple/coreml/runtime/delegate/executorch_operations.mm new file mode 100644 index 00000000000..1206710d0a6 --- /dev/null +++ b/backends/apple/coreml/runtime/delegate/executorch_operations.mm @@ -0,0 +1,29 @@ +#pragma once + +#include "executorch_operations.h" +#import +#import "ETCoreMLStrings.h" +#import "backend_delegate.h" + +#import +#import +#import + +#include +#import + +namespace executorch::core_ml_backend_delegate { + using executorch::runtime::get_backend_class; + +static std::unique_ptr backendInterfaceLazy_; + +void register_backend_coreml() { + auto backendInterface = executorch::runtime::get_backend_class(ETCoreMLStrings.delegateIdentifier.UTF8String); + if (backendInterface == nullptr) { + backendInterfaceLazy_ = std::make_unique(); + executorch::runtime::Backend backend{ETCoreMLStrings.delegateIdentifier.UTF8String, backendInterfaceLazy_.get()}; + std::ignore = register_backend(backend); + } + } + +} // namespace executorch::core_ml_backend_delegate diff --git a/backends/apple/coreml/runtime/delegate/multiarray.mm b/backends/apple/coreml/runtime/delegate/multiarray.mm index 9443f4df73a..447765bbd8d 100644 --- a/backends/apple/coreml/runtime/delegate/multiarray.mm +++ b/backends/apple/coreml/runtime/delegate/multiarray.mm @@ -124,7 +124,10 @@ bool init_bnns_descriptor(BNNSNDArrayDescriptor& bnns_descriptor, const MultiArr bool copy_using_bnns(const MultiArray& src, MultiArray& dst) { if (src.layout().dataType() != dst.layout().dataType()) { - return false; + // Copying from FP16 to FP32 is supported and this is a common use case + if (!(src.layout().dataType() == MultiArray::DataType::Float16 && dst.layout().dataType() == MultiArray::DataType::Float32)) { + return false; + } } if (dst.layout().num_bytes() < src.layout().num_bytes()) { return false; diff --git a/backends/apple/coreml/runtime/include/coreml_backend/delegate.h b/backends/apple/coreml/runtime/include/coreml_backend/delegate.h index ec402e81717..39075e97a75 100644 --- a/backends/apple/coreml/runtime/include/coreml_backend/delegate.h +++ b/backends/apple/coreml/runtime/include/coreml_backend/delegate.h @@ -48,7 +48,7 @@ class CoreMLBackendDelegate final : public ::executorch::runtime::BackendInterfa /// @retval On success, `Error::Ok` otherwise any other `Error` case. executorch::runtime::Error execute(executorch::runtime::BackendExecutionContext& context, executorch::runtime::DelegateHandle* handle, - executorch::runtime::EValue** args) const override; + executorch::runtime::Span args) const override; /// Returns `true` if the delegate is available otherwise `false`. bool is_available() const override; diff --git a/backends/apple/coreml/runtime/inmemoryfs/inmemory_filesystem.cpp b/backends/apple/coreml/runtime/inmemoryfs/inmemory_filesystem.cpp index ec2f0459a47..c5671b1e6c4 100644 --- a/backends/apple/coreml/runtime/inmemoryfs/inmemory_filesystem.cpp +++ b/backends/apple/coreml/runtime/inmemoryfs/inmemory_filesystem.cpp @@ -7,9 +7,11 @@ #include "inmemory_filesystem.hpp" +#include #include #include #include +#include #include #if __has_include() diff --git a/backends/apple/coreml/scripts/install_requirements.sh b/backends/apple/coreml/scripts/install_requirements.sh index e9f73105bcd..5ec1ea6a1de 100755 --- a/backends/apple/coreml/scripts/install_requirements.sh +++ b/backends/apple/coreml/scripts/install_requirements.sh @@ -12,7 +12,7 @@ SCRIPT_DIR_PATH="$( # TODO(jathu): remove the need to fetch coremltools to build deps for coreml_executor_runner. # Keep this version in sync with: pyproject.toml -COREMLTOOLS_VERSION="8.3" +COREMLTOOLS_VERSION="9.0b1" red=`tput setaf 1` green=`tput setaf 2` diff --git a/backends/apple/coreml/test/test_coreml_quantizer.py b/backends/apple/coreml/test/test_coreml_quantizer.py index d5754328796..eb8b9471345 100644 --- a/backends/apple/coreml/test/test_coreml_quantizer.py +++ b/backends/apple/coreml/test/test_coreml_quantizer.py @@ -15,7 +15,7 @@ ) from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer -from torch.export import export_for_training +from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import ( convert_pt2e, prepare_pt2e, @@ -32,9 +32,7 @@ def quantize_and_compare( ) -> None: assert quantization_type in {"PTQ", "QAT"} - pre_autograd_aten_dialect = export_for_training( - model, example_inputs, strict=True - ).module() + pre_autograd_aten_dialect = export(model, example_inputs, strict=True).module() quantization_config = LinearQuantizerConfig.from_dict( { diff --git a/backends/apple/coreml/test/test_coreml_recipes.py b/backends/apple/coreml/test/test_coreml_recipes.py new file mode 100644 index 00000000000..303d8cb78ed --- /dev/null +++ b/backends/apple/coreml/test/test_coreml_recipes.py @@ -0,0 +1,552 @@ +# Copyright © 2025 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + + +import copy +import unittest + +import coremltools as ct +import torch + +from executorch.backends.apple.coreml.recipes import ( + CoreMLRecipeProvider, + CoreMLRecipeType, +) + +from executorch.backends.apple.coreml.test.test_coreml_utils import ( + IS_VALID_TEST_RUNTIME, +) +from executorch.exir.schema import DelegateCall +from executorch.export import export, ExportRecipe, recipe_registry, StageType + +from torch import nn +from torch.testing._internal.common_quantization import TestHelperModules +from torchao.quantization.utils import compute_error + + +class TestCoreMLRecipes(unittest.TestCase): + """Test suite for CoreML recipes focusing on quantization functionality""" + + def setUp(self): + torch._dynamo.reset() + super().setUp() + self.provider = CoreMLRecipeProvider() + # Register the provider for recipe registry tests + recipe_registry.register_backend_recipe_provider(CoreMLRecipeProvider()) + + def tearDown(self): + super().tearDown() + + def check_fully_delegated(self, session) -> None: + """Helper to verify a program is fully delegated to CoreML""" + session.print_delegation_info() + program = session.get_executorch_program() + instructions = program.execution_plan[0].chains[0].instructions + assert instructions is not None + self.assertEqual(len(instructions), 1) + self.assertIsInstance(instructions[0].instr_args, DelegateCall) + + def _compare_eager_quantized_model_outputs(self, session, example_inputs, atol): + """Utility to compare eager quantized model output with session output after coreml lowering""" + if IS_VALID_TEST_RUNTIME: + source_transform_output = session.get_stage_artifacts()[ + StageType.SOURCE_TRANSFORM + ] + eager_quantized_model = source_transform_output.data["forward"] + output = session.run_method("forward", example_inputs[0])[0] + expected = eager_quantized_model(*example_inputs[0]) + self.assertTrue(torch.allclose(output, expected, atol=atol)) + + def _compare_eager_unquantized_model_outputs( + self, session, eager_unquantized_model, example_inputs, sqnr_threshold=20 + ): + """Utility to compare eager unquantized model output with session output using SQNR""" + if IS_VALID_TEST_RUNTIME: + quantized_output = session.run_method("forward", example_inputs[0])[0] + original_output = eager_unquantized_model(*example_inputs[0]) + error = compute_error(original_output, quantized_output) + print(f"SQNR: {error} dB") + self.assertTrue(error > sqnr_threshold) + + def test_fp32_recipe(self): + """Test FP32 recipe functionality""" + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe(CoreMLRecipeType.FP32), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-3) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_fp16_recipe(self): + """Test FP16 recipe functionality""" + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe(CoreMLRecipeType.FP16), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-3) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_fp_recipes_with_custom_parameters(self): + """Test FP recipes with custom deployment target and compute unit""" + test_cases = [ + (CoreMLRecipeType.FP32, {"minimum_deployment_target": ct.target.iOS16}), + (CoreMLRecipeType.FP16, {"compute_unit": ct.ComputeUnit.CPU_ONLY}), + ] + + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + for recipe_type, kwargs in test_cases: + with self.subTest(recipe=recipe_type.value, kwargs=kwargs): + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe(recipe_type, **kwargs), + ) + self.check_fully_delegated(session) + + def test_int4_weight_only_per_channel(self): + """Test INT4 weight-only per-channel quantization""" + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL + ), + ) + self.check_fully_delegated(session) + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-02) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_int4_weight_only_per_group(self): + """Test INT4 weight-only per-group quantization with different group sizes""" + + class CustomTwoLinearModel(nn.Module): + def __init__(self): + super().__init__() + self.layer1 = nn.Linear(32, 32) + self.layer2 = nn.Linear(32, 8) + + def forward(self, x): + x = torch.relu(self.layer1(x)) + x = self.layer2(x) + return x + + model = CustomTwoLinearModel().eval() + example_inputs = [(torch.randn(1, 32),)] + # Test with different group sizes + for group_size in [8, 16, 32]: + with self.subTest(group_size=group_size): + model_to_export = copy.deepcopy(model) + session = export( + model=model_to_export, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, + group_size=group_size, + ), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs( + session, example_inputs, atol=1e-3 + ) + self._compare_eager_unquantized_model_outputs( + session, model, example_inputs, sqnr_threshold=15 + ) + + def test_int4_weight_only_per_group_validation(self): + """Test INT4 per-group parameter validation""" + # Test invalid group size type + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, group_size="32" + ) + self.assertIn("must be an integer", str(cm.exception)) + + # Test negative group size + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, group_size=-1 + ) + self.assertIn("must be positive", str(cm.exception)) + + def test_int8_weight_only_per_channel(self): + """Test INT8 weight-only per-channel quantization""" + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL + ), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-2) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_int8_weight_only_per_group(self): + """Test INT8 weight-only per-group quantization with different group sizes""" + + class SimpleLinearModel(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(64, 2) + + def forward(self, x): + return self.layer(x) + + model = SimpleLinearModel().eval() + example_inputs = [(torch.randn(1, 64),)] + + # Test with different group sizes + for group_size in [16, 32, 64]: + with self.subTest(group_size=group_size): + model_to_export = copy.deepcopy(model) + session = export( + model=model_to_export, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, + group_size=group_size, + ), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs( + session, example_inputs, atol=1e-2 + ) + self._compare_eager_unquantized_model_outputs( + session, model, example_inputs + ) + + def test_codebook_weight_only_recipe(self): + """Test codebook quantization recipe""" + + class SimpleLinearModel(nn.Module): + def __init__(self): + super().__init__() + self.layer = nn.Linear(32, 2) + + def forward(self, x): + return self.layer(x) + + model = SimpleLinearModel().eval() + example_inputs = [(torch.randn(1, 32),)] + + # Test different block sizes + test_cases = [ + {"bits": 3, "block_size": [-1, 8]}, + ] + + for kwargs in test_cases: + with self.subTest(kwargs=kwargs): + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, **kwargs + ), + ) + self.check_fully_delegated(session) + + def test_codebook_parameter_validation(self): + """Test codebook parameter validation""" + # Test invalid bits type + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, bits="3", block_size=[-1, 8] + ) + self.assertIn("must be an integer", str(cm.exception)) + + # Test bits out of range + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, bits=0, block_size=[-1, 8] + ) + self.assertIn("must be between 1 and 8", str(cm.exception)) + + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, bits=9, block_size=[-1, 8] + ) + self.assertIn("must be between 1 and 8", str(cm.exception)) + + # Test invalid block_size type + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, bits=3, block_size="[-1, 16]" + ) + self.assertIn("must be a list", str(cm.exception)) + + def test_int8_static_quantization(self): + """Test INT8 static quantization (weights + activations)""" + + class SimpleLinearModel(nn.Module): + def __init__(self): + super().__init__() + self.layer1 = nn.Linear(32, 16) + self.layer2 = nn.Linear(16, 2) + + def forward(self, x): + x = torch.relu(self.layer1(x)) + x = self.layer2(x) + return x + + model = SimpleLinearModel().eval() + example_inputs = [(torch.randn(1, 32),)] + + recipe = ExportRecipe.get_recipe( + CoreMLRecipeType.PT2E_INT8_STATIC, minimum_deployment_target=ct.target.iOS17 + ) + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=recipe, + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-3) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_int8_weight_only_pt2e(self): + """Test PT2E-based INT8 weight-only quantization""" + model = TestHelperModules.TwoLinearModule().eval() + example_inputs = [(torch.randn(9, 8),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.PT2E_INT8_WEIGHT_ONLY + ), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-2) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_int8_weight_only_pt2e_with_conv(self): + """Test PT2E-based INT8 weight-only quantization with convolution layers""" + + class ConvModel(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(3, 16, 3, padding=1) + self.conv2 = nn.Conv2d(16, 32, 3, padding=1) + self.pool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(32, 10) + + def forward(self, x): + x = torch.relu(self.conv1(x)) + x = torch.relu(self.conv2(x)) + x = self.pool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + return x + + model = ConvModel().eval() + example_inputs = [(torch.randn(1, 3, 32, 32),)] + + session = export( + model=model, + example_inputs=example_inputs, + export_recipe=ExportRecipe.get_recipe( + CoreMLRecipeType.PT2E_INT8_WEIGHT_ONLY + ), + ) + self.check_fully_delegated(session) + + self._compare_eager_quantized_model_outputs(session, example_inputs, atol=1e-2) + self._compare_eager_unquantized_model_outputs(session, model, example_inputs) + + def test_filter_fn_comprehensive(self): + """Comprehensive test for filter_fn parameter functionality""" + + def custom_filter(module, fqn): + return isinstance(module, nn.Linear) and "target" in fqn + + # Test 1: TorchAO recipes accept filter_fn and default to None + torchao_recipes = [ + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, + ] + + for recipe_type in torchao_recipes: + with self.subTest(f"{recipe_type.value}_default"): + # Test default behavior (None) + recipe = self.provider.create_recipe(recipe_type) + config = recipe.quantization_recipe.ao_quantization_configs[0] + self.assertIsNone(config.filter_fn) + + with self.subTest(f"{recipe_type.value}_custom"): + # Test custom filter_fn + recipe = self.provider.create_recipe( + recipe_type, filter_fn=custom_filter + ) + config = recipe.quantization_recipe.ao_quantization_configs[0] + self.assertEqual(config.filter_fn, custom_filter) + + # Test 2: Codebook recipe accepts filter_fn and has sensible default + with self.subTest("codebook_default"): + recipe = self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, bits=3, block_size=[-1, 16] + ) + config = recipe.quantization_recipe.ao_quantization_configs[0] + self.assertIsNotNone(config.filter_fn) + + # Test default filter targets Linear and Embedding layers + linear_module = nn.Linear(10, 5) + embedding_module = nn.Embedding(100, 10) + conv_module = nn.Conv2d(3, 16, 3) + + self.assertTrue(config.filter_fn(linear_module, "linear")) + self.assertTrue(config.filter_fn(embedding_module, "embedding")) + self.assertFalse(config.filter_fn(conv_module, "conv")) + + with self.subTest("codebook_custom"): + recipe = self.provider.create_recipe( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, + filter_fn=custom_filter, + bits=3, + block_size=[-1, 16], + ) + config = recipe.quantization_recipe.ao_quantization_configs[0] + self.assertEqual(config.filter_fn, custom_filter) + + def test_quantization_recipe_structure(self): + """Test that quantization recipes have proper structure""" + quantization_recipes = [ + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, + ] + + for recipe_type in quantization_recipes: + with self.subTest(recipe=recipe_type.value): + kwargs = ( + {"bits": 3, "block_size": [-1, 16]} + if recipe_type == CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY + else {} + ) + recipe = self.provider.create_recipe(recipe_type, **kwargs) + self.assertIsNotNone(recipe) + + # Should have quantization recipe with ao_quantization_configs + self.assertIsNotNone(recipe.quantization_recipe) + self.assertIsNotNone(recipe.quantization_recipe.ao_quantization_configs) + self.assertEqual( + len(recipe.quantization_recipe.ao_quantization_configs), 1 + ) + + # Should have lowering recipe + self.assertIsNotNone(recipe.lowering_recipe) + self.assertIsNotNone(recipe.lowering_recipe.partitioners) + + def test_recipe_creation_with_defaults(self): + """Test that recipes work with default parameters""" + # Test that all recipes can be created without explicit parameters + all_recipes = [ + CoreMLRecipeType.FP32, + CoreMLRecipeType.FP16, + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, # should use default group_size=32 + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL, + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, # should use default group_size=32 + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, # should use default bits=3, block_size=[-1,16] + ] + + for recipe_type in all_recipes: + with self.subTest(recipe=recipe_type.value): + kwargs = ( + {"bits": 3, "block_size": [-1, 16]} + if recipe_type == CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY + else {} + ) + recipe = self.provider.create_recipe(recipe_type, **kwargs) + self.assertIsNotNone(recipe) + self.assertEqual(recipe.name, recipe_type.value) + + def test_minimum_deployment_target_validation(self): + """Test that minimum_deployment_target validation works correctly for quantization recipes""" + test_cases = [ + (CoreMLRecipeType.PT2E_INT8_STATIC, ct.target.iOS17, {}), + (CoreMLRecipeType.PT2E_INT8_WEIGHT_ONLY, ct.target.iOS17, {}), + ( + CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_CHANNEL, + ct.target.iOS18, + {}, + ), + (CoreMLRecipeType.TORCHAO_INT4_WEIGHT_ONLY_PER_GROUP, ct.target.iOS18, {}), + ( + CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_CHANNEL, + ct.target.iOS16, + {}, + ), + (CoreMLRecipeType.TORCHAO_INT8_WEIGHT_ONLY_PER_GROUP, ct.target.iOS18, {}), + ( + CoreMLRecipeType.CODEBOOK_WEIGHT_ONLY, + ct.target.iOS18, + {"bits": 3, "block_size": [-1, 16]}, + ), + ] + + for recipe_type, min_target, kwargs in test_cases: + with self.subTest(recipe=recipe_type.value): + + # Test 1: Providing deployment target below minimum should raise ValueError + too_low_target = ct.target.iOS15 + with self.assertRaises(ValueError) as cm: + self.provider.create_recipe( + recipe_type, minimum_deployment_target=too_low_target, **kwargs + ) + error_msg = str(cm.exception) + self.assertIn( + f"minimum_deployment_target must be {str(min_target)} or higher", + error_msg, + ) + + # Test 2: Providing valid deployment target should work + valid_recipe = self.provider.create_recipe( + recipe_type, minimum_deployment_target=min_target, **kwargs + ) + self.assertIsNotNone(valid_recipe) + + # Test 3: Not providing deployment target should default to minimum + default_recipe = self.provider.create_recipe(recipe_type, **kwargs) + self.assertIsNotNone(default_recipe) + + # Test 4: Providing deployment target higher than minimum should work + higher_target = ( + ct.target.iOS18 + if min_target == ct.target.iOS17 + else ct.target.iOS18 + ) + higher_recipe = self.provider.create_recipe( + recipe_type, minimum_deployment_target=higher_target, **kwargs + ) + self.assertIsNotNone(higher_recipe) diff --git a/backends/apple/coreml/test/test_coreml_utils.py b/backends/apple/coreml/test/test_coreml_utils.py new file mode 100644 index 00000000000..7d9ac7ba5a5 --- /dev/null +++ b/backends/apple/coreml/test/test_coreml_utils.py @@ -0,0 +1,19 @@ +# Copyright © 2025 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + +import platform +import sys + +import torch + + +def is_fbcode(): + return not hasattr(torch.version, "git_version") + + +IS_VALID_TEST_RUNTIME: bool = ( + (sys.platform == "darwin") + and not is_fbcode() + and tuple(map(int, platform.mac_ver()[0].split("."))) >= (15, 0) +) diff --git a/backends/apple/coreml/test/test_enumerated_shapes.py b/backends/apple/coreml/test/test_enumerated_shapes.py new file mode 100644 index 00000000000..d7d209c4ce4 --- /dev/null +++ b/backends/apple/coreml/test/test_enumerated_shapes.py @@ -0,0 +1,112 @@ +# Copyright © 2023 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + +import unittest + +import coremltools as ct + +import executorch.exir + +import torch + +from executorch.backends.apple.coreml.compiler import CoreMLBackend +from executorch.backends.apple.coreml.partition import CoreMLPartitioner +from executorch.backends.apple.coreml.test.test_coreml_utils import ( + IS_VALID_TEST_RUNTIME, +) + +if IS_VALID_TEST_RUNTIME: + from executorch.runtime import Runtime + + +class TestEnumeratedShapes(unittest.TestCase): + def _compare_outputs(self, executorch_program, eager_program, example_inputs): + if not IS_VALID_TEST_RUNTIME: + return + runtime = Runtime.get() + program = runtime.load_program(executorch_program.buffer) + method = program.load_method("forward") + et_outputs = method.execute(example_inputs)[0] + eager_outputs = eager_program(*example_inputs) + self.assertTrue( + torch.allclose(et_outputs, eager_outputs, atol=1e-02, rtol=1e-02) + ) + + def test_e2e(self): + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(10, 5) + self.linear2 = torch.nn.Linear(11, 5) + + def forward(self, x, y): + return self.linear1(x).sum() + self.linear2(y) + + model = Model() + example_inputs = ( + torch.randn((4, 6, 10)), + torch.randn((5, 11)), + ) + enumerated_shapes = {"x": [[1, 5, 10], [4, 6, 10]], "y": [[3, 11], [5, 11]]} + dynamic_shapes = [ + { + 0: torch.export.Dim.AUTO(min=1, max=4), + 1: torch.export.Dim.AUTO(min=5, max=6), + }, + {0: torch.export.Dim.AUTO(min=3, max=5)}, + ] + ep = torch.export.export( + model.eval(), example_inputs, dynamic_shapes=dynamic_shapes + ) + + compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 + ) + compile_specs.append( + CoreMLBackend.generate_enumerated_shapes_compile_spec( + ep, + enumerated_shapes, + ) + ) + partitioner = CoreMLPartitioner( + compile_specs=compile_specs, lower_full_graph=True + ) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[partitioner], + ) + et_prog = delegated_program.to_executorch() + + good_input1 = ( + torch.randn((1, 5, 10)), + torch.randn((3, 11)), + ) + good_input2 = ( + torch.randn((4, 6, 10)), + torch.randn((5, 11)), + ) + bad_input = ( + torch.randn((1, 5, 10)), + torch.randn((5, 11)), + ) + bad_input2 = ( + torch.randn((2, 7, 12)), + torch.randn((3, 11)), + ) + + self._compare_outputs(et_prog, model, good_input1) + self._compare_outputs(et_prog, model, good_input2) + if IS_VALID_TEST_RUNTIME: + self.assertRaises( + RuntimeError, lambda: self._compare_outputs(et_prog, model, bad_input) + ) + self.assertRaises( + RuntimeError, lambda: self._compare_outputs(et_prog, model, bad_input2) + ) + + +if __name__ == "__main__": + test_runner = TestEnumeratedShapes() + test_runner.test_e2e() diff --git a/backends/apple/coreml/test/test_torch_ops.py b/backends/apple/coreml/test/test_torch_ops.py index 323f76afd1b..de54b684ee7 100644 --- a/backends/apple/coreml/test/test_torch_ops.py +++ b/backends/apple/coreml/test/test_torch_ops.py @@ -2,8 +2,6 @@ # # Please refer to the license found in the LICENSE file in the root directory of the source tree. -import platform -import sys import unittest import coremltools as ct @@ -14,40 +12,36 @@ from executorch.backends.apple.coreml.compiler import CoreMLBackend from executorch.backends.apple.coreml.partition import CoreMLPartitioner -from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_ - - -def is_fbcode(): - return not hasattr(torch.version, "git_version") +from executorch.backends.apple.coreml.test.test_coreml_utils import ( + IS_VALID_TEST_RUNTIME, +) +from executorch.exir.backend.utils import format_delegated_graph +from torchao.prototype.quantization.codebook_coreml import CodebookWeightOnlyConfig +from torchao.quantization import IntxWeightOnlyConfig, PerAxis, PerGroup, quantize_ -_TEST_RUNTIME = ( - (sys.platform == "darwin") - and not is_fbcode() - and tuple(map(int, platform.mac_ver()[0].split("."))) >= (15, 0) -) -if _TEST_RUNTIME: +if IS_VALID_TEST_RUNTIME: from executorch.runtime import Runtime class TestTorchOps(unittest.TestCase): edge_compile_config = executorch.exir.EdgeCompileConfig() - def _coreml_partitioner(self): + def _coreml_partitioner(self, *, minimum_deployment_target=ct.target.iOS18): compile_specs = CoreMLBackend.generate_compile_specs( - minimum_deployment_target=ct.target.iOS18 + minimum_deployment_target=minimum_deployment_target ) return CoreMLPartitioner(compile_specs=compile_specs) def _get_test_model(self): model = torch.nn.Sequential( - torch.nn.Embedding(64, 128), torch.nn.Linear(128, 128), torch.nn.ReLU() + torch.nn.Embedding(64, 128), torch.nn.Linear(128, 256), torch.nn.ReLU() ) example_inputs = (torch.LongTensor([0]),) return model, example_inputs def _compare_outputs(self, executorch_program, eager_program, example_inputs): - if not _TEST_RUNTIME: + if not IS_VALID_TEST_RUNTIME: return runtime = Runtime.get() program = runtime.load_program(executorch_program.buffer) @@ -164,6 +158,165 @@ def test_dequantize_affine_c8w_embedding_b4w_linear(self): et_prog = delegated_program.to_executorch() self._compare_outputs(et_prog, model, example_inputs) + def test_dequantize_affine_c8w_embedding_c8w_linear_ios16(self): + model, example_inputs = self._get_test_model() + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0)), + lambda m, fqn: isinstance(m, torch.nn.Embedding), + ) + quantize_( + model, + IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0)), + ) + ep = torch.export.export(model, example_inputs) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[ + self._coreml_partitioner(minimum_deployment_target=ct.target.iOS16) + ], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + + def test_dequantize_codebook_linear_per_grouped_col(self): + model, example_inputs = self._get_test_model() + quantize_( + model, + CodebookWeightOnlyConfig(dtype=torch.uint2, block_size=[-1, 16]), + ) + ep = torch.export.export(model, example_inputs) + assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + + assert ( + "executorch.exir.dialects.edge._ops.quant.dequantize_codebook.default" + in format_delegated_graph(delegated_program.exported_program().graph_module) + ) + + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + + def test_dequantize_codebook_linear_per_grouped_row(self): + model, example_inputs = self._get_test_model() + quantize_( + model, + CodebookWeightOnlyConfig(dtype=torch.uint2, block_size=[16, -1]), + ) + ep = torch.export.export(model, example_inputs) + assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + + assert ( + "executorch.exir.dialects.edge._ops.quant.dequantize_codebook.default" + in format_delegated_graph(delegated_program.exported_program().graph_module) + ) + + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + + def test_dequantize_codebook_embedding_per_grouped_col(self): + model, example_inputs = self._get_test_model() + quantize_( + model, + CodebookWeightOnlyConfig(dtype=torch.uint3, block_size=[-1, 16]), + lambda m, fqn: isinstance(m, torch.nn.Embedding), + ) + ep = torch.export.export(model, example_inputs) + assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + + assert ( + "executorch.exir.dialects.edge._ops.quant.dequantize_codebook.default" + in format_delegated_graph(delegated_program.exported_program().graph_module) + ) + + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + + def test_dequantize_codebook_embedding_per_grouped_row(self): + model, example_inputs = self._get_test_model() + quantize_( + model, + CodebookWeightOnlyConfig(dtype=torch.uint3, block_size=[16, -1]), + lambda m, fqn: isinstance(m, torch.nn.Embedding), + ) + ep = torch.export.export(model, example_inputs) + assert "torch.ops.quant.dequantize_codebook.default" in ep.graph_module.code + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + + assert ( + "executorch.exir.dialects.edge._ops.quant.dequantize_codebook.default" + in format_delegated_graph(delegated_program.exported_program().graph_module) + ) + + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + + def test__clone_dim_order_contiguous(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.ops.dim_order_ops._clone_dim_order( + x, dim_order=[0, 1, 2, 3] + ) + + model, example_inputs = Model(), (torch.randn(1, 3, 8, 8),) + ep = torch.export.export(model, example_inputs) + delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[self._coreml_partitioner()], + ) + for node in delegated_program.exported_program().graph.nodes: + if node.op == "call_function": + assert node.target.__name__ in [ + "executorch_call_delegate", + "getitem", + ], f"Got unexpected node target after delegation: {node.target.__name__}" + et_prog = delegated_program.to_executorch() + self._compare_outputs(et_prog, model, example_inputs) + if __name__ == "__main__": test_runner = TestTorchOps() @@ -172,3 +325,9 @@ def test_dequantize_affine_c8w_embedding_b4w_linear(self): test_runner.test_dequantize_affine_c4w_embedding() test_runner.test_dequantize_affine_c4w_linear() test_runner.test_dequantize_affine_c8w_embedding_b4w_linear() + test_runner.test_dequantize_affine_c8w_embedding_c8w_linear_ios16() + test_runner.test_dequantize_codebook_linear_per_grouped_col() + test_runner.test_dequantize_codebook_linear_per_grouped_row() + test_runner.test_dequantize_codebook_embedding_per_grouped_col() + test_runner.test_dequantize_codebook_embedding_per_grouped_row() + test_runner.test__clone_dim_order_contiguous() diff --git a/backends/apple/mps/runtime/MPSBackend.mm b/backends/apple/mps/runtime/MPSBackend.mm index 261332436d4..3c136e536ec 100644 --- a/backends/apple/mps/runtime/MPSBackend.mm +++ b/backends/apple/mps/runtime/MPSBackend.mm @@ -30,6 +30,7 @@ using executorch::runtime::Error; using executorch::runtime::FreeableBuffer; using executorch::runtime::Result; +using executorch::runtime::Span; class MPSBackend final : public ::executorch::runtime::BackendInterface { public: @@ -72,7 +73,7 @@ bool is_available() const override { Error execute( ET_UNUSED BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const override { + Span args) const override { auto executor = static_cast(handle); std::vector input_pointers; std::vector output_pointers; diff --git a/backends/apple/mps/serialization/mps_graph_serialize.py b/backends/apple/mps/serialization/mps_graph_serialize.py index 6fa46a2f5e5..4e9dabd3d8a 100644 --- a/backends/apple/mps/serialization/mps_graph_serialize.py +++ b/backends/apple/mps/serialization/mps_graph_serialize.py @@ -1,14 +1,16 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 importlib.resources as _resources import json import os import tempfile -import pkg_resources +import executorch.backends.apple.mps.serialization as serialization_package from executorch.backends.apple.mps.serialization.mps_graph_schema import MPSGraph from executorch.exir._serialize._dataclass import _DataclassEncoder from executorch.exir._serialize._flatbuffer import _flatc_compile @@ -19,7 +21,9 @@ def convert_to_flatbuffer(mps_graph: MPSGraph) -> bytes: with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, "schema.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, "schema.fbs") + ) json_path = os.path.join(d, "schema.json") with open(json_path, "wb") as json_file: json_file.write(mps_graph_json.encode("ascii")) diff --git a/backends/apple/mps/setup.md b/backends/apple/mps/setup.md index 0ecb4151e61..a82a4ee2eea 100644 --- a/backends/apple/mps/setup.md +++ b/backends/apple/mps/setup.md @@ -15,8 +15,8 @@ The MPS backend device maps machine learning computational graphs and primitives * [Introduction to ExecuTorch](../../../docs/source/intro-how-it-works.md) * [Setting up ExecuTorch](../../../docs/source/getting-started-setup.rst) * [Building ExecuTorch with CMake](../../../docs/source/using-executorch-cpp.md#building-with-cmake) -* [ExecuTorch iOS Demo App](https://github.com/pytorch-labs/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) -* [ExecuTorch iOS LLaMA Demo App](../../../docs/source/llm/llama-demo-ios.md) +* [ExecuTorch iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) +* [ExecuTorch LLM iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) ::: :::: diff --git a/backends/apple/mps/targets.bzl b/backends/apple/mps/targets.bzl index 74d79448362..99c97d2b318 100644 --- a/backends/apple/mps/targets.bzl +++ b/backends/apple/mps/targets.bzl @@ -3,6 +3,7 @@ # Provided subject to the LICENSE file in the top level directory. # +load("@fbsource//xplat/executorch/build:build_variables.bzl", "MPS_BACKEND_BUCK_SRCS") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") def define_common_targets(is_xplat = False, platforms = []): @@ -37,10 +38,7 @@ def define_common_targets(is_xplat = False, platforms = []): "runtime/*.h", "runtime/operations/*.h", ]), - "srcs": native.glob([ - "runtime/*.mm", - "runtime/operations/*.mm", - ]), + "srcs": MPS_BACKEND_BUCK_SRCS, "visibility": [ "//executorch/backends/apple/...", "//executorch/examples/...", diff --git a/backends/apple/mps/test/test_mps_utils.py b/backends/apple/mps/test/test_mps_utils.py index 674a4b0ba62..5afa604795e 100644 --- a/backends/apple/mps/test/test_mps_utils.py +++ b/backends/apple/mps/test/test_mps_utils.py @@ -206,7 +206,7 @@ def lower_module_and_test_output( expected_output = model(*sample_inputs) - model = torch.export.export_for_training( + model = torch.export.export( model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True ).module() diff --git a/backends/arm/CMakeLists.txt b/backends/arm/CMakeLists.txt index 3830a1b1108..ede7a96a389 100644 --- a/backends/arm/CMakeLists.txt +++ b/backends/arm/CMakeLists.txt @@ -19,69 +19,72 @@ set(_common_include_directories ) add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) - # bare metal backend builds if(EXECUTORCH_BUILD_ARM_BAREMETAL) -add_compile_options("-Wall" "-Werror") + add_compile_options("-Wall" "-Werror") -# Third-party folder and Ethos-U driver inclued -set(THIRD_PARTY_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/third-party") -set(DRIVER_ETHOSU_INCLUDE_DIR "${THIRD_PARTY_ROOT}/ethos-u-core-driver/include") -include_directories(${DRIVER_ETHOSU_INCLUDE_DIR}) + # Third-party folder and Ethos-U driver inclued + set(THIRD_PARTY_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/third-party") + set(DRIVER_ETHOSU_INCLUDE_DIR + "${THIRD_PARTY_ROOT}/ethos-u-core-driver/include" + ) + include_directories(${DRIVER_ETHOSU_INCLUDE_DIR}) -set(_arm_baremetal_sources backends/arm/runtime/EthosUBackend.cpp - backends/arm/runtime/VelaBinStream.cpp -) -list(TRANSFORM _arm_baremetal_sources PREPEND "${EXECUTORCH_ROOT}/") + set(_arm_baremetal_sources backends/arm/runtime/EthosUBackend.cpp + backends/arm/runtime/VelaBinStream.cpp + ) + list(TRANSFORM _arm_baremetal_sources PREPEND "${EXECUTORCH_ROOT}/") -add_library(executorch_delegate_ethos_u STATIC ${_arm_baremetal_sources}) -target_link_libraries( - executorch_delegate_ethos_u PUBLIC executorch_core ethosu_core_driver -) + add_library(executorch_delegate_ethos_u STATIC ${_arm_baremetal_sources}) + target_link_libraries( + executorch_delegate_ethos_u PUBLIC executorch_core ethosu_core_driver + ) -install(TARGETS executorch_delegate_ethos_u EXPORT ExecuTorchTargets) + install(TARGETS executorch_delegate_ethos_u EXPORT ExecuTorchTargets) -# end config for bare metal builds + # end config for bare metal builds endif() - -# VGF backend builds +# VGF backend builds if(EXECUTORCH_BUILD_VGF) -# include libvgf -set(LIBVGF_PATH "${EXECUTORCH_ROOT}/examples/arm/ethos-u-scratch/ml-sdk-for-vulkan-manifest/sw/vgf-lib/") - -set(VULKAN_THIRD_PARTY_PATH ${EXECUTORCH_ROOT}/backends/vulkan/third-party) -set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include) -set(VOLK_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/volk) - -set(LIBVGF_STATIC "${LIBVGF_PATH}/build/src/libvgf.a") -set(LIBVGF_INCLUDE "${LIBVGF_PATH}/include/") - -add_library(vgf STATIC IMPORTED) -set_property( TARGET vgf PROPERTY IMPORTED_LOCATION "${LIBVGF_STATIC}" ) -target_include_directories(vgf INTERFACE "${LIBVGF_INCLUDE}") - -# Add backend delegate for VGF -set(_vgf_backend_sources backends/arm/runtime/VGFBackend.cpp - backends/arm/runtime/VGFSetup.cpp ) - -# vgf backend -list(TRANSFORM _vgf_backend_sources PREPEND "${EXECUTORCH_ROOT}/") -add_library(vgf_backend ${_vgf_backend_sources}) -target_include_directories( - vgf_backend PUBLIC - ${_common_include_directories} - ${VULKAN_HEADERS_PATH} - ${VOLK_HEADERS_PATH} -) -target_compile_options(vgf_backend PRIVATE -DUSE_VULKAN_WRAPPER -DUSE_VULKAN_VOLK) - - -target_link_libraries(vgf_backend PRIVATE executorch_core) -target_link_libraries(vgf_backend PRIVATE vgf) -executorch_target_link_options_shared_lib(vgf_backend) - -# end config for VGF builds + # include libvgf + set(LIBVGF_PATH + "${EXECUTORCH_ROOT}/examples/arm/ethos-u-scratch/ml-sdk-for-vulkan-manifest/sw/vgf-lib/" + ) + + set(VULKAN_THIRD_PARTY_PATH ${EXECUTORCH_ROOT}/backends/vulkan/third-party) + set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include) + set(VOLK_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/volk) + + set(LIBVGF_STATIC "${LIBVGF_PATH}/build/src/libvgf.a") + set(LIBVGF_INCLUDE "${LIBVGF_PATH}/include/") + + add_library(vgf STATIC IMPORTED) + set_property(TARGET vgf PROPERTY IMPORTED_LOCATION "${LIBVGF_STATIC}") + target_include_directories(vgf INTERFACE "${LIBVGF_INCLUDE}") + + # Add backend delegate for VGF + set(_vgf_backend_sources backends/arm/runtime/VGFBackend.cpp + backends/arm/runtime/VGFSetup.cpp + ) + + # vgf backend + list(TRANSFORM _vgf_backend_sources PREPEND "${EXECUTORCH_ROOT}/") + add_library(vgf_backend ${_vgf_backend_sources}) + install(TARGETS vgf_backend EXPORT ExecuTorchTargets) + target_include_directories( + vgf_backend PRIVATE ${_common_include_directories} ${VULKAN_HEADERS_PATH} + ${VOLK_HEADERS_PATH} + ) + target_compile_options( + vgf_backend PRIVATE -DUSE_VULKAN_WRAPPER -DUSE_VULKAN_VOLK + ) + + target_link_libraries(vgf_backend PRIVATE executorch_core) + target_link_libraries(vgf_backend PRIVATE vgf) + executorch_target_link_options_shared_lib(vgf_backend) + + # end config for VGF builds endif() diff --git a/backends/arm/README.md b/backends/arm/README.md index 9fa8ff8f5be..e495a8e40cb 100644 --- a/backends/arm/README.md +++ b/backends/arm/README.md @@ -1,210 +1,226 @@ -# ExecuTorch Arm/TOSA Delegate +# ExecuTorch Arm® Backend -This subtree contains the Arm(R) Delegate implementation for ExecuTorch. +This subtree contains the Arm® Backend implementation for ExecuTorch. +It supports multiple targets using a common infrastructure, that lowers +PyTorch models to a TOSA representation. This representation is used to +deploy to the following targets: -This delegate is structured to, over time, support a number of different Arm devices -through an AoT flow which targets multiple Arm IP using the TOSA standard. +- **Arm® Ethos™-U55/65/85** - Compiled using the Ethos-U Vela compiler. +- **VGF (Vulkan® Graph Format)** – SPIR-V™ representation for Vulkan-capable devices. -The expected flow is: - * torch.nn.module -> TOSA -> command_stream for fully AoT flows e.g. embedded. - * torch.nn.module -> TOSA for flows supporting a JiT compilation step. +The backend provides an ahead-of-time (AOT) flow, that produces a PTE file for your +chosen target. The AOT flow supports the following development operating systems: -Current backend support is being developed for TOSA to Ethos(TM)-U55/65/85 via the -ethos-u-vela compilation stack. which follows the fully AoT flow. +- Linux aarch64 +- Linux x86_64 +- macOS™ with Apple® Silicon -## Layout +In addition, the following deployment paths are supported by this backend: -Export: -- `ethosu_backend.py` - Main entrypoint for the EthosUBackend. For more information see the section on -[Arm Backend Architecture](#arm-backend-architecture). For examples of use see `executorch/examples/arm`. -- `tosa_mapping.py` - utilities for mapping edge dialect to TOSA -- `tosa_quant_utils.py` - utilities for mapping quantization information to TOSA encoding +- Bare metal build of a reference runtime for Arm® Cortex®-M with Ethos-U acceleration: + - Full testing is available in tree using Corstone™ Fixed Virtual Platforms (FVP). +- Linux target support for VGF capable targets, using the executor_runner. -Operators: -- `node_visitor.py` - Base class for edge operator lowering -- `op_*.py` - Edge operator lowering/serialization to TOSA +More information on TOSA can be found here: https://www.mlplatform.org/tosa/tosa_spec.html. -Passes: -- `arm_pass_manager.py` - Pass manager. Will decide which passes need to be applied depending on the compile_spec. -- `*_pass.py` - Compiler passes derived from ExportPass +## Directory Layout -Quantization: -- `arm_quantizer.py` - Quantizers for Arm backend. Contains the EthosUQuantizer which inherits from the TOSAQuantizer -- `arm_quantizer_utils.py` - Utilities for quantization - -Runtime: -- `runtime/ArmEthosUBackend.cpp` - The Arm backend implementation of the ExecuTorch runtime backend (BackendInterface) for Ethos-U - -Other: -- `third-party/` - Dependencies on other code - in particular the TOSA serialization_lib for compiling to TOSA and the ethos-u-core-driver for the bare-metal backend supporting Ethos-U -- `test/` - Unit test and test support functions - -## Testing - -After a setup you can run unit tests with the test_arm_baremetal.sh script. - -To run the pytests suite run - -``` -backends/arm/test/test_arm_baremetal.sh test_pytest -``` - -To run the unit test suite with Corstone3x0 FVP simulator support use +Below is an overview of the key folder and files in this directory: ``` -backends/arm/test/test_arm_baremetal.sh test_pytest_ethosu_fvp +backends/arm/ +│ +├── _passes/ # Graph transformation passes +│ ├── arm_pass_manager.py # Defines ordering of graph transformations +│ └── *_pass.py # Graph transformation implementation +│ +├── common/ # Common functionality used across the backend +│ +├── debug/ # Debugging schema and functionality +│ +├── ethosu/ # Implementations of EthosUPartitioner and EthosUBackend +│ +├── operator_support/ # Checks if operators can be partitioned +│ +├── operators/ # ATen → TOSA serialization +│ ├── node_visitor.py # Defines base class for ATen → TOSA node visitors +│ └── op_*.py # Lowering implementations for individual operators +│ +├── quantizer/ # Quantization-related logic +│ ├── arm_quantizer.py # EthosUQuantizer and VGFQuantizer definitions +│ └── quantization_annotator.py # Defines how operators are annotated for quantization +│ +├── runtime/ # Backends for running inference on target devices +│ ├── ArmEthosUBackend.cpp +│ └── VGFBackend.cpp +│ +├── scripts/ # Auxiliary build, dependency installation and utility scripts +│ +├── test/ # Unit tests for the backend +│ ├── ops/ # Operator level unit tests +│ ├── models/ # Model level unit tests +│ └── tester/ # Testing harnesses and utilities +│ +├── third-party/ # External dependencies +│ +├── tosa/ # Shared TOSA backend implementation and dialect +│ +└── vgf/ # Implementations of VgfPartitioner and VgfBackend ``` -You can test to run some models with the full fvp test flow +## Building -``` -backends/arm/test/test_arm_baremetal.sh test_full_ethosu_fvp -``` - -## Unit tests -This is the structure of the test directory +The Arm backend can be built using the following command: ``` -test # Root test folder -├── misc # Testing of debug features -├── models # Full model tests -├── ops # Single op tests -├── passes # Compiler passes tests -├── tester # Arm Tester class -├── tosautil # Utility functions for TOSA artifacts -├ common.py # Common functions and definitions used by many tests -├ setup_testing.sh # Script to prepare testing for using the Corstone 3x0 FVP -├ test_arm_baremetal.sh # Help script to trigger testing +./install_executorch.sh ``` -Some example commands to run these tests follow. Run a single test: - -``` -pytest -c /dev/null -v -n auto backends/arm/test/ops/test_add.py -k test_add2_tosa_BI -``` +One of the following commands should also be run once to gather the necessary dependencies for your chosen target(s): -Or discover and run many tests: +For the Ethos-U target: ``` -pytest -c /dev/null -v -n auto backends/arm/test/ops/ +./examples/arm/setup.sh --i-agree-to-the-contained-eula ``` - -You can run tests using Corstone3x0 simulators to see how it would work on something more target like -first you need to build and prepare some used target libs +For the VGF target: ``` -examples/arm/run.sh --model_name=add --build_only -backends/arm/test/setup_testing.sh +./examples/arm/setup.sh --disable-ethos-u-deps --enable-mlsdk-deps ``` -The you can run the tests with +For both Ethos-U & VGF targets: ``` -pytest -c /dev/null -v -n auto backends/arm/test +./examples/arm/setup.sh --i-agree-to-the-contained-eula --enable-mlsdk-deps ``` -### Model test dependencies -Some model tests in Arm backend require third-party libraries or packages. To run these tests, you need to install the required dependencies by running the script `examples/arm/setup.sh` with the flag `--setup-test-dependency`. - -Please note that installing model test dependencies is a standalone process. When using the `--setup-test-dependency` flag, the script will install only the necessary dependencies for model tests, skipping all other setup procedures. - -List of models with specific dependencies: -- Stable Diffusion: [diffusers](https://github.com/huggingface/diffusers/tree/main) - -## Passes +**NOTE:** While developing, it can be convenient to use`./install_executorch.sh --editable`, which creates an editable installation of ExecuTorch. -With the default passes in the Arm Ethos-U backend, assuming the model lowers fully to the -Ethos-U, the exported program is composed of a Quantize node, Ethos-U custom delegate -and a Dequantize node. In some circumstances, you may want to feed quantized input to the Neural -Network straight away, e.g. if you have a camera sensor outputting (u)int8 data and keep all the -arithmetic of the application in the int8 domain. For these cases, you can apply the -`exir/passes/quantize_io_pass.py`. See the unit test in `executorch/backends/arm/ -test/passes/test_ioquantization_pass.py`for an example how to feed quantized inputs and -obtain quantized outputs. - - -### Code coverage - -To get code coverage: - -``` -coverage run --source= --rcfile=backends/arm/test/.coveragerc -m pytest \ ---config-file=/dev/null backends/arm/test/ -``` - -All files in `SRC` and its child directories will be analysed for code coverage, -unless explicitly exluded in the .coveragerc file. If using venv this might be -under `env/lib/python/site-packages/executorch/`. To get the -absolute path, run: - -``` -python -c "import executorch; print(executorch.__path__)" -``` +## Testing -This contains a list of paths where the source directory is located. Pick the -one that is located in `env/lib`. If that does not work try the others. Add -`backends/arm` to the path in `--source` to only get code coverage for the Arm -backend. +There are two approaches for running the tests for the Arm backend. This section will explain these two approaches: -### A note on unit tests +### Using test_arm_baremetal.sh -There are currently 3 ways we unit test our code. -1. TOSA main inference. These tests are using non-quantized data and ops. Edge IR representation of the module is lowered to a TOSA flatbuffer, which is tested for numerical correcteness using the ```tosa_reference_model``` tool. -2. TOSA base inference. Same as above, but data and ops are quantized. -3. Ethos-U55. These tests use quantized data and ops (aka TOSA base inference). Edge IR is lowered to a TOSA flatbuffer, which is fed into the Vela compiler. Theses tests are functional tests and do not test numerical correctness, since that should be guaranteed by TOSA. +The backend provides a script `backends/arm/test/test_arm_baremetal.sh`, which is used in the `trunk` CI workflow. +This approach is useful for checking your change against this workflow on your own machine. +These scripts also install the necessary dependencies to run the tests. +Below is an overview of some of the testing options this script provides: -In order to distinguise between the different tests, the following suffixes have been added to the respective test case. -* ```_MI``` for main inference -* ```_BI``` for base inference -* ```_U55_BI``` for base inference on U55 +| Command | Description | +| ---------------------------------------------- | -------------------------------------------- | +| `test_arm_baremetal.sh test_pytest` | Runs all unit tests. | +| `test_arm_baremetal.sh test_pytest_ethosu_fvp` | Same as `test_pytest` but uses Corstone FVP. | +| `test_arm_baremetal.sh test_run_ethosu_fvp` | Runs some models with Corstone FVP. | +| `test_arm_baremetal.sh test_full_ethosu_fvp` | Runs E2E model tests on Corstone FVP. | +| `test_arm_baremetal.sh test_pytest_vkml` | Runs all unit tests with Vulkan ML. | +| `test_arm_baremetal.sh test_full_vkml` | Run E2E models test with Vulkan ML. | -## Help & Improvements -If you have problems or questions, or have suggestions for ways to make -implementation and testing better, please reach out to the Arm team developing this delegate, or -create an issue on [github](https://www.github.com/pytorch/executorch/issues). +For more information, please refer to the `backends/arm/test/test_arm_baremetal.sh` script. -# Arm Backend Architecture +### Using pytest -The broad principle with the Arm backend implemention for ExecuTorch is to support multiple Arm devices and device configurations through a largely Homogeneous flow with maximal sharing of class logic. -The EthosUBackend is currently the one user facing API that target the Ethos-U55 and Ethos-U85 hardware IP. It is using the TOSABackend under the hood to share code and functionality, but also to separate testing possibilities to the TOSA flow itself. +The Arm backend uses `pytest` to run the unit test suite in `backends/arm/test`. +This option offers flexibility, allowing a specific test or a particular subset of the testsuite to be run. +Below provides some examples of how to use it: -In practice for compilation, this means that the flow goes via [Arm TOSA](https://www.mlplatform.org/tosa/tosa_spec.html) to produce a common IR and quantization behaviour compatible with our various IP, and typically, device-specific backends to further lower to a device specific binary which can happen ahead of time (within the Python development flow) or at runtime (during a JIT compilation stage). +- To run all the unit tests run the following command: -In practice for the runtime, this means we will share common runtime backend functionality, with the aim for features like debugging to be available through common tooling. + ``` + pytest -v -n auto backends/arm/test/ + ``` +- To run a specific test in a file: -## Arm Backend Status and Maturity + ``` + pytest -v backends/arm/test/ops/test_add.py -k test_add_tensor_tosa_INT_3 + ``` -The Arm EthosU Backend should be considered a prototype quality at this point, likely subject to significant change and improvement, and with a limited coverage of functionality. We are actively developing this codebase. +#### Testing Dependencies -## Current flows +Some tests, with `u55`, `u85` and `vgf` in the name require external dependencies to run if you use `pytest`: -The EthosUBackend has a two stage process, -- Compile to TOSA to rationalise the graph into known hardware support profiles. Currently this is to v1.0 TOSA INT with specific concern to a subset which gives support on Ethos-U55 and Ethos-U85, the target of the initial prototype efforts. This calls into the TOSABackend. -- Lower via the ethos-u-vela compilation flow which takes TOSA v1.0 as an input and produces a low level commandstream for the hardware which is then passed via the delegate to the ethos-u-core-driver for direct execution. +- When a test contains `u55` or `u85`, you must run the following to setup the executor_runner: + ``` + ./backends/arm/scripts/build_executorch.sh + ./backends/arm/test/setup_testing.sh + ``` +- When a test contains `vgf`, you must run the following to install the ML SDK: + ``` + ./backends/arm/scripts/build_executorch.sh + ./backends/arm/test/setup_testing_vkml.sh + ``` -The EthosUPartitioner is currenly used to ensure the operations converted are Ethos-U compatible, but will be extended to offer spec-correct TOSA Base inference and TOSA Main Inference generation in future. +In addition, some model tests in the Arm backend require third-party libraries or packages. +To run these tests, you need to install the required dependencies by running the script `examples/arm/setup.sh` with the flag `--setup-test-dependency`. -There is also a generic TOSABackend with accompanying TOSAPartitioner and TOSAQuantizer, which are used by the EthosUBackend and friends. The Arm TOSA Backend can be used by it's own to verify the lowering to the TOSA representation of the model (refer to the unit tests in backends/arm/test which uses the TOSA backend in the test suites). +Please note that installing model test dependencies is a standalone process. When using the `--setup-test-dependency` flag, +the script will install only the necessary dependencies for model tests, skipping all other setup procedures. -### Controlling compilation +## Using pre-commit -It is possible to control the compilation flow to aid in development and debug of both networks and the code itself. +A pre-commit script is available in the backend to help developers. Follow the steps below to enable it: -Configuration of the EthosUBackend export flow is controlled by CompileSpec information (essentially used as compilation flags) to determine which of these outputs is produced. In particular this allows for use of the tosa_reference_model to run intermediate output to check for correctness and quantization accuracy without a full loop via hardware implemntation. +``` +cp backends/arm/scripts/pre-commit .git/hooks/ +``` -As this is in active development see the EthosUBackend for accurate information on [compilation flags](https://github.com/pytorch/executorch/blob/29f6dc9353e90951ed3fae3c57ae416de0520067/backends/arm/arm_backend.py#L319-L324) +## Notes on model specific and optional passes -## Model specific and optional passes The current TOSA version does not support int64. However, int64 is commonly used in many models. In order to lower the operators with int64 inputs and/or outputs to TOSA, a few passes have been developed to handle the int64-related issues. The main idea behind these passes is to replace the uses of int64 with int32 where feasible. + - For floating-point models, these passes need to run very early in the lowering process and can be passed in to the to_edge_transform_and_lower() function call as an optional parameter. - For quantized models, these transformations will be automatically handled during annotation before the export stage. List of model specific and optional passes: -- InsertCastForOpsWithInt64InputPass - - Functionality: - - For LLMs such as LLama, some opeartors like aten.embedding have int64 input. In order to lower these operators to TOSA, this pass will insert a casting node that converts the input from int64 to int32. - - Example usage: backends/arm/test/models/test_llama.py - - Supported Ops: - - aten.embedding.default, aten.slice_copy.Tensor + +- ConvertInt64ConstOpsToInt32Pass + - Functionalities: + - Rewrites constant-producing ops that output int64 to instead output int32, when values are within int32 bounds. + - Supported Ops: + - `torch.full`, `torch.arange`, `torch.eye`, `torch.linspace`, `torch.tensor` + - Example usage: + - backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py + - backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py + +- ConvertInt64OutputOpsToInt32Pass + - Overview: + - Rewrites or removes operations that produce int64 outputs, converting them to int32 where possible. + - Overflow checks are applied selectively; for ops without such checks, users need to ensure values fit within the int32 range. + - Functionalities: + 1. Handling casting to int64: + - (1) int32 -> int64: + - Removes the cast and redirect uses of int64 to int32 + - (2) other types -> int64: + - Rewrites the cast to other types -> int32 + - Supported Ops: + - torch.ops.aten.to.\[dtype|dtype_layout\] + - exir_ops.edge.dim_order_ops.\_to_dim_order_copy.default + 2. Post-process argmax outputs: + - Inserts an int64->int32 cast after the argmax operations that produce int64 outputs: + - Supported Ops: + - torch.ops.aten.argmax.default + - exir_ops.edge.aten.argmax.default + - Example usage: + - (Functionality 1) backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py + - (Functionality 2) backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py + +- InsertInt32CastsAfterInt64PlaceholdersPass + - Functionalities: + - Inserts an int64 -> int32 cast immediately after each int64 placeholder (graph input). + - Redirects all uses of each int64 placeholder to its int32 cast output. + - Inserts local int32 -> int64 casts at call sites where an operator requires int64 inputs, e.g. `torch.nn.functional.one_hot` + - Pass ordering: + - When used with `ConvertInt64ConstOpsToInt32Pass` and `ConvertInt64OutputOpsToInt32Pass`, run this pass last. + - Rationale: Those passes may cause retracing to re-infer some int64 placeholders as int32. Running this pass last casts only inputs that remain int64, minimizing inserted casts. + - Example usage: + - backends/arm/test/models/test_llama.py + - backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py + - backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py + +## Help & Improvements + +If you have problems or questions, or have suggestions for ways to improve the Arm backend, please reach out +to the Arm team developing this backend, or create an issue on [here](https://www.github.com/pytorch/executorch/issues) and add the "partner: arm" label. diff --git a/backends/arm/TARGETS b/backends/arm/TARGETS index 810b5c09136..35b16f819e5 100644 --- a/backends/arm/TARGETS +++ b/backends/arm/TARGETS @@ -1,7 +1,12 @@ +# Copyright 2025 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. + # @noautodeps -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "ethosu_partitioner", srcs = [ "ethosu/__init__.py", @@ -12,22 +17,55 @@ python_library( ":arm_partitioner", ] ) -python_library( +runtime.python_library( + name = "vgf_partitioner", + srcs = [ + "vgf/__init__.py", + "vgf/backend.py", + "vgf/partitioner.py" + ], + deps = [ + ":arm_partitioner", + ] +) +runtime.python_library( + name = "constants", + srcs = [ + "constants.py", + ], + deps = [ + "//executorch/exir/dialects:lib", + ], +) +runtime.python_library( + name = "common", + srcs = [ + "common/__init__.py", + "common/debug.py", + ], + deps = [ + "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", + "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", + "//caffe2:torch", + "//executorch/exir:lib", + ], +) +runtime.python_library( name = "arm_partitioner", srcs = [ - "tosa_backend.py", - "tosa_partitioner.py", - "vgf_backend.py", - "vgf_partitioner.py", + "tosa/backend.py", + "tosa/partitioner.py", ], deps = [ ":arm_backend", + ":constants", + "//executorch/backends/arm/debug:schema", "//executorch/backends/arm/operator_support:operator_support", "//executorch/backends/arm/_passes:passes", "//executorch/exir:lib", ], ) -python_library( +runtime.python_library( name = "arm_backend", srcs = [ "arm_backend.py", @@ -46,20 +84,20 @@ python_library( "//executorch/backends/arm/_passes:passes", ], ) -python_library( +runtime.python_library( name = "process_node", srcs = ["process_node.py"], deps = [ "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/tosa:tosa", "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/tosa:tosa", "//executorch/backends/arm/operators:node_visitor", - "//executorch/backends/arm:tosa_mapping", - "//executorch/backends/arm:tosa_quant_utils", - "//executorch/backends/arm:tosa_utils", + "//executorch/backends/arm/tosa:mapping", + "//executorch/backends/arm/tosa:quant_utils", + "//executorch/backends/arm/tosa:utils", "//executorch/exir:lib", ], ) -python_library( +runtime.python_library( name = "arm_vela", srcs = [ "arm_vela.py", @@ -68,54 +106,7 @@ python_library( "fbsource//third-party/pypi/ethos-u-vela:ethos-u-vela", ], ) -python_library( - name = "tosa_mapping", - srcs = [ - "tosa_mapping.py", - ], - deps = [ - "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", - "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", - "//caffe2:torch", - ], -) -python_library( - name = "tosa_quant_utils", - srcs = [ - "tosa_quant_utils.py", - ], - deps = [ - "fbsource//third-party/pypi/numpy:numpy", - "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", - "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", - "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/tosa:tosa", - "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/tosa:tosa", - ":tosa_mapping", - "//executorch/exir/dialects:lib", - ], -) -python_library( - name = "tosa_specification", - srcs = [ - "tosa_specification.py", - ], - deps = [ - "fbsource//third-party/pypi/packaging:packaging", - "//executorch/exir/backend:compile_spec_schema", - ], -) -python_library( - name = "tosa_utils", - srcs = [ - "tosa_utils.py", - ], - deps = [ - "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", - ":tosa_quant_utils", - "//executorch/backends/arm/operators:node_visitor", - ], -) -python_library( +runtime.python_library( name = "arm_model_evaluator", srcs = [ "util/arm_model_evaluator.py", diff --git a/backends/arm/_passes/TARGETS b/backends/arm/_passes/TARGETS index bbb94c1d703..bb4e992ada1 100644 --- a/backends/arm/_passes/TARGETS +++ b/backends/arm/_passes/TARGETS @@ -1,11 +1,13 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "passes", srcs = glob(["*.py"]), deps = [ - "//executorch/backends/arm:tosa_quant_utils", - "//executorch/backends/arm:tosa_utils", + "//executorch/backends/arm:common", + "//executorch/backends/arm:constants", + "//executorch/backends/arm/tosa:quant_utils", + "//executorch/backends/arm/tosa:utils", "//executorch/backends/arm/tosa/dialect:lib", "//executorch/backends/transforms:fuse_view_copy", "//executorch/backends/transforms:remove_getitem_op", diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 986bc4f024b..f9e23f73cc5 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -7,16 +7,19 @@ from . import arm_pass_utils # noqa from .arm_pass import ArmPass # noqa # usort: skip from .add_bias_pass import AddBiasPass # noqa -from .annotate_channels_last_dim_order_pass import AnnotateChannelsLastDimOrder # noqa from .annotate_decomposed_matmul import AnnotateDecomposedMatmulPass # noqa +from .annotate_output_dim_order_pass import AnnotateOutputDimOrderPass # noqa from .broadcast_args_pass import BroadcastArgsPass # noqa from .cast_bool_to_int8_pass import CastBoolToInt8Pass # noqa from .cast_int64_pass import CastInt64BuffersToInt32Pass # noqa from .cast_to_int32_pass import CastToInt32Pass # noqa from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass # noqa from .convert_any_default_dim_dims_pass import ConvertAnyDefaultDimDimsPass # 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 +from .convert_int64_const_ops_to_int32 import ConvertInt64ConstOpsToInt32Pass # noqa +from .convert_int64_output_ops_to_int32 import ConvertInt64OutputOpsToInt32Pass # noqa from .convert_int_pow_to_mul import ConvertIntPowToMuls # noqa from .convert_minmax_pass import ConvertMinMaxPass # noqa from .convert_split_to_slice import ConvertSplitToSlicePass # noqa @@ -25,22 +28,29 @@ from .decompose_acosh_pass import DecomposeAcoshPass # noqa from .decompose_adaptive_avg_pool2d_pass import DecomposeAdaptiveAvgPool2dPass # noqa from .decompose_addmm_pass import DecomposeAddmmPass # noqa -from .decompose_asin_pass import DecomposeAsinPass # noqa +from .decompose_asin_and_acos_pass import DecomposeAsinAndAcosPass # noqa from .decompose_asinh_pass import DecomposeAsinhPass # noqa from .decompose_atan_pass import DecomposeAtanPass # noqa from .decompose_atanh_pass import DecomposeAtanhPass # noqa from .decompose_avg_pool2d import DecomposeAvgPool2d # noqa from .decompose_batch_norm_no_stats import DecomposeBatchNormNoStatsPass # noqa +from .decompose_cosh_pass import DecomposeCoshPass # noqa from .decompose_cosine_similarity_pass import DecomposeCosineSimilarityPass # noqa +from .decompose_cumsum_pass import DecomposeCumsumPass # noqa from .decompose_div_pass import DecomposeDivPass # noqa +from .decompose_div_tensor_mode import DecomposeDivTensorModePass # noqa +from .decompose_elu_pass import DecomposeEluPass # noqa from .decompose_embedding_pass import DecomposeEmbeddingPass # noqa # noqa +from .decompose_expm1_pass import DecomposeExpm1Pass # noqa from .decompose_gelu_pass import DecomposeGeluPass # noqa +from .decompose_glu_pass import DecomposeGluPass # noqa from .decompose_grouped_conv import DecomposeGroupedConv # noqa from .decompose_groupnorm_pass import DecomposeGroupNormPass # noqa from .decompose_layernorm_pass import DecomposeLayerNormPass # noqa from .decompose_leaky_relu_pass import DecomposeLeakyReLUPass # noqa from .decompose_linalg_vector_norm_pass import DecomposeLinearVectorNormPass # noqa from .decompose_linear_pass import DecomposeLinearPass # noqa +from .decompose_logit_pass import DecomposeLogitPass # noqa from .decompose_masked_fill import DecomposeMaskedFill # noqa from .decompose_maxpool2d_with_dilation import DecomposeMaxPool2DPass # noqa from .decompose_meandim_pass import DecomposeMeanDimPass # noqa @@ -65,21 +75,22 @@ from .fuse_constant_ops_pass import ComputeConstantOpsAOT, FuseConstantArgsPass # noqa from .fuse_equal_placeholders_pass import FuseEqualPlaceholdersPass # noqa from .fuse_quantized_activation_pass import FuseQuantizedActivationPass # noqa -from .insert_int64_input_cast_pass import ( # noqa # noqa - InsertCastForOpsWithInt64InputPass, +from .insert_int32_casts_after_int64_placeholders import ( # noqa + InsertInt32CastsAfterInt64PlaceholdersPass, ) from .insert_rescales_pass import InsertRescalePass # noqa from .insert_table_ops import InsertTableOpsPass # noqa from .match_arg_dtype_pass import MatchArgDtypePass # noqa from .match_arg_ranks_pass import MatchArgRanksPass # noqa from .mm_to_bmm_pass import ConvertMmToBmmPass # noqa -from .remove_clone_pass import RemoveClonePass # noqa +from .remove_noop_pass import RemoveNoopPass # noqa from .replace_scalar_with_tensor_pass import ( # noqa ReplaceScalarWithTensorArgPassTOSABI, ReplaceScalarWithTensorArgPassTOSAMI, ) from .scalars_to_attribute_pass import ScalarsToAttributePass # noqa from .size_adjust_input_pass import SizeAdjustInputPass # noqa +from .to_tosa_memory_format_pass import ToTosaMemoryFormatPass # noqa from .unsqueeze_before_repeat_pass import UnsqueezeBeforeRepeatPass # noqa from .unsqueeze_scalar_placeholders_pass import UnsqueezeScalarPlaceholdersPass # noqa from .replace_inf_values_pass import ReplaceInfValues # noqa # usort: skip diff --git a/backends/arm/_passes/annotate_output_dim_order_pass.py b/backends/arm/_passes/annotate_output_dim_order_pass.py new file mode 100644 index 00000000000..08f93383a9c --- /dev/null +++ b/backends/arm/_passes/annotate_output_dim_order_pass.py @@ -0,0 +1,21 @@ +# Copyright 2025 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.arm._passes import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import get_output_dim_orders +from executorch.exir.pass_base import PassResult + + +class AnnotateOutputDimOrderPass(ArmPass): + """ + Stores the current output dim_orders in the meta dict of the output node. This is used + for verifying that the dim order does not change unexpectedly in later passes. + """ + + def call(self, graph_module): + output_node = graph_module.graph.output_node() + output_node.meta["original_dim_orders"] = get_output_dim_orders(graph_module) + + return PassResult(graph_module, True) diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 9bef29cebdc..f49206da67e 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -10,8 +10,8 @@ import executorch.backends.arm.tosa.dialect # noqa: unused from executorch.backends.arm._passes import ( AddBiasPass, - AnnotateChannelsLastDimOrder, AnnotateDecomposedMatmulPass, + AnnotateOutputDimOrderPass, BroadcastArgsPass, CastBoolToInt8Pass, CastInt64BuffersToInt32Pass, @@ -19,8 +19,11 @@ ComputeConstantOpsAOT, Conv1dUnsqueezePass, ConvertAnyDefaultDimDimsPass, + ConvertELUParamsPass, ConvertExpandCopyToRepeatPass, ConvertFullLikeToFullPass, + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, ConvertIntPowToMuls, ConvertMinMaxPass, ConvertMmToBmmPass, @@ -30,22 +33,29 @@ DecomposeAcoshPass, DecomposeAdaptiveAvgPool2dPass, DecomposeAddmmPass, + DecomposeAsinAndAcosPass, DecomposeAsinhPass, - DecomposeAsinPass, DecomposeAtanhPass, DecomposeAtanPass, DecomposeAvgPool2d, DecomposeBatchNormNoStatsPass, + DecomposeCoshPass, DecomposeCosineSimilarityPass, + DecomposeCumsumPass, DecomposeDivPass, + DecomposeDivTensorModePass, + DecomposeEluPass, DecomposeEmbeddingPass, + DecomposeExpm1Pass, DecomposeGeluPass, + DecomposeGluPass, DecomposeGroupedConv, DecomposeGroupNormPass, DecomposeLayerNormPass, DecomposeLeakyReLUPass, DecomposeLinearPass, DecomposeLinearVectorNormPass, + DecomposeLogitPass, DecomposeMaskedFill, DecomposeMaxPool2DPass, DecomposeMeanDimPass, @@ -66,23 +76,25 @@ FuseConstantArgsPass, FuseEqualPlaceholdersPass, FuseQuantizedActivationPass, - InsertCastForOpsWithInt64InputPass, + InsertInt32CastsAfterInt64PlaceholdersPass, InsertRescalePass, InsertTableOpsPass, MatchArgDtypePass, MatchArgRanksPass, QuantizeOperatorArguments, - RemoveClonePass, + RemoveNoopPass, ReplaceInfValues, ReplaceScalarWithTensorArgPassTOSABI, ReplaceScalarWithTensorArgPassTOSAMI, RetraceFoldedDtypesPass, ScalarsToAttributePass, SizeAdjustInputPass, + ToTosaMemoryFormatPass, UnsqueezeBeforeRepeatPass, UnsqueezeScalarPlaceholdersPass, ) -from executorch.backends.arm.tosa_specification import ( + +from executorch.backends.arm.tosa.specification import ( TosaLoweringContext, TosaSpecification, ) @@ -93,6 +105,7 @@ from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass from executorch.exir import ExportedProgram from executorch.exir.pass_manager import PassManager +from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass from torch.fx import GraphModule @@ -107,6 +120,7 @@ def _transform(self, graph_module: GraphModule): return self(graph_module).graph_module def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule: + self.add_pass(AnnotateOutputDimOrderPass()) self.add_pass(FuseQuantizedActivationPass()) self.add_pass(RemoveGetItemPass()) self.add_pass(ConvertSplitToSlicePass()) @@ -127,6 +141,7 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(ReplaceScalarWithTensorArgPassTOSABI()) self.add_pass(AnnotateDecomposedMatmulPass()) self.add_pass(QuantizeOperatorArguments()) + self.add_pass(ConvertELUParamsPass()) self.add_pass(FoldAndAnnotateQParamsPass(exported_program)) # type: ignore[call-arg] self.add_pass(RetraceFoldedDtypesPass()) self.add_pass(UnsqueezeScalarPlaceholdersPass(exported_program)) @@ -139,11 +154,11 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(ComputeConstantOpsAOT(exported_program)) self.add_pass(DecomposeGroupedConv()) - self.add_pass(RemoveClonePass()) self.add_pass(ConvertExpandCopyToRepeatPass()) self.add_pass(UnsqueezeBeforeRepeatPass()) self.add_pass(CastInt64BuffersToInt32Pass(exported_program)) self.add_pass(DecomposeSumPass()) + self.add_pass(DecomposeCumsumPass(exported_program)) self.add_pass(Conv1dUnsqueezePass()) self.add_pass(DecomposeMaxPool2DPass()) self.add_pass(SizeAdjustInputPass()) @@ -156,25 +171,33 @@ def _tosa_INT_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(InsertTableOpsPass(exported_program)) self.add_pass(FuseEqualPlaceholdersPass(exported_program)) - self.add_pass(AnnotateChannelsLastDimOrder()) + self.add_pass(ToTosaMemoryFormatPass(exported_program)) + self.add_pass(RemoveNoopPass()) self.add_pass(InsertRescalePass()) return self._transform(exported_program.graph_module) def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule: + self.add_pass(AnnotateOutputDimOrderPass()) + self.add_pass(DecomposeExpm1Pass()) + self.add_pass(DecomposeLogitPass()) self.add_pass(DecomposeMaskedFill()) self.add_pass(DecomposeRoundPass()) self.add_pass(DecomposeAcoshPass()) - self.add_pass(DecomposeAsinPass()) self.add_pass(DecomposeAsinhPass()) + self.add_pass(DecomposeCoshPass()) + self.add_pass(DecomposeAsinAndAcosPass()) self.add_pass(DecomposeSqrtPass()) self.add_pass(DecomposeAtanPass()) self.add_pass(DecomposeAtanhPass()) self.add_pass(DecomposeAddmmPass()) + self.add_pass(DecomposeEluPass()) + self.add_pass(DecomposeExpm1Pass()) self.add_pass(ConvertIntPowToMuls()) self.add_pass(CastBoolToInt8Pass()) self.add_pass(DecomposeSinhPass()) self.add_pass(DecomposeSignPass()) + self.add_pass(DecomposeDivTensorModePass()) self.add_pass(ReplaceScalarWithTensorArgPassTOSAMI()) self.add_pass(DecomposeEmbeddingPass()) self.add_pass(FuseQuantizedActivationPass()) @@ -182,6 +205,7 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(ConvertSplitToSlicePass()) self.add_pass(FuseBatchnorm2DPass(exported_program)) self.add_pass(ConvertMmToBmmPass()) + self.add_pass(DecomposeGluPass()) self.add_pass(DecomposeLinearPass()) self.add_pass(DecomposeLeakyReLUPass()) self.add_pass(DecomposeGroupNormPass()) @@ -214,11 +238,10 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(ComputeConstantOpsAOT(exported_program)) self.add_pass(DecomposeGroupedConv()) - self.add_pass(RemoveClonePass()) self.add_pass(ConvertExpandCopyToRepeatPass()) self.add_pass(UnsqueezeBeforeRepeatPass()) - self.add_pass(CastInt64BuffersToInt32Pass(exported_program)) self.add_pass(DecomposeSumPass()) + self.add_pass(DecomposeCumsumPass(exported_program)) self.add_pass(Conv1dUnsqueezePass()) self.add_pass(DecomposeMaxPool2DPass()) self.add_pass(SizeAdjustInputPass()) @@ -227,10 +250,12 @@ def _tosa_FP_pipeline(self, exported_program: ExportedProgram) -> GraphModule: self.add_pass(FuseViewCopyTransform()) self.add_pass(FuseConstantArgsPass(exported_program)) + self.add_pass(CastInt64BuffersToInt32Pass(exported_program)) self.add_pass(AddBiasPass(exported_program)) self.add_pass(InsertTableOpsPass(exported_program)) self.add_pass(FuseEqualPlaceholdersPass(exported_program)) - self.add_pass(AnnotateChannelsLastDimOrder()) + self.add_pass(ToTosaMemoryFormatPass(exported_program)) + self.add_pass(RemoveNoopPass()) self.add_pass(InsertRescalePass()) return self._transform(exported_program.graph_module) @@ -247,13 +272,20 @@ def transform_to_backend_pipeline(self, exported_program: ExportedProgram): ) def transform_for_annotation_pipeline(self, graph_module: GraphModule): - self.add_pass(InsertCastForOpsWithInt64InputPass()) + self.add_pass( + RemoveGraphAssertsPass() + ) # ConvertInt64ConstOpsToInt32Pass requires this pass to remove the assertation in Graph + self.add_pass(ConvertInt64ConstOpsToInt32Pass()) + self.add_pass(ConvertInt64OutputOpsToInt32Pass()) + self.add_pass(InsertInt32CastsAfterInt64PlaceholdersPass()) self.add_pass(DecomposeEmbeddingPass()) self.add_pass(DecomposeScaledDotProductAttention()) self.add_pass(DecomposeRoundPass()) + self.add_pass(DecomposeLogitPass()) self.add_pass(CastBoolToInt8Pass()) self.add_pass(DecomposeSignPass()) self.add_pass(DecomposeAddmmPass()) + self.add_pass(DecomposeDivTensorModePass()) self.add_pass(ReplaceScalarWithTensorArgPassTOSABI()) self.add_pass(ScalarsToAttributePass()) self.add_pass(DecomposeGroupNormPass()) @@ -262,6 +294,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): self.add_pass(DecomposeMeanDimPass(graph_module, self.tosa_spec)) self.add_pass(DecomposeNotEqualPass()) self.add_pass(DecomposeCosineSimilarityPass()) + self.add_pass(DecomposeGluPass()) self.add_pass(DecomposeDivPass()) self.add_pass(DecomposeLeakyReLUPass()) self.add_pass(DecomposeLinearVectorNormPass()) diff --git a/backends/arm/_passes/arm_pass_utils.py b/backends/arm/_passes/arm_pass_utils.py index 1e0c21239e2..71e2030958f 100644 --- a/backends/arm/_passes/arm_pass_utils.py +++ b/backends/arm/_passes/arm_pass_utils.py @@ -13,7 +13,7 @@ import torch import torch.fx -from executorch.backends.arm.tosa_utils import get_node_debug_info +from executorch.backends.arm.common.debug import get_node_debug_info from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops @@ -235,3 +235,8 @@ def set_node_arg(node: torch.fx.Node, i: int | str, value): node.kwargs = kwargs else: raise RuntimeError("Invalid type") + + +def get_output_dim_orders(graph_module): + output_node = graph_module.graph.output_node() + return [get_first_fake_tensor(node).dim_order() for node in output_node.args[0]] diff --git a/backends/arm/_passes/convert_elu_params.py b/backends/arm/_passes/convert_elu_params.py new file mode 100644 index 00000000000..7da58ae4bb4 --- /dev/null +++ b/backends/arm/_passes/convert_elu_params.py @@ -0,0 +1,53 @@ +# Copyright 2025 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.arm_pass_utils import create_node +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +class ConvertELUParamsPass(ExportPass): + """ + Pass to convert the input_scale kwarg of ELU operator from float to + int. + + It has been set to 2 as the outputs seem to stay the same regardless of what + the value of input_scale is, as long as that value is not 1. + """ + + def call(self, graph_module: torch.fx.GraphModule): + modified_graph = False + graph = graph_module.graph + node_list = graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.elu.default + ) + for node in node_list: + with graph.inserting_after(node): + replace_node = create_node(graph, exir_ops.edge.aten.elu.default) + old_args = list(node.args) + + alpha = old_args[1] if len(old_args) > 1 else 1.0 + scale = 1.0 + input_scale = 2.0 + + replace_node.args = (old_args[0],) + + updated_kwargs = dict(node.kwargs) + updated_kwargs["alpha"] = int(alpha) + updated_kwargs["scale"] = int(scale) + updated_kwargs["input_scale"] = int(input_scale) + + replace_node.kwargs = updated_kwargs + + node.replace_all_uses_with(replace_node) + graph.erase_node(node) + + modified_graph = True + if modified_graph: + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, modified_graph) diff --git a/backends/arm/_passes/convert_expand_copy_to_repeat.py b/backends/arm/_passes/convert_expand_copy_to_repeat.py index 5632c253437..ee509c7ebb5 100644 --- a/backends/arm/_passes/convert_expand_copy_to_repeat.py +++ b/backends/arm/_passes/convert_expand_copy_to_repeat.py @@ -8,12 +8,43 @@ import logging from typing import cast +import torch + from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass logger = logging.getLogger(__name__) +def calculate_multiples(args): + input_node_or_tensor = args[0] + + if isinstance(input_node_or_tensor, torch.fx.node.Node): + input_data = input_node_or_tensor.meta["val"] + else: + input_data = input_node_or_tensor.data + + input_shape = input_data.shape + + multiples = cast(list[int], args[1]) + expanded_rank = len(multiples) + + # Expanded shape is 'input_shape' front-padded with ones. + padding = expanded_rank - len(input_shape) + extended_shape = [ + input_shape[i] if i >= 0 else 1 for i in range(-padding, len(input_shape)) + ] + + # To convert expand arg to repeat arg, non-repeated dims should have + # multiples[dim] = 1. Passing -1 to expand arg means + # not changing the size of that dimension. + multiples = [ + multiples[i] if multiples[i] != -1 and extended_shape[i] == 1 else 1 + for i in range(expanded_rank) + ] + return multiples + + class ConvertExpandCopyToRepeatPass(ExportPass): """ Replace expand copy with repeat since it is a repeat that can only repeat singleton dimensions. @@ -26,23 +57,7 @@ def call_operator(self, op, args, kwargs, meta): if op != self.expand_copy: return super().call_operator(op, args, kwargs, meta) - input_shape = args[0].data.shape - multiples = cast(list[int], args[1]) - expanded_rank = len(multiples) - - # Expanded shape is 'input_shape' front-padded with ones. - padding = expanded_rank - len(input_shape) - extended_shape = [ - input_shape[i] if i >= 0 else 1 for i in range(-padding, len(input_shape)) - ] - - # To convert expand arg to repeat arg, non-repeated dims should have - # multiples[dim] = 1. Passing -1 to expand arg means - # not changing the size of that dimension. - multiples = [ - multiples[i] if multiples[i] != -1 and extended_shape[i] == 1 else 1 - for i in range(expanded_rank) - ] + multiples = calculate_multiples(args) if all((x == 1 for x in multiples)): # All dimensions/repetitions occur only once. Remove node diff --git a/backends/arm/_passes/convert_int64_const_ops_to_int32.py b/backends/arm/_passes/convert_int64_const_ops_to_int32.py new file mode 100644 index 00000000000..704c89dbd78 --- /dev/null +++ b/backends/arm/_passes/convert_int64_const_ops_to_int32.py @@ -0,0 +1,74 @@ +# Copyright 2025 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-unsafe + + +import logging + +import torch +from executorch.backends.arm._passes.fuse_constant_ops_pass import ComputeConstantOpsAOT +from executorch.exir.pass_base import ExportPass, PassResult + + +logger = logging.getLogger(__name__) +INT32_MIN = torch.iinfo(torch.int32).min +INT32_MAX = torch.iinfo(torch.int32).max + + +class ConvertInt64ConstOpsToInt32Pass(ExportPass): + """ + Rewrite constant ops that produce int64 to int32 where safe. + + List of supported operatos: + 1. `torch.full` + 2. `torch.arange` + 3. `torch.eye` + 4. `torch.linspace` + 5. `torch.tensor` + """ + + torch_ops = [ + torch.ops.aten.full.default, + torch.ops.aten.arange.default, + torch.ops.aten.arange.start, + torch.ops.aten.arange.start_step, + torch.ops.aten.eye.default, + torch.ops.aten.linspace.default, + ] + + def call(self, graph_module: torch.fx.GraphModule): + modified = False + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + + if node.target not in ComputeConstantOpsAOT.targeted_ops + self.torch_ops: + continue + + data = node.target(*node.args, **node.kwargs) + if data.dtype is not torch.int64: + continue + + min_val, max_val = torch.min(data), torch.max(data) + if INT32_MIN <= min_val and max_val <= INT32_MAX: + logger.warning( + f"Casting {node.name} from torch.int64 to torch.int32" + f" defined in {node.meta.get('stack_trace','[no stack trace found]')}" + ) + node.update_kwarg("dtype", torch.int32) + modified = True + else: + logger.warning( + f"[{node.name}] has values: min={min_val}, max={max_val}, which exceeds int32 range " + f"([{INT32_MIN}, {INT32_MAX}]); not converting dtype to int32." + ) + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/convert_int64_output_ops_to_int32.py b/backends/arm/_passes/convert_int64_output_ops_to_int32.py new file mode 100644 index 00000000000..788201be6c8 --- /dev/null +++ b/backends/arm/_passes/convert_int64_output_ops_to_int32.py @@ -0,0 +1,153 @@ +# Copyright 2025 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-unsafe + + +import logging + +import torch +from executorch.backends.arm._passes.arm_pass_utils import ( + create_node, + get_first_fake_tensor, + set_node_arg, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +logger = logging.getLogger(__name__) + + +class ConvertInt64OutputOpsToInt32Pass(ExportPass): + """ + Rewrites or removes operations that produce int64 outputs, converting them + to int32 where possible. + + + Currently, this pass handles casting and argmax operators: + 1. int32 -> int64: + removes the cast and redirects all uses to the original int32 value. + 2. other types -> int64: + rewrites the cast to produce int32 instead of int64. + 3. torch.argmax() + insert an int64->int32 cast after the argmax node + + Future extensions may include operators that return int64 outputs by default + (e.g., `argmin`), rewriting them or inserting an int64 -> int32 cast to yield + int32 results. + + Note: Overflow checks are applied selectively in this pass. For operators without + such checks, it is the user's responsibility to ensure that values fit within + the int32 range. + """ + + aten_cast_ops = ( + torch.ops.aten.to.dtype, + torch.ops.aten.to.dtype_layout, + ) + edge_cast_ops = (exir_ops.edge.dim_order_ops._to_dim_order_copy.default,) + + aten_argmax_ops = (torch.ops.aten.argmax.default,) + edge_argmax_ops = (exir_ops.edge.aten.argmax.default,) + + aten_ops = aten_cast_ops + aten_argmax_ops + edge_ops = edge_cast_ops + edge_argmax_ops + + # dtype is specified in args + cast_ops_args = ( + torch.ops.aten.to.dtype, # to_2: node.args: (gt, torch.int64) node.kwargs: {} + ) + # dtype is specified in kwargs + cast_ops_kwargs = ( + torch.ops.aten.to.dtype_layout, # to_1: node.args: (unsqueeze,) node.kwargs: {'dtype': torch.int64, 'layout': torch.strided, 'device': device(type='cpu')} + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, # node.args: (aten_gt_scalar,) node.kwargs: {'dtype': torch.int64, 'dim_order': [0, 1]} + ) + + def _get_decomposition(self, op): + if op in self.edge_ops: + return exir_ops.edge.dim_order_ops._to_dim_order_copy.default + + if op in self.aten_ops: + return torch.ops.dim_order_ops._to_dim_order_copy.default + + raise RuntimeError( + f"[{self.__class__.__name__}] Can't get decomposition for op {op}" + ) + + def _convert_casting_operators(self, node: torch.fx.Node): + input_node = node.all_input_nodes[0] + input_dtype = get_first_fake_tensor(input_node).dtype + # Case 1: int32 -> int64 - removes the ops + if input_dtype == torch.int32: + users = [user for user in node.users if node != user] + for user in users: + logger.warning( + f"Removing int32->int64 casting node {node.name} defined in" + f" {node.meta.get('stack_trace','[no stack trace found]')}" + ) + user.replace_input_with(node, input_node) + # Case 2: other types -> int64 - rewrites to cast to int32 + else: + if node.target in self.cast_ops_kwargs: + set_node_arg(node, "dtype", torch.int32) + elif node.target in self.cast_ops_args: + set_node_arg(node, 1, torch.int32) + else: + raise RuntimeError(f"Unexpected target {node.target} in {node.name}") + output_dtype = get_first_fake_tensor(node).dtype + logger.warning( + f"Converting casting node {node.name} from {input_dtype}->{output_dtype} to" + f" {input_dtype}->torch.int32 defined in {node.meta.get('stack_trace','[no stack trace found]')}" + ) + + def _convert_argmax_operators(self, node: torch.fx.Node, graph: torch.fx.Graph): + output_tensor = node + to_copy_op = self._get_decomposition(node.target) + with graph.inserting_after(node): + cast_after = create_node( + graph, + to_copy_op, + args=(output_tensor,), + kwargs={ + "dtype": torch.int32, + }, + ) + users = [user for user in node.users if user != cast_after] + for user in users: + user.replace_input_with(output_tensor, cast_after) + logger.warning( + f"Inserting a casting node {cast_after.name} after {node.name} to cast int64 output" + f" to int32 for {node.name} defined in {node.meta.get('stack_trace','[no stack trace found]')}" + ) + + def call(self, graph_module: torch.fx.GraphModule): + modified = False + graph = graph_module.graph + for node in list(graph.nodes): + if node.op != "call_function": + continue + if node.target not in self.aten_ops + self.edge_ops: + continue + output_dtype = get_first_fake_tensor(node).dtype + if output_dtype != torch.int64: + continue + + if node.target in self.aten_cast_ops + self.edge_cast_ops: + self._convert_casting_operators(node) + elif node.target in self.aten_argmax_ops + self.edge_argmax_ops: + # TODO: Add range check based on the input tensor shape before casting the output + self._convert_argmax_operators(node, graph) + else: + raise RuntimeError(f"Unexpected target {node.target} in {node.name}") + + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/decompose_asin_pass.py b/backends/arm/_passes/decompose_asin_and_acos_pass.py similarity index 72% rename from backends/arm/_passes/decompose_asin_pass.py rename to backends/arm/_passes/decompose_asin_and_acos_pass.py index 0c0bcdf7f49..e067f17b0ca 100644 --- a/backends/arm/_passes/decompose_asin_pass.py +++ b/backends/arm/_passes/decompose_asin_and_acos_pass.py @@ -15,10 +15,11 @@ # For MI case edge_asin_op = (exir_ops.edge.aten.asin.default,) +edge_acos_op = (exir_ops.edge.aten.acos.default,) -def get_asin_decomposition(op) -> tuple: - if op in edge_asin_op: +def get_decomposition(op) -> tuple: + if op in (edge_asin_op + edge_acos_op): return ( exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.add.Tensor, @@ -31,25 +32,26 @@ def get_asin_decomposition(op) -> tuple: exir_ops.edge.aten.lt.Scalar, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.full_like.default, - exir_ops.edge.aten.where.self, exir_ops.edge.aten.neg.default, ) - raise RuntimeError(f"Can't get asin decomposition for op {op}") + raise RuntimeError(f"Can't get decomposition for op {op}") -class DecomposeAsinPass(ArmPass): +class DecomposeAsinAndAcosPass(ArmPass): """ - This pass decomposes asin into a rational approximation for small values + This pass decomposes asin and acos into a rational approximation for small values and a transformed rational approximation for large values. - Example: - y = asin(x) - Becomes: + + The decomposition is based on the following mathematical identities: if abs(x) < 0.5: - y = x + P(x^2) / Q(x^2) + asin(x) = x + P(x^2) / Q(x^2) + acos(x) = π/2 - asin(x) else: - y = π/2 - 2 * (s + s^3 * Q(z) / P(z)) - where P and Q are polynomials defined in the function. + asin(x) = π/2 - 2 * (s + s^3 * Q(z) / P(z)) + acos(x) = 2 * (s + s^3 * Q(z) / P(z)) + where P and Q are polynomials defined in the function and s is the square root of z. + """ def _build_polynomial( @@ -84,11 +86,25 @@ def _build_polynomial( ) return result + def _combine_branches( + self, + bool_op, + bool_args: tuple[torch.Tensor, float], + branches: tuple[torch.Tensor, torch.Tensor], + meta: dict[str, str], + ) -> torch.Tensor: + where_op = exir_ops.edge.aten.where.self + mask = super().call_operator(bool_op, bool_args, {}, meta, True) + branch_true, branch_false = branches + return super().call_operator( + where_op, (mask, branch_true, branch_false), {}, meta, True + ) + def call_operator(self, op, args, kwargs, meta): - if op not in edge_asin_op: + if op not in (edge_asin_op + edge_acos_op): return super().call_operator(op, args, kwargs, meta) logging.info( - f"Approximating asin. This may introduce small numerical errors. For details, see {__file__}." + f"Approximating {op}. This may introduce small numerical errors. For details, see {__file__}." ) x = args[0] half = 0.5 @@ -111,9 +127,8 @@ def call_operator(self, op, args, kwargs, meta): lt_op, sub_op, full_like_op, - where_op, neg_op, - ) = get_asin_decomposition(op) + ) = get_decomposition(op) # Coefficients for the rational approximation, calculated with the Minimax (Remez) method p_coefficients = [ @@ -129,7 +144,6 @@ def call_operator(self, op, args, kwargs, meta): x_abs = super().call_operator(abs_op, (x,), {}, meta, True) # Step 1: compute asin_small - rational approximation for [0,0.5] - y = super().call_operator(mul_op, (x_abs, x_abs), {}, meta, True) x3 = super().call_operator(mul_op, (x_abs, y), {}, meta, True) @@ -154,47 +168,40 @@ def call_operator(self, op, args, kwargs, meta): Qz = self._build_polynomial(q_coefficients, z, meta) numer = super().call_operator(mul_op, (s3, Pz), {}, meta, True) + # Calculate r_large = P(z) / Q(z) r_large = super().call_operator(div_op, (numer, Qz), {}, meta, True) # Calculate asin_large = pi/2 - 2 * (s + s^3 * Q(z) / P(z)) t1 = super().call_operator(add_op, (s, r_large), {}, meta, True) t2 = super().call_operator(mul_op_scalar, (t1, two), {}, meta, True) + diff = super().call_operator(sub_op_scalar, (t2, pi_over_2), {}, meta, True) tmp_neg_ones = super().call_operator( full_like_op, (diff, neg_one), {}, meta, True ) asin_large = super().call_operator(mul_op, (diff, tmp_neg_ones), {}, meta, True) - # Combine branches - is_large = super().call_operator(gt_op, (x_abs, half), {}, meta, True) - asin_unsigned = super().call_operator( - where_op, - ( - is_large, - asin_large, - asin_small, - ), - {}, - meta, - True, + asin_unsigned = self._combine_branches( + gt_op, (x_abs, half), (asin_large, asin_small), meta ) # Handle x < 0 - is_neg = super().call_operator(lt_op, (x, zero), {}, meta, True) - # Compute -asin_unsigned negated_asin = super().call_operator(neg_op, (asin_unsigned,), {}, meta, True) - # Combine branches for signed asin - asin_signed = super().call_operator( - where_op, - ( - is_neg, - negated_asin, - asin_unsigned, - ), - {}, - meta, - True, + asin = self._combine_branches( + lt_op, (x, zero), (negated_asin, asin_unsigned), meta ) - return asin_signed + if op in edge_acos_op: + # If x <= 0.5: acos(x) = pi/2 - asin(x) + const_tensor = super().call_operator( + full_like_op, (x, pi_over_2), {}, meta, True + ) + acos_small = super().call_operator( + sub_op, (const_tensor, asin), {}, meta, True + ) + # If x > 0.5, acos(x) = 2 * (s + s^3 * Q(z) / P(z)) = t2 + acos = self._combine_branches(gt_op, (x, half), (t2, acos_small), meta) + return acos + + return asin diff --git a/backends/arm/_passes/decompose_avg_pool2d.py b/backends/arm/_passes/decompose_avg_pool2d.py index 0eb3ce34ecd..21ed6b518c7 100644 --- a/backends/arm/_passes/decompose_avg_pool2d.py +++ b/backends/arm/_passes/decompose_avg_pool2d.py @@ -45,7 +45,10 @@ def call_operator(self, op, args, kwargs, meta): x = args[0] kernel_h, kernel_w = args[1] kernel_size = kernel_h * kernel_w - stride_h, stride_w = args[2] + if len(args) > 2 and args[2] is not None: + stride_h, stride_w = args[2] + else: + stride_h, stride_w = kernel_h, kernel_w pad_h, pad_w = new_pad_h, new_pad_w = args[3] if len(args) > 3 else (0, 0) ceil_mode = args[4] if len(args) > 4 else False count_include_pad = args[5] if len(args) > 5 else True @@ -108,7 +111,14 @@ def call_operator(self, op, args, kwargs, meta): x = super().call_operator(cat_op, (cat_nodes, 2), kwargs, meta) new_pad_h = 0 - avgpool_args = (x, args[1], args[2], [new_pad_h, new_pad_w], ceil_mode, False) + avgpool_args = ( + x, + args[1], + [stride_h, stride_w], + [new_pad_h, new_pad_w], + ceil_mode, + False, + ) x = super().call_operator(avgpool_op, avgpool_args, kwargs, meta) # Multiply by factor (kernel_size / divisor_override) if divisor_override diff --git a/backends/arm/_passes/decompose_cosh_pass.py b/backends/arm/_passes/decompose_cosh_pass.py new file mode 100644 index 00000000000..a94cf9ecff0 --- /dev/null +++ b/backends/arm/_passes/decompose_cosh_pass.py @@ -0,0 +1,48 @@ +# Copyright 2025 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.arm._passes import ArmPass +from executorch.exir.dialects._ops import ops as exir_ops + +# For MI case +edge_cosh = exir_ops.edge.aten.cosh.default + + +class DecomposeCoshPass(ArmPass): + """ + This pass replaces the cosh operator with a sequence of TOSA-equivalent operations that + compute the hyperbolic cosine using the formula: + + cosh(x) = 0.5 * (e^x + e^(-x)) + + """ + + def call_operator(self, op, args, kwargs, meta, updated=False): + if op is not edge_cosh: + return super().call_operator(op, args, kwargs, meta, updated) + + x = args + + exp_op, mul_op, neg_op, add_op = ( + exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.mul.Scalar, + exir_ops.edge.aten.neg.default, + exir_ops.edge.aten.add.Tensor, + ) + + # exp1 = e^x + exp1 = super().call_operator(exp_op, x, {}, meta, updated=True) + + # exp2 = e^(⁻x) + neg_x = super().call_operator(neg_op, x, {}, meta, updated=True) + exp2 = super().call_operator(exp_op, (neg_x,), {}, meta, updated=True) + + # numer = exp1 + exp2 + numer = super().call_operator(add_op, (exp1, exp2), {}, meta, updated=True) + + # out = 0.5 * numer + out = super().call_operator(mul_op, (numer, 0.5), {}, meta, updated=True) + + return out diff --git a/backends/arm/_passes/decompose_cumsum_pass.py b/backends/arm/_passes/decompose_cumsum_pass.py new file mode 100644 index 00000000000..155ccd11594 --- /dev/null +++ b/backends/arm/_passes/decompose_cumsum_pass.py @@ -0,0 +1,142 @@ +# Copyright 2025 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 math import prod + +import torch +from executorch.backends.arm._passes import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.backends.arm._passes.quant_args import QuantArgs + +from executorch.backends.transforms.utils import create_constant_placeholder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.export.graph_signature import InputKind + + +class DecomposeCumsumPass(ArmPass): + """ + Decomposes cumsum into a 1D convolution with a kernel of ones. + + For example, the cumsum of an input tensor [1, 1] is [1, 1 + 1] = [1, 2]. + To decompose this, take the input tensor and pre-padded with len(input)-1 zeros and + slided over with a kernel [1,1], of length len(input): + + Input: [0, 1, 1] + Kernel: [1, 1] = [1] + [1, 1] = [2] + + Since pytorch only supports symmetric padding, in reality the result will have + an additional 1 calculated at the end, which leads to an required extra slice op. + + To extend this to higher dimensions, the input is reshaped to [N, C, H, W] with + N = + C = 1 + H = + W = + And the convolution is applied over dimension H. + """ + + def call(self, graph_module): + graph = graph_module.graph + targets = (exir_ops.edge.aten.cumsum.default, torch.ops.aten.cumsum.default) + modified = False + for node in list(graph.nodes): + if node.op != "call_function" or node.target not in targets: + continue + + if len(node.args) != 2: + raise ValueError( + "Cumsum node should have exactly two arguments: input and dim." + ) + + # Get node data + input_node, dim = node.args + val = node.meta.get("val") + original_shape = list(val.shape) + dtype = input_node.meta.get("val").dtype + dim = dim % len(original_shape) + + # Compute shapes + pre_cumsum_dim = prod(original_shape[:dim]) if dim > 0 else 1 + cumsum_dim = original_shape[dim] + post_cumsum_dim = ( + prod(original_shape[dim + 1 :]) if dim < len(original_shape) - 1 else 1 + ) + conv_shape = [ + pre_cumsum_dim, + 1, + cumsum_dim, + post_cumsum_dim, + ] + pad_shape = [original_shape[dim] - 1, 0] + weight_shape = [1, 1, original_shape[dim], 1] + + # Create convolution weight + with graph.inserting_before(list(graph.nodes)[0]): + weight_data = torch.ones(size=weight_shape, dtype=dtype) + weight_node = create_constant_placeholder( + self.exported_program, + graph, + node.name + "_kernel", + InputKind.PARAMETER, + weight_data, + ) + + # Create decomposed nodes + view_op = exir_ops.edge.aten.view_copy.default + conv_op = exir_ops.edge.aten.convolution.default + slice_op = exir_ops.edge.aten.slice_copy.Tensor + with graph.inserting_before(node): + # Reshape to 4D with + view_args = (input_node, conv_shape) + view_node = create_node(graph, view_op, args=view_args, from_node=node) + + conv_args = ( + view_node, + weight_node, + None, + [1, 1], + pad_shape, + [1, 1], + False, + [0], + 1, + ) + conv_node = create_node(graph, conv_op, args=conv_args, from_node=node) + + # The convolution is inserted after quantization, so we need to set our + # own quantization parameters for the weights here. However since the + # data is ones directly created as int8, they already have correct scale + # and so no scaling needs to be done, i.e. set scale=1.0, zero_point=0.0 + if ( + "input_qparams" in conv_node.meta + and len(conv_node.meta["input_qparams"]) > 0 + ): + qparams = QuantArgs(1.0, 0.0, -128, 127, torch.int8) + conv_node.meta["input_qparams"][1] = qparams + + slice_args = (conv_node, 2, 0, original_shape[dim]) + slice_node = create_node( + graph, slice_op, args=slice_args, from_node=node + ) + + view_original_args = (slice_node, original_shape) + view_original_node = create_node( + graph, view_op, args=view_original_args, from_node=node + ) + + # Replace and remove original + node.replace_all_uses_with(view_original_node) + graph.erase_node(node) + modified = True + + if modified: + # Cleanup + graph.eliminate_dead_code() + graph_module.recompile() + # Apply any operator-level transforms + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/decompose_div_tensor_mode.py b/backends/arm/_passes/decompose_div_tensor_mode.py new file mode 100644 index 00000000000..0e6b40afbb2 --- /dev/null +++ b/backends/arm/_passes/decompose_div_tensor_mode.py @@ -0,0 +1,84 @@ +# Copyright 2025 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-unsafe + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + +edge_div_mode_ops = (exir_ops.edge.aten.div.Tensor_mode,) +aten_div_mode_ops = (torch.ops.aten.div.Tensor_mode,) + +edge_unary = { + "div": exir_ops.edge.aten.div.Tensor, + "floor": exir_ops.edge.aten.floor.default, + "ceil": exir_ops.edge.aten.ceil.default, + "full": exir_ops.edge.aten.full.default, + "lt": exir_ops.edge.aten.lt.Tensor, + "where": exir_ops.edge.aten.where.self, +} + +aten_unary = { + "div": torch.ops.aten.div.Tensor, + "floor": torch.ops.aten.floor.default, + "ceil": torch.ops.aten.ceil.default, + "full": torch.ops.aten.full.default, + "lt": torch.ops.aten.lt.Tensor, + "where": torch.ops.aten.where.self, +} + + +def _get_opset(op): + if op in edge_div_mode_ops: + return edge_unary + if op in aten_div_mode_ops: + return aten_unary + raise RuntimeError(f"div.Tensor_mode not supported for op {op}") + + +class DecomposeDivTensorModePass(ExportPass): + """ + Rewrites aten.div.Tensor_mode into + + rounding_mode=None -> div(a, b) + rounding_mode='floor' -> floor(div(a, b)) + rounding_mode='trunc' -> where(div(a,b) < 0, ceil(div(a,b)), floor(div(a,b))) + """ + + def call_operator(self, op, args, kwargs, meta): + if op not in (edge_div_mode_ops + aten_div_mode_ops): + return super().call_operator(op, args, kwargs, meta) + + opset = _get_opset(op) + + a, b = args[0], args[1] + rounding_mode = kwargs.get("rounding_mode", None) + if rounding_mode is None and len(args) > 2: + rounding_mode = args[2] + + q = super().call_operator(opset["div"], (a, b), {}, meta) + + if rounding_mode is None: + return q + + if rounding_mode == "floor": + return super().call_operator(opset["floor"], (q,), {}, meta) + + if rounding_mode == "trunc": + zero = super().call_operator( + opset["full"], + args=((1,) * len(meta["val"].size()), 0.0), + kwargs={"dtype": torch.float32}, + meta=meta, + ) + lt0 = self.call_operator(opset["lt"], (q, zero), {}, meta) + ceilq = self.call_operator(opset["ceil"], (q,), {}, meta) + floorq = self.call_operator(opset["floor"], (q,), {}, meta) + return self.call_operator(opset["where"], (lt0, ceilq, floorq), {}, meta) + + raise RuntimeError( + f"Unsupported rounding_mode for div.Tensor_mode: {rounding_mode!r}" + ) diff --git a/backends/arm/_passes/decompose_elu_pass.py b/backends/arm/_passes/decompose_elu_pass.py new file mode 100644 index 00000000000..743f1b46f4d --- /dev/null +++ b/backends/arm/_passes/decompose_elu_pass.py @@ -0,0 +1,85 @@ +# Copyright 2025 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.arm._passes import ArmPass +from executorch.exir.dialects._ops import ops as exir_ops + +edge_elu_ops = (exir_ops.edge.aten.elu.default,) + + +def get_elu_decomposition(op) -> tuple: + """ + Returns the decomposition of the given aten.elu operation into + its equivalent TOSA-supported operations + + This handles both edge dialect ops and core PyTorch ops. The decomposition strategy + is: + elu(x, y) → where(greater_or_eq(x, 0), (exp(x)-1), x) + + Returns: + A tuple (expm1_op, ge_op, where_op, mul_op) corresponding to the appropriate operator + overloads for the input op. + + Raises: + RuntimeError: If the provided operator is not a supported elu variant. + """ + + if op in edge_elu_ops: + return ( + exir_ops.edge.aten.expm1.default, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.where.self, + exir_ops.edge.aten.mul.Scalar, + ) + + raise RuntimeError(f"Can't get elu decomposition for op {op}") + + +class DecomposeEluPass(ArmPass): + """ + A transformation pass that decomposes unsupported 'aten.elu' operations + into a combination of supported TOSA-equivalent operations. + + Since TOSA does not provide a native ELU operator, this pass rewrites: + elu(x) → where(greater_or_eq(x, 0), (alpha*(exp(x)-1)), x) + + Supported input ops: + - exir_ops.edge.aten.elu.Tensor(x) + + These are replaced with: + - exir_ops.edge.aten.expm1.default + - exir_ops.edge.aten.ge.Scalar + - exir_ops.edge.aten.where.self + - exir_ops.edge.aten.mul.Scalar + """ + + def call_operator(self, op, args, kwargs, meta): + if op not in edge_elu_ops: + return super().call_operator(op, args, kwargs, meta, updated=False) + + ( + expm1_op, + ge_op, + where_op, + mul_op, + ) = get_elu_decomposition(op) + + input = args[0] + alpha = args[1] if len(args) > 1 else 1.0 + + if alpha == 0: + relu_op = exir_ops.edge.aten.relu.default + return super().call_operator(relu_op, (input,), {}, meta, updated=True) + + expm1_node = super().call_operator(expm1_op, (input,), {}, meta, updated=True) + mul_node = super().call_operator( + mul_op, (expm1_node, alpha), {}, meta, updated=True + ) + ge_node = super().call_operator(ge_op, (input, 0.0), {}, meta, updated=True) + where_node = super().call_operator( + where_op, (ge_node, input, mul_node), {}, meta, updated=True + ) + + return where_node diff --git a/backends/arm/_passes/decompose_expm1_pass.py b/backends/arm/_passes/decompose_expm1_pass.py new file mode 100644 index 00000000000..5b1b90495b5 --- /dev/null +++ b/backends/arm/_passes/decompose_expm1_pass.py @@ -0,0 +1,135 @@ +# Copyright 2025 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.arm._passes import ArmPass +from executorch.exir.dialects._ops import ops as exir_ops + + +edge_expm1_ops = (exir_ops.edge.aten.expm1.default,) # MI case + + +def _get_expm1_decomposition(op) -> tuple: + """ + Returns the decomposition of the given aten.expm1 operation into + its equivalent TOSA-supported operations + + This handles both edge dialect ops and core PyTorch ops. The decomposition strategy + is: + expm1(x) → where(and(ge(x, -0.35), le(x, 0.35)), {taylor_series_expansion}, (exp(x)-1)) + + where {taylor_series_expansion} = x + (x^2/2) + (x^3/6) + (x^4/24) + + Returns: + A tuple (op_pow, op_div, op_add, op_exp, op_sub, op_ge, op_where, op_le, op_and) + corresponding to the appropriate operator overloads for the input op. + + Raises: + RuntimeError: If the provided operator is not a supported elu variant. + """ + if op in edge_expm1_ops: + return ( + exir_ops.edge.aten.pow.Tensor_Scalar, + exir_ops.edge.aten.div.Scalar, + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.sub.Scalar, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.where.self, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.logical_and.default, + ) + + raise RuntimeError(f"Can't get expm1 decomposition for op {op}") + + +class DecomposeExpm1Pass(ArmPass): + """ + A transformation pass that decomposes unsupported 'aten.expm1' operations + into a combination of supported TOSA-equivalent operations. + + Since TOSA does not provide a native expm1 operator, this pass rewrites: + expm1(x) → where(and(ge(x, -0.35), le(x, 0.35)), {taylor_series_expansion}, (exp(x)-1)) + where {taylor_series_expansion} = x + (x^2/2) + (x^3/6) + (x^4/24) + + Supported input ops: + - exir_ops.edge.aten.expm1.default(x) + + These are replaced with: + - exir_ops.edge.aten.pow.Tensor_Scalar, + - exir_ops.edge.aten.div.Scalar, + - exir_ops.edge.aten.add.Tensor, + - exir_ops.edge.aten.exp.default, + - exir_ops.edge.aten.sub.Scalar, + - exir_ops.edge.aten.ge.Scalar, + - exir_ops.edge.aten.where.self, + - exir_ops.edge.aten.le.Scalar, + - exir_ops.edge.aten.logical_and.default + """ + + def call_operator(self, op, args, kwargs, meta): + if op not in edge_expm1_ops: + return super().call_operator(op, args, kwargs, meta, updated=False) + + ( + op_pow, + op_div, + op_add, + op_exp, + op_sub, + op_ge, + op_where, + op_le, + op_and, + ) = _get_expm1_decomposition(op) + + input = args[0] + + cutlo = -0.35 + cuthi = 0.35 + + taylor_term_2_numerator = super().call_operator( + op_pow, (input, 2), {}, meta, updated=False + ) + taylor_term_3_numerator = super().call_operator( + op_pow, (input, 3), {}, meta, updated=False + ) + taylor_term_4_numerator = super().call_operator( + op_pow, (input, 4), {}, meta, updated=False + ) + + taylor_term_2 = super().call_operator( + op_div, (taylor_term_2_numerator, 2), {}, meta, updated=False + ) + taylor_term_3 = super().call_operator( + op_div, (taylor_term_3_numerator, 6), {}, meta, updated=False + ) + taylor_term_4 = super().call_operator( + op_div, (taylor_term_4_numerator, 24), {}, meta, updated=False + ) + + add_terms_1_2 = super().call_operator( + op_add, (input, taylor_term_2), {}, meta, updated=False + ) + add_term_3 = super().call_operator( + op_add, (add_terms_1_2, taylor_term_3), {}, meta, updated=False + ) + taylor_expansion = super().call_operator( + op_add, (add_term_3, taylor_term_4), {}, meta, updated=False + ) + + decomp_exp = super().call_operator(op_exp, (input,), {}, meta, updated=False) + decomp_sub = super().call_operator( + op_sub, (decomp_exp, 1.0), {}, meta, updated=False + ) + + ge = super().call_operator(op_ge, (input, cutlo), {}, meta, updated=False) + le = super().call_operator(op_le, (input, cuthi), {}, meta, updated=False) + + cond_and = super().call_operator(op_and, (ge, le), {}, meta, updated=False) + where = super().call_operator( + op_where, (cond_and, taylor_expansion, decomp_sub), {}, meta, updated=True + ) + + return where diff --git a/backends/arm/_passes/decompose_glu_pass.py b/backends/arm/_passes/decompose_glu_pass.py new file mode 100644 index 00000000000..183dc89cf61 --- /dev/null +++ b/backends/arm/_passes/decompose_glu_pass.py @@ -0,0 +1,75 @@ +# Copyright 2025 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 ArmPass +from executorch.exir.dialects._ops import ops as exir_ops + + +# For FP case +edge_glu = exir_ops.edge.aten.glu.default + +# For INT case +aten_glu = torch.ops.aten.glu.default + + +def get_ops(op): + """Returns the appropriate operator functions based on the input operator.""" + if op == edge_glu: + return ( + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.slice_copy.Tensor, + ) + elif op == aten_glu: + return ( + torch.ops.aten.mul.Tensor, + torch.ops.aten.sigmoid.default, + torch.ops.aten.slice_copy.Tensor, + ) + else: + raise ValueError(f"Unsupported operator: {op}") + + +class DecomposeGluPass(ArmPass): + """Decomposes the GLU operator into hadamard product and sigmoid.""" + + def call_operator(self, op, args, kwargs, meta): + if op not in [edge_glu, aten_glu]: + return super().call_operator(op, args, kwargs, meta) + + hadamard_prod, sigmoid, slice_op = get_ops(op) + X = args[0] + + dim = args[1] if len(args) > 1 else kwargs.get("dim", -1) + + if "val" not in X.node.meta: + raise Exception("Could not get dimension metadata in input.") + + if dim < 0: + dim += X.node.meta["val"].dim() + + n = X.node.meta["val"].size(dim) + + if n % 2: + raise RuntimeError( + f"glu expects an even split along dim={dim}, got size {n}" + ) + + middle = n // 2 + + T1 = super().call_operator( + slice_op, (X, dim, 0, middle), {}, meta, updated=True + ) + + T2 = super().call_operator( + slice_op, (X, dim, middle, n), {}, meta, updated=True + ) + + T2_sigmoid = super().call_operator(sigmoid, (T2,), {}, meta, updated=True) + + return super().call_operator( + hadamard_prod, (T1, T2_sigmoid), {}, meta, updated=True + ) diff --git a/backends/arm/_passes/decompose_grouped_conv.py b/backends/arm/_passes/decompose_grouped_conv.py index 6bfdf4dea5e..ce9fe9c9937 100644 --- a/backends/arm/_passes/decompose_grouped_conv.py +++ b/backends/arm/_passes/decompose_grouped_conv.py @@ -6,7 +6,7 @@ from copy import copy import torch -from executorch.backends.arm.tosa_quant_utils import QuantArgs +from executorch.backends.arm._passes.quant_args import QuantArgs from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass diff --git a/backends/arm/_passes/decompose_linear_pass.py b/backends/arm/_passes/decompose_linear_pass.py index 76207e5849c..3d154d9b81e 100644 --- a/backends/arm/_passes/decompose_linear_pass.py +++ b/backends/arm/_passes/decompose_linear_pass.py @@ -90,6 +90,13 @@ def call(self, graph_module): kwargs={}, from_node=node, ) + # Quantization parameters are inherited from original linear node, but + # output reshape should use the linear node's output qparams for both input + # and output. + if "input_qparams" in output.meta: + output.meta["input_qparams"] = output.meta.get( + "output_qparams", None + ) node.replace_all_uses_with(output) graph_module.graph.erase_node(node) diff --git a/backends/arm/_passes/decompose_logit_pass.py b/backends/arm/_passes/decompose_logit_pass.py new file mode 100644 index 00000000000..40e2b22cb54 --- /dev/null +++ b/backends/arm/_passes/decompose_logit_pass.py @@ -0,0 +1,96 @@ +# Copyright 2025 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 ArmPass +from executorch.exir.dialects._ops import ops as exir_ops + + +# For FP case +edge_logit = exir_ops.edge.aten.logit.default +# For INT case +aten_logit = torch.ops.aten.logit.default + + +def get_ops(op): + """Returns the appropriate operator functions based on the input operator.""" + if op == edge_logit: + return ( + exir_ops.edge.aten.log.default, + exir_ops.edge.aten.add.Scalar, + exir_ops.edge.aten.reciprocal.default, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.mul.Scalar, + exir_ops.edge.aten.clamp.default, + ) + elif op == aten_logit: + return ( + torch.ops.aten.log.default, + torch.ops.aten.add.Scalar, + torch.ops.aten.reciprocal.default, + torch.ops.aten.mul.Tensor, + torch.ops.aten.mul.Scalar, + torch.ops.aten.clamp.default, + ) + else: + raise ValueError(f"Unsupported operator: {op}") + + +class DecomposeLogitPass(ArmPass): + """ + Decomposes the `logit` operator into a sequence of primitive operations. + + If `eps` is provided, the input tensor `x` is first clamped to the range + [eps, 1 - eps]. + + The decomposition follows the identity: + + logit(x) = log(x / (1 - x)) + + Examples: + + logit(x) becomes: + log(x * reciprocal((-1) * x + 1)) + + logit(x, eps) becomes: + y = clamp(x, eps, 1 - eps) + log(y * reciprocal((-1) * y + 1)) + """ + + def call_operator(self, op, args, kwargs, meta): + if op not in [edge_logit, aten_logit]: + return super().call_operator(op, args, kwargs, meta) + + X = args[0] + eps = args[1] if len(args) > 1 else kwargs.get("eps", None) + + ( + log_op, + add_scalar_op, + recip_op, + mul_tensor_op, + mul_scalar_op, + clamp_op, + ) = get_ops(op) + + if eps is not None: + X = super().call_operator( + clamp_op, (X, eps, 1.0 - eps), {}, meta, updated=True + ) + + neg_X = super().call_operator(mul_scalar_op, (X, -1.0), {}, meta, updated=True) + + denom = super().call_operator( + add_scalar_op, (neg_X, 1.0), {}, meta, updated=True + ) + + frac = super().call_operator(recip_op, (denom,), {}, meta, updated=True) + + log_input = super().call_operator( + mul_tensor_op, (X, frac), {}, meta, updated=True + ) + + return super().call_operator(log_op, (log_input,), {}, meta, updated=True) diff --git a/backends/arm/_passes/decompose_meandim_pass.py b/backends/arm/_passes/decompose_meandim_pass.py index fb88b1c90c0..a78514b6af5 100644 --- a/backends/arm/_passes/decompose_meandim_pass.py +++ b/backends/arm/_passes/decompose_meandim_pass.py @@ -9,7 +9,6 @@ import torch from executorch.backends.arm._passes import ArmPass from executorch.backends.arm._passes.arm_pass_utils import get_node_arg -from executorch.backends.arm.operator_support.pool_2d_support import AvgPool2dSupported from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops @@ -67,6 +66,11 @@ def __init__(self, graph_module, tosa_spec): super().__init__() self._graph_module = graph_module self._tosa_spec = tosa_spec + # Lazy import to avoid circular dependency with operator_support + from executorch.backends.arm.operator_support.pool_2d_support import ( + AvgPool2dSupported, + ) + self._avg_pool_checker = AvgPool2dSupported( self._tosa_spec, WhyNoPartitionReporter() ) diff --git a/backends/arm/_passes/decompose_select.py b/backends/arm/_passes/decompose_select.py index 9a25b7c28ae..99c89f474ea 100644 --- a/backends/arm/_passes/decompose_select.py +++ b/backends/arm/_passes/decompose_select.py @@ -7,7 +7,10 @@ # pyre-unsafe import torch -from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.backends.arm._passes.arm_pass_utils import ( + create_node, + get_first_fake_tensor, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -34,8 +37,9 @@ def call(self, graph_module: torch.fx.GraphModule): input_node, dim, index = node.args - rank = len(input_node.meta["val"].size()) - shape = input_node.meta["val"].shape + input_tensor = get_first_fake_tensor(input_node) + rank = len(input_tensor.size()) + shape = input_tensor.shape dim = dim % rank if dim < 0 else dim index = index % shape[dim] if index < 0 else index @@ -44,7 +48,7 @@ def call(self, graph_module: torch.fx.GraphModule): graph_module.graph, slice_op, (input_node, dim, index, index + 1) ) squeeze_node = create_node( - graph_module.graph, squeeze_op, (slice_node, [dim]) + graph_module.graph, squeeze_op, (slice_node, [dim]), from_node=node ) node.replace_all_uses_with(squeeze_node) diff --git a/backends/arm/_passes/decorate_fp32_to_int32_casting_pass.py b/backends/arm/_passes/decorate_fp32_to_int32_casting_pass.py index d6f7ac2ceac..17a682c0a8e 100644 --- a/backends/arm/_passes/decorate_fp32_to_int32_casting_pass.py +++ b/backends/arm/_passes/decorate_fp32_to_int32_casting_pass.py @@ -30,18 +30,17 @@ class DecorateFp32toInt32CastingPass(ArmPass): To lower pytorch fp32 -> int32 casting to TOSA, we need to transform the value with Ceil, Floor, and Where. Before: - output = to_copy(x, dtype=torch.int32) + output = to_dim_order_copy(x, dtype=torch.int32) After: %zero = full((1,), 0.0, dtype=torch.float32) is_non_negative = x >= %zero floor_x = floor(x) ceil_x = ceil(x) decorated_x = where(is_non_negative, floor_x, ceil_x) - output = to_copy(decorated_x, dtype=torch.int32) + output = to_dim_order_copy(decorated_x, dtype=torch.int32) """ targets = [ - exir_ops.edge.aten._to_copy.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, ] diff --git a/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py b/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py index cb9fb8a50c7..491b404f0a4 100644 --- a/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py +++ b/backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py @@ -15,9 +15,9 @@ get_param_tensor, is_param_node, ) -from executorch.backends.arm.constants import DQ_OPS, Q_OPS -from executorch.backends.arm.tosa_quant_utils import QuantArgs +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm.constants import DQ_OPS, Q_OPS from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload diff --git a/backends/arm/_passes/fuse_constant_ops_pass.py b/backends/arm/_passes/fuse_constant_ops_pass.py index d36a15f4c4d..f49565e3c38 100644 --- a/backends/arm/_passes/fuse_constant_ops_pass.py +++ b/backends/arm/_passes/fuse_constant_ops_pass.py @@ -107,7 +107,11 @@ def call(self, graph_module): for node in graph_module.graph.nodes: if node.op != "call_function": continue - if node.target == torch.ops.tosa._table.default: + if node.target in [ + exir_ops.backend.tosa.TABLE.default, + exir_ops.backend.tosa.RESCALE.default, + exir_ops.backend.tosa.TRANSPOSE.default, + ]: continue input_nodes = node.all_input_nodes diff --git a/backends/arm/_passes/fuse_equal_placeholders_pass.py b/backends/arm/_passes/fuse_equal_placeholders_pass.py index 664a0f8ea6c..5631e2f32e9 100644 --- a/backends/arm/_passes/fuse_equal_placeholders_pass.py +++ b/backends/arm/_passes/fuse_equal_placeholders_pass.py @@ -3,6 +3,9 @@ # 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 +from collections import defaultdict + import torch from executorch.backends.arm._passes.arm_pass_utils import ( get_constant_placeholder_kind, @@ -21,7 +24,7 @@ class FuseEqualPlaceholdersPass(ExportPass): """ This pass optimizes memory usage by finding constant placeholders pointing to identical tensors and fusing them to one single placeholder - with multiple users. + with multiple users, using a cache for faster comparison. """ def __init__(self, exported_program: ExportedProgram): @@ -30,58 +33,54 @@ def __init__(self, exported_program: ExportedProgram): def call(self, graph_module: torch.fx.GraphModule) -> PassResult: modified = False - const_placeholder_nodes = [] - for node in graph_module.graph.nodes: - if is_param_node(self.exported_program, node): - const_placeholder_nodes.append(node) - - while const_placeholder_nodes: - # Find equal tensors - node1 = const_placeholder_nodes.pop() - eq_nodes = [node1] - tensor1 = get_param_tensor(self.exported_program, node1) - if tensor1 is None: + # Build a cache of params: mapping hash_key -> list of (node, tensor) + hash_buckets = defaultdict(list) + for node in graph_module.graph.nodes: + if not is_param_node(self.exported_program, node): continue + tensor = get_param_tensor(self.exported_program, node) + if tensor is None: + continue + # Create a lightweight fingerprint: dtype + shape + SHA1 of raw bytes + # Ensure tensor is on CPU and contiguous + t_cpu = tensor.detach().cpu().contiguous() + data_bytes = t_cpu.numpy().tobytes() + key = ( + str(t_cpu.dtype), + tuple(t_cpu.shape), + hashlib.sha1(data_bytes).hexdigest(), + ) + hash_buckets[key].append((node, t_cpu)) - for node2 in const_placeholder_nodes: - tensor2 = get_param_tensor(self.exported_program, node2) - if tensor2 is None: - continue - - if ( - tensor1.dtype == tensor2.dtype - and tensor1.shape == tensor2.shape - and torch.allclose(tensor1, tensor2, atol=1e-08) - ): - eq_nodes.append(node2) + # For each bucket with more than one entry, fuse: + for nodes_tensors in hash_buckets.values(): + if len(nodes_tensors) < 2: + continue - if len(eq_nodes) > 1: - common_name = node1.name + "_common" - common_kind = get_constant_placeholder_kind( - self.exported_program, node1 + # Create a new placeholder from first in list of equal placeholders. + rep_node, rep_tensor = nodes_tensors[0] + common_name = rep_node.name + "_common" + common_kind = get_constant_placeholder_kind(self.exported_program, rep_node) + common_persistent = True + with graph_module.graph.inserting_before(rep_node): + common_node = create_constant_placeholder( + self.exported_program, + graph_module.graph, + common_name, + common_kind, + rep_tensor, + common_persistent, ) - common_persisten_buffer = True - - with graph_module.graph.inserting_before(node1): - common_node = create_constant_placeholder( - self.exported_program, - graph_module.graph, - common_name, - common_kind, - tensor1, - common_persisten_buffer, - ) - - for eq_node in eq_nodes: - eq_node.replace_all_uses_with(common_node) - delete_constant_placeholder(self.exported_program, eq_node) - if eq_node != node1: - const_placeholder_nodes.remove(eq_node) + # Replace uses and delete duplicates + for node, _ in nodes_tensors: + node.replace_all_uses_with(common_node) + delete_constant_placeholder(self.exported_program, node) modified = True if modified: graph_module.recompile() graph_module = super().call(graph_module).graph_module + return PassResult(graph_module=graph_module, modified=modified) diff --git a/backends/arm/_passes/fuse_quantized_activation_pass.py b/backends/arm/_passes/fuse_quantized_activation_pass.py index fb52aab9071..46a7d7f6f98 100644 --- a/backends/arm/_passes/fuse_quantized_activation_pass.py +++ b/backends/arm/_passes/fuse_quantized_activation_pass.py @@ -6,8 +6,8 @@ # pyre-unsafe import torch +from executorch.backends.arm._passes.quant_args import QuantArgs from executorch.backends.arm.constants import Q_OPS -from executorch.backends.arm.tosa_quant_utils import QuantArgs from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult from torch.fx import Node diff --git a/backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py b/backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py new file mode 100644 index 00000000000..4b619af790c --- /dev/null +++ b/backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py @@ -0,0 +1,122 @@ +# Copyright 2025 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-unsafe + + +import logging + +import torch +from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import EdgeOpOverload, ExportPass, PassResult +from torch._subclasses.fake_tensor import FakeTensor + + +logger = logging.getLogger(__name__) + + +class InsertInt32CastsAfterInt64PlaceholdersPass(ExportPass): + """ + Insert an int64->int32 cast after each int64 placeholder. + + Note: Overflow checks are not applied in this pass. It is the user's responsibility to ensure that values fit within + the int32 range. + """ + + # Ops that require i64 inputs → positions of args to upcast. + # Key: op overload; Value: zero-based indices of positional args that must be i64. + I64_INPUT_ARG_POSITIONS = { + torch.ops.aten.one_hot.default: (0,), + } + + def _insert_callsite_i32_to_i64_casts(self, graph_module: torch.fx.GraphModule): + """ + If an operator requires int64 inputs but dtype propagation (via call_operator) + produced int32, insert a local int32→int64 cast at the call site to satisfy + PyTorch's operator input validation. + """ + modified = False + graph = graph_module.graph + for node in graph.nodes: + if node.op != "call_function": + continue + if node.target not in self.I64_INPUT_ARG_POSITIONS: + continue + + with graph.inserting_before(node): + arg_positions = self.I64_INPUT_ARG_POSITIONS.get(node.target) + args_list = list(node.args) + for pos in arg_positions: # type: ignore[union-attr] + input_arg = args_list[pos] + to_copy_op = self._get_decomposition(graph) + cast_node = graph_module.graph.create_node( + "call_function", + to_copy_op, + (input_arg,), + {"dtype": torch.int64}, + ) + cast_node.meta["val"] = node.meta["val"].to(torch.int64) + args_list[pos] = cast_node + node.args = tuple(args_list) + modified = True + return modified + + def _graph_uses_edge_ops(self, graph: torch.fx.Graph) -> bool: + for n in graph.nodes: + if n.op == "call_function": + if isinstance(n.target, EdgeOpOverload): + return True + return False + + def _get_decomposition(self, graph: torch.fx.Graph): + if self._graph_uses_edge_ops(graph): + return exir_ops.edge.dim_order_ops._to_dim_order_copy.default + else: + return torch.ops.dim_order_ops._to_dim_order_copy.default + + def _is_tensor_of_dtype(self, node_val, dtype: torch.dtype) -> bool: + return isinstance(node_val, FakeTensor) and node_val.dtype == dtype + + def _insert_placeholder_i64_to_i32_casts(self, graph_module: torch.fx.GraphModule): + modified = False + graph = graph_module.graph + for node in graph.nodes: + if node.op != "placeholder": + continue + node_val = node.meta["val"] + if not self._is_tensor_of_dtype(node_val, torch.int64): + continue + + to_copy_op = self._get_decomposition(graph) + with graph.inserting_after(node): + cast_after = create_node( + graph, + to_copy_op, + args=(node,), + kwargs={ + "dtype": torch.int32, + }, + ) + users = [user for user in node.users if user != cast_after] + for user in users: + user.replace_input_with(node, cast_after) + logger.warning( + f"Inserting a casting node {cast_after.name} after {node.name} to cast int64 placeholder" + f" to int32 for {node.name} defined in {node.meta.get('stack_trace','[no stack trace found]')}" + ) + modified = True + return modified + + def call(self, graph_module: torch.fx.GraphModule): + modified = False + modified |= self._insert_placeholder_i64_to_i32_casts(graph_module) + modified |= self._insert_callsite_i32_to_i64_casts(graph_module) + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/insert_int64_input_cast_pass.py b/backends/arm/_passes/insert_int64_input_cast_pass.py deleted file mode 100644 index 8864d6bb4eb..00000000000 --- a/backends/arm/_passes/insert_int64_input_cast_pass.py +++ /dev/null @@ -1,110 +0,0 @@ -# Copyright 2025 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-unsafe - - -import logging - -import torch -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass, PassResult - -from .arm_pass_utils import create_node, get_first_fake_tensor - -logger = logging.getLogger(__name__) -logger.setLevel(logging.WARNING) - - -class InsertCastForOpsWithInt64InputPass(ExportPass): - - aten_ops = ( - torch.ops.aten.embedding.default, - torch.ops.aten.slice_copy.Tensor, - ) - edge_ops = ( - exir_ops.edge.aten.embedding.default, - exir_ops.edge.aten.slice_copy.Tensor, - ) - - def get_decomposition(self, op): - if op in self.edge_ops: - return exir_ops.edge.aten._to_copy.default - - if op in self.aten_ops: - return torch.ops.aten._to_copy.default - - raise RuntimeError( - f"[{self.__class__.__name__}] Can't get decomposition for op {op}" - ) - - def _check_aten_embedding_within_int32(self, weights, indices, node: torch.fx.Node): - weights_shape = get_first_fake_tensor(weights).shape - vocab_size = weights_shape[0] - - # Essentially output = weight[indices] which means 0 <= indices[i] < vocab_size - # So should be good if vocab size or number embeddings is below max int32 - if vocab_size >= torch.iinfo(torch.int32).max: - logger.warning( - f"[{node.name}] has size ({vocab_size}) that exceeds int32 limit," - "so aten.embedding will not be lowered to TOSA." - ) - return False - - return True - - def _insert_int32_cast_before_node(self, graph, node, original_input): - to_copy_op = self.get_decomposition(node.target) - with graph.inserting_before(node): - cast_before = create_node( - graph, - to_copy_op, - args=(original_input,), - kwargs={ - "dtype": torch.int32, - "memory_format": torch.preserve_format, - }, - ) - node.replace_input_with(original_input, cast_before) - - def call(self, graph_module): - graph = graph_module.graph - modified_graph = False - - for node in list(graph.nodes): - if node.op != "call_function": - continue - if node.target not in self.aten_ops + self.edge_ops: - continue - - args = node.args - - if node.target in ( - exir_ops.edge.aten.embedding.default, - torch.ops.aten.embedding.default, - ): - weights = args[0] - indices = args[1] - if self._check_aten_embedding_within_int32(weights, indices, node): - self._insert_int32_cast_before_node(graph, node, indices) - modified_graph = True - - elif node.target in ( - exir_ops.edge.aten.slice_copy.Tensor, - torch.ops.aten.slice_copy.Tensor, - ): - # MLETORCH-829: Add range check for slice_copy - input_tensor = args[0] - fake_tensor = input_tensor.meta["val"] - if fake_tensor.dtype != torch.int64: - continue - - self._insert_int32_cast_before_node(graph, node, input_tensor) - modified_graph = True - - if modified_graph: - graph_module.recompile() - graph_module = super().call(graph_module).graph_module - return PassResult(graph_module, modified_graph) diff --git a/backends/arm/_passes/insert_rescales_pass.py b/backends/arm/_passes/insert_rescales_pass.py index 8a2e10b6b2d..7f75aecf24c 100644 --- a/backends/arm/_passes/insert_rescales_pass.py +++ b/backends/arm/_passes/insert_rescales_pass.py @@ -3,70 +3,25 @@ # 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 copy from typing import cast -import torch from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.backends.arm._passes.quant_args import QuantArgs from executorch.backends.arm.constants import DQ_OPS, Q_OPS -from executorch.backends.arm.tosa_quant_utils import QuantArgs +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult -from torch import Tensor from torch.fx import GraphModule, Node -from torch.library import custom_op, register_fake - -logger = logging.getLogger(__name__) - - -@custom_op("tosa::_rescale", mutates_args=()) # type: ignore[misc] -def rescale( - x: Tensor, dtype: torch.dtype, scale: float, in_zp: int, out_zp: int -) -> Tensor: - logger.warning( - "Ran default implementation of tosa::_rescale." - "This op is meant to always be inserted inside a partition and a correct default implementation is not implemented." - ) - # Clone is needed to not return reference when rescaling to same dtype. - # This is a neccessary requirement for non-mutating custom ops. - return x.to(dtype=dtype).clone() - - -@register_fake("tosa::_rescale") # type: ignore[misc] -def rescale_fake( - x: Tensor, dtype: torch.dtype, scale: float, in_zp: int, out_zp: int -) -> Tensor: - """Casts the input tensor to dtype `dtype` to produce the correct tensor meta for a _rescale op. - Additionally validates TOSA constraints of a RESCALE op. - """ - if dtype not in (torch.int32, torch.int8, torch.int16): - raise NotImplementedError( - f"tosa::rescale currently only supports int32, int16 and int8, not {dtype}" - ) - if dtype in (torch.int32, torch.int16) and out_zp != 0: - raise ValueError( - f"TOSA requires output_zp to be zero when the output dtype is {dtype}." - ) - if x.dtype in (torch.int32, torch.int16) and in_zp != 0: - raise ValueError( - f"TOSA requires input_zp to be zero when the input dtype is {dtype}" - ) - if x.dtype == torch.int8 and not -128 <= in_zp <= 127: - raise ValueError(f"{in_zp=} outside valid range (-128,127) for int8.") - if dtype == torch.int8 and not -128 <= out_zp <= 127: - raise ValueError(f"{out_zp=} outside valid range (-128,127) for int8.") - - return x.to(dtype=dtype).clone() class InsertRescalePass(ExportPass): """Finds patterns of dq -> q, and replaces them - with passthrough_to_tosa::rescales. + with backend dialect tosa::RESCALE op. - Does not garantuee that the dtypes and zero points are valid + Does not guarantee that the dtypes and zero points are valid in TOSA, that is the job of the quantization annotator that produced the dq and q nodes. The TOSA constraints are validated - in the fake implementation of passthrough_to_tosa:rescale. + in the fake implementation of. """ def fold_dq_q_to_rescale(self, node: Node, user: Node, graph_module: GraphModule): @@ -77,7 +32,7 @@ def fold_dq_q_to_rescale(self, node: Node, user: Node, graph_module: GraphModule with graph_module.graph.inserting_before(node): rescale_node = create_node( graph_module.graph, - torch.ops.tosa._rescale.default, + exir_ops.backend.tosa.RESCALE.default, ( node.all_input_nodes[0], q_args.dtype, diff --git a/backends/arm/_passes/insert_table_ops.py b/backends/arm/_passes/insert_table_ops.py index cb78ee58cfe..fb5d7de5e12 100644 --- a/backends/arm/_passes/insert_table_ops.py +++ b/backends/arm/_passes/insert_table_ops.py @@ -10,27 +10,18 @@ import torch from executorch.backends.arm._passes.arm_pass_utils import create_node -from executorch.backends.arm.tosa_quant_utils import QuantArgs +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.transforms.utils import create_constant_placeholder + from executorch.exir import ExportedProgram 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.export.graph_signature import InputKind from torch.fx import GraphModule from torch.fx.node import Node -from torch.library import impl, Library - -lib = Library("tosa", "DEF") -lib.define("_table(Tensor self) -> Tensor") - - -@impl(lib, "_table") -def _table_impl(*args, **kwargs): # pyre-ignore - in_dtype = args[0].dtype - if in_dtype == torch.int8: - return args[0] - return args[0].to(dtype=torch.int32) class TableOps: @@ -43,6 +34,7 @@ class TableOps: exir_ops.edge.aten.ceil.default: torch.ceil, exir_ops.edge.aten.erf.default: torch.erf, exir_ops.edge.aten.exp.default: torch.exp, + exir_ops.edge.aten.expm1.default: torch.expm1, exir_ops.edge.aten.floor.default: torch.floor, exir_ops.edge.aten.log.default: torch.log, exir_ops.edge.aten.reciprocal.default: torch.reciprocal, @@ -59,12 +51,15 @@ class TableOps: exir_ops.edge.aten.acosh.default: torch.acosh, exir_ops.edge.aten.asin.default: torch.asin, exir_ops.edge.aten.asinh.default: torch.asinh, + exir_ops.edge.aten.cosh.default: torch.cosh, + exir_ops.edge.aten.acos.default: torch.acos, } # Targets that must be treated explicitly special_table_ops: Set[EdgeOpOverload] = { exir_ops.edge.aten.pow.Tensor_Scalar, exir_ops.edge.aten.gelu.default, + exir_ops.edge.aten.elu.default, } def __init__(self, exported_program: ExportedProgram): @@ -98,6 +93,11 @@ def __getitem__(self, node: Node): return lambda x: torch.nn.functional.gelu( x, approximate=approximate ).flatten() + case exir_ops.edge.aten.elu.default: + input_alpha = cast(int, node.kwargs["alpha"]) + return lambda x: torch.nn.functional.elu( + x, alpha=input_alpha + ).flatten() case _: # Op must be handled if it's inside self.special_ops raise AssertionError("Unhandled table operation") @@ -239,13 +239,8 @@ def call(self, graph_module: GraphModule) -> PassResult: # We only want to replace the node if it's quantized continue # Create table node - with graph_module.graph.inserting_before(node): - table_node = create_node( - graph=graph_module.graph, - op_target=torch.ops.tosa._table.default, - args=(node.args[0],), - ) - output_node = table_node + insert_pos = list(node.graph.nodes)[0] + with graph_module.graph.inserting_before(insert_pos): # Expect exactly one quantization parameter for input and output if len(input_qparams) != 1: raise ValueError( @@ -265,27 +260,37 @@ def call(self, graph_module: GraphModule) -> PassResult: out_quantargs=output_qparams[0], ) # Register buffer in self.exported_program.state_dict - # When the graph is retraced, the implementation _table is used and the suffix _default disappears from the node name - # Remove it here to make it possible to find in the node_visitor - self.register_buffer( - buffer_name=table_node.name.replace("_default", ""), buffer=buffer + const_table_node = create_constant_placeholder( + exp_program=self.exported_program, + graph=node.graph, + kind=InputKind.BUFFER, + name=node.name + "_table_constant", + data=buffer, + persistent_buffer=True, ) + # Create table node + with graph_module.graph.inserting_before(node): + table_op_node = create_node( + graph=graph_module.graph, + op_target=exir_ops.backend.tosa.TABLE.default, + args=(node.args[0], const_table_node), + ) + output_node = table_op_node + if lshift != 0: scale = 2.0**lshift rescale_node = create_node( graph=graph_module.graph, - op_target=torch.ops.tosa._rescale.default, - args=(table_node, output_qparams[0].dtype, scale, 0, 0), + op_target=exir_ops.backend.tosa.RESCALE.default, + args=(table_op_node, output_qparams[0].dtype, scale, 0, 0), ) output_node = rescale_node node.replace_all_uses_with(output_node) - graph_module.graph.erase_node(node) - - output_node.meta["input_qparams"] = input_qparams - output_node.meta["output_qparams"] = output_qparams + table_op_node.meta["input_qparams"] = input_qparams + table_op_node.meta["output_qparams"] = output_qparams modified = True if modified: diff --git a/backends/arm/_passes/quant_args.py b/backends/arm/_passes/quant_args.py new file mode 100644 index 00000000000..974d6dfdbd3 --- /dev/null +++ b/backends/arm/_passes/quant_args.py @@ -0,0 +1,125 @@ +# Copyright 2025 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, cast, NamedTuple + +import torch +from executorch.exir.dialects._ops import ops as exir_ops + +exir_ops = cast(Any, exir_ops) +from executorch.backends.arm.constants import PER_CHANNEL_QDQ_OPS, PER_TENSOR_QDQ_OPS +from torch import Tensor + + +class QuantArgs(NamedTuple): + scale: list[float] | float + zp: list[int] | int + qmin: int + qmax: int + dtype: torch.dtype + axis: int = 0 + per_channel: bool = False + + def quantize_value(self, x: torch.Tensor | float) -> Tensor: + """Quantizes the input tensor or value to a quantized tensor. If the input is + not a tensor, it is converted to a tensor first. If self.per_channel is True, + the quantization is done per channel, otherwise it is done per tensor. + """ + if not isinstance(x, torch.Tensor): + x = torch.Tensor([x]) + x = x.to(torch.float32) + if self.per_channel: + q_op = exir_ops.edge.quantized_decomposed.quantize_per_channel.default + args = ( + x, + torch.tensor(self.scale), + torch.tensor(self.zp), + self.axis, + self.qmin, + self.qmax, + self.dtype, + ) + else: + q_op = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + args = (x, self.scale, self.zp, self.qmin, self.qmax, self.dtype) # type: ignore[assignment] + return q_op(*args) + + def dequantize_value(self, qx: torch.Tensor) -> torch.Tensor: + """Dequantizes the input tensor or value to a dequantized tensor If the input + is not a tensor, it is converted to a tensor first. If self.per_channel is True, + the dequantization is done per channel, otherwise it is done per tensor. + """ + if self.per_channel: + dq_op = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default + args = ( + qx, + torch.tensor(self.scale), + torch.tensor(self.zp), + self.axis, + self.qmin, + self.qmax, + self.dtype, + ) + else: + dq_op = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + args = (qx, self.scale, self.zp, self.qmin, self.qmax, self.dtype) # type: ignore[assignment] + return dq_op(*args) + + @classmethod + def from_operator(cls, op, args): + if op in PER_TENSOR_QDQ_OPS: + return cls( + scale=cast(float, args[1]), + zp=cast(int, args[2]), + qmin=cast(int, args[3]), + qmax=cast(int, args[4]), + dtype=cast(torch.dtype, args[5]), + axis=0, + per_channel=False, + ) + elif op in PER_CHANNEL_QDQ_OPS: + return cls( + scale=cast(list[float], args[1].tolist()), + zp=cast(list[int], args[2].tolist()), + axis=cast(int, args[3]), + qmin=cast(int, args[4]), + qmax=cast(int, args[5]), + dtype=cast(torch.dtype, args[6]), + per_channel=True, + ) + else: + # We're only handling per tensor and per channel quantization + raise NotImplementedError(f"Unsupported quantization operation: {op}") + + def get_scale_per_tensor(self) -> float: + if not isinstance(self.scale, float): + raise TypeError( + f"Expected scale {self.scale} to be a float but found scale of " + f"type {type(self.scale)}" + ) + return self.scale + + def get_zp_per_tensor(self) -> int: + if not isinstance(self.zp, int): + raise TypeError( + f"Expected zero point {self.zp} to be an int but found zp of " + f"type {type(self.zp)}" + ) + return self.zp + + def get_scale_per_channel(self) -> list[float]: + if not isinstance(self.scale, list): + raise TypeError( + f"Expected scale {self.scale} to be a list but found scale of " + f"type {type(self.scale)}" + ) + return self.scale + + def get_zp_per_channel(self) -> list[int]: + if not isinstance(self.zp, list): + raise TypeError( + f"Expected zero point {self.zp} to be a list but found zp of " + f"type {type(self.zp)}" + ) + return self.zp diff --git a/backends/arm/_passes/remove_clone_pass.py b/backends/arm/_passes/remove_clone_pass.py deleted file mode 100644 index a2822c7378e..00000000000 --- a/backends/arm/_passes/remove_clone_pass.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright 2024-2025 Arm Limited and/or its 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-unsafe - -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass - - -class RemoveClonePass(ExportPass): - """Remove all clones from graph_module""" - - def call_operator(self, op, args, kwargs, meta): - if op != exir_ops.edge.aten.clone.default: - return super().call_operator(op, args, kwargs, meta) - - if len(args) != 1: - raise ValueError( - f"clone operator expects exactly one argument, got {len(args)}" - ) - return args[0] diff --git a/backends/arm/_passes/remove_noop_pass.py b/backends/arm/_passes/remove_noop_pass.py new file mode 100644 index 00000000000..623517aac59 --- /dev/null +++ b/backends/arm/_passes/remove_noop_pass.py @@ -0,0 +1,33 @@ +# Copyright 2024-2025 Arm Limited and/or its 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-unsafe + +import logging + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + +logger = logging.getLogger(__name__) + + +class RemoveNoopPass(ExportPass): + """Remove no-ops from graph_module""" + + def call_operator(self, op, args, kwargs, meta): + if op not in ( + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + ): + return super().call_operator(op, args, kwargs, meta) + + input_dtype = args[0].data.dtype + output_dtype = kwargs.get("dtype", input_dtype) + + if input_dtype != output_dtype: + return super().call_operator(op, args, kwargs, meta) + + return args[0] diff --git a/backends/arm/_passes/annotate_channels_last_dim_order_pass.py b/backends/arm/_passes/to_tosa_memory_format_pass.py similarity index 57% rename from backends/arm/_passes/annotate_channels_last_dim_order_pass.py rename to backends/arm/_passes/to_tosa_memory_format_pass.py index f8ead856fbb..e4436d638f4 100644 --- a/backends/arm/_passes/annotate_channels_last_dim_order_pass.py +++ b/backends/arm/_passes/to_tosa_memory_format_pass.py @@ -6,44 +6,34 @@ # pyre-unsafe +import logging + import torch +from executorch.backends.arm._passes import AnnotateOutputDimOrderPass from executorch.backends.arm._passes.arm_pass_utils import ( create_node, get_first_fake_tensor, + get_output_dim_orders, + is_param_node, ) -from executorch.backends.arm.tosa_utils import is_consumer_node_depthwise_conv2d +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.library import impl, Library - -# Define lib with passthrough operators. The operators have no real meaning in edge IR -# except for argument validaiton and a passthrough output. The operators will be used -# when lowering to TOSA, e.g. a passthrough_to_tosa._transpose will not affect -# the edge IR graph but will be lowered to a TOSA-TRANSPOSE. -lib = Library("passthrough_to_tosa", "DEF") -# For certain operators we need the data in a specific data format. Changing tosa_dim_order -# is not sufficient as we also need transpose the data. -# By utilizing an edge IR passthrough operator we can keep the edge program in -# channels-first/contiguous and get the desired behavior in the TOSA lowering. -lib.define("_transpose(Tensor self, int[] dim_order) -> Tensor") - - -@impl(lib, "_transpose") -def _transpose_impl(*args, **kwargs): - # Validate length of dim_order array - dim = args[1] - if len(dim) != 4 and len(dim) != 5: - raise ValueError( - f"Dim order length must be either 4 or 5, got {len(dim)}: {dim}" - ) - # Pass-through in edge-IR - return args[0] + +logger = logging.getLogger(__name__) + + +def _is_input(node: torch.fx.Node, exported_program: ExportedProgram) -> bool: + """ + Returns True if the node is an input node, i.e. a placeholder or a parameter. + """ + return node.op == "placeholder" and not is_param_node(exported_program, node) -class AnnotateChannelsLastDimOrder(ExportPass): +class ToTosaMemoryFormatPass(ExportPass): """ Annotates each node with a tosa_dim_order. tosa_dim_order can be seen as a channels-last dim-order - that in most cases will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts passthrough_to_tosa._transpose + that in most cases will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts backend.tosa.TRANSPOSE when a transition between 3D and 4D/5D tensors happen. The annotated tosa_dim_order is used to permute the node's shape such that it gives a TOSA-compliant shape. """ @@ -54,6 +44,23 @@ class AnnotateChannelsLastDimOrder(ExportPass): NNHWC_order = (0, 1, 3, 4, 2) NNHWC_inverse_order = (0, 1, 4, 2, 3) + def __init__(self, exported_program: ExportedProgram) -> None: + self.exported_program = exported_program + super().__init__() + + @staticmethod + def _is_consumer_node_depthwise_conv2d(node: torch.fx.Node): + consumer_node = list(node.users)[0] + if consumer_node.target == exir_ops.edge.aten.convolution.default: + consumer_node_inputs = consumer_node.all_input_nodes + groups = consumer_node.args[-1] + in_channels = consumer_node_inputs[0].meta["val"].shape[1] + out_channels = consumer_node_inputs[1].meta["val"].shape[0] + if (in_channels == groups) and (out_channels % in_channels) == 0: + return True + + return False + def is_weight_node_for_depthwise_conv2d(self, node: torch.fx.Node): """ returns True for w in the following sequence; @@ -64,7 +71,7 @@ def is_weight_node_for_depthwise_conv2d(self, node: torch.fx.Node): consumer_node = list(node.users)[0] if self.is_weight_node_for_depthwise_conv2d(consumer_node): return True - if is_consumer_node_depthwise_conv2d(node): + if self._is_consumer_node_depthwise_conv2d(node): # Check that node is the weight-argument and not input or bias return consumer_node.args[1] == node @@ -116,54 +123,57 @@ def is_channel_reshape(input_shape, output_shape): @staticmethod def insert_input_transpose(node, input_node, graph_module): + if input_node.target == exir_ops.backend.tosa.TRANSPOSE.default: + pre_permute_node = input_node.all_input_nodes[0] + node.replace_input_with(input_node, pre_permute_node) + return + with graph_module.graph.inserting_before(node): permute_node = create_node( graph_module.graph, - torch.ops.passthrough_to_tosa._transpose.default, + exir_ops.backend.tosa.TRANSPOSE.default, args=( input_node, list( - AnnotateChannelsLastDimOrder.NNHWC_inverse_order + ToTosaMemoryFormatPass.NNHWC_inverse_order if len(get_first_fake_tensor(input_node).size()) == 5 - else AnnotateChannelsLastDimOrder.NHWC_inverse_order + else ToTosaMemoryFormatPass.NHWC_inverse_order ), ), + from_node=node, ) node.replace_input_with(input_node, permute_node) permute_node.meta["tosa_dim_order"] = tuple( range(len(input_node.meta["val"].size())) ) - permute_node.meta["val"] = input_node.meta["val"] @staticmethod def insert_output_transpose(node, graph_module): with graph_module.graph.inserting_after(node): permute_node = create_node( graph_module.graph, - torch.ops.passthrough_to_tosa._transpose.default, + exir_ops.backend.tosa.TRANSPOSE.default, args=( node, list( - AnnotateChannelsLastDimOrder.NNHWC_order + ToTosaMemoryFormatPass.NNHWC_order if len(get_first_fake_tensor(node).size()) == 5 - else AnnotateChannelsLastDimOrder.NHWC_order + else ToTosaMemoryFormatPass.NHWC_order ), ), + from_node=node, ) + permute_node.meta["tosa_dim_order"] = ( - AnnotateChannelsLastDimOrder.NNHWC_order + ToTosaMemoryFormatPass.NNHWC_order if len(get_first_fake_tensor(node).size()) == 5 - else AnnotateChannelsLastDimOrder.NHWC_order - ) - permute_node.meta["val"] = get_first_fake_tensor(node).permute( - AnnotateChannelsLastDimOrder.NNHWC_order - if len(get_first_fake_tensor(node).size()) == 5 - else AnnotateChannelsLastDimOrder.NHWC_order + else ToTosaMemoryFormatPass.NHWC_order ) node.meta["tosa_dim_order"] = tuple( range(len(get_first_fake_tensor(node).size())) ) + users = [user for user in node.users if user != permute_node] for user in users: user.replace_input_with(node, permute_node) @@ -174,20 +184,23 @@ def _insert_view_transpose( ): nchw_to_nhwc = len(input_shape) < 4 and len(output_shape) >= 4 nhwc_to_nchw = len(input_shape) >= 4 and len(output_shape) < 4 - channel_reshape = AnnotateChannelsLastDimOrder.is_channel_reshape( + channel_reshape = ToTosaMemoryFormatPass.is_channel_reshape( output_shape, input_shape ) if ( channel_reshape or nhwc_to_nchw - ) and AnnotateChannelsLastDimOrder.memory_format_differs(input_shape): - AnnotateChannelsLastDimOrder.insert_input_transpose( + ) and ToTosaMemoryFormatPass.memory_format_differs(input_shape): + + ToTosaMemoryFormatPass.insert_input_transpose( node, input_node, graph_module ) + if ( channel_reshape or nchw_to_nhwc - ) and AnnotateChannelsLastDimOrder.memory_format_differs(output_shape): - AnnotateChannelsLastDimOrder.insert_output_transpose(node, graph_module) + ) and ToTosaMemoryFormatPass.memory_format_differs(output_shape): + + ToTosaMemoryFormatPass.insert_output_transpose(node, graph_module) def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule): """ @@ -205,9 +218,10 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule): for node in graph_module.graph.nodes: # call_function and placeholder allowed due to # index.Tensor being able to come in as both - if node.op not in ["call_function", "placeholder"]: + if node.op not in ["call_function", "placeholder", "output"]: continue + # Transpose views elif node.target in ( exir_ops.edge.aten.view_copy.default, exir_ops.edge.aten.index.Tensor, @@ -218,26 +232,67 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule): input_node = node.args[0] input_shape = input_node.meta["val"].shape output_shape = node.meta["val"].shape - self._insert_view_transpose( - input_shape, output_shape, node, input_node, graph_module + input_shape, + output_shape, + node, + input_node, + graph_module, ) + # Transpose inputs + elif _is_input(node, self.exported_program): + input_shape = get_first_fake_tensor(node).size() + if len(input_shape) in (4, 5): + ToTosaMemoryFormatPass.insert_output_transpose(node, graph_module) + + # Transpose outputs + elif node.op == "output": + output_shape = get_first_fake_tensor(node).size() + + if len(output_shape) in (4, 5): + for input_node in node.all_input_nodes: + ToTosaMemoryFormatPass.insert_input_transpose( + node, input_node, graph_module + ) + + def remove_dim_order_kwargs( + self, graph_module: torch.fx.GraphModule, node: torch.fx.Node + ): + if node.op != "call_function": + return + + kwargs = dict(node.kwargs) + + if "dim_order" in kwargs: + logger.warning( + f"Ignoring dim_order kwarg '{kwargs['dim_order']}' for '{node.name}'." + ) + del kwargs["dim_order"] + + node.kwargs = kwargs + def call(self, graph_module: torch.fx.GraphModule): for node in graph_module.graph.nodes: node_data = get_first_fake_tensor(node).data - if node_data.dim() == 4: + self.remove_dim_order_kwargs(graph_module, node) + # Inputs and outputs are always in (N)NCHW format + if _is_input(node, self.exported_program) or node.op == "output": + dim_order = tuple(range(node_data.dim())) + elif node_data.dim() == 4: dim_order = self.NHWC_order if self.is_weight_node_for_depthwise_conv2d(node): # The weights of TOSA DEPTHWISE_CONV2D have shape (H, W, C, M) which corresponds to # dim_order = (2, 3, 0, 1) (https://www.mlplatform.org/tosa/tosa_spec.html#_depthwise_conv2d). dim_order = self.HWCM_order elif node_data.dim() == 5: - dim_order = self.NNHWC_order # type: ignore[assignment] + dim_order = self.NNHWC_order else: dim_order = tuple(range(node_data.dim())) # type: ignore[assignment] + node.meta["tosa_dim_order"] = dim_order + # Insert TOSA transposes to convert between (N)NCHW and (N)NHWC format. # See insert_tosa_transposes for insertion conditions. self.insert_tosa_transposes(graph_module) @@ -245,3 +300,32 @@ def call(self, graph_module: torch.fx.GraphModule): graph_module = super().call(graph_module).graph_module return PassResult(graph_module, True) + + def requires(self, graph_module) -> None: + """ + This is the only pass which handles dim_orders, so verify that the output dim_orders has not changed since the beginning of the lowering pipeline. + """ + + dim_orders = get_output_dim_orders(graph_module) + original_dim_orders = graph_module.graph.output_node().meta.get( + "original_dim_orders" + ) + output_node = graph_module.graph.output_node() + + if original_dim_orders is None: + raise RuntimeError( + f"{AnnotateOutputDimOrderPass.__name__} must be run in the beginning of the pass pipeline to verify that the dim order has not changed unexpectedly during its run." + ) + + if len(dim_orders) != len(original_dim_orders): + raise RuntimeError( + f"The number of outputs has changed since {AnnotateOutputDimOrderPass.__name__} was run." + ) + + for node, dim_order, original_dim_order in zip( + output_node.args[0], dim_orders, original_dim_orders + ): + if dim_order != original_dim_order: + raise RuntimeError( + f"The dim order of output {node.name} has changed from {original_dim_order} to {dim_order} since {AnnotateOutputDimOrderPass.__name__} was run." + ) diff --git a/backends/arm/arm_backend.py b/backends/arm/arm_backend.py index fc638647b46..2e71f91dbb6 100644 --- a/backends/arm/arm_backend.py +++ b/backends/arm/arm_backend.py @@ -10,11 +10,10 @@ # backends. Converts via TOSA as an intermediate form supported by AoT and # JIT compiler flows. # +from enum import Enum from typing import List, Optional -from executorch.backends.arm.tosa_specification import ( # type: ignore[import-not-found] - TosaSpecification, -) +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.backend.compile_spec_schema import ( # type: ignore[import-not-found] CompileSpec, @@ -22,12 +21,17 @@ class ArmCompileSpecBuilder: + class DebugMode(Enum): + JSON = 1 + TOSA = 2 + def __init__(self): self.compile_spec: List[CompileSpec] = [] self.compiler_flags = [] self.output_format = None self.path_for_intermediates = None self.tosa_spec = None + self.tosa_debug_mode = None def vgf_compile_spec( self, @@ -57,7 +61,7 @@ def vgf_compile_spec( f"Invalid TOSA version: {tosa_version}" ) - if not ("FP" or "INT" in tosa_profiles): + if "FP" not in tosa_profiles and "INT" not in tosa_profiles: raise ValueError( "Arm backend only supports converter-backend for FP or INT. " f"Invalid TOSA profile: {tosa_profiles}" @@ -128,7 +132,7 @@ def ethosu_compile_spec( self.compiler_flags.append("--output-format=raw") self.compiler_flags.append("--debug-force-regor") - base_tosa_version = "TOSA-1.0+INT" + base_tosa_version = "TOSA-1.0+INT+int16" if "u55" in target: # Add the Ethos-U55 extension marker base_tosa_version += "+u55" @@ -163,6 +167,13 @@ def dump_intermediate_artifacts_to( self.path_for_intermediates = output_path return self + def dump_debug_info(self, debug_mode: DebugMode) -> "ArmCompileSpecBuilder": + """ + Dump debugging information into the intermediates path + """ + self.tosa_debug_mode = debug_mode.name + return self + def build(self) -> List[CompileSpec]: """ Generate a list of compile spec objects from the builder @@ -188,6 +199,16 @@ def build(self) -> List[CompileSpec]: CompileSpec("debug_artifact_path", self.path_for_intermediates.encode()) ) + if self.tosa_debug_mode is not None: + if not self.path_for_intermediates: + raise ValueError( + "dump_debug_info() must be used in conjunction with dump_intermediate_artifacts_to()" + ) + + self.compile_spec.append( + CompileSpec("dump_debug_info", self.tosa_debug_mode.encode()) + ) + return self.compile_spec @@ -217,13 +238,6 @@ def is_vgf(compile_spec: List[CompileSpec]) -> bool: return False -def get_tosa_spec(compile_spec: List[CompileSpec]) -> TosaSpecification: - for spec in compile_spec: - if spec.key == "tosa_spec": - return TosaSpecification.create_from_string(spec.value.decode()) - raise ValueError("Could not find TOSA version in CompileSpec") - - def get_intermediate_path(compile_spec: List[CompileSpec]) -> Optional[str]: for spec in compile_spec: if spec.key == "debug_artifact_path": diff --git a/backends/arm/common/__init__.py b/backends/arm/common/__init__.py new file mode 100644 index 00000000000..c8d1c683da3 --- /dev/null +++ b/backends/arm/common/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/backends/arm/common/debug.py b/backends/arm/common/debug.py new file mode 100644 index 00000000000..bca6c06d140 --- /dev/null +++ b/backends/arm/common/debug.py @@ -0,0 +1,87 @@ +# Copyright 2025 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 +import os +from typing import Optional + +import serializer.tosa_serializer as ts # type: ignore +import torch +from executorch.exir.print_program import inspect_node + +logger = logging.getLogger(__name__) + + +def debug_node(node: torch.fx.Node, graph_module: torch.fx.GraphModule): + # Debug output of node information + logger.info(get_node_debug_info(node, graph_module)) + + +def get_node_debug_info( + node: torch.fx.Node, graph_module: torch.fx.GraphModule | None = None +) -> str: + output = ( + f" {inspect_node(graph=graph_module.graph, node=node)}\n" + if graph_module + else "" + "-- NODE DEBUG INFO --\n" + f" Op is {node.op}\n" + f" Name is {node.name}\n" + f" Node target is {node.target}\n" + f" Node args is {node.args}\n" + f" Node kwargs is {node.kwargs}\n" + f" Node users is {node.users}\n" + " Node.meta = \n" + ) + for k, v in node.meta.items(): + if k == "stack_trace": + matches = v.split("\n") + output += " 'stack_trace =\n" + for m in matches: + output += f" {m}\n" + else: + output += f" '{k}' = {v}\n" + + if isinstance(v, list): + for i in v: + output += f" {i}\n" + return output + + +# Output TOSA flatbuffer and test harness file +def debug_tosa_dump(tosa_graph: ts.TosaSerializer, path: str, suffix: str = ""): + filename = f"output{suffix}.tosa" + + logger.info(f"Emitting debug output to: {path=}, {suffix=}") + + os.makedirs(path, exist_ok=True) + + fb = tosa_graph.serialize() + js = tosa_graph.writeJson(filename) + + filepath_tosa_fb = os.path.join(path, filename) + with open(filepath_tosa_fb, "wb") as f: + f.write(fb) + if not os.path.exists(filepath_tosa_fb): + raise IOError("Failed to write TOSA flatbuffer") + + filepath_desc_json = os.path.join(path, f"desc{suffix}.json") + with open(filepath_desc_json, "w") as f: + f.write(js) + if not os.path.exists(filepath_desc_json): + raise IOError("Failed to write TOSA JSON") + + +def debug_fail( + node, + graph_module, + tosa_graph: Optional[ts.TosaSerializer] = None, + path: Optional[str] = None, +): + logger.warning("Internal error due to poorly handled node:") + if tosa_graph is not None and path is not None: + debug_tosa_dump(tosa_graph, path) + logger.warning(f"Debug output captured in '{path}'.") + debug_node(node, graph_module) diff --git a/backends/arm/debug/TARGETS b/backends/arm/debug/TARGETS new file mode 100644 index 00000000000..8ddfd9a285c --- /dev/null +++ b/backends/arm/debug/TARGETS @@ -0,0 +1,14 @@ +# @noautodeps +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +runtime.python_library( + name = "schema", + srcs = [ + "__init__.py", + "schema.py", + ], + deps = [ + "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", + "//caffe2:torch", + ], +) diff --git a/backends/arm/debug/__init__.py b/backends/arm/debug/__init__.py new file mode 100644 index 00000000000..c8d1c683da3 --- /dev/null +++ b/backends/arm/debug/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/backends/arm/debug/schema.py b/backends/arm/debug/schema.py new file mode 100644 index 00000000000..82f0fd6bf7e --- /dev/null +++ b/backends/arm/debug/schema.py @@ -0,0 +1,150 @@ +# Copyright 2025 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import json + +from dataclasses import asdict, dataclass +from typing import Any, Optional + +import serializer.tosa_serializer as ts # type: ignore +import torch + +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder + +from torch.fx.traceback import NodeSource + + +@dataclass +class TosaDebugSchema: + node_name: str + operator_name: str + operator_id: int + + +@dataclass +class ATenDebugSchema: + node_name: str + operator_name: str + + @staticmethod + def from_node(node: torch.fx.Node) -> ATenDebugSchema: + # node.target is Union[Callable[..., Any], str], so we need to access this correctly depending on the type + if callable(node.target): + operator_name = node.target.__name__ + else: + operator_name = node.target + + return ATenDebugSchema(node_name=node.name, operator_name=operator_name) + + +@dataclass +class TorchDebugSchema: + stack_trace: list[str] + node_trace: list[dict[str, Any]] | str + nn_module_stack: dict[str, Any] | str + torch_fn: tuple[str, str] | str + + @staticmethod + def serialize_node_trace(node_trace: list[NodeSource]) -> list[dict[str, Any]]: + """Flatten the from_node dictionary to remove nesting.""" + flattened = [] + node_stack = [] + + for n in node_trace: + node_stack.append((n, -1)) + + while len(node_stack) > 0: + node, parent_id = node_stack.pop() + flattened.append( + { + "name": node.name, + "target": node.target, + "graph_id": node.graph_id, + "pass_name": node.pass_name, + "action": node._get_action_string(), + "parent_graph_id": parent_id, + } + ) + + for n in node.from_node: + node_stack.append((n, node.graph_id)) + + return flattened + + @staticmethod + def from_node(node: torch.fx.Node) -> TorchDebugSchema: + node_trace: str | list[dict[str, Any]] = "No node trace available." + + if "from_node" in node.meta: + # Flatten the node_trace dictionary, so there is no nesting + node_trace = TorchDebugSchema.serialize_node_trace(node.meta["from_node"]) + + return TorchDebugSchema( + stack_trace=node.meta.get("stack_trace", "No stack trace available").split( + "\n" + ), + node_trace=node_trace, + nn_module_stack=node.meta.get( + "nn_module_stack", "No module stack trace available" + ), + torch_fn=node.meta.get("torch_fn", "No torch_fn available"), + ) + + +@dataclass +class DebugSchema: + event_id: int + aten_info: ATenDebugSchema + tosa_info: Optional[TosaDebugSchema] + torch_info: TorchDebugSchema + + def to_dict(self) -> dict[str, Any]: + output = asdict(self) + + if self.tosa_info is None: + output.pop("tosa_info") + + return output + + +class DebugHook: + def __init__(self, debug_mode: ArmCompileSpecBuilder.DebugMode) -> None: + self._debug_events: list[DebugSchema] = [] + self.__op_id_to_name = {} + self.mode = debug_mode + + # Build up a mapping from TOSA 1.0 operator IDs to their names + for name, val in vars(ts.Op).items(): + self.__op_id_to_name[val] = name + + def add(self, node: torch.fx.Node, tosa_op: Any, tosa_op_id: int) -> DebugSchema: + tosa_debug_info = None + + # If the debug data is being embedded into the TOSA flatbuffer + # do not collect TOSADebugSchema data, it's redundent + if self.mode != ArmCompileSpecBuilder.DebugMode.TOSA: + tosa_debug_info = TosaDebugSchema( + node_name=str(tosa_op), + operator_name=self.__op_id_to_name[tosa_op_id], + operator_id=tosa_op_id, + ) + + aten_debug_info = ATenDebugSchema.from_node(node) + torch_debug_info = TorchDebugSchema.from_node(node) + + debug_info = DebugSchema( + event_id=len(self._debug_events), + aten_info=aten_debug_info, + tosa_info=tosa_debug_info, + torch_info=torch_debug_info, + ) + self._debug_events.append(debug_info) + + return debug_info + + def serialize(self) -> str: + return json.dumps([event.to_dict() for event in self._debug_events], indent=4) diff --git a/backends/arm/ethosu/backend.py b/backends/arm/ethosu/backend.py index 18fd127f116..c748cf96e93 100644 --- a/backends/arm/ethosu/backend.py +++ b/backends/arm/ethosu/backend.py @@ -16,7 +16,7 @@ from executorch.backends.arm.arm_vela import vela_compile -from executorch.backends.arm.tosa_backend import TOSABackend +from executorch.backends.arm.tosa.backend import TOSABackend from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export.exported_program import ExportedProgram diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index efbd6705615..d76b29eb1d9 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -11,7 +11,7 @@ is_ethosu, ) # usort: skip from executorch.backends.arm.ethosu import EthosUBackend -from executorch.backends.arm.tosa_partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import DelegationSpec from torch.fx.passes.operator_support import OperatorSupportBase diff --git a/backends/arm/operator_support/TARGETS b/backends/arm/operator_support/TARGETS index e14552fd016..8f6721bd911 100644 --- a/backends/arm/operator_support/TARGETS +++ b/backends/arm/operator_support/TARGETS @@ -1,11 +1,12 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "operator_support", srcs = glob(["*.py"]), deps = [ + "//executorch/backends/arm:constants", "//executorch/backends/arm/_passes:passes", - "//executorch/backends/arm:tosa_specification", + "//executorch/backends/arm/tosa:tosa", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/xnnpack/_passes:xnnpack_passes", "//executorch/exir:lib", diff --git a/backends/arm/operator_support/__init__.py b/backends/arm/operator_support/__init__.py index 2075e0f554f..7b73cddad37 100644 --- a/backends/arm/operator_support/__init__.py +++ b/backends/arm/operator_support/__init__.py @@ -6,6 +6,7 @@ # pyre-unsafe from . import ( # noqa + clone_dim_order_support, convolution_support, embedding_support, ethos_u55_support, @@ -17,6 +18,6 @@ right_shift_support, sin_cos_support, slice_copy_support, - to_copy_support, + to_dim_order_copy_support, tosa_supported_operators, ) diff --git a/backends/arm/operator_support/clone_dim_order_support.py b/backends/arm/operator_support/clone_dim_order_support.py new file mode 100644 index 00000000000..1397b74bf38 --- /dev/null +++ b/backends/arm/operator_support/clone_dim_order_support.py @@ -0,0 +1,68 @@ +# Copyright 2025 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 + +import torch +import torch.fx as fx +from executorch.backends.arm.operator_support.tosa_supported_operators import ( + register_tosa_support_check, + SupportedTOSAOperatorCheck, +) +from executorch.backends.arm.tosa import TosaSpecification +from executorch.exir.dialects._ops import ops as exir_ops + +logger = logging.getLogger(__name__) + + +@register_tosa_support_check +class CloneSupported(SupportedTOSAOperatorCheck): + targets = [exir_ops.edge.dim_order_ops._clone_dim_order.default] + + tosa_specs = [ + TosaSpecification.create_from_string("TOSA-1.0+INT"), + TosaSpecification.create_from_string("TOSA-1.0+FP"), + ] + + def is_node_tosa_supported( + self, node: fx.Node, tosa_spec: TosaSpecification + ) -> bool: + if node.target not in self.targets: + self.reporter.report_reject(node, f"Target {node.target} is not supported.") + return False + + input_node = node.args[0] + if not isinstance(input_node, fx.Node): + self.reporter.report_reject(node, "Non tensor clones are not supported") + return False + + # Check input node + if len(node.all_input_nodes) != 1: + self.reporter.report_reject( + node, f"Expected 1 input node, got {len(node.all_input_nodes)}" + ) + return False + + input_val = node.all_input_nodes[0].meta["val"] + if not isinstance(input_val, torch._subclasses.FakeTensor): + self.reporter.report_reject(node, "Expected input to be a FakeTensor.") + return False + + input_dtype = input_val.dtype + + # Check output node + output_val = node.meta["val"] + if not isinstance(output_val, torch._subclasses.FakeTensor): + self.reporter.report_reject(node, "Expected output to be a FakeTensor.") + return False + + if output_val.dtype != input_dtype: + self.reporter.report_reject( + node, + f"Input dtype {input_val.dtype} does not match {output_val.dtype}.", + ) + return False + + return True diff --git a/backends/arm/operator_support/convolution_support.py b/backends/arm/operator_support/convolution_support.py index 692d744025f..6e9d3b3528e 100644 --- a/backends/arm/operator_support/convolution_support.py +++ b/backends/arm/operator_support/convolution_support.py @@ -11,7 +11,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/embedding_support.py b/backends/arm/operator_support/embedding_support.py index 58a3a3e3edb..bf95014e575 100644 --- a/backends/arm/operator_support/embedding_support.py +++ b/backends/arm/operator_support/embedding_support.py @@ -11,7 +11,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index a1b5de85d08..bf9e29d5cb7 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/backends/arm/operator_support/ethos_u55_support.py @@ -14,7 +14,7 @@ from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.insert_table_ops import TableOps from executorch.backends.arm.operators.op_permute import transform_permutation_vector -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa.utils import tosa_shape 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 @@ -149,6 +149,8 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.flip.default, # REVERSE 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.scatter.src, exir_ops.edge.aten.scatter.value, exir_ops.edge.aten.select_scatter.default, diff --git a/backends/arm/operator_support/index_select_support.py b/backends/arm/operator_support/index_select_support.py index 9a48012f603..79f1d154a14 100644 --- a/backends/arm/operator_support/index_select_support.py +++ b/backends/arm/operator_support/index_select_support.py @@ -9,7 +9,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/index_tensor_support.py b/backends/arm/operator_support/index_tensor_support.py index 65ea5755d7e..4b226a9c407 100644 --- a/backends/arm/operator_support/index_tensor_support.py +++ b/backends/arm/operator_support/index_tensor_support.py @@ -12,7 +12,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/minmax_support.py b/backends/arm/operator_support/minmax_support.py index 1c4b0dd6c78..edbf7f61818 100644 --- a/backends/arm/operator_support/minmax_support.py +++ b/backends/arm/operator_support/minmax_support.py @@ -8,7 +8,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/pool_2d_support.py b/backends/arm/operator_support/pool_2d_support.py index 4ce0f7d75e7..ff453741f1f 100644 --- a/backends/arm/operator_support/pool_2d_support.py +++ b/backends/arm/operator_support/pool_2d_support.py @@ -15,7 +15,7 @@ from executorch.backends.arm.operators.operator_validation_utils import ( adjust_pooling_pad_if_needed, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/reduce_sum_support.py b/backends/arm/operator_support/reduce_sum_support.py index 0c614eb2bd5..4ff8f54ad69 100644 --- a/backends/arm/operator_support/reduce_sum_support.py +++ b/backends/arm/operator_support/reduce_sum_support.py @@ -10,7 +10,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/right_shift_support.py b/backends/arm/operator_support/right_shift_support.py index 454a3b525e3..5d3896e3643 100644 --- a/backends/arm/operator_support/right_shift_support.py +++ b/backends/arm/operator_support/right_shift_support.py @@ -13,7 +13,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops logger = logging.getLogger(__name__) diff --git a/backends/arm/operator_support/sin_cos_support.py b/backends/arm/operator_support/sin_cos_support.py index 03ce1da684b..dcdc20f8e4a 100644 --- a/backends/arm/operator_support/sin_cos_support.py +++ b/backends/arm/operator_support/sin_cos_support.py @@ -11,7 +11,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/arm/operator_support/slice_copy_support.py b/backends/arm/operator_support/slice_copy_support.py index ad9b5b250dd..14ca505635c 100644 --- a/backends/arm/operator_support/slice_copy_support.py +++ b/backends/arm/operator_support/slice_copy_support.py @@ -11,7 +11,7 @@ register_tosa_support_check, SupportedTOSAOperatorCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.dialects._ops import ops as exir_ops logger = logging.getLogger(__name__) diff --git a/backends/arm/operator_support/to_copy_support.py b/backends/arm/operator_support/to_copy_support.py deleted file mode 100644 index a10f3acb766..00000000000 --- a/backends/arm/operator_support/to_copy_support.py +++ /dev/null @@ -1,140 +0,0 @@ -# Copyright 2024-2025 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-unsafe -import copy -import logging - -import torch - -import torch.fx as fx - -from executorch.backends.arm.operator_support.tosa_supported_operators import ( - register_tosa_support_check, - SupportedTOSAOperatorCheck, -) -from executorch.backends.arm.tosa_specification import TosaSpecification -from executorch.exir.dialects._ops import ops as exir_ops - -logger = logging.getLogger(__name__) - - -@register_tosa_support_check -class ToCopySupported(SupportedTOSAOperatorCheck): - targets = [ - exir_ops.edge.aten._to_copy.default, - exir_ops.edge.dim_order_ops._to_dim_order_copy.default, - ] - - tosa_specs = [ - TosaSpecification.create_from_string("TOSA-1.0+INT"), - TosaSpecification.create_from_string("TOSA-1.0+FP"), - ] - - SupportedTypeDict = dict[torch.dtype, list[torch.dtype]] - - @staticmethod - def _merge_supported_types( - # pyre-ignore[11] - dtypes1: SupportedTypeDict, - dtypes2: SupportedTypeDict, - ) -> SupportedTypeDict: - merged_dtypes = copy.deepcopy( - dtypes1 - ) # Use deepcopy to avoid unintentionally modifying SUPPORTED_INT_TYPES - for k, v in dtypes2.items(): - merged_dtypes[k] = merged_dtypes.get(k, []) + v - return merged_dtypes - - SUPPORTED_INT_TYPES: SupportedTypeDict = { - torch.bool: [torch.int8, torch.int16, torch.int32], - torch.int8: [torch.bool, torch.int16, torch.int32], - torch.int16: [torch.bool, torch.int8, torch.int32], - torch.int32: [torch.bool, torch.int8, torch.int16], - } - SUPPORTED_FLOAT_TYPES: SupportedTypeDict = { - torch.int8: [torch.float16, torch.bfloat16, torch.float32], - torch.int16: [torch.float16, torch.bfloat16, torch.float32], - torch.int32: [torch.float16, torch.bfloat16, torch.float32], - torch.bfloat16: [torch.int8, torch.int16, torch.int32, torch.float32], - torch.float16: [torch.int8, torch.int16, torch.int32, torch.float32], - torch.float32: [ - torch.int8, - torch.int16, - torch.int32, - torch.bfloat16, - torch.float16, - ], - } - ALL_SUPPORTED_TYPES = _merge_supported_types( - SUPPORTED_INT_TYPES, SUPPORTED_FLOAT_TYPES - ) - POSSIBLE_TYPE_CONVERSIONS = {torch.int64: torch.int32} - - def is_node_tosa_supported( - self, node: fx.Node, tosa_spec: TosaSpecification - ) -> bool: - assert node.target in self.targets - - supported_dtypes = ( - self.ALL_SUPPORTED_TYPES - if tosa_spec.support_float() - else self.SUPPORTED_INT_TYPES - ) - # Take into account possible type conversions - supported_dtypes.update( - (k, supported_dtypes[v]) - for k, v in self.POSSIBLE_TYPE_CONVERSIONS.items() - if v in supported_dtypes - ) - - # Check input type - assert len(node.all_input_nodes) == 1 - input_val = node.all_input_nodes[0].meta["val"] - assert isinstance(input_val, torch._subclasses.FakeTensor) - input_dtype = input_val.dtype - if input_dtype not in supported_dtypes: - self.reporter.report_reject( - node, - f"Input dtype {input_val.dtype} is not supported in {node.target}.", - ) - return False - - # Check output type - output_val = node.meta["val"] - assert isinstance(output_val, torch._subclasses.FakeTensor) - if output_val.dtype not in supported_dtypes[input_dtype]: - self.reporter.report_reject( - node, - f"Output dtype {output_val.dtype} is not supported in " - f"{node.target} for input dtype {input_dtype}. " - f"Supported output types: " - f"{''.join(str(t) for t in supported_dtypes[input_dtype])}", - ) - return False - - # Check memory format (to_copy) - if "memory_format" in node.kwargs: - if node.kwargs["memory_format"] in (torch.preserve_format,): - self.reporter.report_reject( - node, - f"Argument 'memory_format' is not supported for " - f"{node.target} right now.", - ) - return False - - # Check dim_order (to_dim_order_copy) - if "dim_order" in node.kwargs: - dim_order = node.kwargs["dim_order"] - # pyre-ignore[6] - if dim_order != list(range(len(dim_order))): # type: ignore[arg-type] - self.reporter.report_reject( - node, - f"Argument {dim_order=} is not supported for " - f"{node.target} right now.", - ) - return False - - return True diff --git a/backends/arm/operator_support/to_dim_order_copy_support.py b/backends/arm/operator_support/to_dim_order_copy_support.py new file mode 100644 index 00000000000..e21f8a68ad6 --- /dev/null +++ b/backends/arm/operator_support/to_dim_order_copy_support.py @@ -0,0 +1,164 @@ +# Copyright 2024-2025 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-unsafe +import copy +import logging + +import torch + +import torch.fx as fx + +from executorch.backends.arm.operator_support.tosa_supported_operators import ( + register_tosa_support_check, + SupportedTOSAOperatorCheck, +) +from executorch.backends.arm.tosa import TosaSpecification +from executorch.exir.dialects._ops import ops as exir_ops + +logger = logging.getLogger(__name__) + +SupportedTypeDict = dict[torch.dtype, list[torch.dtype]] + + +@register_tosa_support_check +class ToCopySupported(SupportedTOSAOperatorCheck): + targets = [ + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + ] + + tosa_specs = [ + TosaSpecification.create_from_string("TOSA-1.0+INT"), + TosaSpecification.create_from_string("TOSA-1.0+FP"), + ] + + @staticmethod + def _merge_supported_types( + # pyre-ignore[11] + dtypes1: SupportedTypeDict, + dtypes2: SupportedTypeDict, + ) -> SupportedTypeDict: + merged_dtypes = copy.deepcopy( + dtypes1 + ) # Use deepcopy to avoid unintentionally modifying SUPPORTED_INT_TYPES + for k, v in dtypes2.items(): + merged_dtypes[k] = merged_dtypes.get(k, []) + v + return merged_dtypes + + SUPPORTED_INT_TYPES: SupportedTypeDict = { + torch.bool: [torch.bool, torch.int8, torch.int16, torch.int32], + torch.int8: [torch.bool, torch.int8, torch.int16, torch.int32], + torch.int16: [torch.bool, torch.int8, torch.int16, torch.int32], + torch.int32: [torch.bool, torch.int8, torch.int16, torch.int32], + torch.int64: [torch.bool, torch.int8, torch.int16, torch.int32], + } + SUPPORTED_FLOAT_TYPES: SupportedTypeDict = { + torch.int8: [torch.int8, torch.float16, torch.bfloat16, torch.float32], + torch.int16: [torch.int16, torch.float16, torch.bfloat16, torch.float32], + torch.int32: [torch.int32, torch.float16, torch.bfloat16, torch.float32], + # INT64 inputs to casts *should* be ok, since they should be rejected by + # CheckInt64InputsAndOutputs if the cast can't be done AOT. + torch.int64: [ + torch.int8, + torch.int16, + torch.int32, + torch.float16, + torch.bfloat16, + torch.float32, + ], + torch.bfloat16: [ + torch.int8, + torch.int16, + torch.int32, + torch.bfloat16, + torch.float32, + ], + torch.float16: [ + torch.int8, + torch.int16, + torch.int32, + torch.float16, + torch.float32, + ], + torch.float32: [ + torch.float32, + torch.int8, + torch.int16, + torch.int32, + torch.bfloat16, + torch.float16, + ], + } + ALL_SUPPORTED_TYPES = _merge_supported_types( + SUPPORTED_INT_TYPES, SUPPORTED_FLOAT_TYPES + ) + + def is_node_tosa_supported( + self, node: fx.Node, tosa_spec: TosaSpecification + ) -> bool: + + supported_dtypes: SupportedTypeDict = {} + if tosa_spec.support_integer(): + supported_dtypes = self._merge_supported_types( + self.SUPPORTED_INT_TYPES, supported_dtypes + ) + if tosa_spec.support_float(): + supported_dtypes = self._merge_supported_types( + self.SUPPORTED_FLOAT_TYPES, supported_dtypes + ) + + if len(node.all_input_nodes) != 1: + self.reporter.report_reject( + node, + ( + "Expected exactly one input node, " + f"got {len(node.all_input_nodes)} for {node.target}." + ), + ) + return False + input_val = node.all_input_nodes[0].meta["val"] + if not isinstance(input_val, torch._subclasses.FakeTensor): + self.reporter.report_reject( + node, + ( + "Invalid or missing meta: expected FakeTensor input, got " + f"{type(input_val).__name__} for {node.target}." + ), + ) + return False + + # Check input type + input_dtype = input_val.dtype + if input_dtype not in supported_dtypes: + self.reporter.report_reject( + node, + f"Input dtype {input_val.dtype} is not supported in {node.target}.", + ) + return False + + # Check output type + output_val = node.meta["val"] + if not isinstance(output_val, torch._subclasses.FakeTensor): + self.reporter.report_reject( + node, + ( + "Invalid or missing meta: expected FakeTensor output, got " + f"{type(output_val).__name__} for {node.target}." + ), + ) + return False + if output_val.dtype not in supported_dtypes[input_dtype]: + self.reporter.report_reject( + node, + ( + f"Output dtype {output_val.dtype} is not supported in " + f"{node.target} for input dtype {input_dtype}. " + f"Supported output types: " + f"{', '.join(str(t) for t in supported_dtypes[input_dtype])}" + ), + ) + return False + + return True diff --git a/backends/arm/operator_support/tosa_profile_supported_op_lists.py b/backends/arm/operator_support/tosa_profile_supported_op_lists.py new file mode 100644 index 00000000000..d3207c65dff --- /dev/null +++ b/backends/arm/operator_support/tosa_profile_supported_op_lists.py @@ -0,0 +1,247 @@ +# Copyright 2025 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 typing import Final, Set + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops + + +# INT profile: ops supported via native TOSA ops, decompositions/transformations, precompute, TableOps, etc. +TOSA_PRO_INT_SupportList: Final[Set] = { + exir_ops.edge.aten.abs.default, + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.any.default, + exir_ops.edge.aten.any.dim, + exir_ops.edge.aten.any.dims, + exir_ops.edge.aten.bitwise_and.Tensor, + exir_ops.edge.aten.bitwise_or.Tensor, + exir_ops.edge.aten.bitwise_xor.Tensor, + exir_ops.edge.aten.bitwise_and.Scalar, + exir_ops.edge.aten.bitwise_or.Scalar, + exir_ops.edge.aten.bitwise_xor.Scalar, + exir_ops.edge.aten.logical_and.default, + exir_ops.edge.aten.logical_or.default, + exir_ops.edge.aten.logical_xor.default, + exir_ops.edge.aten.logical_not.default, + exir_ops.edge.aten.arange.start_step, + exir_ops.edge.aten.expand_copy.default, + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.ceil.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.aten.cumsum.default, + exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.hardsigmoid.default, + exir_ops.edge.aten.hardtanh.default, + exir_ops.edge.aten.hardswish.default, + exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.div.Tensor_mode, + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.erf.default, + exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.expm1.default, + exir_ops.edge.aten.log.default, + exir_ops.edge.aten.linear.default, + exir_ops.edge.aten.split_with_sizes_copy.default, + exir_ops.edge.aten.floor.default, + exir_ops.edge.aten.full.default, + exir_ops.edge.aten.full_like.default, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.gt.Scalar, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.lt.Tensor, + exir_ops.edge.aten.lt.Scalar, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ne.Scalar, + exir_ops.edge.aten.neg.default, + exir_ops.edge.aten.add.Scalar, + exir_ops.edge.aten.sub.Scalar, + exir_ops.edge.aten.mul.Scalar, + exir_ops.edge.aten.div.Scalar, + exir_ops.edge.aten._native_batch_norm_legit_no_training.default, + exir_ops.edge.aten.native_layer_norm.default, + exir_ops.edge.aten.native_group_norm.default, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.mm.default, + exir_ops.edge.aten.minimum.default, + exir_ops.edge.aten.maximum.default, + exir_ops.edge.aten.repeat.default, + exir_ops.edge.aten.reciprocal.default, + exir_ops.edge.aten.relu.default, + exir_ops.edge.aten.leaky_relu.default, + exir_ops.edge.aten.sqrt.default, + exir_ops.edge.aten.rsqrt.default, + exir_ops.edge.aten.round.default, + exir_ops.edge.aten._softmax.default, + exir_ops.edge.aten.select_copy.int, + exir_ops.edge.aten._log_softmax.default, + exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.tanh.default, + exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.aten.var.correction, + exir_ops.edge.aten.var.dim, + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.pow.Tensor_Scalar, + exir_ops.edge.aten.pow.Tensor_Tensor, + exir_ops.edge.aten.where.self, + operator.getitem, + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, + exir_ops.edge.aten.constant_pad_nd.default, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.eye.default, + exir_ops.edge.aten.linspace.default, + exir_ops.edge.aten.bitwise_left_shift.Tensor, + exir_ops.edge.aten.__lshift__.Scalar, + torch.ops.aten.scalar_tensor.default, + exir_ops.edge.aten.gelu.default, + exir_ops.edge.aten.alias_copy.default, + exir_ops.edge.aten.sinh.default, + exir_ops.edge.aten.atan.default, + exir_ops.edge.aten.acosh.default, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + exir_ops.edge.aten.sign.default, + exir_ops.edge.aten.asin.default, + exir_ops.edge.aten.atanh.default, + exir_ops.edge.aten.addmm.default, + exir_ops.edge.aten.masked_fill.Scalar, + exir_ops.edge.aten.asinh.default, + exir_ops.edge.aten.cosh.default, + exir_ops.edge.aten.glu.default, + exir_ops.edge.aten.logit.default, + exir_ops.edge.aten.acos.default, + exir_ops.edge.aten.elu.default, +} + + +# FP profile: ops supported via native TOSA ops, decompositions/transformations, precompute, etc. +TOSA_PRO_FP_SupportList: Final[Set] = { + exir_ops.edge.aten.abs.default, + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.any.default, + exir_ops.edge.aten.any.dim, + exir_ops.edge.aten.any.dims, + exir_ops.edge.aten.logical_and.default, + exir_ops.edge.aten.logical_or.default, + exir_ops.edge.aten.logical_xor.default, + exir_ops.edge.aten.logical_not.default, + exir_ops.edge.aten.arange.start_step, + exir_ops.edge.aten.expand_copy.default, + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.ceil.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.aten.cumsum.default, + exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.hardsigmoid.default, + exir_ops.edge.aten.hardtanh.default, + exir_ops.edge.aten.hardswish.default, + exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.div.Tensor_mode, + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.erf.default, + exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.expm1.default, + exir_ops.edge.aten.log.default, + exir_ops.edge.aten.linear.default, + exir_ops.edge.aten.split_with_sizes_copy.default, + exir_ops.edge.aten.floor.default, + exir_ops.edge.aten.full.default, + exir_ops.edge.aten.full_like.default, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.ge.Scalar, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.gt.Scalar, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.lt.Tensor, + exir_ops.edge.aten.lt.Scalar, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ne.Scalar, + exir_ops.edge.aten.neg.default, + exir_ops.edge.aten.add.Scalar, + exir_ops.edge.aten.sub.Scalar, + exir_ops.edge.aten.mul.Scalar, + exir_ops.edge.aten.div.Scalar, + exir_ops.edge.aten._native_batch_norm_legit_no_training.default, + exir_ops.edge.aten.native_layer_norm.default, + exir_ops.edge.aten.native_group_norm.default, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.mean.dim, + exir_ops.edge.aten.mm.default, + exir_ops.edge.aten.minimum.default, + exir_ops.edge.aten.maximum.default, + exir_ops.edge.aten.repeat.default, + exir_ops.edge.aten.reciprocal.default, + exir_ops.edge.aten.relu.default, + exir_ops.edge.aten.leaky_relu.default, + exir_ops.edge.aten.sqrt.default, + exir_ops.edge.aten.rsqrt.default, + exir_ops.edge.aten.round.default, + exir_ops.edge.aten._softmax.default, + exir_ops.edge.aten.select_copy.int, + exir_ops.edge.aten._log_softmax.default, + exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.tanh.default, + exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.aten.var.correction, + exir_ops.edge.aten.var.dim, + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.pow.Tensor_Scalar, + exir_ops.edge.aten.pow.Tensor_Tensor, + exir_ops.edge.aten.where.self, + operator.getitem, + exir_ops.edge.aten.constant_pad_nd.default, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.eye.default, + exir_ops.edge.aten.linspace.default, + exir_ops.edge.aten.bitwise_left_shift.Tensor, + exir_ops.edge.aten.__lshift__.Scalar, + torch.ops.aten.scalar_tensor.default, + exir_ops.edge.aten.gelu.default, + exir_ops.edge.aten.alias_copy.default, + exir_ops.edge.aten.sinh.default, + exir_ops.edge.aten.atan.default, + exir_ops.edge.aten.acosh.default, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + exir_ops.edge.aten.sign.default, + exir_ops.edge.aten.asin.default, + exir_ops.edge.aten.atanh.default, + exir_ops.edge.aten.addmm.default, + exir_ops.edge.aten.masked_fill.Scalar, + exir_ops.edge.aten.asinh.default, + exir_ops.edge.aten.cosh.default, + exir_ops.edge.aten.glu.default, + exir_ops.edge.aten.logit.default, + exir_ops.edge.aten.acos.default, + exir_ops.edge.aten.elu.default, +} + + +__all__ = [ + "TOSA_PRO_INT_SupportList", + "TOSA_PRO_FP_SupportList", +] diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index fb38c07b8a6..b580fbb9a9a 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -26,7 +26,11 @@ EthosU55TransposeCheck, EthosU55ViewCheck, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.operator_support.tosa_profile_supported_op_lists import ( + TOSA_PRO_FP_SupportList, + TOSA_PRO_INT_SupportList, +) +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir import ExportedProgram from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops @@ -106,17 +110,21 @@ def tosa_support_factory( Additional checks can be supplied to avoid partitioning additional nodes. """ # Postive checks: Add nodes to partitioning - positive_checks: list[OperatorSupportBase] = [ - BaseTOSASupportList(), - *[ - check(tosa_spec, reporter) - for check in get_registered_tosa_support_checks(tosa_spec) - ], + positive_checks: list[OperatorSupportBase] = [] + + if tosa_spec.support_integer(): + positive_checks.append(TOSAProINTSupportList()) + if tosa_spec.support_float(): + positive_checks.append(TOSAProFPSupportList()) + # TODO: Refactor to use TOSAProSupportLists + negtive checks + positive_checks += [ + check(tosa_spec, reporter) + for check in get_registered_tosa_support_checks(tosa_spec) ] # Negative checks: Remove nodes from partitioning negative_checks: list[OperatorSupportBase] = [ - CheckInt64Inputs(exported_program, reporter), + CheckInt64InputsAndOutputs(exported_program, reporter), CheckFloat64Inputs(exported_program, reporter), RankCheck(reporter, max_rank=5), *[ @@ -143,123 +151,30 @@ def tosa_support_factory( ) -class BaseTOSASupportList(OperatorSupportBase): +class TOSAProINTSupportList(OperatorSupportBase): + """ + TOSA_PRO_INT_SupportList: + Ops supported in INT profile via native TOSA ops, decomposition/transformation, pre-compute, or TableOps + """ def is_node_supported( self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node ) -> bool: - supported = node.op == "call_function" and node.target in [ - exir_ops.edge.aten.abs.default, - exir_ops.edge.aten.add.Tensor, - exir_ops.edge.aten.any.default, - exir_ops.edge.aten.any.dim, - exir_ops.edge.aten.any.dims, - exir_ops.edge.aten.logical_and.default, - exir_ops.edge.aten.logical_or.default, - exir_ops.edge.aten.logical_xor.default, - exir_ops.edge.aten.logical_not.default, - exir_ops.edge.aten.arange.start_step, - exir_ops.edge.aten.bitwise_and.Tensor, - exir_ops.edge.aten.bitwise_or.Tensor, - exir_ops.edge.aten.bitwise_xor.Tensor, - exir_ops.edge.aten.bitwise_and.Scalar, - exir_ops.edge.aten.bitwise_or.Scalar, - exir_ops.edge.aten.bitwise_xor.Scalar, - exir_ops.edge.aten.expand_copy.default, - exir_ops.edge.aten.cat.default, - exir_ops.edge.aten.ceil.default, - exir_ops.edge.aten.clamp.default, - exir_ops.edge.aten.bmm.default, - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.hardsigmoid.default, - exir_ops.edge.aten.hardtanh.default, - exir_ops.edge.aten.hardswish.default, - exir_ops.edge.aten.div.Tensor, - exir_ops.edge.aten.eq.Tensor, - exir_ops.edge.aten.eq.Scalar, - exir_ops.edge.aten.erf.default, - exir_ops.edge.aten.exp.default, - exir_ops.edge.aten.log.default, - exir_ops.edge.aten.linear.default, - exir_ops.edge.aten.split_with_sizes_copy.default, - exir_ops.edge.aten.floor.default, - exir_ops.edge.aten.full.default, - exir_ops.edge.aten.full_like.default, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.ge.Scalar, - exir_ops.edge.aten.gt.Tensor, - exir_ops.edge.aten.gt.Scalar, - exir_ops.edge.aten.le.Tensor, - exir_ops.edge.aten.le.Scalar, - exir_ops.edge.aten.lt.Tensor, - exir_ops.edge.aten.lt.Scalar, - exir_ops.edge.aten.mul.Tensor, - exir_ops.edge.aten.ne.Tensor, - exir_ops.edge.aten.ne.Scalar, - exir_ops.edge.aten.neg.default, - exir_ops.edge.aten.add.Scalar, - exir_ops.edge.aten.sub.Scalar, - exir_ops.edge.aten.mul.Scalar, - exir_ops.edge.aten.div.Scalar, - exir_ops.edge.aten._native_batch_norm_legit_no_training.default, - exir_ops.edge.aten.native_layer_norm.default, - exir_ops.edge.aten.native_group_norm.default, - exir_ops.edge.aten.sigmoid.default, - exir_ops.edge.aten.mean.dim, - exir_ops.edge.aten.mm.default, - exir_ops.edge.aten.minimum.default, - exir_ops.edge.aten.maximum.default, - exir_ops.edge.aten.repeat.default, - exir_ops.edge.aten.reciprocal.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.leaky_relu.default, - exir_ops.edge.aten.sqrt.default, - exir_ops.edge.aten.rsqrt.default, - exir_ops.edge.aten.round.default, - exir_ops.edge.aten._softmax.default, - exir_ops.edge.aten.select_copy.int, - exir_ops.edge.aten._log_softmax.default, - exir_ops.edge.aten.sub.Tensor, - exir_ops.edge.aten.tanh.default, - exir_ops.edge.aten.upsample_bilinear2d.vec, - exir_ops.edge.aten.upsample_nearest2d.vec, - exir_ops.edge.aten.var.correction, - exir_ops.edge.aten.var.dim, - exir_ops.edge.aten.view_copy.default, - exir_ops.edge.aten.clone.default, - exir_ops.edge.aten.unsqueeze_copy.default, - exir_ops.edge.aten.squeeze_copy.dims, - exir_ops.edge.aten.pow.Tensor_Scalar, - exir_ops.edge.aten.pow.Tensor_Tensor, - exir_ops.edge.aten.where.self, - operator.getitem, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.quantize_per_channel.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, - exir_ops.edge.aten.constant_pad_nd.default, - exir_ops.edge.aten.amax.default, - exir_ops.edge.aten.amin.default, - exir_ops.edge.aten.eye.default, - exir_ops.edge.aten.linspace.default, - exir_ops.edge.aten.bitwise_left_shift.Tensor, - exir_ops.edge.aten.__lshift__.Scalar, - torch.ops.aten.scalar_tensor.default, - exir_ops.edge.aten.gelu.default, - exir_ops.edge.aten.alias_copy.default, - exir_ops.edge.aten.sinh.default, - exir_ops.edge.aten.atan.default, - exir_ops.edge.aten.acosh.default, - exir_ops.edge.aten._adaptive_avg_pool2d.default, - exir_ops.edge.aten.sign.default, - exir_ops.edge.aten.asin.default, - exir_ops.edge.aten.atanh.default, - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.masked_fill.Scalar, - exir_ops.edge.aten.asinh.default, - ] - return supported + return node.op == "call_function" and node.target in TOSA_PRO_INT_SupportList + + +class TOSAProFPSupportList(OperatorSupportBase): + """ + TOSA_PRO_FP_SupportList: + Ops supported in FP profile via native TOSA ops, decomposition/transformation, pre-compute + """ + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + + return node.op == "call_function" and node.target in TOSA_PRO_FP_SupportList class NeedsDecompositionCheck(OperatorSupportBase): @@ -298,6 +213,8 @@ def is_node_supported( exir_ops.edge.aten.leaky_relu.default: None, exir_ops.edge.aten.round.default: None, exir_ops.edge.aten.addmm.default: None, + exir_ops.edge.aten.glu.default: None, + exir_ops.edge.aten.logit.default: None, } if node.target in needs_decomp_dict: @@ -446,7 +363,18 @@ def is_node_supported( return True -class CheckInt64Inputs(OperatorSupportBase): +class CheckInt64InputsAndOutputs(OperatorSupportBase): + """TOSA does not support int64 tensors so in general, ops with int64 inputs or outputs should not be partitioned. + There are however some exceptions: + - Nodes with int64 output can be partitioned if they are constant, within int32, + and all users cast to something else. In this case, the int64 tensor can safely be cast to int32 AOT. + - Nodes with int64 output can be partitioned if all users are getitem with non-int64 output. + In this case, there are multiple outputs and the int64 ones are not used. + - Nodes with int64 inputs can be partitioned if the inputs are constant placeholders, or constant + ops fulfilling the criteria above. + Note that we don't check placeholders here, they are partitioned based on whether their users are partitioned + or not. + """ def __init__( self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter @@ -457,27 +385,85 @@ def __init__( if spec.kind == InputKind.USER_INPUT ] self.reporter = reporter + self.int32_min = torch.iinfo(torch.int32).min + self.int32_max = torch.iinfo(torch.int32).max super().__init__() + def inside_int32_bounds(self, node: torch.fx.Node) -> bool: + """Node is assumed to be call_function with int64 output.""" + if isinstance(node.target, str): + return False + data = node.target(*node.args, **node.kwargs) + min_val, max_val = int(torch.min(data)), int(torch.max(data)) + return min_val >= self.int32_min and max_val <= self.int32_max + def is_node_supported( self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node ) -> bool: + vals = node.meta["val"] + tensor_list = vals if isinstance(vals, (list, tuple)) else [vals] + + any_int64 = any(tensor.dtype == torch.int64 for tensor in tensor_list) + # Don't partition nodes with int64 output... + if any_int64: + # ... Except for constant ops that are directly cast to something non-int64. + # This could be an explicit cast, or something like a less than that outputs a different dtype than the input. + users_output_non_int64 = all( + get_first_fake_tensor(output_node).dtype != torch.int64 + for output_node in node.users + ) + if ( + node.target in ComputeConstantOpsAOT.targeted_ops + and users_output_non_int64 + ): + if not self.inside_int32_bounds(node): + self.reporter.report_reject( + node, "Constant node outside int32 range." + ) + return False + # Will never have input nodes, safe to return True + return True + + # ... Or ops with multiple outputs where only non-int64 are used. + users_are_getitem = all( + user.target == operator.getitem for user in node.users + ) + if users_are_getitem and users_output_non_int64: + # Passed output check, go to input check. + pass + else: + self.reporter.report_reject( + node, "Non-constant node with int64 output." + ) + return False + + # Ops with int64 inputs are only partitioned if input nodes are constant and will be partitioned. + # If it is not partitioned, the partition will get an int64 input and fail. for input_node in node.all_input_nodes: - # We can cast constant placeholders and constant ops AOT, such int64 are ok. - # Otherwise, don't partition if one or more inputs are int64. + tensor_in = get_first_fake_tensor(input_node) + if tensor_in.dtype != torch.int64: + continue + # Constant placeholder if ( - input_node.name in self.input_names - or not input_node.op == "placeholder" + input_node.op != "call_function" + and input_node.name not in self.input_names ): - tensor = get_first_fake_tensor(input_node) - if tensor.dtype == torch.int64: - if input_node.target not in ComputeConstantOpsAOT.targeted_ops: - self.reporter.report_reject( - node, - f"Had int64 input {input_node.name} that couldn't be handled.", - ) - return False + continue + # Constant operator + if input_node.op == "call_function": + if input_node.target in ComputeConstantOpsAOT.targeted_ops: + # This is not perfect since the input_node can still be rejected by other checks but + # this should cover the majority of cases. + if self.is_node_supported( + None, input_node # type: ignore[arg-type] #(we don't use 'submodules') + ): + continue + self.reporter.report_reject( + node, f"Non-constant int64 input {input_node.name}" + ) + return False + return True diff --git a/backends/arm/operators/TARGETS b/backends/arm/operators/TARGETS index 7f97450849d..2c255b3c17a 100644 --- a/backends/arm/operators/TARGETS +++ b/backends/arm/operators/TARGETS @@ -1,21 +1,22 @@ # @noautodeps -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "node_visitor", srcs = ["node_visitor.py"], deps = [ - "//executorch/backends/arm:tosa_mapping", - "//executorch/backends/arm:tosa_specification", + "//executorch/backends/arm/debug:schema", + "//executorch/backends/arm/tosa:mapping", + "//executorch/backends/arm/tosa:tosa", ], ) -python_library( +runtime.python_library( name = "operator_validation_utils", srcs = ["operator_validation_utils.py"], ) -python_library( +runtime.python_library( name = "ops", srcs = glob(["op_*.py", "ops_*.py"]), deps = [ @@ -23,15 +24,15 @@ python_library( "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/tosa:tosa", ":node_visitor", ":operator_validation_utils", - "//executorch/backends/arm:tosa_mapping", - "//executorch/backends/arm:tosa_quant_utils", - "//executorch/backends/arm:tosa_utils", + "//executorch/backends/arm/tosa:mapping", + "//executorch/backends/arm/tosa:quant_utils", + "//executorch/backends/arm/tosa:utils", "//executorch/backends/arm/_passes:passes", "//executorch/exir:lib", ], ) -python_library( +runtime.python_library( name = "lib", srcs = ["__init__.py"], deps = [ diff --git a/backends/arm/operators/__init__.py b/backends/arm/operators/__init__.py index 260299d6423..f7a9638254e 100644 --- a/backends/arm/operators/__init__.py +++ b/backends/arm/operators/__init__.py @@ -15,6 +15,7 @@ op_avg_pool2d, op_bmm, op_cat, + op_ceil, op_clamp, op_constant_pad_nd, op_conv2d, @@ -22,12 +23,14 @@ op_eq, op_erf, op_exp, + op_floor, op_ge, op_gt, op_index_select, op_index_tensor, op_le, op_log, + op_logical_not, op_lt, op_max_pool2d, op_maximum, @@ -48,7 +51,6 @@ op_sum, op_table, op_tanh, - op_to_copy, op_to_dim_order_copy, op_transpose, op_upsample_bilinear2d, @@ -57,5 +59,4 @@ op_where, ops_binary, ops_identity, - ops_unary, ) diff --git a/backends/arm/operators/node_visitor.py b/backends/arm/operators/node_visitor.py index afc80bbb849..54a81bdaaff 100644 --- a/backends/arm/operators/node_visitor.py +++ b/backends/arm/operators/node_visitor.py @@ -5,12 +5,15 @@ # pyre-unsafe -from typing import Any, Dict, List +import json +from typing import Any, Dict, List, Optional import torch -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder +from executorch.backends.arm.debug.schema import DebugHook +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.specification import TosaSpecification from torch.export import ExportedProgram @@ -29,9 +32,43 @@ class NodeVisitor: TosaSpecification.create_from_string("TOSA-1.0+FP"), ] - def __init__(self, exported_program: ExportedProgram, tosa_spec: TosaSpecification): + def __init__( + self, + exported_program: ExportedProgram, + tosa_spec: TosaSpecification, + debug_hook: Optional[DebugHook] = None, + ): self._exported_program = exported_program self.tosa_spec = tosa_spec + self.debug_hook = debug_hook + + def _serialize_operator( + self, + node: torch.fx.Node, + tosa_graph: Any, + tosa_op: Any, + inputs: List[str], + outputs: List[str], + attributes: Optional[Any] = None, + ) -> None: + op_location = "" + if self.debug_hook: + debug_info = self.debug_hook.add( + node, + tosa_op=outputs[0], + tosa_op_id=tosa_op, + ) + + if self.debug_hook.mode == ArmCompileSpecBuilder.DebugMode.TOSA: + op_location = json.dumps(debug_info.to_dict()) + + tosa_graph.addOperator( + tosa_op, + inputs=inputs, + outputs=outputs, + attributes=attributes, + location=op_location, + ) def define_node( self, diff --git a/backends/arm/operators/op_abs.py b/backends/arm/operators/op_abs.py index 3000af50ed7..625293d66e0 100644 --- a/backends/arm/operators/op_abs.py +++ b/backends/arm/operators/op_abs.py @@ -6,8 +6,8 @@ # pyre-unsafe from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.quant_utils as tqutils +import executorch.backends.arm.tosa.utils as tutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -123,7 +123,9 @@ def define_node( ) # MI lowering - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().ABS, [inputs[0].name], [output.name], diff --git a/backends/arm/operators/op_add.py b/backends/arm/operators/op_add.py index 7a022b54395..a8f0c3fe14d 100644 --- a/backends/arm/operators/op_add.py +++ b/backends/arm/operators/op_add.py @@ -7,8 +7,8 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.quant_utils as tqutils +import executorch.backends.arm.tosa.utils as tutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -19,8 +19,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -47,20 +47,25 @@ def define_node( validate_num_inputs(self.target, inputs, 2) validate_same_dtype(self.target, [*inputs, output], ts) + valid_dtypes = [] + if self.tosa_spec.support_integer(): + valid_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32]) + if self.tosa_spec.support_float(): + valid_dtypes.extend([ts.DType.INT32]) + validate_valid_dtype( self.target, [*inputs, output], - [ts.DType.INT8, ts.DType.INT32], + valid_dtypes, output.tosa_spec, ) - scale_back = 1.0 if inputs[0].dtype == ts.DType.INT8: - rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32( + rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32_maxscale( tosa_graph, inputs, node, self.tosa_spec ) else: - # input[0].dtype == ts.DType.INT32 + # input[0].dtype == ts.DType.INT16 or ts.DType.INT32 # Non quantized input, natively support by TOSA.ADD rescaled_inputs = inputs @@ -68,13 +73,15 @@ def define_node( broadcasted_shape = tutils.tosa_shape(output.shape, output.dim_order) add_output = tosa_graph.addIntermediate(broadcasted_shape, ts.DType.INT32) else: - # output.dtype == ts.DType.INT32 + # output.dtype == ts.DType.INT16 or ts.DType.INT32 add_output = output input1, input2 = rescaled_inputs # Do the INT32 Add - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().ADD, [input1.name, input2.name], [add_output.name], @@ -85,7 +92,12 @@ def define_node( # Scale output back to 8 bit # pyre-ignore tqutils.insert_rescale_op_to_int8( - tosa_graph, add_output, scale_back, node, self.tosa_spec + tosa_graph, + add_output, + scale_back, + node, + compute_rescale=False, + tosa_spec=self.tosa_spec, ) # type: ignore[possibly-undefined] @@ -111,7 +123,7 @@ def define_node( validate_num_inputs(self.target, inputs, 2) validate_same_dtype(self.target, [*inputs, output], ts) - if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]: + if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32]: # Call the inherited define_node for handling integers super().define_node(node, tosa_graph, inputs, output) else: @@ -123,7 +135,9 @@ def define_node( input1, input2 = inputs # FP lowering - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().ADD, [input1.name, input2.name], [output.name], diff --git a/backends/arm/operators/op_amax.py b/backends/arm/operators/op_amax.py index 526d6ff35ec..1fb751597af 100644 --- a/backends/arm/operators/op_amax.py +++ b/backends/arm/operators/op_amax.py @@ -14,7 +14,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -61,6 +61,11 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.ReduceMaxAttribute(axis=input.dim_order.index(dim), nan_mode=1) - tosa_graph.addOperator( - ts.TosaOp.Op().REDUCE_MAX, [input.name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().REDUCE_MAX, + [input.name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_amin.py b/backends/arm/operators/op_amin.py index 85b0b757c85..9ebe78b946d 100644 --- a/backends/arm/operators/op_amin.py +++ b/backends/arm/operators/op_amin.py @@ -14,7 +14,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -61,6 +61,11 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.ReduceMinAttribute(axis=input.dim_order.index(dim), nan_mode=1) - tosa_graph.addOperator( - ts.TosaOp.Op().REDUCE_MIN, [input.name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().REDUCE_MIN, + [input.name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_any.py b/backends/arm/operators/op_any.py index 0ac307aedd4..0c47d6b190a 100644 --- a/backends/arm/operators/op_any.py +++ b/backends/arm/operators/op_any.py @@ -16,7 +16,7 @@ validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg # type: ignore +from executorch.backends.arm.tosa.mapping import TosaArg # type: ignore from torch.fx import Node @@ -52,6 +52,11 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.ReduceAnyAttribute(inputs[0].dim_order.index(dim)) - tosa_graph.addOperator( - ts.TosaOp.Op().REDUCE_ANY, [inputs[0].name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().REDUCE_ANY, + [inputs[0].name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_avg_pool2d.py b/backends/arm/operators/op_avg_pool2d.py index 9faf8272473..d28f5f27acf 100644 --- a/backends/arm/operators/op_avg_pool2d.py +++ b/backends/arm/operators/op_avg_pool2d.py @@ -22,8 +22,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -100,7 +100,9 @@ def _build_generic_avgpool2d( shape=[1], dtype=output.dtype, vals=[output_zp] ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().AVG_POOL2D, [input_tensor.name, input_zp_tensor.name, output_zp_tensor.name], [output.name], diff --git a/backends/arm/operators/op_bmm.py b/backends/arm/operators/op_bmm.py index c9bb0b003ee..382386ffa26 100644 --- a/backends/arm/operators/op_bmm.py +++ b/backends/arm/operators/op_bmm.py @@ -22,9 +22,9 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_quant_utils import build_rescale -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.quant_utils import build_rescale from tosa.RoundingMode import RoundingMode # type: ignore @@ -55,7 +55,7 @@ def define_node( validate_valid_dtype( self.target, [*inputs, output], - [ts.DType.INT8, ts.DType.FP32], + [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32], output.tosa_spec, ) @@ -78,7 +78,9 @@ def define_node( tosa_graph.addConst([1], inputs[1].dtype, [input1_zp], name=f"{node.name}_B_ZP") # Add the MATMUL to the TOSA graph. - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MATMUL, [ inputs[0].name, diff --git a/backends/arm/operators/op_cat.py b/backends/arm/operators/op_cat.py index 884bfb22a40..65b3e2a9549 100644 --- a/backends/arm/operators/op_cat.py +++ b/backends/arm/operators/op_cat.py @@ -14,7 +14,7 @@ from executorch.backends.arm.operators.operator_validation_utils import ( validate_num_inputs, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -47,7 +47,9 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.ConcatAttribute(dim) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().CONCAT, [tensor.name for tensor in tensors], [output.name], diff --git a/backends/arm/operators/op_ceil.py b/backends/arm/operators/op_ceil.py new file mode 100644 index 00000000000..5cf89710436 --- /dev/null +++ b/backends/arm/operators/op_ceil.py @@ -0,0 +1,54 @@ +# Copyright 2025 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, List + +import torch.fx + +from executorch.backends.arm.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.arm.operators.operator_validation_utils import ( + validate_num_inputs, + validate_same_dtype, + validate_valid_dtype, +) +from executorch.backends.arm.tosa import TosaSpecification + +from executorch.backends.arm.tosa.mapping import TosaArg + + +@register_node_visitor +class CeilVisitor(NodeVisitor): + target = "aten.ceil.default" + + # INT case should be handled by op_table + tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+FP")] + + def __init__(self, *args): + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + tosa_graph: Any, + inputs: List[TosaArg], + output: TosaArg, + ) -> None: + import serializer.tosa_serializer as ts # type: ignore # noqa: F401 + + validate_num_inputs(self.target, inputs, 1) + validate_same_dtype(self.target, [*inputs, output], ts) + validate_valid_dtype( + self.target, + inputs[0], + ts.DType.FP32, + output.tosa_spec, + ) + + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().CEIL, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_clamp.py b/backends/arm/operators/op_clamp.py index 2bdeb89a713..b0bf044a213 100644 --- a/backends/arm/operators/op_clamp.py +++ b/backends/arm/operators/op_clamp.py @@ -20,9 +20,9 @@ validate_same_dtype, validate_valid_dtype, ) +from executorch.backends.arm.tosa import TosaSpecification -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -90,8 +90,13 @@ def define_node( nan_mode=1, ) - tosa_graph.addOperator( - ts.TosaOp.Op().CLAMP, [inputs[0].name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().CLAMP, + [inputs[0].name], + [output.name], + attr, ) @@ -138,6 +143,11 @@ def define_node( nan_mode=1, ) - tosa_graph.addOperator( - ts.TosaOp.Op().CLAMP, [inputs[0].name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().CLAMP, + [inputs[0].name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_constant_pad_nd.py b/backends/arm/operators/op_constant_pad_nd.py index 147a1544ce9..562c4c9ea0e 100644 --- a/backends/arm/operators/op_constant_pad_nd.py +++ b/backends/arm/operators/op_constant_pad_nd.py @@ -21,8 +21,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -100,7 +100,9 @@ def define_node( shape=[1], dtype=pad_const_dtype, vals=[pad_const_val] ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().PAD, [inputs[0].name, padding.name, pad_const.name], [output.name], diff --git a/backends/arm/operators/op_conv2d.py b/backends/arm/operators/op_conv2d.py index 0bbe67c4beb..6bfe0ab21eb 100644 --- a/backends/arm/operators/op_conv2d.py +++ b/backends/arm/operators/op_conv2d.py @@ -20,10 +20,10 @@ from executorch.backends.arm.operators.operator_validation_utils import ( validate_num_inputs, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_quant_utils import build_rescale -from executorch.backends.arm.tosa_specification import TosaSpecification -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.quant_utils import build_rescale +from executorch.backends.arm.tosa.utils import tosa_shape @register_node_visitor @@ -157,7 +157,9 @@ def define_node( reshape_attr = ts.TosaSerializerAttribute() reshape_attr.ReshapeAttribute() - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().RESHAPE, [weight.name, shape.name], [weight_reshaped.name], @@ -188,7 +190,9 @@ def define_node( acc_type=acc_type, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, tosa_op, [ input.name, diff --git a/backends/arm/operators/op_cos.py b/backends/arm/operators/op_cos.py index e37db290d55..0350733190c 100644 --- a/backends/arm/operators/op_cos.py +++ b/backends/arm/operators/op_cos.py @@ -16,8 +16,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -44,4 +44,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().COS, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().COS, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_eq.py b/backends/arm/operators/op_eq.py index eb5b3000d6c..2136fe2e946 100644 --- a/backends/arm/operators/op_eq.py +++ b/backends/arm/operators/op_eq.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -68,9 +68,11 @@ def define_node( input_nodes = rescaled_inputs # Do the equal comparison - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().EQUAL, [input_nodes[0].name, input_nodes[1].name], - output.name, + [output.name], None, ) diff --git a/backends/arm/operators/op_erf.py b/backends/arm/operators/op_erf.py index e238c4fd80a..7797b61e562 100644 --- a/backends/arm/operators/op_erf.py +++ b/backends/arm/operators/op_erf.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -48,4 +48,6 @@ def define_node( ) # MI lowering - tosa_graph.addOperator(ts.TosaOp.Op().ERF, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().ERF, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_exp.py b/backends/arm/operators/op_exp.py index 96c077c838b..f5d5aef2213 100644 --- a/backends/arm/operators/op_exp.py +++ b/backends/arm/operators/op_exp.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -48,4 +48,6 @@ def define_node( output.tosa_spec, ) - tosa_graph.addOperator(ts.TosaOp.Op().EXP, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().EXP, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_floor.py b/backends/arm/operators/op_floor.py new file mode 100644 index 00000000000..77d712096fa --- /dev/null +++ b/backends/arm/operators/op_floor.py @@ -0,0 +1,54 @@ +# Copyright 2025 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, List + +import torch.fx + +from executorch.backends.arm.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.arm.operators.operator_validation_utils import ( + validate_num_inputs, + validate_same_dtype, + validate_valid_dtype, +) +from executorch.backends.arm.tosa import TosaSpecification + +from executorch.backends.arm.tosa.mapping import TosaArg + + +@register_node_visitor +class FloorVisitor(NodeVisitor): + target = "aten.floor.default" + + # INT case should be handled by op_table + tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+FP")] + + def __init__(self, *args): + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + tosa_graph: Any, + inputs: List[TosaArg], + output: TosaArg, + ) -> None: + import serializer.tosa_serializer as ts # type: ignore # noqa: F401 + + validate_num_inputs(self.target, inputs, 1) + validate_same_dtype(self.target, [*inputs, output], ts) + validate_valid_dtype( + self.target, + inputs[0], + ts.DType.FP32, + output.tosa_spec, + ) + + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().FLOOR, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_ge.py b/backends/arm/operators/op_ge.py index 723706702f0..c538e735880 100644 --- a/backends/arm/operators/op_ge.py +++ b/backends/arm/operators/op_ge.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -67,7 +67,9 @@ def define_node( # Update IO input_nodes = rescaled_inputs - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GREATER_EQUAL, [input_nodes[0].name, input_nodes[1].name], [output.name], diff --git a/backends/arm/operators/op_gt.py b/backends/arm/operators/op_gt.py index e79ed009e24..d407e28c1b6 100644 --- a/backends/arm/operators/op_gt.py +++ b/backends/arm/operators/op_gt.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -67,7 +67,9 @@ def define_node( # Update IO input_nodes = rescaled_inputs - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GREATER, [input_nodes[0].name, input_nodes[1].name], [output.name], diff --git a/backends/arm/operators/op_index_select.py b/backends/arm/operators/op_index_select.py index a42f85abc4c..e357416fadb 100644 --- a/backends/arm/operators/op_index_select.py +++ b/backends/arm/operators/op_index_select.py @@ -7,15 +7,15 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils # noqa: F401 +import executorch.backends.arm.tosa.quant_utils as tqutils # noqa: F401 from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, register_node_visitor, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg -from executorch.backends.arm.tosa_utils import build_reshape_tosa_1_0 +from executorch.backends.arm.tosa.utils import build_reshape_tosa_1_0 from torch.fx import Node @@ -86,7 +86,9 @@ def define_node( tosa_graph, indices.name, indices_new_shape, indices_reshaped.name ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GATHER, [weights_reshaped.name, indices_reshaped.name], [output_name], diff --git a/backends/arm/operators/op_index_tensor.py b/backends/arm/operators/op_index_tensor.py index 7afd7fe6612..2ef7eac352b 100644 --- a/backends/arm/operators/op_index_tensor.py +++ b/backends/arm/operators/op_index_tensor.py @@ -8,7 +8,7 @@ import math from typing import Any, List -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.utils as tutils import numpy as np @@ -19,8 +19,8 @@ from executorch.backends.arm.operators.operator_validation_utils import ( validate_same_dtype, ) -from executorch.backends.arm.tosa_mapping import extract_tensor_meta, TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa.mapping import extract_tensor_meta, TosaArg +from executorch.backends.arm.tosa.specification import TosaSpecification from torch.fx import Node @@ -167,7 +167,9 @@ def define_node( data = np.full(index_shape, int(values_strides[i] / C)) mul_const = tosa_graph.addConst(index_shape, index_dtype, data) tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_{i}_shift") - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MUL, [index_name, mul_const.name, f"{node.name}_{i}_shift"], [stride_shifted_indices.name], @@ -194,7 +196,9 @@ def define_node( reshaped_idxs.shape, reshaped_idxs.dtype, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().ADD, [gather_index_name, reshaped_idxs.name], [add_idxs.name], @@ -217,7 +221,9 @@ def define_node( gather_out_shape, output.dtype, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GATHER, [reshaped_input.name, gather_index_name], [gather_out.name], diff --git a/backends/arm/operators/op_le.py b/backends/arm/operators/op_le.py index 9301f91cb4c..403c6c233d3 100644 --- a/backends/arm/operators/op_le.py +++ b/backends/arm/operators/op_le.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -67,7 +67,9 @@ def define_node( # Update IO input_nodes = rescaled_inputs - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GREATER_EQUAL, [input_nodes[1].name, input_nodes[0].name], [output.name], diff --git a/backends/arm/operators/op_log.py b/backends/arm/operators/op_log.py index 8a48fe4fda5..051b10af062 100644 --- a/backends/arm/operators/op_log.py +++ b/backends/arm/operators/op_log.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -45,4 +45,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().LOG, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().LOG, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_logical_not.py b/backends/arm/operators/op_logical_not.py new file mode 100644 index 00000000000..640c3b4e44f --- /dev/null +++ b/backends/arm/operators/op_logical_not.py @@ -0,0 +1,59 @@ +# Copyright 2025 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, List + +import torch.fx + +from executorch.backends.arm.operators.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.arm.operators.operator_validation_utils import ( + validate_num_inputs, + validate_same_dtype, + validate_valid_dtype, +) +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg + + +@register_node_visitor +class LogicalNotVisitor(NodeVisitor): + target = "aten.logical_not.default" + + tosa_specs = [ + TosaSpecification.create_from_string("TOSA-1.0+INT"), + TosaSpecification.create_from_string("TOSA-1.0+FP"), + ] + + def __init__(self, *args): + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + tosa_graph: Any, + inputs: List[TosaArg], + output: TosaArg, + ) -> None: + import serializer.tosa_serializer as ts # type: ignore # noqa: F401 + + validate_num_inputs(self.target, inputs, 1) + validate_same_dtype(self.target, [*inputs, output], ts) + validate_valid_dtype( + self.target, + [*inputs, output], + [ts.DType.BOOL], + output.tosa_spec, + ) + + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().LOGICAL_NOT, + [inputs[0].name], + [output.name], + ) diff --git a/backends/arm/operators/op_lt.py b/backends/arm/operators/op_lt.py index 31083e93590..f5132dd4feb 100644 --- a/backends/arm/operators/op_lt.py +++ b/backends/arm/operators/op_lt.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -67,7 +67,9 @@ def define_node( # Update IO input_nodes = rescaled_inputs - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().GREATER, [input_nodes[1].name, input_nodes[0].name], [output.name], diff --git a/backends/arm/operators/op_max_pool2d.py b/backends/arm/operators/op_max_pool2d.py index 754fcfcd638..39fcbf5cc64 100644 --- a/backends/arm/operators/op_max_pool2d.py +++ b/backends/arm/operators/op_max_pool2d.py @@ -18,8 +18,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -94,7 +94,9 @@ def define_node( kernel=kernel_size, stride=stride, pad=pad_size_list, nan_mode=1 ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MAX_POOL2D, [input_tensor.name], [output.name], diff --git a/backends/arm/operators/op_maximum.py b/backends/arm/operators/op_maximum.py index 27e5fdc2e02..66437f8af1d 100644 --- a/backends/arm/operators/op_maximum.py +++ b/backends/arm/operators/op_maximum.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( get_input_qparams, @@ -22,9 +22,9 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.utils import tosa_shape from torch.fx import Node @@ -87,7 +87,9 @@ def define_node( # Set to PROPOGATE as default attr_maximum.MaximumAttribute(nan_mode=NanPropagationMode.PROPAGATE) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MAXIMUM, [ operand_inputs[0].name, diff --git a/backends/arm/operators/op_minimum.py b/backends/arm/operators/op_minimum.py index 9dfa7d1f394..518366d5463 100644 --- a/backends/arm/operators/op_minimum.py +++ b/backends/arm/operators/op_minimum.py @@ -7,7 +7,7 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils +import executorch.backends.arm.tosa.quant_utils as tqutils from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( get_input_qparams, @@ -21,9 +21,9 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.utils import tosa_shape from torch.fx import Node @@ -86,7 +86,9 @@ def define_node( # Set to PROPOGATE as default attr_minimum.MinimumAttribute(nan_mode=NanPropagationMode.PROPAGATE) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MINIMUM, [ operand_inputs[0].name, diff --git a/backends/arm/operators/op_mul.py b/backends/arm/operators/op_mul.py index 7d9f6eac6aa..9d139c68242 100644 --- a/backends/arm/operators/op_mul.py +++ b/backends/arm/operators/op_mul.py @@ -7,8 +7,8 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.quant_utils as tqutils +import executorch.backends.arm.tosa.utils as tutils import torch from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( @@ -24,8 +24,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -34,6 +34,7 @@ class MulVisitor_INT(NodeVisitor): tosa_specs = [ TosaSpecification.create_from_string("TOSA-1.0+INT"), + TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), ] def define_node( @@ -51,11 +52,11 @@ def define_node( validate_valid_dtype( self.target, [*inputs, output], - [ts.DType.INT8, ts.DType.INT32], + [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32], output.tosa_spec, ) - if inputs[0].dtype == ts.DType.INT8: + if inputs[0].dtype == ts.DType.INT8 or inputs[0].dtype == ts.DType.INT16: input_A = inputs[0] input_B = inputs[1] input_qparams = get_input_qparams(node) @@ -80,20 +81,22 @@ def define_node( tosa_spec=self.tosa_spec, ) else: - # input[0].dtype == ts.DType.INT32 + # input[0].dtype == ts.DType.INT16 or ts.DType.INT32 # Non quantized input, natively support by TOSA.MUL input_A_rescaled, input_B_rescaled = inputs[0], inputs[1] - if output.dtype == ts.DType.INT8: + if output.dtype == ts.DType.INT8 or output.dtype == ts.DType.INT16: output_shape = tutils.tosa_shape(output.shape, output.dim_order) mul_output = tosa_graph.addIntermediate(output_shape, ts.DType.INT32) else: - # output.dtype == ts.DType.INT32 + # output.dtype == ts.DType.INT32 (non-quantized) mul_output = output # Do the INT32 Mul tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_shift") - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MUL, [input_A_rescaled.name, input_B_rescaled.name, f"{node.name}_shift"], [mul_output.name], @@ -108,6 +111,15 @@ def define_node( tqutils.insert_rescale_op_to_int8( tosa_graph, mul_output, output_scale, node, self.tosa_spec ) + elif output.dtype == ts.DType.INT16: + # Scale output back to 16 bit + output_scale = ( + input_A_qargs.get_scale_per_tensor() # type: ignore[possibly-undefined] + * input_B_qargs.get_scale_per_tensor() # type: ignore[possibly-undefined] + ) + tqutils.insert_rescale_op_to_int16( + tosa_graph, mul_output, output_scale, node, self.tosa_spec + ) @register_node_visitor @@ -135,7 +147,9 @@ def define_node( input1, input2 = inputs tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_shift") - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().MUL, [input1.name, input2.name, f"{node.name}_shift"], [output.name], diff --git a/backends/arm/operators/op_neg.py b/backends/arm/operators/op_neg.py index 54f3dafe769..98aeea14bea 100644 --- a/backends/arm/operators/op_neg.py +++ b/backends/arm/operators/op_neg.py @@ -21,7 +21,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg def get_negate_zero_points(node: torch.fx.Node, is_int8: bool) -> tuple[int, int]: @@ -82,7 +82,9 @@ def define_node( (1,), output.dtype, [output_zp], name=output.name + "_output_zp" ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().NEGATE, [inputs[0].name, input_zp_tensor.name, output_zp_tensor.name], [output.name], diff --git a/backends/arm/operators/op_permute.py b/backends/arm/operators/op_permute.py index 0830d8f4504..92cc2b37479 100644 --- a/backends/arm/operators/op_permute.py +++ b/backends/arm/operators/op_permute.py @@ -18,7 +18,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg def permutation_vector_to_matrix(permutation_vector: list[int]) -> torch.Tensor: @@ -135,6 +135,11 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.TransposeAttribute(permutation_vector) - tosa_graph.addOperator( - ts.TosaOp.Op().TRANSPOSE, [inputs[0].name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().TRANSPOSE, + [inputs[0].name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_pow.py b/backends/arm/operators/op_pow.py index 413160c902a..8e7cffc0770 100644 --- a/backends/arm/operators/op_pow.py +++ b/backends/arm/operators/op_pow.py @@ -16,8 +16,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -50,7 +50,9 @@ def define_node( output.tosa_spec, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().POW, [ inputs[0].name, diff --git a/backends/arm/operators/op_reciprocal.py b/backends/arm/operators/op_reciprocal.py index 3838afd9728..5aa45f740c2 100644 --- a/backends/arm/operators/op_reciprocal.py +++ b/backends/arm/operators/op_reciprocal.py @@ -17,8 +17,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -46,6 +46,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator( - ts.TosaOp.Op().RECIPROCAL, [inputs[0].name], [output.name] + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().RECIPROCAL, [inputs[0].name], [output.name] ) diff --git a/backends/arm/operators/op_repeat.py b/backends/arm/operators/op_repeat.py index 3e636e993b7..5db7ce9347c 100644 --- a/backends/arm/operators/op_repeat.py +++ b/backends/arm/operators/op_repeat.py @@ -17,8 +17,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.utils import tosa_shape @register_node_visitor @@ -60,7 +60,9 @@ def define_node( name=node.name + "_multiples", ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().TILE, [inputs[0].name, multiple_shapes.name], [output.name], diff --git a/backends/arm/operators/op_rescale.py b/backends/arm/operators/op_rescale.py index c9ea96baec5..d7be2be737c 100644 --- a/backends/arm/operators/op_rescale.py +++ b/backends/arm/operators/op_rescale.py @@ -15,16 +15,16 @@ from executorch.backends.arm.operators.operator_validation_utils import ( validate_num_inputs, ) -from executorch.backends.arm.tosa_mapping import map_dtype, TosaArg -from executorch.backends.arm.tosa_quant_utils import build_rescale -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import map_dtype, TosaArg +from executorch.backends.arm.tosa.quant_utils import build_rescale from torch.fx import Node @register_node_visitor -class RescaleVisitor_INT(NodeVisitor): - target = "_rescale.default" +class RescaleVisitor(NodeVisitor): + target = "tosa.RESCALE.default" tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+INT")] diff --git a/backends/arm/operators/op_rshift_tensor.py b/backends/arm/operators/op_rshift_tensor.py index 5313f5c8143..2a41d685f5d 100644 --- a/backends/arm/operators/op_rshift_tensor.py +++ b/backends/arm/operators/op_rshift_tensor.py @@ -18,7 +18,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -53,7 +53,9 @@ def define_node( round = True attr.ArithmeticRightShiftAttribute(round=round) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().ARITHMETIC_RIGHT_SHIFT, [inputs[0].name, inputs[1].name], [output.name], diff --git a/backends/arm/operators/op_rsqrt.py b/backends/arm/operators/op_rsqrt.py index df293946ded..362a30f1cf5 100644 --- a/backends/arm/operators/op_rsqrt.py +++ b/backends/arm/operators/op_rsqrt.py @@ -17,8 +17,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -46,4 +46,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().RSQRT, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().RSQRT, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_sigmoid.py b/backends/arm/operators/op_sigmoid.py index dec42ae15f9..2c4673d6b5f 100644 --- a/backends/arm/operators/op_sigmoid.py +++ b/backends/arm/operators/op_sigmoid.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -45,4 +45,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().SIGMOID, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().SIGMOID, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_sin.py b/backends/arm/operators/op_sin.py index 1ea637f960a..76aee063555 100644 --- a/backends/arm/operators/op_sin.py +++ b/backends/arm/operators/op_sin.py @@ -16,8 +16,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -44,4 +44,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().SIN, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().SIN, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_slice.py b/backends/arm/operators/op_slice.py index 56115073ce1..12d38060aa6 100644 --- a/backends/arm/operators/op_slice.py +++ b/backends/arm/operators/op_slice.py @@ -16,7 +16,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -57,7 +57,7 @@ def define_node( validate_valid_dtype( self.target, [inputs[0], output], - [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32], + [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32], output.tosa_spec, ) @@ -117,7 +117,9 @@ def define_node( (sizes_len,), ts.DType.SHAPE, sizes, node.name + "_sizes_shape" ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().SLICE, [input_node.name, start_tensor.name, sizes_tensor.name], [output.name], diff --git a/backends/arm/operators/op_sub.py b/backends/arm/operators/op_sub.py index 18b3c853271..9c27fddf68a 100644 --- a/backends/arm/operators/op_sub.py +++ b/backends/arm/operators/op_sub.py @@ -7,8 +7,8 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.quant_utils as tqutils +import executorch.backends.arm.tosa.utils as tutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -19,8 +19,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -56,7 +56,7 @@ def define_node( scale_back = 1.0 if inputs[0].dtype == ts.DType.INT8: - rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32( + rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32_maxscale( tosa_graph, inputs, node, self.tosa_spec ) else: @@ -72,7 +72,9 @@ def define_node( sub_output = output # Do the INT32 Sub - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().SUB, [ rescaled_inputs[0].name, @@ -86,7 +88,12 @@ def define_node( # Scale output back to 8 bit # pyre-ignore tqutils.insert_rescale_op_to_int8( - tosa_graph, sub_output, scale_back, node, self.tosa_spec + tosa_graph, + sub_output, + scale_back, + node, + compute_rescale=False, + tosa_spec=self.tosa_spec, ) # type: ignore[possibly-undefined] @@ -122,7 +129,9 @@ def define_node( ) # MI lowering - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().SUB, [inputs[0].name, inputs[1].name], [output.name], diff --git a/backends/arm/operators/op_sum.py b/backends/arm/operators/op_sum.py index 54e848a1bef..0bd152a8b8c 100644 --- a/backends/arm/operators/op_sum.py +++ b/backends/arm/operators/op_sum.py @@ -7,8 +7,8 @@ from typing import Any, List -import executorch.backends.arm.tosa_quant_utils as tqutils -import executorch.backends.arm.tosa_utils as tutils +import executorch.backends.arm.tosa.quant_utils as tqutils +import executorch.backends.arm.tosa.utils as tutils from executorch.backends.arm.operators.node_visitor import ( NodeVisitor, @@ -18,8 +18,8 @@ validate_num_inputs, validate_same_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node diff --git a/backends/arm/operators/op_table.py b/backends/arm/operators/op_table.py index 557281f4d2a..41b40268f6d 100644 --- a/backends/arm/operators/op_table.py +++ b/backends/arm/operators/op_table.py @@ -16,14 +16,14 @@ validate_num_inputs, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor class TableVisitor(NodeVisitor): - target = "_table.default" + target = "tosa.TABLE.default" tosa_specs = [TosaSpecification.create_from_string("TOSA-1.0+INT")] @@ -36,7 +36,7 @@ def define_node( ) -> None: import serializer.tosa_serializer as ts # type: ignore - validate_num_inputs(self.target, inputs, 1) + validate_num_inputs(self.target, inputs, 2) validate_valid_dtype( self.target, inputs, [ts.DType.INT8, ts.DType.INT16], output.tosa_spec ) @@ -45,12 +45,12 @@ def define_node( if inputs[0].dtype == ts.DType.INT16: validate_valid_dtype(self.target, output, ts.DType.INT32, output.tosa_spec) - if node.name not in self._exported_program.state_dict.keys(): # type: ignore[union-attr] + if inputs[1].name not in self._exported_program.state_dict.keys(): # type: ignore[union-attr] raise RuntimeError( f"Did not find key {node.name} in state_dict {self._exported_program.state_dict.keys()}." ) - table = self._exported_program.state_dict[node.name] + table = self._exported_program.state_dict[inputs[1].name] # type: ignore[union-attr] table_tensor_name = node.name + "_table" tosa_graph.addConst( @@ -60,7 +60,9 @@ def define_node( name=table_tensor_name, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().TABLE, [inputs[0].name, table_tensor_name], [output.name], diff --git a/backends/arm/operators/op_tanh.py b/backends/arm/operators/op_tanh.py index 0d149397eb6..5837825a6a1 100644 --- a/backends/arm/operators/op_tanh.py +++ b/backends/arm/operators/op_tanh.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -46,4 +46,6 @@ def define_node( self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec ) - tosa_graph.addOperator(ts.TosaOp.Op().TANH, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().TANH, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_to_copy.py b/backends/arm/operators/op_to_copy.py deleted file mode 100644 index 9758a018b87..00000000000 --- a/backends/arm/operators/op_to_copy.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright 2024-2025 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-unsafe -from typing import Any, List - -import torch - -from executorch.backends.arm.operators.node_visitor import ( - NodeVisitor, - register_node_visitor, -) -from executorch.backends.arm.operators.operator_validation_utils import ( - validate_num_inputs, -) -from executorch.backends.arm.tosa_mapping import TosaArg - - -@register_node_visitor -class ToCopyVisitor(NodeVisitor): - """ - Implement the type cast functionality of _to_copy. - - Other features like setting of the memory_format or moving a tensor to a - different device are not supported. - - Also note that the node should not be quantized. - """ - - target = "aten._to_copy.default" - - tosa_specs = NodeVisitor.tosa_specs - - def define_node( - self, - node: torch.fx.Node, - tosa_graph: Any, - inputs: List[TosaArg], - output: TosaArg, - ) -> None: - import serializer.tosa_serializer as ts # type: ignore - - validate_num_inputs(self.target, inputs, 1) - - tosa_graph.addOperator(ts.TosaOp.Op().CAST, [inputs[0].name], [output.name]) diff --git a/backends/arm/operators/op_to_dim_order_copy.py b/backends/arm/operators/op_to_dim_order_copy.py index 74bf1a5ad14..aa5873f698e 100644 --- a/backends/arm/operators/op_to_dim_order_copy.py +++ b/backends/arm/operators/op_to_dim_order_copy.py @@ -15,7 +15,7 @@ from executorch.backends.arm.operators.operator_validation_utils import ( validate_num_inputs, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor @@ -44,4 +44,6 @@ def define_node( validate_num_inputs(self.target, inputs, 1) - tosa_graph.addOperator(ts.TosaOp.Op().CAST, [inputs[0].name], [output.name]) + self._serialize_operator( + node, tosa_graph, ts.TosaOp.Op().CAST, [inputs[0].name], [output.name] + ) diff --git a/backends/arm/operators/op_transpose.py b/backends/arm/operators/op_transpose.py index 0845c3ed61c..48766687a62 100644 --- a/backends/arm/operators/op_transpose.py +++ b/backends/arm/operators/op_transpose.py @@ -18,18 +18,18 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg @register_node_visitor class TransposeVisitor(NodeVisitor): """ - This node visitor targets the _transpose op defined in the - passthrough_to_tosa library. Used when switching between tosa_dim_orders. + This node visitor targets the tosa::TRANSPOSE op defined in the + TOSA backend dialect. Used when switching between tosa_dim_orders. Inserts a TOSA TRANSPOSE. """ - target = "_transpose.default" + target = "tosa.TRANSPOSE.default" tosa_specs = NodeVisitor.tosa_specs @@ -47,7 +47,14 @@ def define_node( validate_valid_dtype( self.target, [inputs[0], output], - [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32], + [ + ts.DType.INT8, + ts.DType.INT16, + ts.DType.INT32, + ts.DType.FP32, + ts.DType.BOOL, + ts.DType.FP16, + ], output.tosa_spec, ) @@ -55,6 +62,11 @@ def define_node( perms = [dim % output_rank for dim in inputs[1].special] attr = ts.TosaSerializerAttribute() attr.TransposeAttribute(perms) - tosa_graph.addOperator( - ts.TosaOp.Op().TRANSPOSE, [inputs[0].name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().TRANSPOSE, + [inputs[0].name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_upsample_bilinear2d.py b/backends/arm/operators/op_upsample_bilinear2d.py index 26927bfcfa2..3cc620727e0 100644 --- a/backends/arm/operators/op_upsample_bilinear2d.py +++ b/backends/arm/operators/op_upsample_bilinear2d.py @@ -17,9 +17,9 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_quant_utils import build_rescale -from executorch.backends.arm.tosa_utils import get_resize_parameters, tosa_shape +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.quant_utils import build_rescale +from executorch.backends.arm.tosa.utils import get_resize_parameters, tosa_shape @register_node_visitor @@ -98,7 +98,9 @@ def in_int16_range(x): [len(border)], ts.DType.SHAPE, border, node.name + "_border" ) if input_dtype == output.dtype == ts.DType.FP32: - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().RESIZE, [ inputs[0].name, @@ -114,7 +116,9 @@ def in_int16_range(x): intermediate = tosa_graph.addIntermediate( tosa_shape(output.shape, output.dim_order), ts.DType.INT32 ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().RESIZE, [ inputs[0].name, diff --git a/backends/arm/operators/op_upsample_nearest2d.py b/backends/arm/operators/op_upsample_nearest2d.py index 46dcc0605e6..3c3ca67c9f5 100644 --- a/backends/arm/operators/op_upsample_nearest2d.py +++ b/backends/arm/operators/op_upsample_nearest2d.py @@ -17,8 +17,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_utils import get_resize_parameters +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.utils import get_resize_parameters from tosa.ResizeMode import ResizeMode # type: ignore @@ -89,7 +89,9 @@ def in_int16_range(x): mode=ResizeMode.NEAREST, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().RESIZE, [ inputs[0].name, diff --git a/backends/arm/operators/op_view.py b/backends/arm/operators/op_view.py index 1e8c06b691f..925404da917 100644 --- a/backends/arm/operators/op_view.py +++ b/backends/arm/operators/op_view.py @@ -17,8 +17,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_utils import tosa_shape +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.utils import tosa_shape @register_node_visitor @@ -44,7 +44,13 @@ def define_node( validate_valid_dtype( self.target, [inputs[0], output], - [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32, ts.DType.BOOL], + [ + ts.DType.INT8, + ts.DType.INT16, + ts.DType.INT32, + ts.DType.FP32, + ts.DType.BOOL, + ], output.tosa_spec, ) @@ -66,6 +72,11 @@ def define_node( attr = ts.TosaSerializerAttribute() attr.ReshapeAttribute() - tosa_graph.addOperator( - ts.TosaOp.Op().RESHAPE, [inputs[0].name, shape.name], [output.name], attr + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().RESHAPE, + [inputs[0].name, shape.name], + [output.name], + attr, ) diff --git a/backends/arm/operators/op_where.py b/backends/arm/operators/op_where.py index e6a87be6387..dbdbbc67944 100644 --- a/backends/arm/operators/op_where.py +++ b/backends/arm/operators/op_where.py @@ -15,8 +15,8 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import TosaArg from torch.fx import Node @@ -33,6 +33,7 @@ def __init__(self, *args): def _add_node_to_tosa_graph( self, + node: Node, tosa_graph: Any, inputs: List[TosaArg], output: TosaArg, @@ -51,7 +52,9 @@ def _add_node_to_tosa_graph( output.tosa_spec, ) - tosa_graph.addOperator( + self._serialize_operator( + node, + tosa_graph, ts.TosaOp.Op().SELECT, [inputs[0].name, inputs[1].name, inputs[2].name], [output.name], @@ -73,7 +76,9 @@ def define_node( ts.DType.INT32, ts.DType.BOOL, ] - self._add_node_to_tosa_graph(tosa_graph, inputs, output, bi_supported_dtypes) + self._add_node_to_tosa_graph( + node, tosa_graph, inputs, output, bi_supported_dtypes + ) @register_node_visitor @@ -103,4 +108,6 @@ def define_node( ts.DType.INT32, ts.DType.BOOL, ] - self._add_node_to_tosa_graph(tosa_graph, inputs, output, mi_supported_dtypes) + self._add_node_to_tosa_graph( + node, tosa_graph, inputs, output, mi_supported_dtypes + ) diff --git a/backends/arm/operators/ops_binary.py b/backends/arm/operators/ops_binary.py index dc9bd446a34..4e8e393732b 100644 --- a/backends/arm/operators/ops_binary.py +++ b/backends/arm/operators/ops_binary.py @@ -19,7 +19,7 @@ validate_same_dtype, validate_valid_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg def binary_operator_factory(bw_target: str, tosa_op): @@ -65,8 +65,12 @@ def define_node( output.tosa_spec, ) - tosa_graph.addOperator( - tosa_op, [inputs[0].name, inputs[1].name], [output.name] + self._serialize_operator( + node, + tosa_graph, + tosa_op, + [inputs[0].name, inputs[1].name], + [output.name], ) register_node_visitor(BinaryOperator) diff --git a/backends/arm/operators/ops_identity.py b/backends/arm/operators/ops_identity.py index 238b033f8eb..62a307f3012 100644 --- a/backends/arm/operators/ops_identity.py +++ b/backends/arm/operators/ops_identity.py @@ -18,7 +18,7 @@ validate_num_inputs, validate_same_dtype, ) -from executorch.backends.arm.tosa_mapping import TosaArg +from executorch.backends.arm.tosa.mapping import TosaArg def identity_operator_factory(identity_target: str): @@ -45,8 +45,12 @@ def define_node( validate_same_dtype(self.target, [*inputs, output], ts) # Simply add an identityOp - tosa_graph.addOperator( - ts.TosaOp.Op().IDENTITY, [inputs[0].name], [output.name] + self._serialize_operator( + node, + tosa_graph, + ts.TosaOp.Op().IDENTITY, + [inputs[0].name], + [output.name], ) register_node_visitor(IdentityOperatorVisitor) diff --git a/backends/arm/operators/ops_unary.py b/backends/arm/operators/ops_unary.py deleted file mode 100644 index 48092e13968..00000000000 --- a/backends/arm/operators/ops_unary.py +++ /dev/null @@ -1,66 +0,0 @@ -# Copyright 2025 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-unsafe -from typing import Any, List - -import torch.fx - -from executorch.backends.arm.operators.node_visitor import ( - NodeVisitor, - register_node_visitor, -) -from executorch.backends.arm.operators.operator_validation_utils import ( - validate_num_inputs, - validate_same_dtype, - validate_valid_dtype, -) - -from executorch.backends.arm.tosa_mapping import TosaArg - - -def unary_operator_factory(unary_target: str, tosa_op): - "Creates and registers NodeVisitors for operations that have one input and map directly into a TOSA op." - - # Some TOSA unary operators only support float - fp_only_ops = ["aten.floor.default"] - - class UnaryOperator(NodeVisitor): - target = unary_target - tosa_specs = NodeVisitor.tosa_specs - - def __init__(self, *args): - super().__init__(*args) - - def define_node( - self, - node: torch.fx.Node, - tosa_graph: Any, - inputs: List[TosaArg], - output: TosaArg, - ) -> None: - import serializer.tosa_serializer as ts # type: ignore # noqa: F401 - - validate_num_inputs(self.target, inputs, 1) - validate_same_dtype(self.target, [*inputs, output], ts) - - if self.target in fp_only_ops: - validate_valid_dtype( - self.target, - inputs[0], - ts.DType.FP32, - output.tosa_spec, - ) - - tosa_graph.addOperator(tosa_op, [inputs[0].name], [output.name]) - - register_node_visitor(UnaryOperator) - - -import serializer.tosa_serializer as ts # type: ignore - -unary_operator_factory("aten.ceil.default", ts.TosaOp.Op().CEIL) -unary_operator_factory("aten.floor.default", ts.TosaOp.Op().FLOOR) -unary_operator_factory("aten.logical_not.default", ts.TosaOp.Op().LOGICAL_NOT) diff --git a/backends/arm/process_node.py b/backends/arm/process_node.py index edbd2ca2a29..9ca435c60c5 100644 --- a/backends/arm/process_node.py +++ b/backends/arm/process_node.py @@ -8,12 +8,13 @@ from typing import Any, cast, Dict import numpy as np +import serializer.tosa_serializer as ts import torch import torch.fx from executorch.backends.arm.operators.node_visitor import NodeVisitor -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification -from executorch.backends.arm.tosa_utils import getNodeArgs, tosa_shape +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.arm.tosa.utils import tosa_shape from torch._export.utils import ( get_buffer, get_lifted_tensor_constant, @@ -32,7 +33,10 @@ def process_call_function( tosa_spec: TosaSpecification, ): # Unpack arguments and convert - inputs = getNodeArgs(node, tosa_spec) + try: + inputs = [TosaArg(arg, tosa_spec) for arg in node.args] + except ValueError as e: + raise ValueError(f"Failed processing args to op:\n{node}") from e # Convert output (this node itself) try: @@ -81,11 +85,6 @@ def process_inputs( "Is the original torch function supported?" ) from e - if isinstance(tosa_spec, Tosa_1_00): - import serializer.tosa_serializer as ts - else: - raise ValueError(f"Unsupported TOSA spec: {tosa_spec}") - input_shape = tosa_arg.shape input_dim_order = tosa_arg.dim_order tensor = ts.TosaSerializerTensor( diff --git a/backends/arm/quantizer/TARGETS b/backends/arm/quantizer/TARGETS index 7ee44b3751a..1a02340f92b 100644 --- a/backends/arm/quantizer/TARGETS +++ b/backends/arm/quantizer/TARGETS @@ -1,7 +1,7 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") # Exposed through __init__.py -python_library( +runtime.python_library( name = "quantization_config", srcs = ["quantization_config.py"], deps = [ @@ -11,7 +11,7 @@ python_library( ) # Exposed through __init__.py -python_library( +runtime.python_library( name = "arm_quantizer", srcs = ["arm_quantizer.py"], deps = [ @@ -23,7 +23,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "quantization_annotator", srcs = ["quantization_annotator.py"], deps = [ @@ -34,7 +34,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "arm_quantizer_utils", srcs = ["arm_quantizer_utils.py"], deps = [ @@ -43,7 +43,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "lib", srcs = ["__init__.py"], deps = [ diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index 734ddec4359..ae7c8255428 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -14,18 +14,17 @@ from __future__ import annotations import functools -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Union import torch -from executorch.backends.arm._passes import ArmPassManager from executorch.backends.arm.quantizer import QuantizationConfig -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.specification import get_tosa_spec from .arm_quantizer_utils import is_annotated, mark_node_as_annotated from .quantization_annotator import annotate_graph from executorch.backends.arm.arm_backend import ( - get_tosa_spec, is_ethosu, is_vgf, ) # usort: skip @@ -102,18 +101,20 @@ def get_symmetric_quantization_config( weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = ( MinMaxObserver ) + # Determine the right observer/fake-quant constructor if is_qat: - # Set plain fake-quant with true min/max - weight_observer_or_fake_quant_ctr = FakeQuantize + if is_per_channel: + weight_observer_or_fake_quant_ctr = PerChannelMinMaxObserver + else: + # Set plain fake-quant with true min/max + weight_observer_or_fake_quant_ctr = FakeQuantize else: # PTQ: set min/max observer weight_observer_or_fake_quant_ctr = ( PerChannelMinMaxObserver if is_per_channel else MinMaxObserver ) - extra_args = {"eps": 2**-12} - weight_quantization_spec = QuantizationSpec( dtype=torch.int8, quant_min=weight_qmin, @@ -144,6 +145,86 @@ def get_symmetric_quantization_config( return quantization_config +@functools.lru_cache +def 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, +): + """ + 16A8W quantization config: 16-bit activations, 8-bit weights. + + This configuration provides better accuracy than 8A8W while maintaining + reasonable memory usage through 8-bit weights. + + Args: + is_per_channel: Whether to use per-channel quantization for weights + is_qat: Whether this is for Quantization Aware Training + is_dynamic: Whether to use dynamic quantization + weight_qmin: Minimum quantization value for weights + weight_qmax: Maximum quantization value for weights + + Returns: + QuantizationConfig with 16-bit activations and 8-bit weights + """ + extra_args: Dict[str, Any] = {"eps": 2**-12} + + # Setup observer/fake-quant for 16-bit activations + if is_qat: + if is_dynamic: + act_observer_or_fake_quant_ctr = FakeQuantize + dynamic_quant_observer = MovingAverageMinMaxObserver.with_args( + averaging_constant=1 + ) + extra_args["observer"] = dynamic_quant_observer + else: + act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize # type: ignore[assignment] + else: + if is_dynamic: + act_observer_or_fake_quant_ctr = PlaceholderObserver # type: ignore[assignment] + else: + # HistogramObserver works well for 16-bit range + act_observer_or_fake_quant_ctr = HistogramObserver # type: ignore[assignment] + + # 16-bit activation quantization spec + act_quantization_spec = QuantizationSpec( + dtype=torch.int16, + quant_min=torch.iinfo(torch.int16).min, # -32768 + quant_max=torch.iinfo(torch.int16).max, # 32767 + qscheme=torch.per_tensor_symmetric, + is_dynamic=is_dynamic, + observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args( + **extra_args, + ), + ) + + # Instead of reconstructing quantization_config, just clone and update as needed + # Clone the quantization_config from get_symmetric_quantization_config and update activation spec + base_config = get_symmetric_quantization_config( + is_per_channel=is_per_channel, + is_qat=is_qat, + is_dynamic=is_dynamic, + ) + # Replace activation quantization spec with 16-bit version + if is_dynamic: + quantization_config = QuantizationConfig( + act_quantization_spec, # 16-bit input activations + None, + base_config.weight, # 8-bit weights from base config + None, + ) + else: + quantization_config = QuantizationConfig( + act_quantization_spec, # 16-bit input activations + act_quantization_spec, # 16-bit output activations + base_config.weight, # 8-bit weights from base config + None, + ) + return quantization_config + + NodeFilterType = Callable[[Node], bool] """Type for a Node Filter used by annotators. A Node filter is a function that takes a Node and returns whether the node should be annotated or not. @@ -218,9 +299,35 @@ def not_module_type_or_name_filter(n: Node) -> bool: class TOSAQuantizer(Quantizer): - def __init__(self, tosa_spec: TosaSpecification) -> None: + def __init__( + self, compile_spec_or_tosa_spec: Union[TosaSpecification, List[CompileSpec]] + ) -> None: + super().__init__() - self.tosa_spec = tosa_spec + if isinstance(compile_spec_or_tosa_spec, TosaSpecification): + self.tosa_spec = compile_spec_or_tosa_spec + self.compile_spec = None + elif isinstance(compile_spec_or_tosa_spec, list): + self.compile_spec = compile_spec_or_tosa_spec + # find entry that is 'tosa_spec' + for cs in compile_spec_or_tosa_spec: + if cs.key == "tosa_spec": + spec_val = ( + cs.value.decode() if isinstance(cs.value, bytes) else cs.value + ) + self.tosa_spec = TosaSpecification.create_from_string(spec_val) + break + else: + raise ValueError( + "compile_spec list did not contain a 'tosa_spec' entry" + ) + else: + raise TypeError( + f"TOSAQuantizer constructor expects " + f"a TosaSpecification or compile_spec list, " + f"got {type(compile_spec_or_tosa_spec)}" + ) + self.global_config: Optional[QuantizationConfig] = None self.io_config: Optional[QuantizationConfig] = None self.module_type_config: Dict[Callable, Optional[QuantizationConfig]] = {} @@ -264,6 +371,9 @@ def transform_for_annotation(self, model: GraphModule) -> GraphModule: Currently transforms scalar values to tensor attributes. """ + # TODO: Fix the need to lazily import this. + from executorch.backends.arm._passes import ArmPassManager + return ArmPassManager(self.tosa_spec).transform_for_annotation_pipeline( # type: ignore[arg-type] graph_module=model ) diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index fbc1f8f7cb8..ff1ad50e517 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -11,8 +11,8 @@ import torch import torch.fx import torch.nn.functional as F +from executorch.backends.arm.common.debug import get_node_debug_info from executorch.backends.arm.quantizer import QuantizationConfig -from executorch.backends.arm.tosa_utils import get_node_debug_info from torch._subclasses import FakeTensor from torch.fx import Node @@ -265,6 +265,8 @@ def _match_pattern( torch.ops.aten.ceil.default, torch.ops.aten.erf.default, torch.ops.aten.exp.default, + torch.ops.aten.expm1.default, + torch.ops.aten.elu.default, torch.ops.aten.floor.default, torch.ops.aten.log.default, torch.ops.aten.reciprocal.default, @@ -287,6 +289,9 @@ def _match_pattern( torch.ops.aten.asin.default, torch.ops.aten.atanh.default, torch.ops.aten.asinh.default, + torch.ops.aten.cosh.default, + torch.ops.aten.acos.default, + torch.ops.aten.cumsum.default, ] _one_to_one_shared_input_qspec = [ @@ -335,6 +340,10 @@ def _match_pattern( torch.ops.aten.unflatten.int, torch.ops.aten.index_select.default, torch.ops.aten.index.Tensor, + # Neg operator flips the range, but keps the magnitude the same. + # That is why we force it to use the same qparams and avoid + # dequant -> neg -> requant chain. + torch.ops.aten.neg.default, ] _one_to_one_shared_input_or_input_act_qspec = [ @@ -464,6 +473,10 @@ def any_or_hardtanh_min_zero(n: Node): ] quant_properties.quant_output = _QuantProperty(0, output_act_qspec) elif node.target in ( + torch.ops.aten.add.Tensor, + torch.ops.aten.add_.Tensor, + torch.ops.aten.sub.Tensor, + torch.ops.aten.sub_.Tensor, torch.ops.aten.matmul.default, torch.ops.aten.mm.default, torch.ops.aten.bmm.default, @@ -476,10 +489,6 @@ def any_or_hardtanh_min_zero(n: Node): ] quant_properties.quant_output = _QuantProperty(0, output_act_qspec) elif node.target in ( - torch.ops.aten.add.Tensor, - torch.ops.aten.add_.Tensor, - torch.ops.aten.sub.Tensor, - torch.ops.aten.sub_.Tensor, torch.ops.aten.minimum.default, torch.ops.aten.maximum.default, ): @@ -536,9 +545,6 @@ def any_or_hardtanh_min_zero(n: Node): ) ] quant_properties.quant_output = _QuantProperty(0, shared_qspec) # type: ignore[arg-type] - elif node.target in (torch.ops.aten.neg.default,): - quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)] - quant_properties.quant_output = _QuantProperty(0, input_act_qspec) elif node.target in _one_to_one: quant_properties.quant_inputs = [_QuantProperty(0, input_act_qspec)] quant_properties.quant_output = _QuantProperty(0, output_act_qspec) diff --git a/backends/arm/requirements-arm-ethos-u.txt b/backends/arm/requirements-arm-ethos-u.txt new file mode 100644 index 00000000000..5fad9d2fe94 --- /dev/null +++ b/backends/arm/requirements-arm-ethos-u.txt @@ -0,0 +1,6 @@ +# Copyright 2025 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. + +ethos-u-vela @ git+https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela@d37febc1715edf0d236c2ff555739a8a9aadcf9a diff --git a/backends/arm/requirements-arm-models-test.txt b/backends/arm/requirements-arm-models-test.txt new file mode 100644 index 00000000000..ac4e1d9bad7 --- /dev/null +++ b/backends/arm/requirements-arm-models-test.txt @@ -0,0 +1,6 @@ +# Copyright 2025 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. + +diffusers[torch] == 0.33.1 diff --git a/backends/arm/requirements-arm-tosa.txt b/backends/arm/requirements-arm-tosa.txt new file mode 100644 index 00000000000..4b7a3ec0273 --- /dev/null +++ b/backends/arm/requirements-arm-tosa.txt @@ -0,0 +1,9 @@ +# Copyright 2025 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. + +ml_dtypes == 0.5.1 +flatbuffers == 24.3.25 + +tosa-tools @ git+https://git.gitlab.arm.com/tosa/tosa-reference-model.git@v2025.07.0 diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index d29c32b02f3..bff5ff69284 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -70,6 +70,7 @@ using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::MemoryAllocator; using executorch::runtime::Result; +using executorch::runtime::Span; #define ETHOSU_NUM_BASE_ADDRS 3 @@ -140,7 +141,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { Error execute( BackendExecutionContext& context, DelegateHandle* input_handle, - EValue** args) const override { + Span args) const override { #if defined(ET_EVENT_TRACER_ENABLED) EventTracer* event_tracer = context.event_tracer(); EventTracerEntry event_tracer_local_scope; @@ -191,8 +192,9 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { // Use a temporary allocator for the intermediate tensors of the // computation. The allocator is released in runtime/executor/method.cpp at // the end of the execution of the Ethos-U custom delegate - char* ethosu_scratch = - static_cast(temp_allocator->allocate(handles.scratch_data_size)); + // Ethos-U driver requires 16 bit alignment. + char* ethosu_scratch = static_cast( + temp_allocator->allocate(handles.scratch_data_size, 16UL)); if (ethosu_scratch == nullptr) { ET_LOG( Error, @@ -259,9 +261,6 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { // Select a compatible copy routine including checking for input layouts // which require permutation. - bool permuted_input_shape; - ET_CHECK_OK_OR_RETURN_ERROR(check_requires_permute( - i, tensor_in, &handles.inputs->io[i], &permuted_input_shape)); bool both_int = tensor_in.scalar_type() == ScalarType::Int && handles.inputs->io[i].elem_size == 4; bool both_char = tensor_in.scalar_type() == ScalarType::Char && @@ -271,19 +270,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { bool both_bool = tensor_in.scalar_type() == ScalarType::Bool && (handles.inputs->io[i].elem_size == 1); - // Select a compatible copy routine - if ((both_char || both_bool) && permuted_input_shape) { - EXECUTORCH_PROF_SCOPE( - event_tracer, - "+EthosUBackend::execute()handles.input.permute_CHW_to_HWC()"); - // permuted byte copy CHW to HWC - permute_CHW_to_HWC( - tensor_in.mutable_data_ptr(), - scratch_addr, - tensor_in.size(1), - tensor_in.size(2), - tensor_in.size(3)); - } else if (both_char || both_int || both_short || both_bool) { + if (both_char || both_int || both_short || both_bool) { EXECUTORCH_PROF_SCOPE( event_tracer, "+EthosUBackend::execute()handles.input.memcpy()"); // Sizes match and elt size matches so memcpy @@ -295,18 +282,16 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { ET_LOG(Error, "No matching input copy routine"); return Error::InvalidProgram; } - if (!permuted_input_shape) { - calculate_dimensions( - tensor_in, &handles.inputs->io[i], &tensor_count, &io_count); - if (tensor_count != io_count) { - ET_LOG(Error, "Input tensor sizes do not match"); - ET_LOG( - Error, - "Program expects %d elements but got %d", - io_count, - tensor_count); - return Error::InvalidProgram; - } + calculate_dimensions( + tensor_in, &handles.inputs->io[i], &tensor_count, &io_count); + if (tensor_count != io_count) { + ET_LOG(Error, "Input tensor sizes do not match"); + ET_LOG( + Error, + "Program expects %d elements but got %d", + io_count, + tensor_count); + return Error::InvalidProgram; } } @@ -367,34 +352,13 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { tensor_dim = tensor_dim + tensor_count; io_dim = io_dim + io_count; - bool permuted_output_shape; - ET_CHECK_OK_OR_RETURN_ERROR(check_requires_permute( - i, tensor_out, &handles.outputs->io[i], &permuted_output_shape)); + EXECUTORCH_PROF_SCOPE( + event_tracer, "+EthosUBackend::execute()handles.output.memcpy()"); - if ((tensor_out.scalar_type() == ScalarType::Char || - tensor_out.scalar_type() == ScalarType::Bool) && - permuted_output_shape) { - EXECUTORCH_PROF_SCOPE( - event_tracer, - "+EthosUBackend::execute()handles.output.permute_HWC_to_CHW()"); - - const char* output_address = static_cast(output_addr); - - permute_HWC_to_CHW( - output_address, - tensor_out.mutable_data_ptr(), - tensor_out.size(1), - tensor_out.size(2), - tensor_out.size(3)); - } else { - EXECUTORCH_PROF_SCOPE( - event_tracer, "+EthosUBackend::execute()handles.output.memcpy()"); - - memcpy( - tensor_out.mutable_data_ptr(), - static_cast(output_addr), - tensor_out.nbytes()); - } + memcpy( + tensor_out.mutable_data_ptr(), + static_cast(output_addr), + tensor_out.nbytes()); } if (tensor_dim != io_dim) { ET_LOG(Error, "Total output tensor sizes do not match"); @@ -424,46 +388,6 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { *io_count = *io_count * io->shape[i]; } } - - Error check_requires_permute( - int index, - const executorch::aten::Tensor tensor, - VelaIO* io, - bool* is_permuted) const { - bool permuted_shape = false; - - if (tensor.dim() == 4) { - // special case for NHWC workaround in AOT; as the compilation has - // permuted to channel last in an undetectable way, we assume here - // that the application has similarly permuted any input/output tensors. - permuted_shape = tensor.size(0) == io->shape[0] && - tensor.size(1) == io->shape[3] && tensor.size(2) == io->shape[1] && - tensor.size(3) == io->shape[2]; - if (permuted_shape) { - ET_LOG(Debug, "Tensor input/output %d will be permuted", index); - } - } - *is_permuted = permuted_shape; - return Error::Ok; - } - - void permute_CHW_to_HWC(const char* input, char* output, int C, int H, int W) - const { - for (int i = 0; i != H * W; ++i) { - for (int j = 0; j < C; ++j) { - output[i * C + j] = input[i + j * W * H]; - } - } - } - - void permute_HWC_to_CHW(const char* input, char* output, int C, int H, int W) - const { - for (int i = 0; i != H * W; ++i) { - for (int j = 0; j < C; ++j) { - output[i + j * W * H] = input[i * C + j]; - } - } - } }; namespace { @@ -474,4 +398,4 @@ static auto registered = register_backend(backend_id); } // namespace arm } // namespace backends -} // namespace executorch +} // namespace executorch \ No newline at end of file diff --git a/backends/arm/runtime/VGFBackend.cpp b/backends/arm/runtime/VGFBackend.cpp index ea4f4286eb9..ce182055b6d 100644 --- a/backends/arm/runtime/VGFBackend.cpp +++ b/backends/arm/runtime/VGFBackend.cpp @@ -25,6 +25,7 @@ using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::MemoryAllocator; using executorch::runtime::Result; +using executorch::runtime::Span; // We use the platform and runtime environment provided by the Vulkan delegate #include @@ -152,7 +153,7 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface { Error execute( ET_UNUSED BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const override { + Span args) const override { VgfRepr* repr = static_cast(handle); // Copy all inputs from EValue to VkDeviceMemory @@ -248,7 +249,6 @@ VkResult vkml_allocate_basics( VkDevice* device, VkQueue* queue, VkCommandPool* command_pool) { - const char* dev_exts[] = {"VK_ARM_tensors", "VK_ARM_data_graph"}; VkResult result; if (VK_SUCCESS != volkInitialize()) { @@ -264,15 +264,60 @@ VkResult vkml_allocate_basics( .engineVersion = 0, .apiVersion = VK_API_VERSION_1_3, }; + + std::vector requested_extensions; + VkInstanceCreateFlags instance_flags = 0; + +#ifdef __APPLE__ + instance_flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR; + + uint32_t extension_count = 0; + result = vkEnumerateInstanceExtensionProperties( + nullptr, &extension_count, nullptr); + + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to enumerate instance extensions"); + return result; + } + + std::vector extension_properties(extension_count); + result = vkEnumerateInstanceExtensionProperties( + nullptr, &extension_count, extension_properties.data()); + + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to enumerate instance extensions"); + return result; + } + + if (std::any_of( + extension_properties.begin(), + extension_properties.end(), + [](const auto& extension) { + return strcmp( + VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME, + extension.extensionName) == 0; + })) { + requested_extensions.push_back( + VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME); + } + + if (requested_extensions.empty()) { + ET_LOG(Error, "VK_KHR_portability_enumeration not found"); + } + +#endif + VkInstanceCreateInfo instance_info{ .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO, .pNext = nullptr, - .flags = 0, + .flags = instance_flags, .pApplicationInfo = &app_info, - 0, - nullptr, - 0, - nullptr}; + .enabledLayerCount = 0, + .ppEnabledLayerNames = nullptr, + .enabledExtensionCount = + static_cast(requested_extensions.size()), + .ppEnabledExtensionNames = requested_extensions.data(), + }; result = vkCreateInstance(&instance_info, nullptr, instance); if (result != VK_SUCCESS) { ET_LOG(Error, "Failed to create VkInstance"); @@ -326,11 +371,103 @@ VkResult vkml_allocate_basics( .pQueuePriorities = &qp, }; + // Query features + VkPhysicalDeviceVulkan12Features available_12 = { + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, + .pNext = NULL, + }; + VkPhysicalDeviceVulkan11Features available_11 = { + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, + .pNext = &available_12, + }; + VkPhysicalDeviceFeatures2 available_2 = { + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, + .pNext = &available_11, + }; + vkGetPhysicalDeviceFeatures2(*physical_device, &available_2); + + // Select features + VkPhysicalDeviceShaderReplicatedCompositesFeaturesEXT features_c{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_REPLICATED_COMPOSITES_FEATURES_EXT, + nullptr}; + features_c.shaderReplicatedComposites = true; + VkPhysicalDeviceVulkan13Features features_13{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, nullptr}; + features_13.synchronization2 = true; + features_13.maintenance4 = true; + features_13.pipelineCreationCacheControl = true; + features_13.pNext = &features_c; + VkPhysicalDeviceVulkan12Features features_12{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, nullptr}; + features_12.hostQueryReset = true; + features_12.storageBuffer8BitAccess = true; + features_12.uniformAndStorageBuffer8BitAccess = + available_12.uniformAndStorageBuffer8BitAccess; + features_12.shaderInt8 = true; + features_12.shaderFloat16 = available_12.shaderFloat16; + features_12.vulkanMemoryModel = true; + features_12.vulkanMemoryModelDeviceScope = + available_12.vulkanMemoryModelDeviceScope; + features_12.pNext = &features_13; + VkPhysicalDeviceVulkan11Features features_11{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, nullptr}; + features_11.storageBuffer16BitAccess = available_11.storageBuffer16BitAccess; + features_11.uniformAndStorageBuffer16BitAccess = + available_11.uniformAndStorageBuffer16BitAccess; + features_11.pNext = &features_12; + VkPhysicalDeviceTensorFeaturesARM features_tensor{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TENSOR_FEATURES_ARM, nullptr}; + 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; + + VkPhysicalDeviceFeatures device_features = {}; + device_features.shaderInt16 = VK_TRUE; + device_features.shaderInt64 = VK_TRUE; + + // Extension strings to enable + auto dev_exts = { + "VK_ARM_tensors", + "VK_ARM_data_graph", + "VK_KHR_maintenance4", + "VK_KHR_maintenance5", + "VK_KHR_deferred_host_operations", + "VK_EXT_shader_replicated_composites"}; + + uint32_t exts = 0; + vkEnumerateDeviceExtensionProperties( + *physical_device, nullptr, &exts, nullptr); + vector available(exts); + vkEnumerateDeviceExtensionProperties( + *physical_device, nullptr, &exts, available.data()); + + vector requested_exts; + for (auto& ext : dev_exts) { + bool found = false; + for (auto const& ext_avail : available) { + if (strcmp(ext, ext_avail.extensionName) == 0) { + found = true; + requested_exts.push_back(ext); + } + } + if (found == false) { + ET_LOG(Info, "Failed to find extension %s", ext); + } + } + + // Create the device with our subset of features VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, nullptr}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &queue_info; - dci.enabledExtensionCount = 2; - dci.ppEnabledExtensionNames = dev_exts; + dci.enabledExtensionCount = requested_exts.size(); + dci.ppEnabledExtensionNames = requested_exts.data(); + ; + dci.pEnabledFeatures = &device_features; + dci.pNext = &features_graph; result = vkCreateDevice(*physical_device, &dci, nullptr, device); if (result != VK_SUCCESS) { ET_LOG(Error, "Failed to create VkDevice"); diff --git a/backends/arm/runtime/VGFSetup.cpp b/backends/arm/runtime/VGFSetup.cpp index 18c9dbc9727..abb4c50d8be 100644 --- a/backends/arm/runtime/VGFSetup.cpp +++ b/backends/arm/runtime/VGFSetup.cpp @@ -75,6 +75,24 @@ void free_tensor( vkFreeMemory(device, memory, nullptr); } +uint32_t get_memory_index( + VkPhysicalDevice vk_physical, + VkMemoryRequirements2 memory_requirements, + VkMemoryPropertyFlags aims) { + VkPhysicalDeviceMemoryProperties mem_properties; + vkGetPhysicalDeviceMemoryProperties(vk_physical, &mem_properties); + + uint32_t memory_type = 0; + for (size_t i = 0; i < 31; ++i) { + if (memory_requirements.memoryRequirements.memoryTypeBits & (0x1 << i)) { + memory_type = i; + if ((mem_properties.memoryTypes[i].propertyFlags & aims) == aims) + break; + } + } + return memory_type; +} + /** * Tensor allocation helper function */ @@ -135,26 +153,18 @@ VkResult allocate_tensor( vkGetTensorMemoryRequirementsARM( device, &memory_requirements_info, &memory_requirements); - VkPhysicalDeviceMemoryProperties memProps; - vkGetPhysicalDeviceMemoryProperties(physical, &memProps); + VkMemoryPropertyFlags aims = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + uint32_t memory_index = get_memory_index(physical, memory_requirements, aims); // Allocate memory - uint32_t memory_type = 0; - for (size_t j = 0; j < 31; ++j) { - if (memory_requirements.memoryRequirements.memoryTypeBits & (0x1 << j)) { - memory_type = j; - uint32_t aims = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | - VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; - if ((memProps.memoryTypes[j].propertyFlags & aims) == aims) - break; - } - } const VkMemoryAllocateInfo allocate_info = { .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, .pNext = nullptr, .allocationSize = memory_requirements.memoryRequirements.size, - .memoryTypeIndex = memory_type}; + .memoryTypeIndex = memory_index, + }; vkAllocateMemory(device, &allocate_info, nullptr, memory); @@ -213,8 +223,8 @@ static void debug_print_sequence( ET_LOG( Info, " %d: %s", - i, - string(sequence_decoder->getName(input_names, i)).c_str()); + j, + string(sequence_decoder->getName(input_names, j)).c_str()); } } } @@ -323,20 +333,23 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { // Parse the sequences in the VGF (while there can be multiple sequences of // COMPUTE and GRAPH segments in the sequence, we currently expect a single // GRAPH segment to be present. + const int segment_id = 0; + debug_print_sequence(sequence_decoder); if (sequence_decoder->modelSequenceTableSize() != 1) { ET_LOG(Error, "Expected sequence length 1"); return false; } - if (sequence_decoder->getSegmentType(0) != vgflib::ModuleType::GRAPH) { + if (sequence_decoder->getSegmentType(segment_id) != + vgflib::ModuleType::GRAPH) { ET_LOG(Error, "Expected segment to be of type GRAPH"); return false; } // Extract first segment and it's associated module debug_print_modules(module_decoder); - auto segment_name = string(sequence_decoder->getSegmentName(0)); - auto segment_module = sequence_decoder->getSegmentModuleIndex(0); + auto segment_name = string(sequence_decoder->getSegmentName(segment_id)); + auto segment_module = sequence_decoder->getSegmentModuleIndex(segment_id); auto segment_m_name = string(module_decoder->getModuleName(segment_module)); auto segment_m_entrypoint = @@ -454,13 +467,15 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { } // Constants table - mapping of shader bindings to MRT's and their descriptors - for (int i = 0; i < constant_decoder->size(); i++) { + auto constant_indexes = + sequence_decoder->getSegmentConstantIndexes(segment_id); + for (uint32_t i : constant_indexes) { auto mrt_i = constant_decoder->getConstantMrtIndex(i); auto constant_data = constant_decoder->getConstant(i); constants.push_back(VkDataGraphPipelineConstantARM{ .sType = VK_STRUCTURE_TYPE_DATA_GRAPH_PIPELINE_CONSTANT_ARM, .pNext = &descriptors[mrt_i], - .id = mrt_i, + .id = i, .pConstantData = constant_data.begin(), }); } @@ -469,9 +484,11 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { vector layout_bindings; vector data_graph_resources; - auto set_count = sequence_decoder->getSegmentDescriptorSetInfosSize(0); + auto set_count = + sequence_decoder->getSegmentDescriptorSetInfosSize(segment_id); for (uint32_t d_idx = 0; d_idx < set_count; d_idx++) { - auto handle = sequence_decoder->getDescriptorBindingSlotsHandle(0, d_idx); + auto handle = + sequence_decoder->getDescriptorBindingSlotsHandle(segment_id, d_idx); auto binding_count = sequence_decoder->getBindingsSize(handle); for (int binding = 0; binding < binding_count; binding++) { auto binding_index = @@ -517,14 +534,30 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { return false; } + std::vector poolSizes; + poolSizes.reserve(layout_bindings.size()); + for (const auto& b : layout_bindings) { + bool found = false; + for (size_t idx = 0; idx < poolSizes.size(); ++idx) { + if (poolSizes[idx].type == b.descriptorType) { + poolSizes[idx].descriptorCount += b.descriptorCount; + found = true; + break; + } + } + if (!found) { + poolSizes.push_back({b.descriptorType, b.descriptorCount}); + } + } + // Create descriptor pool and descriptors for pipeline const VkDescriptorPoolCreateInfo descriptor_pool_info = { .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO, .pNext = nullptr, .flags = 0, .maxSets = static_cast(set_count), - .poolSizeCount = 0, - .pPoolSizes = nullptr, + .poolSizeCount = static_cast(poolSizes.size()), + .pPoolSizes = poolSizes.data(), }; result = vkCreateDescriptorPool( vk_device, &descriptor_pool_info, nullptr, &vk_descriptor_pool); @@ -544,8 +577,7 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { // Alloc descriptor sets // currently, as we require modelSequenceTableSize to == 1 // we can only get one descriptor set. - vector descriptor_sets; - descriptor_sets.resize(1); + descriptor_sets.resize(layout_bindings.size()); result = vkAllocateDescriptorSets( vk_device, &descriptor_set_info, descriptor_sets.data()); if (result != VK_SUCCESS) { @@ -554,7 +586,8 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { } // write descriptor updates for every input - auto input_slots = sequence_decoder->getSegmentInputBindingSlotsHandle(0); + auto input_slots = + sequence_decoder->getSegmentInputBindingSlotsHandle(segment_id); auto input_size = sequence_decoder->getBindingsSize(input_slots); for (uint32_t i = 0; i < input_size; i++) { auto binding = sequence_decoder->getBindingSlotBinding(input_slots, i); @@ -582,7 +615,8 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { } // write descriptor updates for every output - auto output_slots = sequence_decoder->getSegmentOutputBindingSlotsHandle(0); + auto output_slots = + sequence_decoder->getSegmentOutputBindingSlotsHandle(segment_id); auto output_size = sequence_decoder->getBindingsSize(output_slots); for (uint32_t i = 0; i < output_size; i++) { auto binding = sequence_decoder->getBindingSlotBinding(output_slots, i); @@ -678,23 +712,144 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { } // Allocate command buffer - VkCommandBufferAllocateInfo allocate_info{ + VkCommandBufferAllocateInfo buffer_allocate_info{ .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO, .pNext = nullptr, .commandPool = vk_command_pool, .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY, .commandBufferCount = 1}; - result = vkAllocateCommandBuffers(vk_device, &allocate_info, &vk_execute_cmd); + result = vkAllocateCommandBuffers( + vk_device, &buffer_allocate_info, &vk_execute_cmd); if (result != VK_SUCCESS) { ET_LOG(Error, "Failed to allocate command buffers"); return result; } + // Allocate intermediates memory based on the pipeline requirements provided + // by the driver + VkDataGraphPipelineSessionBindPointRequirementsInfoARM + bind_point_requirements_info = { + .sType = + VK_STRUCTURE_TYPE_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_REQUIREMENTS_INFO_ARM, + .pNext = nullptr, + .session = vk_session, + }; + + uint32_t bind_point_count = 0; + result = vkGetDataGraphPipelineSessionBindPointRequirementsARM( + vk_device, &bind_point_requirements_info, &bind_point_count, nullptr); + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to get session bind point count"); + return result; + } + + vector + bind_point_requirements; + bind_point_requirements.resize(bind_point_count); + result = vkGetDataGraphPipelineSessionBindPointRequirementsARM( + vk_device, + &bind_point_requirements_info, + &bind_point_count, + bind_point_requirements.data()); + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to get session bind point requirements"); + return result; + } + + // Given the bind points, just make individual allocations and bind them + for (const auto& bind_point_requirement : bind_point_requirements) { + // These are the only allowed type and bindpoint with the current spec + if (bind_point_requirement.bindPointType != + VK_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_TYPE_MEMORY_ARM) { + ET_LOG( + Error, + "Expected VK_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_TYPE_MEMORY_ARM"); + return VK_ERROR_UNKNOWN; + } + if (bind_point_requirement.bindPoint != + VK_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_TRANSIENT_ARM) { + ET_LOG( + Error, + "Expected VK_DATA_GRAPH_PIPELINE_SESSION_BIND_POINT_TRANSIENT_ARM"); + return VK_ERROR_UNKNOWN; + } + if (bind_point_requirement.numObjects != 1) { + ET_LOG(Error, "Expected only one object for the bindpoint"); + return VK_ERROR_UNKNOWN; + } + + VkDataGraphPipelineSessionMemoryRequirementsInfoARM memory_requirements_info = { + .sType = + VK_STRUCTURE_TYPE_DATA_GRAPH_PIPELINE_SESSION_MEMORY_REQUIREMENTS_INFO_ARM, + .pNext = nullptr, + .session = vk_session, + .bindPoint = bind_point_requirement.bindPoint, + .objectIndex = 0, // NOTE: tied to numObjects assert above + }; + VkMemoryRequirements2 memory_requirements; + vkGetDataGraphPipelineSessionMemoryRequirementsARM( + vk_device, &memory_requirements_info, &memory_requirements); + + VkMemoryPropertyFlags aims = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | + VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | + VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + uint32_t memory_index = + get_memory_index(vk_physical, memory_requirements, aims); + + VkMemoryAllocateInfo memory_allocate_info = { + .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, + .pNext = nullptr, + .allocationSize = memory_requirements.memoryRequirements.size, + .memoryTypeIndex = memory_index, + }; + + VkDeviceMemory memory; + result = + vkAllocateMemory(vk_device, &memory_allocate_info, nullptr, &memory); + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to allocate memory for intermediates"); + return result; + } + // so we can free this object in destructor + intermediates.push_back(memory); + + VkBindDataGraphPipelineSessionMemoryInfoARM bind_info = { + .sType = + VK_STRUCTURE_TYPE_BIND_DATA_GRAPH_PIPELINE_SESSION_MEMORY_INFO_ARM, + .pNext = nullptr, + .session = vk_session, + .bindPoint = bind_point_requirement.bindPoint, + .objectIndex = 0, // NOTE: tied to numObjects assert above + .memory = memory, + .memoryOffset = 0, + }; + result = vkBindDataGraphPipelineSessionMemoryARM(vk_device, 1, &bind_info); + if (result != VK_SUCCESS) { + ET_LOG(Error, "Failed to bind intermediates memory"); + return result; + } + } + // Populate command once with our dispatch information VkCommandBufferBeginInfo beginInfo{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; vkBeginCommandBuffer(vk_execute_cmd, &beginInfo); + // Sync what will be the data coming in from host + VkMemoryBarrier2 barrier = { + .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER_2, + .srcStageMask = VK_PIPELINE_STAGE_2_HOST_BIT, + .srcAccessMask = VK_ACCESS_2_HOST_WRITE_BIT, + .dstStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .dstAccessMask = VK_ACCESS_2_SHADER_READ_BIT, + }; + VkDependencyInfo dependency_info = { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .memoryBarrierCount = 1, + .pMemoryBarriers = &barrier, + }; + vkCmdPipelineBarrier2(vk_execute_cmd, &dependency_info); + // bind pipeline + descriptor set vkCmdBindPipeline( vk_execute_cmd, VK_PIPELINE_BIND_POINT_DATA_GRAPH_ARM, vk_pipeline); @@ -713,6 +868,21 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef specs) { // Dispatch the graph command vkCmdDispatchDataGraphARM(vk_execute_cmd, vk_session, nullptr); + // Sync data back + VkMemoryBarrier2 barrier_2 = { + .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER_2, + .srcStageMask = VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT, + .srcAccessMask = VK_ACCESS_2_SHADER_WRITE_BIT, + .dstStageMask = VK_PIPELINE_STAGE_2_HOST_BIT, + .dstAccessMask = VK_ACCESS_2_HOST_READ_BIT, + }; + VkDependencyInfo dependency_info_2 = { + .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO, + .memoryBarrierCount = 1, + .pMemoryBarriers = &barrier_2, + }; + vkCmdPipelineBarrier2(vk_execute_cmd, &dependency_info_2); + // end the command buffer vkEndCommandBuffer(vk_execute_cmd); @@ -748,6 +918,9 @@ void VgfRepr::free_vgf() { free_tensor( vk_device, IOs[i].tensor_view, IOs[i].tensor, IOs[i].tensor_memory); } + for (auto memory : intermediates) { + vkFreeMemory(vk_device, memory, nullptr); + } } static uint32_t get_format_size(VkFormat format) { diff --git a/backends/arm/runtime/VGFSetup.h b/backends/arm/runtime/VGFSetup.h index 29fc287865e..3faa7ca620a 100644 --- a/backends/arm/runtime/VGFSetup.h +++ b/backends/arm/runtime/VGFSetup.h @@ -75,6 +75,7 @@ class VgfRepr { * with the EValues coming the backend execute call */ vector IOs; + vector intermediates; bool map_io(IO* io, void** handle) { VkResult result = @@ -112,6 +113,7 @@ class VgfRepr { VkDescriptorSetLayout vk_layout; VkShaderModule vk_shader; // Note: the vector of tensor memory is stored in IOs above + vector descriptor_sets; }; } // namespace vgf diff --git a/backends/arm/scripts/TOSA_minimal_example.ipynb b/backends/arm/scripts/TOSA_minimal_example.ipynb new file mode 100644 index 00000000000..785affc657b --- /dev/null +++ b/backends/arm/scripts/TOSA_minimal_example.ipynb @@ -0,0 +1,281 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Copyright 2025 Arm Limited and/or its affiliates.\n", + "#\n", + "# This source code is licensed under the BSD-style license found in the\n", + "# LICENSE file in the root directory of this source tree." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# TOSA delegate flow example\n", + "\n", + "This guide walks through the complete process of running a module on Arm TOSA using ExecuTorch, with a focus on TOSA lowering exploration. \n", + "This workflow is intended for validating and experimenting with model lowering to TOSA, and is aimed at contributors and developers, rather than production deployment.\n", + "It’s important to note that the compilation flow and passes applied can vary based on the target, so this flow does not necessarily produce TOSA flatbuffers and PTE files which are optimal (or even compatible) with any one target.\n", + "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 TOSA dependencies using `examples/arm/setup.sh --disable-ethos-u-deps`\n", + "\n", + "With all commands executed from the base `executorch` folder.\n", + "\n", + "\n", + "\n", + "*Some scripts in this notebook produces long output logs: Configuring the 'Customizing Notebook Layout' settings to enable 'Output:scrolling' and setting 'Output:Text Line Limit' makes this more manageable*" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## AOT Flow\n", + "\n", + "The first step is creating the PyTorch module and exporting it. Exporting converts the python code in the module into a graph structure. The result is still runnable python code, which can be displayed by printing the `graph_module` of the exported program. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "print(torch.__version__)\n", + "\n", + "class Add(torch.nn.Module):\n", + " def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n", + " return x + y\n", + "\n", + "example_inputs = (torch.ones(1,1,1,1),torch.ones(1,1,1,1))\n", + "\n", + "model = Add()\n", + "model = model.eval()\n", + "exported_program = torch.export.export(model, example_inputs)\n", + "graph_module = exported_program.module()\n", + "\n", + "_ = graph_module.print_readable()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## TOSA backend supports both INT and FP targets.\n", + "\n", + "To lower the graph_module for FP targets using the TOSA backend, we run it through the default FP lowering pipeline.\n", + "\n", + "FP lowering can be customized for different subgraphs; the sequence shown here is the recommended workflow for TOSA. Because we are staying in floating-point precision, no calibration with example inputs is required.\n", + "\n", + "If you print the module again, you will see that nodes are left in FP form (or annotated with any necessary casts) without any quantize/dequantize wrappers." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from executorch.backends.arm.arm_backend import (\n", + " ArmCompileSpecBuilder,\n", + ")\n", + "from executorch.backends.arm.tosa.specification import TosaSpecification\n", + "from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e\n", + "from pathlib import Path\n", + "\n", + "target = \"TOSA-1.0+FP\"\n", + "base_name = \"tosa_simple_example\"\n", + "cwd_dir = Path.cwd()\n", + "\n", + "# Create a compilation spec describing the target for configuring the quantizer\n", + "# Dump intermediate artifacts (in this case TOSA flat buffers) to specified location\n", + "tosa_spec = TosaSpecification.create_from_string(target)\n", + "spec_builder = (ArmCompileSpecBuilder()\n", + " .tosa_compile_spec(tosa_spec)\n", + " .dump_intermediate_artifacts_to(str(cwd_dir / base_name)))\n", + "compile_spec = spec_builder.build()\n", + "\n", + "_ = graph_module.print_readable()\n", + "\n", + "# Create a new exported program using the quantized_graph_module\n", + "lowered_exported_program = torch.export.export(graph_module, example_inputs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To lower the graph_module for INT targets using the TOSA backend, we apply the arm_quantizer.\n", + "\n", + "Quantization can be performed in various ways and tailored to different subgraphs; the sequence shown here represents the recommended workflow for TOSA.\n", + "\n", + "This step also requires calibrating the module with representative inputs.\n", + "\n", + "If you print the module again, you’ll see that each node is now wrapped in quantization/dequantization nodes that embed the calculated quantization parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from executorch.backends.arm.arm_backend import (\n", + " ArmCompileSpecBuilder,\n", + " get_tosa_spec,\n", + ")\n", + "from executorch.backends.arm.quantizer import (\n", + " TOSAQuantizer,\n", + " get_symmetric_quantization_config,\n", + ")\n", + "from executorch.backends.arm.tosa.specification import TosaSpecification\n", + "from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e\n", + "from pathlib import Path\n", + "\n", + "target = \"TOSA-1.0+INT\"\n", + "base_name = \"tosa_simple_example\"\n", + "cwd_dir = Path.cwd()\n", + "\n", + "# Create a compilation spec describing the target for configuring the quantizer\n", + "# Dump intermediate artifacts (in this case TOSA flat buffers) to specified location\n", + "tosa_spec = TosaSpecification.create_from_string(target)\n", + "spec_builder = (ArmCompileSpecBuilder()\n", + " .tosa_compile_spec(tosa_spec)\n", + " .dump_intermediate_artifacts_to(str(cwd_dir / base_name)))\n", + "compile_spec = spec_builder.build()\n", + "\n", + "# Create and configure quantizer to use a symmetric quantization config globally on all nodes\n", + "quantizer = TOSAQuantizer(get_tosa_spec(compile_spec))\n", + "operator_config = get_symmetric_quantization_config()\n", + "quantizer.set_global(operator_config)\n", + "\n", + "# Post training quantization\n", + "quantized_graph_module = prepare_pt2e(graph_module, quantizer)\n", + "quantized_graph_module(*example_inputs) # Calibrate the graph module with the example input\n", + "quantized_graph_module = convert_pt2e(quantized_graph_module)\n", + "\n", + "_ = quantized_graph_module.print_readable()\n", + "\n", + "# Create a new exported program using the quantized_graph_module\n", + "lowered_exported_program = torch.export.export(quantized_graph_module, example_inputs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The lowering in the TOSABackend happens in four steps:\n", + "\n", + "1. **Lowering to core Aten operator set**: Transform module to use a subset of operators applicable to edge devices. \n", + "2. **Partitioning**: Find subgraphs which are supported for running on TOSA\n", + "3. **Lowering to TOSA compatible operator set**: Perform transforms to make the TOSA subgraph(s) compatible with TOSA operator set\n", + "4. **Serialization to TOSA**: Compiles the graph module into a TOSA graph \n", + "Step 4 also prints a Network summary for each processed subgraph.\n", + "\n", + "All of this happens behind the scenes in `to_edge_transform_and_lower`. Printing the graph module shows that what is left in the graph is two quantization nodes for `x` and `y` going into an `executorch_call_delegate` node, followed by a dequantization node." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from executorch.backends.arm.tosa.partitioner import TOSAPartitioner\n", + "from executorch.exir import (\n", + " EdgeCompileConfig,\n", + " ExecutorchBackendConfig,\n", + " to_edge_transform_and_lower,\n", + ")\n", + "from executorch.extension.export_util.utils import save_pte_program\n", + "\n", + "# Create partitioner from compile spec\n", + "partitioner = TOSAPartitioner(compile_spec)\n", + "\n", + "# Lower the exported program to the TOSA backend\n", + "edge_program_manager = to_edge_transform_and_lower(\n", + " lowered_exported_program,\n", + " partitioner=[partitioner],\n", + " compile_config=EdgeCompileConfig(\n", + " _check_ir_validity=False,\n", + " ),\n", + " )\n", + "\n", + "# Convert edge program to executorch\n", + "executorch_program_manager = edge_program_manager.to_executorch(\n", + " config=ExecutorchBackendConfig(extract_delegate_segments=False)\n", + " )\n", + "\n", + "executorch_program_manager.exported_program().module().print_readable()\n", + "\n", + "# Save pte file\n", + "pte_name = base_name + \".pte\"\n", + "pte_path = cwd_dir / base_name / pte_name\n", + "save_pte_program(executorch_program_manager, str(pte_path))\n", + "assert pte_path.exists(), \"Build failed; no .pte-file found\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Use TOSA reference model to verify TOSA graph\n", + "\n", + "After the AOT compilation flow is done, the resulting lowered TOSA graph can be verified using the TOSA reference model tool." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "import tosa_reference_model as reference_model\n", + "from executorch.backends.arm.test.runner_utils import TosaReferenceModelDispatch\n", + "\n", + "# Run TOSA graph through reference model using sample inputs\n", + "with TosaReferenceModelDispatch():\n", + " executorch_program_manager.exported_program().module()(*example_inputs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/backends/arm/scripts/build_executor_runner.sh b/backends/arm/scripts/build_executor_runner.sh index 974c5ca1ff7..a05287ac4bf 100755 --- a/backends/arm/scripts/build_executor_runner.sh +++ b/backends/arm/scripts/build_executor_runner.sh @@ -25,6 +25,7 @@ output_folder_set=false output_folder="." et_build_root="${et_root_dir}/arm_test" ethosu_tools_dir=${et_root_dir}/examples/arm/ethos-u-scratch +select_ops_list="" build_bundleio_flags=" -DET_BUNDLE_IO=OFF " build_with_etdump_flags=" -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF " @@ -32,21 +33,26 @@ build_with_etdump_flags=" -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF " help() { echo "Usage: $(basename $0) [options]" echo "Options:" - echo " --pte= pte file (genrated by the aot_arm_compier from the model to include in the elf" - echo " --target= Target to build and run for Default: ${target}" - echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" - echo " --bundleio Support both pte and Bundle IO bpte using Devtools BundelIO with Input/RefOutput included" - echo " --system_config= System configuration to select from the Vela configuration file (see vela.ini). Default: Ethos_U55_High_End_Embedded for EthosU55 targets, Ethos_U85_SYS_DRAM_Mid for EthosU85 targets." - echo " NOTE: If given, this option must match the given target. This option along with the memory_mode sets timing adapter values customized for specific hardware, see ./executor_runner/CMakeLists.txt." - echo " --memory_mode= Vela memory mode, used for setting the Timing Adapter parameters of the Corstone platforms." - echo " Valid values are Shared_Sram(for Ethos-U55, Ethos-U65, Ethos-85), Sram_Only(for Ethos-U55, Ethos-U65, Ethos-U85) or Dedicated_Sram(for Ethos-U65, Ethos-U85)." - echo " Default: Shared_Sram for the Ethos-U55 and Sram_Only for the Ethos-U85" - echo " --etdump Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log" - echo " --extra_build_flags= Extra flags to pass to cmake like -DET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE=60000 Default: none " - echo " --output= Output folder Default: /_.pte" - echo " --et_build_root= Build output root folder to use, defaults to ${et_build_root}" - echo " --ethosu_tools_dir= Path to your Ethos-U tools dir if you not using default: ${ethosu_tools_dir}" - echo " --toolchain= Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc" + echo " --pte=||semihosting Set to a pte file (generated by the aot_arm_compier) to include the model in the elf." + echo " Or a hex address in the format of 0x00000000 if placed in memory you need to place it on this ADDR on your target, with your flash tool or other means." + echo " Or specify the word 'semihosting' to supply pte at runtime." + echo " --target= Target to build and run for Default: ${target}" + echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" + echo " --bundleio Support both pte and Bundle IO bpte using Devtools BundelIO with Input/RefOutput included" + echo " --system_config= System configuration to select from the Vela configuration file (see vela.ini). Default: Ethos_U55_High_End_Embedded for EthosU55 targets, Ethos_U85_SYS_DRAM_Mid for EthosU85 targets." + echo " NOTE: If given, this option must match the given target. This option along with the memory_mode sets timing adapter values customized for specific hardware, see ./executor_runner/CMakeLists.txt." + echo " --memory_mode= Vela memory mode, used for setting the Timing Adapter parameters of the Corstone platforms." + echo " Valid values are Shared_Sram(for Ethos-U55, Ethos-U65, Ethos-85), Sram_Only(for Ethos-U55, Ethos-U65, Ethos-U85) or Dedicated_Sram(for Ethos-U65, Ethos-U85)." + echo " Default: Shared_Sram for the Ethos-U55 and Sram_Only for the Ethos-U85" + echo " --etdump Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log" + echo " --extra_build_flags= Extra flags to pass to cmake like -DET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE=60000 Default: none " + echo " --output= Output folder Default: /_.pte" + echo " --et_build_root= Build output root folder to use, defaults to ${et_build_root}" + echo " --ethosu_tools_dir= Path to your Ethos-U tools dir if you not using default: ${ethosu_tools_dir}" + echo " --toolchain= Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc Default: ${toolchain}" + echo " --select_ops_list= Comma separated list of portable (non delagated) kernels to include Default: ${select_ops_list}" + echo " NOTE: This is used when select_ops_model is not possible to use, e.g. for semihosting or bundleio." + echo " See https://docs.pytorch.org/executorch/stable/kernel-library-selective-build.html for more information." exit 0 } @@ -65,6 +71,7 @@ for arg in "$@"; do --et_build_root=*) et_build_root="${arg#*=}";; --ethosu_tools_dir=*) ethosu_tools_dir="${arg#*=}";; --toolchain=*) toolchain="${arg#*=}";; + --select_ops_list=*) select_ops_list="${arg#*=}";; *) ;; esac @@ -75,7 +82,7 @@ if [[ ${toolchain} == "arm-none-eabi-gcc" ]]; then elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then toolchain_cmake=${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake else - echo "Error: Invalid toolchain selection, provided: ${tolchain}" + echo "Error: Invalid toolchain selection, provided: ${toolchain}" echo " Valid options are {arm-none-eabi-gcc, arm-zephyr-eabi-gcc}" exit 1; fi @@ -88,20 +95,35 @@ toolchain_cmake=$(realpath ${toolchain_cmake}) source ${setup_path_script} -pte_file=$(realpath ${pte_file}) +if [[ ${pte_file} == "semihosting" ]]; then + pte_data="-DSEMIHOSTING=ON" +else + if [[ "$pte_file" =~ ^0x[0-9a-fA-F]{1,16}$ ]]; then + echo "PTE in memory at ${pte_file}, make sure to put it there on your target before starting." + pte_data="-DET_MODEL_PTE_ADDR=${pte_file}" + if [ "$output_folder_set" = false ] ; then + # Not locked down to a PTE use + output_folder=${et_build_root}/${target}_${pte_file}/cmake-out + fi + else + echo "PTE included in elf from file ${pte_file}" + pte_file=$(realpath ${pte_file}) + pte_data="-DET_PTE_FILE_PATH:PATH=${pte_file}" + if [ "$output_folder_set" = false ] ; then + # remove file ending + output_folder=${pte_file%.*}/cmake-out + fi + fi +fi ethosu_tools_dir=$(realpath ${ethosu_tools_dir}) ethos_u_root_dir="$ethosu_tools_dir/ethos-u" mkdir -p "${ethos_u_root_dir}" ethosu_tools_dir=$(realpath ${ethos_u_root_dir}) et_build_dir=${et_build_root}/cmake-out +mkdir -p ${et_build_dir} et_build_dir=$(realpath ${et_build_dir}) -if [ "$output_folder_set" = false ] ; then - # remove file ending - output_folder=${pte_file%.*} -fi - if [[ ${system_config} == "" ]] then system_config="Ethos_U55_High_End_Embedded" @@ -129,7 +151,7 @@ else target_cpu=cortex-m85 fi echo "--------------------------------------------------------------------------------" -echo "Build Arm ${toolchain/-gcc/} executor_runner for ${target} with ${pte_file} using ${system_config} ${memory_mode} ${extra_build_flags} to '${output_folder}/cmake-out'" +echo "Build Arm ${toolchain/-gcc/} executor_runner for ${target} PTE: ${pte_file} using ${system_config} ${memory_mode} ${extra_build_flags} to '${output_folder}'" echo "--------------------------------------------------------------------------------" cd ${et_root_dir}/examples/arm/executor_runner @@ -149,23 +171,24 @@ cmake \ -DTARGET_CPU=${target_cpu} \ -DET_DIR_PATH:PATH=${et_root_dir} \ -DET_BUILD_DIR_PATH:PATH=${et_build_dir} \ - -DET_PTE_FILE_PATH:PATH="${pte_file}" \ -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir} \ -DETHOSU_TARGET_NPU_CONFIG=${target} \ + ${pte_data} \ ${build_bundleio_flags} \ ${build_with_etdump_flags} \ -DPYTHON_EXECUTABLE=$(which python3) \ -DSYSTEM_CONFIG=${system_config} \ -DMEMORY_MODE=${memory_mode} \ + -DEXECUTORCH_SELECT_OPS_LIST="${select_ops_list}" \ ${extra_build_flags} \ - -B ${output_folder}/cmake-out + -B ${output_folder} echo "[${BASH_SOURCE[0]}] Configured CMAKE" -cmake --build ${output_folder}/cmake-out -j$(nproc) -- arm_executor_runner +cmake --build ${output_folder} -j$(nproc) -- arm_executor_runner echo "[${BASH_SOURCE[0]}] Generated ${toolchain} elf file:" -find ${output_folder}/cmake-out -name "arm_executor_runner" -echo "executable_text: $(find ${output_folder}/cmake-out -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $1}') bytes" -echo "executable_data: $(find ${output_folder}/cmake-out -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $2}') bytes" -echo "executable_bss: $(find ${output_folder}/cmake-out -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $3}') bytes" +find ${output_folder} -name "arm_executor_runner" +echo "executable_text: $(find ${output_folder} -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $1}') bytes" +echo "executable_data: $(find ${output_folder} -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $2}') bytes" +echo "executable_bss: $(find ${output_folder} -name arm_executor_runner -exec ${toolchain/-gcc/-size} {} \; | grep -v filename | awk '{print $3}') bytes" diff --git a/backends/arm/scripts/build_executor_runner_vkml.sh b/backends/arm/scripts/build_executor_runner_vkml.sh new file mode 100755 index 00000000000..1df63acc425 --- /dev/null +++ b/backends/arm/scripts/build_executor_runner_vkml.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# Copyright 2025 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 -eu + +script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +et_root_dir=$(cd ${script_dir}/../../.. && pwd) +et_root_dir=$(realpath ${et_root_dir}) +setup_path_script=${et_root_dir}/examples/arm/ethos-u-scratch/setup_path.sh +_setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." + +build_type="Release" +build_with_etdump=false +extra_build_flags="" +output_folder="cmake-out-vkml" + +build_with_etdump_flags=" -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF " + +help() { + echo "Usage: $(basename $0) [options]" + echo "Options:" + echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" + echo " --etdump Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log" + echo " --extra_build_flags= Extra flags to pass to cmake. Default: none " + echo " --output= Output folder Default: $(output_folder)" + exit 0 +} + +for arg in "$@"; do + case $arg in + -h|--help) help ;; + --build_type=*) build_type="${arg#*=}";; + --etdump) build_with_etdump=true ;; + --extra_build_flags=*) extra_build_flags="${arg#*=}";; + --output=*) output_folder="${arg#*=}";; + --select_ops_list=*) select_ops_list="${arg#*=}";; + *) + ;; + esac +done + +# Source the tools +# This should be prepared by the setup.sh +[[ -f ${setup_path_script} ]] \ + || { echo "Missing ${setup_path_script}. ${_setup_msg}"; exit 1; } + +source ${setup_path_script} + +mkdir -p "${output_folder}" +output_folder=$(realpath ${output_folder}) + +echo "--------------------------------------------------------------------------------" +echo "Build Arm VKML executor runner: '${output_folder}' with extra build flags: ${extra_build_flags}" +echo "--------------------------------------------------------------------------------" + +cd ${et_root_dir}/examples/arm/executor_runner + +if [ "$build_with_etdump" = true ] ; then + build_with_etdump_flags=" -DEXECUTORCH_ENABLE_EVENT_TRACER=ON " +fi + +echo "Building with extra flags: ${build_with_etdump_flags} ${extra_build_flags}" +cmake \ + -Wall \ + -Werror \ + -DCMAKE_BUILD_TYPE=${build_type} \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_XNNPACK=OFF \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DEXECUTORCH_BUILD_VGF=ON \ + -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DPYTHON_EXECUTABLE=$(which python3) \ + ${extra_build_flags} \ + -B ${output_folder} ${et_root_dir} + +echo "[${BASH_SOURCE[0]}] Configured CMAKE" + +cmake --build ${output_folder} -j$(nproc) + +echo "[${BASH_SOURCE[0]}] Built VKML runner: " +find ${output_folder} -name "executor_runner" diff --git a/backends/arm/scripts/build_executorch.sh b/backends/arm/scripts/build_executorch.sh index 1e2ac6ad055..84c675ddb4a 100755 --- a/backends/arm/scripts/build_executorch.sh +++ b/backends/arm/scripts/build_executorch.sh @@ -29,7 +29,7 @@ help() { echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" echo " --devtools Build Devtools libs" echo " --etdump Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log" - echo " --toolchain= Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc" + echo " --toolchain= Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc Default: ${toolchain}" exit 0 } @@ -51,7 +51,7 @@ if [[ ${toolchain} == "arm-none-eabi-gcc" ]]; then elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then toolchain_cmake=${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake else - echo "Error: Invalid toolchain selection, provided: ${tolchain}" + echo "Error: Invalid toolchain selection, provided: ${toolchain}" echo " Valid options are {arm-none-eabi-gcc, arm-zephyr-eabi-gcc}" exit 1; fi diff --git a/backends/arm/scripts/build_portable_kernels.sh b/backends/arm/scripts/build_portable_kernels.sh index 0d06b59dd03..cfa008c80d5 100755 --- a/backends/arm/scripts/build_portable_kernels.sh +++ b/backends/arm/scripts/build_portable_kernels.sh @@ -4,92 +4,4 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# Optional parameter: -# --build_type= "Release" | "Debug" | "RelWithDebInfo" -# --etdump build with devtools-etdump support - -set -eu - -script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -et_root_dir=$(cd ${script_dir}/../../.. && pwd) -et_root_dir=$(realpath ${et_root_dir}) -toolchain=arm-none-eabi-gcc -setup_path_script=${et_root_dir}/examples/arm/ethos-u-scratch/setup_path.sh -_setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." - - -et_build_root="${et_root_dir}/arm_test" -build_type="Release" -portable_kernels="aten::_softmax.out" - -help() { - echo "Usage: $(basename $0) [options]" - echo "Options:" - echo " --et_build_root= Build output root folder to use, defaults to ${et_build_root}" - echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" - echo " --portable_kernels= Comma separated list of portable (non delagated) kernels to include Default: ${portable_kernels}" - echo " --toolchain= Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc" - exit 0 -} - -for arg in "$@"; do - case $arg in - -h|--help) help ;; - --et_build_root=*) et_build_root="${arg#*=}";; - --build_type=*) build_type="${arg#*=}";; - --portable_kernels=*) portable_kernels="${arg#*=}";; - --toolchain=*) toolchain="${arg#*=}";; - *) - ;; - esac -done - -if [[ ${toolchain} == "arm-none-eabi-gcc" ]]; then - toolchain_cmake=${et_root_dir}/examples/arm/ethos-u-setup/${toolchain}.cmake -elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then - toolchain_cmake=${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake -else - echo "Error: Invalid toolchain selection, provided: ${tolchain}" - echo " Valid options are {arm-none-eabi-gcc, arm-zephyr-eabi-gcc}" - exit 1; -fi -toolchain_cmake=$(realpath ${toolchain_cmake}) - -# Source the tools -# This should be prepared by the setup.sh -[[ -f ${setup_path_script} ]] \ - || { echo "Missing ${setup_path_script}. ${_setup_msg}"; exit 1; } - -source ${setup_path_script} - -et_build_dir=${et_build_root}/cmake-out - -cd "${et_root_dir}" - -echo "--------------------------------------------------------------------------------" ; -echo "Build ExecuTorch Libraries ${build_type} portable kernels: ${portable_kernels} into '${et_build_dir}'" ; -echo "--------------------------------------------------------------------------------" - -if ! [[ $portable_kernels =~ ^((^|,)aten::[a-zA-Z0-9_]+\.[a-zA-Z0-9_]*out)*$ ]]; then - echo " ERROR: specified argument --portable_kernels=${portable_kernels}" - echo " is in the wrong format please use \"aten::.out,aten::.out,...\"" - echo " e.g. \"aten::_softmax.out,aten::add.out\"" - exit 1 -fi - -set -x - -cmake \ - -DCMAKE_INSTALL_PREFIX=${et_build_dir} \ - -DCMAKE_BUILD_TYPE=${build_type} \ - -DCMAKE_TOOLCHAIN_FILE="${toolchain_cmake}" \ - -DEXECUTORCH_SELECT_OPS_LIST=${portable_kernels} \ - -B"${et_build_dir}/examples/arm" \ - "${et_root_dir}/examples/arm" - -cmake --build "${et_build_dir}/examples/arm" -j$(nproc) --config ${build_type} -- - -set +x - -echo "[$(basename $0)] Generated static libraries for ExecuTorch:" -find "${et_build_dir}/examples/arm" -name "*.a" -exec ls -al {} \; +echo "DEPRECATED: build_portable_kernels.sh is deprecated and will be removed. The kernel registration library is now built directly with the arm_executor_runner." diff --git a/backends/arm/scripts/corstone_utils.cmake b/backends/arm/scripts/corstone_utils.cmake index af5f866c461..8253f3985ca 100644 --- a/backends/arm/scripts/corstone_utils.cmake +++ b/backends/arm/scripts/corstone_utils.cmake @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. function(fetch_ethos_u_content ETHOS_SDK_PATH ET_DIR_PATH) + message(STATUS "Fetching Ethos-U content into ${ETHOS_SDK_PATH}") + file(MAKE_DIRECTORY ${ETHOS_SDK_PATH}/../ethos_u) include(FetchContent) set(ethos_u_base_tag "25.05") @@ -55,7 +57,6 @@ function(fetch_ethos_u_content ETHOS_SDK_PATH ET_DIR_PATH) "pwd && source backends/arm/scripts/utils.sh && patch_repo ${ETHOS_SDK_PATH}/core_platform ${core_platform_base_rev} ${patch_dir}" WORKING_DIRECTORY ${ET_DIR_PATH} COMMAND_ECHO STDOUT ) - endfunction() function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH) diff --git a/backends/arm/scripts/fvp_utils.sh b/backends/arm/scripts/fvp_utils.sh new file mode 100644 index 00000000000..cf0774e8706 --- /dev/null +++ b/backends/arm/scripts/fvp_utils.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# Copyright 2025 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. + +# The purpose of this file is to isolate all the functions to download +# and setup FVP. The reasons for behind this are multiple. +# The FVP needs a EULA acceptance and since the software download +# differs for arch and os it becomes quite messy to try and handle +# inside setup.sh. + +# Important to check for unset variables since this script is always sourced from setup.sh +set -u + +# Check if the script is being sourced +(return 0 2>/dev/null) +if [[ $? -ne 0 ]]; then + echo "Error: This script must be sourced." + exit 1 +fi + +if [[ "${ARCH}" == "x86_64" ]]; then + # FVPs + corstone300_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-300/FVP_Corstone_SSE-300_11.22_20_Linux64.tgz?rev=018659bd574f4e7b95fa647e7836ccf4&hash=22A79103C6FA5FFA7AFF3BE0447F3FF9" + corstone300_model_dir="Linux64_GCC-9.3" + corstone300_md5_checksum="98e93b949d0fbac977292d8668d34523" + + corstone320_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-320/FVP_Corstone_SSE-320_11.27_25_Linux64.tgz?rev=a507bffc219a4d5792f1192ab7002d89&hash=D9A824AA8227D2E679C9B9787FF4E8B6FBE3D7C6" + corstone320_model_dir="Linux64_GCC-9.3" + corstone320_md5_checksum="3deb3c68f9b2d145833f15374203514d" +elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then + # FVPs + corstone300_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-300/FVP_Corstone_SSE-300_11.22_20_Linux64_armv8l.tgz?rev=9cc6e9a32bb947ca9b21fa162144cb01&hash=7657A4CF27D42E892E3F08D452AAB073" + corstone300_model_dir="Linux64_armv8l_GCC-9.3" + corstone300_md5_checksum="cbbabbe39b07939cff7a3738e1492ef1" + + corstone320_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-320/FVP_Corstone_SSE-320_11.27_25_Linux64_armv8l.tgz?rev=b6ebe0923cb84f739e017385fd3c333c&hash=8965C4B98E2FF7F792A099B08831FE3CB6120493" + corstone320_model_dir="Linux64_armv8l_GCC-9.3" + corstone320_md5_checksum="3889f1d80a6d9861ea4aa6f1c88dd0ae" +else + echo "[main] Error: only x86-64 & aarch64/arm64 architecture is supported for now!"; exit 1; +fi + +function install_fvp() { + # Download and install the Corstone 300 FVP simulator platform + fvps=("corstone300" "corstone320") + + for fvp in "${fvps[@]}"; do + cd "${root_dir}" + if [[ ! -e "FVP_${fvp}.tgz" ]]; then + echo "[${FUNCNAME[0]}] Downloading FVP ${fvp}..." + url_variable=${fvp}_url + fvp_url=${!url_variable} + curl --output "FVP_${fvp}.tgz" "${fvp_url}" + md5_variable=${fvp}_md5_checksum + fvp_md5_checksum=${!md5_variable} + verify_md5 ${fvp_md5_checksum} FVP_${fvp}.tgz || exit 1 + fi + + echo "[${FUNCNAME[0]}] Installing FVP ${fvp}..." + rm -rf FVP-${fvp} + mkdir -p FVP-${fvp} + cd FVP-${fvp} + tar xf ../FVP_${fvp}.tgz + + # Install the FVP + case ${fvp} in + corstone300) + ./FVP_Corstone_SSE-300.sh --i-agree-to-the-contained-eula --force --destination ./ --quiet --no-interactive + ;; + corstone320) + ./FVP_Corstone_SSE-320.sh --i-agree-to-the-contained-eula --force --destination ./ --quiet --no-interactive + ;; + *) + echo "[${FUNCNAME[0]}] Error: Unknown FVP model ${fvp}. Exiting." + exit 1 + ;; + esac + done +} + +function check_fvp_eula () { + # Mandatory user arg --i-agree-to-the-contained-eula + eula_acceptance_by_variable="${ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA:-False}" + + if [[ "${eula_acceptance}" -eq 0 ]]; then + if [[ ${eula_acceptance_by_variable} != "True" ]]; then + echo "Must pass argument '--i-agree-to-the-contained-eula' to agree to EULA associated with downloading the FVP." + echo "Alternativly set environment variable ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA=True." + echo "Exiting!" + exit 1 + else + echo "Arm EULA for FVP agreed to with ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA=True environment variable" + fi + fi +} + +function setup_fvp() { + if [[ "${OS}" != "Linux" ]]; then + # Check if FVP is callable + if command -v FVP_Corstone_SSE-300_Ethos-U55 &> /dev/null; then + echo "[${FUNCNAME[0]}] Info: FVP for MacOS seem to be installed. Continuing..." + return 0 # If true exit gracefully and proceed with setup + else + echo "[${FUNCNAME[0]}] Warning: FVP only supported with Linux OS, skipping FVP setup..." + echo "[${FUNCNAME[0]}] Warning: For MacOS, using https://github.com/Arm-Examples/FVPs-on-Mac is recommended." + echo "[${FUNCNAME[0]}] Warning: Follow the instructions and make sure the path is set correctly." + return 1 # Throw error. User need to install FVP according to ^^^ + fi + fi +} + +function setup_path_fvp() { + fvps=("corstone300" "corstone320") + for fvp in "${fvps[@]}"; do + model_dir_variable=${fvp}_model_dir + fvp_model_dir=${!model_dir_variable} + fvp_bin_path="${root_dir}/FVP-${fvp}/models/${fvp_model_dir}" + append_env_in_setup_path PATH ${fvp_bin_path} + done + + # 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 +} diff --git a/backends/arm/scripts/install_models_for_test.sh b/backends/arm/scripts/install_models_for_test.sh index 9c8b034909e..001d733a014 100644 --- a/backends/arm/scripts/install_models_for_test.sh +++ b/backends/arm/scripts/install_models_for_test.sh @@ -5,6 +5,4 @@ # LICENSE file in the root directory of this source tree. set -e - -# Install diffusers for Stable Diffusion model test -pip install "diffusers[torch]==0.33.1" +pip install -r backends/arm/requirements-arm-models-test.txt diff --git a/backends/arm/scripts/install_reference_model.sh b/backends/arm/scripts/install_reference_model.sh deleted file mode 100755 index 089eab899db..00000000000 --- a/backends/arm/scripts/install_reference_model.sh +++ /dev/null @@ -1,35 +0,0 @@ -#!/usr/bin/env bash -# Copyright 2025 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 - -# Installation script to manage transition to 1.0 - -# TOSA reference model -tosa_reference_model_url="https://git.gitlab.arm.com/tosa/tosa-reference-model.git" -tosa_reference_model_1_0_rev="1e6e4526df3391e1d6bc41562596bb18b3153bf3" - -script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) - -source ${script_dir}/utils.sh - - -function setup_tosa_reference_model() { - local work_dir="$1" - - if [[ -z "$work_dir" ]]; then - echo "Error: work_dir parameter is required." - return 1 - fi - - mkdir -p "$work_dir" - pushd "$work_dir" || exit 1 - - # Install the 1.0 branch from upstream - CMAKE_POLICY_VERSION_MINIMUM=3.5 BUILD_PYBIND=1 pip install "tosa-tools@git+${tosa_reference_model_url}@${tosa_reference_model_1_0_rev}" ml_dtypes==0.5.1 --no-dependencies flatbuffers -} - -setup_tosa_reference_model $1 diff --git a/backends/arm/scripts/mlsdk_utils.sh b/backends/arm/scripts/mlsdk_utils.sh index ed6d78c900a..3129676ec4b 100755 --- a/backends/arm/scripts/mlsdk_utils.sh +++ b/backends/arm/scripts/mlsdk_utils.sh @@ -6,8 +6,7 @@ set -euo pipefail -# TODO -mlsdk_manifest_url="" +mlsdk_manifest_url="https://github.com/arm/ai-ml-sdk-manifest.git" script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) @@ -55,8 +54,9 @@ function download_ai_mlsdk_manifest() { function setup_model_converter() { local work_dir="$1" local manifest_dir="$2" - local enable_vgf_lib="$3" - local enable_emulation_layer="$4" + local enable_model_converter="$3" + local enable_vgf_lib="$4" + local enable_emulation_layer="$5" if [[ -z "$work_dir" ]]; then echo "Error: work_dir parameter is required." @@ -76,29 +76,34 @@ function setup_model_converter() { pushd "$manifest_dir" # model-converter - # TODO: Remove macOS patch after mlsdk fully supports macOS - if [[ "$(uname)" == "Darwin" ]]; then + if [[ "${enable_model_converter}" -eq 1 ]]; then + # TODO: Remove this workaround once MLSDK has full Darwin support + # Do not indent sed command, the whitespace is significant for the patch to work. + if [[ "$(uname)" == "Darwin" ]]; then sed -i '' '/^ *print(f"Unsupported host platform/ i\ if system == "Darwin":\ - # Use default Apple toolchain (Clang) on macOS\ return True\ \ ' sw/model-converter/scripts/build.py + fi + python sw/model-converter/scripts/build.py -j$(nproc) fi - python sw/model-converter/scripts/build.py -j$(nproc) # libvgf if [[ "${enable_vgf_lib}" -eq 1 ]]; then - # TODO: Remove macOS patch after mlsdk fully supports macOS + # TODO: Remove this workaround once MLSDK has full Darwin support + # Do not indent sed command, the whitespace is significant for the patch to work. if [[ "$(uname)" == "Darwin" ]]; then sed -i '' '/^ *print(f"ERROR: Unsupported host platform/ i\ if system == "Darwin":\ - # Use default Apple toolchain (Clang) on macOS\ return True\ \ ' sw/vgf-lib/scripts/build.py fi - python sw/vgf-lib/scripts/build.py -j$(nproc) + pushd sw/vgf-lib + python scripts/build.py -j$(nproc) + cmake --install build --prefix deploy + popd fi # emu layer @@ -110,12 +115,38 @@ function setup_model_converter() { -DSPIRV_HEADERS_PATH=../../dependencies/SPIRV-Headers \ -DSPIRV_TOOLS_PATH=../../dependencies/SPIRV-Tools \ -DVULKAN_HEADERS_PATH=../../dependencies/Vulkan-Headers + cmake --build build + cmake --install build --prefix deploy popd fi popd } +function setup_path_model_converter() { + cd "${root_dir}" + model_converter_bin_path="$(cd ${mlsdk_manifest_dir}/sw/model-converter/build && pwd)" + append_env_in_setup_path PATH ${model_converter_bin_path} +} + +function setup_path_vgf_lib() { + cd "${root_dir}" + model_vgf_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/deploy && pwd)" + append_env_in_setup_path PATH ${model_vgf_path}/bin + append_env_in_setup_path LD_LIBRARY_PATH "${model_vgf_path}/lib" + append_env_in_setup_path DYLD_LIBRARY_PATH "${model_vgf_path}/lib" +} + +function setup_path_emulation_layer() { + cd "${root_dir}" + model_emulation_layer_path="$(cd ${mlsdk_manifest_dir}/sw/emulation-layer/ && pwd)" + prepend_env_in_setup_path LD_LIBRARY_PATH "${model_emulation_layer_path}/deploy/lib" + prepend_env_in_setup_path DYLD_LIBRARY_PATH "${model_emulation_layer_path}/deploy/lib" + prepend_env_in_setup_path VK_INSTANCE_LAYERS VK_LAYER_ML_Tensor_Emulation + prepend_env_in_setup_path VK_INSTANCE_LAYERS VK_LAYER_ML_Graph_Emulation + prepend_env_in_setup_path VK_LAYER_PATH "${model_emulation_layer_path}/deploy/share/vulkan/explicit_layer.d" +} + #setup_model_converter() $1 # `"$manifest_dir"' diff --git a/backends/arm/scripts/parse_test_names.py b/backends/arm/scripts/parse_test_names.py index b966cc1e8ca..c6eaafa597b 100644 --- a/backends/arm/scripts/parse_test_names.py +++ b/backends/arm/scripts/parse_test_names.py @@ -8,6 +8,7 @@ CUSTOM_EDGE_OPS = [ "linspace.default", "eye.default", + "expm1.default", "vector_norm.default", "hardsigmoid.default", "hardswish.default", @@ -18,10 +19,13 @@ "bitwise_right_shift.Tensor", "bitwise_left_shift.Tensor", "native_group_norm.default", + "silu.default", + "sdpa.default", "unbind.int", "unflatten.int", "_native_batch_norm_legit_no_training.default", "_native_batch_norm_legit.no_stats", + "alias_copy.default", ] ALL_EDGE_OPS = SAMPLE_INPUT.keys() | CUSTOM_EDGE_OPS diff --git a/backends/arm/scripts/run_fvp.sh b/backends/arm/scripts/run_fvp.sh index e4dfa1d7339..769b2e30282 100755 --- a/backends/arm/scripts/run_fvp.sh +++ b/backends/arm/scripts/run_fvp.sh @@ -18,6 +18,7 @@ _setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly ins elf_file="" +data_file="" target="ethos-u55-128" timeout="600" @@ -25,6 +26,7 @@ help() { echo "Usage: $(basename $0) [options]" echo "Options:" echo " --elf= elf file to run" + echo " --data=@
Place a file in memory at this address, useful to emulate a PTE flashed into memory instead as part of the code." echo " --target= Target to build and run for Default: ${target}" echo " --timeout= Maximum target runtime, used to detect hanging, might need to be higer on large models Default: ${timeout}" exit 0 @@ -34,6 +36,7 @@ for arg in "$@"; do case $arg in -h|--help) help ;; --elf=*) elf_file="${arg#*=}";; + --data=*) data_file="--data ${arg#*=}";; --target=*) target="${arg#*=}";; --timeout=*) timeout="${arg#*=}";; *) @@ -80,34 +83,49 @@ fi log_file=$(mktemp) - - if [[ ${target} == *"ethos-u55"* ]]; then - ${nobuf} ${fvp_model} \ + ${nobuf} ${fvp_model} \ -C ethosu.num_macs=${num_macs} \ -C mps3_board.visualisation.disable-visualisation=1 \ -C mps3_board.telnetterminal0.start_telnet=0 \ -C mps3_board.uart0.out_file='-' \ -C mps3_board.uart0.shutdown_on_eot=1 \ - -a "${elf_file}" \ - --timelimit ${timeout} 2>&1 | tee ${log_file} || true # seconds + -a "${elf_file}" \ + ${data_file} \ + --timelimit ${timeout} 2>&1 | sed 's/\r$//' | tee ${log_file} || true # seconds echo "[${BASH_SOURCE[0]}] Simulation complete, $?" elif [[ ${target} == *"ethos-u85"* ]]; then - ${nobuf} ${fvp_model} \ + ${nobuf} ${fvp_model} \ -C mps4_board.subsystem.ethosu.num_macs=${num_macs} \ -C mps4_board.visualisation.disable-visualisation=1 \ -C vis_hdlcd.disable_visualisation=1 \ -C mps4_board.telnetterminal0.start_telnet=0 \ -C mps4_board.uart0.out_file='-' \ -C mps4_board.uart0.shutdown_on_eot=1 \ - -a "${elf_file}" \ - --timelimit ${timeout} 2>&1 | tee ${log_file} || true # seconds + -a "${elf_file}" \ + ${data_file} \ + --timelimit ${timeout} 2>&1 | sed 's/\r$//' | tee ${log_file} || true # seconds echo "[${BASH_SOURCE[0]}] Simulation complete, $?" else echo "Running ${elf_file} for ${target} is not supported" exit 1 fi +echo "Checking for a etdump in log" +! grep "#\[RUN THIS\]" ${log_file} >/dev/null +if [ $? != 0 ]; then + echo "Found ETDump in log!" + echo "#!/bin/sh" > etdump_script.sh + sed -n '/^#\[RUN THIS\]$/,/^#\[END\]$/p' ${log_file} >> etdump_script.sh + # You can run etdump_script.sh if you do + # $ chmod a+x etdump_script.sh + # $ ./etdump_script.sh + # But lets not trust the script as a bad patch would run bad code on your machine + grep ">etdump.bin" etdump_script.sh | cut -d\" -f2- | cut -d\" -f1 >etdump.base64 + base64 -d etdump.base64 >etdump.bin + python3 -m devtools.inspector.inspector_cli --etdump_path etdump.bin --source_time_scale cycles --target_time_scale cycles +fi + echo "Checking for problems in log:" ! grep -E "^(F|E|\\[critical\\]|Hard fault.|Info: Simulation is stopping. Reason: CPU time has been exceeded.).*$" ${log_file} if [ $? != 0 ]; then diff --git a/backends/arm/scripts/run_vkml.sh b/backends/arm/scripts/run_vkml.sh new file mode 100755 index 00000000000..8a64a937638 --- /dev/null +++ b/backends/arm/scripts/run_vkml.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# Copyright 2025 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. + +# Optional parameter: +# --build_type= "Release" | "Debug" | "RelWithDebInfo" +# --etdump build with devtools-etdump support + +set -eu +set -o pipefail + +script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +et_root_dir=$(cd ${script_dir}/../../.. && pwd) +et_root_dir=$(realpath ${et_root_dir}) +setup_path_script=${et_root_dir}/examples/arm/ethos-u-scratch/setup_path.sh +_setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." + + +model="" +build_path="cmake-out-vkml" +converter="model-converter" + +help() { + echo "Usage: $(basename $0) [options]" + echo "Options:" + echo " --model= .pte model file to run" + echo " --build_path= Path to executor_runner build. for Default: ${build_path}" + exit 0 +} + +for arg in "$@"; do + case $arg in + -h|--help) help ;; + --model=*) model="${arg#*=}";; + --build_path=*) build_path="${arg#*=}";; + *) + ;; + esac +done + +if [[ -z ${model} ]]; then echo "Model name needs to be provided"; exit 1; fi + + +# Source the tools +# This should be prepared by the setup.sh +[[ -f ${setup_path_script} ]] \ + || { echo "Missing ${setup_path_script}. ${_setup_msg}"; exit 1; } + +source ${setup_path_script} + +# basic checks before we get started +hash ${converter} \ + || { echo "Could not find ${converter} on PATH, ${_setup_msg}"; exit 1; } + + + +runner="${build_path}/executor_runner" + +echo "--------------------------------------------------------------------------------" +echo "Running ${model} with ${runner}" +echo "WARNING: The VK_ML layer driver will not provide accurate performance information" +echo "--------------------------------------------------------------------------------" + +# Check if stdbuf is intalled and use stdbuf -oL together with tee below to make the output +# go all the way to the console more directly and not be buffered + +if hash stdbuf 2>/dev/null; then + nobuf="stdbuf -oL" +else + nobuf="" +fi + +log_file=$(mktemp) + + +${nobuf} ${runner} -model_path ${model} | tee ${log_file} +echo "[${BASH_SOURCE[0]}] execution complete, $?" + +# Most of these can happen for bare metal or linx executor_runner runs. +echo "Checking for problems in log:" +! grep -E "^(F|E|\\[critical\\]|Hard fault.|Info: Simulation is stopping. Reason: CPU time has been exceeded.).*$" ${log_file} +if [ $? != 0 ]; then + echo "Found ERROR" + rm "${log_file}" + exit 1 +fi +echo "No problems found!" +rm "${log_file}" diff --git a/backends/arm/scripts/toolchain_utils.sh b/backends/arm/scripts/toolchain_utils.sh new file mode 100644 index 00000000000..161bfd29cd1 --- /dev/null +++ b/backends/arm/scripts/toolchain_utils.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# Copyright 2025 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. + +# Important to check for unset variables since this script is always sourced from setup.sh +set -u + +# Check if the script is being sourced +(return 0 2>/dev/null) +if [[ $? -ne 0 ]]; then + echo "Error: This script must be sourced." + exit 1 +fi + +function gcc_select_toolchain() { + if [[ "${ARCH}" == "x86_64" ]] ; then + toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi" + toolchain_md5_checksum="0601a9588bc5b9c99ad2b56133b7f118" + elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]] ; then + if [[ "${OS}" == "Darwin" ]]; then + toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi" + toolchain_md5_checksum="f1c18320bb3121fa89dca11399273f4e" + elif [[ "${OS}" == "Linux" ]]; then + toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi.tar.xz" + toolchain_dir="arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi" + toolchain_md5_checksum="303102d97b877ebbeb36b3158994b218" + fi + else + # This should never happen, it should be covered by setup.sh but catch it anyway + echo "[gcc_select_toolchain]: Unsupported arch!"; exit 1 + fi +} + +function zephyr_select_toolchain() { + if [[ "${OS}" != "Linux" ]] ; then + echo "[zephyr_select_toolchain] Error: Linux is only supported for zephyr!"; exit 1; + fi + + if [[ "${ARCH}" == "x86_64" ]] ; then + toolchain_url="https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/toolchain_linux-x86_64_arm-zephyr-eabi.tar.xz" + toolchain_dir="arm-zephyr-eabi" + toolchain_md5_checksum="93128be0235cf5cf5f1ee561aa6eac5f" + elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]] ; then + toolchain_url="https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/toolchain_linux-aarch64_arm-zephyr-eabi.tar.xz" + toolchain_dir="arm-zephyr-eabi" + toolchain_md5_checksum="ef4ca56786204439a75270ba800cc64b" + else + # This should never happen, it should be covered by setup.sh but catch it anyway + echo "[zephyr_select_toolchain]: Unsupported arch!"; exit 1 + fi +} + +function select_toolchain() { + if [[ "${target_toolchain}" == "zephyr" ]]; then + zephyr_select_toolchain + else + gcc_select_toolchain + fi + echo "[main] Info selected ${toolchain_dir} for ${ARCH} - ${OS} platform" +} + +function setup_toolchain() { + # Download and install the arm toolchain (default is arm-none-eabi) + # setting --target-toolchain to zephyr sets this to arm-zephyr-eabi + cd "${root_dir}" + if [[ ! -e "${toolchain_dir}.tar.xz" ]]; then + echo "[${FUNCNAME[0]}] Downloading ${toolchain_dir} toolchain ..." + curl --output "${toolchain_dir}.tar.xz" -L "${toolchain_url}" + verify_md5 ${toolchain_md5_checksum} "${toolchain_dir}.tar.xz" || exit 1 + fi + + echo "[${FUNCNAME[0]}] Installing ${toolchain_dir} toolchain ..." + rm -rf "${toolchain_dir}" + tar xf "${toolchain_dir}.tar.xz" +} + +function setup_path_toolchain() { + toolchain_bin_path="$(cd ${toolchain_dir}/bin && pwd)" + append_env_in_setup_path PATH ${toolchain_bin_path} +} diff --git a/backends/arm/scripts/utils.sh b/backends/arm/scripts/utils.sh index 8b4c8d4f96f..d0c1dadbb3e 100644 --- a/backends/arm/scripts/utils.sh +++ b/backends/arm/scripts/utils.sh @@ -4,6 +4,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# Important to check for unset variables since this script is always sourced from setup.sh +set -u + +# Check if the script is being sourced +(return 0 2>/dev/null) +if [[ $? -ne 0 ]]; then + echo "Error: This script must be sourced." + exit 1 +fi + function verify_md5() { # Compare the md5 of a file with a provided expected value. @@ -56,3 +66,53 @@ function patch_repo() { echo -e "[${FUNCNAME[0]}] Patched ${name} @ $(git describe --all --long 2> /dev/null) in ${repo_dir} dir.\n" popd } + +function check_platform_support() { + # No args + # Exits with return code 1 if the platform is unsupported + + # Make sure we are on a supported platform + if [[ "${ARCH}" != "x86_64" ]] && [[ "${ARCH}" != "aarch64" ]] \ + && [[ "${ARCH}" != "arm64" ]]; then + echo "[main] Error: only x86-64 & aarch64 architecture is supported for now!" + exit 1 + fi +} + +function check_os_support() { + # No args + # Exits with return code 1 if invalid combination of platform and os + + # Check valid combinations of OS and platform + + # Linux on x86_64 + if [[ "${ARCH}" == "x86_64" ]] && [[ "${OS}" != "Linux" ]]; then + echo "Error: Only Linux is supported on x86_64" + exit 1 + fi + + # Linux on arm64/aarch64 + # Darwin on arm64/aarch64 + if [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then + if [[ "${OS}" != "Darwin" ]] && [[ "${OS}" != "Linux" ]]; then + echo "Error: Only Linux and Darwin are supported on arm64" + exit 1 + fi + fi +} + +function prepend_env_in_setup_path() { + echo "export $1=$2:\${$1-}" >> ${setup_path_script}.sh + echo "set --path -pgx $1 $2" >> ${setup_path_script}.fish +} + +function append_env_in_setup_path() { + echo "export $1=\${$1-}:$2" >> ${setup_path_script}.sh + echo "set --path -agx $1 $2" >> ${setup_path_script}.fish +} + +function clear_setup_path() { + # Clear setup_path_script + echo "" > "${setup_path_script}.sh" + echo "" > "${setup_path_script}.fish" +} diff --git a/backends/arm/scripts/vulkan_utils.sh b/backends/arm/scripts/vulkan_utils.sh new file mode 100644 index 00000000000..c22d8f26a0c --- /dev/null +++ b/backends/arm/scripts/vulkan_utils.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +# Copyright 2025 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. + +# Important to check for unset variables since this script is always sourced from setup.sh +set -u + +# Check if the script is being sourced +(return 0 2>/dev/null) +if [[ $? -ne 0 ]]; then + echo "Error: This script must be sourced." + exit 1 +fi + +vulkan_sdk_version="1.4.321.1" +vulkan_sdk_base_dir="vulkan_sdk" + +# MLSDK dependencies +mlsdk_manifest_dir="ml-sdk-for-vulkan-manifest" +vulkan_sdk_bin_dir="${vulkan_sdk_base_dir}/${vulkan_sdk_version}/${ARCH}/bin" + + +if [[ "${ARCH}" == "x86_64" ]]; then + # Vulkan SDK + vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/linux/vulkansdk-linux-x86_64-${vulkan_sdk_version}.tar.xz" + vulkan_sdk_sha256="f22a3625bd4d7a32e7a0d926ace16d5278c149e938dac63cecc00537626cbf73" + +elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then + # Vulkan SDK + vulkan_sdk_url="https://github.com/jakoch/vulkan-sdk-arm/releases/download/1.4.321.1/vulkansdk-ubuntu-22.04-arm-1.4.321.1.tar.xz" + vulkan_sdk_sha256="c57e318d0940394d3a304034bb7ddabda788b5b0b54638e80e90f7264efe9f84" +else + echo "[main] Error: only x86-64 & aarch64/arm64 architecture is supported for now!"; exit 1; +fi + +function setup_vulkan_sdk() { + if command -v glslc > /dev/null 2>&1; then + echo "[${FUNCNAME[0]}] GLSL already installed, no need to get Vulkan SDK..." + enable_vulkan_sdk=0 + return + fi + + cd "${root_dir}" + + vulkan_sdk_tar_file="${vulkan_sdk_url##*/}" + if [[ ! -e "${vulkan_sdk_tar_file}" ]]; then + echo "[${FUNCNAME[0]}] Downloading Vulkan SDK - ${vulkan_sdk_url}.." + curl -L --output "${vulkan_sdk_tar_file}" "${vulkan_sdk_url}" + echo "${vulkan_sdk_sha256} ${vulkan_sdk_tar_file}" | sha256sum -c - + rm -fr ${vulkan_sdk_base_dir} + fi + + mkdir -p ${vulkan_sdk_base_dir} + tar -C ${vulkan_sdk_base_dir} -xJf "${vulkan_sdk_tar_file}" + + vulkan_sdk_bin_path="$(cd ${vulkan_sdk_bin_dir} && pwd)" + if ${vulkan_sdk_bin_path}/glslc --version > /dev/null 2>&1; then + echo "[${FUNCNAME[0]}] Vulkan SDK install (GLSL) OK" + else + echo "[${FUNCNAME[0]}] Vulkan SDK install NOK - glslc returned error" + ${vulkan_sdk_bin_path}/glslc --version + exit 1 + fi +} + +function setup_path_vulkan() { + cd "${root_dir}" + vulkan_sdk_bin_path="$(cd ${vulkan_sdk_bin_dir} && pwd)" + append_env_in_setup_path PATH ${vulkan_sdk_bin_path} +} diff --git a/backends/arm/test/TARGETS b/backends/arm/test/TARGETS index 9443547879d..c27d00590f3 100644 --- a/backends/arm/test/TARGETS +++ b/backends/arm/test/TARGETS @@ -1,10 +1,10 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load(":targets.bzl", "define_arm_tests") oncall("executorch") -python_library( +runtime.python_library( name = "conftest", srcs = ["conftest.py"], deps = [ @@ -14,7 +14,7 @@ python_library( ] ) -python_library( +runtime.python_library( name = "runner_utils", srcs = ["runner_utils.py"], deps = [ @@ -25,17 +25,17 @@ python_library( ] ) -python_library( +runtime.python_library( name = "common", srcs = ["common.py"], deps = [ ":runner_utils", - "//executorch/backends/arm:tosa_specification", + "//executorch/backends/arm/tosa:tosa", "fbsource//third-party/pypi/pytest:pytest", ] ) -python_library( +runtime.python_library( name = "arm_tester", srcs = glob(["tester/*.py"]), deps = [ @@ -43,7 +43,8 @@ python_library( "//executorch/backends/xnnpack/test/tester:tester", "//executorch/backends/arm:ethosu_partitioner", "//executorch/backends/arm/quantizer:lib", - "//executorch/backends/arm:tosa_mapping", + "//executorch/backends/arm/tosa:mapping", + "//executorch/backends/arm:vgf_partitioner", "//executorch/devtools/backend_debug:delegation_info", "//executorch/exir/backend:operator_support", "fbsource//third-party/pypi/tabulate:tabulate", diff --git a/backends/arm/test/common.py b/backends/arm/test/common.py index 462098c9b77..608c273b2ef 100644 --- a/backends/arm/test/common.py +++ b/backends/arm/test/common.py @@ -19,8 +19,9 @@ corstone300_installed, corstone320_installed, model_converter_installed, + vkml_emulation_layer_installed, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -62,16 +63,24 @@ def maybe_get_tosa_collate_path() -> str | None: def get_tosa_compile_spec( - tosa_spec: str | TosaSpecification, custom_path=None + tosa_spec: str | TosaSpecification, + custom_path: Optional[str] = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, ) -> list[CompileSpec]: """ Default compile spec for TOSA tests. """ - return get_tosa_compile_spec_unbuilt(tosa_spec, custom_path).build() + return get_tosa_compile_spec_unbuilt( + tosa_spec, + custom_path, + tosa_debug_mode, + ).build() def get_tosa_compile_spec_unbuilt( - tosa_spec: str | TosaSpecification, custom_path=None + tosa_spec: str | TosaSpecification, + custom_path: Optional[str], + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode], ) -> ArmCompileSpecBuilder: """Get the ArmCompileSpecBuilder for the default TOSA tests, to modify the compile spec before calling .build() to finalize it. @@ -81,44 +90,15 @@ def get_tosa_compile_spec_unbuilt( if custom_path is not None: os.makedirs(custom_path, exist_ok=True) + compile_spec_builder = ( ArmCompileSpecBuilder() .tosa_compile_spec(tosa_spec) .dump_intermediate_artifacts_to(custom_path) ) - return compile_spec_builder - - -def get_vgf_compile_spec( - tosa_spec: str | TosaSpecification, - compiler_flags: Optional[str] = "", - custom_path=None, -) -> list[CompileSpec]: - """ - Default compile spec for VGF tests. - """ - return get_vgf_compile_spec_unbuilt(tosa_spec, compiler_flags, custom_path).build() - - -def get_vgf_compile_spec_unbuilt( - tosa_spec: str | TosaSpecification, - compiler_flags: Optional[str] = "", - custom_path=None, -) -> ArmCompileSpecBuilder: - """Get the ArmCompileSpecBuilder for the default VGF tests, to modify - the compile spec before calling .build() to finalize it. - """ - if not custom_path: - custom_path = maybe_get_tosa_collate_path() - - if custom_path is not None: - os.makedirs(custom_path, exist_ok=True) - compile_spec_builder = ( - ArmCompileSpecBuilder() - .vgf_compile_spec(tosa_spec, compiler_flags) - .dump_intermediate_artifacts_to(custom_path) - ) + if tosa_debug_mode is not None: + compile_spec_builder.dump_debug_info(tosa_debug_mode) return compile_spec_builder @@ -129,6 +109,7 @@ def get_u55_compile_spec( memory_mode: str = "Shared_Sram", extra_flags: str = "--debug-force-regor --output-format=raw", custom_path: Optional[str] = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, config: Optional[str] = "Arm/vela.ini", ) -> list[CompileSpec]: """ @@ -140,16 +121,18 @@ def get_u55_compile_spec( memory_mode=memory_mode, extra_flags=extra_flags, custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, config=config, ).build() def get_u85_compile_spec( macs: int = 128, - system_config="Ethos_U85_SYS_DRAM_Mid", - memory_mode="Shared_Sram", - extra_flags="--output-format=raw", - custom_path=None, + system_config: str = "Ethos_U85_SYS_DRAM_Mid", + memory_mode: str = "Shared_Sram", + extra_flags: str = "--output-format=raw", + custom_path: Optional[str] = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, config: Optional[str] = "Arm/vela.ini", ) -> list[CompileSpec]: """ @@ -161,16 +144,32 @@ def get_u85_compile_spec( memory_mode=memory_mode, extra_flags=extra_flags, custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, config=config, ).build() +def get_vgf_compile_spec( + tosa_spec: str | TosaSpecification, + compiler_flags: Optional[str] = "", + custom_path: Optional[str] = "", + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, +) -> list[CompileSpec]: + """ + Default compile spec for VGF tests. + """ + return get_vgf_compile_spec_unbuilt( + tosa_spec, compiler_flags, custom_path, tosa_debug_mode + ).build() + + def get_u55_compile_spec_unbuilt( macs: int, system_config: str, memory_mode: str, extra_flags: str, custom_path: Optional[str], + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode], config: Optional[str], ) -> ArmCompileSpecBuilder: """Get the ArmCompileSpecBuilder for the Ethos-U55 tests, to modify @@ -194,6 +193,10 @@ def get_u55_compile_spec_unbuilt( ) .dump_intermediate_artifacts_to(artifact_path) ) + + if tosa_debug_mode is not None: + compile_spec.dump_debug_info(tosa_debug_mode) + return compile_spec @@ -203,6 +206,7 @@ def get_u85_compile_spec_unbuilt( memory_mode: str, extra_flags: str, custom_path: Optional[str], + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode], config: Optional[str], ) -> list[CompileSpec]: """Get the ArmCompileSpecBuilder for the Ethos-U85 tests, to modify @@ -225,9 +229,44 @@ def get_u85_compile_spec_unbuilt( ) .dump_intermediate_artifacts_to(artifact_path) ) + + if tosa_debug_mode is not None: + compile_spec.dump_debug_info(tosa_debug_mode) + return compile_spec # type: ignore[return-value] +def get_vgf_compile_spec_unbuilt( + tosa_spec: str | TosaSpecification, + compiler_flags: Optional[str], + custom_path: Optional[str], + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode], +) -> ArmCompileSpecBuilder: + """Get the ArmCompileSpecBuilder for the default VGF tests, to modify + the compile spec before calling .build() to finalize it. + """ + if "FP" in repr(tosa_spec): + artifact_path = custom_path or tempfile.mkdtemp(prefix="arm_vgf_fp_") + elif "INT" in repr(tosa_spec): + artifact_path = custom_path or tempfile.mkdtemp(prefix="arm_vgf_int_") + else: + raise ValueError(f"Unsupported vgf compile_spec: {repr(tosa_spec)}") + + if not os.path.exists(artifact_path): + os.makedirs(artifact_path, exist_ok=True) + + compile_spec_builder = ( + ArmCompileSpecBuilder() + .vgf_compile_spec(tosa_spec, compiler_flags) + .dump_intermediate_artifacts_to(artifact_path) + ) + + if tosa_debug_mode is not None: + compile_spec_builder.dump_debug_info(tosa_debug_mode) + + return compile_spec_builder + + XfailIfNoCorstone300 = pytest.mark.xfail( condition=not ( corstone300_installed() and arm_executor_runner_exists("corstone-300") @@ -251,7 +290,14 @@ def get_u85_compile_spec_unbuilt( raises=FileNotFoundError, reason="Did not find model-converter on path", ) -"""Xfails a test if model-converter is not installed""" +"""Skips a test if model-converter is not installed""" + +XfailfNoVKMLEmulationLayer = pytest.mark.xfail( + condition=not (vkml_emulation_layer_installed()), + raises=TypeError, + reason="VKML environment is not set properly or executor_runner path is misused", +) +"""Xfails a test if VKML Emulation Layer is not installed""" xfail_type = str | tuple[str, type[Exception]] diff --git a/backends/arm/test/misc/test_bn_relu_folding_qat.py b/backends/arm/test/misc/test_bn_relu_folding_qat.py index c39c1694d0a..c88c38e869d 100644 --- a/backends/arm/test/misc/test_bn_relu_folding_qat.py +++ b/backends/arm/test/misc/test_bn_relu_folding_qat.py @@ -40,13 +40,17 @@ def forward(self, x: torch.Tensor): models = { - "conv_bn_relu": ConvModule(batch_norm=True), - "conv_relu": ConvModule(batch_norm=False), + # name : (model, is_per_channel) + "conv_bn_relu_per_channel": (ConvModule(batch_norm=True), True), + "conv_relu_per_channel": (ConvModule(batch_norm=False), True), + "conv_bn_relu_per_tensor": (ConvModule(batch_norm=True), False), + "conv_relu_per_tensor": (ConvModule(batch_norm=False), False), } -@common.parametrize("model", models) -def test_qat_tosa_INT(model: torch.nn.Module): +@common.parametrize("test_data", models) +def test_qat_tosa_INT(test_data): + model, per_channel = test_data pipeline = TosaPipelineINT[input_t1](model, model.test_data, [], [], qtol=1) tosa_version = conftest.get_option("tosa_version") tosa_profiles = { @@ -59,7 +63,7 @@ def test_qat_tosa_INT(model: torch.nn.Module): Quantize( quantizer=quantizer, quantization_config=get_symmetric_quantization_config( - is_qat=True, is_per_channel=False + is_qat=True, is_per_channel=per_channel ), is_qat=True, ), diff --git a/backends/arm/test/misc/test_conv_relu_residual_add.py b/backends/arm/test/misc/test_conv_relu_residual_add.py new file mode 100644 index 00000000000..fdd6ec972a6 --- /dev/null +++ b/backends/arm/test/misc/test_conv_relu_residual_add.py @@ -0,0 +1,110 @@ +# Copyright 2025 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 +import torch.nn as nn +from executorch.backends.arm.test import common + +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, +) + + +# Model with Conv1D - ReLU sequence and a residual add. +# Testing the annotation of Conv1D-ReLU(to be fused) and annotation of add. +# ReLU outputs positive numbers and linear outputs positive and negative numbers, so they +# should have different quantisation parameters. If the ReLU gets wrong quantisation parameters(e.g. qmin!=zp) +# because of a shared observer of a following operators(e.g. add), the Conv1D-ReLU sequence is not fused +# and is left in FP32. As a result, the test fails. +class AddDifferentRanges(torch.nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, input_dim): + super().__init__() + self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size) + self.relu = torch.nn.ReLU() + self.linear = nn.Linear(out_channels, out_channels) + + def forward(self, x): + # Permute: (N, T, C) -> (N, C, T) + x = x.permute(0, 2, 1) + x = self.conv1(x) + x = self.relu(x) + x = x.permute(0, 2, 1) + out = x + self.linear(x) + return out + + +input_t = Tuple[torch.Tensor] +model = AddDifferentRanges(in_channels=3, out_channels=16, kernel_size=3, input_dim=10) +model_inputs = (torch.randn(1, 10, 3),) +quant_test_data = { + "per_channel_quantization=true": True, + "per_channel_quantization=false": False, +} + + +def test_tosa_FP(): + pipeline = TosaPipelineFP[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +@common.parametrize("per_channel_quantization", quant_test_data) +def test_tosa_INT(per_channel_quantization): + pipeline = TosaPipelineINT[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + qtol=0, + ) + pipeline.run() + + +@pytest.mark.slow +@common.XfailIfNoCorstone300 +@common.parametrize("per_channel_quantization", quant_test_data) +def test_tosa_u55_INT(per_channel_quantization): + pipeline = EthosU55PipelineINT[input_t]( + model, + model_inputs, + [], + [], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + qtol=0, + ) + pipeline.run() + + +@pytest.mark.slow +@common.XfailIfNoCorstone320 +@common.parametrize("per_channel_quantization", quant_test_data) +def test_tosa_u85_INT(per_channel_quantization): + pipeline = EthosU85PipelineINT[input_t]( + model, + model_inputs, + [], + [], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + qtol=0, + ) + pipeline.run() diff --git a/backends/arm/test/misc/test_debug_feats.py b/backends/arm/test/misc/test_debug_feats.py index 288d5b41615..3e10a9336f9 100644 --- a/backends/arm/test/misc/test_debug_feats.py +++ b/backends/arm/test/misc/test_debug_feats.py @@ -3,15 +3,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import os import shutil import tempfile +from pathlib import Path from typing import Tuple import pytest import torch +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, @@ -188,6 +191,58 @@ def test_collate_tosa_INT_tests(test_data: input_t1): shutil.rmtree("test_collate_tosa_tests", ignore_errors=True) +@common.parametrize("test_data", Linear.inputs) +def test_dump_tosa_debug_json(test_data: input_t1): + with tempfile.TemporaryDirectory() as tmpdir: + pipeline = TosaPipelineINT[input_t1]( + module=Linear(), + test_data=test_data, + aten_op=[], + exir_op=[], + custom_path=tmpdir, + tosa_debug_mode=ArmCompileSpecBuilder.DebugMode.JSON, + ) + + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + json_output_path = Path(tmpdir) / "debug.json" + + # The file should exist + assert json_output_path.exists() + + # Check the file is valid JSON and can be loaded + with json_output_path.open("r") as file: + try: + data = json.load(file) + + # Check it's not empty + assert data + except json.JSONDecodeError: + pytest.fail("Failed to load debug JSON file") + + +@common.parametrize("test_data", Linear.inputs) +def test_dump_tosa_debug_tosa(test_data: input_t1): + with tempfile.TemporaryDirectory() as tmpdir: + pipeline = TosaPipelineINT[input_t1]( + module=Linear(), + test_data=test_data, + aten_op=[], + exir_op=[], + custom_path=tmpdir, + tosa_debug_mode=ArmCompileSpecBuilder.DebugMode.TOSA, + ) + + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + json_output_path = Path(tmpdir) / "debug.json" + + # A JSON file should not be created when TOSA mode used + assert not json_output_path.exists() + + @common.parametrize("test_data", Linear.inputs) def test_dump_tosa_ops(caplog, test_data: input_t1): pipeline = TosaPipelineINT[input_t1](Linear(), test_data, [], []) diff --git a/backends/arm/test/misc/test_debug_hook.py b/backends/arm/test/misc/test_debug_hook.py new file mode 100644 index 00000000000..935f3984403 --- /dev/null +++ b/backends/arm/test/misc/test_debug_hook.py @@ -0,0 +1,183 @@ +# Copyright 2025 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 dataclasses import dataclass +from types import SimpleNamespace + +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder +from executorch.backends.arm.debug.schema import DebugHook, DebugSchema +from executorch.backends.arm.test import common + + +@dataclass +class DebugHookTestCase: + mock_node: SimpleNamespace + tosa_op: str + op_id: int + expected_events: int + num_nodes_traced: int + + +def create_mock_node_1(): + def _get_action_str() -> str: + return "create" + + from_node_2 = SimpleNamespace( + name="convolution", + target="aten.convolution.default", + graph_id=6052414368, + pass_name="ExportedProgram.module()", + action="create", + from_node=[], + _get_action_string=_get_action_str, + ) + + from_node_1 = SimpleNamespace( + name="convolution", + target="aten.convolution.default", + graph_id=5705954832, + pass_name="Interpreter_PropagateUnbackedSymInts", + action="create", + from_node=[from_node_2], + _get_action_string=_get_action_str, + ) + + fx_node_mock = SimpleNamespace( + name="aten_convolution_default", + target="aten.convolution.default", + meta={ + "stack_trace": 'File "models/model.py", line 221, in forward\nreturn self.features(x)', + "nn_module_stack": {"__self__": ["", "model.Model"]}, + "torch_fn": ("conv2d", "builtin_function_or_method.conv2d"), + "from_node": [from_node_1], + }, + ) + + return fx_node_mock + + +def create_mock_node_2(): + def _get_action_str() -> str: + return "create" + + from_node_1 = SimpleNamespace( + name="convolution", + target="aten.convolution.default", + graph_id=5705954832, + pass_name="Interpreter_PropagateUnbackedSymInts", + action="create", + from_node=[], + _get_action_string=_get_action_str, + ) + + fx_node_mock = SimpleNamespace( + name="aten_convolution_default", + target="aten.convolution.default", + meta={ + "from_node": [from_node_1], + }, + ) + + return fx_node_mock + + +def create_mock_node_3(): + fx_node_mock = SimpleNamespace( + name="aten_convolution_default", + target="aten.convolution.default", + meta={ + "from_node": [], + }, + ) + + return fx_node_mock + + +def _compare_tosa_and_schema(debug_event: DebugSchema, tosa_op): + tosa_info = debug_event.tosa_info + + assert tosa_info.node_name == tosa_op + + # The mapping between op_ids to operator names could change + # So just check operator_name is a string + assert isinstance(tosa_info.operator_name, str) + + +def _compare_node_and_schema(debug_event: DebugSchema, mocked_node): + # Check aten info + aten_info = debug_event.aten_info + + assert aten_info.node_name == mocked_node.name + assert aten_info.operator_name == mocked_node.target + + # Check torch info + torch_info = debug_event.torch_info + + if "nn_module_stack" in mocked_node.meta: + assert torch_info.nn_module_stack == mocked_node.meta["nn_module_stack"] + else: + assert torch_info.nn_module_stack == "No module stack trace available" + + if "stack_trace" in mocked_node.meta: + assert torch_info.stack_trace == mocked_node.meta["stack_trace"].split("\n") + else: + assert torch_info.stack_trace == ["No stack trace available"] + + if "torch_fn" in mocked_node.meta: + assert torch_info.torch_fn == mocked_node.meta["torch_fn"] + else: + assert torch_info.torch_fn == "No torch_fn available" + + +TESTCASES = { + "mocked_node": DebugHookTestCase( + mock_node=create_mock_node_1(), + tosa_op="layer-1", + op_id=3, + expected_events=1, + num_nodes_traced=2, + ), + "mocked_node_partially_empty": DebugHookTestCase( + mock_node=create_mock_node_2(), + tosa_op="layer-1", + op_id=1, + expected_events=1, + num_nodes_traced=1, + ), + "mocked_node_all_empty": DebugHookTestCase( + mock_node=create_mock_node_3(), + tosa_op="layer-2", + op_id=1, + expected_events=1, + num_nodes_traced=0, + ), +} + + +@common.parametrize("test_data", TESTCASES) +def test_debug_hook_add_json(test_data: DebugHookTestCase): + hook = DebugHook(ArmCompileSpecBuilder.DebugMode.JSON) + hook.add(test_data.mock_node, test_data.tosa_op, test_data.op_id) + + debug_events = hook._debug_events + assert len(debug_events) == test_data.expected_events + assert len(debug_events[0].torch_info.node_trace) == test_data.num_nodes_traced + + _compare_tosa_and_schema(debug_events[0], test_data.tosa_op) + _compare_node_and_schema(debug_events[0], test_data.mock_node) + + +@common.parametrize("test_data", TESTCASES) +def test_debug_hook_add_tosa(test_data: DebugHookTestCase): + hook = DebugHook(ArmCompileSpecBuilder.DebugMode.TOSA) + hook.add(test_data.mock_node, test_data.tosa_op, test_data.op_id) + + debug_events = hook._debug_events + assert len(debug_events) == test_data.expected_events + assert len(debug_events[0].torch_info.node_trace) == test_data.num_nodes_traced + + assert debug_events[0].tosa_info is None + + _compare_node_and_schema(debug_events[0], test_data.mock_node) diff --git a/backends/arm/test/misc/test_dim_order_guards.py b/backends/arm/test/misc/test_dim_order_guards.py index b291aaa52cf..80a3c014abc 100644 --- a/backends/arm/test/misc/test_dim_order_guards.py +++ b/backends/arm/test/misc/test_dim_order_guards.py @@ -22,7 +22,7 @@ class Conv2D(torch.nn.Module): inputs: dict[str, input_t1] = { - "randn": (torch.randn(1, 2, 20, 20),), + "randn": (torch.randn(1, 2, 20, 20).to(memory_format=torch.channels_last),), } def __init__(self): @@ -30,7 +30,7 @@ def __init__(self): self.conv2d = torch.nn.Conv2d(in_channels=2, out_channels=3, kernel_size=(3, 3)) def forward(self, x): - return self.conv2d(x.to(memory_format=torch.channels_last)) + return self.conv2d(x) @common.parametrize("test_data", Conv2D.inputs) diff --git a/backends/arm/test/misc/test_extract_io_params_tosa.py b/backends/arm/test/misc/test_extract_io_params_tosa.py new file mode 100644 index 00000000000..da471b0bb74 --- /dev/null +++ b/backends/arm/test/misc/test_extract_io_params_tosa.py @@ -0,0 +1,92 @@ +# Copyright 2025 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 copy + +import pytest +import torch +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder +from executorch.backends.arm.quantizer import VgfQuantizer +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_quantization_config, + TOSAQuantizer, +) + +from executorch.backends.arm.test.common import SkipIfNoModelConverter +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.vgf import VgfPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.passes.quantize_io_pass import extract_io_quant_params +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +class SimpleAdd(torch.nn.Module): + def forward(self, x, y): + return x + y + + +@pytest.mark.parametrize( + "builder_method, quantizer_cls, partitioner_cls", + [ + ("tosa_compile_spec", TOSAQuantizer, TOSAPartitioner), + pytest.param( + "vgf_compile_spec", + VgfQuantizer, + VgfPartitioner, + marks=SkipIfNoModelConverter, + id="VGF", + ), + ], +) +def test_roundtrip_extracts_io_params(builder_method, quantizer_cls, partitioner_cls): + """ + Validates that IO quantization parameters round-trip for both flows. + """ + example_inputs = ( + torch.ones(1, 5), + torch.full((1, 5), 2.0), + ) + mod = SimpleAdd().eval() + + base_spec = TosaSpecification.create_from_string("TOSA-1.0+INT") + compile_spec = getattr(ArmCompileSpecBuilder(), builder_method)( + tosa_spec=base_spec + ).build() + + quantizer = quantizer_cls(compile_spec) + operator_config = get_symmetric_quantization_config(is_qat=True) + quantizer.set_global(operator_config) + + exported = torch.export.export(mod, copy.deepcopy(example_inputs), strict=True) + prepared = prepare_pt2e(exported.module(), quantizer) + _ = prepared(*example_inputs) + + converted = convert_pt2e(prepared) + final_export = torch.export.export(converted, example_inputs, strict=True) + partitioner = partitioner_cls(compile_spec) + edge_prog = to_edge_transform_and_lower(final_export, partitioner=[partitioner]) + + # Extract IO quantization parameters + q = extract_io_quant_params( + edge_prog, + input_idxs=(0, 1), + output_idxs=(0,), + ) + + assert "inputs" in q + assert "outputs" in q + assert len(q["inputs"]) == 2 + assert len(q["outputs"]) == 1 + + for name, params in q["inputs"].items(): + assert isinstance(name, str) + assert isinstance(params["scale"], float) + assert isinstance(params["zero_point"], int) + + out_name, out_params = next(iter(q["outputs"].items())) + assert isinstance(out_name, str) + assert isinstance(out_params["scale"], float) + assert isinstance(out_params["zero_point"], int) diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py new file mode 100644 index 00000000000..d6d6d6cb39c --- /dev/null +++ b/backends/arm/test/misc/test_int64.py @@ -0,0 +1,116 @@ +# Copyright 2025 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 + + +class ConstAdd(torch.nn.Module): + def __init__(self, dtype: torch.dtype, bias=0): + super().__init__() + self.dtype = dtype + self.bias = bias + + def forward(self, x: torch.Tensor): + c = torch.arange(self.bias, self.bias + 10, 1, dtype=self.dtype) + # Add explicit float cast to make quantization work, will be inserted by type promotion otherwise. + return x + c.to(torch.float32) + + +class BufferAdd(torch.nn.Module): + def __init__(self, dtype: torch.dtype, bias=0): + super().__init__() + self.dtype = dtype + self.buffer = torch.arange(0, 10, 1, dtype=self.dtype) + bias + self.bias = bias + + def forward(self, x: torch.Tensor): + c = self.buffer + # Add explicit float cast to make quantization work, will be inserted by type promotion otherwise. + return x + c.to(torch.float32) + + +class ConstChainAdd(torch.nn.Module): + def __init__(self, dtype: torch.dtype): + super().__init__() + self.dtype = dtype + + def forward(self, x: torch.Tensor): + c = torch.arange(0, 10, 1, dtype=self.dtype).reshape((2, 5)).unsqueeze(-1) + # Add explicit float cast to make quantization work, will be inserted by type promotion otherwise. + return x + c.to(torch.float32) + + +class BufferChainAdd(torch.nn.Module): + def __init__(self, dtype: torch.dtype): + super().__init__() + self.dtype = dtype + self.buffer = torch.arange(0, 10, 1, dtype=self.dtype) + + def forward(self, x: torch.Tensor): + c = self.buffer.reshape((2, 5)).unsqueeze(-1) + # Add explicit float cast to make quantization work, will be inserted by type promotion otherwise. + return x + c.to(torch.float32) + + +test_data_suite = { + "fp32_in+int64_buffer": (BufferAdd(torch.int64), (torch.rand(10) - 0.5,)), + "fp32_in+int64_buffer_overflow": ( + BufferAdd(torch.int64, 2**40), + (torch.rand(10) - 0.5,), + ), + "fp32_in+int64_const": (ConstAdd(torch.int64), (torch.rand(10) - 0.5,)), + "fp32_in+int64_const_overflow": ( + ConstAdd(torch.int64, 2**40), + (torch.rand(10) - 0.5,), + ), + "int64_in+float_const": ( + ConstAdd(torch.float32), + (torch.randint(0, 10, (10,)),), + ), + "fp32_in+int64_buffer_chain": ( + BufferChainAdd(torch.int64), + (torch.rand(2, 5, 3) - 0.5,), + ), + "fp32_in+int64_const_chain": ( + ConstChainAdd(torch.int64), + (torch.rand(2, 5, 3) - 0.5,), + ), + "int64_in+float_const_chain": ( + ConstChainAdd(torch.float32), + (torch.randint(0, 10, (2, 5, 3)),), + ), +} + + +@common.parametrize("test_data", test_data_suite) +def test_int64_tosa_FP(test_data: Tuple): + model, inputs = test_data + ( + ArmTester( + model, + inputs, + common.get_tosa_compile_spec("TOSA-1.0+FP", custom_path="tosa/int64"), + ) + .export() + .to_edge_transform_and_lower() + .to_executorch() + .run_method_and_compare_outputs(inputs) + ) + + +@common.parametrize("test_data", test_data_suite) +def test_int64_tosa_INT(test_data: Tuple): + model, inputs = test_data + ( + ArmTester(model, inputs, common.get_tosa_compile_spec("TOSA-1.0+INT")) + .quantize() + .export() + .to_edge_transform_and_lower() + .to_executorch() + .run_method_and_compare_outputs(inputs) + ) diff --git a/backends/arm/test/misc/test_outputs_order.py b/backends/arm/test/misc/test_outputs_order.py new file mode 100644 index 00000000000..43d35b6d13c --- /dev/null +++ b/backends/arm/test/misc/test_outputs_order.py @@ -0,0 +1,119 @@ +# Copyright 2025 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-unsafe +import tempfile +from pathlib import Path + +import pytest +import torch +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.exir import to_edge_transform_and_lower +from torch import nn +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from tosa import TosaGraph + + +class Network(nn.Module): + def __init__(self, batch_norm=False): + super().__init__() + self.conv2d_0 = nn.Sequential( + nn.Conv2d(1, 8, 3, padding=1, bias=False), + nn.BatchNorm2d(8) if batch_norm else nn.Identity(), + nn.ReLU(), + ) + self.conv2d_1 = nn.Sequential( + nn.Conv2d(8, 8, 3, padding=1, bias=False), + nn.BatchNorm2d(8) if batch_norm else nn.Identity(), + nn.ReLU(), + ) + self.conv2d_2 = nn.Sequential( + nn.Conv2d(8, 8, 3, padding=1, bias=False), + nn.BatchNorm2d(8) if batch_norm else nn.Identity(), + nn.ReLU(), + ) + self.out_0 = nn.Sequential(nn.Conv2d(8, 1, 3, padding=1, bias=False), nn.ReLU()) + self.out_1 = nn.Sequential(nn.Conv2d(8, 2, 3, padding=1, bias=False), nn.ReLU()) + self.out_2 = nn.Sequential(nn.Conv2d(8, 3, 3, padding=1, bias=False), nn.ReLU()) + + def forward(self, x): + x = self.conv2d_0(x) + x = self.conv2d_1(x) + x = self.conv2d_2(x) + out0 = self.out_0(x) + out1 = self.out_1(x) + out2 = self.out_2(x) + return out0, out1, out2 + + +def _read_tosa_outputs(tosa_path: Path): + # Find output tensor names in order and return shapes + buf = tosa_path.read_bytes() + buf_arr = bytearray(buf) + graph = TosaGraph.TosaGraph.GetRootAsTosaGraph(buf_arr, 0) + region = graph.Regions(0) + block = region.Blocks(0) + # Build a dict name - tensor‑shape + tensors = {} + for i in range(block.TensorsLength()): + t = block.Tensors(i) + name = t.Name().decode() + # NHWC + shape = [t.Shape(j) for j in range(t.ShapeLength())] + tensors[name] = shape + shapes = [] + for i in range(block.OutputsLength()): + out_name = block.Outputs(i).decode() + shapes.append(tensors[out_name]) + return shapes + + +@pytest.mark.parametrize("batch_size", [1, 4]) +def test_network_output_order_and_restore(tmp_path, batch_size): + model = Network(batch_norm=True).eval() + # Prepare spec + spec = TosaSpecification.create_from_string("TOSA-1.0+INT") + compile_spec = ArmCompileSpecBuilder().tosa_compile_spec(tosa_spec=spec).build() + # Setup quantizer + quantizer = TOSAQuantizer(compile_spec) + quantizer.set_global( + get_symmetric_quantization_config(is_qat=True, is_per_channel=False) + ) + # Trace the model + dummy = torch.randn(batch_size, 1, 28, 28) + fx_mod = torch.export.export_for_training(model, (dummy,)).module() + model = prepare_pt2e(fx_mod, quantizer) + model(dummy) + model = convert_pt2e(model) + # Export to aten dialect + aten_gm = torch.export.export(model, args=(dummy,), strict=True) + with tempfile.TemporaryDirectory() as tmpdir: + art_dir = Path(tmpdir) + part = TOSAPartitioner( + ArmCompileSpecBuilder() + .tosa_compile_spec(spec) + .dump_intermediate_artifacts_to(str(art_dir)) + .build() + ) + _ = to_edge_transform_and_lower(aten_gm, partitioner=[part]) + # Expect exactly one .tosa file in the artefact dir + tosa_files = list(art_dir.glob("*.tosa")) + assert ( + len(tosa_files) == 1 + ), f"Expected 1 .tosa artefact, found {len(tosa_files)} in {art_dir}" + out_shapes = _read_tosa_outputs(tosa_files[0]) + # We use shape that is unique to output to check + # that we preserve output order + channel_dims = [s[1] for s in reversed(out_shapes)] + assert channel_dims == [1, 2, 3], ( + "Outputs in .tosa do not keep author order: " + f"expected [1, 2, 3], got {channel_dims}" + ) diff --git a/backends/arm/test/misc/test_partition_decomposed_quantized_ops.py b/backends/arm/test/misc/test_partition_decomposed_quantized_ops.py index 1aaa2950337..04ecd57e7b1 100644 --- a/backends/arm/test/misc/test_partition_decomposed_quantized_ops.py +++ b/backends/arm/test/misc/test_partition_decomposed_quantized_ops.py @@ -38,7 +38,7 @@ ] linear_residual_exir_op: list[str] = [ "executorch_exir_dialects_edge__ops_aten_gelu_default", - "executorch_exir_dialects_edge__ops_aten_clone_default", + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", "executorch_exir_dialects_edge__ops_aten_linear_default", "executorch_exir_dialects_edge__ops_aten_add_Tensor", ] diff --git a/backends/arm/test/misc/test_tosa_spec.py b/backends/arm/test/misc/test_tosa_spec.py index 66f7dcf0745..968512f54c6 100644 --- a/backends/arm/test/misc/test_tosa_spec.py +++ b/backends/arm/test/misc/test_tosa_spec.py @@ -5,9 +5,11 @@ import unittest -from executorch.backends.arm.arm_backend import get_tosa_spec - -from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification +from executorch.backends.arm.tosa.specification import ( + get_tosa_spec, + Tosa_1_00, + TosaSpecification, +) from executorch.exir.backend.compile_spec_schema import CompileSpec from parameterized import parameterized # type: ignore[import-untyped] diff --git a/backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py b/backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py index 9561e2132ee..0e99f3f5bfa 100644 --- a/backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py +++ b/backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py @@ -6,8 +6,13 @@ import unittest +import pytest import torch -from executorch.backends.arm._passes import InsertCastForOpsWithInt64InputPass +from executorch.backends.arm._passes import ( + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, + InsertInt32CastsAfterInt64PlaceholdersPass, +) from executorch.backends.arm.test import common from executorch.backends.arm.test.models.stable_diffusion.stable_diffusion_module_test_configs import ( @@ -28,13 +33,9 @@ class TestCLIPTextModelWithProjection(unittest.TestCase): # for that is some assert ops are removed by passes in the # .to_executorch step, i.e. after Arm partitioner. ops_after_partitioner = { - "executorch_exir_dialects_edge__ops_aten__to_copy_default": 3, "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, - "executorch_exir_dialects_edge__ops_aten_index_Tensor": 1, - "executorch_exir_dialects_edge__ops_aten_lt_Tensor": 1, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 2, - "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, - "torch.ops.higher_order.executorch_call_delegate": 3, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, } def _prepare_inputs( @@ -60,7 +61,7 @@ def prepare_model_and_inputs(self): return text_encoder_model, text_encoder_model_inputs - def test_CLIPTextModelWithProjection_tosa_MI(self): + def test_CLIPTextModelWithProjection_tosa_FP(self): text_encoder_model, text_encoder_model_inputs = self.prepare_model_and_inputs() with torch.no_grad(): ( @@ -68,7 +69,11 @@ def test_CLIPTextModelWithProjection_tosa_MI(self): text_encoder_model, example_inputs=text_encoder_model_inputs, compile_spec=common.get_tosa_compile_spec(tosa_spec="TOSA-1.0+FP"), - transform_passes=[InsertCastForOpsWithInt64InputPass()], + transform_passes=[ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + InsertInt32CastsAfterInt64PlaceholdersPass(), + ], ) .export() .to_edge_transform_and_lower() @@ -80,9 +85,7 @@ def test_CLIPTextModelWithProjection_tosa_MI(self): ) ) - # MLETORCH-867, MLETORCH-1059 - # Failures: "Fatal Python error: Aborted, Dependency cycles, KeyError in CastInt64BuffersToInt32Pass") - @unittest.expectedFailure + @pytest.mark.xfail(raises=AssertionError, reason="Output difference.") def test_CLIPTextModelWithProjection_tosa_INT(self): text_encoder_model, text_encoder_model_inputs = self.prepare_model_and_inputs() with torch.no_grad(): diff --git a/backends/arm/test/models/stable_diffusion/test_SD3Transformer2DModel.py b/backends/arm/test/models/stable_diffusion/test_SD3Transformer2DModel.py index 880dc17166d..f9d814d044b 100644 --- a/backends/arm/test/models/stable_diffusion/test_SD3Transformer2DModel.py +++ b/backends/arm/test/models/stable_diffusion/test_SD3Transformer2DModel.py @@ -22,11 +22,8 @@ class TestSD3Transformer2DModel(unittest.TestCase): SD3Transformer2DModel is the transformer model used by Stable Diffusion 3.5 Medium """ - # Adjust nbr below as we increase op support. Note: most of the delegates - # calls are directly consecutive to each other in the .pte. The reason - # for that is some assert ops are removed by passes in the - # .to_executorch step, i.e. after Arm partitioner. - ops_after_partitioner = { + # Adjust nbr below as we increase op support. + ops_after_partitioner_FP = { "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 2, @@ -34,6 +31,13 @@ class TestSD3Transformer2DModel(unittest.TestCase): "torch.ops.higher_order.executorch_call_delegate": 1, } + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 2, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, + } + def _prepare_inputs( self, batch_size=2, @@ -102,7 +106,7 @@ def test_SD3Transformer2DModel_tosa_FP(self): ) .export() .to_edge_transform_and_lower() - .check_count(self.ops_after_partitioner) + .check_count(self.ops_after_partitioner_FP) .to_executorch() .run_method_and_compare_outputs( inputs=sd35_transformer2D_model_inputs, @@ -125,7 +129,7 @@ def test_SD3Transformer2DModel_tosa_INT(self): .quantize() .export() .to_edge_transform_and_lower() - .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_count(self.ops_after_partitioner_INT) .to_executorch() .run_method_and_compare_outputs( inputs=sd35_transformer2D_model_inputs, diff --git a/backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py b/backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py index aba58379a92..22a47042eb1 100644 --- a/backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py +++ b/backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py @@ -7,7 +7,11 @@ import unittest import torch -from executorch.backends.arm._passes import InsertCastForOpsWithInt64InputPass +from executorch.backends.arm._passes import ( + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, + InsertInt32CastsAfterInt64PlaceholdersPass, +) from executorch.backends.arm.test import common from executorch.backends.arm.test.models.stable_diffusion.stable_diffusion_module_test_configs import ( @@ -23,22 +27,13 @@ class TestT5EncoderModel(unittest.TestCase): T5EncoderModel is one of the text_encoder used by Stable Diffusion 3.5 Medium """ - # Adjust nbr below as we increase op support. Note: most of the delegates - # calls are directly consecutive to each other in the .pte. The reason - # for that is some assert ops are removed by passes in the - # .to_executorch step, i.e. after Arm partitioner. - ops_after_partitioner = { - "executorch_exir_dialects_edge__ops_aten__to_copy_default": 2, - "executorch_exir_dialects_edge__ops_aten_abs_default": 1, - "executorch_exir_dialects_edge__ops_aten_add_Tensor": 3, - "executorch_exir_dialects_edge__ops_aten_full_like_default": 1, - "executorch_exir_dialects_edge__ops_aten_gt_Scalar": 1, - "executorch_exir_dialects_edge__ops_aten_lt_Scalar": 1, - "executorch_exir_dialects_edge__ops_aten_minimum_default": 1, - "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1, - "executorch_exir_dialects_edge__ops_aten_sub_Tensor": 1, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_aten_where_self": 1, + # Adjust nbr below as we increase op support. + ops_after_partitioner_FP = { + "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_INT = { "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 3, "torch.ops.higher_order.executorch_call_delegate": 3, } @@ -66,7 +61,7 @@ def prepare_model_and_inputs(self): return t5_encoder_model, t5_encoder_model_inputs - def test_T5EncoderModel_tosa_MI(self): + def test_T5EncoderModel_tosa_FP(self): t5_encoder_model, t5_encoder_model_inputs = self.prepare_model_and_inputs() with torch.no_grad(): ( @@ -74,12 +69,16 @@ def test_T5EncoderModel_tosa_MI(self): t5_encoder_model, example_inputs=t5_encoder_model_inputs, compile_spec=common.get_tosa_compile_spec(tosa_spec="TOSA-1.0+FP"), - transform_passes=[InsertCastForOpsWithInt64InputPass()], + transform_passes=[ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + InsertInt32CastsAfterInt64PlaceholdersPass(), + ], ) .export() .to_edge_transform_and_lower() .dump_operator_distribution() - .check_count(self.ops_after_partitioner) + .check_count(self.ops_after_partitioner_FP) .to_executorch() .run_method_and_compare_outputs( inputs=t5_encoder_model_inputs, @@ -99,6 +98,7 @@ def test_T5EncoderModel_tosa_INT(self): .export() .to_edge_transform_and_lower() .dump_operator_distribution() + .check_count(self.ops_after_partitioner_INT) .to_executorch() .run_method_and_compare_outputs( inputs=t5_encoder_model_inputs, diff --git a/backends/arm/test/models/stable_diffusion/test_vae_AutoencoderKL.py b/backends/arm/test/models/stable_diffusion/test_vae_AutoencoderKL.py index cab4ca53d9c..ab0f4892fb8 100644 --- a/backends/arm/test/models/stable_diffusion/test_vae_AutoencoderKL.py +++ b/backends/arm/test/models/stable_diffusion/test_vae_AutoencoderKL.py @@ -41,7 +41,7 @@ def forward(self, *args, **kwargs): return auto_encoder_model, auto_encoder_model_inputs - def test_AutoencoderKL_tosa_MI(self): + def test_AutoencoderKL_tosa_FP(self): auto_encoder_model, auto_encoder_model_inputs = self.prepare_model_and_inputs() with torch.no_grad(): ( diff --git a/backends/arm/test/models/test_conformer.py b/backends/arm/test/models/test_conformer.py index e3b9bc21ebf..3119145aef1 100644 --- a/backends/arm/test/models/test_conformer.py +++ b/backends/arm/test/models/test_conformer.py @@ -15,6 +15,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from torchaudio.models import Conformer @@ -64,7 +65,7 @@ def test_conformer_tosa_INT(): pipeline = TosaPipelineINT[input_t]( TestConformer.conformer, TestConformer.model_example_inputs, - aten_op=TestConformer.aten_ops, + aten_op=[], # RemoveGraphAssertsPass is added in transform_for_annotation_pipeline to remove the assert ops exir_op=[], use_to_edge_transform_and_lower=True, ) @@ -124,3 +125,40 @@ def test_conformer_u85_INT(): atol=5.0, ) pipeline.run() + + +@common.SkipIfNoModelConverter +def test_conformer_vgf_INT(): + pipeline = VgfPipeline[input_t]( + TestConformer.conformer, + TestConformer.model_example_inputs, + aten_op=[], # RemoveGraphAssertsPass is added in transform_for_annotation_pipeline to remove the assert ops + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", + # get_test_inputs( + # TestConformer.dim, TestConformer.lengths, TestConformer.num_examples + # ), + # rtol=1.0, + # atol=3.0, + # ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_conformer_vgf_FP(): + pipeline = VgfPipeline[input_t]( + TestConformer.conformer, + TestConformer.model_example_inputs, + aten_op=TestConformer.aten_ops, + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + 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 4d7f8c925f2..22685a079bd 100644 --- a/backends/arm/test/models/test_deit_tiny_arm.py +++ b/backends/arm/test/models/test_deit_tiny_arm.py @@ -11,9 +11,12 @@ import torch +from executorch.backends.arm.test import common + from executorch.backends.arm.test.tester.test_pipeline import ( TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD @@ -56,3 +59,31 @@ def test_deit_tiny_tosa_INT(): qtol=1, ) pipeline.run() + + +@common.SkipIfNoModelConverter +def test_deit_tiny_vgf_INT(): + pipeline = VgfPipeline[input_t]( + deit_tiny, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + atol=1.5, + qtol=1, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_deit_tiny_vgf_FP(): + pipeline = VgfPipeline[input_t]( + deit_tiny, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_dl3_arm.py b/backends/arm/test/models/test_dl3_arm.py index 433948d15b0..2000ac34794 100644 --- a/backends/arm/test/models/test_dl3_arm.py +++ b/backends/arm/test/models/test_dl3_arm.py @@ -16,6 +16,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from executorch.examples.models import deeplab_v3 @@ -87,3 +88,37 @@ def test_dl3_u85_INT(): "run_method_and_compare_outputs", rtol=1.0, atol=1.0 ) # TODO: MLETORCH-1036 decrease tolerance pipeline.run() + + +@common.SkipIfNoModelConverter +def test_dl3_vgf_INT(): + pipeline = VgfPipeline[input_t]( + TestDl3.dl3, + TestDl3.model_example_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", rtol=1.0, atol=1.0 + # ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_dl3_vgf_FP(): + pipeline = VgfPipeline[input_t]( + TestDl3.dl3, + TestDl3.model_example_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", rtol=1.0, atol=1.0 + # ) + pipeline.run() diff --git a/backends/arm/test/models/test_inception_v3_arm.py b/backends/arm/test/models/test_inception_v3_arm.py new file mode 100644 index 00000000000..f973521c1fa --- /dev/null +++ b/backends/arm/test/models/test_inception_v3_arm.py @@ -0,0 +1,121 @@ +# Copyright 2025 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 common +import pytest + +import torch + +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +from torchvision import models, transforms + +ic3 = models.inception_v3(weights=models.Inception_V3_Weights) +ic3 = ic3.eval() + +# Normalization values referenced from here: +# https://docs.pytorch.org/vision/main/models/generated/torchvision.models.quantization.inception_v3.html +normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + +model_inputs = (normalize(torch.rand(1, 3, 224, 224)),) +input_t = Tuple[torch.Tensor] + + +@pytest.mark.slow +def test_ic3_tosa_FP(): + pipeline = TosaPipelineFP[input_t]( + ic3, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +@pytest.mark.slow +def test_ic3_tosa_BI(): + pipeline = TosaPipelineINT[input_t]( + ic3, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=0.65, + qtol=1, + ) + pipeline.run() + + +@pytest.mark.slow +@pytest.mark.skip(reason="Takes too long to run on CI") +@common.XfailIfNoCorstone300 +def test_ic3_u55_BI(): + pipeline = EthosU55PipelineINT[input_t]( + ic3, + model_inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + atol=0.6, + qtol=1, + ) + pipeline.run() + + +@pytest.mark.slow +@pytest.mark.skip(reason="Takes too long to run on CI") +@common.XfailIfNoCorstone320 +def test_ic3_u85_BI(): + pipeline = EthosU85PipelineINT[input_t]( + ic3, + model_inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + atol=0.6, + qtol=1, + ) + pipeline.run() + + +@pytest.mark.slow +@pytest.mark.skip(reason="Takes too long to run on CI") +@common.SkipIfNoModelConverter +def test_ic3_vgf_FP(): + pipeline = VgfPipeline[input_t]( + ic3, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +@pytest.mark.slow +@pytest.mark.skip(reason="Takes too long to run on CI") +@common.SkipIfNoModelConverter +def test_ic3_vgf_INT(): + pipeline = VgfPipeline[input_t]( + ic3, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_llama.py b/backends/arm/test/models/test_llama.py index ee9750f853c..d47398be3b0 100644 --- a/backends/arm/test/models/test_llama.py +++ b/backends/arm/test/models/test_llama.py @@ -15,12 +15,13 @@ import pytest import torch -from executorch.backends.arm._passes import InsertCastForOpsWithInt64InputPass +from executorch.backends.arm._passes import InsertInt32CastsAfterInt64PlaceholdersPass -from executorch.backends.arm.test import conftest +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from executorch.examples.models.llama.export_llama_lib import ( build_args_parser, @@ -111,7 +112,7 @@ def test_llama_tosa_FP(): aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, - transform_passes=[InsertCastForOpsWithInt64InputPass()], + transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) pipeline.run() @@ -131,3 +132,42 @@ def test_llama_tosa_INT(): use_to_edge_transform_and_lower=True, ) pipeline.run() + + +@common.SkipIfNoModelConverter +def test_llama_vgf_FP(): + llama_model, llama_inputs, llama_meta = TestLlama().prepare_model() + + if llama_model is None or llama_inputs is None: + pytest.skip("Missing model and/or input files") + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + llama_model, + llama_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_llama_vgf_INT(): + llama_model, llama_inputs, llama_meta = TestLlama().prepare_model() + + if llama_model is None or llama_inputs is None: + pytest.skip("Missing model and/or input files") + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + llama_model, + llama_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_lstm_arm.py b/backends/arm/test/models/test_lstm_arm.py index bb9b92a0f7d..1e63472f5f4 100644 --- a/backends/arm/test/models/test_lstm_arm.py +++ b/backends/arm/test/models/test_lstm_arm.py @@ -13,6 +13,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from torch.nn.quantizable.modules import rnn @@ -98,3 +99,37 @@ def test_lstm_u85_INT(): "run_method_and_compare_outputs", get_test_inputs(), atol=3e-1, qtol=1.0 ) pipeline.run() + + +@common.SkipIfNoModelConverter +def test_lstm_vgf_INT(): + pipeline = VgfPipeline[input_t]( + TestLSTM.lstm, + TestLSTM.model_example_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", get_test_inputs(), atol=3e-1, qtol=1.0 + # ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_lstm_vgf_FP(): + pipeline = VgfPipeline[input_t]( + TestLSTM.lstm, + TestLSTM.model_example_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", get_test_inputs(), atol=3e-1, qtol=1.0 + # ) + pipeline.run() diff --git a/backends/arm/test/models/test_mobilenet_v2_arm.py b/backends/arm/test/models/test_mobilenet_v2_arm.py index 090d7f849d3..d4e3bbc8e28 100644 --- a/backends/arm/test/models/test_mobilenet_v2_arm.py +++ b/backends/arm/test/models/test_mobilenet_v2_arm.py @@ -16,6 +16,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from torchvision import models, transforms # type: ignore[import-untyped] @@ -94,3 +95,41 @@ def test_mv2_u85_INT(per_channel_quantization): qtol=1, ) pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("per_channel_quantization", quant_test_data) +def test_mv2_vgf_INT(per_channel_quantization): + pipeline = VgfPipeline[input_t]( + mv2, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + atol=0.25, + qtol=1, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", get_test_inputs(), atol=3e-1, qtol=1.0 + # ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_mv2_vgf_FP(): + pipeline = VgfPipeline[input_t]( + mv2, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + # TODO: MLETORCH-1167 Create Vulkan backend e2e tests + # pipeline.change_args( + # "run_method_and_compare_outputs", get_test_inputs(), atol=3e-1, qtol=1.0 + # ) # TODO: MLETORCH-1036 decrease tolerance + pipeline.run() diff --git a/backends/arm/test/models/test_mobilenet_v3_arm.py b/backends/arm/test/models/test_mobilenet_v3_arm.py index c43f20b2884..0dcbd9757ac 100644 --- a/backends/arm/test/models/test_mobilenet_v3_arm.py +++ b/backends/arm/test/models/test_mobilenet_v3_arm.py @@ -15,6 +15,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from torchvision import models, transforms @@ -82,3 +83,32 @@ def test_mv3_u85_INT(): qtol=1, ) pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.slow +def test_mv3_vgf_INT(): + pipeline = VgfPipeline[input_t]( + mv3, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + atol=0.5, + qtol=1, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_mv3_vgf_FP(): + pipeline = VgfPipeline[input_t]( + mv3, + model_inputs, + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_nn_functional.py b/backends/arm/test/models/test_nn_functional.py index 651f9585459..4896074b544 100644 --- a/backends/arm/test/models/test_nn_functional.py +++ b/backends/arm/test/models/test_nn_functional.py @@ -81,9 +81,6 @@ def forward(self, *args): @parametrize( "test_data", module_tests, - xfails={ - "affine_grid": "Int64 input. Partition handling fails since arange int64 output is split between 2 partitions.", - }, ) def test_nn_functional_FP(test_data): module, inputs = test_data @@ -91,7 +88,6 @@ def test_nn_functional_FP(test_data): module, inputs, "", use_to_edge_transform_and_lower=False ) pipeline.pop_stage("check.aten") - pipeline.dump_artifact("to_edge") pipeline.pop_stage("check_count.exir") try: pipeline.run() @@ -103,14 +99,11 @@ def test_nn_functional_FP(test_data): raise e -x_fails = { - "normalize": "MLETORCH-852: Support aten.index_put.default", - "unfold": "Int64 input && MLETORCH-827: Support aten.index.Tensor", - "fold": "Int64 input && MLETORCH-827: Support aten.index_put.default", -} - - -@parametrize("test_data", module_tests, x_fails, strict=False) +@parametrize( + "test_data", + module_tests, + {"normalize": "MLETORCH-1255: Unsupported dtype in InsertTableOpsPass"}, +) def test_nn_functional_INT(test_data): module, inputs = test_data pipeline = TosaPipelineINT[input_t]( diff --git a/backends/arm/test/models/test_resnet18.py b/backends/arm/test/models/test_resnet18.py new file mode 100644 index 00000000000..6e965daeb8b --- /dev/null +++ b/backends/arm/test/models/test_resnet18.py @@ -0,0 +1,99 @@ +# Copyright 2025 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.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, +) + +from torchvision import transforms # type: ignore[import-untyped] +from torchvision.models import resnet18, ResNet18_Weights + +model = resnet18(weights=ResNet18_Weights) +model = model.eval() +normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + +model_inputs = (normalize(torch.randn((1, 3, 224, 224))),) + +input_t = Tuple[torch.Tensor] + + +quant_test_data = { + "per_channel_quantization=true": True, + "per_channel_quantization=false": False, +} + + +def test_resnet_tosa_FP(): + pipeline = TosaPipelineFP[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +@common.parametrize("per_channel_quantization", quant_test_data) +def test_resnet_tosa_INT(per_channel_quantization): + pipeline = TosaPipelineINT[input_t]( + model, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + atol=0.5, + qtol=1, + ) + pipeline.run() + + +@pytest.mark.slow +@common.XfailIfNoCorstone300 +@common.parametrize("per_channel_quantization", quant_test_data) +def test_resnet_u55_INT(per_channel_quantization): + pipeline = EthosU55PipelineINT[input_t]( + model, + model_inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + atol=0.5, + qtol=1, + ) + pipeline.run() + + +@pytest.mark.slow +@pytest.mark.xfail( + reason="For resnet18 for Ethos-U85, the SRAM memory footprint is very high. The compiler team is investigating." +) +@common.XfailIfNoCorstone320 +@common.parametrize("per_channel_quantization", quant_test_data) +def test_resnet_u85_INT(per_channel_quantization): + pipeline = EthosU85PipelineINT[input_t]( + model, + model_inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=True, + use_to_edge_transform_and_lower=True, + per_channel_quantization=per_channel_quantization, + atol=0.5, + qtol=1, + ) + pipeline.run() diff --git a/backends/arm/test/models/test_w2l_arm.py b/backends/arm/test/models/test_w2l_arm.py index fa19a3b97e4..c627cd7f887 100644 --- a/backends/arm/test/models/test_w2l_arm.py +++ b/backends/arm/test/models/test_w2l_arm.py @@ -17,6 +17,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from torchaudio import models @@ -35,7 +36,6 @@ class TestW2L(unittest.TestCase): input_frames = 400 num_features = 1 - w2l = models.Wav2Letter(num_features=num_features).eval() model_example_inputs = get_test_inputs(batch_size, num_features, input_frames) all_operators = [ @@ -44,11 +44,17 @@ class TestW2L(unittest.TestCase): "executorch_exir_dialects_edge__ops_aten_relu_default", ] + @staticmethod + def create_model(input_type: str = "waveform"): + return models.Wav2Letter( + num_features=TestW2L.num_features, input_type=input_type + ).eval() + @pytest.mark.slow # about 3min on std laptop def test_w2l_tosa_FP(): pipeline = TosaPipelineFP[input_t]( - TestW2L.w2l, + TestW2L.create_model(), TestW2L.model_example_inputs, aten_op=[], exir_op=TestW2L.all_operators, @@ -61,7 +67,7 @@ def test_w2l_tosa_FP(): @pytest.mark.flaky def test_w2l_tosa_INT(): pipeline = TosaPipelineINT[input_t]( - TestW2L.w2l, + TestW2L.create_model(), TestW2L.model_example_inputs, aten_op=[], exir_op=TestW2L.all_operators, @@ -73,12 +79,14 @@ def test_w2l_tosa_INT(): @pytest.mark.slow @common.XfailIfNoCorstone300 @pytest.mark.xfail( - reason="MLETORCH-1009: Wav2Letter fails on U55 due to unsupported conditions", - strict=False, + reason="Wav2Letter fails on U55 due to insufficient memory", + strict=True, ) def test_w2l_u55_INT(): pipeline = EthosU55PipelineINT[input_t]( - TestW2L.w2l, + # Use "power_spectrum" variant because the default ("waveform") has a + # conv1d layer with an unsupported stride size. + TestW2L.create_model("power_spectrum"), TestW2L.model_example_inputs, aten_ops=[], exir_ops=[], @@ -93,7 +101,7 @@ def test_w2l_u55_INT(): @pytest.mark.skip(reason="Intermittent timeout issue: MLETORCH-856") def test_w2l_u85_INT(): pipeline = EthosU85PipelineINT[input_t]( - TestW2L.w2l, + TestW2L.create_model(), TestW2L.model_example_inputs, aten_ops=[], exir_ops=[], @@ -101,3 +109,30 @@ def test_w2l_u85_INT(): run_on_fvp=True, ) pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.slow +def test_w2l_vgf_INT(): + pipeline = VgfPipeline[input_t]( + TestW2L.create_model(), + TestW2L.model_example_inputs, + aten_op=[], + exir_op=TestW2L.all_operators, + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_w2l_vgf_FP(): + pipeline = VgfPipeline[input_t]( + TestW2L.create_model(), + TestW2L.model_example_inputs, + aten_op=[], + exir_op=TestW2L.all_operators, + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_acos.py b/backends/arm/test/ops/test_acos.py new file mode 100644 index 00000000000..28dadcf95be --- /dev/null +++ b/backends/arm/test/ops/test_acos.py @@ -0,0 +1,128 @@ +# Copyright 2025 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.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +input_t = Tuple[torch.Tensor] +aten_op = "torch.ops.aten.acos.default" +exir_op = "executorch_exir_dialects_edge__ops_aten__acos_default" + + +test_data_suite = { + "ones": lambda: torch.ones(1, 7, 10, 12), + "rand_in_range": lambda: (torch.rand(10, 10) - 0.5) * 2, # Uniform in [-1, 1) + "ramp_valid": lambda: torch.linspace(-1.0, 1.0, steps=160), + "edge_cases": lambda: torch.tensor([-1.0, 0.0, 1.0]), + "1d_tensor": lambda: torch.linspace(-1.0, 1.0, steps=10), # Shape: [10] + "2d_batch": lambda: torch.tensor( + [[-1.0, -0.5, 0.0, 0.5, 1.0], [0.9, -0.9, 0.3, -0.3, 0.0]] + ), # Shape: [2, 5] + "3d_batch": lambda: torch.rand(4, 5, 6) * 2 - 1, # Shape: [4, 5, 6] in [-1, 1) + "3d_mixed_shape": lambda: (torch.rand(7, 15, 2) - 0.5) * 2, + "4d_mixed": lambda: torch.linspace(-1, 1, steps=1 * 3 * 4 * 5).reshape( + 1, 3, 4, 5 + ), # Shape: [2, 3, 4, 5] + "4d_random": lambda: (torch.rand(1, 5, 10, 7) - 0.5) * 2, + "bool_casted": lambda: torch.ones(3, 3, dtype=torch.bool).to( + dtype=torch.float32 + ), # All 1.0 (edge case) +} + + +class Acos(torch.nn.Module): + + def forward(self, x: torch.Tensor): + return torch.acos(x) + + +@common.parametrize("test_data", test_data_suite) +def test_acos_tosa_FP(test_data: Tuple): + pipeline = TosaPipelineFP[input_t]( + Acos(), + (test_data(),), + aten_op, + exir_op=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +def test_acos_tosa_INT(test_data: Tuple): + pipeline = TosaPipelineINT[input_t]( + Acos(), + (test_data(),), + aten_op=aten_op, + exir_op=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +def test_acos_u55_INT(test_data: Tuple): + pipeline = EthosU55PipelineINT[input_t]( + Acos(), + (test_data(),), + aten_ops=aten_op, + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +def test_acos_u85_INT(test_data: Tuple): + pipeline = EthosU85PipelineINT[input_t]( + Acos(), + (test_data(),), + aten_ops=aten_op, + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_acos_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t]( + Acos(), + (test_data(),), + [], + [], + tosa_version="TOSA-1.0+FP", + run_on_vulkan_runtime=True, + ) + try: + pipeline.run() + except FileNotFoundError as e: + pytest.skip(f"VKML executor_runner not found - not built - skip {e}") + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_acos_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t]( + Acos(), + (test_data(),), + [], + [], + tosa_version="TOSA-1.0+INT", + run_on_vulkan_runtime=True, + ) + try: + pipeline.run() + except FileNotFoundError as e: + pytest.skip(f"VKML executor_runner not found - not built - skip {e}") diff --git a/backends/arm/test/ops/test_add.py b/backends/arm/test/ops/test_add.py index 421ec0adc61..2eabd302df6 100644 --- a/backends/arm/test/ops/test_add.py +++ b/backends/arm/test/ops/test_add.py @@ -7,9 +7,13 @@ from typing import Tuple +import pytest import torch -from executorch.backends.arm.arm_backend import get_tosa_spec from executorch.backends.arm.quantizer import arm_quantizer +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, @@ -18,7 +22,8 @@ TosaPipelineINT, VgfPipeline, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.specification import get_tosa_spec from executorch.backends.xnnpack.test.tester import Quantize from torchao.quantization.pt2e import HistogramObserver from torchao.quantization.pt2e.quantizer import QuantizationSpec @@ -58,13 +63,17 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): "4d_randn_1": lambda: (torch.randn(1, 1, 4, 4), torch.ones(1, 1, 4, 1)), "4d_randn_2": lambda: (torch.randn(1, 3, 4, 4), torch.randn(1, 3, 4, 4)), "4d_randn_big": lambda: ( - 10000 * torch.randn(1, 1, 4, 4), + (1 << 30) * torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1), ), "4d_randn_1_mutltiple_broadcasts": lambda: ( torch.randn(1, 4, 4, 1), torch.ones(1, 1, 4, 4), ), + "4d_big_small": lambda: ( + (10e10) * torch.randn(1, 10, 20, 30), + torch.randn(1, 10, 20, 30), + ), } @@ -87,7 +96,7 @@ def test_add_tensor_tosa_FP(test_data: input_t1): @common.parametrize("test_data", Add.test_data) def test_add_tensor_tosa_INT(test_data: input_t1): - pipeline = TosaPipelineINT[input_t1](Add(), test_data(), aten_op, exir_op) + pipeline = TosaPipelineINT[input_t1](Add(), test_data(), aten_op, exir_op, qtol=0) pipeline.run() @@ -112,9 +121,16 @@ def test_add_tensor_tosa_INT_i32(test_data: input_t1): quant_max=2**31 - 1, quant_min=-(2**31), ) + output_act_qspec = QuantizationSpec( + torch.int32, + observer, + qscheme=torch.per_tensor_symmetric, + quant_max=2**31 - 1, + quant_min=-(2**31), + ) # This quantization_config will be set as global config. quantization_config = arm_quantizer.QuantizationConfig( - input_act_qspec, None, None, None + input_act_qspec, output_act_qspec, None, None ) quantize_stage = Quantize(quantizer, quantization_config) pipeline.change_args("quantize", quantize_stage) @@ -158,13 +174,13 @@ def test_add_tensor_tosa_FP_3(test_data: input_t2): @common.parametrize("test_data", Add3.test_data) def test_add_tensor_tosa_INT_3(test_data: input_t2): - pipeline = TosaPipelineINT[input_t2](Add3(), test_data(), aten_op, exir_op) + pipeline = TosaPipelineINT[input_t2](Add3(), test_data(), aten_op, exir_op, qtol=0) pipeline.run() @common.parametrize("test_data", Add2.test_data) def test_add_tensor_tosa_INT_2(test_data: input_t2): - pipeline = TosaPipelineINT[input_t2](Add2(), test_data(), aten_op, exir_op) + pipeline = TosaPipelineINT[input_t2](Add2(), test_data(), aten_op, exir_op, qtol=0) pipeline.run() @@ -190,9 +206,17 @@ def test_add_tensor_u85_INT_2(test_data: input_t2): @common.SkipIfNoModelConverter def test_add_tensor_vgf_FP(test_data: input_t1): pipeline = VgfPipeline[input_t1]( - Add(), test_data(), aten_op, exir_op, tosa_version="TOSA-1.0+FP" + Add(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + run_on_vulkan_runtime=True, ) - pipeline.run() + try: + pipeline.run() + except FileNotFoundError as e: + pytest.skip(f"VKML executor_runner not found - not built - skip {e}") @common.parametrize("test_data", Add.test_data) @@ -204,5 +228,111 @@ def test_add_tensor_vgf_INT(test_data: input_t1): aten_op, exir_op, tosa_version="TOSA-1.0+INT", + run_on_vulkan_runtime=True, + ) + try: + pipeline.run() + except FileNotFoundError as e: + pytest.skip(f"VKML executor_runner not found - not built - skip {e}") + + +def get_symmetric_a16w8_add_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", Add.test_data) +@pytest.mark.xfail( + reason="missing int16 add ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13730" +) +def test_add_tensor_16a8w_tosa_INT(test_data: input_t1): + """Test add operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Add(), + test_data(), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_add_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", Add.test_data) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 add operations. See: https://github.com/pytorch/executorch/issues/13730" +) +def test_add_tensor_16a8w_u55_INT16(test_data: input_t1): + """Test add operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Add(), + test_data(), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_add_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", Add.test_data) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 add operations. See: https://github.com/pytorch/executorch/issues/13730" +) +def test_add_tensor_16a8w_u85_INT16(test_data: input_t1): + """Test add operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Add(), + test_data(), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_add_quantizer( + per_channel_quantization=per_channel_quantization + ), ) pipeline.run() diff --git a/backends/arm/test/ops/test_addmm.py b/backends/arm/test/ops/test_addmm.py index cfe324ab0af..753cb599b2b 100644 --- a/backends/arm/test/ops/test_addmm.py +++ b/backends/arm/test/ops/test_addmm.py @@ -5,9 +5,14 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -15,6 +20,8 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.addmm.default" @@ -182,3 +189,105 @@ def test_addmm_vgf_INT(test_data: input_t1): tosa_version="TOSA-1.0+INT", ) pipeline.run() + + +def get_symmetric_a16w8_addmm_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +@pytest.mark.xfail( + reason="missing int16 addmm ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13979" +) +def test_addmm_16a8w_tosa_INT(test_data: input_t1): + """Test addmm (FC layer) operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Addmm(), + (*test_data,), + aten_op=[], + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_addmm_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 addmm operations" +) +def test_addmm_16a8w_u55_INT16(test_data: input_t1): + """Test addmm (FC layer) operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Addmm(), + (*test_data,), + aten_ops=[], + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_addmm_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 addmm operations" +) +def test_addmm_16a8w_u85_INT16(test_data: input_t1): + """Test addmm (FC layer) operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Addmm(), + (*test_data,), + aten_ops=[], + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_addmm_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_alias_copy.py b/backends/arm/test/ops/test_alias_copy.py index cf8caca02c4..8b951a4d856 100644 --- a/backends/arm/test/ops/test_alias_copy.py +++ b/backends/arm/test/ops/test_alias_copy.py @@ -41,7 +41,9 @@ def __init__(self): super().__init__() def forward(self, x: torch.Tensor): - return torch.alias_copy(x) + return ( + torch.alias_copy(x) * 1 + ) # Multiply by one to make sure it is partitioned. @common.parametrize("test_data", AliasCopy.test_data) diff --git a/backends/arm/test/ops/test_amax.py b/backends/arm/test/ops/test_amax.py index 3600c34c94c..080dddda92e 100644 --- a/backends/arm/test/ops/test_amax.py +++ b/backends/arm/test/ops/test_amax.py @@ -95,10 +95,7 @@ def test_amax_u55_INT_not_delegated(): pipeline.run() -fvp_xfails = {"rank_4_mult_batches": "MLETORCH-517 : Multiple batches not supported"} - - -@common.parametrize("test_data", Amax.test_data, fvp_xfails, strict=False) +@common.parametrize("test_data", Amax.test_data) @common.XfailIfNoCorstone320 def test_amax_u85_INT(test_data: Amax.input_t): data, dim, keep_dims = test_data() diff --git a/backends/arm/test/ops/test_amin.py b/backends/arm/test/ops/test_amin.py index 3ae94fe3c6e..a24da9e1ba0 100644 --- a/backends/arm/test/ops/test_amin.py +++ b/backends/arm/test/ops/test_amin.py @@ -104,10 +104,7 @@ def test_amin_u55_INT_not_delegated(): pipeline.run() -fvp_xfails = {"rank_4_mult_batches": "MLETORCH-517 : Multiple batches not supported"} - - -@common.parametrize("test_data", Amin.test_data, fvp_xfails, strict=False) +@common.parametrize("test_data", Amin.test_data) @common.XfailIfNoCorstone320 def test_amin_u85_INT(test_data: Amin.input_t): data, dim, keep_dims = test_data() diff --git a/backends/arm/test/ops/test_arange.py b/backends/arm/test/ops/test_arange.py index ede00768f52..33cca542922 100644 --- a/backends/arm/test/ops/test_arange.py +++ b/backends/arm/test/ops/test_arange.py @@ -12,6 +12,7 @@ from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, @@ -46,6 +47,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: lambda: (torch.randint(0, 10, [10], dtype=torch.int32),), (0.0, 10.0, 1.0, torch.int32), ), + } + test_reject: dict[str, test_data_t] = { "int32_int64": ( lambda: (torch.randint(0, 10, [10], dtype=torch.int32),), (0.0, 10.0, 1.0, torch.int64), @@ -77,6 +80,15 @@ def test_arange_start_step_tosa_FP_dtypes(test_data: test_data_t): pipeline.run() +@common.parametrize("test_data", ArangeAdd.test_reject) +def test_arange_start_step_tosa_FP_not_delegated(test_data: test_data_t): + input_data, init_data = test_data + pipeline = OpNotSupportedPipeline[input_t]( + ArangeAdd(*init_data), input_data(), non_delegated_ops={ArangeAdd.exir_op: 1} + ) + pipeline.run() + + @common.parametrize("test_data", ArangeAdd.test_data) def test_arange_start_step_tosa_INT(test_data: test_data_t): input_data, init_data = test_data diff --git a/backends/arm/test/ops/test_bitwise.py b/backends/arm/test/ops/test_bitwise.py index 1c0f0e36a6a..218f2290cab 100644 --- a/backends/arm/test/ops/test_bitwise.py +++ b/backends/arm/test/ops/test_bitwise.py @@ -4,19 +4,19 @@ # LICENSE file in the root directory of this source tree. +from copy import copy from typing import Tuple +import pytest import torch from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( EthosU85PipelineINT, OpNotSupportedPipeline, - TosaPipelineFP, TosaPipelineINT, VgfPipeline, ) - input_t2 = Tuple[torch.Tensor, torch.Tensor] # Input x, y @@ -56,6 +56,9 @@ class BitwiseBinary(torch.nn.Module): ), } + test_data_u85 = copy(test_data) + del test_data_u85["zeros"] + class BitwiseBinaryScalar(torch.nn.Module): test_data = { @@ -77,6 +80,9 @@ class BitwiseBinaryScalar(torch.nn.Module): ), } + test_data_u85 = copy(test_data) + del test_data_u85["zeros"] + class And(BitwiseBinary): aten_op = "torch.ops.aten.bitwise_and.Tensor" @@ -106,6 +112,7 @@ class AndScalar(BitwiseBinaryScalar): aten_op = "torch.ops.aten.bitwise_and.Scalar" # Tensor because it gets converted from Scalar -> Tensor in lowering exir_op = "executorch_exir_dialects_edge__ops_aten_bitwise_and_Tensor" + exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_bitwise_and_Scalar" def forward(self, tensor: torch.Tensor, scalar: int): return tensor.bitwise_and(scalar) @@ -115,6 +122,7 @@ class XorScalar(BitwiseBinaryScalar): aten_op = "torch.ops.aten.bitwise_xor.Scalar" # Tensor because it gets converted from Scalar -> Tensor in lowering exir_op = "executorch_exir_dialects_edge__ops_aten_bitwise_xor_Tensor" + exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_bitwise_xor_Scalar" def forward(self, tensor: torch.Tensor, scalar: int): return tensor.bitwise_xor(scalar) @@ -124,6 +132,7 @@ class OrScalar(BitwiseBinaryScalar): aten_op = "torch.ops.aten.bitwise_or.Scalar" # Tensor because it gets converted from Scalar -> Tensor in lowering exir_op = "executorch_exir_dialects_edge__ops_aten_bitwise_or_Tensor" + exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_bitwise_or_Scalar" def forward(self, tensor: torch.Tensor, scalar: int): return tensor.bitwise_or(scalar) @@ -136,28 +145,20 @@ def forward(self, tensor: torch.Tensor, scalar: int): @common.parametrize("test_data", And().test_data) def test_bitwise_and_tensor_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( And(), test_data(), - And().aten_op, - And().exir_op, - atol=0, - rtol=0, - qtol=0, + {And.exir_op: 1}, ) pipeline.run() @common.parametrize("test_data", AndScalar.test_data) def test_bitwise_and_scalar_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( AndScalar(), test_data(), - AndScalar.aten_op, - AndScalar.exir_op, - atol=0, - rtol=0, - qtol=0, + {AndScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -200,7 +201,7 @@ def test_bitwise_and_tensor_u55_INT(test_data: input_t2): pipeline = OpNotSupportedPipeline[input_t2]( And(), test_data(), - {And().exir_op: 1}, + {And.exir_op: 1}, quantize=True, u55_subset=True, ) @@ -226,7 +227,7 @@ def test_bitwise_and_scalar_u55_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", AndScalar.test_data) +@common.parametrize("test_data", AndScalar.test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_and_scalar_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -244,7 +245,7 @@ def test_bitwise_and_scalar_u85_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", And().test_data) +@common.parametrize("test_data", And().test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_and_tensor_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -265,15 +266,10 @@ def test_bitwise_and_tensor_u85_INT(test_data: input_t2): @common.parametrize("test_data", And().test_data) @common.SkipIfNoModelConverter def test_bitwise_and_tensor_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( And(), test_data(), - And().aten_op, - And().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {And.exir_op: 1}, ) pipeline.run() @@ -281,15 +277,10 @@ def test_bitwise_and_tensor_vgf_FP(test_data: input_t2): @common.parametrize("test_data", AndScalar().test_data) @common.SkipIfNoModelConverter def test_bitwise_and_scalar_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( AndScalar(), test_data(), - AndScalar().aten_op, - AndScalar().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {AndScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -337,28 +328,20 @@ def test_bitwise_and_scalar_vgf_INT(test_data: input_t2): @common.parametrize("test_data", Xor().test_data) def test_bitwise_xor_tensor_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( Xor(), test_data(), - Xor().aten_op, - Xor().exir_op, - atol=0, - rtol=0, - qtol=0, + {Xor.exir_op: 1}, ) pipeline.run() @common.parametrize("test_data", XorScalar.test_data) def test_bitwise_xor_scalar_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( XorScalar(), test_data(), - XorScalar.aten_op, - XorScalar.exir_op, - atol=0, - rtol=0, - qtol=0, + {XorScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -427,7 +410,7 @@ def test_bitwise_xor_scalar_u55_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", Xor().test_data) +@common.parametrize("test_data", Xor().test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_xor_tensor_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -445,7 +428,7 @@ def test_bitwise_xor_tensor_u85_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", XorScalar.test_data) +@common.parametrize("test_data", XorScalar.test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_xor_scalar_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -466,15 +449,10 @@ def test_bitwise_xor_scalar_u85_INT(test_data: input_t2): @common.parametrize("test_data", Xor().test_data) @common.SkipIfNoModelConverter def test_bitwise_xor_tensor_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( Xor(), test_data(), - Xor().aten_op, - Xor().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {Xor.exir_op: 1}, ) pipeline.run() @@ -482,15 +460,10 @@ def test_bitwise_xor_tensor_vgf_FP(test_data: input_t2): @common.parametrize("test_data", XorScalar().test_data) @common.SkipIfNoModelConverter def test_bitwise_xor_scalar_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( XorScalar(), test_data(), - XorScalar().aten_op, - XorScalar().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {XorScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -538,28 +511,20 @@ def test_bitwise_xor_scalar_vgf_INT(test_data: input_t2): @common.parametrize("test_data", Or().test_data) def test_bitwise_or_tensor_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( Or(), test_data(), - Or().aten_op, - Or().exir_op, - atol=0, - rtol=0, - qtol=0, + {Or.exir_op: 1}, ) pipeline.run() @common.parametrize("test_data", OrScalar.test_data) def test_bitwise_or_scalar_tosa_FP(test_data: input_t2): - pipeline = TosaPipelineFP[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( OrScalar(), test_data(), - OrScalar.aten_op, - OrScalar.exir_op, - atol=0, - rtol=0, - qtol=0, + {OrScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -628,7 +593,7 @@ def test_bitwise_or_scalar_u55_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", Or().test_data) +@common.parametrize("test_data", Or().test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_or_tensor_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -646,7 +611,7 @@ def test_bitwise_or_tensor_u85_INT(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", OrScalar.test_data) +@common.parametrize("test_data", OrScalar.test_data_u85) @common.XfailIfNoCorstone320 def test_bitwise_or_scalar_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -667,15 +632,10 @@ def test_bitwise_or_scalar_u85_INT(test_data: input_t2): @common.parametrize("test_data", Or().test_data) @common.SkipIfNoModelConverter def test_bitwise_or_tensor_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( Or(), test_data(), - Or().aten_op, - Or().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {Or.exir_op: 1}, ) pipeline.run() @@ -683,15 +643,10 @@ def test_bitwise_or_tensor_vgf_FP(test_data: input_t2): @common.parametrize("test_data", OrScalar().test_data) @common.SkipIfNoModelConverter def test_bitwise_or_scalar_vgf_FP(test_data: input_t2): - pipeline = VgfPipeline[input_t2]( + pipeline = OpNotSupportedPipeline[input_t2]( OrScalar(), test_data(), - OrScalar().aten_op, - OrScalar().exir_op, - atol=0, - rtol=0, - qtol=0, - tosa_version="TOSA-1.0+FP", + {OrScalar.exir_op_scalar: 1}, ) pipeline.run() @@ -730,3 +685,30 @@ def test_bitwise_or_scalar_vgf_INT(test_data: input_t2): pipeline.pop_stage("quantize") pipeline.pop_stage("check.quant_nodes") pipeline.run() + + +@pytest.mark.xfail( + reason="MLBEDSW-11029: Fatal Python floating point error in Vela for rank 4 bitwse ops with int32 dtype." +) +def test_bitwise_or_tensor_u85_INT_zeros(): + raise RuntimeError( + "Dummy test to xfail mark u85 zeros test case since running the actual test causes a fatal crash." + ) + + +@pytest.mark.xfail( + reason="MLBEDSW-11029: Fatal Python floating point error in Vela for rank 4 bitwse ops with int32 dtype." +) +def test_bitwise_and_tensor_u85_INT_zeros(): + raise RuntimeError( + "Dummy test to xfail mark u85 zeros test case since running the actual test causes a fatal crash." + ) + + +@pytest.mark.xfail( + reason="MLBEDSW-11029: Fatal Python floating point error in Vela for rank 4 bitwse ops with int32 dtype." +) +def test_bitwise_xor_tensor_u85_INT_zeros(): + raise RuntimeError( + "Dummy test to xfail mark u85 zeros test case since running the actual test causes a fatal crash." + ) diff --git a/backends/arm/test/ops/test_cat.py b/backends/arm/test/ops/test_cat.py index 826689622fb..84ecd8641b5 100644 --- a/backends/arm/test/ops/test_cat.py +++ b/backends/arm/test/ops/test_cat.py @@ -8,8 +8,13 @@ from typing import Tuple +import pytest import torch -from executorch.backends.arm.test import common +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, @@ -18,6 +23,8 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize input_t1 = Tuple[torch.Tensor] # Input x @@ -105,15 +112,7 @@ def test_cat_tosa_INT(test_data: Tuple): pipeline.run() -x_fails = { - "cat_rand_two_tensors_dim_0": "MLETORCH-630: AssertionError: Output 0 does not match reference output.", - "cat_rand_two_tensors_dim_0": "MLETORCH-630: AssertionError: Output 0 does not match reference output.", - "cat_rand_two_tensors_dim_3": "MLETORCH-630: AssertionError: Output 0 does not match reference output.", - "cat_rand_large": "MLETORCH-630: AssertionError: Output 0 does not match reference output.", -} - - -@common.parametrize("test_data", Cat.test_parameters, x_fails) +@common.parametrize("test_data", Cat.test_parameters) @common.XfailIfNoCorstone300 def test_cat_u55_INT(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( @@ -126,7 +125,7 @@ def test_cat_u55_INT(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Cat.test_parameters, x_fails) +@common.parametrize("test_data", Cat.test_parameters) @common.XfailIfNoCorstone320 def test_cat_u85_INT(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( @@ -159,3 +158,105 @@ def test_cat_vgf_INT(test_data: Tuple): tosa_version="TOSA-1.0+INT", ) pipeline.run() + + +def get_symmetric_a16w8_cat_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", Cat.test_parameters) +@pytest.mark.xfail( + reason="missing int16 cat ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13978" +) +def test_cat_16a8w_tosa_INT(test_data: Tuple): + """Test cat operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Cat(), + test_data(), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_cat_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", Cat.test_parameters) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 cat operations" +) +def test_cat_16a8w_u55_INT16(test_data: Tuple): + """Test cat operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Cat(), + test_data(), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_cat_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", Cat.test_parameters) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 cat operations" +) +def test_cat_16a8w_u85_INT16(test_data: Tuple): + """Test cat operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Cat(), + test_data(), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_cat_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_clone.py b/backends/arm/test/ops/test_clone.py index 7a24848697e..b240fb1ea07 100644 --- a/backends/arm/test/ops/test_clone.py +++ b/backends/arm/test/ops/test_clone.py @@ -3,13 +3,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# -# Tests the clone op which copies the data of the input tensor (possibly with new data format) -# from typing import Tuple -import pytest import torch from executorch.backends.arm.test import common @@ -23,62 +19,87 @@ ) aten_op = "torch.ops.aten.clone.default" -exir_op = "executorch_exir_dialects_edge__ops_aten_clone_default" +exir_op = "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default" input_t = Tuple[torch.Tensor] -class Clone(torch.nn.Module): - """A simple module that clones an input tensor.""" +class CloneFirstArg(torch.nn.Module): + def forward(self, x): + return x.clone() + x - def forward(self, x: torch.Tensor): - return x.clone() +class CloneSecondArg(torch.nn.Module): + def forward(self, x): + return x * x.clone() -test_data_suite = { - "ones_1D_10": lambda: (torch.ones(10),), - "ones_1D_50": lambda: (torch.ones(50),), - "rand_1D_20": lambda: (torch.rand(20),), - "rand_2D_10x10": lambda: (torch.rand(10, 10),), - "rand_3D_5x5x5": lambda: (torch.rand(5, 5, 5),), - "rand_4D_2x3x4x5": lambda: (torch.rand(2, 3, 4, 5),), - "large_tensor": lambda: (torch.rand(1000),), -} + +class CloneOutput(torch.nn.Module): + def forward(self, x): + return (x / x).clone() + + +class CloneBothArgs(torch.nn.Module): + def forward(self, x): + return x.clone() + x.clone() + + +class CloneAfterOtherOp(torch.nn.Module): + def forward(self, x): + x = x * 2 + return x.clone() + x + + +class CloneParallelToOtherOp(torch.nn.Module): + def forward(self, x): + return x * 2 + x.clone() -@common.parametrize("test_data", test_data_suite) -def test_clone_tosa_FP(test_data: Tuple[torch.Tensor]): +delegated_clones = { + "clone_first_arg": lambda: (CloneFirstArg, (torch.rand(1, 2, 3, 4),)), + "clone_second_arg": lambda: (CloneSecondArg, (torch.rand(1, 2, 3, 4),)), + "clone_output": lambda: (CloneOutput, (torch.rand(1, 2, 3, 4),)), + "clone_both_args": lambda: (CloneBothArgs, (torch.rand(1, 2, 3, 4),)), + "clone_after_other_op": lambda: (CloneAfterOtherOp, (torch.rand(1, 2, 3, 4),)), + "clone_parallel_to_other_op": lambda: ( + CloneParallelToOtherOp, + (torch.rand(1, 2, 3, 4),), + ), +} + +@common.parametrize("input_data", delegated_clones) +def test_clone_tosa_FP(input_data): + module, input_tensor = input_data() pipeline = TosaPipelineFP[input_t]( - Clone(), - test_data(), - aten_op, - exir_op, + module(), + input_tensor, + [], ) - pipeline.run() -@common.parametrize("test_data", test_data_suite) -def test_clone_tosa_INT(test_data): +@common.parametrize("input_data", delegated_clones) +def test_clone_tosa_INT(input_data): + module, input_tensor = input_data() + pipeline = TosaPipelineINT[input_t]( - Clone(), - test_data(), + module(), + input_tensor, aten_op, exir_op, ) pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("input_data", delegated_clones) @common.XfailIfNoCorstone300 -@pytest.mark.xfail( - reason="Empty subgraph leads to Vela compilation failure. See: https://jira.arm.com/browse/MLBEDSW-10477" -) -def test_clone_u55_INT(test_data): +def test_clone_u55_INT(input_data): + module, input_tensor = input_data() + pipeline = EthosU55PipelineINT[input_t]( - Clone(), - test_data(), + module(), + input_tensor, aten_op, exir_op, run_on_fvp=True, @@ -87,15 +108,14 @@ def test_clone_u55_INT(test_data): pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("input_data", delegated_clones) @common.XfailIfNoCorstone320 -@pytest.mark.xfail( - reason="Empty subgraph leads to Vela compilation failure. See: https://jira.arm.com/browse/MLBEDSW-10477" -) -def test_clone_u85_INT(test_data): +def test_clone_u85_INT(input_data): + module, input_tensor = input_data() + pipeline = EthosU85PipelineINT[input_t]( - Clone(), - test_data(), + module(), + input_tensor, aten_op, exir_op, run_on_fvp=True, @@ -104,21 +124,23 @@ def test_clone_u85_INT(test_data): pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", delegated_clones) @common.SkipIfNoModelConverter def test_clone_vgf_FP(test_data): + module, input_tensor = test_data() pipeline = VgfPipeline[input_t]( - Clone(), test_data(), aten_op, exir_op, tosa_version="TOSA-1.0+FP" + module(), input_tensor, aten_op, exir_op, tosa_version="TOSA-1.0+FP" ) pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", delegated_clones) @common.SkipIfNoModelConverter def test_clone_vgf_INT(test_data): + module, input_tensor = test_data() pipeline = VgfPipeline[input_t]( - Clone(), - test_data(), + module(), + input_tensor, aten_op, exir_op, tosa_version="TOSA-1.0+INT", diff --git a/backends/arm/test/ops/test_conv2d.py b/backends/arm/test/ops/test_conv2d.py index 0d23d2a6c7e..0300f7c2049 100644 --- a/backends/arm/test/ops/test_conv2d.py +++ b/backends/arm/test/ops/test_conv2d.py @@ -388,15 +388,6 @@ def forward(self, x): for q in [True, False] } -fvp_xfails = { - f"{k},per_channel_quant={q}": reason - for k, reason in { - "2x2_3x2x40x40_nobias": "MLETORCH-520: Numerical issues on FVP.", - "5x5_3x2x128x128_st1": "MLETORCH-520: Numerical issues on FVP.", - }.items() - for q in [True, False] -} - input_t = Tuple[torch.Tensor] @@ -426,7 +417,7 @@ def test_convolution_2d_tosa_INT(test_data): pipeline.run() -@common.parametrize("test_data", test_data_INT, fvp_xfails) +@common.parametrize("test_data", test_data_INT) @common.XfailIfNoCorstone300 def test_convolution_2d_u55_INT(test_data): model, per_channel_quantization = test_data() @@ -441,7 +432,7 @@ def test_convolution_2d_u55_INT(test_data): pipeline.run() -@common.parametrize("test_data", test_data_INT, fvp_xfails) +@common.parametrize("test_data", test_data_INT) @common.XfailIfNoCorstone320 def test_convolution_u85_INT(test_data): model, per_channel_quantization = test_data() diff --git a/backends/arm/test/ops/test_conv_combos.py b/backends/arm/test/ops/test_conv_combos.py index 76502daf45c..a7a031468ea 100644 --- a/backends/arm/test/ops/test_conv_combos.py +++ b/backends/arm/test/ops/test_conv_combos.py @@ -47,28 +47,28 @@ def __init__(self): # (t, c, n, s) = (6, 96, 1, 1) # 1. 1x1 CONV2d + ReLU6 (Pointwise) self.pointwise_conv2d = torch.nn.Conv2d( - in_channels=32, out_channels=128, kernel_size=1, stride=1, groups=1 + in_channels=16, out_channels=96, kernel_size=1, stride=1, groups=1 ) ## (1, 128, 81, 81) - self.batch_norm2d_16 = torch.nn.BatchNorm2d(128, affine=False) + self.batch_norm2d_16 = torch.nn.BatchNorm2d(96, affine=False) self.relu6 = torch.nn.ReLU6() # 2. 3x3 DepthwiseConv2d + ReLu6 self.depthwise_conv2d = torch.nn.Conv2d( - in_channels=128, - out_channels=128, + in_channels=96, + out_channels=96, kernel_size=3, padding=1, stride=1, - groups=128, + groups=96, ) ## (1, 128, H, W) # 3. Linear 1x1 Conv2d self.pointwise_conv2d_linear = torch.nn.Conv2d( - in_channels=128, out_channels=32, kernel_size=1, stride=1, groups=1 + in_channels=96, out_channels=16, kernel_size=1, stride=1, groups=1 ) ## (1, 32, 81, 81) def get_inputs(self) -> Tuple[torch.Tensor]: - return (torch.randn(1, 32, 81, 81),) + return (torch.randn(1, 16, 81, 81),) def forward(self, x): input = x diff --git a/backends/arm/test/ops/test_cosh.py b/backends/arm/test/ops/test_cosh.py new file mode 100644 index 00000000000..60920d03f94 --- /dev/null +++ b/backends/arm/test/ops/test_cosh.py @@ -0,0 +1,114 @@ +# Copyright 2025 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 ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +aten_op = "torch.ops.aten.cosh.default" +exir_op = "executorch_exir_dialects_edge__ops_aten__cosh_default" + +input_t1 = Tuple[torch.Tensor] # Input x + +test_data_suite = { + # (test_name, test_data) + "zeros": torch.zeros(10, 10, 10), + "zeros_4D": torch.zeros(1, 10, 32, 7), + "zeros_alt_shape": torch.zeros(10, 3, 5), + "ones": torch.ones(15, 10, 7), + "ones_4D": torch.ones(1, 3, 32, 16), + "rand": torch.rand(10, 10) - 0.5, + "rand_alt_shape": torch.rand(10, 3, 5) - 0.5, + "rand_4D": torch.rand(1, 6, 5, 7) - 0.5, + "randn_pos": torch.randn(10) + 10, + "randn_neg": torch.randn(10) - 10, + "ramp": torch.arange(-16, 16, 0.2), + "large": 100 * torch.ones(1, 1), + "small": 0.000001 * torch.ones(1, 1), + "small_rand": torch.rand(100) * 0.01, + "biggest": torch.tensor([700.0, 710.0, 750.0]), +} + + +class Cosh(torch.nn.Module): + def forward(self, x: torch.Tensor): + return torch.cosh(x) + + +@common.parametrize("test_data", test_data_suite) +def test_cosh_tosa_FP(test_data: Tuple): + pipeline = TosaPipelineFP[input_t1]( + Cosh(), + (test_data,), + aten_op, + exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +def test_cosh_tosa_INT(test_data: Tuple): + pipeline = TosaPipelineINT[input_t1]( + Cosh(), (test_data,), aten_op=aten_op, exir_op=exir_op + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +@common.parametrize("test_data", test_data_suite) +def test_cosh_u55_INT(test_data: Tuple): + pipeline = EthosU55PipelineINT[input_t1]( + Cosh(), (test_data,), aten_ops=aten_op, exir_ops=exir_op + ) + pipeline.run() + + +@common.XfailIfNoCorstone320 +@common.parametrize( + "test_data", + test_data_suite, + xfails={ + "ones_4D": "MLBEDSW-11046 - Incorrect output for TABLE followed by RESHAPE" + }, + strict=False, +) +def test_cosh_u85_INT(test_data: Tuple): + pipeline = EthosU85PipelineINT[input_t1]( + Cosh(), (test_data,), aten_ops=aten_op, exir_ops=exir_op + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_cosh_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Cosh(), + (test_data,), + [], + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_cosh_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Cosh(), + (test_data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_cumsum.py b/backends/arm/test/ops/test_cumsum.py new file mode 100644 index 00000000000..ce175fb37c0 --- /dev/null +++ b/backends/arm/test/ops/test_cumsum.py @@ -0,0 +1,122 @@ +# Copyright 2025 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 ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +input_t1 = Tuple[torch.Tensor, int] +aten_op = "torch.ops.aten.cumsum.default" + +""" +Tests the aten.cumsum operator by decomposing it into a convolution and +verifying results across various dims and pipelines. +""" + + +class CumsumModule(torch.nn.Module): + test_parameters = { + "1d_dim0": lambda: (torch.rand(10), 0), + "1d_dim_neg1": lambda: (torch.rand(10), -1), + "2d_dim1": lambda: (torch.rand(5, 6), 1), + "3d_dim2": lambda: (torch.rand(2, 3, 4), 2), + "3d_dim0": lambda: (torch.rand(2, 3, 4), 0), + "4d_dim3": lambda: (torch.rand(1, 2, 3, 4), 3), + "4d_dim1": lambda: (torch.rand(1, 2, 3, 4), 1), + } + + def forward(self, x: torch.Tensor, dim: int) -> torch.Tensor: + return torch.cumsum(x, dim) + + +@common.parametrize("test_data", CumsumModule.test_parameters) +def test_cumsum_tosa_FP(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = TosaPipelineFP[input_t1]( + module, + args, + aten_op, + exir_op=[], + ) + pipeline.run() + + +@common.parametrize("test_data", CumsumModule.test_parameters) +def test_cumsum_tosa_INT(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = TosaPipelineINT[input_t1]( + module, + args, + aten_op, + exir_op=[], + ) + pipeline.run() + + +@common.parametrize("test_data", CumsumModule.test_parameters) +@common.SkipIfNoModelConverter +def test_cumsum_vgf_FP(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = VgfPipeline[input_t1]( + module, + args, + aten_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", CumsumModule.test_parameters) +@common.SkipIfNoModelConverter +def test_cumsum_vgf_INT(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = VgfPipeline[input_t1]( + module, + args, + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", CumsumModule.test_parameters) +@common.XfailIfNoCorstone300 +def test_cumsum_u55_INT(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + module, + args, + aten_ops=aten_op, + exir_ops=[], + ) + pipeline.run() + + +@common.parametrize("test_data", CumsumModule.test_parameters) +@common.XfailIfNoCorstone320 +def test_cumsum_u85_INT(test_data: input_t1): + module = CumsumModule() + args = test_data() + pipeline = EthosU85PipelineINT[input_t1]( + module, + args, + aten_ops=aten_op, + exir_ops=[], + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_depthwise_conv.py b/backends/arm/test/ops/test_depthwise_conv.py index bf6aad840ac..0f8b34d3d47 100644 --- a/backends/arm/test/ops/test_depthwise_conv.py +++ b/backends/arm/test/ops/test_depthwise_conv.py @@ -199,7 +199,7 @@ @common.parametrize("test_data", test_data_conv1d_FP | test_data_conv2d_FP) -def test_depthwise_convolution_2d_tosa_FP(test_data: torch.nn.Module): +def test_convolution_2d_tosa_FP_depthwise(test_data: torch.nn.Module): pipeline = TosaPipelineFP[input_t]( test_data(), test_data().get_inputs(), @@ -211,7 +211,7 @@ def test_depthwise_convolution_2d_tosa_FP(test_data: torch.nn.Module): @pytest.mark.flaky(reruns=5) # TODO: Investigate flakyness (MLTORCH-307) @common.parametrize("test_data", test_data_conv1d_INT | test_data_conv2d_INT) -def test_depthwise_convolution_2d_tosa_INT(test_data): +def test_convolution_2d_tosa_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = TosaPipelineINT[input_t]( model, @@ -225,7 +225,7 @@ def test_depthwise_convolution_2d_tosa_INT(test_data): @common.parametrize("test_data", test_data_conv1d_FP | test_data_conv2d_FP) @common.SkipIfNoModelConverter -def test_depthwise_convolution_2d_vgf_FP(test_data: torch.nn.Module): +def test_convolution_2d_vgf_FP_depthwise(test_data: torch.nn.Module): model = test_data() pipeline = VgfPipeline[input_t]( model, @@ -239,7 +239,7 @@ def test_depthwise_convolution_2d_vgf_FP(test_data: torch.nn.Module): @common.parametrize("test_data", test_data_conv1d_INT | test_data_conv2d_INT) @common.SkipIfNoModelConverter -def test_depthwise_convolution_2d_vgf_INT(test_data): +def test_convolution_2d_vgf_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = VgfPipeline[input_t]( model, @@ -251,19 +251,9 @@ def test_depthwise_convolution_2d_vgf_INT(test_data): pipeline.run() -x_fails = { - f"{k},per_channel_quant={q}": reason - for k, reason in { - "3x3_2x8x198x198_gp8_st3": "MLETORCH-517: Operators fail with batches > 1", - "two_dw_conv2d": "MLETORCH-517: Operators fail with batches > 1", - }.items() - for q in [True, False] -} - - @common.XfailIfNoCorstone300 # TODO: MLETORCH-516 -@common.parametrize("test_data", test_data_conv2d_INT, x_fails) -def test_depthwise_convolution_2d_u55_INT(test_data): +@common.parametrize("test_data", test_data_conv2d_INT) +def test_convolution_2d_u55_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = EthosU55PipelineINT[input_t]( model, @@ -278,7 +268,7 @@ def test_depthwise_convolution_2d_u55_INT(test_data): @common.XfailIfNoCorstone300 # TODO: MLETORCH-516 @common.parametrize("test_data", test_data_conv1d_INT) -def test_depthwise_convolution_1d_u55_INT(test_data): +def test_convolution_1d_u55_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = EthosU55PipelineINT[input_t]( model, @@ -292,8 +282,8 @@ def test_depthwise_convolution_1d_u55_INT(test_data): @common.XfailIfNoCorstone320 # TODO: MLETORCH-516 -@common.parametrize("test_data", test_data_conv2d_INT, x_fails) -def test_depthwise_convolution_2d_u85_INT(test_data): +@common.parametrize("test_data", test_data_conv2d_INT) +def test_convolution_2d_u85_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = EthosU85PipelineINT[input_t]( model, @@ -307,8 +297,8 @@ def test_depthwise_convolution_2d_u85_INT(test_data): @common.XfailIfNoCorstone320 # TODO: MLETORCH-516 -@common.parametrize("test_data", test_data_conv1d_INT, x_fails) -def test_depthwise_convolution_1d_u85_INT(test_data): +@common.parametrize("test_data", test_data_conv1d_INT) +def test_convolution_1d_u85_INT_depthwise(test_data): model, per_channel_quantization = test_data() pipeline = EthosU85PipelineINT[input_t]( model, diff --git a/backends/arm/test/ops/test_div.py b/backends/arm/test/ops/test_div.py index 026939758a0..5bacac1c962 100644 --- a/backends/arm/test/ops/test_div.py +++ b/backends/arm/test/ops/test_div.py @@ -101,17 +101,7 @@ def test_div_tensor_tosa_INT(test_data: Tuple): pipeline.run() -x_fails = { - "op_div_rank4_ones": "MLETORCH-521: Numerical issues on FVP likely due to mul op", - "op_div_rank4_negative_ones": "MLETORCH-521: Numerical issues on FVP likely due to mul op", - "op_div_rank4_ones_div_negative": "MLETORCH-521: Numerical issues on FVP likely due to mul op", - "op_div_rank4_large_rand": "MLETORCH-521: Numerical issues on FVP likely due to mul op", - "op_div_rank4_negative_large_rand": "MLETORCH-521: Numerical issues on FVP likely due to mul op", - "op_div_rank4_large_randn": "MLETORCH-521: Numerical issues on FVP likely due to mul op", -} - - -@common.parametrize("test_data", test_data_suite, xfails=x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone300 def test_div_tensor_u55_INT(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( @@ -124,7 +114,7 @@ def test_div_tensor_u55_INT(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", test_data_suite, xfails=x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone320 def test_div_tensor_u85_INT(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( diff --git a/backends/arm/test/ops/test_div_tensor_mode.py b/backends/arm/test/ops/test_div_tensor_mode.py new file mode 100644 index 00000000000..909b83bd97f --- /dev/null +++ b/backends/arm/test/ops/test_div_tensor_mode.py @@ -0,0 +1,153 @@ +# Copyright 2025 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 ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +input_tt = Tuple[torch.Tensor, torch.Tensor] + + +class DivTensorModeFloat(torch.nn.Module): + """ + torch.div(x, y, rounding_mode=mode) with + mode from {None, "floor", "trunc"}. + """ + + aten_ops = ["aten.div.Tensor_mode"] + aten_ops_int = ["aten.mul.Tensor", "aten.reciprocal.default"] + + def __init__(self, mode=None): + super().__init__() + assert mode in (None, "floor", "trunc") + self.mode = mode + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.div(x, y, rounding_mode=self.mode) + + +test_data = { + "mode_none": lambda: (None, (torch.randn(4, 8), torch.randn(4, 8).abs() + 1e-3)), + "mode_floor": lambda: ( + "floor", + (torch.randn(4, 8), torch.randn(4, 8).abs() + 1e-3), + ), + "mode_trunc": lambda: ( + "trunc", + (torch.randn(4, 8), torch.randn(4, 8).abs() + 1e-3), + ), + "int_denominator": lambda: (None, (torch.randn(4, 8), 2)), +} + + +@common.parametrize("data", test_data) +def test_div_tensor_mode_tosa_FP(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = TosaPipelineFP[input_tt]( + model, + inputs, + aten_op=model.aten_ops, + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + pipeline.run() + + +@common.parametrize("data", test_data) +def test_div_tensor_mode_tosa_INT(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = TosaPipelineINT[input_tt]( + model, + inputs, + aten_op=model.aten_ops_int, + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + pipeline.run() + + +@common.XfailIfNoCorstone300 +@common.parametrize( + "data", test_data, xfails={"mode_trunc": "CPU op missing in unittests"} +) +def test_div_tensor_mode_u55_INT(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = EthosU55PipelineINT[input_tt]( + model, + inputs, + aten_ops=model.aten_ops_int, + exir_ops=[], + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + pipeline.run() + + +@common.XfailIfNoCorstone320 +@common.parametrize("data", test_data) +def test_div_tensor_mode_u85_INT(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = EthosU85PipelineINT[input_tt]( + model, + inputs, + aten_ops=model.aten_ops_int, + exir_ops=[], + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("data", test_data) +def test_div_tensor_mode_vgf_INT(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = VgfPipeline[input_tt]( + model, + inputs, + aten_op=model.aten_ops_int, + exir_op=[], + tosa_version="TOSA-1.0+INT", + use_to_edge_transform_and_lower=True, + ) + pipeline.pop_stage("check_count.exir") + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("data", test_data) +def test_div_tensor_mode_vgf_FP(data): + mode, inputs = data() + model = DivTensorModeFloat(mode) + + pipeline = VgfPipeline[input_tt]( + model, + inputs, + aten_op=model.aten_ops, + exir_op=[], + tosa_version="TOSA-1.0+FP", + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_elu.py b/backends/arm/test/ops/test_elu.py new file mode 100644 index 00000000000..884f54c0202 --- /dev/null +++ b/backends/arm/test/ops/test_elu.py @@ -0,0 +1,133 @@ +# Copyright 2025 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 +import torch.nn as nn + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +test_data_suite = { + # (test_name, test_data) + "zeros_default": lambda: (1.0, torch.zeros(1, 10, 10, 10)), + "ones_default": lambda: (1.0, torch.ones(10, 10, 10)), + "rand_default": lambda: (1.0, torch.rand(10, 10) - 0.5), + "randn_pos_default": lambda: (1.0, torch.randn(1, 2, 3, 3) + 10), + "randn_neg_default": lambda: (1.0, torch.randn(2, 4, 3) - 10), + "ramp_default": lambda: (1.0, torch.arange(-16, 16, 0.2)), + "large_pos_default": lambda: (1.0, torch.randn(3, 3) * 1e6 + 1e7), + "large_neg_default": lambda: (1.0, -torch.empty(5).uniform_(1e5, 1e8)), + "small_pos_default": lambda: (1.0, torch.empty(5).uniform_(1e-8, 1e-5)), + "small_neg_default": lambda: (1.0, -torch.empty(5).uniform_(1e-8, 1e-5)), + "zeros_custom": lambda: (2.0, torch.zeros(1, 10, 10, 10)), + "ones_custom": lambda: (2.0, torch.ones(10, 10, 10)), + "rand_custom": lambda: (2.0, torch.rand(10, 10) - 0.5), + "randn_pos_custom": lambda: (2.0, torch.randn(1, 3, 3) + 10), + "randn_neg_custom": lambda: (2.0, torch.randn(1, 2, 4, 3) - 10), + "ramp_custom": lambda: (2.0, torch.arange(-16, 16, 0.2)), + "large_pos_custom": lambda: (2.0, torch.randn(3, 3) * 1e6 + 1e7), + "large_neg_custom": lambda: (2.0, -torch.empty(5).uniform_(1e5, 1e8)), + "small_pos_custom": lambda: (2.0, torch.empty(5).uniform_(1e-8, 1e-5)), + "small_neg_custom": lambda: (2.0, -torch.empty(5).uniform_(1e-8, 1e-5)), + "zeros_zero": lambda: (0.0, torch.zeros(1, 10, 10, 10)), + "ones_zero": lambda: (0.0, torch.ones(10, 10, 10)), + "rand_zero": lambda: (0.0, torch.rand(10, 10) - 0.5), + "randn_pos_zero": lambda: (0.0, torch.randn(1, 3, 3) + 10), + "randn_neg_zero": lambda: (0.0, torch.randn(1, 2, 4, 3) - 10), + "ramp_zero": lambda: (0.0, torch.arange(-16, 16, 0.2)), + "large_pos_zero": lambda: (0.0, torch.randn(3, 3) * 1e6 + 1e7), + "large_neg_zero": lambda: (0.0, -torch.empty(5).uniform_(1e5, 1e8)), + "small_pos_zero": lambda: (0.0, torch.empty(5).uniform_(1e-8, 1e-5)), + "small_neg_zero": lambda: (0.0, -torch.empty(5).uniform_(1e-8, 1e-5)), +} + + +class Elu(nn.Module): + aten_op = "torch.ops.aten.elu.default" + exir_op = "executorch_exir_dialects_edge__ops_aten__elu_default" + + def __init__(self, input_alpha: float = 1.0): + super().__init__() + self.elu = torch.nn.ELU(alpha=input_alpha) + + def forward(self, input_: torch.Tensor): + return self.elu(input_) + + +input_t1 = Tuple[torch.Tensor] + + +@common.parametrize("test_module", test_data_suite) +def test_elu_tosa_FP(test_module: input_t1): + alpha, test_data = test_module() + pipeline = TosaPipelineFP[input_t1]( + Elu(alpha), (test_data,), aten_op=Elu.aten_op, exir_op=Elu.exir_op + ) + pipeline.run() + + +@common.parametrize("test_module", test_data_suite) +def test_elu_tosa_INT(test_module: input_t1): + alpha, test_data = test_module() + pipeline = TosaPipelineINT[input_t1]( + Elu(alpha), (test_data,), aten_op=Elu.aten_op, exir_op=Elu.exir_op + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +@common.parametrize("test_module", test_data_suite) +def test_elu_u55_INT(test_module: input_t1): + alpha, test_data = test_module() + pipeline = EthosU55PipelineINT[input_t1]( + Elu(alpha), (test_data,), aten_ops=Elu.aten_op, exir_ops=Elu.exir_op + ) + pipeline.run() + + +@common.XfailIfNoCorstone320 +@common.parametrize("test_module", test_data_suite) +def test_elu_u85_INT(test_module: input_t1): + alpha, test_data = test_module() + pipeline = EthosU85PipelineINT[input_t1]( + Elu(alpha), (test_data,), aten_ops=Elu.aten_op, exir_ops=Elu.exir_op + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("test_module", test_data_suite) +def test_elu_vgf_FP(test_module: input_t1): + alpha, test_data = test_module() + pipeline = VgfPipeline[input_t1]( + Elu(alpha), + (test_data,), + aten_op=Elu.aten_op, + exir_op=Elu.exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@common.parametrize("test_module", test_data_suite) +def test_elu_vgf_INT(test_module: input_t1): + alpha, test_data = test_module() + pipeline = VgfPipeline[input_t1]( + Elu(alpha), + (test_data,), + aten_op=Elu.aten_op, + exir_op=Elu.exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_embedding.py b/backends/arm/test/ops/test_embedding.py index b0a4647c3ae..901fbbc0916 100644 --- a/backends/arm/test/ops/test_embedding.py +++ b/backends/arm/test/ops/test_embedding.py @@ -6,8 +6,9 @@ from typing import Tuple +import pytest import torch -from executorch.backends.arm._passes import InsertCastForOpsWithInt64InputPass +from executorch.backends.arm._passes import InsertInt32CastsAfterInt64PlaceholdersPass from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -57,6 +58,7 @@ def forward(self, weights: torch.Tensor, indices: torch.Tensor): } +@pytest.mark.skip(reason="MLETORCH-1274 Improve data type checks during partitioning") @common.parametrize("test_input", test_input) def test_embedding_tosa_FP(test_input: input_params): op = Embedding() @@ -66,7 +68,7 @@ def test_embedding_tosa_FP(test_input: input_params): op.aten_op, op.exir_op, use_to_edge_transform_and_lower=True, - transform_passes=[InsertCastForOpsWithInt64InputPass()], + transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) pipeline.run() @@ -87,6 +89,7 @@ def test_embedding_tosa_INT(test_input: input_params): pipeline.run() +@pytest.mark.skip("reason=MLETORCH-1274 Improve data type checks during partitioning") @common.parametrize("test_input", test_input) @common.SkipIfNoModelConverter def test_embedding_vgf_FP(test_input: input_params): @@ -98,7 +101,7 @@ def test_embedding_vgf_FP(test_input: input_params): op.exir_op, tosa_version="TOSA-1.0+FP", use_to_edge_transform_and_lower=True, - transform_passes=[InsertCastForOpsWithInt64InputPass()], + transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) pipeline.run() diff --git a/backends/arm/test/ops/test_expand.py b/backends/arm/test/ops/test_expand.py index 607d8650946..b5784c9ff93 100644 --- a/backends/arm/test/ops/test_expand.py +++ b/backends/arm/test/ops/test_expand.py @@ -10,20 +10,21 @@ from typing import Sequence, Tuple -import pytest - import torch from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, ) aten_op = "torch.ops.aten.expand.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_expand_copy_default" + input_t1 = Tuple[torch.Tensor, torch.Tensor] # Input x, Input y @@ -48,7 +49,7 @@ def forward(self, x: torch.Tensor, m: Sequence): return x.expand(m) -@common.parametrize("test_data", Expand.test_parameters | Expand.test_reject_set) +@common.parametrize("test_data", Expand.test_parameters) def test_expand_tosa_FP(test_data: Tuple): pipeline = TosaPipelineFP[input_t1]( Expand(), @@ -59,7 +60,7 @@ def test_expand_tosa_FP(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Expand.test_parameters | Expand.test_reject_set) +@common.parametrize("test_data", Expand.test_parameters) def test_expand_tosa_INT(test_data: Tuple): pipeline = TosaPipelineINT[input_t1]( Expand(), @@ -70,14 +71,7 @@ def test_expand_tosa_INT(test_data: Tuple): pipeline.run() -x_fails = { - "rand_batch_2": "AssertionError: Output 0 does not match reference output.", - "rand_mix_neg": "AssertionError: Output 0 does not match reference output.", - "rand_small_neg": "AssertionError: Output 0 does not match reference output.", -} - - -@common.parametrize("test_data", Expand.test_parameters, x_fails) +@common.parametrize("test_data", Expand.test_parameters) @common.XfailIfNoCorstone300 def test_expand_u55_INT(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( @@ -90,7 +84,7 @@ def test_expand_u55_INT(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Expand.test_parameters, x_fails) +@common.parametrize("test_data", Expand.test_parameters) @common.XfailIfNoCorstone320 def test_expand_u85_INT(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( @@ -103,7 +97,7 @@ def test_expand_u85_INT(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Expand.test_parameters | Expand.test_reject_set) +@common.parametrize("test_data", Expand.test_parameters) @common.SkipIfNoModelConverter def test_expand_vgf_FP(test_data: Tuple): pipeline = VgfPipeline[input_t1]( @@ -116,7 +110,7 @@ def test_expand_vgf_FP(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", Expand.test_parameters | Expand.test_reject_set) +@common.parametrize("test_data", Expand.test_parameters) @common.SkipIfNoModelConverter def test_expand_vgf_INT(test_data: Tuple): pipeline = VgfPipeline[input_t1]( @@ -130,32 +124,8 @@ def test_expand_vgf_INT(test_data: Tuple): @common.parametrize("test_data", Expand.test_reject_set) -@common.XfailIfNoCorstone300 -@pytest.mark.xfail( - reason="MLETORCH-716: Node will be optimized away and Vela can't handle empty graphs" -) -def test_expand_u55_INT_failure_set(test_data: Tuple): - pipeline = EthosU55PipelineINT[input_t1]( - Expand(), - test_data(), - aten_op, - exir_ops=[], - run_on_fvp=True, - ) - pipeline.run() - - -@common.parametrize("test_data", Expand.test_reject_set) -@common.XfailIfNoCorstone320 -@pytest.mark.xfail( - reason="MLETORCH-716: Node will be optimized away and Vela can't handle empty graphs" -) -def test_expand_u85_INT_failure_set(test_data: Tuple): - pipeline = EthosU85PipelineINT[input_t1]( - Expand(), - test_data(), - aten_op, - exir_ops=[], - run_on_fvp=True, +def test_expand_u55_INT_not_delegated(test_data: Tuple): + pipeline = OpNotSupportedPipeline[input_t1]( + Expand(), test_data(), {exir_op: 1}, n_expected_delegates=0 ) pipeline.run() diff --git a/backends/arm/test/ops/test_expm1.py b/backends/arm/test/ops/test_expm1.py new file mode 100644 index 00000000000..dad95b24f7b --- /dev/null +++ b/backends/arm/test/ops/test_expm1.py @@ -0,0 +1,113 @@ +# Copyright 2025 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 ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +aten_op = "torch.ops.aten.expm1.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_expm1_default" + +input_t1 = Tuple[torch.Tensor] + +test_data_suite = { + "zeroes": torch.zeros(1, 10, 10, 10), + "ones": torch.ones(10, 2, 3), + "rand": torch.rand(10, 10) - 0.5, + "near_zero": torch.randn(100) * 0.01, + "taylor_small": torch.empty(5).uniform_( + -0.35, 0.35 + ), # test cases for taylor series expansion + "randn_large_pos": torch.randn(10) + 10, + "randn_large_neg": torch.randn(10) - 10, + "ramp": torch.arange(-16, 16, 0.2), +} + + +class Expm1(torch.nn.Module): + + def forward(self, x: torch.Tensor): + return torch.expm1(x) + + +@common.parametrize("test_data", test_data_suite) +def test_expm1_tosa_FP(test_data: Tuple): + pipeline = TosaPipelineFP[input_t1]( + Expm1(), + (test_data,), + aten_op=aten_op, + exir_op=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +def test_expm1_tosa_INT(test_data: Tuple): + pipeline = TosaPipelineINT[input_t1]( + Expm1(), + (test_data,), + aten_op=aten_op, + exir_op=exir_op, + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +@common.parametrize("test_data", test_data_suite) +def test_expm1_u55_INT(test_data: Tuple): + pipeline = EthosU55PipelineINT[input_t1]( + Expm1(), + (test_data,), + aten_ops=aten_op, + exir_ops=exir_op, + ) + pipeline.run() + + +@common.XfailIfNoCorstone320 +@common.parametrize("test_data", test_data_suite) +def test_expm1_u85_INT(test_data: Tuple): + pipeline = EthosU85PipelineINT[input_t1]( + Expm1(), + (test_data,), + aten_ops=aten_op, + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_expm1_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Expm1(), + (test_data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_expm1_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Expm1(), + (test_data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_eye.py b/backends/arm/test/ops/test_eye.py index 48f93379fc0..eef32259c10 100644 --- a/backends/arm/test/ops/test_eye.py +++ b/backends/arm/test/ops/test_eye.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. +import pytest import torch from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -38,16 +39,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ), } - test_data_not_delegated: dict[str, test_data_t] = { + # Mixed dtypes - the eye op is delegated, but it leads to a non-delegated add op. + test_data_mixed_dtypes: dict[str, test_data_t] = { "fp32_int64": (lambda: (torch.randn(10),), (10, torch.int64)), "fp32_int32": (lambda: (torch.randn(10),), (10, torch.int32)), - "int32_int64": ( - lambda: (torch.randint(0, 10, [10], dtype=torch.int32),), - (10, torch.int64), - ), } +# skip test since int32 isn't support on FP profile +# "int32_int32": "view/RESHAPE of integer tensor is not supported for +FP profile" +@pytest.mark.skip(reason="MLETORCH-1274 Improve data type checks during partitioning") @common.parametrize("test_data", EyeAdd.test_data) def test_eye_tosa_FP(test_data: test_data_t): input_data, init_data = test_data @@ -59,7 +60,7 @@ def test_eye_tosa_FP(test_data: test_data_t): pipeline.run() -@common.parametrize("test_data", EyeAdd.test_data) +@common.parametrize("test_data", EyeAdd.test_data | EyeAdd.test_data_mixed_dtypes) def test_eye_tosa_INT(test_data: test_data_t): input_data, init_data = test_data pipeline = TosaPipelineINT[input_t]( @@ -99,6 +100,9 @@ def test_eye_u85_INT(test_data: test_data_t): pipeline.run() +# skip since int32 isn't support on FP profile +# "int32_int32": "view/RESHAPE of integer tensor is not supported for +FP profile" +@pytest.mark.skip(reason="MLETORCH-1274 Improve data type checks during partitioning") @common.parametrize( "test_data", EyeAdd.test_data, @@ -134,16 +138,15 @@ def test_eye_vgf_INT(test_data: test_data_t): @common.parametrize( "test_data", - EyeAdd.test_data_not_delegated, - xfails={ - "fp32_int32": "MLETORCG-716: Do not delegate empty networks to vela", - "fp32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - "int32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - }, + EyeAdd.test_data_mixed_dtypes, ) def test_eye_tosa_INT_not_delegated(test_data: test_data_t): input_data, init_data = test_data pipeline = OpNotSupportedPipeline[input_t]( - EyeAdd(*init_data), input_data(), non_delegated_ops={}, quantize=True + EyeAdd(*init_data), + input_data(), + non_delegated_ops={"executorch_exir_dialects_edge__ops_aten_add_Tensor": 1}, + n_expected_delegates=1, + quantize=True, ) pipeline.run() diff --git a/backends/arm/test/ops/test_ge.py b/backends/arm/test/ops/test_ge.py index c66f6d164b9..94f33d28630 100644 --- a/backends/arm/test/ops/test_ge.py +++ b/backends/arm/test/ops/test_ge.py @@ -153,7 +153,6 @@ def test_ge_scalar_u55_INT(test_module): @common.parametrize( "test_module", test_data_tensor, - xfails={"ge_tensor_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85"}, ) @common.XfailIfNoCorstone320 def test_ge_tensor_u85_INT(test_module): @@ -170,7 +169,6 @@ def test_ge_tensor_u85_INT(test_module): @common.parametrize( "test_module", test_data_scalar, - xfails={"ge_scalar_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85"}, ) @common.XfailIfNoCorstone320 def test_ge_scalar_u85_INT(test_module): diff --git a/backends/arm/test/ops/test_glu.py b/backends/arm/test/ops/test_glu.py new file mode 100644 index 00000000000..c19fb892c92 --- /dev/null +++ b/backends/arm/test/ops/test_glu.py @@ -0,0 +1,130 @@ +# Copyright 2025 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 +import torch.nn.functional as F +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +aten_op = "torch.ops.aten.glu.default" +exir_op = "executorch_exir_dialects_edge__ops_aten__glu_default" + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite = { + "zeros": [torch.zeros(10, 10, 2), -1], + "ones": [torch.ones(10, 10, 2), -1], + "rand": [torch.rand(10, 10, 2) - 0.5, -1], + "randn_pos": [torch.randn(10, 2) + 10, -1], + "randn_neg": [torch.randn(10, 2) - 10, -1], + "ramp": [torch.linspace(-16, 15.8, 160).reshape(-1, 2), -1], + "zeros_custom_dim": [torch.zeros(7, 10, 5), 1], + "rand_custom_dim": [torch.rand(10, 3, 3) - 0.5, 0], +} + + +class Glu(torch.nn.Module): + + def forward(self, a: torch.Tensor, dim: int) -> torch.Tensor: + return F.glu(a, dim=dim) + + +@common.parametrize( + "test_data", + test_data_suite, +) +def test_glu_tosa_FP(test_data: Tuple): + pipeline = TosaPipelineFP[input_t1]( + Glu(), + (*test_data,), + aten_op, + exir_op, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +def test_glu_tosa_INT(test_data: Tuple): + pipeline = TosaPipelineINT[input_t1]( + Glu(), + (*test_data,), + aten_op=[], + exir_op=exir_op, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.XfailIfNoCorstone300 +def test_glu_u55_INT(test_data: Tuple): + pipeline = EthosU55PipelineINT[input_t1]( + Glu(), + (*test_data,), + aten_ops=[], + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.XfailIfNoCorstone320 +def test_glu_u85_INT(test_data: Tuple): + pipeline = EthosU85PipelineINT[input_t1]( + Glu(), + (*test_data,), + aten_ops=[], + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.SkipIfNoModelConverter +def test_glu_vgf_FP(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Glu(), + (*test_data,), + [], + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.SkipIfNoModelConverter +def test_glu_vgf_INT(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Glu(), + (*test_data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_group_norm.py b/backends/arm/test/ops/test_group_norm.py index 5fa4cd328de..0f314064548 100644 --- a/backends/arm/test/ops/test_group_norm.py +++ b/backends/arm/test/ops/test_group_norm.py @@ -102,6 +102,9 @@ def test_native_group_norm_tosa_INT(test_data): "test_data", test_data_suite, xfails={ + "rand_4_6_8_groups_2_eps_no_affine": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", + "rand_4_6_groups_1": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", + "rand_4_6_groups_2": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", "randn_1_12_8_6_groups_12": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", "rand_6_8_10_12_groups_1": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", "rand_6_8_10_12_groups_4_no_affine": "MLETORCH-925: Fix numerical issue for aten.native_group_norm", diff --git a/backends/arm/test/ops/test_gt.py b/backends/arm/test/ops/test_gt.py index 83c85e5f9fc..41229397eb5 100644 --- a/backends/arm/test/ops/test_gt.py +++ b/backends/arm/test/ops/test_gt.py @@ -154,9 +154,6 @@ def test_gt_scalar_u55_INT(test_module): @common.parametrize( "test_module", test_data_tensor, - xfails={ - "gt_tensor_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85", - }, ) @common.XfailIfNoCorstone320 def test_gt_tensor_u85_INT(test_module): @@ -173,9 +170,6 @@ def test_gt_tensor_u85_INT(test_module): @common.parametrize( "test_module", test_data_scalar, - xfails={ - "gt_scalar_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85", - }, ) @common.XfailIfNoCorstone320 def test_gt_scalar_u85_INT(test_module): diff --git a/backends/arm/test/ops/test_index_select.py b/backends/arm/test/ops/test_index_select.py index bb28d66f7cf..95ebaa62a38 100644 --- a/backends/arm/test/ops/test_index_select.py +++ b/backends/arm/test/ops/test_index_select.py @@ -12,6 +12,7 @@ from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, @@ -120,6 +121,20 @@ def test_index_select_tosa_INT_rand(test_data: input_params): pipeline.run() +@pytest.mark.parametrize("test_data", list(test_data.values())[-1:]) +def test_index_select_u55_INT_not_delegated(test_data: input_params): + op, test_input = test_data + + pipeline = OpNotSupportedPipeline[input_params]( + op, + test_input, + {op.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + @pytest.mark.parametrize("test_data", list(test_data.values())) @common.SkipIfNoModelConverter def test_index_select_vgf_FP(test_data: input_params): diff --git a/backends/arm/test/ops/test_index_tensor.py b/backends/arm/test/ops/test_index_tensor.py index 37ed0e131a4..557846922b8 100644 --- a/backends/arm/test/ops/test_index_tensor.py +++ b/backends/arm/test/ops/test_index_tensor.py @@ -10,6 +10,7 @@ import torch from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, ) @@ -460,3 +461,18 @@ def test_index_tensor_tosa_INT_none(test_data: input_params): IndexTensorTestCommon.exir_op, ).run() ) + + +@common.parametrize("test_data", IndexTensor.test_data) +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_not_delegated(test_data: input_params): + """Ethos-U55 backend BI pipeline test for index.Tensor""" + test_input = test_data + with torch.no_grad(): + OpNotSupportedPipeline[input_params]( + IndexTensor(), + test_input, + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ).run() diff --git a/backends/arm/test/ops/test_le.py b/backends/arm/test/ops/test_le.py index 6cb185ecb92..31422302a2d 100644 --- a/backends/arm/test/ops/test_le.py +++ b/backends/arm/test/ops/test_le.py @@ -155,9 +155,6 @@ def test_le_scalar_u55_INT_not_delegated(test_module): @common.parametrize( "test_module", test_data_tensor, - xfails={ - "le_tensor_rank4_randn": "4D fails because boolean Tensors can't be subtracted" - }, ) @common.XfailIfNoCorstone320 def test_le_tensor_u85_INT(test_module): @@ -175,9 +172,6 @@ def test_le_tensor_u85_INT(test_module): @common.parametrize( "test_module", test_data_scalar, - xfails={ - "le_scalar_rank4_randn": "4D fails because boolean Tensors can't be subtracted" - }, ) @common.XfailIfNoCorstone320 def test_le_scalar_u85_INT(test_module): diff --git a/backends/arm/test/ops/test_leaky_relu.py b/backends/arm/test/ops/test_leaky_relu.py index c18255a73c0..432c4da7ecc 100644 --- a/backends/arm/test/ops/test_leaky_relu.py +++ b/backends/arm/test/ops/test_leaky_relu.py @@ -30,10 +30,10 @@ def forward(self, x: torch.Tensor): test_data: dict[str, input_t1] = { "zeros": lambda: ((torch.zeros(1, 1, 5, 5),), 0.01), - "ones": lambda: ((torch.ones(1, 32, 112, 112),), 0.01), - "rand": lambda: ((torch.rand(1, 96, 56, 56),), 0.2), + "ones": lambda: ((torch.ones(1, 16, 96, 96),), 0.01), + "rand": lambda: ((torch.rand(1, 64, 56, 56),), 0.2), "3Dtensor": lambda: ((torch.rand(5, 5, 5),), 0.001), - "negative_slope": lambda: ((torch.rand(1, 16, 128, 128),), -0.002), + "negative_slope": lambda: ((torch.rand(1, 16, 96, 96),), -0.002), } diff --git a/backends/arm/test/ops/test_linear.py b/backends/arm/test/ops/test_linear.py index 57ce490dae8..e5d00c83e9f 100644 --- a/backends/arm/test/ops/test_linear.py +++ b/backends/arm/test/ops/test_linear.py @@ -9,9 +9,12 @@ from typing import Tuple import pytest - import torch -from executorch.backends.arm.test import common +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, @@ -20,6 +23,8 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.linear.default" @@ -143,7 +148,6 @@ def test_linear_tosa_FP(test_data: torch.Tensor): pipeline.run() -@pytest.mark.flaky(reruns=5) # TODO: Investigate flakyness. @common.parametrize("test_data", test_data_rank1_INT | test_data_rank4_INT) def test_linear_tosa_INT(test_data: torch.Tensor): test_data, out_features, has_bias, per_channel_quantization = test_data() @@ -183,24 +187,9 @@ def test_linear_u55_INT(test_data: torch.Tensor): ).run() -x_fail = { - f"{k},per_channel_quant={q}": reason - for k, reason in { - "model_linear_rank4_zeros": "AssertionError: Output 0 does not match reference output.", - "model_linear_rank4_ones": "AssertionError: Output 0 does not match reference output.", - "model_linear_rank4_negative_ones": "AssertionError: Output 0 does not match reference output.", - "model_linear_rank4_rand": "AssertionError: Output 0 does not match reference output.", - "model_linear_rank4_negative_large_rand": "AssertionError: Output 0 does not match reference output.", - "model_linear_rank4_large_randn": "AssertionError: Output 0 does not match reference output.", - }.items() - for q in [True, False] -} - - @common.parametrize( "test_data", test_data_rank1_INT | test_data_rank4_INT, - x_fail, ) @common.XfailIfNoCorstone320 def test_linear_u85_INT(test_data: torch.Tensor): @@ -258,3 +247,64 @@ def test_linear_vgf_INT(test_data: torch.Tensor): per_channel_quantization=per_channel_quantization, ) pipeline.run() + + +def get_symmetric_a16w8_linear_quantizer( + u55_config=False, per_channel_quantization=False +): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + quantizer.set_module_type( + torch.nn.Linear, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_rank1_INT | test_data_rank4_INT) +@pytest.mark.xfail( + reason="missing int16 linear ops support; fails at TOSA reference model run with Invalid TOSA graph" +) +def test_linear_16a8w_tosa_INT(test_data: torch.Tensor): + """Test linear operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + test_data, out_features, has_bias, per_channel_quantization = test_data() + in_features = test_data.shape[-1] + + # Create pipeline with custom 16A8W quantization config + pipeline = TosaPipelineINT[input_t1]( + Linear( + in_features=in_features, + out_features=out_features, + bias=has_bias, + ), + (test_data,), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_linear_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + # Run the pipeline + pipeline.run() diff --git a/backends/arm/test/ops/test_logical.py b/backends/arm/test/ops/test_logical.py index 2b160ce7b50..bb7c5773342 100644 --- a/backends/arm/test/ops/test_logical.py +++ b/backends/arm/test/ops/test_logical.py @@ -86,6 +86,9 @@ def forward(self, tensor: torch.Tensor): ################# +xfails = {"rand_rank4": "MLBEDSW-11031: Output diff on u85 bool transpose."} + + @common.parametrize("test_data", And().test_data) def test_logical_and_tosa_FP(test_data: input_t2): pipeline = TosaPipelineFP[input_t2]( @@ -129,7 +132,7 @@ def test_logical_and_u55_INT_not_delegated(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", And().test_data) +@common.parametrize("test_data", And().test_data, xfails=xfails) @common.XfailIfNoCorstone320 def test_logical_and_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -223,7 +226,7 @@ def test_logical_xor_u55_INT_not_delegated(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", Xor().test_data) +@common.parametrize("test_data", Xor().test_data, xfails=xfails) @common.XfailIfNoCorstone320 def test_logical_xor_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -317,7 +320,7 @@ def test_logical_or_u55_INT_not_delegated(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", Or().test_data) +@common.parametrize("test_data", Or().test_data, xfails=xfails) @common.XfailIfNoCorstone320 def test_logical_or_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( @@ -411,7 +414,7 @@ def test_logical_not_u55_INT_not_delegated(test_data: input_t2): pipeline.run() -@common.parametrize("test_data", Not().test_data) +@common.parametrize("test_data", Not().test_data, xfails=xfails) @common.XfailIfNoCorstone320 def test_logical_not_u85_INT(test_data: input_t2): pipeline = EthosU85PipelineINT[input_t2]( diff --git a/backends/arm/test/ops/test_logit.py b/backends/arm/test/ops/test_logit.py new file mode 100644 index 00000000000..8915c151bb9 --- /dev/null +++ b/backends/arm/test/ops/test_logit.py @@ -0,0 +1,119 @@ +# Copyright 2025 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 ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +aten_op = "torch.ops.aten.logit.default" +exir_op = "executorch_exir_dialects_edge__ops_aten__logit_default" + +input_t1 = Tuple[torch.Tensor] + +test_data_suite = { + "zeros": [torch.zeros((10, 10, 10)), None], + "ones": [torch.ones((10, 10, 10)), None], + "uniform_valid": [torch.rand((10, 10, 10)), None], + "near_zero": [torch.full((10, 10), 1e-8), None], + "near_one": [torch.full((10, 10), 1 - 1e-8), None], + "mixed": [torch.tensor([0.0, 1e-5, 0.5, 1 - 1e-5, 1.0]), None], + "multi_dim": [torch.rand((2, 3, 4)), None], + "eps": [torch.zeros((10, 10, 10)), 1e-6], + "invalid_neg": [torch.full((5,), -0.1), 1e-6], + "invalid_gt1": [torch.full((5,), 1.1), 1e-6], +} + + +class Logit(torch.nn.Module): + + def forward(self, x: torch.Tensor, eps: torch.float32): + return torch.logit(x, eps=eps) + + +@common.parametrize("test_data", test_data_suite) +def test_logit_tosa_FP(test_data: Tuple): + pipeline = TosaPipelineFP[input_t1]( + Logit(), + (*test_data,), + aten_op=aten_op, + exir_op=exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +def test_logit_tosa_INT(test_data: Tuple): + pipeline = TosaPipelineINT[input_t1]( + Logit(), + (*test_data,), + aten_op=[], + exir_op=exir_op, + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +@common.parametrize("test_data", test_data_suite) +def test_logit_u55_INT(test_data: Tuple): + pipeline = EthosU55PipelineINT[input_t1]( + Logit(), + (*test_data,), + aten_ops=[], + exir_ops=exir_op, + ) + pipeline.run() + + +@common.XfailIfNoCorstone320 +@common.parametrize("test_data", test_data_suite) +def test_logit_u85_INT(test_data: Tuple): + pipeline = EthosU85PipelineINT[input_t1]( + Logit(), + (*test_data,), + aten_ops=[], + exir_ops=exir_op, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.SkipIfNoModelConverter +def test_logit_vgf_FP(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Logit(), + (*test_data,), + [], + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + test_data_suite, +) +@common.SkipIfNoModelConverter +def test_logit_vgf_INT(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Logit(), + (*test_data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_logsoftmax.py b/backends/arm/test/ops/test_logsoftmax.py index b1b934fbcc8..791069aa4b0 100644 --- a/backends/arm/test/ops/test_logsoftmax.py +++ b/backends/arm/test/ops/test_logsoftmax.py @@ -68,7 +68,7 @@ def test_log_softmax_tosa_INT(test_data): "test_data", LogSoftmax.test_data, xfails={ - "randn_mult_batches": "MLETORCH-433: Multiple batches not supported on FVP" + "randn_neg_dim": "MLBEDSW-11032: ILLEGAL_OFM_BASE error: Base addresses must be aligned to brick depth on u55." }, ) @common.XfailIfNoCorstone300() @@ -85,13 +85,7 @@ def test_log_softmax_u55_INT(test_data): pipeline.run() -@common.parametrize( - "test_data", - LogSoftmax.test_data, - xfails={ - "randn_mult_batches": "MLETORCH-433: Multiple batches not supported on FVP" - }, -) +@common.parametrize("test_data", LogSoftmax.test_data) @common.XfailIfNoCorstone320 def test_log_softmax_u85_INT(test_data): data, dim = test_data() diff --git a/backends/arm/test/ops/test_lt.py b/backends/arm/test/ops/test_lt.py index 86d903e3f88..98d0298b195 100644 --- a/backends/arm/test/ops/test_lt.py +++ b/backends/arm/test/ops/test_lt.py @@ -154,9 +154,6 @@ def test_lt_scalar_u55_INT_not_delegated(test_module): @common.parametrize( "test_module", test_data_tensor, - xfails={ - "lt_tensor_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85", - }, ) @common.XfailIfNoCorstone320 def test_lt_tensor_u85_INT(test_module): @@ -173,9 +170,6 @@ def test_lt_tensor_u85_INT(test_module): @common.parametrize( "test_module", test_data_scalar, - xfails={ - "lt_scalar_rank4_randn": "MLETORCH-847: Boolean eq result unstable on U85", - }, ) @common.XfailIfNoCorstone320 def test_lt_scalar_u85_INT(test_module): diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py index 6b75c2b7d0a..7db56311837 100644 --- a/backends/arm/test/ops/test_max_pool.py +++ b/backends/arm/test/ops/test_max_pool.py @@ -65,12 +65,10 @@ torch.rand(1, 16, 54, 54), [3, (1, 3), 1], ), -} - -test_data_suite_mult_batches = { "randn": lambda: (torch.randn(5, 16, 50, 32), [4, 2, 0]), } + test_data_suite_dilation = [ # Simple dilation=2 on 8x8 input, kernel=3, stride=1, no padding ("dilation2", torch.rand(1, 1, 8, 8), [3, 1, 0, 2]), @@ -161,61 +159,6 @@ def test_max_pool2d_u85_INT(test_data: torch.Tensor): ).run() -@common.parametrize("test_data", test_data_suite_mult_batches) -def test_max_pool2d_tosa_FP_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - pipeline = TosaPipelineFP[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_op, - ) - pipeline.run() - - -@common.parametrize("test_data", test_data_suite_mult_batches) -def test_max_pool2d_tosa_INT_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - pipeline = TosaPipelineINT[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_op, - ) - pipeline.run() - - -x_fail = {"randn": "MLETORCH-986: Numerical issues with mutli batches."} - - -@common.parametrize("test_data", test_data_suite_mult_batches, x_fail) -@common.XfailIfNoCorstone300 -def test_max_pool2d_u55_INT_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - EthosU55PipelineINT[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_ops=[], - run_on_fvp=True, - use_to_edge_transform_and_lower=True, - ).run() - - -@common.parametrize("test_data", test_data_suite_mult_batches, x_fail) -@common.XfailIfNoCorstone320 -def test_max_pool2d_u85_INT_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - EthosU85PipelineINT[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_op, - run_on_fvp=True, - use_to_edge_transform_and_lower=True, - ).run() - - reject_data_suite = { "reject_1": lambda: (MaxPool2d(1, 4, 0), torch.rand(1, 10, 10, 10)), "reject_2": lambda: (MaxPool2d((1, 257), 1, 0), torch.rand(1, 16, 5, 300)), @@ -306,34 +249,6 @@ def test_max_pool2d_vgf_INT(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite_mult_batches) -@common.SkipIfNoModelConverter -def test_max_pool2d_vgf_FP_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - pipeline = VgfPipeline[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_op, - tosa_version="TOSA-1.0+FP", - ) - pipeline.run() - - -@common.parametrize("test_data", test_data_suite_mult_batches) -@common.SkipIfNoModelConverter -def test_max_pool2d_vgf_INT_mult_batches(test_data: torch.Tensor): - test_data, model_params = test_data() - pipeline = VgfPipeline[input_t1]( - MaxPool2d(*model_params), - (test_data,), - aten_op, - exir_op, - tosa_version="TOSA-1.0+INT", - ) - pipeline.run() - - @common.parametrize("test_data", dilation_test_data) @common.SkipIfNoModelConverter def test_max_pool2d_vgf_FP_dilation(test_data: torch.Tensor): diff --git a/backends/arm/test/ops/test_mean_dim.py b/backends/arm/test/ops/test_mean_dim.py index 1483b5d82b6..061e8da14f1 100644 --- a/backends/arm/test/ops/test_mean_dim.py +++ b/backends/arm/test/ops/test_mean_dim.py @@ -280,6 +280,7 @@ def test_mean_dim_tosa_INT(test_data): (test_data,), [], # Might be sum, avgpool, or both symmetric_io_quantization=True, + custom_path="MEANDIM", ) pipeline.run() diff --git a/backends/arm/test/ops/test_mul.py b/backends/arm/test/ops/test_mul.py index b0b7f5f4b7d..b2db55d90fd 100644 --- a/backends/arm/test/ops/test_mul.py +++ b/backends/arm/test/ops/test_mul.py @@ -8,9 +8,14 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -18,6 +23,9 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification + +from executorch.backends.xnnpack.test.tester import Quantize input_t1 = Tuple[torch.Tensor, torch.Tensor] # Input x aten_op = "torch.ops.aten.mul.Tensor" @@ -129,7 +137,17 @@ def test_mul_tensor_tosa_FP_diff_input_ranks(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite_int32) +# MLETORCH-1274 Improve data type checks during partitioning +# view/RESHAPE of integer tensor is not supported for +FP profile which causes issues +# with view_copy (RESHAPE) which isn't supported in FP so removing the int32 tests +# to allow for the dtype validation patches to land. +# filter out the 'op_mul_rank4_randn_int32' only +test_data_int32_without_broadcasting = { + k: v for k, v in test_data_suite_int32.items() if k != "op_mul_rank4_randn_int32" +} + + +@common.parametrize("test_data", test_data_int32_without_broadcasting) def test_mul_tensor_tosa_FP_int32(test_data: torch.Tensor): pipeline = TosaPipelineFP[input_t1]( Mul(), @@ -200,15 +218,7 @@ def test_mul_tensor_u85_INT(test_data: torch.Tensor): pipeline.run() -@common.parametrize( - "test_data", - test_data_suite_int32, - xfails={ - # TODO: MLETORCH-1132 Investigate why tests with inputs that require broadcasting fail on u55/u85 - "op_mul_rank4_randn_mutltiple_broadcasts_int32": "RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Int", - "op_mul_rank4_randn_broadcast_int32": "RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Int", - }, -) +@common.parametrize("test_data", test_data_suite_int32) @common.XfailIfNoCorstone300 def test_mul_tensor_u55_INT_int32(test_data: torch.Tensor): pipeline = EthosU55PipelineINT[input_t1]( @@ -222,15 +232,7 @@ def test_mul_tensor_u55_INT_int32(test_data: torch.Tensor): pipeline.run() -@common.parametrize( - "test_data", - test_data_suite_int32, - xfails={ - # TODO: MLETORCH-1132 Investigate why tests with inputs that require broadcasting fail on u55/u85 - "op_mul_rank4_randn_mutltiple_broadcasts_int32": "RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Int", - "op_mul_rank4_randn_broadcast_int32": "RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: Int", - }, -) +@common.parametrize("test_data", test_data_suite_int32) @common.XfailIfNoCorstone320 def test_mul_tensor_u85_INT_int32(test_data: torch.Tensor): pipeline = EthosU85PipelineINT[input_t1]( @@ -244,8 +246,14 @@ def test_mul_tensor_u85_INT_int32(test_data: torch.Tensor): pipeline.run() +# view/RESHAPE of integer tensor is not supported for +FP profile which causes issues +# with view_copy (RESHAPE) which isn't supported in FP so removing the int32 tests +# to allow for the dtype validation patches to land. + + @common.parametrize( - "test_data", test_data_suite | test_data_suite_2 | test_data_suite_int32 + "test_data", + test_data_suite | test_data_suite_2 | test_data_int32_without_broadcasting, ) @common.SkipIfNoModelConverter def test_mul_tensor_vgf_FP(test_data: torch.Tensor): @@ -284,3 +292,105 @@ def test_mul_tensor_vgf_INT_int32(test_data: torch.Tensor): ) pipeline.pop_stage("check.quant_nodes") pipeline.run() + + +def get_symmetric_a16w8_mul_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +@pytest.mark.xfail( + reason="missing int16 mul ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13947" +) +def test_mul_tensor_16a8w_tosa_INT(test_data: input_t1): + """Test mul operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Mul(), + test_data(), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_mul_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 mul operations. See: https://github.com/pytorch/executorch/issues/13947" +) +def test_mul_tensor_16a8w_u55_INT16(test_data: input_t1): + """Test mul operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Mul(), + test_data(), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_mul_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 mul operations. See: https://github.com/pytorch/executorch/issues/13947" +) +def test_mul_tensor_16a8w_u85_INT16(test_data: input_t1): + """Test mul operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Mul(), + test_data(), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_mul_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_ones.py b/backends/arm/test/ops/test_ones.py index 18204a8eaaa..f4dafca5e10 100644 --- a/backends/arm/test/ops/test_ones.py +++ b/backends/arm/test/ops/test_ones.py @@ -39,17 +39,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ), } - test_data_not_delegated: dict[str, test_data_t] = { + # Mixed dtypes - the ones op is delegated, but it leads to a non-delegated add op. + test_data_mixed_dtypes: dict[str, test_data_t] = { "fp32_int64": (lambda: (torch.randn(10),), (10, torch.int64)), "fp32_int32": (lambda: (torch.randn(10),), (10, torch.int32)), - "int32_int64": ( - lambda: (torch.randint(0, 10, [10], dtype=torch.int32),), - (10, torch.int64), - ), } -@common.parametrize("test_data", OnesAdd.test_data) +@common.parametrize("test_data", OnesAdd.test_data | OnesAdd.test_data_mixed_dtypes) def test_ones_tosa_FP(test_data: test_data_t): input_data, init_data = test_data pipeline = TosaPipelineFP[input_t]( @@ -60,7 +57,7 @@ def test_ones_tosa_FP(test_data: test_data_t): pipeline.run() -@common.parametrize("test_data", OnesAdd.test_data) +@common.parametrize("test_data", OnesAdd.test_data | OnesAdd.test_data_mixed_dtypes) def test_ones_tosa_INT(test_data: test_data_t): input_data, init_data = test_data pipeline = TosaPipelineINT[input_t]( @@ -102,17 +99,16 @@ def test_ones_u85_INT(test_data: test_data_t): @common.parametrize( "test_data", - OnesAdd.test_data_not_delegated, - xfails={ - "fp32_int32": "MLETORCG-716: Do not delegate empty networks to vela", - "fp32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - "int32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - }, + OnesAdd.test_data_mixed_dtypes, ) def test_ones_tosa_INT_not_delegated(test_data: test_data_t): input_data, init_data = test_data pipeline = OpNotSupportedPipeline[input_t]( - OnesAdd(*init_data), input_data(), non_delegated_ops={}, quantize=True + OnesAdd(*init_data), + input_data(), + non_delegated_ops={"executorch_exir_dialects_edge__ops_aten_add_Tensor": 1}, + n_expected_delegates=1, + quantize=True, ) pipeline.run() diff --git a/backends/arm/test/ops/test_permute.py b/backends/arm/test/ops/test_permute.py index 57f7f9603a1..eb482bcee54 100644 --- a/backends/arm/test/ops/test_permute.py +++ b/backends/arm/test/ops/test_permute.py @@ -72,13 +72,11 @@ def test_permute_tosa_INT(test_data: torch.Tensor): pipeline.run() -x_fails = { - "rank_4_2": "AssertionError: Output 0 does not match reference output.", - "rank_4_3": "AssertionError: Output 0 does not match reference output.", -} - - -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize( + "test_data", + test_data_suite, + xfails={"rank_4_3": "MLETORCH-955 : Permutation numerical diff for u55"}, +) @common.XfailIfNoCorstone300 def test_permute_u55_INT(test_data): test_data, dims = test_data() @@ -92,8 +90,7 @@ def test_permute_u55_INT(test_data): pipeline.run() -# Fails since on FVP since N > 1 is not supported. MLETORCH-517 -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone320 def test_permute_u85_INT(test_data: torch.Tensor): test_data, dims = test_data() diff --git a/backends/arm/test/ops/test_scalar_tensor.py b/backends/arm/test/ops/test_scalar_tensor.py index cf3d0818dbc..22c1cc0373d 100644 --- a/backends/arm/test/ops/test_scalar_tensor.py +++ b/backends/arm/test/ops/test_scalar_tensor.py @@ -11,6 +11,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) float_test_data_suite = { @@ -99,3 +100,30 @@ def test_scalar_tensor_u85_INT(test_data): ScalarTensor.aten_op, run_on_fvp=True, ).run() + + +@common.parametrize("test_data", float_test_data_suite) +@common.SkipIfNoModelConverter +def test_scalar_tensor_vgf_FP(test_data): + scalar, dtype, data = test_data() + pipeline = VgfPipeline( + ScalarTensor(scalar, dtype), + tuple(data), + ScalarTensor.aten_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", int_test_data_suite) +@common.SkipIfNoModelConverter +def test_scalar_tensor_vgf_INT(test_data): + scalar, dtype, data = test_data() + pipeline = VgfPipeline( + ScalarTensor(scalar, dtype), + tuple(data), + ScalarTensor.aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() diff --git a/backends/arm/test/ops/test_scalars.py b/backends/arm/test/ops/test_scalars.py index 1243a522526..c4f371a1a14 100644 --- a/backends/arm/test/ops/test_scalars.py +++ b/backends/arm/test/ops/test_scalars.py @@ -417,6 +417,9 @@ def test_div_scalar_u85_INT(): # SHIFT ETHOS-U ------------------------------------------------------ +@pytest.mark.skip( + reason="integer operations (shift and sub) are not supported on FP profile" +) def test_bitwise_right_shift_tensor_tosa_FP_inplace(): pipeline = TosaPipelineFP[input_t1]( ShiftInplaceSub(), diff --git a/backends/arm/test/ops/test_sdpa.py b/backends/arm/test/ops/test_sdpa.py index c4b05972f76..009e4b2ad70 100644 --- a/backends/arm/test/ops/test_sdpa.py +++ b/backends/arm/test/ops/test_sdpa.py @@ -8,9 +8,11 @@ import torch +from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) @@ -27,14 +29,14 @@ def forward(self, query, key, value): input_t = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] -def test_sdpa_FP(): +def test_sdpa_tosa_FP(): test_input = tuple(torch.randn(1, 3, 197, 64) for x in range(3)) pipeline = TosaPipelineFP[input_t](SDPA(), test_input, [], []) pipeline.pop_stage("check_count.exir") pipeline.run() -def test_sdpa_INT(): +def test_sdpa_tosa_INT(): test_input = tuple(torch.randn(1, 3, 197, 64) for x in range(3)) pipeline = TosaPipelineINT[input_t](SDPA(), test_input, [], []) pipeline.pop_stage("check.quant_nodes") @@ -43,3 +45,25 @@ def test_sdpa_INT(): "run_method_and_compare_outputs" ) # TODO: reference is not quantized pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sdpa_vgf_FP(): + test_input = tuple(torch.randn(1, 3, 197, 64) for _ in range(3)) + pipeline = VgfPipeline[input_t]( + SDPA(), test_input, [], [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sdpa_vgf_INT(): + test_input = tuple(torch.randn(1, 3, 197, 64) for _ in range(3)) + pipeline = VgfPipeline[input_t]( + SDPA(), + test_input, + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_select.py b/backends/arm/test/ops/test_select.py index 9cd3cf6f3b7..4c3887f1e18 100644 --- a/backends/arm/test/ops/test_select.py +++ b/backends/arm/test/ops/test_select.py @@ -16,6 +16,7 @@ OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t1 = Tuple[torch.Tensor, int, int] @@ -101,12 +102,7 @@ def test_select_int_tosa_INT(test_data: Tuple): pipeline.run() -x_fails = { - "select4d_0_dim_2_index": "AssertionError: Output 0 does not match reference output." -} - - -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone300 def test_select_int_u55_INT_copy(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( @@ -120,7 +116,7 @@ def test_select_int_u55_INT_copy(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone300 def test_select_int_u55_INT(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( @@ -147,7 +143,7 @@ def test_select_int_u55_INT_not_delegated(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone320 def test_select_int_u85_INT_copy(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( @@ -161,7 +157,7 @@ def test_select_int_u85_INT_copy(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", test_data_suite, x_fails) +@common.parametrize("test_data", test_data_suite) @common.XfailIfNoCorstone320 def test_select_int_u85_INT(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( @@ -173,3 +169,47 @@ def test_select_int_u85_INT(test_data: Tuple): use_to_edge_transform_and_lower=True, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_select_int_vgf_FP_copy(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SelectCopy(), test_data(), aten_op_copy, [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_select_int_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SelectInt(), test_data(), aten_op_int, [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_select_int_vgf_INT_copy(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SelectCopy(), + test_data(), + aten_op_copy, + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_select_int_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SelectInt(), + test_data(), + aten_op_int, + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sigmoid.py b/backends/arm/test/ops/test_sigmoid.py index b4f8458574e..aac2ee1c9b1 100644 --- a/backends/arm/test/ops/test_sigmoid.py +++ b/backends/arm/test/ops/test_sigmoid.py @@ -8,14 +8,22 @@ from typing import Tuple +import pytest import torch -from executorch.backends.arm.test import common +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.sigmoid.default" # Used for checking that we do not have softmax in the graph after decompose exir_op = "executorch_exir_dialects_edge__ops_aten_sigmoid_default" @@ -154,3 +162,203 @@ def test_sigmoid_u85_INT(test_data: Tuple): run_on_fvp=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sigmoid(), + (test_data(),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sigmoid(), + (test_data(),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_FP_add(): + pipeline = VgfPipeline[input_t1]( + AddSigmoid(), + (test_data_suite["zeros"](),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_INT_add(): + pipeline = VgfPipeline[input_t1]( + AddSigmoid(), + (test_data_suite["ramp"](),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_FP_add_2(): + pipeline = VgfPipeline[input_t1]( + SigmoidAdd(), + (test_data_suite["zeros"](),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_INT_add_2(): + pipeline = VgfPipeline[input_t1]( + SigmoidAdd(), + (test_data_suite["zeros"](),), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_FP_add_3(): + pipeline = VgfPipeline[input_t1]( + SigmoidAddSigmoid(), + (test_data_suite["randn_neg"](), test_data_suite["randn_pos"]()), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sigmoid_vgf_INT_add_3(): + pipeline = VgfPipeline[input_t1]( + SigmoidAddSigmoid(), + (test_data_suite["randn_neg"](), test_data_suite["randn_pos"]()), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +def get_symmetric_a16w8_sigmoid_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +@pytest.mark.xfail( + reason="missing int16 sigmoid ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13974" +) +def test_sigmoid_16a8w_tosa_INT(test_data: torch.Tensor): + """Test sigmoid operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Sigmoid(), + (test_data(),), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_sigmoid_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 sigmoid operations" +) +def test_sigmoid_16a8w_u55_INT16(test_data: torch.Tensor): + """Test sigmoid operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Sigmoid(), + (test_data(),), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_sigmoid_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 sigmoid operations" +) +def test_sigmoid_16a8w_u85_INT16(test_data: torch.Tensor): + """Test sigmoid operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Sigmoid(), + (test_data(),), + aten_op, + exir_op, + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_sigmoid_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sigmoid_16bit.py b/backends/arm/test/ops/test_sigmoid_16bit.py index a41681675ce..ad8c49b234c 100644 --- a/backends/arm/test/ops/test_sigmoid_16bit.py +++ b/backends/arm/test/ops/test_sigmoid_16bit.py @@ -16,7 +16,7 @@ OpNotSupportedPipeline, TosaPipelineINT, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.backends.xnnpack.test.tester import Quantize from torchao.quantization.pt2e import HistogramObserver from torchao.quantization.pt2e.quantizer import QuantizationSpec @@ -41,7 +41,7 @@ def get_16bit_sigmoid_quantizer(u55_config=False): tosa_version = conftest.get_option("tosa_version") tosa_profiles = { "1.0": TosaSpecification.create_from_string( - "TOSA-1.0+INT" + ("+u55" if u55_config else "") + "TOSA-1.0+INT+int16" + ("+u55" if u55_config else "") ), } @@ -94,6 +94,7 @@ def test_sigmoid_tosa_INT(test_data): Sigmoid.aten_op, Sigmoid.exir_op, qtol=1, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_16bit_sigmoid_quantizer()) pipeline.run() @@ -114,18 +115,12 @@ def test_sigmoid_tosa_INT_add_sigmoid(test_data): Sigmoid.aten_op, Sigmoid.exir_op, qtol=1, + tosa_extensions=["int16"], ) + pipeline.change_args("quantize", get_16bit_sigmoid_quantizer()) pipeline.run() -xfails = { - "ones": "AssertionError: Output 0 does not match reference output. MLETORCH-787", - "rand": "AssertionError: Output 0 does not match reference output. MLETORCH-787", - "rand_4d": "AssertionError: Output 0 does not match reference output. MLETORCH-787", - "ramp": "AssertionError: Output 0 does not match reference output. MLETORCH-787", -} - - @common.parametrize( "test_data", test_data_suite, @@ -154,6 +149,7 @@ def test_sigmoid_u55_INT_add_sigmoid(test_data): n_expected_delegates=1, quantize=True, u55_subset=True, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_16bit_sigmoid_quantizer(True)) pipeline.run() @@ -180,7 +176,7 @@ def test_sigmoid_u85_INT(test_data): "ramp": "AssertionError: Output 0 does not match reference output. MLETORCH-787" }, ) -@pytest.mark.flaky(reruns=5) # MLETORCH-787: Investigate int16-int8 rescaling precision +@pytest.mark.xfail # MLETORCH-787: Investigate int16-int8 rescaling precision @common.XfailIfNoCorstone320 def test_sigmoid_u85_INT_add_sigmoid(test_data): pipeline = EthosU85PipelineINT( diff --git a/backends/arm/test/ops/test_sigmoid_32bit.py b/backends/arm/test/ops/test_sigmoid_32bit.py index 7d2e649bcd8..70863cd4757 100644 --- a/backends/arm/test/ops/test_sigmoid_32bit.py +++ b/backends/arm/test/ops/test_sigmoid_32bit.py @@ -12,7 +12,7 @@ OpNotSupportedPipeline, TosaPipelineINT, ) -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification from executorch.backends.xnnpack.test.tester import Quantize from torchao.quantization.pt2e import HistogramObserver from torchao.quantization.pt2e.quantizer import QuantizationSpec @@ -57,7 +57,7 @@ def get_32bit_sigmoid_quantizer(u55_config=False): tosa_version = conftest.get_option("tosa_version") tosa_profiles = { "1.0": TosaSpecification.create_from_string( - "TOSA-1.0+INT" + ("+u55" if u55_config else "") + "TOSA-1.0+INT+int16" + ("+u55" if u55_config else "") ), } @@ -110,6 +110,7 @@ def test_sigmoid_tosa_INT(test_data): Sigmoid.aten_op, Sigmoid.exir_op, qtol=1, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_32bit_sigmoid_quantizer()) pipeline.run() @@ -123,6 +124,7 @@ def test_sigmoid_tosa_INT_add_sigmoid(test_data): Sigmoid.aten_op, Sigmoid.exir_op, qtol=1, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_32bit_sigmoid_quantizer()) pipeline.run() @@ -136,6 +138,7 @@ def test_sigmoid_u55_INT(test_data): {Sigmoid.exir_op: 1}, quantize=True, u55_subset=True, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_32bit_sigmoid_quantizer(True)) pipeline.run() @@ -150,6 +153,7 @@ def test_sigmoid_u55_INT_add_sigmoid(test_data): n_expected_delegates=1, quantize=True, u55_subset=True, + tosa_extensions=["int16"], ) pipeline.change_args("quantize", get_32bit_sigmoid_quantizer(True)) pipeline.run() diff --git a/backends/arm/test/ops/test_sign.py b/backends/arm/test/ops/test_sign.py index 5e9a5c679b6..35ea9fc3e45 100644 --- a/backends/arm/test/ops/test_sign.py +++ b/backends/arm/test/ops/test_sign.py @@ -13,6 +13,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.sign.default" @@ -84,3 +85,29 @@ def test_sign_u85_INT(test_data: Tuple): exir_ops=exir_op, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sign_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sign(), + (test_data,), + aten_op=aten_op, + exir_op=exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sign_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sign(), + (test_data,), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_silu.py b/backends/arm/test/ops/test_silu.py index c938d2b707e..edc7d769be1 100644 --- a/backends/arm/test/ops/test_silu.py +++ b/backends/arm/test/ops/test_silu.py @@ -15,6 +15,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) @@ -111,3 +112,49 @@ def test_silu_u85_INT_inplace(test_data: input_t): Silu(), silu_data, Silu.aten_op_INT, run_on_fvp=True ) pipeline.run() + + +@common.parametrize("test_data", Silu.test_data) +@common.SkipIfNoModelConverter +def test_silu_vgf_FP(test_data: input_t): + silu_data = (test_data(), False) + pipeline = VgfPipeline[input_t]( + Silu(), silu_data, Silu.aten_op_FP, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", Silu.test_data) +@common.SkipIfNoModelConverter +def test_silu_vgf_FP_inplace(test_data: input_t): + silu_data = (test_data(), True) + pipeline = VgfPipeline[input_t]( + Silu(), silu_data, Silu.aten_op_inplace_FP, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", Silu.test_data) +@common.SkipIfNoModelConverter +def test_silu_vgf_INT(test_data: input_t): + silu_data = (test_data(), False) + pipeline = VgfPipeline[input_t]( + Silu(), + silu_data, + Silu.aten_op_INT, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", Silu.test_data) +@common.SkipIfNoModelConverter +def test_silu_vgf_INT_inplace(test_data: input_t): + silu_data = (test_data(), True) + pipeline = VgfPipeline[input_t]( + Silu(), + silu_data, + Silu.aten_op_INT, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sin.py b/backends/arm/test/ops/test_sin.py index 6f9037e1021..3ca593ad608 100644 --- a/backends/arm/test/ops/test_sin.py +++ b/backends/arm/test/ops/test_sin.py @@ -14,6 +14,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.sin.default" @@ -60,7 +61,7 @@ def test_sin_tosa_INT(test_data: Tuple): @common.parametrize("test_data", test_data_suite) -def test_sin_tosa_u55_INT(test_data: Tuple): +def test_sin_u55_INT(test_data: Tuple): pipeline = EthosU55PipelineINT[input_t1]( Sin(), (test_data,), @@ -72,7 +73,7 @@ def test_sin_tosa_u55_INT(test_data: Tuple): @common.parametrize("test_data", test_data_suite) -def test_sin_tosa_u85_INT(test_data: Tuple): +def test_sin_u85_INT(test_data: Tuple): pipeline = EthosU85PipelineINT[input_t1]( Sin(), (test_data,), @@ -81,3 +82,24 @@ def test_sin_tosa_u85_INT(test_data: Tuple): run_on_fvp=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sin_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sin(), (test_data,), aten_op, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sin_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sin(), + (test_data,), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sinh.py b/backends/arm/test/ops/test_sinh.py index ff486e6a4b8..a059ce0ad26 100644 --- a/backends/arm/test/ops/test_sinh.py +++ b/backends/arm/test/ops/test_sinh.py @@ -12,6 +12,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.sinh.default" @@ -76,3 +77,24 @@ def test_sinh_u85_INT(test_data: Tuple): Sinh(), (test_data,), aten_ops=aten_op, exir_ops=exir_op ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sinh_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sinh(), (test_data,), aten_op, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_sinh_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Sinh(), + (test_data,), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_slice.py b/backends/arm/test/ops/test_slice.py index 8fcf343dd57..eafeb04320e 100644 --- a/backends/arm/test/ops/test_slice.py +++ b/backends/arm/test/ops/test_slice.py @@ -7,16 +7,24 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.slice.Tensor" exir_op = "executorch_exir_dialects_edge__ops_aten_slice_copy" @@ -92,3 +100,131 @@ def test_slice_tensor_u85_INT(test_data: torch.Tensor): run_on_fvp=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_slice_tensor_vgf_FP(test_data: torch.Tensor): + pipeline = VgfPipeline[input_t1]( + Slice(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_slice_tensor_vgf_INT(test_data: torch.Tensor): + pipeline = VgfPipeline[input_t1]( + Slice(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +def get_symmetric_a16w8_slice_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +@pytest.mark.xfail( + reason="missing int16 slice ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13976" +) +def test_slice_tensor_16a8w_tosa_INT(test_data: torch.Tensor): + """Test slice operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Slice(), + test_data(), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_slice_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 slice operations" +) +def test_slice_tensor_16a8w_u55_INT16(test_data: torch.Tensor): + """Test slice operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Slice(), + test_data(), + aten_ops=[], + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_slice_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 slice operations" +) +def test_slice_tensor_16a8w_u85_INT16(test_data: torch.Tensor): + """Test slice operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Slice(), + test_data(), + aten_ops=[], + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_slice_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_softmax.py b/backends/arm/test/ops/test_softmax.py index db309ca1ab9..dc258f20ec4 100644 --- a/backends/arm/test/ops/test_softmax.py +++ b/backends/arm/test/ops/test_softmax.py @@ -14,6 +14,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.softmax.default" # Used for checking that we do not have softmax in the graph after decompose @@ -63,8 +64,8 @@ def test_softmax_tosa_INT(test_data): @common.parametrize( "test_data", Softmax.test_data, - xfails={ - "randn_mult_batches": "MLETORCH-433: Multiple batches not supported on FVP" + { + "randn_neg_dim": "MLBEDSW-11032: ILLEGAL_OFM_BASE error: Base addresses must be aligned to brick depth on u55." }, ) @common.XfailIfNoCorstone300 @@ -76,13 +77,7 @@ def test_softmax_u55_INT(test_data): pipeline.run() -@common.parametrize( - "test_data", - Softmax.test_data, - xfails={ - "randn_mult_batches": "MLETORCH-433: Multiple batches not supported on FVP" - }, -) +@common.parametrize("test_data", Softmax.test_data) @common.XfailIfNoCorstone320 def test_softmax_u85_INT(test_data): data, dim = test_data() @@ -90,3 +85,35 @@ def test_softmax_u85_INT(test_data): pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op]) pipeline.change_args("run_method_and_compare_outputs", qtol=1) pipeline.run() + + +@common.parametrize("test_data", Softmax.test_data) +@common.SkipIfNoModelConverter +def test_softmax_vgf_FP(test_data): + data, dim = test_data() + pipeline = VgfPipeline[input_t1]( + Softmax(dim), + data, + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.add_stage_after( + "to_edge_transform_and_lower", pipeline.tester.check_not, [exir_op] + ) + pipeline.run() + + +@common.parametrize("test_data", Softmax.test_data) +@common.SkipIfNoModelConverter +def test_softmax_vgf_INT(test_data): + data, dim = test_data() + pipeline = VgfPipeline[input_t1]( + Softmax(dim), + data, + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.add_stage_after("quantize", pipeline.tester.check_not, [aten_op]) + # TODO: MLETORCH-1136 Change args of run_method_and_compare_outputs of the vgf tests + # pipeline.change_args("run_method_and_compare_outputs", qtol=1) + pipeline.run() diff --git a/backends/arm/test/ops/test_split.py b/backends/arm/test/ops/test_split.py index 330f37b35e6..388e85762af 100644 --- a/backends/arm/test/ops/test_split.py +++ b/backends/arm/test/ops/test_split.py @@ -14,6 +14,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) exir_op = "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default" @@ -101,6 +102,21 @@ def test_split_with_sizes_tosa_FP_one_out(test_data: input_t1): pipeline.run() +@common.parametrize( + "test_data", + (Split.test_data | Split.test_data_list), +) +def test_split_with_sizes_tosa_FP_two_out(test_data: input_t1): + + pipeline = TosaPipelineFP[input_t1]( + SplitTwoOut(), + test_data(), + aten_op=[], + exir_op=exir_op, + ) + pipeline.run() + + @common.parametrize( "test_data", (Split.test_data | Split.test_data_list), @@ -145,3 +161,84 @@ def test_split_with_sizes_u85_INT(test_data: input_t1): run_on_fvp=False, ) pipeline.run() + + +@common.parametrize( + "test_data", + (Split.test_data | Split.test_data_list), +) +@common.SkipIfNoModelConverter +def test_split_with_sizes_vgf_FP(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Split(), + test_data(), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", Split.test_data_list) +@common.SkipIfNoModelConverter +def test_split_with_sizes_vgf_FP_2(test_data: input_t1): + + pipeline = VgfPipeline[input_t1]( + SplitWithSizes(), + test_data(), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + (Split.test_data | Split.test_data_list), +) +@common.SkipIfNoModelConverter +def test_split_with_sizes_vgf_FP_one_out(test_data: input_t1): + + pipeline = VgfPipeline[input_t1]( + SplitSingleOut(), + test_data(), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + (Split.test_data | Split.test_data_list), +) +@common.SkipIfNoModelConverter +def test_split_with_sizes_vgf_FP_two_out(test_data: input_t1): + + pipeline = VgfPipeline[input_t1]( + SplitTwoOut(), + test_data(), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize( + "test_data", + (Split.test_data | Split.test_data_list), +) +@common.SkipIfNoModelConverter +def test_split_with_sizes_vgf_INT(test_data: input_t1): + + pipeline = VgfPipeline[input_t1]( + Split(), + test_data(), + aten_op=[], + exir_op=exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sqrt.py b/backends/arm/test/ops/test_sqrt.py index ee554ce4fd2..15e2dd45322 100644 --- a/backends/arm/test/ops/test_sqrt.py +++ b/backends/arm/test/ops/test_sqrt.py @@ -13,6 +13,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) @@ -39,11 +40,6 @@ def forward(self, x): } -fvp_xfails = { - "sqrt_tensor_rank4_multibatch": "MLETORCH-517 : Multiple batches not supported", -} - - @common.parametrize("test_data", Sqrt.test_data) def test_sqrt_tosa_FP(test_data: Sqrt.input_t): pipeline = TosaPipelineFP[Sqrt.input_t]( @@ -66,7 +62,7 @@ def test_sqrt_tosa_INT(test_data: Sqrt.input_t): pipeline.run() -@common.parametrize("test_data", Sqrt.test_data, fvp_xfails) +@common.parametrize("test_data", Sqrt.test_data) @common.XfailIfNoCorstone300 def test_sqrt_u55_INT(test_data: Sqrt.input_t): pipeline = EthosU55PipelineINT[Sqrt.input_t]( @@ -79,7 +75,7 @@ def test_sqrt_u55_INT(test_data: Sqrt.input_t): pipeline.run() -@common.parametrize("test_data", Sqrt.test_data, fvp_xfails) +@common.parametrize("test_data", Sqrt.test_data) @common.XfailIfNoCorstone320 def test_sqrt_u85_INT(test_data: Sqrt.input_t): pipeline = EthosU85PipelineINT[Sqrt.input_t]( @@ -90,3 +86,29 @@ def test_sqrt_u85_INT(test_data: Sqrt.input_t): run_on_fvp=True, ) pipeline.run() + + +@common.parametrize("test_data", Sqrt.test_data) +@common.SkipIfNoModelConverter +def test_sqrt_vgf_FP(test_data: Sqrt.input_t): + pipeline = VgfPipeline[Sqrt.input_t]( + Sqrt(), + test_data(), + Sqrt.aten_op_FP, + Sqrt.exir_op_FP, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", Sqrt.test_data) +@common.SkipIfNoModelConverter +def test_sqrt_vgf_INT(test_data: Sqrt.input_t): + pipeline = VgfPipeline[Sqrt.input_t]( + Sqrt(), + test_data(), + Sqrt.aten_op_INT, + Sqrt.exir_op_INT, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_squeeze.py b/backends/arm/test/ops/test_squeeze.py index 10600169441..5c9f031deec 100644 --- a/backends/arm/test/ops/test_squeeze.py +++ b/backends/arm/test/ops/test_squeeze.py @@ -18,6 +18,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t1 = Tuple[torch.Tensor] # Input x @@ -56,6 +57,11 @@ def forward(self, x: torch.Tensor): return x.squeeze() +############## +## Squeeze ### +############## + + @common.parametrize("test_data", Squeeze.test_parameters) def test_squeeze_dim_tosa_FP(test_data: Tuple): pipeline = TosaPipelineFP[input_t1]( @@ -104,6 +110,37 @@ def test_squeeze_dim_u85_INT(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", Squeeze.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dim_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Squeeze(), + test_data(), + "torch.ops.aten.squeeze.default", + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", Squeeze.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dim_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Squeeze(), + test_data(), + "torch.ops.aten.squeeze.default", + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +################# +## SqueezeDim ### +################# + + @common.parametrize("test_data", SqueezeDim.test_parameters) def test_squeeze_dim_tosa_FP_2(test_data: Tuple): pipeline = TosaPipelineFP[input_t1]( @@ -152,6 +189,37 @@ def test_squeeze_dim_u85_INT_2(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", SqueezeDim.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dim_vgf_FP_2(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SqueezeDim(), + test_data(), + "torch.ops.aten.squeeze.dim", + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", SqueezeDim.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dim_vgf_INT_2(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SqueezeDim(), + test_data(), + "torch.ops.aten.squeeze.dim", + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +################## +## SqueezeDims ### +################## + + @common.parametrize("test_data", SqueezeDims.test_parameters) def test_squeeze_dims_tosa_FP(test_data: Tuple): pipeline = TosaPipelineFP[input_t1]( @@ -198,3 +266,29 @@ def test_squeeze_dims_u85_INT(test_data: Tuple): run_on_fvp=True, ) pipeline.run() + + +@common.parametrize("test_data", SqueezeDims.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dims_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SqueezeDims(), + test_data(), + "torch.ops.aten.squeeze.dims", + [], + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", SqueezeDims.test_parameters) +@common.SkipIfNoModelConverter +def test_squeeze_dims_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + SqueezeDims(), + test_data(), + "torch.ops.aten.squeeze.dims", + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sub.py b/backends/arm/test/ops/test_sub.py index 09f5884b1c4..c691506beb2 100644 --- a/backends/arm/test/ops/test_sub.py +++ b/backends/arm/test/ops/test_sub.py @@ -5,15 +5,19 @@ # 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 Tuple import torch + from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.sub.Tensor" @@ -30,13 +34,15 @@ "zeros": lambda: (torch.zeros(10),), } -fvp_sub_xfails = {"rand_4D_2x3x4x5": "MLETORCH-517 : Multiple batches not supported"} - # Two-input subtraction (x - y) sub2_test_data = { "rand_2D_4x4": lambda: (torch.rand(4, 4), torch.rand(4, 4)), "rand_3D_4x4x4": lambda: (torch.rand(4, 2, 2), torch.rand(4, 2, 2)), "rand_4D_2x2x4x4": lambda: (torch.rand(2, 2, 4, 4), torch.rand(2, 2, 4, 4)), + "rand_4D_big_small": lambda: ( + (10e30) * torch.randn(1, 20, 30, 40), + torch.randn(1, 20, 30, 40), + ), "zeros": lambda: (torch.rand(4, 4), torch.zeros(4, 4)), "randn_4D_mutltiple_broadcasts": lambda: ( torch.randn(1, 4, 4, 1), @@ -45,7 +51,16 @@ "rand_3d_rand_Scalar": lambda: (torch.rand(1, 6, 2), torch.rand(1)), "rand_3d_Scalar": lambda: (torch.rand(1, 6, 2), 1), } -fvp_sub2_xfails = {"rand_4D_2x2x4x4": "MLETORCH-517 : Multiple batches not supported"} + +# Sub and tan - the tan has a really steep curve just before Pi/2 and a point of discontinuity at Pi/2 +# so if the sub result is inaccurate, the error will be amplified by the tan +sub_tan_test_data = { + "rand_4D_pi": lambda: ( + torch.randn(1, 10, 20, 30) * math.pi / 2, + torch.randn(1, 10, 20, 30) * math.pi / 2, + ), + "rand_3D_pi": lambda: (torch.randn(1, 30, 40) * math.pi / 2, torch.rand(1, 30, 40)), +} class Sub(torch.nn.Module): @@ -58,6 +73,14 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): return x - y +class SubTan(torch.nn.Module): + + def forward(self, x: torch.Tensor, y: torch.Tensor): + z = x - y + t = torch.tan(z) + return t + + input_t1 = Tuple[torch.Tensor] # Input x input_t2 = Tuple[torch.Tensor, torch.Tensor] # Input x, y @@ -89,28 +112,27 @@ def test_sub_tensor_tosa_FP_2(test_data: Tuple[torch.Tensor, torch.Tensor]): @common.parametrize("test_data", sub_test_data) def test_sub_tensor_tosa_INT(test_data): """Test Subtraction (TOSA INT)""" - pipeline = TosaPipelineINT[input_t1]( - Sub(), - test_data(), - aten_op, - exir_op, - ) + pipeline = TosaPipelineINT[input_t1](Sub(), test_data(), aten_op, exir_op, qtol=0) pipeline.run() @common.parametrize("test_data", sub2_test_data) def test_sub_tensor_tosa_INT_2(test_data: Tuple[torch.Tensor, torch.Tensor]): + """Test Two-Operand Subtraction (TOSA INT)""" + pipeline = TosaPipelineINT[input_t2](Sub2(), test_data(), aten_op, exir_op, qtol=0) + pipeline.run() + + +@common.parametrize("test_data", sub_tan_test_data) +def test_sub_tensor_tosa_INT_3(test_data: Tuple[torch.Tensor, torch.Tensor]): """Test Two-Operand Subtraction (TOSA INT)""" pipeline = TosaPipelineINT[input_t2]( - Sub2(), - test_data(), - aten_op, - exir_op, + SubTan(), test_data(), aten_op, exir_op, qtol=0 ) pipeline.run() -@common.parametrize("test_data", sub_test_data, fvp_sub_xfails) +@common.parametrize("test_data", sub_test_data) @common.XfailIfNoCorstone300 def test_sub_tensor_u55_INT(test_data): """Test Subtraction on Ethos-U55 (FVP Mode)""" @@ -124,7 +146,7 @@ def test_sub_tensor_u55_INT(test_data): pipeline.run() -@common.parametrize("test_data", sub2_test_data, fvp_sub2_xfails) +@common.parametrize("test_data", sub2_test_data) @common.XfailIfNoCorstone300 def test_sub_tensor_u55_INT_2(test_data: Tuple[torch.Tensor, torch.Tensor]): """Test Two-Operand Subtraction on Ethos-U55 (FVP Mode)""" @@ -138,7 +160,7 @@ def test_sub_tensor_u55_INT_2(test_data: Tuple[torch.Tensor, torch.Tensor]): pipeline.run() -@common.parametrize("test_data", sub_test_data, fvp_sub_xfails) +@common.parametrize("test_data", sub_test_data) @common.XfailIfNoCorstone320 def test_sub_tensor_u85_INT_2(test_data): """Test Subtraction on Ethos-U85 (FVP Mode)""" @@ -152,7 +174,7 @@ def test_sub_tensor_u85_INT_2(test_data): pipeline.run() -@common.parametrize("test_data", sub2_test_data, fvp_sub2_xfails) +@common.parametrize("test_data", sub2_test_data) @common.XfailIfNoCorstone320 def test_sub_tensor_u85_INT(test_data: Tuple[torch.Tensor, torch.Tensor]): """Test Two-Operand Subtraction on Ethos-U85 (FVP Mode)""" @@ -164,3 +186,59 @@ def test_sub_tensor_u85_INT(test_data: Tuple[torch.Tensor, torch.Tensor]): run_on_fvp=True, ) pipeline.run() + + +@common.parametrize("test_data", sub_test_data) +@common.SkipIfNoModelConverter +def test_sub_tensor_vgf_FP(test_data: Tuple[torch.Tensor]): + """Test Subtraction (VGF FP)""" + pipeline = VgfPipeline[input_t1]( + Sub(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", sub2_test_data) +@common.SkipIfNoModelConverter +def test_sub_tensor_vgf_FP_2(test_data: Tuple[torch.Tensor, torch.Tensor]): + """Test Two-Operand Subtraction (VGF FP)""" + pipeline = VgfPipeline[input_t2]( + Sub2(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", sub_test_data) +@common.SkipIfNoModelConverter +def test_sub_tensor_vgf_INT(test_data: Tuple[torch.Tensor]): + """Test Subtraction (VGF INT)""" + pipeline = VgfPipeline[input_t1]( + Sub(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", sub2_test_data) +@common.SkipIfNoModelConverter +def test_sub_tensor_vgf_INT_2(test_data: Tuple[torch.Tensor, torch.Tensor]): + """Test Two-Operand Subtraction (VGF INT)""" + pipeline = VgfPipeline[input_t2]( + Sub2(), + test_data(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_sum.py b/backends/arm/test/ops/test_sum.py index 13e92fabb9b..9308315f76d 100644 --- a/backends/arm/test/ops/test_sum.py +++ b/backends/arm/test/ops/test_sum.py @@ -13,6 +13,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.sum.dim_IntList" @@ -59,6 +60,7 @@ def test_sum_dim_intlist_tosa_INT(test_data: input_t1): aten_op, exir_op=[], ) + pipeline.dump_artifact("export") pipeline.run() @@ -88,6 +90,27 @@ def test_view_u85_INT_1_0(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", Sum.test_parameters) +@common.SkipIfNoModelConverter +def test_sum_dim_intlist_vgf_FP(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Sum(), test_data(), aten_op, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", Sum.test_parameters) +@common.SkipIfNoModelConverter +def test_sum_dim_intlist_vgf_INT(test_data: input_t1): + pipeline = VgfPipeline[input_t1]( + Sum(), + test_data(), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + reject_inputs = { "reject_large_0_dim": lambda: (torch.rand((65537, 1, 1)), 0, False), "reject_large_2_dim": lambda: (torch.rand((800, 90, 1)), 2, False), diff --git a/backends/arm/test/ops/test_tanh.py b/backends/arm/test/ops/test_tanh.py index 1bd746d7b24..0e74618fd2f 100644 --- a/backends/arm/test/ops/test_tanh.py +++ b/backends/arm/test/ops/test_tanh.py @@ -6,15 +6,23 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.tanh.default" input_t1 = Tuple[torch.Tensor] # Input x @@ -83,3 +91,126 @@ def test_tanh_u85_INT(test_data: Tuple): run_on_fvp=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_tanh_vgf_FP(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Tanh(), (test_data(),), aten_op, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_tanh_vgf_INT(test_data: Tuple): + pipeline = VgfPipeline[input_t1]( + Tanh(), + (test_data(),), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +def get_symmetric_a16w8_tanh_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +@pytest.mark.xfail( + reason="missing int16 tanh ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13975" +) +def test_tanh_16a8w_tosa_INT(test_data: torch.Tensor): + """Test tanh operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Tanh(), + (test_data(),), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_tanh_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 tanh operations" +) +def test_tanh_16a8w_u55_INT16(test_data: torch.Tensor): + """Test tanh operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU55PipelineINT[input_t1]( + Tanh(), + (test_data(),), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_tanh_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 tanh operations" +) +def test_tanh_16a8w_u85_INT16(test_data: torch.Tensor): + """Test tanh operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = EthosU85PipelineINT[input_t1]( + Tanh(), + (test_data(),), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_tanh_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_to_copy.py b/backends/arm/test/ops/test_to_copy.py index f63909c41d0..5c01788c805 100644 --- a/backends/arm/test/ops/test_to_copy.py +++ b/backends/arm/test/ops/test_to_copy.py @@ -15,6 +15,8 @@ from executorch.backends.arm.test.tester.test_pipeline import ( OpNotSupportedPipeline, TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, ) input_t1 = Tuple[torch.Tensor] # Input x @@ -29,6 +31,15 @@ def forward(self, x: torch.Tensor): return x.to(dtype=self.target_dtype) +class CastAdd(torch.nn.Module): + def __init__(self, target_dtype): + super().__init__() + self.target_dtype = target_dtype + + def forward(self, x: torch.Tensor): + return x.to(dtype=self.target_dtype) + x.to(dtype=self.target_dtype) + + """ Tests the _to_copy operation. @@ -60,7 +71,7 @@ def forward(self, x: torch.Tensor): @common.parametrize("test_data", _TO_COPY_TEST_DATA_FP) -def test_copy_tosa_FP(test_data: Tuple): +def test_to_tosa_FP(test_data: Tuple): test_tensor, new_dtype = test_data() pipeline = TosaPipelineFP[input_t1]( @@ -69,6 +80,39 @@ def test_copy_tosa_FP(test_data: Tuple): aten_op=[], exir_op=[], ) + # int to int cast is not supported in TOSA+FP profile + if not new_dtype.is_floating_point and not torch.is_floating_point(test_tensor): + pipeline.change_args( + "check_count.exir", + { + "torch.ops.higher_order.executorch_call_delegate": 0, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + }, + ) + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +@common.parametrize("test_data", _TO_COPY_TEST_DATA_FP) +@common.SkipIfNoModelConverter +def test_to_vgf_FP(test_data: Tuple): + test_tensor, new_dtype = test_data() + pipeline = VgfPipeline[input_t1]( + Cast(new_dtype), + (test_tensor,), + aten_op=[], + exir_op=[], + tosa_version="TOSA-1.0+FP", + ) + # int to int cast is not supported in TOSA+FP profile + if not new_dtype.is_floating_point and not torch.is_floating_point(test_tensor): + pipeline.change_args( + "check_count.exir", + { + "torch.ops.higher_order.executorch_call_delegate": 0, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + }, + ) pipeline.run() @@ -104,7 +148,7 @@ def test_copy_tosa_FP(test_data: Tuple): @common.parametrize("test_data", _TO_COPY_TEST_DATA_INT) -def test_copy_tosa_INT(test_data: Tuple): +def test_to_tosa_INT_not_delegated(test_data: Tuple): test_tensor, new_dtype = test_data() pipeline = OpNotSupportedPipeline[input_t1]( @@ -116,3 +160,87 @@ def test_copy_tosa_INT(test_data: Tuple): quantize=True, ) pipeline.run() + + +@common.parametrize("test_data", _TO_COPY_TEST_DATA_INT) +@common.SkipIfNoModelConverter +def test_to_vgf_INT(test_data: Tuple): + # Op not supported + pass + + +_TO_COPY_TEST_DATA_REDUNDANT_CAST = { + "rand_fp16_fp16": lambda: ( + torch.rand((1, 2, 3, 4), dtype=torch.float16), + torch.float16, + ), + "rand_fp32_fp32": lambda: ( + torch.rand((1, 2, 3, 4), dtype=torch.float32), + torch.float32, + ), + "rand_int8_int8": lambda: ( + torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int8), + torch.int8, + ), + "rand_int16_int16": lambda: ( + torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int16), + torch.int16, + ), + "rand_int32_int32": lambda: ( + torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int32), + torch.int32, + ), +} + +redundant_xfails_FP = { + "rand_fp16_fp16": "FP16 is not supported", + "rand_int8_int8": "Tracing graph with quantized input is not supported.", + "rand_int16_int16": "Tracing graph with quantized input is not supported.", +} + +redundant_xfails_INT = { + "rand_fp16_fp16": "FP16 is not supported", + "rand_int8_int8": "Tracing graph with quantized input is not supported.", +} + + +@common.parametrize( + "test_data", _TO_COPY_TEST_DATA_REDUNDANT_CAST, xfails=redundant_xfails_FP +) +def test_to_tosa_FP_REDUNDANT_CAST(test_data: Tuple): + test_tensor, new_dtype = test_data() + pipeline = TosaPipelineFP[input_t1]( + CastAdd(new_dtype), + (test_tensor,), + aten_op=[], + exir_op=[], + ) + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +@common.parametrize( + "test_data", _TO_COPY_TEST_DATA_REDUNDANT_CAST, xfails=redundant_xfails_INT +) +def test_to_tosa_INT_REDUNDANT_CAST(test_data: Tuple): + test_tensor, new_dtype = test_data() + pipeline = TosaPipelineINT[input_t1]( + CastAdd(new_dtype), + (test_tensor,), + aten_op=[], + exir_op=[], + ) + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +@common.parametrize("test_data", _TO_COPY_TEST_DATA_REDUNDANT_CAST) +def test_to_tosa_INT_not_delegated_REDUNDANT_CAST(test_data: Tuple): + test_tensor, new_dtype = test_data() + pipeline = OpNotSupportedPipeline[input_t1]( + Cast(new_dtype), + (test_tensor,), + non_delegated_ops={}, # These are removed outside of the Arm backend so the graph is empty + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_unary_combos.py b/backends/arm/test/ops/test_unary_combos.py new file mode 100644 index 00000000000..db442d2d8d0 --- /dev/null +++ b/backends/arm/test/ops/test_unary_combos.py @@ -0,0 +1,134 @@ +# Copyright 2025 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.tester.test_pipeline import ( + EthosU55PipelineINT, + EthosU85PipelineINT, + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + +Tensor1 = Tuple[torch.Tensor] + + +class NegAdd(torch.nn.Module): + # neg(x) + 1 + edge_op_list = [ + "executorch_exir_dialects_edge__ops_aten_neg_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + ] + + def get_inputs(self) -> Tensor1: + return (torch.rand(10, 10, 10),) + + def forward(self, x): + return torch.neg(x) + 1.0 + + +class MinAddZero(torch.nn.Module): + # min(x, 0) + 1 + edge_op_list = [ + "executorch_exir_dialects_edge__ops_aten_full_like_default", + "executorch_exir_dialects_edge__ops_aten_minimum_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + ] + + # range [-1, 1] + def get_inputs(self) -> Tensor1: + return (torch.rand(10, 10, 10) * 2 - 1,) + + def forward(self, x): + # We want Tensor-Tensor minimum + z = torch.full_like(x, 0.0) + return torch.minimum(x, z) + 1.0 + + +class MaxAddZero(torch.nn.Module): + # max(x, 0) + 1.0 + edge_op_list = [ + "executorch_exir_dialects_edge__ops_aten_full_like_default", + "executorch_exir_dialects_edge__ops_aten_maximum_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + ] + + # range [-1, 1] + def get_inputs(self) -> Tensor1: + return (torch.rand(10, 10, 10) * 2 - 1,) + + def forward(self, x): + z = torch.full_like(x, 0.0) + return torch.maximum(x, z) + 1.0 + + +class AbsAdd(torch.nn.Module): + # abs(x) + 1.0 + edge_op_list = [ + "executorch_exir_dialects_edge__ops_aten_abs_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + ] + + def get_inputs(self) -> Tensor1: + return (torch.rand(10, 10, 10),) + + def forward(self, x): + return torch.abs(x) + 1.0 + + +MODELS = [NegAdd, AbsAdd, MaxAddZero, MinAddZero] + + +def _build(model_cls): + m = model_cls() + return m, m.get_inputs(), model_cls.edge_op_list + + +@pytest.mark.parametrize("model_cls", MODELS, ids=lambda c: c.__name__) +def test_unary_combos_tosa_FP(model_cls): + m, inputs, exir = _build(model_cls) + p = TosaPipelineFP[Tensor1](m, inputs, aten_op=[], exir_op=exir) + p.run() + + +@pytest.mark.parametrize("model_cls", MODELS, ids=lambda c: c.__name__) +def test_unary_combos_tosa_INT(model_cls): + m, inputs, exir = _build(model_cls) + p = TosaPipelineINT[Tensor1](m, inputs, aten_op=[], exir_op=exir, qtol=1) + p.run() + + +@common.XfailIfNoCorstone300 +@pytest.mark.parametrize("model_cls", MODELS, ids=lambda c: c.__name__) +def test_unary_combos_u55_INT(model_cls): + m, inputs, exir = _build(model_cls) + p = EthosU55PipelineINT[Tensor1]( + m, inputs, aten_ops=[], exir_ops=exir, run_on_fvp=True + ) + p.run() + + +@common.XfailIfNoCorstone320 +@pytest.mark.parametrize("model_cls", MODELS, ids=lambda c: c.__name__) +def test_unary_combos_u85_INT(model_cls): + m, inputs, exir = _build(model_cls) + p = EthosU85PipelineINT[Tensor1]( + m, inputs, aten_ops=[], exir_ops=exir, run_on_fvp=True + ) + p.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.parametrize("model_cls", MODELS, ids=lambda c: c.__name__) +def test_unary_combos_vgf_INT(model_cls): + m, inputs, exir = _build(model_cls) + p = VgfPipeline[Tensor1]( + m, inputs, aten_op=[], exir_op=exir, tosa_version="TOSA-1.0+INT" + ) + p.run() diff --git a/backends/arm/test/ops/test_unbind.py b/backends/arm/test/ops/test_unbind.py index d1425719b0b..cd33f8217df 100644 --- a/backends/arm/test/ops/test_unbind.py +++ b/backends/arm/test/ops/test_unbind.py @@ -11,6 +11,7 @@ from executorch.backends.arm.test.tester.test_pipeline import ( TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t = tuple[torch.Tensor] @@ -53,3 +54,29 @@ def test_unbind_int_tosa_INT(test_data: test_data_t): Unbind.aten_op, ) pipeline.run() + + +@common.parametrize("test_data", Unbind.test_data) +@common.SkipIfNoModelConverter +def test_unbind_int_vgf_FP(test_data: test_data_t): + input_data, init_data = test_data + pipeline = VgfPipeline[input_t]( + Unbind(*init_data), + input_data(), + Unbind.aten_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", Unbind.test_data) +@common.SkipIfNoModelConverter +def test_unbind_int_vgf_INT(test_data: test_data_t): + input_data, init_data = test_data + pipeline = VgfPipeline[input_t]( + Unbind(*init_data), + input_data(), + Unbind.aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_unflatten.py b/backends/arm/test/ops/test_unflatten.py index e3bcb32375d..95c68b2940d 100644 --- a/backends/arm/test/ops/test_unflatten.py +++ b/backends/arm/test/ops/test_unflatten.py @@ -11,6 +11,7 @@ from executorch.backends.arm.test.tester.test_pipeline import ( TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t = tuple[torch.Tensor] @@ -54,3 +55,29 @@ def test_unflatten_int_tosa_INT(test_data: test_data_t): Unflatten.aten_op, ) pipeline.run() + + +@common.parametrize("test_data", Unflatten.test_data) +@common.SkipIfNoModelConverter +def test_unflatten_int_vgf_FP(test_data: test_data_t): + module, inputs = test_data() + pipeline = VgfPipeline[input_t]( + module, + inputs, + Unflatten.aten_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", Unflatten.test_data) +@common.SkipIfNoModelConverter +def test_unflatten_int_vgf_INT(test_data: test_data_t): + module, inputs = test_data() + pipeline = VgfPipeline[input_t]( + module, + inputs, + Unflatten.aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_unsqueeze.py b/backends/arm/test/ops/test_unsqueeze.py index d192d5289fd..54e1b0dd0ce 100644 --- a/backends/arm/test/ops/test_unsqueeze.py +++ b/backends/arm/test/ops/test_unsqueeze.py @@ -17,6 +17,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.unsqueeze.default" @@ -80,3 +81,26 @@ def test_unsqueeze_u85_INT(test_tensor: torch.Tensor): run_on_fvp=True, ) pipeline.run() + + +@common.parametrize("test_tensor", Unsqueeze.test_parameters) +@common.SkipIfNoModelConverter +def test_unsqueeze_vgf_FP(test_tensor: torch.Tensor): + for i in range(-test_tensor[0].dim() - 1, test_tensor[0].dim() + 1): + pipeline = VgfPipeline[input_t1]( + Unsqueeze(), (*test_tensor, i), aten_op, tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_tensor", Unsqueeze.test_parameters) +@common.SkipIfNoModelConverter +def test_unsqueeze_vgf_INT(test_tensor: torch.Tensor): + for i in range(-test_tensor[0].dim() - 1, test_tensor[0].dim() + 1): + pipeline = VgfPipeline[input_t1]( + Unsqueeze(), + (*test_tensor, i), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_upsample_bilinear2d.py b/backends/arm/test/ops/test_upsample_bilinear2d.py index d3b3ce1e303..95e69bc5204 100644 --- a/backends/arm/test/ops/test_upsample_bilinear2d.py +++ b/backends/arm/test/ops/test_upsample_bilinear2d.py @@ -13,6 +13,7 @@ OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.upsample_bilinear2d.vec" @@ -305,3 +306,99 @@ def test_upsample_bilinear2d_vec_U85_INT_UpsamplingBilinear2d( if not compare_outputs: pipeline.pop_stage(-1) pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_FP_UpsamplingBilinear2d(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + UpsamplingBilinear2d(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_FP_Upsample(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + Upsample(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_FP_Interpolate(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + Interpolate(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_INT_UpsamplingBilinear2d(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + UpsamplingBilinear2d(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_INT_Upsample(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + Upsample(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_tosa) +@common.SkipIfNoModelConverter +def test_upsample_bilinear2d_vgf_INT_Interpolate(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data + pipeline = VgfPipeline[input_t1]( + Interpolate(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() diff --git a/backends/arm/test/ops/test_upsample_nearest2d.py b/backends/arm/test/ops/test_upsample_nearest2d.py index d0a13b3036d..a39adefc168 100644 --- a/backends/arm/test/ops/test_upsample_nearest2d.py +++ b/backends/arm/test/ops/test_upsample_nearest2d.py @@ -12,6 +12,7 @@ OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) aten_op = "torch.ops.aten.upsample_nearest2d.vec" @@ -150,7 +151,7 @@ def test_upsample_nearest2d_vec_tosa_FP_interpolate(test_data: torch.Tensor): @common.parametrize("test_data", test_data_suite) -def test_upsample_nearest2d_vec_tosa_INT_interpolate(test_data: torch.Tensor): +def test_upsample_nearest2d_vec_tosa_INT(test_data: torch.Tensor): test_data, size, scale_factor, compare_outputs = test_data() pipeline = TosaPipelineINT[input_t1]( @@ -179,6 +180,117 @@ def test_upsample_nearest2d_vec_tosa_INT_nearest(test_data: torch.Tensor): pipeline.run() +@common.parametrize("test_data", test_data_suite) +def test_upsample_nearest2d_vec_tosa_INT_interpolate(test_data: torch.Tensor): + test_data, size, scale_factor, compare_outputs = test_data() + + pipeline = TosaPipelineINT[input_t1]( + Interpolate(size, scale_factor), + (test_data,), + aten_op, + exir_op=[], + ) + if not compare_outputs: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_FP(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + UpsamplingNearest2d(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_FP_nearest(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + Upsample(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_FP_interpolate(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + Interpolate(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_INT(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + UpsamplingNearest2d(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_INT_nearest(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + Upsample(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.SkipIfNoModelConverter +def test_upsample_nearest2d_vgf_INT_interpolate(test_data: torch.Tensor): + data, size, scale_factor, compare = test_data() + pipeline = VgfPipeline[input_t1]( + Interpolate(size, scale_factor), + (data,), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + ) + if not compare: + pipeline.pop_stage(-1) + pipeline.run() + + @common.parametrize("test_data", test_data_u55) @common.XfailIfNoCorstone300 def test_upsample_nearest2d_vec_U55_INT_Upsample_not_delegated( diff --git a/backends/arm/test/ops/test_var.py b/backends/arm/test/ops/test_var.py index 6e71dca557a..9567f90c480 100644 --- a/backends/arm/test/ops/test_var.py +++ b/backends/arm/test/ops/test_var.py @@ -14,6 +14,7 @@ EthosU85PipelineINT, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t1 = Tuple[torch.Tensor] # Input x @@ -155,6 +156,11 @@ def forward( return x.var(dim=self.dim, keepdim=self.keepdim, correction=self.correction) +########## +## Var ### +########## + + @common.parametrize("test_data", Var.test_parameters) def test_var_dim_tosa_FP_no_dim(test_data: Tuple): test_data, keepdim, correction = test_data() @@ -207,6 +213,35 @@ def test_var_dim_u85_INT_no_dim(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", Var.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_FP_no_dim(test_data: Tuple): + data, keepdim, correction = test_data() + pipeline = VgfPipeline[input_t1]( + Var(keepdim, correction), (data,), [], [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", Var.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_INT_no_dim(test_data: Tuple): + data, keepdim, correction = test_data() + pipeline = VgfPipeline[input_t1]( + Var(keepdim, correction), + (data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +############# +## VarDim ### +############# + + @common.parametrize("test_data", VarDim.test_parameters) def test_var_dim_tosa_FP(test_data: Tuple): test_data, dim, keepdim, unbiased = test_data() @@ -260,6 +295,35 @@ def test_var_dim_u85_INT(test_data: Tuple): pipeline.run() +@common.parametrize("test_data", VarDim.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_FP(test_data: Tuple): + data, dim, keepdim, unbiased = test_data() + pipeline = VgfPipeline[input_t1]( + VarDim(dim, keepdim, unbiased), (data,), [], [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", VarDim.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_INT(test_data: Tuple): + data, dim, keepdim, unbiased = test_data() + pipeline = VgfPipeline[input_t1]( + VarDim(dim, keepdim, unbiased), + (data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +#################### +## VarCorrection ### +#################### + + @common.parametrize("test_data", VarCorrection.test_parameters) def test_var_dim_tosa_FP_correction(test_data: Tuple): test_data, dim, keepdim, correction = test_data() @@ -310,3 +374,27 @@ def test_var_dim_u85_INT_correction(test_data: Tuple): run_on_fvp=True, ) pipeline.run() + + +@common.parametrize("test_data", VarCorrection.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_FP_correction(test_data: Tuple): + data, dim, keepdim, corr = test_data() + pipeline = VgfPipeline[input_t1]( + VarCorrection(dim, keepdim, corr), (data,), [], [], tosa_version="TOSA-1.0+FP" + ) + pipeline.run() + + +@common.parametrize("test_data", VarCorrection.test_parameters) +@common.SkipIfNoModelConverter +def test_var_dim_vgf_INT_correction(test_data: Tuple): + data, dim, keepdim, corr = test_data() + pipeline = VgfPipeline[input_t1]( + VarCorrection(dim, keepdim, corr), + (data,), + [], + [], + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_view.py b/backends/arm/test/ops/test_view.py index 0f8024c32dc..fb0ba54436e 100644 --- a/backends/arm/test/ops/test_view.py +++ b/backends/arm/test/ops/test_view.py @@ -9,16 +9,24 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) +from executorch.backends.arm.tosa.specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize aten_op = "torch.ops.aten.view.default" @@ -81,22 +89,7 @@ def test_view_tosa_INT(test_data: Tuple): pipeline.run() -xfails = { - "rand_4d_neg": "MLETORCH-517: Multiple batches not supported", - "rand_4d_4d_small": "MLETORCH-517: Multiple batches not supported", - "rand_4d_4d": "MLETORCH-517: Multiple batches not supported", - "rand_4d_2d": "MLETORCH-517: Multiple batches not supported", - "rand_4d_3d": "MLETORCH-517: Multiple batches not supported", - "rand_4d_1": "MLETORCH-517: Multiple batches not supported", - "rand_4d_2": "MLETORCH-517: Multiple batches not supported", - "rand_4d_2_4_big": "MLETORCH-517: Multiple batches not supported", - "rand_4d_4_3": "MLETORCH-517: Multiple batches not supported", - "rand_4d_4_2": "MLETORCH-517: Multiple batches not supported", - "rand_4d_2_4_same": "MLETORCH-517: Multiple batches not supported", -} - - -@common.parametrize("test_data", View.needs_transpose_tests, xfails=xfails) +@common.parametrize("test_data", View.needs_transpose_tests) @common.XfailIfNoCorstone300 def test_view_u55_INT(test_data: Tuple): test_tensor, new_shape = test_data() @@ -109,7 +102,33 @@ def test_view_u55_INT(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", View.rank_product_too_large, xfails=xfails) +@common.parametrize("test_data", View.needs_transpose_tests) +@common.SkipIfNoModelConverter +def test_view_vgf_FP(test_data: Tuple): + test_tensor, new_shape = test_data() + pipeline = VgfPipeline[input_t1]( + View(new_shape), + (test_tensor,), + aten_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_data", View.needs_transpose_tests) +@common.SkipIfNoModelConverter +def test_view_vgf_INT(test_data: Tuple): + test_tensor, new_shape = test_data() + pipeline = VgfPipeline[input_t1]( + View(new_shape), + (test_tensor,), + aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.run() + + +@common.parametrize("test_data", View.rank_product_too_large) @common.XfailIfNoCorstone300 def test_view_u55_INT_not_delegated(test_data: Tuple): test_tensor, new_shape = test_data() @@ -124,7 +143,7 @@ def test_view_u55_INT_not_delegated(test_data: Tuple): pipeline.run() -@common.parametrize("test_data", View.needs_transpose_tests, xfails=xfails) +@common.parametrize("test_data", View.needs_transpose_tests) @common.XfailIfNoCorstone320 def test_view_u85_INT(test_data: Tuple): test_tensor, new_shape = test_data() @@ -135,3 +154,108 @@ def test_view_u85_INT(test_data: Tuple): exir_ops=[], ) pipeline.run() + + +def get_symmetric_a16w8_view_quantizer(per_channel_quantization=False): + tosa_version = conftest.get_option("tosa_version") + tosa_profiles = { + "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"), + } + + quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) + quantizer.set_global( + get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", View.needs_transpose_tests) +@pytest.mark.xfail( + reason="missing int16 view ops support; fails at TOSA reference model with Unsupported operation type or rank. See: https://github.com/pytorch/executorch/issues/13977" +) +def test_view_16a8w_tosa_INT(test_data: Tuple): + """Test view operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + test_tensor, new_shape = test_data() + + pipeline = TosaPipelineINT[input_t1]( + View(new_shape), + (test_tensor,), + aten_op, + exir_op=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_view_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", View.needs_transpose_tests) +@common.XfailIfNoCorstone300 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 view operations" +) +def test_view_16a8w_u55_INT16(test_data: Tuple): + """Test view operation with 16A8W quantization on U55 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + test_tensor, new_shape = test_data() + + pipeline = EthosU55PipelineINT[input_t1]( + View(new_shape), + (test_tensor,), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_view_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() + + +@common.parametrize("test_data", View.needs_transpose_tests) +@common.XfailIfNoCorstone320 +@pytest.mark.xfail( + reason="Vela compilation fails with 'Invalid arguments' for int16 view operations" +) +def test_view_16a8w_u85_INT16(test_data: Tuple): + """Test view operation with 16A8W quantization on U85 (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + test_tensor, new_shape = test_data() + + pipeline = EthosU85PipelineINT[input_t1]( + View(new_shape), + (test_tensor,), + aten_op, + exir_ops=[], + per_channel_quantization=per_channel_quantization, + use_to_edge_transform_and_lower=True, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_view_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_where.py b/backends/arm/test/ops/test_where.py index c6b65612d59..ea036d26361 100644 --- a/backends/arm/test/ops/test_where.py +++ b/backends/arm/test/ops/test_where.py @@ -18,6 +18,7 @@ OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) from executorch.backends.xnnpack.test.tester.tester import Quantize @@ -215,3 +216,30 @@ def test_where_self_u85_INT(test_module): symmetric_io_quantization=True, ) pipeline.run() + + +@common.parametrize("test_module", test_modules_FP) +@common.SkipIfNoModelConverter +def test_where_self_vgf_FP(test_module): + pipeline = VgfPipeline[input_t]( + test_module(), + test_module().get_inputs(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+FP", + ) + pipeline.run() + + +@common.parametrize("test_module", test_modules_INT) +@common.SkipIfNoModelConverter +def test_where_self_vgf_INT(test_module): + pipeline = VgfPipeline[input_t]( + test_module(), + test_module().get_inputs(), + aten_op, + exir_op, + tosa_version="TOSA-1.0+INT", + symmetric_io_quantization=True, + ) + pipeline.run() diff --git a/backends/arm/test/ops/test_zeros.py b/backends/arm/test/ops/test_zeros.py index c93ba0802f1..caee678282a 100644 --- a/backends/arm/test/ops/test_zeros.py +++ b/backends/arm/test/ops/test_zeros.py @@ -12,6 +12,7 @@ OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, + VgfPipeline, ) input_t = tuple[torch.Tensor] @@ -38,17 +39,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ), } - test_data_not_delegated: dict[str, test_data_t] = { + # Mixed dtypes - the zeros op is delegated, but it leads to a non-delegated add op. + test_data_mixed_dtypes: dict[str, test_data_t] = { "fp32_int64": (lambda: (torch.randn(10),), (10, torch.int64)), "fp32_int32": (lambda: (torch.randn(10),), (10, torch.int32)), - "int32_int64": ( - lambda: (torch.randint(0, 10, [10], dtype=torch.int32),), - (10, torch.int64), - ), } -@common.parametrize("test_data", ZerosAdd.test_data) +@common.parametrize("test_data", ZerosAdd.test_data | ZerosAdd.test_data_mixed_dtypes) def test_zeros_tosa_FP(test_data: test_data_t): input_data, init_data = test_data pipeline = TosaPipelineFP[input_t]( @@ -59,7 +57,7 @@ def test_zeros_tosa_FP(test_data: test_data_t): pipeline.run() -@common.parametrize("test_data", ZerosAdd.test_data) +@common.parametrize("test_data", ZerosAdd.test_data | ZerosAdd.test_data_mixed_dtypes) def test_zeros_tosa_INT(test_data: test_data_t): input_data, init_data = test_data pipeline = TosaPipelineINT[input_t]( @@ -101,16 +99,45 @@ def test_zeros_u85_INT(test_data: test_data_t): @common.parametrize( "test_data", - ZerosAdd.test_data_not_delegated, - xfails={ - "fp32_int32": "MLETORCG-716: Do not delegate empty networks to vela", - "fp32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - "int32_int64": "MLETORCG-716: Do not delegate empty networks to vela", - }, + ZerosAdd.test_data_mixed_dtypes, ) def test_zeros_tosa_INT_not_delegated(test_data: test_data_t): input_data, init_data = test_data pipeline = OpNotSupportedPipeline[input_t]( - ZerosAdd(*init_data), input_data(), non_delegated_ops={}, quantize=True + ZerosAdd(*init_data), + input_data(), + non_delegated_ops={"executorch_exir_dialects_edge__ops_aten_add_Tensor": 1}, + n_expected_delegates=1, + quantize=True, + ) + pipeline.run() + + +@common.parametrize( + "test_data", + ZerosAdd.test_data, +) +@common.SkipIfNoModelConverter +def test_zeros_vgf_FP(test_data: test_data_t): + input_data, init_data = test_data + pipeline = VgfPipeline[input_t]( + ZerosAdd(*init_data), input_data(), ZerosAdd.aten_op, tosa_version="TOSA-1.0+FP" ) pipeline.run() + + +@common.parametrize( + "test_data", + ZerosAdd.test_data, +) +@common.SkipIfNoModelConverter +def test_zeros_vgf_INT(test_data: test_data_t): + input_data, init_data = test_data + pipeline = VgfPipeline[input_t]( + ZerosAdd(*init_data), + input_data(), + ZerosAdd.aten_op, + tosa_version="TOSA-1.0+INT", + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() diff --git a/backends/arm/test/passes/test_convert_int64_const_ops_to_int32.py b/backends/arm/test/passes/test_convert_int64_const_ops_to_int32.py new file mode 100644 index 00000000000..ddb31625849 --- /dev/null +++ b/backends/arm/test/passes/test_convert_int64_const_ops_to_int32.py @@ -0,0 +1,511 @@ +# Copyright 2025 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, Union + +import pytest + +import torch +from executorch.backends.arm._passes import ( + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, +) + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, +) + +input_t1 = Tuple[torch.Tensor] # Input x +input_t2 = Tuple[torch.Tensor, torch.Tensor] # Input x, y + + +##################################################### +## Test arange(dtype=int64) -> arange(dtype=int32) ## +##################################################### + + +class ArangeDefaultIncrementViewLessThan(torch.nn.Module): + + def forward(self, x: torch.Tensor): + return (torch.arange(10, dtype=torch.int64) + 1).view(-1, 1) < x + + test_data = { + "randint": ( + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + ), + } + + +@common.parametrize("test_data", ArangeDefaultIncrementViewLessThan.test_data) +def test_convert_arange_default_int64_dtype_to_int32_pass_tosa_FP(test_data: input_t1): + module = ArangeDefaultIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineFP[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + transform_passes=[ConvertInt64ConstOpsToInt32Pass()], + ) + pipeline.run() + + +@common.parametrize("test_data", ArangeDefaultIncrementViewLessThan.test_data) +def test_convert_arange_default_int64_dtype_to_int32_pass_tosa_INT(test_data: input_t1): + module = ArangeDefaultIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineINT[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +class ArangeStartIncrementViewLessThan(torch.nn.Module): + + def forward(self, x: torch.Tensor): + return (torch.arange(0, 10, dtype=torch.int64) + 1).view(-1, 1) < x + + test_data = { + "randint": ( + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + ), + } + + +@common.parametrize("test_data", ArangeStartIncrementViewLessThan.test_data) +def test_convert_arange_start_int64_dtype_to_int32_pass_tosa_FP(test_data: input_t1): + module = ArangeStartIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineFP[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + transform_passes=[ConvertInt64ConstOpsToInt32Pass()], + ) + pipeline.run() + + +@common.parametrize("test_data", ArangeStartIncrementViewLessThan.test_data) +def test_convert_arange_start_int64_dtype_to_int32_pass_tosa_INT(test_data: input_t1): + module = ArangeStartIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineINT[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +class ArangeStartStepIncrementViewLessThan(torch.nn.Module): + + def forward(self, x: torch.Tensor): + return (torch.arange(0, 10, 2, dtype=torch.int64) + 1).view(-1, 1) < x + + test_data = { + "randint": ( + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + ), + } + + +@common.parametrize("test_data", ArangeStartStepIncrementViewLessThan.test_data) +def test_convert_arange_start_step_int64_dtype_to_int32_pass_tosa_FP( + test_data: input_t1, +): + module = ArangeStartStepIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineFP[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + transform_passes=[ConvertInt64ConstOpsToInt32Pass()], + ) + pipeline.run() + + +@common.parametrize("test_data", ArangeStartStepIncrementViewLessThan.test_data) +def test_convert_arange_start_step_int64_dtype_to_int32_pass_tosa_INT( + test_data: input_t1, +): + module = ArangeStartStepIncrementViewLessThan() + aten_ops_checks = [ + "torch.ops.aten.lt.Tensor", + "torch.ops.aten.view.default", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + ] + pipeline = TosaPipelineINT[input_t1]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +######################################################### +## Test arange(dtype=None) -> arange(dtype=None/int32) ## +######################################################### + + +class ArangeAddDtypeNone(torch.nn.Module): + aten_op: str = "torch.ops.aten.arange.start_step" + exir_op: str = "executorch_exir_dialects_edge__ops_aten_arange_start_step" + + def __init__(self, start: float, stop: float, step: float): + super().__init__() + self.args = (start, stop, step) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.arange(*self.args) + x + + test_data = { + "int64": (lambda: (torch.randn(10, 1),), (0, 10, 1)), + "float32_start": (lambda: (torch.randn(10, 1),), (0.0, 10, 1)), + "float32_stop": (lambda: (torch.randn(10, 1),), (0, 10.0, 1)), + "float32_step": (lambda: (torch.randn(10, 1),), (0, 10, 1.0)), + "int64_bool_0": (lambda: (torch.randn(10, 1),), (False, True, True)), + "int64_bool_1": (lambda: (torch.randn(10, 1),), (False, True, True * 10)), + "float32_bool_0": (lambda: (torch.randn(10, 1),), (0.0, True, True)), + "float32_bool_1": (lambda: (torch.randn(10, 1),), (False, True, True * 10.0)), + } + + +@common.parametrize("test_data", ArangeAddDtypeNone.test_data) +def test_arange_dtype_none_tosa_FP(test_data): + input_data, init_data = test_data + pipeline = TosaPipelineFP[input_t1]( + ArangeAddDtypeNone(*init_data), + input_data(), + ArangeAddDtypeNone.aten_op, + ArangeAddDtypeNone.exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", ArangeAddDtypeNone.test_data) +def test_arange_dtype_none_tosa_INT(test_data): + input_data, init_data = test_data + pipeline = TosaPipelineINT[input_t1]( + ArangeAddDtypeNone(*init_data), + input_data(), + ArangeAddDtypeNone.aten_op, + ArangeAddDtypeNone.exir_op, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +################################################# +## Test full(dtype=int64) -> full(dtype=int32) ## +################################################# + + +class FullIncrementViewMulXLessThanY(torch.nn.Module): + + def forward(self, x: torch.Tensor, y: torch.Tensor): + return ( + ( + torch.full( + ( + 1, + 3, + 5, + ), + 10, + dtype=torch.int64, + ) + + 1 + ).view(-1, 1) + * x + ) < y + + test_data = { + "randint": ( + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + ), + } + + +@common.parametrize("test_data", FullIncrementViewMulXLessThanY.test_data) +def test_convert_full_int64_dtype_to_int32_pass_tosa_FP(test_data: input_t1): + """ + There are four int64 placeholders in the original graph: + 1. _lifted_tensor_constant0: 1 + 2. x + 3. y + Ideally, after applying ConvertInt64ConstOpsToInt32Pass to convert the aten.full from int64 to int32, + the int32 type should propagate throughout the graph, and no int64 values should remain. + However, due to unexpected retracing behavior, a cast from int32 → int64 for x was reintroducedh. + + Applying ConvertInt64OutputOpsToInt32Pass afterward resolves this issue, + removing the int64 cast and producing a fully delegated int32 graph. + """ + module = FullIncrementViewMulXLessThanY() + aten_ops_checks = [ + "torch.ops.aten.full.default", + "torch.ops.aten.add.Tensor", + "torch.ops.aten.view.default", + "torch.ops.aten.mul.Tensor", + "torch.ops.aten.lt.Tensor", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_full_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + ] + pipeline = TosaPipelineFP[input_t2]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + transform_passes=[ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + ], + ) + pipeline.run() + + +@common.parametrize("test_data", FullIncrementViewMulXLessThanY.test_data) +def test_convert_full_int64_dtype_to_int32_pass_tosa_INT(test_data: input_t1): + """ + For INT profile, _lifted_tensor_constant0 is still int64 after applying ConvertInt64ConstOpsToInt32Pass(). + And an int64->int32 cast is inserted at the beginning of the graph. + TODO: Explore why _lifted_tensor_constant0 is handled in different ways in FP and INT profile. + Find a way to optimize out the int64->int32 cast. + """ + module = FullIncrementViewMulXLessThanY() + aten_ops_checks = [ + "torch.ops.aten.full.default", + "torch.ops.aten.add.Tensor", + "torch.ops.aten.view.default", + "torch.ops.aten.mul.Tensor", + "torch.ops.aten.lt.Tensor", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_full_default", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + "executorch_exir_dialects_edge__ops_aten_view_copy_default", + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_lt_Tensor", + ] + pipeline = TosaPipelineINT[input_t2]( + module, + test_data, + aten_ops_checks, + exir_ops_checks, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() + + +class RejectFullIncrementViewMulXLessThanY(torch.nn.Module): + + def forward(self, x: torch.Tensor, y: torch.Tensor): + return ( + ( + torch.full( + ( + 1, + 3, + 5, + ), + torch.iinfo(torch.int32).max + 1, + dtype=torch.int64, + ) + + 1 + ).view(-1, 1) + * x + ) < y + + test_data = { + "randint": ( + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + torch.randint( + 0, + 10, + (1,), + dtype=torch.int32, + ), + ), + } + + +@common.parametrize("test_data", RejectFullIncrementViewMulXLessThanY.test_data) +@pytest.mark.xfail( + reason="MLETORCH-1254: Add operator support check for aten.arange and aten.full" +) +def test_reject_convert_full_int64_dtype_to_int32_pass_tosa_FP(test_data: input_t1): + module = RejectFullIncrementViewMulXLessThanY() + aten_ops_checks = [ + "torch.ops.aten.full.default", + "torch.ops.aten.add.Tensor", + "torch.ops.aten.view.default", + "torch.ops.aten.mul.Tensor", + "torch.ops.aten.lt.Tensor", + ] + pipeline = TosaPipelineFP[input_t2]( + module, + test_data, + aten_ops_checks, + exir_op=[], + transform_passes=[ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + ], + ) + pipeline.run() + + +##################################################### +## Test full(dtype=None) -> full(dtype=None/int32) ## +##################################################### + + +class AddConstFullDtypeNone(torch.nn.Module): + # Input + a full with constant value. + exir_op = "executorch_exir_dialects_edge__ops_aten_full_default" + + def __init__(self, size: tuple, fill_value: Union[bool, float, int]): + super().__init__() + self.args = (size, fill_value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.full(*self.args) + x + + test_data = { + "int64": (lambda: (torch.randn(1),), ((1, 2, 3), 10)), + "float32": (lambda: (torch.randn(1),), ((1, 2, 3), 10.0)), + } + + test_data_bool = { + "bool": (lambda: (torch.randn(1),), ((1, 2, 3), True)), + } + + +@common.parametrize("test_data", AddConstFullDtypeNone.test_data) +def test_full_dtype_none_tosa_FP(test_data): + input_data, init_data = test_data + pipeline = TosaPipelineFP[input_t1]( + AddConstFullDtypeNone(*init_data), + input_data(), + aten_op=[], + exir_op=AddConstFullDtypeNone.exir_op, + ) + pipeline.run() + + +@common.parametrize("test_data", AddConstFullDtypeNone.test_data_bool) +def test_full_dtype_none_tosa_FP_bool(test_data): + input_data, init_data = test_data + pipeline = TosaPipelineFP[input_t1]( + AddConstFullDtypeNone(*init_data), + input_data(), + aten_op=[], + exir_op=AddConstFullDtypeNone.exir_op, + ) + pipeline.change_args( + "check_count.exir", + {"torch.ops.higher_order.executorch_call_delegate": 2}, + ) + pipeline.run() + + +@common.parametrize( + "test_data", AddConstFullDtypeNone.test_data | AddConstFullDtypeNone.test_data_bool +) +def test_full_dtype_none_tosa_INT(test_data): + input_data, init_data = test_data + pipeline = TosaPipelineINT[input_t1]( + AddConstFullDtypeNone(*init_data), + input_data(), + aten_op=[], + exir_op=AddConstFullDtypeNone.exir_op, + ) + pipeline.pop_stage("check.quant_nodes") + pipeline.run() diff --git a/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py b/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py new file mode 100644 index 00000000000..ea7e03f8e21 --- /dev/null +++ b/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py @@ -0,0 +1,122 @@ +# Copyright 2025 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 ConvertInt64OutputOpsToInt32Pass + +from executorch.backends.arm.test import common + +from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP + +input_t1 = Tuple[torch.Tensor] # Input x + + +######################################### +## Test [int32 | other types] -> int64 ## +######################################### + + +class CastingToInt64Model(torch.nn.Module): + def __init__(self, target_dtype): + super().__init__() + self.target_dtype = target_dtype + + def forward(self, x: torch.Tensor): + return x.to(dtype=self.target_dtype) + + +test_data_suite_convert = { + "fp32_input": lambda: (torch.rand((1, 2, 3, 4), dtype=torch.float32), torch.int64), + "fp16_input": lambda: (torch.rand((1, 2, 3, 4), dtype=torch.float16), torch.int64), +} + +test_data_suite_remove = { + "int32_input": lambda: ( + torch.randint(-127, 128, (1, 2, 3, 4), dtype=torch.int32), + torch.int64, + ), +} + + +@common.parametrize("test_data", test_data_suite_convert) +def test_convert_or_remove_casting_to_int64_convert_tosa_FP(test_data: Tuple): + test_tensor, target_dtype = test_data() + module = CastingToInt64Model(target_dtype) + + pipeline = TosaPipelineFP[input_t1]( + module, + (test_tensor,), + aten_op="torch.ops.aten.to.dtype", + exir_op=[], + transform_passes=[ConvertInt64OutputOpsToInt32Pass()], + ) + pipeline.pop_stage( + "run_method_and_compare_outputs" + ) # As expected: RuntimeError: Int did not match Long + pipeline.run() + + +@common.parametrize("test_data", test_data_suite_remove) +def test_convert_or_remove_casting_to_int64_remove_tosa_FP(test_data: Tuple): + test_tensor, target_dtype = test_data() + module = CastingToInt64Model(target_dtype) + + pipeline = TosaPipelineFP[input_t1]( + module, + (test_tensor,), + aten_op=[], + exir_op=[], + transform_passes=[ConvertInt64OutputOpsToInt32Pass()], + ) + pipeline.change_args( + "check_count.exir", {"torch.ops.higher_order.executorch_call_delegate": 0} + ) # Empty graph without nodes + pipeline.pop_stage("run_method_and_compare_outputs") + pipeline.run() + + +##################################################### +## Test arange(dtype=int64) -> arange(dtype=int32) ## +##################################################### + + +class Int64OutputModel(torch.nn.Module): + + def forward(self, x: torch.Tensor): + # return torch.argmax(x) # RuntimeError: Int did not match Long; But this is expected as we expect _argmax_i32 to generate int32 output + # return (10 * torch.argmax(x) + 10).to(dtype=torch.int32) # [1]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (function _resize_output_check) + return (10 * torch.argmax(x, dim=-1) + 10) + 1.5 + + def get_inputs(self) -> input_t1: + return ( + torch.randint( + 0, + 10, + (2, 4, 6, 8), + ), + ) + + +def test_insert_int64_output_to_int32_cast_tosa_FP(): + module = Int64OutputModel() + aten_ops_checks = [ + "torch.ops.aten.argmax.default", + "torch.ops.aten.mul.Tensor", + "torch.ops.aten.add.Tensor", + ] + exir_ops_checks = [ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_add_Tensor", + ] + pipeline = TosaPipelineFP[input_t1]( + module, + module.get_inputs(), + aten_op=aten_ops_checks, + exir_op=exir_ops_checks, + transform_passes=[ConvertInt64OutputOpsToInt32Pass()], + ) + pipeline.run() diff --git a/backends/arm/test/passes/test_decompose_avg_pool2d_pass.py b/backends/arm/test/passes/test_decompose_avg_pool2d_pass.py new file mode 100644 index 00000000000..4d686039456 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_avg_pool2d_pass.py @@ -0,0 +1,75 @@ +# Copyright 2025 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.decompose_avg_pool2d import DecomposeAvgPool2d +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import PassPipeline + +input_t = Tuple[torch.Tensor] # Input x + + +class AvgPool2dWithStride(torch.nn.Module): + """ + avg_pool2d model with explicit stride parameter + """ + + def get_inputs(self) -> input_t: + return (torch.rand(1, 3, 8, 8),) + + def forward(self, x): + return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + + +class AvgPool2dWithoutStride(torch.nn.Module): + """ + avg_pool2d model without stride parameter (should default to kernel_size) + """ + + def get_inputs(self) -> input_t: + return (torch.rand(1, 3, 8, 8),) + + def forward(self, x): + return torch.nn.functional.avg_pool2d(x, kernel_size=3) + + +class AvgPool2dListKernel(torch.nn.Module): + """ + avg_pool2d model with list kernel_size and no stride + """ + + def get_inputs(self) -> input_t: + return (torch.rand(1, 3, 8, 8),) + + def forward(self, x): + return torch.nn.functional.avg_pool2d(x, kernel_size=[2, 3]) + + +modules = { + "avg_pool2d_with_stride": AvgPool2dWithStride(), + "avg_pool2d_without_stride": AvgPool2dWithoutStride(), + "avg_pool2d_list_kernel": AvgPool2dListKernel(), +} + + +@common.parametrize("module", modules) +def test_decompose_avg_pool2d_tosa_MI(module): + """Test that DecomposeAvgPool2d pass works correctly with and without stride parameters.""" + pipeline = PassPipeline[input_t]( + module, + module.get_inputs(), + quantize=False, + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1, + }, + ops_after_pass={ + # After decomposition, we should still see avg_pool2d (transformed) + "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1, + }, + pass_list=[DecomposeAvgPool2d], + ) + pipeline.run() diff --git a/backends/arm/test/passes/test_fuse_constant_ops_pass.py b/backends/arm/test/passes/test_fuse_constant_ops_pass.py index 1a318c5cd42..417ad7bff2a 100644 --- a/backends/arm/test/passes/test_fuse_constant_ops_pass.py +++ b/backends/arm/test/passes/test_fuse_constant_ops_pass.py @@ -12,7 +12,11 @@ FuseConstantArgsPass, ) from executorch.backends.arm.test import common -from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.arm.test.tester.test_pipeline import ( + PassPipeline, + TosaPipelineFP, + TosaPipelineINT, +) input_t = Tuple[torch.Tensor] # Input x input_t2 = Tuple[torch.Tensor, torch.Tensor] @@ -103,6 +107,22 @@ def forward(self, a, b): return torch.cat((a, b), dim=0) +class LinearConst(torch.nn.Module): + """A linear layer that can be computed AOT""" + + def __init__(self, in_out_features: int = 3, bias: bool = True): + super().__init__() + self.linear = torch.nn.Linear(in_out_features, in_out_features, bias=bias) + self.example_input = torch.rand(in_out_features, in_out_features) + + def forward(self, x: torch.Tensor): + y = torch.full_like(x, 1.0) + return self.linear(y) + x + + def get_example_input(self): + return self.example_input + + modules = { "fuse_parameter": FuseParameter(), "fuse_buffer": FuseBuffer(), @@ -152,3 +172,30 @@ def test_fuse_const_ops_tosa_BI_cat(module: torch.nn.Module): passes_with_exported_program=[ComputeConstantOpsAOT, FuseConstantArgsPass], ) pipeline.run() + + +def test_linear_const_tosa_FP(): + model = LinearConst() + example_input = model.get_example_input() + pipeline = TosaPipelineFP[input_t]( + model, + (example_input,), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.run() + + +def test_linear_const_tosa_INT(): + model = LinearConst() + example_input = model.get_example_input() + pipeline = TosaPipelineINT[input_t]( + model, + (example_input,), + aten_op=[], + exir_op=[], + per_channel_quantization=False, + use_to_edge_transform_and_lower=True, + ) + pipeline.run() diff --git a/backends/arm/test/passes/test_insert_int64_to_int32_cast_pass.py b/backends/arm/test/passes/test_insert_int32_casts_after_int64_placeholders_pass.py similarity index 83% rename from backends/arm/test/passes/test_insert_int64_to_int32_cast_pass.py rename to backends/arm/test/passes/test_insert_int32_casts_after_int64_placeholders_pass.py index da6eeb59459..efc1bebb610 100644 --- a/backends/arm/test/passes/test_insert_int64_to_int32_cast_pass.py +++ b/backends/arm/test/passes/test_insert_int32_casts_after_int64_placeholders_pass.py @@ -6,7 +6,7 @@ from typing import Tuple import torch -from executorch.backends.arm._passes import InsertCastForOpsWithInt64InputPass +from executorch.backends.arm._passes import InsertInt32CastsAfterInt64PlaceholdersPass from executorch.backends.arm.test.tester.test_pipeline import PassPipeline @@ -31,7 +31,7 @@ def test_int64_model_tosa_FP(): "executorch_exir_dialects_edge__ops_aten_embedding_default": 1, } op_checks_after = { - "executorch_exir_dialects_edge__ops_aten__to_copy_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, "executorch_exir_dialects_edge__ops_aten_embedding_default": 1, } @@ -40,7 +40,7 @@ def test_int64_model_tosa_FP(): module.get_inputs(), ops_before_pass=op_checks_before, ops_after_pass=op_checks_after, - pass_list=[InsertCastForOpsWithInt64InputPass], + pass_list=[InsertInt32CastsAfterInt64PlaceholdersPass], ) pipeline.pop_stage(-1) # Do not compare output pipeline.run() diff --git a/backends/arm/test/passes/test_insert_table_ops_pass.py b/backends/arm/test/passes/test_insert_table_ops_pass.py index 029942dd659..5e695c237a0 100644 --- a/backends/arm/test/passes/test_insert_table_ops_pass.py +++ b/backends/arm/test/passes/test_insert_table_ops_pass.py @@ -33,13 +33,13 @@ def test_insert_table_tosa_INT(test_data: input_t): module, test_data, quantize=True, - ops_before_pass={}, + ops_before_pass={"executorch_exir_dialects_edge__ops_aten_sigmoid_default": 1}, ops_after_pass={ "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1, - "tosa._table": 1, + "backend__ops_tosa_TABLE_default": 1, }, - ops_not_after_pass=["aten_sigmoid_default"], + ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_sigmoid_default"], pass_list=[FoldAndAnnotateQParamsPass], passes_with_exported_program=[InsertTableOpsPass], ) diff --git a/backends/arm/test/passes/test_remove_clone_pass.py b/backends/arm/test/passes/test_remove_clone_pass.py index dea0bb06f5e..180f9c2ffe5 100755 --- a/backends/arm/test/passes/test_remove_clone_pass.py +++ b/backends/arm/test/passes/test_remove_clone_pass.py @@ -6,7 +6,7 @@ from typing import Tuple import torch -from executorch.backends.arm._passes.remove_clone_pass import RemoveClonePass +from executorch.backends.arm._passes.remove_noop_pass import RemoveNoopPass from executorch.backends.arm.test.tester.test_pipeline import PassPipeline @@ -15,7 +15,7 @@ class Clone(torch.nn.Module): """ - Basic remove layer model to test RemoveClonePass + Basic remove layer model to test RemoveNoopePass """ def __init__(self): @@ -35,9 +35,11 @@ def test_remove_clone_tosa_INT(): module.get_inputs(), quantize=True, ops_before_pass={ - "executorch_exir_dialects_edge__ops_aten_clone_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, }, - ops_not_after_pass=["executorch_exir_dialects_edge__ops_aten_clone_default"], - pass_list=[RemoveClonePass], + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default" + ], + pass_list=[RemoveNoopPass], ) pipeline.run() diff --git a/backends/arm/test/passes/test_rescale_pass.py b/backends/arm/test/passes/test_rescale_pass.py index 0fe72f6d1fe..0959a0eaa25 100644 --- a/backends/arm/test/passes/test_rescale_pass.py +++ b/backends/arm/test/passes/test_rescale_pass.py @@ -9,13 +9,18 @@ import pytest import torch -import torch.library from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, TosaPipelineINT, ) +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 input_t = Tuple[torch.Tensor, torch.Tensor] # Input x @@ -45,8 +50,19 @@ def test_rescale_op(): 127, ), ] - for sample_input in sample_inputs[1:2]: - torch.library.opcheck(torch.ops.tosa._rescale, sample_input) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.0+INT") + ), FakeTensorMode() as mode: + for sample_input in sample_inputs: + exir_ops.backend.tosa.RESCALE.default( + *tuple( + [ + mode.from_tensor(i) if isinstance(i, torch.Tensor) else i + for i in sample_input + ] + ) + ) def test_nonzero_zp_for_int32(): @@ -67,9 +83,22 @@ def test_nonzero_zp_for_int32(): 1, # Should be 0, expect error ), ] - for sample_input in sample_inputs: - with pytest.raises(Exception, match="opcheck"): - torch.library.opcheck(torch.ops.tosa._rescale, sample_input) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.0+INT") + ), FakeTensorMode() as mode: + for sample_input in sample_inputs: + with pytest.raises( + ValueError, match="TOSA requires (output|input)_zp to be zero" + ): + exir_ops.backend.tosa.RESCALE.default( + *tuple( + [ + mode.from_tensor(i) if isinstance(i, torch.Tensor) else i + for i in sample_input + ] + ) + ) def test_zp_outside_range(): @@ -90,9 +119,21 @@ def test_zp_outside_range(): -129, # Should be >-129m expect error ), ] - for sample_input in sample_inputs: - with pytest.raises(Exception, match="opcheck"): - torch.library.opcheck(torch.ops.tosa._rescale, sample_input) + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.0+INT") + ), FakeTensorMode() as mode: + for sample_input in sample_inputs: + with pytest.raises( + Exception, match="(in_zp|out_zp)=-?[0-9]* outside valid range" + ): + exir_ops.backend.tosa.RESCALE.default( + *tuple( + [ + mode.from_tensor(i) if isinstance(i, torch.Tensor) else i + for i in sample_input + ] + ) + ) class RescaleNetwork(torch.nn.Module): @@ -131,7 +172,14 @@ def test_quantized_rescale_tosa_bi(test_data: tuple[torch.Tensor, torch.Tensor]) pipeline.run() -@common.parametrize("test_data", RescaleNetwork.test_data) +u55_xfails = { + "ones": "MLBEDSW-11032: ILLEGAL_OFM_BASE error: Base addresses must be aligned to brick depth on u55.", + "randn_ones": "MLBEDSW-11032: ILLEGAL_OFM_BASE error: Base addresses must be aligned to brick depth on u55.", + "randn_large": "MLBEDSW-11032: ILLEGAL_OFM_BASE error: Base addresses must be aligned to brick depth on u55.", +} + + +@common.parametrize("test_data", RescaleNetwork.test_data, xfails=u55_xfails) @common.XfailIfNoCorstone300 def test_quantized_rescale_u55(test_data: tuple[torch.Tensor, torch.Tensor]): """Tests a model with many ops that requires rescales. As more ops are quantized to int32 and diff --git a/backends/arm/test/passes/test_to_tosa_memory_format.py b/backends/arm/test/passes/test_to_tosa_memory_format.py new file mode 100644 index 00000000000..1e9b8ffc63d --- /dev/null +++ b/backends/arm/test/passes/test_to_tosa_memory_format.py @@ -0,0 +1,192 @@ +# Copyright 2025 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 ToTosaMemoryFormatPass + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + PassPipeline, + TosaPipelineINT, +) +from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass + +input_t = Tuple[torch.Tensor] # Input x + + +class NoNHWC(torch.nn.Module): + """ + Test-module with no ops requiring NHWC mermory format. + """ + + ops_after_pass = {"executorch_exir_dialects_backend__ops_tosa_TRANSPOSE_default": 2} + ops_not_after_pass = [] + + def forward(self, x): + x = x + x + return x + + def get_inputs(self): + return (torch.rand(1, 2, 2, 2),) + + +class ParallelClusters(torch.nn.Module): + """ + Test-module with multiple parallel clusters of nodes requiring different memory formats. + """ + + ops_after_pass = {"executorch_exir_dialects_backend__ops_tosa_TRANSPOSE_default": 2} + ops_not_after_pass = [] + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels=2, + out_channels=2, + kernel_size=1, + bias=True, + ) + self.maxpool = torch.nn.MaxPool2d(1, 1) + self.avgpool = torch.nn.AvgPool2d(1, 1) + + def forward(self, x): + x1 = self.conv(x) + x2 = self.maxpool(x) + x3 = self.avgpool(x) + x4 = x * x + return x1 + x2 + x3 + x4 + + def get_inputs(self): + return (torch.rand(1, 2, 2, 2),) + + +class SerialClusters(torch.nn.Module): + """ + Test-module with multiple serial clusters of nodes requring different memory formats. + """ + + ops_before_pass = {} + ops_after_pass = {"executorch_exir_dialects_backend__ops_tosa_TRANSPOSE_default": 4} + ops_not_after_pass = [] + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels=2, + out_channels=2, + kernel_size=1, + bias=True, + ) + self.fc = torch.nn.Linear( + in_features=2, + out_features=2, + bias=True, + ) + + def forward(self, x): + x = self.conv(x) + x = x * x + x = self.conv(x) + x = x.view((2, 1, 2, 4)) + x = x * 2 + x = x.view((2, 2, 2, 2)) + x = self.conv(x) + return x + + def get_inputs(self): + return (torch.rand(2, 2, 2, 2),) + + +class Reshapes(torch.nn.Module): + """ + Test-module with different configurations of views requiring different memory formats. + """ + + ops_before_pass = {} + ops_after_pass = { + "executorch_exir_dialects_backend__ops_tosa_TRANSPOSE_default": 16 + } + ops_not_after_pass = [] + + def __init__(self): + super().__init__() + self.maxpool = torch.nn.MaxPool2d(1, 1) # Use maxpool to force NHWC format + + def forward(self, x): + + x = self.maxpool(x) + x = x.view((2, 2, 4, 16, 1)) # N-C-HW-invariant intact, no transposes needed + x = x * 2 # Add op to avoid views merging + x = x.view((4, 4, 4, 4)) + x = x / 2 # Add op to avoid views merging + x = self.maxpool(x) + + x = x.view((256)) # Break N-C-HW invariant + x = x * 2 + x = x.view((4, 4, 4, 4)) + x = x / 2 + x = self.maxpool(x) + + x = x.view((16, 16)) # Break N-C-HW invariant + x = x * 2 + x = x.view((4, 4, 4, 4)) + x = x / 2 + x = self.maxpool(x) + + x = x.view((16, 4, 4)) # Break N-C-HW invariant + x = x * 2 + x = x.view((4, 4, 4, 4)) + x = x / 2 + x = self.maxpool(x) + + x = x.view((2, 4, 4, 8)) # Break N-C-HW invariant + x = x * 2 + x = x.view((4, 4, 4, 4)) + x = x / 2 + x = self.maxpool(x) + + x = x.view((8, 1, 2, 4, 4)) # Break N-C-HW invariant + x = x * 2 + x = x.view((4, 4, 4, 4)) + x = self.maxpool(x) + + return x + + def get_inputs(self): + return (torch.rand(4, 4, 4, 4),) + + +modules = { + "no_nhwc": NoNHWC(), + "parallel_clusters": ParallelClusters(), + "serial_clusters": SerialClusters(), + "reshapes": Reshapes(), +} + + +@common.parametrize("module", modules) +def test_to_tosa_memory_format_tosa_INT(module): + # We cannot check op counts after a specific pass with the full pipeline + pipeline = PassPipeline[input_t]( + module, + module.get_inputs(), + ops_after_pass=module.ops_after_pass, + ops_not_after_pass=module.ops_not_after_pass, + pass_list=[RemoveGetItemPass], + passes_with_exported_program=[ToTosaMemoryFormatPass], + ) + pipeline.pop_stage( + "run_method_and_compare_outputs" + ) # Eager execution is not possible after introducing tosa.TRANSPOSE + pipeline.run() + + +@common.parametrize("module", modules) +def test_to_tosa_memory_format_tosa_INT_functional(module): + # Also run the actual pass pipeline to ensure functional correctness. + pipeline = TosaPipelineINT[input_t](module, module.get_inputs(), []) + pipeline.run() diff --git a/backends/arm/test/runner_utils.py b/backends/arm/test/runner_utils.py index bd06e817d8f..aeb0e3a56bd 100644 --- a/backends/arm/test/runner_utils.py +++ b/backends/arm/test/runner_utils.py @@ -18,10 +18,13 @@ import numpy as np import torch -from executorch.backends.arm.arm_backend import get_tosa_spec, is_tosa +from executorch.backends.arm.arm_backend import is_tosa, is_vgf from executorch.backends.arm.test.conftest import is_option_enabled -from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification - +from executorch.backends.arm.tosa.specification import ( + get_tosa_spec, + Tosa_1_00, + TosaSpecification, +) from executorch.exir import ExecutorchProgramManager, ExportedProgram from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.lowered_backend_module import LoweredBackendModule @@ -54,6 +57,8 @@ torch.complex128: np.complex128, } +VALID_TARGET = {"corstone-300", "corstone-320", "vkml_emulation_layer"} + class QuantizationParams: __slots__ = ["node_name", "zp", "scale", "qmin", "qmax", "dtype"] @@ -125,28 +130,8 @@ def get_input_quantization_params( return quant_params -def get_output_nodes(program: ExportedProgram) -> list[Node]: - """ - Get output node to this model. - - Args: - program (ExportedProgram): The program to get the output nodes from. - Returns: - The nodes that are the outputs of the 'program'. - """ - output_nodes = [] - for node in program.graph.nodes: - if node.op == "output": - for output in node.args[0]: - output_nodes.append(output) - if len(output_nodes) == 0: - raise RuntimeError("No output nodes found.") - else: - return output_nodes - - def get_output_quantization_params( - output_nodes: list[Node], + output_node: Node, ) -> dict[Node, QuantizationParams | None]: """ Get output QuantizationParams from a program. @@ -159,7 +144,7 @@ def get_output_quantization_params( RuntimeError if no output quantization parameters are found. """ quant_params = {} - for node in output_nodes: + for node in output_node.args[0]: if node.target == torch.ops.quantized_decomposed.dequantize_per_tensor.default: quant_params[node] = QuantizationParams( node_name=node.args[0].name, @@ -215,6 +200,117 @@ def __torch_function__(self, func, types, args=..., kwargs=None): return func(*args, **kwargs) +def run_target( + executorch_program_manager: ExecutorchProgramManager, + inputs: Tuple[torch.Tensor], + intermediate_path: str | Path, + target_board: Literal["corestone-300", "corestone-320", "vkml_emulation_layer"], + elf_path: str | Path, + timeout: int = 120, # s +): + if target_board not in VALID_TARGET: + raise ValueError(f"Unsupported target: {target_board}") + + if target_board in ("corstone-300", "corstone-320"): + return run_corstone( + executorch_program_manager, + inputs, + intermediate_path, + target_board, + elf_path, + timeout, + ) + elif target_board == "vkml_emulation_layer": + return run_vkml_emulation_layer( + executorch_program_manager, + inputs, + intermediate_path, + elf_path, + ) + + +def save_inputs_to_file( + exported_program: ExportedProgram, + inputs: Tuple[torch.Tensor], + intermediate_path: str | Path, +): + input_file_paths = [] + input_names = get_input_names(exported_program) + for input_name, input_ in zip(input_names, inputs): + input_path = save_bytes(intermediate_path, input_, input_name) + input_file_paths.append(input_path) + + return input_file_paths + + +def get_output_from_file( + exported_program: ExportedProgram, + intermediate_path: str | Path, + output_base_name: str, +): + output_np = [] + output_node = exported_program.graph_module.graph.output_node() + for i, node in enumerate(output_node.args[0]): + output_shape = node.meta["val"].shape + output_dtype = node.meta["val"].dtype + tosa_ref_output = np.fromfile( + os.path.join(intermediate_path, f"{output_base_name}-{i}.bin"), + _torch_to_numpy_dtype_dict[output_dtype], + ) + + output_np.append(torch.from_numpy(tosa_ref_output).reshape(output_shape)) + return tuple(output_np) + + +def run_vkml_emulation_layer( + executorch_program_manager: ExecutorchProgramManager, + inputs: Tuple[torch.Tensor], + intermediate_path: str | Path, + elf_path: str | Path, +): + """Executes an inference of the exported_program on ML Emulation Layer for Vulkan + Args: + `executorch_program_manager`: The executorch program to run. + `intermediate_path`: Directory to save the .pte and capture outputs. + `elf_path`: Path to the Vulkan-capable executor_runner binary. + """ + exported_program = executorch_program_manager.exported_program() + intermediate_path = Path(intermediate_path) + intermediate_path.mkdir(exist_ok=True) + elf_path = Path(elf_path) + if not elf_path.exists(): + raise FileNotFoundError(f"Did not find elf file {elf_path}") + + # Save pte to file + pte_path = os.path.join(intermediate_path, "program.pte") + with open(pte_path, "wb") as f: + f.write(executorch_program_manager.buffer) + + output_base_name = "out" + out_path = os.path.join(intermediate_path, output_base_name) + + cmd_line = f"{elf_path} -model_path {pte_path} -output_file {out_path}" + + input_string = None + input_paths = save_inputs_to_file(exported_program, inputs, intermediate_path) + for input_path in input_paths: + if input_string is None: + input_string = f" -inputs={input_path}" + else: + input_string += f",{input_path}" + if input_string is not None: + cmd_line += input_string + cmd_line = cmd_line.split() + + result = _run_cmd(cmd_line) + + # TODO: MLETORCH-1234: Support VGF e2e tests in VgfPipeline + # TODO: Add regex to check for error or fault messages in stdout from Emulation Layer + result_stdout = result.stdout.decode() # noqa: F841 + + return get_output_from_file(exported_program, intermediate_path, output_base_name) + + def run_corstone( executorch_program_manager: ExecutorchProgramManager, inputs: Tuple[torch.Tensor], @@ -226,7 +322,7 @@ def run_corstone( """Executes an inference of the exported_program on FVP. Returns a list of tensors with the output. Args: - `executorch_program_manager`: the executorch program to run. + `executorch_program_manager`: The executorch program to run. The output of a EdgeProgramManager.to_executorch() call. `inputs`: A list of tensors with the inputs of the inference. `dump_path`: A directory where the .pte and inputs are saved to file. @@ -254,14 +350,10 @@ def run_corstone( with open(pte_path, "wb") as f: f.write(executorch_program_manager.buffer) - # Save inputs to file - input_names = get_input_names(exported_program) - input_paths = [] - for input_name, input_ in zip(input_names, inputs): - input_path = save_bytes(intermediate_path, input_, input_name) - input_paths.append(input_path) + input_paths = save_inputs_to_file(exported_program, inputs, intermediate_path) - out_path = os.path.join(intermediate_path, "out") + output_base_name = "out" + out_path = os.path.join(intermediate_path, output_base_name) cmd_line = f"executor_runner -m {pte_path} -o {out_path}" for input_path in input_paths: @@ -343,18 +435,7 @@ def run_corstone( f"Corstone simulation failed:\ncmd: {' '.join(command_args)}\nlog: \n {result_stdout}\n{result.stderr.decode()}" ) - output_nodes = get_output_nodes(exported_program) - output_np = [] - for i, node in enumerate(output_nodes): - output_shape = node.meta["val"].shape - output_dtype = node.meta["val"].dtype - tosa_ref_output = np.fromfile( - os.path.join(intermediate_path, f"out-{i}.bin"), - _torch_to_numpy_dtype_dict[output_dtype], - ) - - output_np.append(torch.from_numpy(tosa_ref_output).reshape(output_shape)) - return tuple(output_np) + return get_output_from_file(exported_program, intermediate_path, output_base_name) def prep_data_for_save( @@ -555,18 +636,53 @@ def model_converter_installed() -> bool: return True -def get_elf_path(target_board): - elf_path = os.path.join( - "arm_test", - f"arm_semihosting_executor_runner_{target_board}", - "arm_executor_runner", - ) +def vkml_emulation_layer_installed() -> bool: + # Check VK_INSTANCE_LAYERS + vk_instance_layers = os.environ.get("VK_INSTANCE_LAYERS", "") + required_layers = { + "VK_LAYER_ML_Graph_Emulation", + "VK_LAYER_ML_Tensor_Emulation", + } + existing_layers = set(vk_instance_layers.split(":")) + layers_exists = required_layers.issubset(existing_layers) + + # Check LD_LIBRARY_PATH for "emulation-layer/deploy" + ld_library_path = os.environ.get("LD_LIBRARY_PATH", "") + deploy_exists = False + for path in ld_library_path.split(os.path.pathsep): + if "emulation-layer/deploy" in path and os.path.isdir(path): + deploy_exists = True + + return layers_exists and deploy_exists + + +def assert_elf_path_exists(elf_path): if not os.path.exists(elf_path): raise FileNotFoundError( - f"Did not find build arm_executor_runner in path {elf_path}, run setup_testing.sh?" + f"Did not find build arm_executor_runner or executor_runner in path {elf_path}, \ + run setup_testing.sh or setup_testing_vkml.sh?" ) - else: - return elf_path + + +def get_elf_path(target_board): + if target_board not in VALID_TARGET: + raise ValueError(f"Unsupported target: {target_board}") + + if target_board in ("corstone-300", "corstone-320"): + elf_path = os.path.join( + "arm_test", + f"arm_semihosting_executor_runner_{target_board}", + "arm_executor_runner", + ) + assert_elf_path_exists(elf_path) + elif target_board == "vkml_emulation_layer": + elf_path = os.path.join( + "arm_test/arm_executor_runner_vkml", + "executor_runner", + ) + assert_elf_path_exists(elf_path) + + return elf_path def arm_executor_runner_exists(target_board): @@ -585,7 +701,6 @@ def run_tosa_graph( ) -> list[torch.Tensor]: """Runs the TOSA reference model with inputs and returns the result.""" inputs_np = [input.numpy() for input in inputs] - transpose_data_format(inputs_np, to="NHWC") if isinstance(tosa_version, Tosa_1_00): import tosa_reference_model as reference_model @@ -607,25 +722,12 @@ def run_tosa_graph( status == reference_model.GraphStatus.TOSA_VALID ), "Non-valid TOSA given to reference model." - transpose_data_format(outputs_np, to="NCHW") return [torch.from_numpy(output) for output in outputs_np] -def transpose_data_format(data: list[np.ndarray], to: Literal["NHWC", "NCHW"]): - for i in range(len(data)): - if hasattr(data[i], "shape") and data[i].ndim in (4, 5): - match to: - case "NCHW": - dim_order = (0, 3, 1, 2) if data[i].ndim == 4 else (0, 1, 4, 2, 3) - case "NHWC": - dim_order = (0, 2, 3, 1) if data[i].ndim == 4 else (0, 1, 3, 4, 2) - case _: - raise NotImplementedError(f"Cant transpose to dim order {to}") - # Copy is needed to force actual data conversion, not setting stride. - data[i] = np.transpose(data[i], dim_order).copy() - - def get_target_board(compile_spec: list[CompileSpec]) -> str | None: + if is_vgf(compile_spec): + return "vkml_emulation_layer" for spec in compile_spec: if spec.key == "compile_flags": flags = spec.value.decode() diff --git a/backends/arm/test/setup_testing.sh b/backends/arm/test/setup_testing.sh index fd47a6bb464..d1e4725d93b 100755 --- a/backends/arm/test/setup_testing.sh +++ b/backends/arm/test/setup_testing.sh @@ -1,58 +1,15 @@ #!/usr/bin/env bash # Copyright 2024-2025 Arm Limited and/or its 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. set -eu -script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -et_root_dir=$(cd ${script_dir}/../../.. && pwd) -ethos_u_root_dir=${et_root_dir}/examples/arm/ethos-u-scratch/ethos-u - -toolchain_cmake=${et_root_dir}/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake -et_build_dir=${et_root_dir}/arm_test/cmake-out +script_dir=$(realpath "$(dirname "${BASH_SOURCE[0]}")") +et_root_dir=$(realpath "${script_dir}/../../..") +build_executor_runner=${et_root_dir}/backends/arm/scripts/build_executor_runner.sh build_root_test_dir=${et_root_dir}/arm_test/arm_semihosting_executor_runner -# Build Arm Baremetal executor_runner in semihosting mode. -# Put in backends/arm/test/res to be used by unit tests. -function build_semihosting_executorch_runner() { - target_board=$1 - system_config=$2 - build_test_dir=${build_root_test_dir}_${target_board} - echo "[${FUNCNAME[0]}] Configuring ${target_board} with system config ${system_config}" - if [[ ${target_board} == "corstone-300" ]]; then - local target_cpu=cortex-m55 - elif [[ ${target_board} == "corstone-320" ]]; then - local target_cpu=cortex-m85 - else - echo "[${FUNCNAME[0]}] ERROR: Invalid target_board specified!" - exit 1 - fi - cd ${et_root_dir}/examples/arm/executor_runner - pwd - mkdir -p ${build_test_dir} - cmake -DCMAKE_TOOLCHAIN_FILE=${toolchain_cmake} \ - -DCMAKE_BUILD_TYPE=RelWithDebInfo \ - -DTARGET_CPU=${target_cpu} \ - -DSEMIHOSTING=ON \ - -DCMAKE_RUNTIME_OUTPUT_DIRECTORY=${build_test_dir} \ - -DETHOS_SDK_PATH:PATH=${ethos_u_root_dir} \ - -DET_DIR_PATH:PATH=${et_root_dir} \ - -DET_BUILD_DIR_PATH:PATH=${et_build_dir} \ - -DPYTHON_EXECUTABLE=$(which python3) \ - -DSYSTEM_CONFIG=${system_config} \ - -B ${build_test_dir} - echo "[${FUNCNAME[0]}] Configured CMAKE" - - n=$(nproc) - cmake --build ${build_test_dir} -j"$((n - 5))" -- arm_executor_runner - echo "[${FUNCNAME[0]}] Generated baremetal elf file: with semihosting enabled" - find ${build_test_dir} -name "arm_executor_runner" -} - -# Use most optimal system_configs for testing -build_semihosting_executorch_runner corstone-300 Ethos_U55_High_End_Embedded - -build_semihosting_executorch_runner corstone-320 Ethos_U85_SYS_DRAM_Mid +${build_executor_runner} --pte=semihosting --target=ethos-u55-128 --output="${build_root_test_dir}_corstone-300" +${build_executor_runner} --pte=semihosting --target=ethos-u85-128 --output="${build_root_test_dir}_corstone-320" diff --git a/backends/arm/test/setup_testing_vkml.sh b/backends/arm/test/setup_testing_vkml.sh new file mode 100755 index 00000000000..6edb2dca361 --- /dev/null +++ b/backends/arm/test/setup_testing_vkml.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Copyright 2024-2025 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 -eu + +script_dir=$(realpath "$(dirname "${BASH_SOURCE[0]}")") +et_root_dir=$(realpath "${script_dir}/../../..") +build_executor_runner=${et_root_dir}/backends/arm/scripts/build_executor_runner_vkml.sh +build_root_test_dir=${et_root_dir}/arm_test/arm_executor_runner_vkml + +${build_executor_runner} --output="${build_root_test_dir}" diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index acb27f13798..a6181cf34ce 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -13,11 +13,16 @@ def define_arm_tests(): # Operators test_files += [ + "ops/test_add.py", + "ops/test_addmm.py", "ops/test_avg_pool2d.py", + "ops/test_cat.py", "ops/test_linear.py", + "ops/test_mul.py", "ops/test_slice.py", "ops/test_sigmoid.py", "ops/test_tanh.py", + "ops/test_view.py", "ops/test_cos.py", ] diff --git a/backends/arm/test/test_arm_baremetal.sh b/backends/arm/test/test_arm_baremetal.sh index 609a8430522..53c707cad28 100755 --- a/backends/arm/test/test_arm_baremetal.sh +++ b/backends/arm/test/test_arm_baremetal.sh @@ -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. +# TODO: Rename this script + set -e script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) @@ -17,7 +19,6 @@ _setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly ins TEST_SUITE=$1 -TOSA_VERSION="${2:-TOSA-1.0+INT}" # Source the tools # This should be prepared by the setup.sh @@ -101,7 +102,7 @@ test_pytest_models() { # Test ops and other things source backends/arm/scripts/install_models_for_test.sh # Run arm baremetal pytest tests without FVP - pytest --verbose --color=yes --durations=0 backends/arm/test/models + pytest --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models echo "${TEST_SUITE_NAME}: PASS" } @@ -117,7 +118,6 @@ test_pytest_ops_ethosu_fvp() { # Same as test_pytest but also sometime verify us # Prepare Corstone-3x0 FVP for pytest backends/arm/scripts/build_executorch.sh - backends/arm/scripts/build_portable_kernels.sh # Build semihosting version of the runner used by pytest testing. This builds: # arm_test/arm_semihosting_executor_runner_corstone-300 # arm_test/arm_semihosting_executor_runner_corstone-320 @@ -133,7 +133,6 @@ test_pytest_models_ethosu_fvp() { # Same as test_pytest but also sometime verify # Prepare Corstone-3x0 FVP for pytest backends/arm/scripts/build_executorch.sh - backends/arm/scripts/build_portable_kernels.sh # Build semihosting version of the runner used by pytest testing. This builds: # arm_test/arm_semihosting_executor_runner_corstone-300 # arm_test/arm_semihosting_executor_runner_corstone-320 @@ -143,7 +142,7 @@ test_pytest_models_ethosu_fvp() { # Same as test_pytest but also sometime verify source backends/arm/scripts/install_models_for_test.sh # Run arm baremetal pytest tests with FVP - pytest --verbose --color=yes --durations=0 backends/arm/test/models + pytest --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models echo "${TEST_SUITE_NAME}: PASS" } @@ -154,32 +153,111 @@ test_pytest_ethosu_fvp() { # Same as test_pytest but also sometime verify using echo "${TEST_SUITE_NAME}: PASS" } + +test_pytest_ops_vkml() { # Same as test_pytest but also sometime verify using VKML runtime + echo "${TEST_SUITE_NAME}: Run pytest with VKML" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing_vkml.sh + + pytest --verbose --color=yes --numprocesses=auto --durations=10 backends/arm/test/ --ignore=backends/arm/test/models + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_models_vkml() { # Same as test_pytest but also sometime verify VKML runtime + echo "${TEST_SUITE_NAME}: Run pytest with VKML" + + backends/arm/scripts/build_executorch.sh + backends/arm/test/setup_testing_vkml.sh + + # Install model dependencies for pytest + source backends/arm/scripts/install_models_for_test.sh + + pytest --verbose --color=yes --numprocesses=auto --durations=0 backends/arm/test/models + echo "${TEST_SUITE_NAME}: PASS" +} + +test_pytest_vkml() { # Same as test_pytest but also sometime verify VKML runtime + echo "${TEST_SUITE_NAME}: Run pytest with VKML" + test_pytest_ops_vkml + test_pytest_models_vkml + echo "${TEST_SUITE_NAME}: PASS" +} + +test_run_vkml() { # End to End model tests using run.sh + echo "${TEST_SUITE_NAME}: Test VKML delegate examples with run.sh" + + echo "${TEST_SUITE_NAME}: Test VKML" + out_folder="arm_test/test_run" + examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=add --output=${out_folder}/runner + examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=mul --output=${out_folder}/runner + + examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=qadd --output=${out_folder}/runner + examples/arm/run.sh --et_build_root=${out_folder} --target=vgf --model_name=qops --output=${out_folder}/runner + + echo "${TEST_SUITE_NAME}: PASS" +} + test_run_ethosu_fvp() { # End to End model tests using run.sh echo "${TEST_SUITE_NAME}: Test ethos-u delegate examples with run.sh" # TOSA quantized - echo "${TEST_SUITE_NAME}: Test ethos-u target TOSA" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=${TOSA_VERSION} --model_name=add - examples/arm/run.sh --et_build_root=arm_test/test_run --target=${TOSA_VERSION} --model_name=mul + echo "${TEST_SUITE_NAME}: Test target TOSA" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=add + examples/arm/run.sh --et_build_root=arm_test/test_run --target=TOSA-1.0+INT --model_name=mul # Ethos-U55 - echo "${TEST_SUITE_NAME}: Test ethos-u target Ethos-U55" - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add + echo "${TEST_SUITE_NAME}: Test target Ethos-U55" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-64 --model_name=add + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --bundleio --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --etdump examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-256 --model_name=add --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=mul --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/add.py --bundleio # Ethos-U85 - echo "${TEST_SUITE_NAME}: Test ethos-u target Ethos-U85" + echo "${TEST_SUITE_NAME}: Test target Ethos-U85" examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=add - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=add --bundleio + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-512 --model_name=add --bundleio --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=add --etdump + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-2048 --model_name=mul --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=mul --bundleio --pte_placement=elf + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-256 --model_name=mul --bundleio --pte_placement=0x38000000 + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=examples/arm/example_modules/add.py + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-1024 --model_name=examples/arm/example_modules/add.py --bundleio # Cortex-M op tests + echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on a Ethos-U)" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --no_delegate --select_ops_list="aten::add.out" examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qadd --bundleio examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio - examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --portable_kernels="aten::sub.out,aten::add.out,aten::mul.out" + examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --select_ops_list="aten::sub.out,aten::add.out,aten::mul.out" examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u85-128 --model_name=qops --bundleio echo "${TEST_SUITE_NAME}: PASS" - } +} + +test_models_vkml() { # End to End model tests using model_test.py + echo "${TEST_SUITE_NAME}: Test VKML delegated models with test_model.py" + + # Build common libs once + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --build_libs + + # VKML + echo "${TEST_SUITE_NAME}: Test target VKML" + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=vgf --model=mv2 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=vgf --no_quantize --model=mv2 + + echo "${TEST_SUITE_NAME}: PASS" +} test_models_tosa() { # End to End model tests using model_test.py echo "${TEST_SUITE_NAME}: Test TOSA delegated models with test_model.py" @@ -189,20 +267,20 @@ test_models_tosa() { # End to End model tests using model_test.py # TOSA quantized echo "${TEST_SUITE_NAME}: Test ethos-u target TOSA" - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=mv2 - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=mv3 - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=lstm - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=edsr - # python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=emformer_transcribe # Takes long time to run - # python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=emformer_join # Takes long time to run - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=w2l - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=ic3 - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=ic4 - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=resnet18 - python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=${TOSA_VERSION} --model=resnet50 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=mv2 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=mv3 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=lstm + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=edsr + # python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=emformer_transcribe # Takes long time to run + # python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=emformer_join # Takes long time to run + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=w2l + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=ic3 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=ic4 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=resnet18 + python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=TOSA-1.0+INT --model=resnet50 echo "${TEST_SUITE_NAME}: PASS" - } +} test_models_ethos-u55() { # End to End model tests using model_test.py echo "${TEST_SUITE_NAME}: Test Ethos-U55 delegated models with test_model.py" @@ -221,7 +299,7 @@ test_models_ethos-u55() { # End to End model tests using model_test.py #python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u55-128 --model=resnet50 --extra_flags="-DET_ATOL=6.2 -DET_RTOL=6.2" echo "${TEST_SUITE_NAME}: PASS" - } +} test_models_ethos-u85() { # End to End model tests using model_test.py echo "${TEST_SUITE_NAME}: Test Ethos-U85 delegated models with test_model.py" @@ -240,7 +318,7 @@ test_models_ethos-u85() { # End to End model tests using model_test.py python3 backends/arm/test/test_model.py --test_output=arm_test/test_model --target=ethos-u85-128 --model=resnet50 --extra_flags="-DET_ATOL=0.2 -DET_RTOL=0.2" echo "${TEST_SUITE_NAME}: PASS" - } +} test_full_ethosu_fvp() { # All End to End model tests @@ -251,8 +329,41 @@ test_full_ethosu_fvp() { # All End to End model tests test_models_ethos-u55 test_models_ethos-u85 echo "${TEST_SUITE_NAME}: PASS" - } +} + +test_full_vkml() { # All End to End model tests + echo "${TEST_SUITE_NAME}: Test VGF delegate models and examples with VKML" + + test_run_vkml + test_models_vkml + echo "${TEST_SUITE_NAME}: PASS" +} +test_smaller_stories_llama() { + echo "${TEST_SUITE_NAME}: Test smaller_stories_llama" + + backends/arm/scripts/build_executorch.sh + + mkdir -p stories110M + pushd stories110M + wget -N https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt + echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json + popd + + # Get path to source directory + pytest \ + -c /dev/null \ + --verbose \ + --color=yes \ + --numprocesses=auto \ + --log-level=DEBUG \ + --junit-xml=stories110M/test-reports/unittest.xml \ + -s \ + backends/arm/test/models/test_llama.py \ + --llama_inputs stories110M/stories110M.pt stories110M/params.json stories110m + + echo "${TEST_SUITE_NAME}: PASS" +} ${TEST_SUITE} diff --git a/backends/arm/test/test_model.py b/backends/arm/test/test_model.py index 5e53da4a0ef..8833b7050e7 100755 --- a/backends/arm/test/test_model.py +++ b/backends/arm/test/test_model.py @@ -56,6 +56,13 @@ def get_args(): default=False, help="Don't save temporary files during compilation", ) + parser.add_argument( + "--no_quantize", + action="store_true", + required=False, + default=False, + help="Don't quantize model", + ) parser.add_argument( "--extra_flags", required=False, @@ -110,15 +117,6 @@ def build_libs(et_build_root: str, script_path: str): "--etdump", ] ) - run_external_cmd( - [ - "bash", - os.path.join(script_path, "build_portable_kernels.sh"), - f"--et_build_root={et_build_root}", - "--build_type=Release", - "--portable_kernels=aten::_softmax.out", - ] - ) def build_pte( @@ -129,30 +127,37 @@ def build_pte( memory_mode: str, build_output: str, no_intermediate: bool, + no_quantize: bool, ): + pte_file_ending = "pte" + command_list = [ + "python3", + "-m", + "examples.arm.aot_arm_compiler", + "--delegate", + f"--model_name={model_name}", + f"--target={target}", + f"--output={build_output}", + ] + + if "vgf" != target: + pte_file_ending = "bpte" + command_list.append("--bundleio") + command_list.append(f"--system_config={system_config}") + command_list.append(f"--memory_mode={memory_mode}") + + if not no_quantize: + command_list.append("--quantize") - intermediate = "" if not no_intermediate: - intermediate = f"--intermediate={output}" + command_list.append(f"--intermediate={output}") - run_external_cmd( - [ - "python3", - "-m", - "examples.arm.aot_arm_compiler", - "--delegate", - "--quantize", - "--bundleio", - intermediate, - f"--model_name={model_name}", - f"--target={target}", - f"--output={build_output}", - f"--system_config={system_config}", - f"--memory_mode={memory_mode}", - ] + run_external_cmd(command_list) + + pte_file = os.path.join( + output, f"{model_name}_arm_delegate_{args.target}.{pte_file_ending}" ) - pte_file = os.path.join(output, f"{model_name}_arm_delegate_{args.target}.bpte") return pte_file @@ -166,6 +171,7 @@ def build_ethosu_runtime( extra_flags: str, elf_build_path: str, ): + elf_build_path = os.path.join(elf_build_path, "cmake-out") run_external_cmd( [ "bash", @@ -183,7 +189,7 @@ def build_ethosu_runtime( ] ) - elf_file = os.path.join(elf_build_path, "cmake-out", "arm_executor_runner") + elf_file = os.path.join(elf_build_path, "arm_executor_runner") return elf_file @@ -199,6 +205,39 @@ def run_elf_with_fvp(script_path: str, elf_file: str, target: str, timeout: int) ) +def build_vkml_runtime( + et_build_root: str, + script_path: str, + extra_flags: str, + build_path: str, +): + run_external_cmd( + [ + "bash", + os.path.join(script_path, "build_executor_runner_vkml.sh"), + f"--et_build_root={et_build_root}", + "--etdump", + "--build_type=Release", + f"--extra_build_flags=-DET_DUMP_OUTPUT=OFF {extra_flags}", + f"--output={build_path}", + ] + ) + + runner = os.path.join(build_path, "executor_runner") + return runner + + +def run_vkml(script_path: str, pte_file: str, runner_build_path: str): + run_external_cmd( + [ + "bash", + os.path.join(script_path, "run_vkml.sh"), + f"--model={pte_file}", + f"--build_path={runner_build_path}", + ] + ) + + if __name__ == "__main__": total_start_time = time.perf_counter() args = get_args() @@ -232,13 +271,38 @@ def run_elf_with_fvp(script_path: str, elf_file: str, target: str, timeout: int) args.memory_mode, output, args.no_intermediate, + args.no_quantize, ) end_time = time.perf_counter() print( f"[Test model: {end_time - start_time:.2f} s] PTE file created: {pte_file}" ) - if "ethos-u" in args.target: + if "vgf" == args.target: + build_path = os.path.join( + output, f"{model_name}_arm_delegate_{args.target}" + ) + + start_time = time.perf_counter() + vkml_runner = build_vkml_runtime( + args.test_output, + script_path, + args.extra_flags, + build_path, + ) + end_time = time.perf_counter() + print( + f"[Test model: {end_time - start_time:.2f} s] ELF file created: {vkml_runner}" + ) + + start_time = time.perf_counter() + run_vkml(script_path, pte_file, build_path) + end_time = time.perf_counter() + print( + f"[Test model: {end_time - start_time:.2f} s] Tested VKML runner: {vkml_runner}" + ) + + elif "ethos-u" in args.target: elf_build_path = os.path.join( output, f"{model_name}_arm_delegate_{args.target}" ) diff --git a/backends/arm/test/tester/analyze_output_utils.py b/backends/arm/test/tester/analyze_output_utils.py index 96060b7b563..82d4f5d9837 100644 --- a/backends/arm/test/tester/analyze_output_utils.py +++ b/backends/arm/test/tester/analyze_output_utils.py @@ -10,7 +10,6 @@ from executorch.backends.arm.arm_backend import get_intermediate_path from executorch.backends.arm.test.runner_utils import ( get_input_quantization_params, - get_output_nodes, get_output_quantization_params, ) @@ -254,9 +253,9 @@ def dump_error_output( export_stage = tester.stages.get(StageType.EXPORT, None) quantize_stage = tester.stages.get(StageType.QUANTIZE, None) if export_stage is not None and quantize_stage is not None: - output_nodes = get_output_nodes(export_stage.artifact) + output_node = export_stage.artifact.graph_module.graph.output_node() qp_input = get_input_quantization_params(export_stage.artifact) - qp_output = get_output_quantization_params(output_nodes) + qp_output = get_output_quantization_params(output_node) logger.error(f"Input QuantArgs: {qp_input}") logger.error(f"Output QuantArgs: {qp_output}") diff --git a/backends/arm/test/tester/arm_tester.py b/backends/arm/test/tester/arm_tester.py index b848af2d25c..fe17bd3f448 100644 --- a/backends/arm/test/tester/arm_tester.py +++ b/backends/arm/test/tester/arm_tester.py @@ -34,7 +34,6 @@ from executorch.backends.arm.arm_backend import ( get_intermediate_path, - get_tosa_spec, is_ethosu, is_tosa, is_vgf, @@ -49,10 +48,9 @@ from executorch.backends.arm.test.runner_utils import ( dbg_tosa_fb_to_json, get_elf_path, - get_output_nodes, get_output_quantization_params, get_target_board, - run_corstone, + run_target, TosaReferenceModelDispatch, ) @@ -60,11 +58,12 @@ dump_error_output, print_error_diffs, ) -from executorch.backends.arm.tosa_mapping import extract_tensor_meta -from executorch.backends.arm.tosa_partitioner import TOSAPartitioner -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.mapping import extract_tensor_meta +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.specification import get_tosa_spec -from executorch.backends.arm.vgf_partitioner import VgfPartitioner +from executorch.backends.arm.vgf import VgfPartitioner from executorch.backends.test.harness.stages import Stage, StageType from executorch.backends.xnnpack.test.tester import Tester @@ -172,7 +171,9 @@ def dump_artifact(self, path_to_dump: Optional[str]): super().dump_artifact(path_to_dump) _dump_lowered_modules_artifact(path_to_dump, self.artifact, self.graph_module) - def run(self, artifact: ExportedProgram, inputs=None) -> None: + def run( + self, artifact: ExportedProgram, inputs=None, generate_etrecord: bool = False + ) -> None: artifact_to_run = copy.deepcopy(artifact) self.edge_dialect_program = to_edge_transform_and_lower( artifact_to_run, @@ -180,6 +181,7 @@ def run(self, artifact: ExportedProgram, inputs=None) -> None: compile_config=self.edge_compile_conf, partitioner=self.partitioners, constant_methods=self.constant_methods, + generate_etrecord=generate_etrecord, ) @@ -210,7 +212,7 @@ def run_artifact(self, inputs): f"Did not find build arm_executor_runner in path {elf_path}, run setup_testing.sh?" ) - return run_corstone( + return run_target( self.executorch_program_manager, inputs_flattened, intermediate_path, @@ -482,9 +484,8 @@ def run_method_and_compare_outputs( reference_stage = self.stages[StageType.INITIAL_MODEL] exported_program = self.stages[StageType.EXPORT].artifact - output_nodes = get_output_nodes(exported_program) - - output_qparams = get_output_quantization_params(output_nodes) + output_node = exported_program.graph_module.graph.output_node() + output_qparams = get_output_quantization_params(output_node) quantization_scales = [] for node in output_qparams: diff --git a/backends/arm/test/tester/test_pipeline.py b/backends/arm/test/tester/test_pipeline.py index fb9f05444e5..102ccd209e9 100644 --- a/backends/arm/test/tester/test_pipeline.py +++ b/backends/arm/test/tester/test_pipeline.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. import logging +import warnings as _warnings from typing import ( Any, @@ -21,6 +22,7 @@ import torch +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder from executorch.backends.arm.quantizer import ( EthosUQuantizer, get_symmetric_quantization_config, @@ -29,7 +31,10 @@ ) from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.arm_tester import ArmTester, RunPasses -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) from executorch.backends.xnnpack.test.tester.tester import Quantize from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -226,6 +231,12 @@ def find_pos(self, stage_id: str): raise Exception(f"Stage id {stage_id} not found in pipeline") + def has_stage(self, stage_id: str): + try: + return self.find_pos(stage_id) >= 0 + except: + return False + def add_stage_after(self, stage_id: str, func: Callable, *args, **kwargs): """Adds a stage after the given stage id.""" pos = self.find_pos(stage_id) + 1 @@ -271,7 +282,34 @@ def run(self): raise e -class TosaPipelineINT(BasePipelineMaker, Generic[T]): +class TOSAPipelineMaker(BasePipelineMaker, Generic[T]): + + @staticmethod + def is_tosa_ref_model_available(): + """Checks if the TOSA reference model is available.""" + # Not all deployments of ET have the TOSA reference model available. + # Make sure we don't try to use it if it's not available. + try: + import tosa_reference_model + + # Check if the module has content + return bool(dir(tosa_reference_model)) + except ImportError: + return False + + def run(self): + if ( + self.has_stage("run_method_and_compare_outputs") + and not self.is_tosa_ref_model_available() + ): + _warnings.warn( + "Warning: Skipping run_method_and_compare_outputs stage. TOSA reference model is not available." + ) + self.pop_stage("run_method_and_compare_outputs") + super().run() + + +class TosaPipelineINT(TOSAPipelineMaker, Generic[T]): """ Lowers a graph to INT TOSA spec (with quantization) and tests it with the TOSA reference model. @@ -302,18 +340,26 @@ def __init__( per_channel_quantization: bool = True, use_to_edge_transform_and_lower: bool = True, custom_path: str = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 1, dynamic_shapes: Optional[Tuple[Any]] = None, + tosa_extensions: Optional[List[str]] = None, ): + if tosa_extensions is None: + tosa_extensions = [] tosa_profiles = { - "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT"), + "1.0": TosaSpecification.create_from_string( + "TOSA-1.0+INT" + "".join([f"+{ext}" for ext in tosa_extensions]) + ), } tosa_version = conftest.get_option("tosa_version") compile_spec = common.get_tosa_compile_spec( - tosa_profiles[tosa_version], custom_path=custom_path + tosa_profiles[tosa_version], + custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, ) quantizer = TOSAQuantizer(tosa_profiles[tosa_version]) @@ -370,7 +416,7 @@ def __init__( ) -class TosaPipelineFP(BasePipelineMaker, Generic[T]): +class TosaPipelineFP(TOSAPipelineMaker, Generic[T]): """ Lowers a graph to FP TOSA spec and tests it with the TOSA reference model. @@ -399,6 +445,7 @@ def __init__( run_on_tosa_ref_model: bool = True, use_to_edge_transform_and_lower: bool = True, custom_path: str = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 0, @@ -406,14 +453,21 @@ def __init__( transform_passes: Optional[ Union[Sequence[PassType], Dict[str, Sequence[PassType]]] ] = None, + tosa_extensions: Optional[List[str]] = None, ): + if tosa_extensions is None: + tosa_extensions = [] tosa_profiles = { - "1.0": TosaSpecification.create_from_string("TOSA-1.0+FP"), + "1.0": TosaSpecification.create_from_string( + "TOSA-1.0+FP" + "".join([f"+{ext}" for ext in tosa_extensions]) + ), } tosa_version = conftest.get_option("tosa_version") compile_spec = common.get_tosa_compile_spec( - tosa_profiles[tosa_version], custom_path=custom_path + tosa_profiles[tosa_version], + custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, ) super().__init__( module, @@ -472,11 +526,15 @@ def __init__( per_channel_quantization: bool = True, use_to_edge_transform_and_lower: bool = True, custom_path: str = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 1, ): - compile_spec = common.get_u55_compile_spec(custom_path=custom_path) + compile_spec = common.get_u55_compile_spec( + custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, + ) quantizer = EthosUQuantizer(compile_spec) quantization_config = get_symmetric_quantization_config( is_per_channel=per_channel_quantization @@ -559,11 +617,15 @@ def __init__( per_channel_quantization: bool = True, use_to_edge_transform_and_lower: bool = True, custom_path: str = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 1, ): - compile_spec = common.get_u85_compile_spec(custom_path=custom_path) + compile_spec = common.get_u85_compile_spec( + custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, + ) quantizer = EthosUQuantizer(compile_spec) quantization_config = get_symmetric_quantization_config( is_per_channel=per_channel_quantization @@ -619,7 +681,7 @@ def __init__( ) -class PassPipeline(BasePipelineMaker, Generic[T]): +class PassPipeline(TOSAPipelineMaker, Generic[T]): """ Runs single passes directly on an edge_program and checks operators before/after. @@ -655,16 +717,22 @@ def __init__( pass_functions: Optional[List[Callable]] = None, passes_with_exported_program: Optional[List[Type[ExportPass]]] = None, custom_path: str = None, + tosa_extensions: Optional[List[str]] = None, ): + if tosa_extensions is None: + tosa_extensions = [] tosa_profiles = { "1.0": TosaSpecification.create_from_string( - "TOSA-1.0+" + ("INT" if quantize else "FP") + "TOSA-1.0+" + + ("INT" if quantize else "FP") + + "".join([f"+{ext}" for ext in tosa_extensions]), ), } tosa_version = conftest.get_option("tosa_version") + self.tosa_spec = tosa_profiles[tosa_version] compile_spec = common.get_tosa_compile_spec( - tosa_profiles[tosa_version], custom_path=custom_path + self.tosa_spec, custom_path=custom_path ) super().__init__( module, @@ -703,8 +771,12 @@ def __init__( self.add_stage(self.tester.check_not, ops_not_after_pass, suffix="after") self.add_stage(self.tester.run_method_and_compare_outputs) + def run(self): + with TosaLoweringContext(self.tosa_spec): + super().run() + -class TransformAnnotationPassPipeline(BasePipelineMaker, Generic[T]): +class TransformAnnotationPassPipeline(TOSAPipelineMaker, Generic[T]): """ Runs transform_for_annotation_pipeline passes directly on an exported program and checks output. @@ -721,9 +793,14 @@ def __init__( module: torch.nn.Module, test_data: T, custom_path: str = None, + tosa_extensions: Optional[List[str]] = None, ): + if tosa_extensions is None: + tosa_extensions = [] tosa_profiles = { - "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT"), + "1.0": TosaSpecification.create_from_string( + "TOSA-1.0+INT" + "".join([f"+{ext}" for ext in tosa_extensions]), + ), } tosa_version = conftest.get_option("tosa_version") @@ -755,7 +832,7 @@ def __init__( ) -class OpNotSupportedPipeline(BasePipelineMaker, Generic[T]): +class OpNotSupportedPipeline(TOSAPipelineMaker, Generic[T]): """ Runs the partitioner on a module and checks that ops are not delegated to test SupportedTOSAOperatorChecks. @@ -779,18 +856,23 @@ def __init__( custom_path: str = None, quantize: Optional[bool] = False, u55_subset: Optional[bool] = False, + tosa_extensions: Optional[List[str]] = None, ): + if tosa_extensions is None: + tosa_extensions = [] tosa_profiles = { - "1.0": "TOSA-1.0+" + ("INT" if quantize else "FP"), + "1.0": TosaSpecification.create_from_string( + "TOSA-1.0+" + + ("INT" if quantize else "FP") + + ("+u55" if u55_subset and quantize else "") + + "".join([f"+{ext}" for ext in tosa_extensions]), + ), } - tosa_version = tosa_profiles[conftest.get_option("tosa_version")] + tosa_version = conftest.get_option("tosa_version") - if u55_subset and quantize: - tosa_version = f"{tosa_version}+u55" + tosa_spec = tosa_profiles[tosa_version] - compile_spec = common.get_tosa_compile_spec( - tosa_version, custom_path=custom_path - ) + compile_spec = common.get_tosa_compile_spec(tosa_spec, custom_path=custom_path) super().__init__( module, test_data, @@ -799,7 +881,7 @@ def __init__( [], ) - if "INT" in tosa_version: + if tosa_spec.support_integer(): self.add_stage(self.tester.quantize, pos=0) self.change_args("check_not.exir", []) @@ -825,7 +907,7 @@ class VgfPipeline(BasePipelineMaker, Generic[T]): exir_ops: Exir dialect ops expected to be found in the graph after to_edge. if not using use_edge_to_transform_and_lower. - run_on_vulkan_runtime: Not yet supported. + run_on_vulkan_runtime: Set to true to test VGF output on VKML runtime. vgf_compiler_flags: Optional compiler flags. @@ -848,6 +930,7 @@ def __init__( per_channel_quantization: bool = True, use_to_edge_transform_and_lower: bool = True, custom_path: str = None, + tosa_debug_mode: Optional[ArmCompileSpecBuilder.DebugMode] = None, atol: float = 1e-03, rtol: float = 1e-03, qtol: int = 1, @@ -855,11 +938,19 @@ def __init__( transform_passes: Optional[ Union[Sequence[PassType], Dict[str, Sequence[PassType]]] ] = None, + tosa_extensions: Optional[List[str]] = None, ): - tosa_profile = TosaSpecification.create_from_string(tosa_version) + if tosa_extensions is None: + tosa_extensions = [] + tosa_spec = TosaSpecification.create_from_string( + tosa_version + "".join([f"+{ext}" for ext in tosa_extensions]) + ) compile_spec = common.get_vgf_compile_spec( - tosa_profile, compiler_flags=vgf_compiler_flags, custom_path=custom_path + tosa_spec, + compiler_flags=vgf_compiler_flags, + custom_path=custom_path, + tosa_debug_mode=tosa_debug_mode, ) super().__init__( @@ -873,7 +964,7 @@ def __init__( transform_passes=transform_passes, ) - if "INT" in tosa_version: + if tosa_spec.support_integer(): quantizer = VgfQuantizer(compile_spec) quantization_config = get_symmetric_quantization_config( is_per_channel=per_channel_quantization @@ -920,4 +1011,11 @@ def __init__( ) if run_on_vulkan_runtime: - pass + self.add_stage(self.tester.serialize) + self.add_stage( + self.tester.run_method_and_compare_outputs, + atol=atol, + rtol=rtol, + qtol=qtol, + inputs=self.test_data, + ) diff --git a/backends/arm/third-party/serialization_lib b/backends/arm/third-party/serialization_lib deleted file mode 160000 index 187af0d41fe..00000000000 --- a/backends/arm/third-party/serialization_lib +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 187af0d41fe75d08d2a7ec84c1b4d24b9b641ed2 diff --git a/backends/arm/tosa/TARGETS b/backends/arm/tosa/TARGETS new file mode 100644 index 00000000000..b1df4f37c53 --- /dev/null +++ b/backends/arm/tosa/TARGETS @@ -0,0 +1,61 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +runtime.python_library( + name = "mapping", + srcs = [ + "mapping.py", + ], + deps = [ + "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", + "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", + "//caffe2:torch", + ":specification", + ], +) +runtime.python_library( + name = "quant_utils", + srcs = [ + "quant_utils.py", + ], + deps = [ + "fbsource//third-party/pypi/numpy:numpy", + "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", + "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/serializer:serializer", + "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/tosa:tosa", + "fbsource//third-party/tosa_tools/v1.00/serialization_lib/python/tosa:tosa", + "//executorch/backends/arm:constants", + ":mapping", + "//executorch/exir/dialects:lib", + ], +) +runtime.python_library( + name = "specification", + srcs = [ + "specification.py", + ], + deps = [ + "fbsource//third-party/pypi/packaging:packaging", + "//executorch/exir/backend:compile_spec_schema", + ], +) +runtime.python_library( + name = "utils", + srcs = [ + "utils.py", + ], + deps = [ + "fbsource//third-party/tosa_tools/v0.80/serialization_lib/python/serializer:serializer", + ":quant_utils", + "//executorch/backends/arm/operators:node_visitor", + ], +) + +runtime.python_library( + name = "tosa", + srcs = [ + "__init__.py", + ], + deps = [ + ":specification", + ], +) diff --git a/backends/arm/tosa/__init__.py b/backends/arm/tosa/__init__.py new file mode 100644 index 00000000000..132d3563a43 --- /dev/null +++ b/backends/arm/tosa/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2025 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-unsafe + +from .specification import TosaSpecification + +__all__ = ["TosaSpecification"] diff --git a/backends/arm/tosa_backend.py b/backends/arm/tosa/backend.py similarity index 64% rename from backends/arm/tosa_backend.py rename to backends/arm/tosa/backend.py index 00eda8ed2df..ce2b7a27487 100644 --- a/backends/arm/tosa_backend.py +++ b/backends/arm/tosa/backend.py @@ -11,30 +11,55 @@ # JIT compiler flows. # import logging -from typing import cast, final, List - -import executorch.backends.arm.tosa_specification as tosa_specification - -from executorch.backends.arm.arm_backend import get_tosa_spec -from executorch.backends.arm.operators.node_visitor import get_node_visitors -from executorch.backends.arm._passes import ( - ArmPassManager, -) # usort: skip +from collections import deque +from itertools import count +from typing import cast, Dict, final, List, Set + +import serializer.tosa_serializer as ts # type: ignore +from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder +from executorch.backends.arm.common.debug import debug_fail, debug_tosa_dump +from executorch.backends.arm.debug.schema import DebugHook from executorch.backends.arm.process_node import ( process_call_function, process_output, process_placeholder, ) -from executorch.backends.arm.tosa_utils import dbg_fail, dbg_tosa_dump +from executorch.backends.arm.tosa.specification import get_tosa_spec from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec from torch.export.exported_program import ExportedProgram -from torch.fx import Node +from torch.fx import Graph, Node # TOSA backend debug functionality logger = logging.getLogger(__name__) +def _annotate_external_ids(ep_graph: Graph) -> Dict[str, int]: + """ + Returns dictionary: node name -> external ids + + Assign id to an output node of the model so we can trace it. + """ + node2external_id = {} + + def bfs_mark(start_nodes: List[Node], idx: int, seen: Set[Node]): + q = deque(start_nodes) + while q: + n = q.popleft() + if n in seen: + continue + seen.add(n) + node2external_id[n.name] = idx + # Walk backwards so we touch every producer + q.extend(n.all_input_nodes) + + out = next(n for n in ep_graph.nodes if n.op == "output") + seen: Set[Node] = set() + for idx, val in enumerate(out.args[0]): + bfs_mark([val], idx, seen) + return node2external_id + + def arm_get_first_delegation_tag(graph_module) -> str: """Get the first delegation tag from the graph_module or return empty string.""" for node in graph_module.graph.nodes: @@ -63,6 +88,7 @@ def preprocess( # noqa: C901 artifact_path = None output_format = "" compile_flags = [] + dump_debug_info = None for spec in compile_spec: if spec.key == "debug_artifact_path": artifact_path = spec.value.decode() @@ -70,11 +96,16 @@ def preprocess( # noqa: C901 output_format = spec.value.decode() if spec.key == "compile_flags": compile_flags.append(spec.value.decode()) + if spec.key == "dump_debug_info": + dump_debug_info = spec.value.decode() # Check that the output format is set correctly in the compile spec if output_format != "tosa": raise ValueError(f'Invalid output format {output_format}, must be "tosa"') + # Assign to every node external id + node_2_id = _annotate_external_ids(edge_program.graph) + tosa_spec = get_tosa_spec(compile_spec) if tosa_spec is None: raise ValueError( @@ -85,13 +116,6 @@ def preprocess( # noqa: C901 # Converted output for this subgraph, serializer needs path early as it emits # const data directly. Path created and data written only in debug builds. - if isinstance(tosa_spec, tosa_specification.Tosa_1_00): - import serializer.tosa_serializer as ts # type: ignore - else: - raise RuntimeError( - f"Unknown TOSA version {tosa_spec}, no pip package installed to handle serialization to that version." - ) - tosa_graph = ts.TosaSerializer(artifact_path) assert ( @@ -99,11 +123,44 @@ def preprocess( # noqa: C901 and tosa_spec.version.minor == ts.TOSA_VERSION_MINOR ), f"TOSA serializer version ({ts.TOSA_VERSION_MAJOR}.{ts.TOSA_VERSION_MINOR}) doesn't match specification {tosa_spec}" + # TODO: Fix the need to lazily import this. + from executorch.backends.arm._passes import ArmPassManager + graph_module = ArmPassManager(tosa_spec).transform_to_backend_pipeline( # type: ignore exported_program=edge_program ) - node_visitors = get_node_visitors(edge_program, tosa_spec) + debug_hook = None + if dump_debug_info is not None: + debug_hook = DebugHook(ArmCompileSpecBuilder.DebugMode[dump_debug_info]) + + # TODO: Fix the need to lazily import this. + from executorch.backends.arm.operators.node_visitor import get_node_visitors + + node_visitors = get_node_visitors(edge_program, tosa_spec, debug_hook) + + # Re-shuffle output nodes to preserve author's order + def _external_id(n: Node, node_2_id, fallback: int) -> int: + return node_2_id.get(n.name, fallback) + + out_node = next(n for n in graph_module.graph.nodes if n.op == "output") + _counter = count() + + # sort nodes by the key that is id + def _sort_key(t: Node) -> int: + return _external_id(t, node_2_id, next(_counter)) + + orig_ord = tuple(sorted(out_node.args[0], key=_sort_key)) + + current_order = tuple(out_node.args[0]) + if orig_ord != current_order: + replacement = ( + list(orig_ord) if isinstance(out_node.args[0], list) else orig_ord + ) + out_node.args = (replacement,) + graph_module.graph.lint() + graph_module.recompile() + input_count = 0 for node in graph_module.graph.nodes: node = cast(Node, node) @@ -123,17 +180,23 @@ def preprocess( # noqa: C901 # any checking of compatibility. raise RuntimeError(f"{node.name} is unsupported op {node.op}") except Exception: - dbg_fail(node, graph_module, tosa_graph, artifact_path) + debug_fail(node, graph_module, tosa_graph, artifact_path) raise if artifact_path: tag = arm_get_first_delegation_tag(graph_module) - dbg_tosa_dump( + debug_tosa_dump( tosa_graph, artifact_path, suffix="{}".format(f"_{tag}" if tag else "") + (f"_{tosa_spec}"), ) + if debug_hook is not None: + if debug_hook.mode == ArmCompileSpecBuilder.DebugMode.JSON: + json_output = debug_hook.serialize() + with open(f"{artifact_path}/debug.json", "w") as f: + f.write(json_output) + # Serialize and return the TOSA flatbuffer. binary = bytes(tosa_graph.serialize()) diff --git a/backends/arm/tosa/dialect/TARGETS b/backends/arm/tosa/dialect/TARGETS index bbb40cbd5b1..398d519fee6 100644 --- a/backends/arm/tosa/dialect/TARGETS +++ b/backends/arm/tosa/dialect/TARGETS @@ -1,6 +1,36 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( + name = "core", + srcs = [ + "lib.py", + "ops_registration.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/tosa:tosa", + "//executorch/exir/dialects:lib", + ], +) + +runtime.python_library( + name = "ops", + srcs = glob(["ops/*.py"]), + deps = [ + ":core", + "//caffe2:torch", + "//executorch/backends/arm/tosa:tosa", + ], +) + +runtime.python_library( name = "lib", - srcs = glob(["*.py"]), + srcs = ["__init__.py"], + deps = [ + ":core", + ":ops", + "//caffe2:torch", + "//executorch/backends/arm/tosa:tosa", + "//executorch/exir/dialects:lib", + ], ) diff --git a/backends/arm/tosa/dialect/__init__.py b/backends/arm/tosa/dialect/__init__.py new file mode 100644 index 00000000000..136f59beb62 --- /dev/null +++ b/backends/arm/tosa/dialect/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2025 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.arm.tosa.dialect.ops import ( # noqa F401 + rescale, + table, + transpose, +) diff --git a/backends/arm/tosa/dialect/lib.py b/backends/arm/tosa/dialect/lib.py index 3c965418c72..4a807d682dc 100644 --- a/backends/arm/tosa/dialect/lib.py +++ b/backends/arm/tosa/dialect/lib.py @@ -51,9 +51,9 @@ def not_callable(): class TosaValueError(ValueError): - def __init__(self, message="A TOSA value error occurred", *args, **kwargs): - super().__init__(message, *args, **kwargs) - self.op = kwargs.get("op", None) + def __init__(self, message="A TOSA value error occurred", *args, op=None): + super().__init__(message, *args) + self.op = op def __str__(self): base_message = super().__str__() diff --git a/backends/arm/tosa/dialect/ops/rescale.py b/backends/arm/tosa/dialect/ops/rescale.py new file mode 100644 index 00000000000..5f0cf9d15dc --- /dev/null +++ b/backends/arm/tosa/dialect/ops/rescale.py @@ -0,0 +1,51 @@ +# Copyright 2025 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_registration import register_fake_tosa_op + +from executorch.backends.arm.tosa.specification import ( + get_context_spec, + TosaSpecification, +) + + +@register_fake_tosa_op( + "RESCALE(Tensor input1, ScalarType dtype, float scale, int in_zp, int out_zp) -> Tensor", # schema + ( + TosaSpecification.create_from_string("TOSA-1.0+INT"), + ), # target TOSA specifications +) +def RESCALE( + x: torch.Tensor, dtype: torch.dtype, scale: float, in_zp: int, out_zp: int +) -> torch.Tensor: + tosa_spec = get_context_spec() + """Casts the input tensor to dtype `dtype` to produce the correct tensor meta for a _rescale op. + Additionally validates TOSA constraints of a RESCALE op. + """ + if not tosa_spec.support_integer(): + raise TosaValueError( + f"TOSA spec {tosa_spec} doesn't support integers", op="RESCALE" + ) + + if dtype not in (torch.int32, torch.int8, torch.int16): + raise NotImplementedError( + f"tosa::rescale currently only supports int32, int16 and int8, not {dtype}" + ) + if dtype in (torch.int32, torch.int16) and out_zp != 0: + raise ValueError( + f"TOSA requires output_zp to be zero when the output dtype is {dtype}." + ) + if x.dtype in (torch.int32, torch.int16) and in_zp != 0: + raise ValueError( + f"TOSA requires input_zp to be zero when the input dtype is {dtype}" + ) + if x.dtype == torch.int8 and not -128 <= in_zp <= 127: + raise ValueError(f"{in_zp=} outside valid range (-128,127) for int8.") + if dtype == torch.int8 and not -128 <= out_zp <= 127: + raise ValueError(f"{out_zp=} outside valid range (-128,127) for int8.") + + return torch.empty_like(x, dtype=dtype) diff --git a/backends/arm/tosa/dialect/ops/table.py b/backends/arm/tosa/dialect/ops/table.py new file mode 100644 index 00000000000..40f592e952b --- /dev/null +++ b/backends/arm/tosa/dialect/ops/table.py @@ -0,0 +1,53 @@ +# Copyright 2025 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_registration import register_fake_tosa_op + +from executorch.backends.arm.tosa.specification import ( + get_context_spec, + TosaSpecification, +) + + +@register_fake_tosa_op( + "TABLE(Tensor input1, Tensor table) -> Tensor", # schema + ( + TosaSpecification.create_from_string("TOSA-1.0+INT"), + ), # target TOSA specifications +) +def TABLE(a, table): + tosa_spec = get_context_spec() + # verifiy input types according to the spec + if not tosa_spec.support_integer(): + raise TosaValueError( + f"TOSA spec {tosa_spec} doesn't support integers", op="TABLE" + ) + + if a.dtype == torch.int8: + if table.shape != torch.Size((256,)): + raise TosaValueError( + f"Table of wrong size ({table.shape}!={torch.Size((256,))}", op="TABLE" + ) + if table.dtype != torch.int8: + raise TosaValueError(f"Table dtype {table.dtype} is not int8", op="TABLE") + return_dtype = torch.int8 + elif a.dtype == torch.int16: + if not tosa_spec.support_extension("int16"): + raise TosaValueError( + f"Context TOSA spec {tosa_spec} doesn't support int16", op="TABLE" + ) + if table.shape != torch.Size((513,)): + raise TosaValueError( + f"Table of wrong size ({table.shape}!={torch.Size((513,))})", op="TABLE" + ) + if table.dtype != torch.int16: + raise TosaValueError(f"Table dtype {table.dtype} is not int32", op="TABLE") + return_dtype = torch.int32 + else: + raise TosaValueError(f"Unsupported dtype {a.dtype} for {tosa_spec}", op="TABLE") + + return torch.empty_like(a, dtype=return_dtype) diff --git a/backends/arm/tosa/dialect/ops/transpose.py b/backends/arm/tosa/dialect/ops/transpose.py new file mode 100644 index 00000000000..9c5aba05394 --- /dev/null +++ b/backends/arm/tosa/dialect/ops/transpose.py @@ -0,0 +1,35 @@ +# Copyright 2025 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_registration import register_fake_tosa_op + +from executorch.backends.arm.tosa.specification import TosaSpecification + + +@register_fake_tosa_op( + "TRANSPOSE(Tensor input, int[] perms) -> Tensor", # schema + ( + TosaSpecification.create_from_string("TOSA-1.0+FP"), + TosaSpecification.create_from_string("TOSA-1.0+INT"), + ), # target TOSA specifications +) +def TRANSPOSE(a, perms): + # The TOSA TRANSPOSE only do the transpose in the TOSA serialized world, + # so just return the same shape and type. + + # For certain operators we need the data in a specific data format. Changing tosa_dim_order + # is not sufficient as we also need transpose the data. + # By utilizing an edge IR passthrough operator we can keep the edge program in + # channels-first/contiguous and get the desired behavior in the TOSA lowering. + + if len(perms) not in (4, 5): + raise TosaValueError( + f"Only 4D and 5D tensors are supported, got {len(perms)}: {perms}", + op="TRANSPOSE", + ) + + return torch.empty_like(a, dtype=a.dtype) diff --git a/backends/arm/tosa/dialect/ops_registration.py b/backends/arm/tosa/dialect/ops_registration.py index 865eca6b21b..84aa5cbc6e9 100644 --- a/backends/arm/tosa/dialect/ops_registration.py +++ b/backends/arm/tosa/dialect/ops_registration.py @@ -8,7 +8,7 @@ from executorch.backends.arm.tosa.dialect.lib import register_tosa_dialect_op -from executorch.backends.arm.tosa_specification import ( +from executorch.backends.arm.tosa.specification import ( get_context_spec, TosaSpecification, ) @@ -26,7 +26,7 @@ _registered_tosa_ops_by_func: dict[Callable, Callable] = {} -def register_tosa_op( +def register_fake_tosa_op( op_schema: str, tosa_specs: Iterable[TosaSpecification] ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ diff --git a/backends/arm/tosa_mapping.py b/backends/arm/tosa/mapping.py similarity index 90% rename from backends/arm/tosa_mapping.py rename to backends/arm/tosa/mapping.py index 6704499a0d3..60ef98a37c0 100644 --- a/backends/arm/tosa_mapping.py +++ b/backends/arm/tosa/mapping.py @@ -13,8 +13,10 @@ from typing import Any, Optional, Sequence +import serializer.tosa_serializer as ts # type: ignore + import torch -from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification +from executorch.backends.arm.tosa.specification import TosaSpecification UNSUPPORTED_DTYPES = ( torch.float64, @@ -32,10 +34,6 @@ def map_dtype(data_type: torch.dtype, tosa_spec: TosaSpecification) -> Any: if data_type in UNSUPPORTED_DTYPES: raise ValueError(f"Unsupported type: {data_type}") - if isinstance(tosa_spec, Tosa_1_00): - import serializer.tosa_serializer as ts # type: ignore - else: - raise RuntimeError(f"Unsupported tosa_spec: {tosa_spec}") dtype_map = { torch.float32: ts.DType.FP32, @@ -134,10 +132,6 @@ def __repr__(self): if self.name is not None: attrs.append(f"name={self.name!r}") if self.dtype is not None: - if isinstance(self.tosa_spec, Tosa_1_00): - import serializer.tosa_serializer as ts # type: ignore - else: - raise RuntimeError(f"Unsupported tosa_spec: {self.tosa_spec}") attrs.append(f"dtype={ts.DTypeNames[self.dtype]}") if self.shape is not None: attrs.append(f"shape={self.shape!r}") diff --git a/backends/arm/tosa_partitioner.py b/backends/arm/tosa/partitioner.py similarity index 75% rename from backends/arm/tosa_partitioner.py rename to backends/arm/tosa/partitioner.py index 8c923568265..c0f546fe50a 100644 --- a/backends/arm/tosa_partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -9,16 +9,16 @@ from typing import Callable, List, Optional, Sequence, Tuple import torch -from executorch.backends.arm.constants import DQ_OPS, Q_OPS -from executorch.backends.arm.arm_backend import ( - get_tosa_spec, - is_tosa, -) # usort: skip 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, +) +from executorch.backends.arm.constants import DQ_OPS, Q_OPS from executorch.backends.arm.operator_support.tosa_supported_operators import ( tosa_support_factory, ) -from executorch.backends.arm.tosa_backend import TOSABackend +from executorch.backends.arm.tosa.backend import TOSABackend +from executorch.backends.arm.tosa.specification import get_tosa_spec from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import ( DelegationSpec, @@ -26,20 +26,45 @@ PartitionResult, ) from executorch.exir.backend.utils import tag_constant_data, WhyNoPartitionReporter +from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import ExportedProgram from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase - logger = logging.getLogger(__name__) +def is_noop_clone(node: torch.fx.node.Node) -> bool: + return node.target == exir_ops.edge.dim_order_ops._clone_dim_order.default + + +def is_noop_alias_copy(node: torch.fx.node.Node) -> bool: + return node.target == exir_ops.edge.aten.alias_copy.default + + +def is_noop_to_dim_order_copy(node: torch.fx.node.Node) -> bool: + if node.target != exir_ops.edge.dim_order_ops._to_dim_order_copy.default: + return False + else: + return node.meta.get("dtype") == get_first_fake_tensor(node.args[0]).dtype # type: ignore[arg-type] + + +def is_noop_expand(node: torch.fx.node.Node) -> bool: + if node.target != exir_ops.edge.aten.expand_copy.default: + return False + else: + multiples = calculate_multiples(node.args) + return all(m == 1 for m in multiples) + + class TOSAPartitioner(Partitioner): def __init__( self, compile_spec: List[CompileSpec], additional_checks: Optional[Sequence[OperatorSupportBase]] = None, ) -> None: + from executorch.backends.arm.arm_backend import is_tosa + if not is_tosa(compile_spec): raise RuntimeError("compile spec is not targeting TOSA") self.delegation_spec = DelegationSpec(TOSABackend.__name__, compile_spec) @@ -50,7 +75,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: # no # subgraphs containing the nodes with the tags logger.info("TOSAPartitioner::partition") - partition_tags = {} + partition_tags: dict[str, DelegationSpec] = {} tosa_spec = get_tosa_spec(self.delegation_spec.compile_specs) @@ -66,6 +91,17 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: # no allows_single_node_partition=True, ) partition_list = capability_partitioner.propose_partitions() + + def reject_partition(reason: str, partition, tag) -> None: + for node in partition.nodes: + if "delegation_tag" in node.meta: + del node.meta["delegation_tag"] + reporter.report_reject( + node, + reason, + ) + partition_tags.pop(tag, None) + for partition in partition_list: tag = f"tag{partition.id}" @@ -112,6 +148,22 @@ def is_partitioned(node: torch.fx.Node, tag=tag) -> bool: del node.meta["delegation_tag"] break + is_noop_partition = all( + is_noop_clone(node) + or is_noop_alias_copy(node) + or is_noop_expand(node) + or is_noop_to_dim_order_copy(node) + or node.target in Q_OPS + or node.target in DQ_OPS + for node in partition.nodes + ) + if is_noop_partition: + reject_partition( + "Partition contained only ops which are removed in the TOSA lowering, leading to an empty partition.", + partition, + tag, + ) + tag_constant_data(exported_program) logger.info(f"The following nodes were rejected for {tosa_spec}:") logger.info("\n" + reporter.get_table_report()) @@ -160,6 +212,7 @@ def filter_fn(node: torch.fx.Node) -> bool: torch.ops.aten.linear.default, torch.ops.aten.eye.default, torch.ops.aten.linspace.default, + torch.ops.aten.logit.default, ] + ops_to_not_decompose_if_quant_op tosa_spec = get_tosa_spec(self.delegation_spec.compile_specs) diff --git a/backends/arm/tosa/quant_utils.py b/backends/arm/tosa/quant_utils.py new file mode 100644 index 00000000000..86e8e5bad8b --- /dev/null +++ b/backends/arm/tosa/quant_utils.py @@ -0,0 +1,445 @@ +# Copyright 2023-2025 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-unsafe + +# Utility functions for TOSA quantized lowerings + +import math + +from typing import Any, Tuple + +import serializer.tosa_serializer as ts # type: ignore + +from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, + get_output_qparams, +) + +from executorch.backends.arm.tosa.mapping import TosaArg +from torch.fx import Node +from tosa.RoundingMode import RoundingMode # type: ignore + + +def insert_rescale_ops_to_int32_maxscale( + tosa_graph: Any, inputs: list[TosaArg], node: Node, tosa_spec=None +) -> tuple[list[Any], float]: + """For ADD and SUB, we rescale to int32 using a different common scale(2*max(left scale,right scale)) + compared to all the other cases. We also multiply the left and right scales by 1<<20 giving us extra precision + for the computation without overflowing. + + Returns a list of the rescaled nodes and the scale factor used, + needed by insert_rescale_op_to_int8. + """ + + if len(inputs) > 2: + raise ValueError("More than two inputs not supported") + + tensors = inputs.copy() + # Reshape tensor according to TOSA dim order + for tensor in tensors: + dim_order = tensor.dim_order + tensor.shape = [tensor.shape[i] for i in dim_order] + + input_qparams = get_input_qparams(node) + lhs_qparams, rhs_qparams = input_qparams.values() + lhs_scale = lhs_qparams.get_scale_per_tensor() + rhs_scale = rhs_qparams.get_scale_per_tensor() + # Common scale for the two numbers + max_scale_2x = 2 * max(lhs_scale, rhs_scale) + SHIFT_INT8 = 20 + # We are adding two int8 numbers. If the zero point is non-null, the result will be in the range [-255;255], therefore we need 9 bits for the result. + # We have a 32-bit accumulator, so we can shift to the left by 20 bits and not overflow. In reality, because we divide by the 2*max(lhs_scale,rhs_scale) + # we are shifting to the left by 19. + lhs_factor = (1 << SHIFT_INT8) * lhs_scale / max_scale_2x + rhs_factor = (1 << SHIFT_INT8) * rhs_scale / max_scale_2x + rescaled_lhs = build_rescale_to_int32( + tosa_graph, + tensors[0], + lhs_qparams.get_zp_per_tensor(), + lhs_factor, + tosa_spec=tosa_spec, + ) + rescaled_rhs = build_rescale_to_int32( + tosa_graph, + tensors[1], + rhs_qparams.get_zp_per_tensor(), + rhs_factor, + tosa_spec=tosa_spec, + ) + out_qparam = get_output_qparams(node)[0] + out_scale = out_qparam.get_scale_per_tensor() + back_scale = max_scale_2x / (out_scale * (1 << SHIFT_INT8)) + + return [rescaled_lhs, rescaled_rhs], back_scale + + +def insert_rescale_ops_to_int32( + tosa_graph: Any, + inputs: list[TosaArg], + node: Node, + tosa_spec=None, +) -> tuple[list[Any], float]: + """Rescales all 'nodes' to int32, adding suitable RESCALE ops to 'tosa_graph'. + The scales are adjusted using the smallest scale of all 'nodes'. + + Returns a list of the rescaled nodes and the scale factor used, + needed by insert_rescale_op_to_int8. + + This functions is used in serialization to TOSA for target ops that are + handled by the DQ/D folding pass, which stores the quantization parameters + in the node meta dict. + """ + + from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, + ) + + tensors = inputs.copy() + + # Reshape tensor according to TOSA dim order + for tensor in tensors: + dim_order = tensor.dim_order + tensor.shape = [tensor.shape[i] for i in dim_order] + + input_qparams = get_input_qparams(node) + qargs = input_qparams.values() + + # Scale the int8 quantized input to a common scale in the integer + # domain + min_scale = min([qarg.get_scale_per_tensor() for qarg in qargs]) + scales = [qarg.get_scale_per_tensor() / min_scale for qarg in qargs] + + rescaled_nodes: list[Any] = [] + for tensor, qarg, scale in zip(tensors, qargs, scales): + rescaled_nodes.append( + build_rescale_to_int32( + tosa_graph, tensor, qarg.get_zp_per_tensor(), scale, tosa_spec=tosa_spec + ) + ) + return rescaled_nodes, min_scale + + +def insert_rescale_op_to_int8( + tosa_graph: Any, + last_tensor: TosaArg, + scale: float, + node: Node, + compute_rescale=True, + tosa_spec=None, +) -> None: + """Rescales the node back to int8, adding a suitable RESCALE op to 'tosa_graph'. + Parameters: + node: The original node that is being handled by the rescales. + last_tensor:the tosa tensor to rescale back. + scale: the scaling factor used to rescale to int32, from the function 'insert_rescale_ops_to_int32' + compute_rescale: boolean indicating whether we need to divide the output scale by the original scale. + tosa_graph: the tosa_graph to manipulate. + + This functions is used in serialization to TOSA for target ops that are + handled by the DQ/D folding pass, which stores the quantization parameters + in the node meta dict. + """ + _insert_rescale_op_to_dtype( + tosa_graph, last_tensor, scale, node, ts.DType.INT8, compute_rescale, tosa_spec + ) + + +def insert_rescale_op_to_int16( + tosa_graph: Any, + last_tensor: TosaArg, + scale: float, + node: Node, + compute_rescale=True, + tosa_spec=None, +) -> None: + """Rescales the node back to int16, adding a suitable RESCALE op to 'tosa_graph'. + Parameters: + node: The original node that is being handled by the rescales. + last_tensor:the tosa tensor to rescale back. + scale: the scaling factor used to rescale to int32, from the function 'insert_rescale_ops_to_int32' + compute_rescale: boolean indicating whether we need to divide the output scale by the original scale. + tosa_graph: the tosa_graph to manipulate. + + This functions is used in serialization to TOSA for target ops that are + handled by the DQ/D folding pass, which stores the quantization parameters + in the node meta dict. + """ + _insert_rescale_op_to_dtype( + tosa_graph, last_tensor, scale, node, ts.DType.INT16, compute_rescale, tosa_spec + ) + + +def _insert_rescale_op_to_dtype( + tosa_graph: Any, + last_tensor: TosaArg, + scale: float, + node: Node, + output_dtype: Any, + compute_rescale=True, + tosa_spec=None, +) -> None: + """Common implementation for rescaling nodes back to a specific dtype. + Parameters: + node: The original node that is being handled by the rescales. + last_tensor:the tosa tensor to rescale back. + scale: the scaling factor used to rescale to int32, from the function 'insert_rescale_ops_to_int32' + output_dtype: The target dtype (ts.DType.INT8 or ts.DType.INT16) + compute_rescale: boolean indicating whether we need to divide the output scale by the original scale. + tosa_graph: the tosa_graph to manipulate. + + This functions is used in serialization to TOSA for target ops that are + handled by the DQ/D folding pass, which stores the quantization parameters + in the node meta dict. + """ + from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_output_qparams, + ) + + output_qparams = get_output_qparams(node) + if len(output_qparams) != 1: + raise ValueError("More than one output not supported") + + qargs_out = output_qparams[0] + if compute_rescale: + output_rescale_scale = scale / qargs_out.get_scale_per_tensor() + else: + output_rescale_scale = scale + + # Rescale Back to the specified dtype + build_rescale_from_int32_to_dtype( + tosa_graph, + last_tensor, + node.name, + qargs_out.get_zp_per_tensor(), + output_rescale_scale, + output_dtype, + tosa_spec=tosa_spec, + ) + + +# TOSA uses the RESCALE operation to scale between values with differing precision. +# The RESCALE operator is defined using an integer multiply, add, and shift. +# This utility function is for calculating the multiplier and shift given a scale. +# Ref: https://www.mlplatform.org/tosa/tosa_spec.html#_precision_scaling +def compute_multiplier_and_shift( + scales: list[float], scaleWidth: int = 32 +) -> Tuple[list[int], list[int]]: + if scaleWidth == 16: + offset = 15 + elif scaleWidth == 32: + offset = 31 + else: + raise ValueError( + f"Unsupported scale width: {scaleWidth}, only 16 and 32 are valid values." + ) + + multipliers = [] + shifts = [] + for scale in scales: + mantissa, exponent = math.frexp(scale) + shift = exponent + + const_2_power_15_or_31 = 1 << offset + shifted_mantissa = round(mantissa * const_2_power_15_or_31) + + assert shifted_mantissa <= const_2_power_15_or_31 + + if shifted_mantissa == const_2_power_15_or_31: + shifted_mantissa = shifted_mantissa // 2 + shift += 1 + + # TOSA expects right shift to be positive, and embed (1 << offset) into right shift bits. + shift = offset - shift + + # INT32_MAX, 2^31 - 1 + assert shifted_mantissa <= (const_2_power_15_or_31 - 1) + + multiplier = shifted_mantissa + + if shift > 62: + multiplier = multiplier >> min(31, shift - 62) + shift = 62 + multipliers.append(multiplier) + shifts.append(shift) + return multipliers, shifts + + +# For TOSA spec v1.0 RESCALE operator requires multiplier, shifts, input_zp and output_zp to be +# const inputs. Create constant operators from the data already initialized. +def create_const_ops_for_rescale( + tosa_fb, + scale_32, + input_dtype, + node_name, + multipliers, + shifts, + input_zp, + output_zp, + output_dtype, + ts, +): + + multipliers = tosa_fb.addConst( + (len(multipliers),), + ts.DType.INT32 if scale_32 else ts.DType.INT16, + multipliers, + name=node_name + "_multipliers", + ) + shifts = tosa_fb.addConst( + (len(shifts),), ts.DType.INT8, shifts, name=node_name + "_shifts" + ) + input_zp = tosa_fb.addConst( + [1], input_dtype, input_zp, name=node_name + "_input_zp" + ) + output_zp = tosa_fb.addConst( + [1], output_dtype, output_zp, name=node_name + "_output_zp" + ) + + return [multipliers.name, shifts.name, input_zp.name, output_zp.name] + + +def build_rescale( + tosa_fb: Any, + scale: list[float], + input_node: Any, + output_name: str, + output_type: Any, + input_zp: list[int], + output_zp: list[int], + rounding_mode: RoundingMode, + per_channel=False, +): + import serializer.tosa_serializer as ts # type: ignore + import tosa.Op as TosaOp # type: ignore + + scaleWidth = 32 + is_scale32 = True + multipliers, shifts = compute_multiplier_and_shift(scale, scaleWidth) + rescale_inputs = create_const_ops_for_rescale( + tosa_fb, + is_scale32, + input_node.dtype, + output_name, + multipliers, + shifts, + input_zp, + output_zp, + output_type, + ts, + ) + attr_rescale = ts.TosaSerializerAttribute() + attr_rescale.RescaleAttribute( + scale32=is_scale32, + rounding_mode=rounding_mode, + per_channel=per_channel, + input_unsigned=False, + output_unsigned=False, + ) + + tosa_fb.addOperator( + TosaOp.Op().RESCALE, + [input_node.name, *rescale_inputs], + [output_name], + attr_rescale, + ) + + return + + +def build_rescale_to_int32( + tosa_fb: Any, + input_arg: TosaArg, + input_zp: int, + rescale_scale: float, + is_scale32: bool = True, + is_double_round: bool = False, + per_channel: bool = False, + tosa_spec=None, +) -> Any: + input_A_rescaled_to_int32 = None + + input_A_rescaled_to_int32 = tosa_fb.addIntermediate(input_arg.shape, ts.DType.INT32) + + build_rescale( + tosa_fb, + [rescale_scale], + input_arg, + input_A_rescaled_to_int32.name, + ts.DType.INT32, + [input_zp], + [0], + rounding_mode=RoundingMode.SINGLE_ROUND, + ) # type: ignore[call-arg] + + return input_A_rescaled_to_int32 + + +def build_rescale_from_int32( + tosa_fb: Any, + input_node: TosaArg, + output_name: str, + output_zp: int, + rescale_scale: float, + is_scale32: bool = True, + is_double_round: bool = False, + per_channel: bool = False, + tosa_spec=None, +) -> None: + # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs + # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale + build_rescale_from_int32_to_dtype( + tosa_fb, + input_node, + output_name, + output_zp, + rescale_scale, + ts.DType.INT8, + is_scale32, + is_double_round, + per_channel, + tosa_spec, + ) + + return + + +def build_rescale_from_int32_to_dtype( + tosa_fb: Any, + input_node: TosaArg, + output_name: str, + output_zp: int, + rescale_scale: float, + output_dtype: Any, + is_scale32: bool = True, + is_double_round: bool = False, + per_channel: bool = False, + tosa_spec=None, +) -> None: + """Common implementation for rescaling from INT32 to a specific dtype (INT8 or INT16). + + Parameters: + tosa_fb: The TOSA serializer + input_node: Input tensor (should be INT32) + output_name: Name for the output tensor + output_zp: Output zero point + rescale_scale: Rescaling factor + output_dtype: Target dtype (ts.DType.INT8 or ts.DType.INT16) + Other parameters: Standard rescale parameters + """ + # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs + # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale + build_rescale( + tosa_fb, + [rescale_scale], + input_node, + output_name=output_name, + output_type=output_dtype, + input_zp=[0], + output_zp=[output_zp], + rounding_mode=RoundingMode.SINGLE_ROUND, + ) # type: ignore[call-arg] + + return diff --git a/backends/arm/tosa_specification.py b/backends/arm/tosa/specification.py similarity index 91% rename from backends/arm/tosa_specification.py rename to backends/arm/tosa/specification.py index 5f16605aa56..92b68955cdd 100644 --- a/backends/arm/tosa_specification.py +++ b/backends/arm/tosa/specification.py @@ -15,6 +15,10 @@ import re from typing import List +from executorch.exir.backend.compile_spec_schema import ( # type: ignore[import-not-found] + CompileSpec, +) + from packaging.version import Version @@ -158,6 +162,13 @@ def support_integer(self): def support_float(self): return "FP" in self.profiles + def support_extension(self, extension: str) -> bool: + for p in self.profiles: + if extension in self.valid_extensions[p] and extension in self.extensions: + return True + + return False + class TosaLoweringContext: """ @@ -188,3 +199,10 @@ def get_context_spec() -> TosaSpecification: return TosaLoweringContext.tosa_spec_var.get() except LookupError: raise RuntimeError("Function must be executed within a TosaLoweringContext") + + +def get_tosa_spec(compile_spec: List[CompileSpec]) -> TosaSpecification: + for spec in compile_spec: + if spec.key == "tosa_spec": + return TosaSpecification.create_from_string(spec.value.decode()) + raise ValueError("Could not find TOSA version in CompileSpec") diff --git a/backends/arm/tosa_utils.py b/backends/arm/tosa/utils.py similarity index 65% rename from backends/arm/tosa_utils.py rename to backends/arm/tosa/utils.py index 1ac47ce8c03..15c8612d33f 100644 --- a/backends/arm/tosa_utils.py +++ b/backends/arm/tosa/utils.py @@ -6,8 +6,7 @@ # pyre-unsafe import logging -import os -from typing import Any, Optional +from typing import Any import numpy as np import serializer.tosa_serializer as ts # type: ignore @@ -16,11 +15,8 @@ import torch -from executorch.backends.arm.tosa_mapping import extract_tensor_meta, TosaArg - -from executorch.backends.arm.tosa_specification import Tosa_1_00, TosaSpecification -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.print_program import inspect_node +from executorch.backends.arm.tosa.mapping import extract_tensor_meta +from executorch.backends.arm.tosa.specification import TosaSpecification from torch._subclasses.fake_tensor import FakeTensor from torch.fx import Node @@ -28,84 +24,6 @@ logger = logging.getLogger(__name__) -def dbg_node(node: torch.fx.Node, graph_module: torch.fx.GraphModule): - # Debug output of node information - logger.info(get_node_debug_info(node, graph_module)) - - -def get_node_debug_info( - node: torch.fx.Node, graph_module: torch.fx.GraphModule | None = None -) -> str: - output = ( - f" {inspect_node(graph=graph_module.graph, node=node)}\n" - if graph_module - else "" - "-- NODE DEBUG INFO --\n" - f" Op is {node.op}\n" - f" Name is {node.name}\n" - f" Node target is {node.target}\n" - f" Node args is {node.args}\n" - f" Node kwargs is {node.kwargs}\n" - f" Node users is {node.users}\n" - " Node.meta = \n" - ) - for k, v in node.meta.items(): - if k == "stack_trace": - matches = v.split("\n") - output += " 'stack_trace =\n" - for m in matches: - output += f" {m}\n" - else: - output += f" '{k}' = {v}\n" - - if isinstance(v, list): - for i in v: - output += f" {i}\n" - return output - - -# Output TOSA flatbuffer and test harness file -def dbg_tosa_dump(tosa_graph: ts.TosaSerializer, path: str, suffix: str = ""): - filename = f"output{suffix}.tosa" - - logger.info(f"Emitting debug output to: {path=}, {suffix=}") - - os.makedirs(path, exist_ok=True) - - fb = tosa_graph.serialize() - js = tosa_graph.writeJson(filename) - - filepath_tosa_fb = os.path.join(path, filename) - with open(filepath_tosa_fb, "wb") as f: - f.write(fb) - assert os.path.exists(filepath_tosa_fb), "Failed to write TOSA flatbuffer" - - filepath_desc_json = os.path.join(path, f"desc{suffix}.json") - with open(filepath_desc_json, "w") as f: - f.write(js) - assert os.path.exists(filepath_desc_json), "Failed to write TOSA JSON" - - -def dbg_fail( - node, - graph_module, - tosa_graph: Optional[ts.TosaSerializer] = None, - path: Optional[str] = None, -): - logger.warning("Internal error due to poorly handled node:") - if tosa_graph is not None and path is not None: - dbg_tosa_dump(tosa_graph, path) - logger.warning(f"Debug output captured in '{path}'.") - dbg_node(node, graph_module) - - -def getNodeArgs(node: Node, tosa_spec: TosaSpecification) -> list[TosaArg]: - try: - return [TosaArg(arg, tosa_spec) for arg in node.args] - except ValueError as e: - raise ValueError(f"Failed processing args to op:\n{node}") from e - - def are_fake_tensors_broadcastable( fake_tensors: list[FakeTensor], ) -> tuple[bool, list[int]]: @@ -169,14 +87,6 @@ def broadcast_tensors( for broadcast. However this function also performs the broadcast and does not have a limit on only two input tensors. """ - - if isinstance(tosa_spec, Tosa_1_00): - import serializer.tosa_serializer as ts - - reshape_helper = build_reshape_tosa_1_0 - else: - raise ValueError(f"Unsupported TOSA spec: {tosa_spec}") - index_fake_tensors = [node.meta["val"] for node in nodes] broadcastable, common_shape = are_fake_tensors_broadcastable(index_fake_tensors) if not broadcastable: @@ -198,26 +108,25 @@ def broadcast_tensors( tens_dtype, ) - reshape_helper(tosa_fb, node.name, new_shape, reshaped.name) + build_reshape_tosa_1_0(tosa_fb, node.name, new_shape, reshaped.name) tiled = tosa_fb.addIntermediate(common_shape, tens_dtype) multipliers = [ comm if curr == 1 else 1 for comm, curr in zip(common_shape, new_shape) ] - if isinstance(tosa_spec, Tosa_1_00): - multiple_shapes = tosa_fb.addConst( - (len(multipliers),), - ts.DType.SHAPE, - multipliers, - name=f"{node.name}_multiples", - ) + multiple_shapes = tosa_fb.addConst( + (len(multipliers),), + ts.DType.SHAPE, + multipliers, + name=f"{node.name}_multiples", + ) - tosa_fb.addOperator( - ts.TosaOp.Op().TILE, - [reshaped.name, multiple_shapes.name], - [tiled.name], - None, - ) + tosa_fb.addOperator( + ts.TosaOp.Op().TILE, + [reshaped.name, multiple_shapes.name], + [tiled.name], + None, + ) broadcast_tensors.append(tiled) @@ -227,38 +136,23 @@ def broadcast_tensors( def build_reshape_tosa_1_0( tosa_graph, input_name, new_shape, output_name, shape_name_override="" ): - import serializer.tosa_serializer as ts_ # type: ignore - shape = tosa_graph.addConst( np.array(new_shape).shape, - ts_.DType.SHAPE, + ts.DType.SHAPE, np.array(new_shape), name=shape_name_override if shape_name_override else output_name + "_shape", ) - attr = ts_.TosaSerializerAttribute() + attr = ts.TosaSerializerAttribute() attr.ReshapeAttribute() tosa_graph.addOperator( - ts_.TosaOp.Op().RESHAPE, + ts.TosaOp.Op().RESHAPE, [input_name, shape.name], [output_name], attr, ) -def is_consumer_node_depthwise_conv2d(node: Node): - consumer_node = list(node.users)[0] - if consumer_node.target == exir_ops.edge.aten.convolution.default: - consumer_node_inputs = consumer_node.all_input_nodes - groups = consumer_node.args[-1] - in_channels = consumer_node_inputs[0].meta["val"].shape[1] - out_channels = consumer_node_inputs[1].meta["val"].shape[0] - if (in_channels == groups) and (out_channels % in_channels) == 0: - return True - - return False - - def tosa_shape(shape, dim_order): reordered = tuple([shape[dim] for dim in dim_order]) # Dynamic shapes in executorch are represented with torch.SymInt objects in the shapes, diff --git a/backends/arm/tosa_quant_utils.py b/backends/arm/tosa_quant_utils.py deleted file mode 100644 index d6a2d7bbe59..00000000000 --- a/backends/arm/tosa_quant_utils.py +++ /dev/null @@ -1,454 +0,0 @@ -# Copyright 2023-2025 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-unsafe - -# Utiliy functions for TOSA quantized lowerings - -import math - -from typing import Any, cast, NamedTuple, Tuple - -import executorch.backends.arm.tosa_specification as tosa_specification - -import torch.fx -import torch.fx.node -from executorch.backends.arm.constants import PER_CHANNEL_QDQ_OPS, PER_TENSOR_QDQ_OPS - -from executorch.backends.arm.tosa_mapping import TosaArg -from executorch.exir.dialects._ops import ops as exir_ops -from torch import Tensor -from torch.fx import Node -from tosa.RoundingMode import RoundingMode # type: ignore - - -def insert_rescale_ops_to_int32( - tosa_graph: Any, - inputs: list[TosaArg], - node: Node, - tosa_spec=None, -) -> tuple[list[Any], float]: - """Rescales all 'nodes' to int32, adding suitable RESCALE ops to 'tosa_graph'. - The scales are adjusted using the smallest scale of all 'nodes'. - - Returns a list of the rescaled nodes and the scale factor used, - needed by rescale_node_back_to_int8. - - This functions is used in serialization to TOSA for target ops that are - handled by the DQ/D folding pass, which stores the quantization parameters - in the node meta dict. - """ - - from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( - get_input_qparams, - ) - - tensors = inputs.copy() - - # Reshape tensor according to TOSA dim order - for tensor in tensors: - dim_order = tensor.dim_order - tensor.shape = [tensor.shape[i] for i in dim_order] - - input_qparams = get_input_qparams(node) - qargs = input_qparams.values() - - # Scale the int8 quantized input to a common scale in the integer - # domain - min_scale = min([qarg.get_scale_per_tensor() for qarg in qargs]) - scales = [qarg.get_scale_per_tensor() / min_scale for qarg in qargs] - - rescaled_nodes: list[Any] = [] - for tensor, qarg, scale in zip(tensors, qargs, scales): - rescaled_nodes.append( - build_rescale_to_int32( - tosa_graph, tensor, qarg.get_zp_per_tensor(), scale, tosa_spec=tosa_spec - ) - ) - return rescaled_nodes, min_scale - - -def insert_rescale_op_to_int8( - tosa_graph: Any, - last_tensor: TosaArg, - scale: float, - node: Node, - tosa_spec=None, -) -> None: - """Rescales the node back to int8, adding a suitable RESCALE op to 'tosa_graph'. - Parameters: - node: The original node that is being handled by the rescales. - last_tensor:the tosa tensor to rescale back. - scale: the scaling factor used to rescale to int32, from the function 'insert_rescale_op_to_int32' - tosa_graph: the tosa_graph to manipulate. - - This functions is used in serialization to TOSA for target ops that are - handled by the DQ/D folding pass, which stores the quantization parameters - in the node meta dict. - """ - from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( - get_output_qparams, - ) - - output_qparams = get_output_qparams(node) - assert len(output_qparams) == 1, "More than one output not supported" - - qargs_out = output_qparams[0] - output_rescale_scale = scale / qargs_out.get_scale_per_tensor() - - # Rescale Back to INT8 - build_rescale_from_int32( - tosa_graph, - last_tensor, - node.name, - qargs_out.get_zp_per_tensor(), - output_rescale_scale, - tosa_spec=tosa_spec, - ) - - -class QuantArgs(NamedTuple): - scale: list[float] | float - zp: list[int] | int - qmin: int - qmax: int - dtype: torch.dtype - axis: int = 0 - per_channel: bool = False - - def quantize_value(self, x: torch.Tensor | float) -> Tensor: - """Quantizes the input tensor or value to a quantized tensor. If the input is - not a tensor, it is converted to a tensor first. If self.per_channel is True, - the quantization is done per channel, otherwise it is done per tensor. - """ - if not isinstance(x, torch.Tensor): - x = torch.Tensor([x]) - x = x.to(torch.float32) - if self.per_channel: - q_op = exir_ops.edge.quantized_decomposed.quantize_per_channel.default - args = ( - x, - torch.tensor(self.scale), - torch.tensor(self.zp), - self.axis, - self.qmin, - self.qmax, - self.dtype, - ) - else: - q_op = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - args = (x, self.scale, self.zp, self.qmin, self.qmax, self.dtype) # type: ignore[assignment] - - return q_op(*args) - - def dequantize_value(self, qx: torch.Tensor) -> torch.Tensor: - """Dequantizes the input tensor or value to a dequantized tensor If the input - is not a tensor, it is converted to a tensor first. If self.per_channel is True, - the dequantization is done per channel, otherwise it is done per tensor. - """ - if self.per_channel: - dq_op = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default - args = ( - qx, - torch.tensor(self.scale), - torch.tensor(self.zp), - self.axis, - self.qmin, - self.qmax, - self.dtype, - ) - else: - dq_op = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - args = (qx, self.scale, self.zp, self.qmin, self.qmax, self.dtype) # type: ignore[assignment] - - return dq_op(*args) - - @classmethod - def from_operator(cls, op, args): - if op in PER_TENSOR_QDQ_OPS: - return cls( - scale=cast(float, args[1]), - zp=cast(int, args[2]), - qmin=cast(int, args[3]), - qmax=cast(int, args[4]), - dtype=cast(torch.dtype, args[5]), - axis=0, - per_channel=False, - ) - elif op in PER_CHANNEL_QDQ_OPS: - return cls( - scale=cast(list[float], args[1].tolist()), - zp=cast(list[int], args[2].tolist()), - axis=cast(int, args[3]), - qmin=cast(int, args[4]), - qmax=cast(int, args[5]), - dtype=cast(torch.dtype, args[6]), - per_channel=True, - ) - - else: - # We're only handling per tensor and per channel quantization - raise NotImplementedError(f"Unsupported quantization operation: {op}") - - def get_scale_per_tensor(self) -> float: - if not isinstance(self.scale, float): - raise TypeError( - f"Expected scale {self.scale} to be a float but found scale of " - f"type {type(self.scale)}" - ) - return self.scale - - def get_zp_per_tensor(self) -> int: - if not isinstance(self.zp, int): - raise TypeError( - f"Expected zero point {self.zp} to be an int but found zp of " - f"type {type(self.zp)}" - ) - return self.zp - - def get_scale_per_channel(self) -> list[float]: - if not isinstance(self.scale, list): - raise TypeError( - f"Expected scale {self.scale} to be a list but found scale of " - f"type {type(self.scale)}" - ) - return self.scale - - def get_zp_per_channel(self) -> list[int]: - if not isinstance(self.zp, list): - raise TypeError( - f"Expected zero point {self.zp} to be a list but found zp of " - f"type {type(self.zp)}" - ) - return self.zp - - -# TOSA uses the RESCALE operation to scale between values with differing precision. -# The RESCALE operator is defined using an integer multiply, add, and shift. -# This utility function is for calculating the multier and shift given a scale. -# Ref: https://www.mlplatform.org/tosa/tosa_spec.html#_precision_scaling -def compute_multiplier_and_shift( - scales: list[float], scaleWidth: int = 32 -) -> Tuple[list[int], list[int]]: - if scaleWidth == 16: - offset = 15 - elif scaleWidth == 32: - offset = 31 - else: - raise ValueError( - f"Unsupported scale width: {scaleWidth}, only 16 and 32 are valid values." - ) - - multipliers = [] - shifts = [] - for scale in scales: - mantissa, exponent = math.frexp(scale) - shift = exponent - - const_2_power_15_or_31 = 1 << offset - shifted_mantissa = round(mantissa * const_2_power_15_or_31) - - assert shifted_mantissa <= const_2_power_15_or_31 - - if shifted_mantissa == const_2_power_15_or_31: - shifted_mantissa = shifted_mantissa // 2 - shift += 1 - - # TOSA expects right shift to be positive, and embed (1 << offset) into right shift bits. - shift = offset - shift - - # INT32_MAX, 2^31 - 1 - assert shifted_mantissa <= (const_2_power_15_or_31 - 1) - - multiplier = shifted_mantissa - - if shift > 62: - multiplier = multiplier >> min(31, shift - 62) - shift = 62 - multipliers.append(multiplier) - shifts.append(shift) - return multipliers, shifts - - -# For TOSA spec v1.0 RESCALE operator requires multipler, shifts, input_zp and output_zp to be -# const inputs. Create constant operators from the data already initialized. -def create_const_ops_for_rescale( - tosa_fb, - scale_32, - input_dtype, - node_name, - multipliers, - shifts, - input_zp, - output_zp, - output_dtype, - ts, -): - - multipliers = tosa_fb.addConst( - (len(multipliers),), - ts.DType.INT32 if scale_32 else ts.DType.INT16, - multipliers, - name=node_name + "_multipliers", - ) - shifts = tosa_fb.addConst( - (len(shifts),), ts.DType.INT8, shifts, name=node_name + "_shifts" - ) - input_zp = tosa_fb.addConst( - [1], input_dtype, input_zp, name=node_name + "_input_zp" - ) - output_zp = tosa_fb.addConst( - [1], output_dtype, output_zp, name=node_name + "_output_zp" - ) - - return [multipliers.name, shifts.name, input_zp.name, output_zp.name] - - -def build_rescale( - tosa_fb: Any, - scale: list[float], - input_node: Any, - output_name: str, - output_type: Any, - input_zp: list[int], - output_zp: list[int], - rounding_mode: RoundingMode, - per_channel=False, -): - import serializer.tosa_serializer as ts # type: ignore - import tosa.Op as TosaOp # type: ignore - - scaleWidth = 32 - is_scale32 = True - multipliers, shifts = compute_multiplier_and_shift(scale, scaleWidth) - rescale_inputs = create_const_ops_for_rescale( - tosa_fb, - is_scale32, - input_node.dtype, - output_name, - multipliers, - shifts, - input_zp, - output_zp, - output_type, - ts, - ) - attr_rescale = ts.TosaSerializerAttribute() - attr_rescale.RescaleAttribute( - scale32=is_scale32, - rounding_mode=rounding_mode, - per_channel=per_channel, - input_unsigned=False, - output_unsigned=False, - ) - - tosa_fb.addOperator( - TosaOp.Op().RESCALE, - [input_node.name, *rescale_inputs], - [output_name], - attr_rescale, - ) - - return - - -def build_rescale_to_int32( - tosa_fb: Any, - input_arg: TosaArg, - input_zp: int, - rescale_scale: float, - is_scale32: bool = True, - is_double_round: bool = False, - per_channel: bool = False, - tosa_spec=None, -) -> Any: - input_A_rescaled_to_int32 = None - - if isinstance(tosa_spec, tosa_specification.Tosa_1_00): - # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs - # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale - import serializer.tosa_serializer as ts # type: ignore - - input_A_rescaled_to_int32 = tosa_fb.addIntermediate( - input_arg.shape, ts.DType.INT32 - ) - - build_rescale( - tosa_fb, - [rescale_scale], - input_arg, - input_A_rescaled_to_int32.name, - ts.DType.INT32, - [input_zp], - [0], - rounding_mode=RoundingMode.SINGLE_ROUND, - ) # type: ignore[call-arg] - - return input_A_rescaled_to_int32 - - -def build_rescale_from_int32( - tosa_fb: Any, - input_node: TosaArg, - output_name: str, - output_zp: int, - rescale_scale: float, - is_scale32: bool = True, - is_double_round: bool = False, - per_channel: bool = False, - tosa_spec=None, -) -> None: - if isinstance(tosa_spec, tosa_specification.Tosa_1_00): - import serializer.tosa_serializer as ts # type: ignore - - # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs - # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale - build_rescale( - tosa_fb, - [rescale_scale], - input_node, - output_name=output_name, - output_type=ts.DType.INT8, - input_zp=[0], - output_zp=[output_zp], - rounding_mode=RoundingMode.SINGLE_ROUND, - ) # type: ignore[call-arg] - return - - -""" Creates a TOSA rescale op based on conv2d parameters. """ - - -def build_rescale_conv_output( - tosa_fb: Any, - op: Any, - output_name: str, - output_type: Any, - input_scale: list[float], - weight_scale: list[float], - output_scale: list[float], - output_zp: list[int], - tosa_spec=None, -): - # TODO add check to verify if this is a Per-channel quantization. - post_conv2d_scale = [ - (inp * w) / out for inp, w, out in zip(input_scale, weight_scale, output_scale) - ] - - if isinstance(tosa_spec[0], tosa_specification.Tosa_1_00): - # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs - # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale - build_rescale( - tosa_fb=tosa_fb, - scale=post_conv2d_scale, - input_node=op, - output_name=output_name, - output_type=output_type, - input_zp=[0], - output_zp=output_zp, - rounding_mode=RoundingMode.SINGLE_ROUND, - per_channel=isinstance(weight_scale, torch.Tensor), - ) # type: ignore[call-arg] - return diff --git a/backends/arm/vgf/__init__.py b/backends/arm/vgf/__init__.py new file mode 100644 index 00000000000..4ab8144cbd6 --- /dev/null +++ b/backends/arm/vgf/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025 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-unsafe + +from .backend import VgfBackend # noqa: F401 +from .partitioner import VgfPartitioner # noqa: F401 + +__all__ = [ + "VgfBackend", + "VgfPartitioner", +] diff --git a/backends/arm/vgf_backend.py b/backends/arm/vgf/backend.py similarity index 98% rename from backends/arm/vgf_backend.py rename to backends/arm/vgf/backend.py index 475df41308b..7c408748529 100644 --- a/backends/arm/vgf_backend.py +++ b/backends/arm/vgf/backend.py @@ -18,7 +18,7 @@ import tempfile from typing import final, List -from executorch.backends.arm.tosa_backend import ( +from executorch.backends.arm.tosa.backend import ( arm_get_first_delegation_tag, TOSABackend, ) diff --git a/backends/arm/vgf_partitioner.py b/backends/arm/vgf/partitioner.py similarity index 89% rename from backends/arm/vgf_partitioner.py rename to backends/arm/vgf/partitioner.py index 0943c01c511..f6dab597487 100644 --- a/backends/arm/vgf_partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -10,8 +10,8 @@ from executorch.backends.arm.arm_backend import ( is_vgf, ) # usort: skip -from executorch.backends.arm.tosa_partitioner import TOSAPartitioner -from executorch.backends.arm.vgf_backend import VgfBackend +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.vgf import VgfBackend from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import DelegationSpec from torch.fx.passes.operator_support import OperatorSupportBase diff --git a/backends/cadence/CMakeLists.txt b/backends/cadence/CMakeLists.txt index d541fafe957..47183bed21d 100644 --- a/backends/cadence/CMakeLists.txt +++ b/backends/cadence/CMakeLists.txt @@ -22,8 +22,9 @@ endif() include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. - ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) @@ -38,52 +39,58 @@ if(EXECUTORCH_CADENCE_CPU_RUNNER) executorch_target_link_options_shared_lib(executorch) executorch_target_link_options_shared_lib(portable_ops_lib) - target_include_directories(executorch INTERFACE ${_common_include_directories}) + target_include_directories( + executorch INTERFACE ${_common_include_directories} + ) find_package( - gflags REQUIRED PATHS ${CMAKE_CURRENT_BINARY_DIR}/../../third-party + gflags REQUIRED PATHS ${CMAKE_CURRENT_BINARY_DIR}/../../third-party ) - add_executable(cadence_runner - ${EXECUTORCH_ROOT}/examples/devtools/example_runner/example_runner.cpp + add_executable( + cadence_runner + ${EXECUTORCH_ROOT}/examples/devtools/example_runner/example_runner.cpp ) target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) target_include_directories( - etdump INTERFACE ${CMAKE_CURRENT_BINARY_DIR}/../../devtools/include - ${EXECUTORCH_ROOT}/third-party/flatcc/include + etdump INTERFACE ${CMAKE_CURRENT_BINARY_DIR}/../../devtools/include + ${EXECUTORCH_ROOT}/third-party/flatcc/include ) target_include_directories( - cadence_runner PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} - ${_common_include_directories} + cadence_runner PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} + ${_common_include_directories} ) target_link_libraries( - cadence_runner - executorch - gflags - etdump - extension_data_loader - bundled_program - cadence_ops_lib - flatccrt + cadence_runner + executorch + gflags + etdump + extension_data_loader + bundled_program + cadence_ops_lib + flatccrt ) endif() if(EXECUTORCH_NNLIB_OPT) set(TARGET_DIR hifi) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/third-party/nnlib - ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) + add_subdirectory( + ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/third-party/nnlib + ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 + ) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/kernels) elseif(EXECUTORCH_FUSION_G3_OPT) set(TARGET_DIR fusion_g3) - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/third-party/nnlib - ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) + add_subdirectory( + ${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/third-party/nnlib + ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 + ) else() set(TARGET_DIR reference) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/kernels) endif() - add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/${TARGET_DIR}/operators) diff --git a/backends/cadence/aot/TARGETS b/backends/cadence/aot/TARGETS index 8492bb55877..0c39fd3d38e 100644 --- a/backends/cadence/aot/TARGETS +++ b/backends/cadence/aot/TARGETS @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. load("@fbcode_macros//build_defs:export_files.bzl", "export_file") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") load( "@fbsource//tools/build_defs:default_platform_defs.bzl", @@ -16,7 +16,7 @@ load("@fbcode_macros//build_defs:cpp_python_extension.bzl", "cpp_python_extensio oncall("odai_jarvis") -python_library( +runtime.python_library( name = "utils", srcs = [ "utils.py", @@ -31,7 +31,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "compiler", srcs = [ "compiler.py", @@ -53,7 +53,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "export_example", srcs = [ "export_example.py", @@ -75,7 +75,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "pass_utils", srcs = [ "pass_utils.py", @@ -90,7 +90,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "passes", srcs = [ "passes.py", @@ -101,6 +101,7 @@ python_library( ":reorder_ops", ":replace_ops", ":simplify_ops", + ":type_dispatch", ":utils", "//caffe2:torch", "//executorch/exir:pass_base", @@ -110,7 +111,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "ops_registrations", srcs = [ "ops_registrations.py", @@ -121,7 +122,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "ref_implementations", srcs = [ "ref_implementations.py", @@ -161,7 +162,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "compiler_utils", srcs = [ "compiler_utils.py", @@ -173,7 +174,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "graph_builder", srcs = [ "graph_builder.py", @@ -185,7 +186,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "program_builder", srcs = [ "program_builder.py", @@ -213,7 +214,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "fuse_ops", srcs = [ "fuse_ops.py", @@ -233,7 +234,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "simplify_ops", srcs = [ "simplify_ops.py", @@ -248,7 +249,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "remove_ops", srcs = [ "remove_ops.py", @@ -266,7 +267,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "reorder_ops", srcs = [ "reorder_ops.py", @@ -284,7 +285,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_ops", srcs = [ "replace_ops.py", @@ -305,7 +306,7 @@ python_library( "//executorch/exir/passes:spec_prop_pass", ], ) -python_library( +runtime.python_library( name = "decompose_ops", srcs = [ "decompose_ops.py", @@ -322,7 +323,38 @@ python_library( ], ) -python_library( +runtime.python_library( + name = "type_dispatch", + srcs = [ + "type_dispatch.py", + ], + typing = True, + deps = [ + "//caffe2:torch", + "//executorch/backends/cadence/aot:pass_utils", + "//executorch/exir:pass_base", + ], +) + +python_unittest( + name = "test_type_dispatch_passes", + srcs = [ + "tests/test_type_dispatch_passes.py", + ], + supports_static_listing = False, + typing = True, + deps = [ + ":ops_registrations", + ":type_dispatch", + "//caffe2:torch", + "//executorch/backends/cadence/aot:graph_builder", + "//executorch/backends/cadence/aot:pass_utils", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], +) + +runtime.python_library( name = "typing_stubs", srcs = [ "typing_stubs.py", @@ -333,7 +365,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "compiler_funcs", srcs = [ "compiler_funcs.py", @@ -487,7 +519,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "memory_planning_algo", srcs = [ "memory_planning_algo.py", @@ -503,7 +535,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory_planning", srcs = [ "memory_planning.py", @@ -521,7 +553,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory_constraints", srcs = [ "memory_constraints.py", @@ -572,3 +604,18 @@ python_unittest( "//later:lib", ], ) + +python_unittest( + name = "test_ref_implementations", + srcs = [ + "tests/test_ref_implementations.py", + ], + supports_static_listing = False, + typing = True, + deps = [ + ":typing_stubs", + "//executorch/backends/cadence/aot:ops_registrations", + "//executorch/backends/cadence/aot:ref_implementations", + "//caffe2:torch", + ] +) diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py index eaabc6589b5..6c497d5bec4 100644 --- a/backends/cadence/aot/compiler.py +++ b/backends/cadence/aot/compiler.py @@ -13,7 +13,6 @@ import executorch.backends.cadence.aot.ops_registrations # noqa import torch from executorch.backends.cadence.aot.compiler_funcs import ( - convert as convert_fn, prepare as prepare_fn, trace as trace_fn, ) @@ -42,6 +41,7 @@ from executorch.exir.program._program import to_edge from torch.export.exported_program import ExportedProgram +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e from .passes import apply_exir_ops_passes, apply_torch_ops_passes @@ -139,7 +139,7 @@ def convert_pt2( Returns a GraphModule with the converted model. """ - converted_model = convert_fn(graph_module) + converted_model = convert_pt2e(graph_module) if dump_graphs: logging.info("Graph after convert:") @@ -388,9 +388,6 @@ def quantize_and_export_to_cadence( ) -# Export the model and lower it to an EdgeProgramManager (in edge IR), and -# apply passes specific to Cadence DSP execution. Return both to print the -# differences. def export_to_executorch_gen_etrecord( model: torch.nn.Module, inputs: tuple[object, ...], @@ -402,7 +399,33 @@ def export_to_executorch_gen_etrecord( memory_config: Optional[MemoryConfig] = None, dump_graphs: bool = False, ) -> ExecutorchProgramManager: - edge_prog_manager = export_to_edge(model, inputs, dump_graphs) + ep = torch.export.export(model, inputs, strict=True) + return _lower_ep_to_cadence_gen_etrecord( + ep, + output_dir=output_dir, + opt_level=opt_level, + mem_algo=mem_algo, + alloc_graph_input=alloc_graph_input, + alloc_graph_output=alloc_graph_output, + memory_config=memory_config, + dump_graphs=dump_graphs, + ) + + +# Export the model and lower it to an EdgeProgramManager (in edge IR), and +# apply passes specific to Cadence DSP execution. Return both to print the +# differences. +def _lower_ep_to_cadence_gen_etrecord( + ep: ExportedProgram, + output_dir: Optional[str] = None, + opt_level: int = 1, + mem_algo: int = 0, + alloc_graph_input: bool = True, + alloc_graph_output: bool = True, + memory_config: Optional[MemoryConfig] = None, + dump_graphs: bool = False, +) -> ExecutorchProgramManager: + edge_prog_manager = _lower_ep_to_edge(ep, dump_graphs) cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager) # Print some information to terminal diff --git a/backends/cadence/aot/compiler_funcs.py b/backends/cadence/aot/compiler_funcs.py index 5d5523ba31d..6ff6057255c 100644 --- a/backends/cadence/aot/compiler_funcs.py +++ b/backends/cadence/aot/compiler_funcs.py @@ -11,11 +11,7 @@ import torch from torch._inductor.decomposition import remove_decompositions -from torchao.quantization.pt2e.quantize_pt2e import ( - convert_pt2e, - prepare_pt2e, - prepare_qat_pt2e, -) +from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, prepare_qat_pt2e from torchao.quantization.pt2e.quantizer import Quantizer @@ -56,8 +52,3 @@ def prepare( prepared_model = prepare_pt2e(traced_model, quantizer) return prepared_model - - -def convert(prepared_model: torch.fx.GraphModule) -> torch.fx.GraphModule: - converted_model = convert_pt2e(prepared_model) - return converted_model diff --git a/backends/cadence/aot/compiler_utils.py b/backends/cadence/aot/compiler_utils.py index cabfb120341..b55d388691f 100644 --- a/backends/cadence/aot/compiler_utils.py +++ b/backends/cadence/aot/compiler_utils.py @@ -201,13 +201,6 @@ def contains_node_with_matching_target( return any(node.target == op_target for node in nodes) -def is_quantized_tensor(x: torch.Tensor) -> bool: - """ - Return true if the tensor x is quantized - """ - return x.is_quantized - - def get_scale(x: torch.Tensor) -> torch.Tensor: """ Return the scale of a quantized tensor as a float32 tensor. diff --git a/backends/cadence/aot/decompose_ops.py b/backends/cadence/aot/decompose_ops.py index 60514c52902..7ee1bb36fef 100644 --- a/backends/cadence/aot/decompose_ops.py +++ b/backends/cadence/aot/decompose_ops.py @@ -7,9 +7,7 @@ # This file contains all the functions that decompose one op into simpler ops in the -# graph. The functions decomposing ops for models deployed with Jarvis are grouped -# together in class 'DecomposeOpsInGraph'. Some examples of functions in the class are -# 1. functions that decompose an ATen gelu op into an equivalent series of simpler ops +# graph. # pyre-strict diff --git a/backends/cadence/aot/export_example.py b/backends/cadence/aot/export_example.py index 14d100ea1f8..6af7a88fdc2 100644 --- a/backends/cadence/aot/export_example.py +++ b/backends/cadence/aot/export_example.py @@ -11,12 +11,14 @@ import logging import tempfile +import torch + from executorch.backends.cadence.aot.ops_registrations import * # noqa -from typing import Any, Tuple +from typing import Any, Optional, Tuple from executorch.backends.cadence.aot.compiler import ( + _lower_ep_to_cadence_gen_etrecord, convert_pt2, - export_to_executorch_gen_etrecord, fuse_pt2, prepare_pt2, ) @@ -38,13 +40,12 @@ def export_model( model: nn.Module, example_inputs: Tuple[Any, ...], file_name: str = "CadenceDemoModel", - run_and_compare: bool = True, - eps_error: float = 1e-1, - eps_warn: float = 1e-5, + working_dir: Optional[str] = None, ): # create work directory for outputs and model binary - working_dir = tempfile.mkdtemp(dir="/tmp") - logging.debug(f"Created work directory {working_dir}") + if working_dir is None: + working_dir = tempfile.mkdtemp(dir="/tmp") + logging.debug(f"Created work directory {working_dir}") # Instantiate the quantizer quantizer = CadenceDefaultQuantizer() @@ -66,9 +67,11 @@ def export_model( # the one used in prepare_and_convert_pt2) quantized_model = fuse_pt2(converted_model, quantizer) + ep = torch.export.export(quantized_model, example_inputs, strict=True) + # Get edge program after Cadence specific passes - exec_prog: ExecutorchProgramManager = export_to_executorch_gen_etrecord( - quantized_model, example_inputs, output_dir=working_dir + exec_prog: ExecutorchProgramManager = _lower_ep_to_cadence_gen_etrecord( + ep, output_dir=working_dir ) logging.info("Final exported graph:\n") @@ -92,13 +95,24 @@ def export_model( f"Executorch bundled program buffer saved to {file_name} is {len(buffer)} total bytes" ) - # TODO: move to test infra - if run_and_compare: - runtime.run_and_compare( - executorch_prog=exec_prog, - inputs=example_inputs, - ref_outputs=ref_outputs, - working_dir=working_dir, - eps_error=eps_error, - eps_warn=eps_warn, - ) + +def export_and_run_model( + model: nn.Module, + example_inputs: Tuple[Any, ...], + file_name: str = "CadenceDemoModel", + eps_error: float = 1e-1, + eps_warn: float = 1e-5, +): + # create work directory for outputs and model binary + working_dir = tempfile.mkdtemp(dir="/tmp") + logging.debug(f"Created work directory {working_dir}") + exec_prog = export_model(model, example_inputs, file_name, working_dir) + ref_outputs = model(*example_inputs) + runtime.run_and_compare( + executorch_prog=exec_prog, + inputs=example_inputs, + ref_outputs=ref_outputs, + working_dir=working_dir, + eps_error=eps_error, + eps_warn=eps_warn, + ) diff --git a/backends/cadence/aot/functions.yaml b/backends/cadence/aot/functions.yaml index 9dbf28f3114..196480931e0 100644 --- a/backends/cadence/aot/functions.yaml +++ b/backends/cadence/aot/functions.yaml @@ -190,10 +190,15 @@ - arg_meta: null kernel_name: impl::reference::dequantize_per_tensor_out -- func: cadence::quantized_conv.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) +- func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null - kernel_name: impl::reference::quantized_conv_out + kernel_name: impl::reference::quantized_conv_nchw_out + +- func: cadence::quantized_conv_nhwc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_out - func: cadence::quantized_layer_norm.out(Tensor input, Tensor in_scale, Tensor in_zero_point, int[] normalized_shape, Tensor weight, Tensor bias, float eps, float output_scale, int output_zero_point, *, Tensor(a!) out) -> Tensor(a!) kernels: @@ -209,6 +214,21 @@ - arg_meta: null kernel_name: impl::reference::quantized_linear_out +- func: cadence::quantized_linear.per_tensor_out(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_linear_per_tensor_out + +- func: cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_linear_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_linear_asym8uxasym8u_asym8u_per_tensor_out + - func: cadence::quantized_relu.out(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null @@ -219,15 +239,45 @@ - arg_meta: null kernel_name: impl::reference::quantized_relu_per_tensor_out +- func: cadence::quantized_relu_asym8s_asym8s.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_relu_asym8s_asym8s_per_tensor_out + +- func: cadence::quantized_relu_asym8u_asym8u.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_relu_asym8u_asym8u_per_tensor_out + +- func: cadence::quantized_add.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_add_per_tensor_out + +- func: cadence::quantized_add_asym8sxasym8s_asym8s.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_add_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_add_asym8uxasym8u_asym8u.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_add_asym8uxasym8u_asym8u_per_tensor_out + - func: cadence::quantized_matmul.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null kernel_name: impl::reference::quantized_matmul_out -- func: cadence::quantized_linear.per_tensor_out(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) +- func: cadence::quantized_matmul_asym8sxasym8s_asym8s.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null - kernel_name: impl::reference::quantized_linear_per_tensor_out + kernel_name: impl::reference::quantized_matmul_asym8sxasym8s_asym8s_out + +- func: cadence::quantized_matmul_asym8uxasym8u_asym8u.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_matmul_asym8uxasym8u_asym8u_out - func: cadence::im2row.out(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, Tensor in_zero_point, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) kernels: @@ -239,10 +289,75 @@ - arg_meta: null kernel_name: impl::reference::im2row_per_tensor_out -- func: cadence::quantized_conv.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) +- func: cadence::quantized_conv_nchw.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_per_tensor_out + +- func: cadence::quantized_conv_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_per_tensor_out + +- func: cadence::quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null - kernel_name: impl::reference::quantized_conv_per_tensor_out + kernel_name: impl::reference::quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out - func: cadence::quantized_fully_connected.out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, Tensor weight_zero_point, Tensor out_multiplier, Tensor out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) kernels: @@ -254,6 +369,16 @@ - arg_meta: null kernel_name: impl::reference::quantized_fully_connected_per_tensor_out +- func: cadence::quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: impl::reference::quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out + - func: cadence::requantize.out(Tensor input, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point, ScalarType out_dtype, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null diff --git a/backends/cadence/aot/functions_hifi.yaml b/backends/cadence/aot/functions_hifi.yaml index 04228f40be7..cf4c5a8fffb 100644 --- a/backends/cadence/aot/functions_hifi.yaml +++ b/backends/cadence/aot/functions_hifi.yaml @@ -290,10 +290,85 @@ - arg_meta: null kernel_name: cadence::impl::HiFi::dequantize_per_tensor_out -- func: cadence::quantized_conv.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!) +- func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null - kernel_name: cadence::impl::HiFi::quantized_conv_out + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_out + +- func: cadence::quantized_conv_nhwc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_out + +- func: cadence::quantized_conv_nchw.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_per_tensor_out + +- func: cadence::quantized_conv_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_per_tensor_out + +- func: cadence::quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out + +- func: cadence::quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out + +- func: cadence::quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out - func: cadence::quantized_layer_norm.out(Tensor input, Tensor in_scale, Tensor in_zero_point, int[] normalized_shape, Tensor weight, Tensor bias, float eps, float output_scale, int output_zero_point, *, Tensor(a!) out) -> Tensor(a!) kernels: @@ -314,6 +389,16 @@ - arg_meta: null kernel_name: cadence::impl::HiFi::quantized_linear_per_tensor_out +- func: cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_linear_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_linear_asym8uxasym8u_asym8u_per_tensor_out + - func: cadence::quantized_relu_per_tensor.out(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null @@ -329,17 +414,57 @@ - arg_meta: null kernel_name: cadence::impl::HiFi::quantized_relu_per_tensor_out -- func: cadence::quantized_fully_connected.out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, Tensor weight_zero_point, Tensor out_multiplier, Tensor out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) +- func: cadence::quantized_relu_asym8s_asym8s.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null - kernel_name: cadence::impl::HiFi::quantized_fully_connected_out + kernel_name: cadence::impl::HiFi::quantized_relu_asym8s_asym8s_per_tensor_out + +- func: cadence::quantized_relu_asym8u_asym8u.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_relu_asym8u_asym8u_per_tensor_out + +- func: cadence::quantized_add_asym8sxasym8s_asym8s.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_add_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_add_asym8uxasym8u_asym8u.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_add_asym8uxasym8u_asym8u_per_tensor_out - func: cadence::quantized_matmul.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null kernel_name: cadence::impl::HiFi::quantized_matmul_out +- func: cadence::quantized_matmul_asym8sxasym8s_asym8s.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_matmul_asym8sxasym8s_asym8s_out + +- func: cadence::quantized_matmul_asym8uxasym8u_asym8u.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_matmul_asym8uxasym8u_asym8u_out + +- func: cadence::quantized_fully_connected.out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, Tensor weight_zero_point, Tensor out_multiplier, Tensor out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_fully_connected_out + - func: cadence::quantized_fully_connected.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) kernels: - arg_meta: null kernel_name: cadence::impl::HiFi::quantized_fully_connected_per_tensor_out + +- func: cadence::quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out + +- func: cadence::quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!) + kernels: + - arg_meta: null + kernel_name: cadence::impl::HiFi::quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out diff --git a/backends/cadence/aot/fuse_ops.py b/backends/cadence/aot/fuse_ops.py index 16d4dbde32b..dbd19e1d3af 100644 --- a/backends/cadence/aot/fuse_ops.py +++ b/backends/cadence/aot/fuse_ops.py @@ -72,11 +72,13 @@ def fuse_mm_with_add(self, graph_module: torch.fx.GraphModule): fuse it with mm. """ graph = graph_module.graph - for node in graph.nodes: + for node in graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.mm.default + ): # We want to discover a chain of mm -> add, or mm -> view -> add. # Only proceed if the current node is an mm node, and has only one # user/successor. - if node.target != exir_ops.edge.aten.mm.default or len(node.users) != 1: + if len(node.users) != 1: continue # Our addmm implementation computes (mat1 * mat2 + bias). So the @@ -128,6 +130,7 @@ def fuse_mm_with_add(self, graph_module: torch.fx.GraphModule): mm_arg_shape is None or bias_arg_shape is None or not broadcastable(mm_arg_shape, bias_arg_shape) + or len(bias_arg_shape) > 2 ): continue diff --git a/backends/cadence/aot/memory_constraints.py b/backends/cadence/aot/memory_constraints.py index 8e784cd2779..0eaaa8987c6 100644 --- a/backends/cadence/aot/memory_constraints.py +++ b/backends/cadence/aot/memory_constraints.py @@ -654,6 +654,37 @@ def compute_slice_and_select_loc_constraints( ] +@register_cadence_pass(CadencePassAttribute(opt_level=0)) +class GenerateIdmaConstraints(PassBase): + """Generate constraints for idma ops.""" + + def __init__(self, constraint: MemConstraints) -> None: + self.constraint = constraint + + def call(self, graph_module: torch.fx.GraphModule) -> Optional[PassResult]: + for node in graph_module.graph.find_nodes( + op="call_function", target=torch.ops.cadence.idma_wait.out + ): + # This is just an alias op. + self.constraint.add_relative_placement_constraint(node.args[0], node) + + for node in graph_module.graph.find_nodes( + op="call_function", target=torch.ops.cadence.idma_load.out + ): + # TODO: set correct dtcm bank here. + mem_id = 1 + self.constraint.add_absolute_placement_constraint(node, mem_id, None) + + for node in graph_module.graph.find_nodes( + op="call_function", target=torch.ops.cadence.idma_store.out + ): + # TODO: set correct dtcm bank here. + mem_id = 1 + self.constraint.add_absolute_placement_constraint( + node.args[0], mem_id, None + ) + + # The class to generate all the constraints that will be passed on to the memory # planning algorithm. class GenerateMemConstraints: @@ -671,6 +702,7 @@ def __call__(self, graph_module: torch.fx.GraphModule) -> PassResult: constraint_gen_passes: Sequence[ConstraintsGenPass] = cast( list[ConstraintsGenPass], [ + GenerateIdmaConstraints, GenerateMemoryViewConstraints, GenerateSliceAndSelectNopConstraints, GenerateCatNopConstraints, diff --git a/backends/cadence/aot/memory_planning.py b/backends/cadence/aot/memory_planning.py index 67da42a9d3c..5ff39580a23 100644 --- a/backends/cadence/aot/memory_planning.py +++ b/backends/cadence/aot/memory_planning.py @@ -9,7 +9,7 @@ import collections import itertools import logging -from typing import Iterable, Optional, Sequence +from typing import Callable, Iterable, Optional, Sequence, TypeAlias import torch from executorch.backends.cadence.aot.memory_constraints import MemConstraints @@ -26,6 +26,8 @@ from executorch.exir import ExecutorchProgramManager from executorch.exir.memory_planning import collect_specs_from_nodes, Verifier +from executorch.exir.pass_base import PassBase +from executorch.exir.pass_manager import PassManager from executorch.exir.passes import MemoryPlanningPass from executorch.exir.tensor import TensorSpec from tabulate import tabulate @@ -116,6 +118,9 @@ def plan_spec( Greedily place the spec in the first memory that can fit it. """ for spec.mem_id in range(1, self.get_num_memories()): + if placement_constraints.is_mem_id_in_blocklist(spec, spec.mem_id): + # Skip placement for blocked memory id. + continue prev_offset, smallest_gap = 0, float("inf") for allocated_spec in state.allocated_buffers[spec.mem_id]: if not Verifier.lifetime_overlap(spec, allocated_spec): @@ -141,11 +146,11 @@ def plan_spec( ) if spec.mem_offset is None: spec.mem_offset = prev_offset - if not self.is_valid_placement(spec, placement_constraints): - spec.mem_offset = None - continue - else: - spec.mem_offset = prev_offset + + if not self.is_valid_placement(spec, placement_constraints): + # Skip placement for invalid memory id. + spec.mem_offset = None + continue state.place_spec(spec) # A data structure used for maintaining the tensor order @@ -356,6 +361,35 @@ def print_memory_planning_info( ) +class SimplifyIdmaOpsPass(PassBase): + """Replace idma_load and idma_store with idma_copy.""" + + def call(self, graph_module: torch.fx.GraphModule) -> Optional[PassResult]: + modified = False + for node in graph_module.graph.find_nodes( + op="call_function", target=torch.ops.cadence.idma_load.out + ): + modified = True + node.target = torch.ops.cadence.idma_copy.out + node.args = (node.args[0], *node.args[2:]) + + for node in graph_module.graph.find_nodes( + op="call_function", target=torch.ops.cadence.idma_store.out + ): + modified = True + node.target = torch.ops.cadence.idma_copy.out + + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + +ConstraintGenPassType: TypeAlias = Callable[ + [MemConstraints], + Callable[[torch.fx.GraphModule], Optional[PassResult]], +] + + class CadenceMemoryPlanning: def __init__( self, @@ -420,10 +454,16 @@ def run( # True. mem_planning = MemoryPlanningPass( self.algo, - allow_lifetime_and_storage_overlap=(self.opt_level >= 2), + # Always allow lifetime and storage overlap. + # At opt level 0, we need overlap for idma wait. + allow_lifetime_and_storage_overlap=True, alloc_graph_input=self.alloc_graph_input, alloc_graph_output=self.alloc_graph_output, ) mem_planning.run(graph_module, graph_signature) + graph_module = PassManager(passes=[SimplifyIdmaOpsPass()])( + graph_module + ).graph_module + return PassResult(graph_module, True) diff --git a/backends/cadence/aot/memory_planning_algo.py b/backends/cadence/aot/memory_planning_algo.py index 8193b73c9fd..672f48a55fd 100644 --- a/backends/cadence/aot/memory_planning_algo.py +++ b/backends/cadence/aot/memory_planning_algo.py @@ -204,7 +204,7 @@ def _place_memory_id_pinned_specs( for spec, c in spec_with_abs_constraint.items() if c is not None and c.pinned_memory_id == mem_id and c.offset is None } - logging.error(f"Placing specs {mem_id_pinned_specs} for {mem_id=}") + logging.debug(f"Placing specs {mem_id_pinned_specs} for {mem_id=}") with self.block_memories_except(mem_id): self.plan( @@ -220,7 +220,7 @@ def _place_memory_id_pinned_specs( if constraint is None: continue - logging.error(f"Placing spec {spec} with {constraint}") + logging.debug(f"Placing spec {spec} with {constraint}") if not state.is_placed(spec): raise MemoryError( diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index 5713861103c..35b4cbf3902 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -16,6 +16,7 @@ get_im2row_output_size, ) from executorch.exir.scalar_type import ScalarType +from torch._meta_registrations import _linalg_svd_meta from torch.library import Library, register_fake lib = Library("cadence", "DEF") @@ -56,10 +57,26 @@ lib.define( "quantized_linear.per_tensor_out(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_linear_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_linear_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_linear.per_tensor(Tensor src, Tensor weight, Tensor bias, SymInt src_zero_point, " "SymInt weight_zero_point, SymInt out_multiplier, SymInt out_shift, SymInt out_zero_point, Tensor? offset) -> Tensor" ) +lib.define( + "quantized_linear_asym8sxasym8s_asym8s.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)" +) +lib.define( + "quantized_linear_asym8uxasym8u_asym8u.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)" +) lib.define( "quantized_relu(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Y)" @@ -69,24 +86,119 @@ ) lib.define( - "quantized_conv(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, bool channel_last=False) -> (Tensor Z)" + "quantized_conv_nhwc(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)" ) lib.define( - "quantized_conv.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)" + "quantized_conv_nhwc.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)" ) lib.define( - "quantized_conv.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False) -> (Tensor Z)" + "quantized_conv_nhwc.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" ) lib.define( - "quantized_conv.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)" + "quantized_conv_nhwc.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" ) - lib.define( "quantized_matmul(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False) -> (Tensor Z)" ) lib.define( "quantized_matmul.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_matmul_asym8sxasym8s_asym8s(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False) -> (Tensor Z)" +) +lib.define( + "quantized_matmul_asym8sxasym8s_asym8s.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift) -> (Tensor Z)" +) +lib.define( + "quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor_out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, int weight_zero_point, float bias_scale, float out_scale, int out_zero_point, int out_multiplier, int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_matmul_asym8uxasym8u_asym8u(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False) -> (Tensor Z)" +) +lib.define( + "quantized_matmul_asym8uxasym8u_asym8u.out(Tensor X, int X_zero_point, Tensor Y, int Y_zero_point, Tensor? bias, int out_multiplier, int out_shift, int out_zero_point, bool transposed=False, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "convolution(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, " @@ -139,6 +251,12 @@ "int in_zero_point, bool channel_last=False) -> (Tensor out)" ) lib.define("linalg_vector_norm(Tensor X) -> (Tensor Y)") +lib.define( + "linalg_svd(Tensor A, bool full_matrices=False, bool compute_uv=True, str? driver=None) -> (Tensor U, Tensor S, Tensor Vh)" +) +lib.define( + "linalg_svd.out(Tensor A, bool full_matrices=False, bool compute_uv=True, str? driver=None, *, Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh) -> (Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh)" +) lib.define( "transposed_im2row(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, " "int[2] output_padding, Tensor in_zero_point, bool channel_last=False) -> (Tensor out)" @@ -162,6 +280,14 @@ "quantized_fully_connected.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)" ) +lib.define( + "quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)" +) +lib.define( + "quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset) -> (Tensor Z)" +) lib.define("where_Scalar(Tensor condition, float self, float other) -> (Tensor Z)") lib.define( "where_Scalar.out(Tensor condition, float self, float other, *, Tensor(a!) out) -> Tensor(a!)" @@ -178,7 +304,13 @@ # Post memory planning, we check that outputs/inputs for the load/store are in # DTCM and replace idma_load/idma_store with idma_copy. lib.define("idma_load(Tensor src, int task_num=0, int channel=0) -> Tensor") +lib.define( + "idma_load.out(Tensor src, int task_num=0, int channel=0, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define("idma_store(Tensor src, int task_num=0, int channel=0) -> Tensor") +lib.define( + "idma_store.out(Tensor src, int task_num=0, int channel=0, *, Tensor(a!) out) -> Tensor(a!)" +) # Non-blocking iDMA copy. lib.define("idma_copy(Tensor src, int task_num=0, int channel=0) -> Tensor") @@ -208,6 +340,20 @@ "quantized_relu.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, " "int out_shift, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_relu_asym8s_asym8s.per_tensor(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift) -> Tensor" +) +lib.define( + "quantized_relu_asym8s_asym8s.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, " + "int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_relu_asym8u_asym8u.per_tensor(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, int out_shift) -> Tensor" +) +lib.define( + "quantized_relu_asym8u_asym8u.per_tensor_out(Tensor X, int X_zero_point, int out_zero_point, int out_multiplier, " + "int out_shift, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_add.out(Tensor X, Tensor X_scale, Tensor X_zero_point, Tensor Y, Tensor Y_scale, " "Tensor Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" @@ -216,6 +362,22 @@ "quantized_add.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, " "int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_add_asym8sxasym8s_asym8s.per_tensor(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, " + "int Y_zero_point, float out_scale, int out_zero_point) -> Tensor" +) +lib.define( + "quantized_add_asym8sxasym8s_asym8s.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, " + "int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_add_asym8uxasym8u_asym8u.per_tensor(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, " + "int Y_zero_point, float out_scale, int out_zero_point) -> Tensor" +) +lib.define( + "quantized_add_asym8uxasym8u_asym8u.per_tensor_out(Tensor X, float X_scale, int X_zero_point, Tensor Y, float Y_scale, " + "int Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_mul.out(Tensor X, Tensor X_scale, Tensor X_zero_point, Tensor Y, Tensor Y_scale, " "Tensor Y_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" @@ -240,6 +402,14 @@ "quantized_fully_connected.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" +) +lib.define( + "quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor_out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, " + "int weight_zero_point, int out_multiplier, int out_shift, int out_zero_point, Tensor? offset, *, Tensor(a!) out) -> Tensor(a!)" +) lib.define( "quantized_embedding_byte.out(Tensor weight, Tensor weight_scales, Tensor weight_zero_points, " "Tensor indices, bool pruned_weights=False, *, Tensor(a!) out) -> Tensor(a!)" @@ -284,6 +454,12 @@ "roi_align_box_processor(Tensor rois, int output_size_h, int output_size_w, " "int sampling_ratio, bool aligned) -> (Tensor out)" ) +lib.define( + "_softmax_f32_f32(Tensor self, int dim, bool? half_to_float) -> (Tensor out)" +) +lib.define( + "_softmax_f32_f32.out(Tensor self, int dim, bool? half_to_float, *, Tensor(a!) out) -> Tensor(a!)" +) # Custom ops with aten namespace. Need to specify the lib var as FRAGMENT type as aten library is already defined aten_lib = Library("aten", "FRAGMENT") @@ -382,56 +558,796 @@ def quantized_add_per_tensor_meta( out_zero_point: int, ) -> torch.Tensor: - out_size = torch.broadcast_shapes(X.size(), Y.size()) - return X.new_empty(out_size, dtype=X.dtype) + out_size = torch.broadcast_shapes(X.size(), Y.size()) + return X.new_empty(out_size, dtype=X.dtype) + + +@register_fake("cadence::quantized_add_asym8sxasym8s_asym8s.per_tensor") +def quantized_add_asym8sxasym8s_asym8s_per_tensor_meta( + X: torch.Tensor, + X_scale: float, + X_zero_point: int, + Y: torch.Tensor, + Y_scale: float, + Y_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + out_size = torch.broadcast_shapes(X.size(), Y.size()) + return X.new_empty(out_size, dtype=X.dtype) + + +@register_fake("cadence::quantized_add_asym8uxasym8u_asym8u.per_tensor") +def quantized_add_asym8uxasym8u_asym8u_per_tensor_meta( + X: torch.Tensor, + X_scale: float, + X_zero_point: int, + Y: torch.Tensor, + Y_scale: float, + Y_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + out_size = torch.broadcast_shapes(X.size(), Y.size()) + return X.new_empty(out_size, dtype=X.dtype) + + +@register_fake("cadence::quantized_linear") +def quantized_linear_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: torch.Tensor, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, + out_zero_point: int, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_linear.per_tensor") +def quantized_linear_per_tensor_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: torch.SymInt, + weight_zero_point: torch.SymInt, + out_multiplier: torch.SymInt, + out_shift: torch.SymInt, + out_zero_point: torch.SymInt, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_linear_asym8sxasym8s_asym8s.per_tensor") +def quantized_linear_asym8sxasym8s_asym8s_per_tensor_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: int, + out_multiplier: int, + out_shift: int, + out_zero_point: int, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_linear_asym8uxasym8u_asym8u.per_tensor") +def quantized_linear_asym8uxasym8u_asym8u_per_tensor_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: int, + out_multiplier: int, + out_shift: int, + out_zero_point: int, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_conv_nhwc") +def quantized_conv_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw") +def quantized_conv_nchw_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: torch.Tensor, + bias_scale: torch.Tensor, + output_scale: float, + output_zero_point: int, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, +) -> torch.Tensor: + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw.per_tensor") +def quantized_conv_nchw_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nhwc.per_tensor") +def quantized_conv_nhwc_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + True, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, True + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) + + +@register_fake("cadence::quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 + + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) -@register_fake("cadence::quantized_linear") -def quantized_linear_meta( - src: torch.Tensor, +@register_fake("cadence::quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_meta( + input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, + stride: Tuple[int], + padding: Tuple[int], + dilation: Tuple[int], + groups: int, in_zero_point: int, - weight_zero_point: torch.Tensor, - out_multiplier: torch.Tensor, - out_shift: torch.Tensor, - out_zero_point: int, - offset: Optional[torch.Tensor], + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, ) -> torch.Tensor: - # src comes in shape [leading_dims, in_dim] - # weight comes in shape [out_dim, in_dim] - # output comes in empty with shape [leading_dims, out_dim] - out_size = list(src.size()) - weight_size = list(weight.size()) - assert len(weight_size) == 2 - out_size[-1] = weight_size[0] - return src.new_empty(out_size, dtype=src.dtype) + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, _, *kernel_size = weight.shape + in_size = input.shape + # Assert that the input tensor has at least 3 dimensions, and at most 6 + assert len(in_size) > 2 + assert len(in_size) < 6 -@register_fake("cadence::quantized_linear.per_tensor") -def quantized_linear_per_tensor_meta( - src: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - in_zero_point: torch.SymInt, - weight_zero_point: torch.SymInt, - out_multiplier: torch.SymInt, - out_shift: torch.SymInt, - out_zero_point: torch.SymInt, - offset: Optional[torch.Tensor], -) -> torch.Tensor: - # src comes in shape [leading_dims, in_dim] - # weight comes in shape [out_dim, in_dim] - # output comes in empty with shape [leading_dims, out_dim] - out_size = list(src.size()) - weight_size = list(weight.size()) - assert len(weight_size) == 2 - out_size[-1] = weight_size[0] - return src.new_empty(out_size, dtype=src.dtype) + # Compute the output tensor size + output_size = ( + get_conv1d_output_size( + in_size, + out_channels, + stride[1], + padding[1], + dilation[1], + kernel_size[0], + False, + ) + if len(in_size) == 3 + else get_conv2d_output_size( + in_size, out_channels, stride, padding, dilation, kernel_size, False + ) + ) + + return input.new_empty(output_size, dtype=input.dtype) -@register_fake("cadence::quantized_conv") -def quantized_conv_meta( +@register_fake("cadence::quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor") +def quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_meta( input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, @@ -440,18 +1356,19 @@ def quantized_conv_meta( dilation: Tuple[int], groups: int, in_zero_point: int, - weight_zero_point: torch.Tensor, - bias_scale: torch.Tensor, + weight_zero_point: int, + bias_scale: float, output_scale: float, output_zero_point: int, - out_multiplier: torch.Tensor, - out_shift: torch.Tensor, - channel_last: bool = False, + out_multiplier: int, + out_shift: int, ) -> torch.Tensor: - if channel_last: - out_channels, *kernel_size, _ = weight.shape - else: - out_channels, _, *kernel_size = weight.shape + assert ( + input.dtype == torch.int8 + and weight.dtype == torch.int8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape in_size = input.shape # Assert that the input tensor has at least 3 dimensions, and at most 6 @@ -467,19 +1384,19 @@ def quantized_conv_meta( padding[1], dilation[1], kernel_size[0], - channel_last, + True, ) if len(in_size) == 3 else get_conv2d_output_size( - in_size, out_channels, stride, padding, dilation, kernel_size, channel_last + in_size, out_channels, stride, padding, dilation, kernel_size, True ) ) return input.new_empty(output_size, dtype=input.dtype) -@register_fake("cadence::quantized_conv.per_tensor") -def quantized_conv_per_tensor_meta( +@register_fake("cadence::quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor") +def quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_meta( input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, @@ -494,12 +1411,13 @@ def quantized_conv_per_tensor_meta( output_zero_point: int, out_multiplier: int, out_shift: int, - channel_last: bool = False, ) -> torch.Tensor: - if channel_last: - out_channels, *kernel_size, _ = weight.shape - else: - out_channels, _, *kernel_size = weight.shape + assert ( + input.dtype == torch.uint8 + and weight.dtype == torch.uint8 + and bias.dtype == torch.int32 + ) + out_channels, *kernel_size, _ = weight.shape in_size = input.shape # Assert that the input tensor has at least 3 dimensions, and at most 6 @@ -515,11 +1433,11 @@ def quantized_conv_per_tensor_meta( padding[1], dilation[1], kernel_size[0], - channel_last, + True, ) if len(in_size) == 3 else get_conv2d_output_size( - in_size, out_channels, stride, padding, dilation, kernel_size, channel_last + in_size, out_channels, stride, padding, dilation, kernel_size, True ) ) @@ -531,7 +1449,7 @@ def quantized_layer_norm_meta( input: torch.Tensor, X_scale: torch.Tensor, X_zero_point: torch.Tensor, - normalized_shape: int, + normalized_shape: list[int], weight: torch.Tensor, bias: torch.Tensor, eps: float, @@ -546,7 +1464,7 @@ def quantized_layer_norm_per_tensor_meta( input: torch.Tensor, X_scale: float, X_zero_point: int, - normalized_shape: int, + normalized_shape: list[int], weight: torch.Tensor, bias: torch.Tensor, eps: float, @@ -610,6 +1528,92 @@ def quantized_matmul_meta( return X.new_empty(out_size, dtype=X.dtype) +@register_fake("cadence::quantized_matmul_asym8sxasym8s_asym8s") +def quantized_matmul_asym8sxasym8s_asym8s_meta( + X: torch.Tensor, + X_zero_point: int, + Y: torch.Tensor, + Y_zero_point: int, + bias: Optional[torch.Tensor], + out_multiplier: int, + out_shift: int, + out_zero_point: int, + transposed: bool = False, +) -> torch.Tensor: + X_size = list(X.size()) + Y_size = list(Y.size()) + + # Get the batch dimensions for both tensors + X_batch_dims = X_size[:-2] + Y_batch_dims = Y_size[:-2] + + # If they don't match, check that they're compatible + if X_batch_dims != Y_batch_dims: + assert prod(X_batch_dims) == prod( + Y_batch_dims + ), f"Batch dimensions of X and Y do not match: {X_batch_dims} vs {Y_batch_dims}" + + # Get the matmul output size + if transposed: + assert X_size[-1] == Y_size[-1], "matrices cannot be multiplied" + mat_size = [X_size[-2], Y_size[-2]] + else: + assert X_size[-1] == Y_size[-2], "matrices cannot be multiplied" + mat_size = [X_size[-2], Y_size[-1]] + + # Combine the larger batch dimensions with the matmul output size + out_size = ( + X_batch_dims + mat_size + if len(X_batch_dims) > len(Y_batch_dims) + else Y_batch_dims + mat_size + ) + + return X.new_empty(out_size, dtype=X.dtype) + + +@register_fake("cadence::quantized_matmul_asym8uxasym8u_asym8u") +def quantized_matmul_asym8uxasym8u_asym8u_meta( + X: torch.Tensor, + X_zero_point: int, + Y: torch.Tensor, + Y_zero_point: int, + bias: Optional[torch.Tensor], + out_multiplier: int, + out_shift: int, + out_zero_point: int, + transposed: bool = False, +) -> torch.Tensor: + X_size = list(X.size()) + Y_size = list(Y.size()) + + # Get the batch dimensions for both tensors + X_batch_dims = X_size[:-2] + Y_batch_dims = Y_size[:-2] + + # If they don't match, check that they're compatible + if X_batch_dims != Y_batch_dims: + assert prod(X_batch_dims) == prod( + Y_batch_dims + ), f"Batch dimensions of X and Y do not match: {X_batch_dims} vs {Y_batch_dims}" + + # Get the matmul output size + if transposed: + assert X_size[-1] == Y_size[-1], "matrices cannot be multiplied" + mat_size = [X_size[-2], Y_size[-2]] + else: + assert X_size[-1] == Y_size[-2], "matrices cannot be multiplied" + mat_size = [X_size[-2], Y_size[-1]] + + # Combine the larger batch dimensions with the matmul output size + out_size = ( + X_batch_dims + mat_size + if len(X_batch_dims) > len(Y_batch_dims) + else Y_batch_dims + mat_size + ) + + return X.new_empty(out_size, dtype=X.dtype) + + @register_fake("cadence::im2row") def im2row_meta( input: torch.Tensor, @@ -651,6 +1655,26 @@ def linalg_vector_norm_meta( return X.new_empty([], dtype=X.dtype) +@register_fake("cadence::linalg_svd") +def linalg_svd_meta( + A: torch.Tensor, + full_matrices: bool = False, + compute_uv: bool = True, + driver: Optional[str] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # Based on the _linalg_svd meta implementation, but ensuring contiguous strides + + # Get the shapes from the original meta function + U, S, Vh = _linalg_svd_meta(A, full_matrices, compute_uv, driver) + + # Create new tensors with contiguous strides to fix the non-contiguous issue + U_contiguous = A.new_empty(U.shape, dtype=A.dtype).contiguous() + S_contiguous = A.new_empty(S.shape, dtype=A.dtype).contiguous() + Vh_contiguous = A.new_empty(Vh.shape, dtype=A.dtype).contiguous() + + return U_contiguous, S_contiguous, Vh_contiguous + + @register_fake("cadence::requantize") def requantize_meta( input: torch.Tensor, @@ -694,6 +1718,28 @@ def quantized_relu_per_tensor_meta( return input.new_empty(input.size(), dtype=input.dtype) +@register_fake("cadence::quantized_relu_asym8s_asym8s.per_tensor") +def quantized_relu_asym8s_asym8s_per_tensor_meta( + input: torch.Tensor, + in_zero_point: int, + out_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + return input.new_empty(input.size(), dtype=input.dtype) + + +@register_fake("cadence::quantized_relu_asym8u_asym8u.per_tensor") +def quantized_relu_asym8u_asym8u_per_tensor_meta( + input: torch.Tensor, + in_zero_point: int, + out_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + return input.new_empty(input.size(), dtype=input.dtype) + + @register_fake("cadence::fully_connected") def fully_connected_meta( src: torch.Tensor, @@ -725,6 +1771,7 @@ def quantized_fully_connected_meta( # src comes in shape [leading_dims, in_dim] # weight comes in shape [out_dim, in_dim] # output comes in empty with shape [leading_dims, out_dim] + assert src.shape[0] == 1 out_size = list(src.size()) weight_size = list(weight.size()) assert len(weight_size) == 2 @@ -747,6 +1794,53 @@ def quantized_fully_connected_per_tensor_meta( # src comes in shape [leading_dims, in_dim] # weight comes in shape [out_dim, in_dim] # output comes in empty with shape [leading_dims, out_dim] + assert src.shape[0] == 1 + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor") +def quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: int, + out_multiplier: int, + out_shift: int, + out_zero_point: int, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + assert src.shape[0] == 1 + out_size = list(src.size()) + weight_size = list(weight.size()) + assert len(weight_size) == 2 + out_size[-1] = weight_size[0] + return src.new_empty(out_size, dtype=src.dtype) + + +@register_fake("cadence::quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor") +def quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_meta( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: int, + out_multiplier: int, + out_shift: int, + out_zero_point: int, + offset: Optional[torch.Tensor], +) -> torch.Tensor: + # src comes in shape [leading_dims, in_dim] + # weight comes in shape [out_dim, in_dim] + # output comes in empty with shape [leading_dims, out_dim] + assert src.shape[0] == 1 out_size = list(src.size()) weight_size = list(weight.size()) assert len(weight_size) == 2 @@ -808,7 +1902,7 @@ def transposed_convolution_meta( ) -> torch.Tensor: # The native definition of torch transposed conv will have weight shape as # (in_channels, out_channels/groups, *kernel_size). - # However, the two channel position is flipped in the Jarvis pass of replacing it + # However, the two channel position is flipped in the Cadence pass of replacing it # with cadence::transposed_convolution here: https://fburl.com/code/d2s7pkyy out_channels, _input_channels, *kernel_size = weight.shape out_channels *= groups @@ -1057,3 +2151,13 @@ def roi_align_box_processor_meta( aligned: bool, ) -> torch.Tensor: return rois.new_empty((rois.shape[0], 80), dtype=torch.uint8) + + +@register_fake("cadence::_softmax_f32_f32") +def softmax_f32_f32_meta( + self: torch.Tensor, + dim: int, + dtype: torch.dtype, + half_to_float: Optional[bool] = None, +) -> torch.Tensor: + return self.new_empty(self.size(), dtype=self.dtype) diff --git a/backends/cadence/aot/passes.py b/backends/cadence/aot/passes.py index d7c692f12e9..bb4a8f065d5 100644 --- a/backends/cadence/aot/passes.py +++ b/backends/cadence/aot/passes.py @@ -33,6 +33,7 @@ ReplaceMulTensorWithMulAndFullOpsPass, ) from executorch.backends.cadence.aot.simplify_ops import CadenceSimplifyOpsInGraph +from executorch.backends.cadence.aot.type_dispatch import CompileTimeTypeDispatchPass from executorch.exir import EdgeProgramManager from executorch.exir.pass_base import ExportPass, PassResult from executorch.exir.pass_manager import PassManager, PassType @@ -90,6 +91,7 @@ def get_passes_in_default_order() -> list[Type[ExportPass]]: FuseFullThenReshapePass, FuseTransposeOrPermuteOpPairsPass, RemoveNopSliceOrViewOpPass, + CompileTimeTypeDispatchPass, ] return pytree.tree_flatten(passes)[0] diff --git a/backends/cadence/aot/program_builder.py b/backends/cadence/aot/program_builder.py index 0bb71c95a4a..862ba4e977c 100644 --- a/backends/cadence/aot/program_builder.py +++ b/backends/cadence/aot/program_builder.py @@ -2,15 +2,18 @@ # pyre-strict +from enum import auto, Enum from typing import Optional from executorch.backends.cadence.aot.graph_builder import GraphBuilder from executorch.exir import EdgeCompileConfig, EdgeProgramManager from executorch.exir.pass_base import ProxyValue from executorch.exir.verification.verifier import EXIREdgeDialectVerifier - from torch import Tensor +from torch._export.verifier import Verifier +from torch._ops import OpOverload from torch.export import ExportedProgram +from torch.export.exported_program import ModuleCallEntry, ModuleCallSignature from torch.export.graph_signature import ( ExportGraphSignature, InputKind, @@ -19,16 +22,30 @@ OutputSpec, TensorArgument, ) +from torch.utils import _pytree as pytree + + +class IrMode(Enum): + EXIR = auto() + ATEN = auto() class ProgramBuilder(GraphBuilder): """Utility class to build a program from a graph module.""" - def __init__(self) -> None: + def __init__( + self, + mode: Optional[IrMode] = None, + _core_aten_ops_exception_list: Optional[list[OpOverload]] = None, + ) -> None: self.input_specs: list[InputSpec] = [] self.output_specs: list[OutputSpec] = [] self.constants: dict[str, Tensor] = {} self.state_dict: dict[str, Tensor] = {} + self.mode: IrMode = mode or IrMode.EXIR + self._core_aten_ops_exception_list: list[OpOverload] = ( + _core_aten_ops_exception_list or [] + ) super().__init__() def insert_input_spec( @@ -68,26 +85,47 @@ def output( ) return super().output(results) + def get_verifiers(self) -> Optional[list[Verifier]]: + if self.mode == IrMode.ATEN: + return None + return [ + EXIREdgeDialectVerifier( + edge_compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=self._core_aten_ops_exception_list, + ), + core_aten_ops_exception_list=self._core_aten_ops_exception_list, + class_only=True, + ) + ] + def get_program(self) -> ExportedProgram: gm = self.get_graph_module() + graph_signature = ExportGraphSignature(self.input_specs, self.output_specs) + in_spec = pytree.tree_flatten((tuple(graph_signature.user_inputs), {}))[1] + out_spec = pytree.tree_flatten(graph_signature.user_outputs)[1] return ExportedProgram( root=gm, graph=gm.graph, - graph_signature=ExportGraphSignature( - input_specs=self.input_specs, output_specs=self.output_specs - ), + graph_signature=graph_signature, # pyre-ignore[6]: Incompatible parameter type. constants=self.constants, state_dict=self.state_dict, range_constraints={}, - module_call_graph=[], - verifiers=[ - EXIREdgeDialectVerifier( - edge_compile_config=EdgeCompileConfig(_check_ir_validity=False), - class_only=True, + module_call_graph=[ + ModuleCallEntry( + "", + ModuleCallSignature( + inputs=[], outputs=[], in_spec=in_spec, out_spec=out_spec + ), ) ], + # pyre-ignore[6]: Incompatible parameter type. + verifiers=self.get_verifiers(), ) def get_edge_program(self) -> EdgeProgramManager: - return EdgeProgramManager(self.get_program()) + return EdgeProgramManager( + self.get_program(), + core_aten_ops_exception_list=self._core_aten_ops_exception_list, + ) diff --git a/backends/cadence/aot/quantizer/TARGETS b/backends/cadence/aot/quantizer/TARGETS index c612a2e19fb..d13a69bc79f 100644 --- a/backends/cadence/aot/quantizer/TARGETS +++ b/backends/cadence/aot/quantizer/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("odai_jarvis") -python_library( +runtime.python_library( name = "utils", srcs = [ "utils.py", @@ -13,7 +13,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "patterns", srcs = [ "patterns.py", @@ -25,7 +25,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "quantizer", srcs = [ "quantizer.py", @@ -38,7 +38,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "fusion_pass", srcs = [ "fusion_pass.py", diff --git a/backends/cadence/aot/quantizer/fusion_pass.py b/backends/cadence/aot/quantizer/fusion_pass.py index a726f6c7fba..729056ea2c8 100644 --- a/backends/cadence/aot/quantizer/fusion_pass.py +++ b/backends/cadence/aot/quantizer/fusion_pass.py @@ -331,7 +331,6 @@ def get_args_and_kwargs_conv( "out_zero_point": quant_node.args[2], "out_multiplier": out_multiplier_, "out_shift": out_shift_, - "channel_last": False, } return args, kwargs diff --git a/backends/cadence/aot/quantizer/patterns.py b/backends/cadence/aot/quantizer/patterns.py index 88c16139733..74987f8b38d 100644 --- a/backends/cadence/aot/quantizer/patterns.py +++ b/backends/cadence/aot/quantizer/patterns.py @@ -247,7 +247,7 @@ def get_anchors( ) def replacement_op(self) -> OpOverload: - return torch.ops.cadence.quantized_conv.default + return torch.ops.cadence.quantized_conv_nchw.default class Conv2dPattern(QuantizationPattern): @@ -286,7 +286,7 @@ def get_anchors( ) def replacement_op(self) -> OpOverload: - return torch.ops.cadence.quantized_conv.default + return torch.ops.cadence.quantized_conv_nchw.default class LayerNormPattern(QuantizationPattern): diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index 9eaac004bcf..40ae6d23085 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -6,7 +6,11 @@ # pyre-strict + +from typing import Callable + import torch + from executorch.exir.scalar_type import ScalarType from torch.library import impl, Library @@ -20,6 +24,836 @@ } +@impl(m, "quantize_per_tensor") +def quantize_per_tensor( + input_tensor: torch.Tensor, + scale: float, + zero_point: int, + quant_min: int, + quant_max: int, + dtype: torch.dtype, +) -> torch.Tensor: + """ + Quantizes a floating-point tensor to an integral tensor. + + Args: + - input_tensor (Tensor): input tensor + - scale (float): Inverse of quantization scale. Derived from the ratio + between the min/max of the floating-point tensor and the + min/max of the quantized range, and then inverted. + - zero_point (int): The point which represents 0 in the quantized + range. For example, consider the floating point range [-1., 2.] and + quantized integer range [-7, 7]. In this case, 0 is 1/3 of way from + -1. to 2. So, the point that represents 0 in the quantized range should + be 1/3 of the way from [-7, 7]. This ends up being -2 in the integer space. + - quant_min (int): The smallest value in the quantized domain. Unused since scale + is already provided. + - quant_max (int): The largest value in the quantized domain. Unused since scale + is already provided. + - dtype (torch.dtype): The type of the output tensor + """ + supported_quant_types = [ + torch.int8, + torch.int16, + torch.int32, + torch.uint8, + torch.uint16, + ] + if dtype not in supported_quant_types: + raise ValueError( + f"Unsupported dtype to quantize to. Supported dtypes must be one of {supported_quant_types}" + ) + + quantized = torch.round(input_tensor * scale + zero_point).to(dtype) + return torch.max( + torch.min(quantized, torch.tensor(quant_max)), + torch.tensor(quant_min), + ) + + +@impl(m, "dequantize_per_tensor") +def dequantize_per_tensor( + input_tensor: torch.Tensor, + scale: float, + zero_point: int, + quant_min: int, + quant_max: int, + dtype: torch.dtype, +) -> torch.Tensor: + """ + Dequantizes an integral tensor to a floating-point tensor. + + Args: + - input (Tensor): input tensor + - scale (float): Quantization scale. Derived from the ratio + between the min/max of the floating-point tensor and the + min/max of the quantized range. + - zero_point (int): The point which represents 0 in the quantized + range. For example, consider the floating point range [-1., 2.] and + quantized integer range [-7, 7]. In this case, 0 is 1/3 of way from + -1. to 2. So, the point that represents 0 in the quantized range should + be 1/3 of the way from [-7, 7]. This ends up being -2 in the integer space. + - quant_min (int): The smallest value in the quantized domain. Unused since scale + is already provided. + - quant_max (int): The largest value in the quantized domain. Unused since scale + is already provided. + - dtype (torch.dtype): The type of the output tensor. Must be a floating point type. + """ + supported_quant_types = [ + torch.int8, + torch.int16, + torch.int32, + torch.uint8, + torch.uint16, + ] + if input_tensor.dtype not in supported_quant_types: + raise ValueError(f"Input dtype must be one of {supported_quant_types}") + supported_dequant_types = [ + torch.float, + torch.float32, + torch.float16, + torch.bfloat16, + ] + if dtype not in supported_dequant_types: + raise ValueError( + f"Unsupported dtype to dequantize to. Supported dtypes must be one of {supported_dequant_types}" + ) + + # Needed to prevent underflow in cases where the zero_point is larger than + # the quantized value. + if not input_tensor.dtype.is_signed: + input_tensor = input_tensor.to(torch.int32) + + return (input_tensor - zero_point).to(dtype) * scale + + +@impl(m, "quantized_add") +def quantized_add( + X: torch.Tensor, + X_scale: torch.Tensor, + X_zero_point: torch.Tensor, + Y: torch.Tensor, + Y_scale: torch.Tensor, + Y_zero_point: torch.Tensor, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + """ + Sums up two quantized tensors and returns another quantized tensor. The intuition + is that we want dequant(out) ~= dequant(X) + dequant(Y) + + If we do that math, we get + out_scale(out - out_zero_point) = X_scale(X - X_zero_point) + Y_scale(Y - Y_zero_point) + + Rearranging, we get + out = (X_scale(X - X_zero_point) + Y_scale(Y - Y_zero_point)) / out_scale + out_zero_point + + Args: + - X (Tensor): The first operand + - X_scale (Tensor): The ratio between the sizes of X's floating point and quantized + ranges + - X_zero_point (Tensor): The quantized mapping of zero for X + - Y (Tensor): The second operand + - Y_scale (Tensor): The ratio between the sizes of Y's floating point and quantized + ranges + - Y_zero_point (Tensor): The quantized mapping of zero for Y + - out_scale (float): The ratio between the sizes of the output's floating point and + quantized ranges + - out_zero_point (int): The quantized mapping of zero for the output + """ + supported_dtypes = [torch.int8, torch.uint8] + if X.dtype != Y.dtype: + raise ValueError("X and Y dtypes need to match") + + dtype = X.dtype + if dtype not in supported_dtypes: + raise ValueError( + f"X and Y dtypes need to be in {supported_dtypes}. Got {dtype}" + ) + + if dtype == torch.uint8: + X = X.to(torch.int8) + Y = Y.to(torch.int8) + + # TODO(agrebenisan): This should be done in fixed point arithmetic, but to match the quantized_add_out.cpp + # reference implementation, we'll do it in floating point. + dequant_X = X_scale * (X - X_zero_point) + dequant_Y = Y_scale * (Y - Y_zero_point) + + out_scale_inv = 1 / out_scale + + # q_min/q_max are unused args + return quantize_per_tensor( + dequant_X + dequant_Y, + out_scale_inv, + out_zero_point, + torch.iinfo(dtype).min, + torch.iinfo(dtype).max, + dtype, + ) + + +def quantized_linear_common( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: torch.Tensor | int, + out_multiplier: torch.Tensor | int, + out_shift: int, + out_zero_point: int, +) -> torch.Tensor: + """ + Quantized linear (transposed matmul) operation. + + Args: + - src (Tensor): The activations tensor + - weight (Tensor): The weight tensor + - bias (Tensor): The bias tensor + - in_zero_point (int): The quantized mapping of zero for the input + - weight_zero_point (Tensor): The quantized mapping of zero for the weight + - out_multiplier (Tensor): The multiplier used to scale the output + - out_shift (Tensor): The shift used to scale the output + - out_zero_point (int): The quantized mapping of zero for the output + - offset (Tensor): Unused + """ + out_scale = -out_multiplier * (1 / (1 << 31)) * (2**out_shift) + out_scale_inv = 1 / out_scale + + N, K = weight.shape + + leading_dims = src.shape[:-1] + src = src.view(-1, K) + + dtype = src.dtype + supported_dtypes = [torch.int8, torch.uint8, torch.int32] + if dtype not in supported_dtypes: + raise ValueError( + f"Unsupported dtype to quantize to. Supported dtypes must be one of {supported_dtypes}" + ) + + out = torch.nn.functional.linear( + (src - in_zero_point).float(), + (weight - weight_zero_point).float(), + bias.float(), + ) + return quantize_per_tensor( + out, + out_scale_inv, + out_zero_point, + torch.iinfo(dtype).min, + torch.iinfo(dtype).max, + dtype, + ).reshape(*leading_dims, N) + + +def quantized_linear_variant( + per_tensor: bool, + fully_connected: bool, + src_dtype: torch.dtype | None = None, + weight_dtype: torch.dtype | None = None, +) -> Callable[[Callable[..., torch.Tensor]], Callable[..., torch.Tensor]]: + + def decorator(_: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + def variant( + src: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + in_zero_point: int, + weight_zero_point: torch.Tensor | int, + out_multiplier: torch.Tensor | int, + out_shift: torch.Tensor | int, + out_zero_point: int, + offset: torch.Tensor | None = None, + ) -> torch.Tensor: + if fully_connected and src.shape[0] != 1: + raise ValueError( + "Fully connected quantized linear only supports batch size of 1" + ) + if src_dtype and src.dtype != src_dtype: + raise ValueError( + f"src dtype must be {src_dtype}. Got {src.dtype} instead" + ) + if weight_dtype and weight.dtype != weight_dtype: + raise ValueError( + f"weight dtype must be {weight_dtype}. Got {weight.dtype} instead" + ) + if bias.dtype != torch.int32: + raise ValueError( + f"bias dtype must be torch.int32. Got {bias.dtype} instead" + ) + + if per_tensor: + assert isinstance(weight_zero_point, int) + assert isinstance(out_multiplier, int) + assert isinstance(out_shift, int) + return quantized_linear_common( + src, + weight, + bias, + in_zero_point, + weight_zero_point, + out_multiplier, + out_shift, + out_zero_point, + ) + else: + assert isinstance(out_shift, torch.Tensor) + if out_shift.numel() != 1: + raise ValueError("out_shift must be a scalar") + + if out_shift.dtype != torch.int64: + raise ValueError("out_shift must be an int64") + + return quantized_linear_common( + src, + weight, + bias, + in_zero_point, + weight_zero_point, + out_multiplier, + int(out_shift.item()), + out_zero_point, + ) + + return variant + + return decorator + + +@impl(m, "quantized_linear") +@quantized_linear_variant(False, False) +def quantized_linear() -> torch.Tensor: ... + + +@impl(m, "quantized_linear.per_tensor") +@quantized_linear_variant(True, False) +def quantized_linear_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_linear_asym8sxasym8s_asym8s.per_tensor") +@quantized_linear_variant(True, False, torch.int8, torch.int8) +def quantized_linear_asym8sxasym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_linear_asym8uxasym8u_asym8u.per_tensor") +@quantized_linear_variant(True, False, torch.uint8, torch.uint8) +def quantized_linear_asym8uxasym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_fully_connected") +@quantized_linear_variant(False, True) +def quantized_fully_connected() -> torch.Tensor: ... + + +@impl(m, "quantized_fully_connected.per_tensor") +@quantized_linear_variant(True, True) +def quantized_fully_connected_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor") +@quantized_linear_variant(True, True, torch.int8, torch.int8) +def quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor") +@quantized_linear_variant(True, True, torch.uint8, torch.uint8) +def quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_layer_norm.per_tensor") +def quantized_layer_norm_per_tensor( + input_tensor: torch.Tensor, + X_scale: float, + X_zero_point: int, + normalized_shape: list[int], + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + output_scale: float, + output_zero_point: int, +) -> torch.Tensor: + """ + Quantized layer norm operation. + + Args: + - input_tensor (Tensor): The activations tensor + - X_scale (float): The scale of the input + - X_zero_point (int): The zero point of the input + - normalized_shape (int): The shape of the input + - weight (Tensor): The weight tensor + - bias (Tensor): The bias tensor + - eps (float): The epsilon value + - output_scale (float): The scale of the output + - output_zero_point (int): The zero point of the output + """ + supported_dtypes = [torch.int8, torch.uint8] + if input_tensor.dtype not in supported_dtypes: + raise ValueError( + f"Input dtype must be one of {supported_dtypes}. Got {input_tensor.dtype}" + ) + + float_input_tensor = dequantize_per_tensor( + input_tensor, X_scale, X_zero_point, -128, 127, torch.float32 + ) + out = torch.nn.functional.layer_norm( + float_input_tensor, normalized_shape, weight, bias, eps=eps + ) + + return quantize_per_tensor( + out, + 1 / output_scale, + output_zero_point, + torch.iinfo(input_tensor.dtype).min, + torch.iinfo(input_tensor.dtype).max, + input_tensor.dtype, + ) + + +def quantized_conv_per_tensor( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int, int], + padding: tuple[int, int], + dilation: tuple[int, int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + """ + Quantized convolution operation. + + Args: + - input_tensor (Tensor): The activations tensor + - weight (Tensor): The weight tensor + - bias (Tensor): The bias tensor + - stride (Tuple[int]): The stride of the convolution + - padding (Tuple[int]): The padding of the convolution + - dilation (Tuple[int]): The dilation of the convolution + - groups (int): The number of groups + - in_zero_point (int): The quantized mapping of zero for the input + - weight_zero_point (int): The quantized mapping of zero for the weight + - bias_scale (float): The quantized bias scale + - output_scale (float): The scale of the output + - output_zero_point (int): The zero point of the output + - out_multiplier (int): Unused + - out_shift (int): Unused + """ + if len(input_tensor.shape) == 3: + float_out = torch.nn.functional.conv1d( + (input_tensor - in_zero_point).float(), + (weight - weight_zero_point).float(), + (bias * bias_scale).float(), + stride[1], + padding[1], + dilation[1], + groups, + ) + + elif len(input_tensor.shape) == 4: + float_out = torch.nn.functional.conv2d( + (input_tensor - in_zero_point).float(), + (weight - weight_zero_point).float(), + (bias * bias_scale).float(), + stride, + padding, + dilation, + groups, + ) + else: + raise ValueError("Input tensor must be 3D or 4D") + + return quantize_per_tensor( + float_out, + 1.0 / output_scale, + output_zero_point, + torch.iinfo(input_tensor.dtype).min, + torch.iinfo(input_tensor.dtype).max, + input_tensor.dtype, + ) + + +@impl(m, "quantized_conv_nchw.per_tensor") +def quantized_conv_nchw_per_tensor( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int, int], + padding: tuple[int, int], + dilation: tuple[int, int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + """ + Quantized convolution operation. + + Args: + - input_tensor (Tensor): The activations tensor + - weight (Tensor): The weight tensor + - bias (Tensor): The bias tensor + - stride (Tuple[int]): The stride of the convolution + - padding (Tuple[int]): The padding of the convolution + - dilation (Tuple[int]): The dilation of the convolution + - groups (int): The number of groups + - in_zero_point (int): The quantized mapping of zero for the input + - weight_zero_point (int): The quantized mapping of zero for the weight + - bias_scale (float): The quantized bias scale + - output_scale (float): The scale of the output + - output_zero_point (int): The zero point of the output + - out_multiplier (int): Unused + - out_shift (int): Unused + """ + if not input_tensor.is_contiguous(memory_format=torch.contiguous_format): + raise ValueError("Input tensor must be in NCHW format") + return quantized_conv_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + + +@impl(m, "quantized_conv_nhwc.per_tensor") +def quantized_conv_nhwc_per_tensor( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int, int], + padding: tuple[int, int], + dilation: tuple[int, int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + """ + Quantized convolution operation. + + Args: + - input_tensor (Tensor): The activations tensor + - weight (Tensor): The weight tensor + - bias (Tensor): The bias tensor + - stride (Tuple[int]): The stride of the convolution + - padding (Tuple[int]): The padding of the convolution + - dilation (Tuple[int]): The dilation of the convolution + - groups (int): The number of groups + - in_zero_point (int): The quantized mapping of zero for the input + - weight_zero_point (int): The quantized mapping of zero for the weight + - bias_scale (float): The quantized bias scale + - output_scale (float): The scale of the output + - output_zero_point (int): The zero point of the output + - out_multiplier (int): Unused + - out_shift (int): Unused + """ + + # Convert to NCHW format to reuse the existing implementation + conv_is_1d = False + if len(input_tensor.shape) == 3: + conv_is_1d = True + input_tensor = input_tensor.movedim(-1, 1).contiguous() + if len(weight.shape) != 3: + raise ValueError("Weight tensor must be 3D if input is 3D") + weight = weight.movedim(-1, 1).contiguous() + else: + input_tensor = input_tensor.movedim(-1, -3) + if len(weight.shape) != 4: + raise ValueError("Weight tensor must be 4D if input is nd > 3") + weight = torch.permute(weight, (0, -1, 1, 2)).contiguous() + + nchw_out = quantized_conv_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + + if conv_is_1d: + return nchw_out.movedim(1, -1).contiguous() + else: + return nchw_out.movedim(-3, -1).contiguous() + + +def quantized_conv_variant( + layout: str, + input_dtype: torch.dtype, + weight_dtype: torch.dtype, +) -> Callable[[Callable[..., torch.Tensor]], Callable[..., torch.Tensor]]: + """Create a quantized conv variant with type checking.""" + + def decorator(_: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + def variant( + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int, int], + padding: tuple[int, int], + dilation: tuple[int, int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, + ) -> torch.Tensor: + assert ( + input_tensor.dtype == input_dtype + ), f"Expected input dtype {input_dtype}, got {input_tensor.dtype}" + assert ( + weight.dtype == weight_dtype + ), f"Expected weight dtype {weight_dtype}, got {weight.dtype}" + + assert ( + bias.dtype == torch.int32 + ), f"Expected bias dtype int32, got {bias.dtype}" + + # Call the appropriate base function + match layout: + case "nchw": + return quantized_conv_nchw_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + case "nhwc": + return quantized_conv_nhwc_per_tensor( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + case _: + raise ValueError(f"Unknown layout {layout}") + + return variant + + return decorator + + +@impl(m, "quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nchw", torch.int8, torch.int8) +def quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nchw", torch.uint8, torch.uint8) +def quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nhwc", torch.int8, torch.int8) +def quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nhwc", torch.uint8, torch.uint8) +def quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nchw", torch.int8, torch.int8) +def quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nchw", torch.uint8, torch.uint8) +def quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nhwc", torch.int8, torch.int8) +def quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nhwc", torch.uint8, torch.uint8) +def quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nchw", torch.int8, torch.int8) +def quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nchw", torch.uint8, torch.uint8) +def quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor") +@quantized_conv_variant("nhwc", torch.int8, torch.int8) +def quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor") +@quantized_conv_variant("nhwc", torch.uint8, torch.uint8) +def quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor() -> torch.Tensor: ... + + +def quantized_relu_common( + X: torch.Tensor, + X_zero_point: torch.Tensor | int, + out_zero_point: int, + out_multiplier: int, + out_shift: int, +) -> torch.Tensor: + """ + Quantized ReLU operation followed by requantization. + + Args: + - X (Tensor): The input tensor + - X_zero_point (Tensor): The quantized mapping of zero for the input + - out_zero_point (int): The quantized mapping of zero for the output + - out_multiplier (Tensor): The multiplier used to scale the output + - out_shift (Tensor): The shift used to scale the output + """ + supported_dtypes = [torch.int8, torch.int16, torch.uint8, torch.uint16] + if X.dtype not in supported_dtypes: + raise ValueError(f"X dtype must be one of {supported_dtypes}. Got {X.dtype}") + + out_scale = -out_multiplier * (1 / (1 << 31)) * (2**out_shift) + dequantized_X = torch.where(X > X_zero_point, X - X_zero_point, torch.zeros_like(X)) + return quantize_per_tensor( + dequantized_X, + out_scale, + out_zero_point, + torch.iinfo(X.dtype).min, + torch.iinfo(X.dtype).max, + X.dtype, + ) + + +def quantized_relu_variant( + per_tensor: bool, + dtype: torch.dtype | None = None, +) -> Callable[[Callable[..., torch.Tensor]], Callable[..., torch.Tensor]]: + """Create a quantized relu variant with type checking.""" + + def decorator(_: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + def variant( + X: torch.Tensor, + X_zero_point: torch.Tensor | int, + out_zero_point: int, + out_multiplier: torch.Tensor | int, + out_shift: torch.Tensor | int, + ) -> torch.Tensor: + if per_tensor: + if dtype and X.dtype != dtype: + raise ValueError(f"X dtype must be {dtype}. Got {X.dtype}") + + assert isinstance(out_shift, int) + assert isinstance(out_multiplier, int) + _out_shift = out_shift + _out_multiplier = out_multiplier + else: + assert isinstance(out_multiplier, torch.Tensor) + if out_multiplier.numel() > 1: + raise ValueError("Only scalar out_multiplier is supported") + + assert isinstance(out_shift, torch.Tensor) + if out_shift.numel() > 1: + raise ValueError("Only scalar out_shift is supported") + + assert isinstance(X_zero_point, torch.Tensor) + if X_zero_point.shape != X.shape: + raise ValueError( + f"X_zero_point shape must be {X.shape}. Got {X_zero_point.shape}" + ) + + _out_multiplier = int(out_multiplier.item()) + _out_shift = int(out_shift.item()) + + return quantized_relu_common( + X, + X_zero_point, + out_zero_point, + _out_multiplier, + _out_shift, + ) + + return variant + + return decorator + + +@impl(m, "quantized_relu") +@quantized_relu_variant(False) +def quantized_relu() -> torch.Tensor: ... + + +@impl(m, "quantized_relu.per_tensor") +@quantized_relu_variant(True) +def quantized_relu_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_relu_asym8s_asym8s.per_tensor") +@quantized_relu_variant(True, torch.int8) +def quantized_relu_asym8s_asym8s_per_tensor() -> torch.Tensor: ... + + +@impl(m, "quantized_relu_asym8u_asym8u.per_tensor") +@quantized_relu_variant(True, torch.uint8) +def quantized_relu_asym8u_asym8u_per_tensor() -> torch.Tensor: ... + + @impl(m, "requantize") def requantize( input: torch.Tensor, diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index 4721e5a1926..663c5825e52 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/backends/cadence/aot/remove_ops.py @@ -7,16 +7,6 @@ # pyre-strict -# This file contains functions to remove operators from the graph. The removed -# ops should belong to either of the following categories: -# 1. The op should be redundant for inference (e.g., dropout). Such ops are grouped -# together in 'RemoveRedundantOps'. Anyone running inference can add this class -# in their pass list, and it should semantic-preserving transformation. -# 2. The op should be redundant for Jarvis (e.g., contiguous). Such ops are grouped -# together in 'CadenceRemoveNops'. The ops removed in this class might not be nop -# in a context outside of Jarvis', so exercise caution while invoking this in a -# pass list outside of Jarvis. - import logging from dataclasses import dataclass, field from typing import cast, List, Optional, Sequence, Set @@ -152,7 +142,7 @@ def call_operator( @register_cadence_pass(CadencePassAttribute(opt_level=0)) class RemoveToOpsPass(ExportPass): - # aten.to.* as of now are all nops for Jarvis + # aten.to.* as of now are all nops def call_operator( self, op, # pyre-ignore @@ -413,7 +403,7 @@ def call_operator( class RemoveAliasCopyOpPass(ExportPass): """ - alias_copy is a no-op for Jarvis and can be removed. + alias_copy is a no-op and can be removed. """ def call_operator( @@ -936,10 +926,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: return super().call(graph_module) -# The following class consolidates functions to remove ops that are redundant -# in Jarvis. Currently, each function in this class iterates over each node of -# the graph module once. In future, we could consolidate them into a monolithic -# function. class CadenceRemoveNops: passes = [ SimplifySliceOpPass, diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index 61ab7b4c40f..001ab95d629 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -7,12 +7,7 @@ # This file contains all the functions that replace one op with another in the -# graph. The functions replacing ops for models deployed with Jarvis are grouped -# together in class 'ReplaceOpsInGraph'. Some examples of functions in the class are -# 1. functions that replace an ATen op with a custom op that accepts extra arguments -# 2. functions that replace in-place variants of ATen ops with out-of-place version. -# 3. functions that replace an ATen op with another semantically equivalent ATen op. -# 4. functions that concretize optional args. +# graph. # pyre-unsafe @@ -20,17 +15,15 @@ import math import operator from operator import neg -from typing import cast, Dict, Iterable, Optional, Sequence, Set, Tuple +from typing import cast, Dict, Iterable, Optional, Sequence, Tuple import torch import torch.fx from executorch.backends.cadence.aot.compiler_utils import ( get_shape, get_tensor_from_attr, - get_transposed_dims, get_zero_point, is_node_with_op, - is_quantized_tensor, quantize_tensor_multiplier, ) from executorch.backends.cadence.aot.fuse_ops import ( @@ -54,7 +47,7 @@ from torch.fx.node import Argument # A map to represent ops that: -# (a) are functionally equivalent wrt. Jarvis; and +# (a) are functionally equivalent; and # (b) have identical arguments # An op whose target is 'key' in this dict can be replaced by the functionally euivalent # op whose target is 'value'. The replacement would just involve changing the op target. @@ -650,7 +643,7 @@ def call_operator(self, op, args, kwargs, meta): # Make that pass runnable standalone at opt level 0. @register_cadence_pass(CadencePassAttribute(opt_level=0)) -class ReplaceAtenConvolutionWithJarvisConvolutionPass(ExportPass): +class ReplaceAtenConvolutionWithCadenceConvolutionPass(ExportPass): """ Replace aten convolution op with jarvis-specific convolution op, since the aten version is not supported by jarvis. @@ -777,186 +770,6 @@ def call_operator(self, op, args, kwargs, meta): return super().call_operator(target, new_args, kwargs, meta) -# TODO(matthiascremon): this is a fuse op, not a replace op -class ReplaceConvWithChannelLastConv: - """ - Convolution op in pytorch expects NCHW layout for input, weight, and output - tensors. However, if the input and output to the convolution op are originally - in NWHC layout, and are then permuted to conform to NCHW layout, we can fuse - the two permute ops with the convolution op, and call the NHWC layout - convolution op in Jarvis. - """ - - def __init__(self): - self.counter = 0 - self.graph_module = None - - def __call__(self, graph_module: torch.fx.GraphModule): - self.replace_conv_with_nhwc_conv(graph_module) - - def conv_layout_is_nhwc(self, node: torch.fx.Node) -> bool: - """ - Return true if the convolution input and output are connected to permute - ops, and the input/output to/from the permute ops is NHWC layout tensor. - """ - # There must only be a single user of the output node (which must be a - # permute/tranpsose op). The input of the convolution must be connected - # to a permute op, and that permute op should have a single user. - conv_inp = node.args[0] - assert isinstance(conv_inp, torch.fx.Node) - if len(node.users) != 1 or len(conv_inp.users) != 1: - return False - - # Get the input and output (permute/transpose) nodes of the convolution - conv_user = list(node.users.keys())[0] - assert isinstance(conv_user, torch.fx.Node) - pt_nodes: Set[torch.fx.Node] = {conv_inp, conv_user} - - # Any node in pt_nodes must not be a placeholder. - if contains_placeholder_or_param(pt_nodes): - return False - - # Determine if the convolution is 1d or 2d. The output tensor must be - # 3- or 4-dimensional - out_shape = get_shape(self.graph_module, node) - assert out_shape is not None - out_dims = len(out_shape) - assert out_dims in {3, 4}, "Jarvis only supports conv1d and conv2d" - conv1d = out_dims == 3 - - # Get the possible targets for the nodes in pt_nodes. Since conv1d has - # 3-dimensional input and output tensors, the nodes in pt_nodes could - # be either permute or transpose op. For conv2d, the nodes in pt_nodes - # must be permute ops. - p_target = exir_ops.edge.aten.permute_copy.default - t_target = exir_ops.edge.aten.transpose_copy.int - pt_targets = [p_target] + ([t_target] if conv1d else []) - - # If any node in pt_nodes is not permute op (or tranpose op for conv1d), - # bail. - if any(x.target not in pt_targets for x in pt_nodes): - return False - - # Now we need to determine the dimension permutations: - # If the input had NHWC layout, which was then permuted/transposed - # by a permute/transpose op to NCHW layout, the permutation must be - # [0, 3, 2, 1] (or [0, 2, 1] for conv1d). - # If the output had NCHW layout, and was then permuted to NHWC layout, - # the permutation must be [0, 2, 3, 1] (or [0, 2, 1] for conv1d). - nhwc_permute_order = { - node.args[0]: [0, 2, 1] if conv1d else [0, 3, 1, 2], - list(node.users.keys())[0]: [0, 2, 1] if conv1d else [0, 2, 3, 1], - } - for x in pt_nodes: - order = ( - x.args[1] - if x.target == p_target - else get_transposed_dims(x, list(range(out_dims))) - ) - if order != nhwc_permute_order[x]: - return False - - return True - - def replace_conv_with_nhwc_conv(self, graph_module: torch.fx.GraphModule): - self.graph_module = graph_module - graph = graph_module.graph - for node in graph.nodes: - # We are only interested in convolution nodes that have NHWC layout - if node.target not in { - exir_ops.edge.cadence.quantized_conv.default, - exir_ops.edge.cadence.convolution.default, - exir_ops.edge.cadence.quantized_transposed_conv.default, - exir_ops.edge.cadence.transposed_convolution.default, - } or not self.conv_layout_is_nhwc(node): - continue - - # Get the args of convolution op - args = list(node.args) - # The input is connected to a permute/transpose op that converts the - # NHWC layout to NCHW layout. The input of the permute op will become - # this convolution op's input. - in_tp = args[0] - args[0] = in_tp.args[0] - # The weight is in NHWC layout. Permute it to NHWC layout. - weight_tensor = get_tensor_from_attr(graph_module, args[1]) - assert isinstance(weight_tensor, torch.Tensor) - # We cannot directly permute a per-channel quantized tensor. We will - # dequantize it, permute the fp32 tensor, and then requantize the - # permuted tensor. - if ( - is_quantized_tensor(weight_tensor) - and weight_tensor.qscheme() == torch.per_channel_affine - ): - # We have already asserted during quantizing conv op that the - # quantization axis is 0. - dequant_weight = weight_tensor.dequantize() - dequant_weight = ( - dequant_weight.permute([0, 2, 1]) - if dequant_weight.dim() == 3 - else dequant_weight.permute([0, 2, 3, 1]) - ) - weight_tensor = torch.quantize_per_channel( - dequant_weight.contiguous(), - weight_tensor.q_per_channel_scales(), - weight_tensor.q_per_channel_zero_points(), - 0, - weight_tensor.dtype, - ) - else: - weight_tensor = ( - weight_tensor.permute([0, 2, 1]) - if weight_tensor.dim() == 3 - else weight_tensor.permute([0, 2, 3, 1]) - ) - # Make the weight tensor contiguous, since we have permuted it. - weight_tensor = weight_tensor.contiguous() - # Add the permuted weight into the graph, and update the weight in - # args. - with graph.inserting_before(node): - weight_name = f"_weight_nhwc_{self.counter}" - graph_module.register_buffer(weight_name, weight_tensor) - weight = graph.get_attr(weight_name) - args[1] = weight - - # The 'channel_last' arg is True. It is the last arg. - args[-1] = True - # Now update the convolution node args to mark it as NHWC convolution - node.args = tuple(args) - - # Replace all the uses of the permute op connected to the output op - # with this convolution. - out_tp = list(node.users.keys())[0] - out_tp.replace_all_uses_with(node) - node.meta = out_tp.meta - - # Erase the permute ops connected to the input and output of the - # convolution op. - graph.erase_node(in_tp) - graph.erase_node(out_tp) - self.counter += 1 - - graph_module.recompile() - - -# This pass needs to be reworked to be compatible with PT2. It is an optimization -# pass anyway, so move it to opt level 2. -# TODO: T213724613 update and improve this pass. -# @register_cadence_pass(CadencePassAttribute(opt_level=2)) -class ReplaceConvWithChannelLastConvPass(ExportPass): - """ - Replace the ATen convolution op with custom conv op with NCHW or NHWC layout - input tensors, depending on the presence of permute/transpose ops connected - to the input tensor. - """ - - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - result = ReplaceAtenConvolutionWithJarvisConvolutionPass()(graph_module) - assert result is not None - ReplaceConvWithChannelLastConv()(result.graph_module) - return result - - @register_cadence_pass(CadencePassAttribute(opt_level=2)) class ReplaceTrivialConvWithLinear(ExportPass): """ @@ -974,7 +787,8 @@ class ReplaceTrivialConvWithLinear(ExportPass): trivial_conv_op_to_linear_op: Dict[EdgeOpOverload, EdgeOpOverload] = { exir_ops.edge.cadence.convolution.default: exir_ops.edge.aten.linear.default, - exir_ops.edge.cadence.quantized_conv.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv_nchw.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv_nhwc.default: exir_ops.edge.cadence.quantized_linear.default, } def call_operator(self, op, args, kwargs, meta): @@ -985,7 +799,10 @@ def call_operator(self, op, args, kwargs, meta): # and quantized_conv have the same first 8 args. The quantized op has # extra args holding at least the zero point and scale of input, weight, bias, # and output tensor. - quantized_op = op == exir_ops.edge.cadence.quantized_conv.default + quantized_op = ( + op == exir_ops.edge.cadence.quantized_conv_nchw.default + or op == exir_ops.edge.cadence.quantized_conv_nhwc.default + ) assert (len(args) == 8 and not quantized_op) or ( len(args) >= 12 and quantized_op ), "Inconsistent args for convolution" @@ -1132,7 +949,7 @@ def transpose_dims( @register_cadence_pass(CadencePassAttribute(opt_level=3)) -class ForceChannelLastForConvPass(ExportPassWithTransposeHelper): +class ReplaceConvWithChannelLastConvPass(ExportPassWithTransposeHelper): def change_nchw_to_nhwc(self, proxy: ProxyValue, meta: NodeMetadata) -> ProxyValue: shape = proxy.to_tensor().shape if len(shape) == 3: @@ -1162,35 +979,38 @@ def call_operator( ) -> ProxyValue: if op not in { exir_ops.edge.cadence.convolution.default, - exir_ops.edge.cadence.quantized_conv.default, + exir_ops.edge.cadence.quantized_conv_nchw.default, }: return super().call_operator(op, args, kwargs, meta) - quantized_op = op == exir_ops.edge.cadence.quantized_conv.default - channel_last_arg_index = 14 if quantized_op else 7 - channel_last = ( - args[channel_last_arg_index] - if len(args) > channel_last_arg_index - # Default is false (NCHW). - else False - ) - if channel_last: + quantized_op = op == exir_ops.edge.cadence.quantized_conv_nchw.default + + if not quantized_op and len(args) == 8 and args[-1] is True: + # Already in NHWC layout. return super().call_operator(op, args, kwargs, meta) + new_op = ( + exir_ops.edge.cadence.quantized_conv_nhwc.default + if quantized_op + else exir_ops.edge.cadence.convolution.default + ) + input_proxy = cast(ProxyValue, args[0]) weight_proxy = cast(ProxyValue, args[1]) input_proxy = self.change_nchw_to_nhwc(input_proxy, meta) weight_proxy = self.change_nchw_to_nhwc(weight_proxy, meta) + # Non-quantized ops still need to set the last optional argument to True. + channel_last_arg = [] if quantized_op else [True] + new_args = ( # Transposed input/weights. (input_proxy, weight_proxy) # All other args (bias, quant params, etc) - + tuple(args[2:channel_last_arg_index]) - # Channel last. - + (True,) + + tuple(args[2:]) + + tuple(channel_last_arg) ) - output_proxy = super().call_operator(op, new_args, kwargs, meta) + output_proxy = super().call_operator(new_op, new_args, kwargs, meta) nchw_proxy = self.change_nhwc_to_nchw(output_proxy, meta) return nchw_proxy @@ -1247,7 +1067,8 @@ class ReplaceConvWithIm2RowAndLinear(ExportPass): # decompose to. conv_op_to_linear_op: Dict[EdgeOpOverload, EdgeOpOverload] = { exir_ops.edge.cadence.convolution.default: exir_ops.edge.aten.linear.default, - exir_ops.edge.cadence.quantized_conv.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv_nchw.default: exir_ops.edge.cadence.quantized_linear.default, + exir_ops.edge.cadence.quantized_conv_nhwc.default: exir_ops.edge.cadence.quantized_linear.default, } def call_operator(self, op, args, kwargs, meta): @@ -1255,7 +1076,10 @@ def call_operator(self, op, args, kwargs, meta): return super().call_operator(op, args, kwargs, meta) # Get the relevant args from convolution node. - quantized_op = op == exir_ops.edge.cadence.quantized_conv.default + quantized_op = ( + op == exir_ops.edge.cadence.quantized_conv_nchw.default + or op == exir_ops.edge.cadence.quantized_conv_nhwc.default + ) assert (len(args) == 8 and not quantized_op) or ( len(args) >= 12 and quantized_op ), "Inconsistent args for convolution" @@ -1286,9 +1110,7 @@ def call_operator(self, op, args, kwargs, meta): # channel_last layout is specified by the channel_last arg of conv # op, which is either the last argument (15th) or implicitely False # if the op is quantized, or the last argument if not. - channel_last = ( - (args[14] if len(args) == 15 else False) if quantized_op else args[-1] - ) + channel_last = op == exir_ops.edge.cadence.quantized_conv_nhwc.default # The weight tensor is [out_channels, in_channels, X] for NCHW layout, # and [out_channels, X, in_channels] for NHWC layout. Here, X is the # kernel_width for conv1d, and X = kernel_height * kernel_width for @@ -1700,7 +1522,6 @@ def call_operator(self, op, args, kwargs, meta): ) -# pyre-ignore[6]: Incompatible parameter type (doesn't get the inheritance) register_cadence_pass(CadencePassAttribute(opt_level=0))(ReplaceScalarWithTensorArgPass) @@ -1801,8 +1622,12 @@ class ReplaceSingleElementTensorArgumentsFromFullOpWithScalarPass(ExportPass): exir_ops.edge.cadence.quantized_add.per_tensor, [1, 2, 4, 5], ), - exir_ops.edge.cadence.quantized_conv: ( - exir_ops.edge.cadence.quantized_conv.per_tensor, + exir_ops.edge.cadence.quantized_conv_nchw: ( + exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + [8, 9, 12, 13], + ), + exir_ops.edge.cadence.quantized_conv_nhwc: ( + exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, [8, 9, 12, 13], ), exir_ops.edge.cadence.quantized_fully_connected: ( @@ -1871,9 +1696,9 @@ def call_operator(self, op, args, kwargs, meta): @register_cadence_pass(CadencePassAttribute(opt_level=0)) -class ReplaceAtenAvgPoolWithJarvisAvgPoolPass(ExportPass): +class ReplaceAtenAvgPoolWithCadenceAvgPoolPass(ExportPass): """ - Replace the aten avg_pool op with the jarvis custom avg_pool2d op. + Replace the aten avg_pool op with the cadence custom avg_pool2d op. """ def call_operator(self, op, args, kwargs, meta): @@ -2327,10 +2152,16 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # Cast the const_arg to the dtype of the x_arg full_arg = self.resolve_full_arg(x_arg, const_arg) + full_output_dtype = ( + torch.int32 if isinstance(full_arg, int) else torch.float32 + ) + # Extract an argument to a separate full op. with graph_module.graph.inserting_before(mul_node): full_node = graph_module.graph.call_function( - torch.ops.aten.full.default, args=([1], full_arg) + torch.ops.aten.full.default, + args=([1], full_arg), + kwargs={"dtype": full_output_dtype}, ) full_node.meta = mul_node.meta full_node.meta["val"] = [1] @@ -2411,26 +2242,48 @@ def call_operator(self, op, args, kwargs, meta): ) +class CommonReplacePasses: + passes = [ + ReplaceSqueezeAndUnsqueezeWithViewPass, + ReplaceSplitWithSlicePass, + ReplaceSelectWithViewOpPass, + ReplaceMMWithAddMMPass, + ReplaceRepeatWithCatPass, + ReplaceFullLikeWithFullPass, + ] + + +@register_cadence_pass(CadencePassAttribute(opt_level=0)) +class ReplaceAtenLinalgSvdWithCadenceLinalgSvdPass(ExportPass): + """ + Replace aten linalg svd op with cadence custom op. + """ + + def call_operator(self, op, args, kwargs, meta): + if op != exir_ops.edge.aten._linalg_svd.default: + return super().call_operator(op, args, kwargs, meta) + + return super().call_operator( + exir_ops.edge.cadence.linalg_svd.default, args, kwargs, meta + ) + + # This class encapsulates all the functions that replace/switch one op in the # graph with another. class CadenceReplaceOpsInGraph: - passes = [ + passes = CommonReplacePasses.passes + [ + ReplaceAtenLinalgSvdWithCadenceLinalgSvdPass, ReplaceEmptyTensorsWithFullPass, ReplaceFunctionallyEquivalentOpTargets, ReplacePermuteWithTransposePass, ReplaceScalarWithTensorArgPass, ReplaceConvolutionOptionalArgsWithConcreteArgsPass, - ReplaceMMWithAddMMPass, - ReplaceSqueezeAndUnsqueezeWithViewPass, ReplaceAddMMWithLinearPass, RemoveNopSelectOpPass, - ReplaceSelectWithViewOpPass, - ReplaceRepeatWithCatPass, ReplacePadWithCatPass, ReplaceConstantPadNdWithSlicePass, + ReplaceAtenConvolutionWithCadenceConvolutionPass, ReplaceConvWithChannelLastConvPass, - ReplaceAtenConvolutionWithJarvisConvolutionPass, - ForceChannelLastForConvPass, ReplaceTrivialConvWithLinear, ReplaceConvWithIm2RowAndLinear, ReplaceTransposedConvWithLinearPass, @@ -2441,17 +2294,15 @@ class CadenceReplaceOpsInGraph: ReplaceNopTransposeOrPermuteWithViewPass, ReplaceLinearWithFullyConnectedOpPass, ReplaceScalarTensorWithFullPass, - ReplaceFullLikeWithFullPass, ReplaceInfArgInFullWithValuePass, ReplaceLogicalNotBooleanWhereWithWherePass, ReplacePT2QuantWithCadenceQuantPass, ReplacePT2DequantWithCadenceDequantPass, ReplaceSingleElementTensorArgumentsFromFullOpWithScalarPass, ReplaceAdaptiveAvgPoolWithAtenAvgPoolPass, - ReplaceAtenAvgPoolWithJarvisAvgPoolPass, + ReplaceAtenAvgPoolWithCadenceAvgPoolPass, ReplaceWhereWithFullArgsWithWhereScalar, ReplaceAtenApproxGeluWithApproxGeluPass, - ReplaceSplitWithSlicePass, ReplacePowWithMulPass, ReplaceMulTensorWithMulAndFullOpsPass, ] diff --git a/backends/cadence/aot/tests/test_fusion_ops_passes.py b/backends/cadence/aot/tests/test_fusion_ops_passes.py index 556c227b38d..d160a02721a 100644 --- a/backends/cadence/aot/tests/test_fusion_ops_passes.py +++ b/backends/cadence/aot/tests/test_fusion_ops_passes.py @@ -40,7 +40,29 @@ def check_op_counts( self.assertTrue(op_counts_match(graph_module, expected_op_counts)) -class TestFusionPasses(TestFusionPassesBase): +class TestFuseMMWithAddPass(TestFusionPassesBase): + def test_no_fuse_for_3d_bias(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(4, 3, dtype=torch.float32)) + y = builder.placeholder("y", torch.randn(3, 5, dtype=torch.float32)) + z = builder.placeholder("z", torch.randn(1, 4, 5, dtype=torch.float32)) + mm = builder.call_operator( + op=exir_ops.edge.aten.mm.default, + args=(x, y), + ) + output = builder.call_operator(op=exir_ops.edge.aten.add.Tensor, args=(mm, z)) + builder.output([output]) + original_graph = builder.get_graph_module() + + p = FuseMMWithAdd() + converted_graph = cast(PassResult, p(original_graph)).graph_module + converted_graph.graph.eliminate_dead_code() + self.assertEqual( + count_node(converted_graph, exir_ops.edge.aten.addmm.default), 0 + ) + self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.mm.default), 1) + self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.add.Tensor), 1) + def test_fuse_mm_with_add(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(3, 5, dtype=torch.float32)) @@ -176,6 +198,8 @@ def test_keep_mm_add_with_multiple_users(self) -> None: self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.mm.default), 1) self.assertEqual(count_node(converted_graph, exir_ops.edge.aten.add.Tensor), 3) + +class TestFusionPasses(TestFusionPassesBase): def test_permute_transpose_fusion(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(3, 1, 3, 1, 4, dtype=torch.float32)) diff --git a/backends/cadence/aot/tests/test_memory_passes.py b/backends/cadence/aot/tests/test_memory_passes.py index a1da8ede61e..41f903ccf06 100644 --- a/backends/cadence/aot/tests/test_memory_passes.py +++ b/backends/cadence/aot/tests/test_memory_passes.py @@ -1044,7 +1044,7 @@ class DummyMemIdBlockConstraintGen(PassBase): mul: blocks 1, 3 """ - def __init__(self, memory_constraints: MemoryConfig): + def __init__(self, memory_constraints: MemConstraints): self.memory_constraints = memory_constraints def call(self, graph_module: torch.fx.GraphModule) -> PassResult: diff --git a/backends/cadence/aot/tests/test_program_builder.py b/backends/cadence/aot/tests/test_program_builder.py index f2c138dce80..a16d42e2378 100644 --- a/backends/cadence/aot/tests/test_program_builder.py +++ b/backends/cadence/aot/tests/test_program_builder.py @@ -1,10 +1,11 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # pyre-strict - import torch -from executorch.backends.cadence.aot.program_builder import ProgramBuilder +from executorch.backends.cadence.aot.program_builder import IrMode, ProgramBuilder +from executorch.exir.dialects._ops import ops as exir_ops from later.unittest import TestCase +from torch._export.verifier import SpecViolationError from torch.export.graph_signature import InputKind, OutputKind @@ -120,3 +121,102 @@ def test_user_input_mutation(self) -> None: self.assertEqual( program.graph_signature.output_specs[0].kind, OutputKind.USER_INPUT_MUTATION ) + + def test_get_verifier_exir_mode(self) -> None: + """Test that get_verifier returns EXIREdgeDialectVerifier for EXIR mode.""" + builder = ProgramBuilder(mode=IrMode.EXIR) + verifiers = builder.get_verifiers() + self.assertIsNotNone(verifiers) + self.assertEqual(len(verifiers), 1) + + def test_get_verifier_aten_mode(self) -> None: + """Test that get_verifier returns None for ATEN mode.""" + builder = ProgramBuilder(mode=IrMode.ATEN) + verifiers = builder.get_verifiers() + self.assertIsNone(verifiers) + + def test_get_verifier_default_mode(self) -> None: + """Test that get_verifier returns EXIREdgeDialectVerifier for default mode.""" + builder = ProgramBuilder() # Should default to EXIR + self.assertEqual(builder.mode, IrMode.EXIR) + verifiers = builder.get_verifiers() + self.assertIsNotNone(verifiers) + self.assertEqual(len(verifiers), 1) + + def test_aten_add_tensor_exir_mode(self) -> None: + """Test using torch.ops.aten.add.Tensor with EXIR mode.""" + inp = torch.randn([3, 5]) + buffer = torch.randn([5]) + + builder = ProgramBuilder(mode=IrMode.EXIR) + inp_proxy = builder.placeholder("inp", inp) + buffer_proxy = builder.placeholder( + "buffer", buffer, input_kind=InputKind.BUFFER + ) + add = builder.call_operator( + torch.ops.aten.add.Tensor, (inp_proxy, buffer_proxy) + ) + builder.output([add]) + builder.get_program() + + def test_aten_add_tensor_aten_mode(self) -> None: + """Test using torch.ops.aten.add.Tensor with ATEN mode.""" + inp = torch.randn([3, 5]) + buffer = torch.randn([5]) + + builder = ProgramBuilder(mode=IrMode.ATEN) + inp_proxy = builder.placeholder("inp", inp) + buffer_proxy = builder.placeholder( + "buffer", buffer, input_kind=InputKind.BUFFER + ) + add = builder.call_operator( + torch.ops.aten.add.Tensor, (inp_proxy, buffer_proxy) + ) + builder.output([add]) + program = builder.get_program() + + # Verify the program was created successfully + self.assertEqual(len(program.graph_signature.input_specs), 2) + self.assertEqual(len(program.graph_signature.output_specs), 1) + self.assertEqual(builder.mode, IrMode.ATEN) + + def test_exir_edge_aten_add_tensor_exir_mode(self) -> None: + """Test using exir_ops.edge.aten.add.Tensor with EXIR mode.""" + inp = torch.randn([3, 5]) + buffer = torch.randn([5]) + + builder_exir = ProgramBuilder(mode=IrMode.EXIR) + inp_proxy_exir = builder_exir.placeholder("inp", inp) + buffer_proxy_exir = builder_exir.placeholder( + "buffer", buffer, input_kind=InputKind.BUFFER + ) + add_exir = builder_exir.call_operator( + exir_ops.edge.aten.add.Tensor, (inp_proxy_exir, buffer_proxy_exir) + ) + builder_exir.output([add_exir]) + program_exir = builder_exir.get_program() + + # Verify the program was created successfully + self.assertEqual(len(program_exir.graph_signature.input_specs), 2) + self.assertEqual(len(program_exir.graph_signature.output_specs), 1) + self.assertEqual(builder_exir.mode, IrMode.EXIR) + + def test_exir_edge_aten_add_tensor_aten_mode(self) -> None: + """Test using exir_ops.edge.aten.add.Tensor with ATEN mode.""" + inp = torch.randn([3, 5]) + buffer = torch.randn([5]) + + builder_aten = ProgramBuilder(mode=IrMode.ATEN) + inp_proxy_aten = builder_aten.placeholder("inp", inp) + buffer_proxy_aten = builder_aten.placeholder( + "buffer", buffer, input_kind=InputKind.BUFFER + ) + add_aten = builder_aten.call_operator( + exir_ops.edge.aten.add.Tensor, (inp_proxy_aten, buffer_proxy_aten) + ) + builder_aten.output([add_aten]) + + with self.assertRaises( + SpecViolationError, msg="Operator '" + ): + builder_aten.get_program() diff --git a/backends/cadence/aot/tests/test_ref_implementations.py b/backends/cadence/aot/tests/test_ref_implementations.py new file mode 100644 index 00000000000..04b3e8e75ba --- /dev/null +++ b/backends/cadence/aot/tests/test_ref_implementations.py @@ -0,0 +1,1014 @@ +# 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 typing +import unittest + +import executorch.backends.cadence.aot.ops_registrations # noqa +import executorch.backends.cadence.aot.ref_implementations # noqa + +import numpy as np +import torch +from executorch.backends.cadence.aot.typing_stubs import expand + + +class TestRefImplementations(unittest.TestCase): + @expand( + [ + ("basic_int8", 0.42, -1.0, 2.0, -7, 7, torch.int8, 0), + ("basic_int16", 0.42, -1.0, 5.0, -6, 7, torch.int16, -3), + ] + ) + def test_quantize_per_tensor( + self, + name: str, + input_value: float, + f_min: float, + f_max: float, + q_min: int, + q_max: int, + target_dtype: torch.dtype, + expected_value: int, + ) -> None: + input_tensor = torch.tensor([input_value]) + scale = (f_max - f_min) / (q_max - q_min) + inv_scale = 1.0 / scale + zero_point = round(-f_min * inv_scale) + q_min + expected_output = torch.tensor([expected_value], dtype=target_dtype) + + output = torch.ops.cadence.quantize_per_tensor( + input_tensor, inv_scale, zero_point, q_min, q_max, target_dtype + ) + + self.assertEqual( + output.dtype, expected_output.dtype, f"Dtype mismatch in {name}" + ) + self.assertTrue( + torch.equal(output, expected_output), + f"Values don't match in {name}: got {output}, expected {expected_output}", + ) + + @expand( + [ + # Signed quantization ranges + ("signed_range0_int8", 0, -1.0, 2.0, -7, 7, torch.int8, 0.428), + ("signed_range0_int16", 0, -1.0, 2.0, -7, 7, torch.int16, 0.428), + ("signed_range0_int32", 0, -1.0, 2.0, -7, 7, torch.int32, 0.428), + ("signed_range1_int8", -3, -1.0, 5.0, -6, 7, torch.int8, 0.461), + ("signed_range1_int16", -3, -1.0, 5.0, -6, 7, torch.int16, 0.461), + ("signed_range1_int32", -3, -1.0, 5.0, -6, 7, torch.int32, 0.461), + # Unsigned quantization ranges + ("unsigned_range0_uint8", 3, -1.0, 2.0, 0, 7, torch.uint8, 0.428), + ("unsigned_range0_uint16", 3, -1.0, 2.0, 0, 7, torch.uint16, 0.428), + ("unsigned_range1_uint8", 4, -1.0, 5.0, 3, 7, torch.uint8, 0.0), + ("unsigned_range1_uint16", 4, -1.0, 5.0, 3, 7, torch.uint16, 0.0), + ] + ) + def test_dequantize_per_tensor( + self, + name: str, + input_value: int, + f_min: float, + f_max: float, + q_min: int, + q_max: int, + input_dtype: torch.dtype, + expected_value: int, + ) -> None: + input_tensor = torch.tensor([input_value], dtype=input_dtype) + scale = (f_max - f_min) / (q_max - q_min) + zero_point = round(-f_min / scale) + q_min + expected_output = torch.tensor([expected_value], dtype=torch.float32) + + output = torch.ops.cadence.dequantize_per_tensor( + input_tensor, scale, zero_point, q_min, q_max, torch.float32 + ) + + self.assertEqual( + output.dtype, expected_output.dtype, f"Dtype mismatch in {name}" + ) + self.assertTrue( + torch.allclose(output, expected_output, rtol=0.001, atol=0.001), + f"Values don't match in {name}: got {output}, expected {expected_output}", + ) + + @expand( + [ + # Only these types need to be tested as per ET_FORALL_JARVIS_QUANTIZED_TYPES in + # on_device_ai/Assistant/Jarvis/min_runtime/operators/generic/operators.h + ("int16", 5, 0.8, 4, 5, 0.8, 4, 0.8, 4, 6, torch.int8), + ("uint8", 5, 0.8, 4, 5, 0.8, 4, 0.8, 4, 6, torch.uint8), + ] + ) + def test_quantized_add( + self, + name: str, + X: int, + X_scale: float, + X_zero_point: int, + Y: int, + Y_scale: float, + Y_zero_point: int, + out_scale: float, + out_zero_point: int, + expected_value: int, + dtype: torch.dtype, + ) -> None: + X_tensor = torch.tensor([X], dtype=dtype) + Y_tensor = torch.tensor([Y], dtype=dtype) + expected_output = torch.tensor([expected_value], dtype=dtype) + + output = torch.ops.cadence.quantized_add( + X_tensor, + torch.tensor(X_scale), + torch.tensor(X_zero_point, dtype=dtype), + Y_tensor, + torch.tensor(Y_scale), + torch.tensor(Y_zero_point, dtype=dtype), + out_scale, + out_zero_point, + ) + + self.assertTrue( + torch.equal(output, expected_output), + f"Values don't match in {name}: got {output}, expected {expected_output}", + ) + + @expand( + [ + # Test case 1: 1x2 input, 1x2 weight (1 output feature) + *[ + ( + torch.Size([1, 2]), # src_shape: 1 sample, 2 input features + torch.Size( + [1, 2] + ), # weight_shape: 1 output feature, 2 input features + 0, # in_zero_point + torch.tensor([0, 0], dtype=dtype), # weight_zero_point + torch.tensor( + [1073741824], dtype=torch.int32 + ), # out_multiplier (0.5 * 2^31) + torch.tensor([0], dtype=torch.int64), # out_shift + 0, # out_zero_point + torch.tensor([[-2]], dtype=dtype), # expected_output + per_tensor, + ) + for (per_tensor, dtype) in ( + (False, torch.int8), + (True, torch.int8), + (True, torch.uint8), + ) + ], + # Test case 2: 1x3 input, 2x3 weight (2 output features) + *[ + ( + torch.Size([1, 3]), # src_shape: 1 sample, 3 input features + torch.Size( + [2, 3] + ), # weight_shape: 2 output features, 3 input features + 0, # in_zero_point + torch.tensor([0, 0, 0], dtype=dtype), # weight_zero_point + torch.tensor( + [1073741824], dtype=torch.int32 + ), # out_multiplier (0.5 * 2^31) + torch.tensor([0], dtype=torch.int64), # out_shift + 0, # out_zero_point + torch.tensor([[-10, -30]], dtype=dtype), # expected_output + per_tensor, + ) + for (per_tensor, dtype) in ( + (False, torch.int8), + (True, torch.int8), + (True, torch.uint8), + ) + ], + # Test case 3: Batch case with different dimensions + *[ + ( + torch.Size([1, 2, 2]), # src_shape: batch=1, seq=2, features=2 + torch.Size( + [3, 2] + ), # weight_shape: 3 output features, 2 input features + 0, # in_zero_point + torch.tensor([0, 0], dtype=dtype), # weight_zero_point + torch.tensor( + [1073741824], dtype=torch.int32 + ), # out_multiplier (0.5 * 2^31) + torch.tensor([0], dtype=torch.int64), # out_shift + 0, # out_zero_point + torch.tensor( + [[[-2, -8, -14], [-6, -28, -50]]], dtype=dtype + ), # expected_output + per_tensor, + ) + for (per_tensor, dtype) in ( + (False, torch.int8), + (True, torch.int8), + (True, torch.uint8), + ) + ], + # Test case 4: Non-zero zero points + *[ + ( + torch.Size([1, 2]), # src_shape: 1 sample, 2 input features + torch.Size( + [2, 2] + ), # weight_shape: 2 output feature, 1 input feature + 2, # in_zero_point + torch.tensor([1, 1], dtype=dtype), # weight_zero_point + torch.tensor( + [268435456], dtype=torch.int32 + ), # out_multiplier (1.0 * 2^31) + torch.tensor([0], dtype=torch.int64), # out_shift + 1, # out_zero_point + torch.tensor([[-15, 25]], dtype=dtype), # expected_output + per_tensor, + ) + for (per_tensor, dtype) in ( + (False, torch.int8), + (True, torch.int8), + (True, torch.uint8), + ) + ], + # Test case 5: Non-uniform weight zero points + *[ + ( + torch.Size([1, 2]), # src_shape: 1 sample, 2 input features + torch.Size( + [2, 2] + ), # weight_shape: 2 output feature, 1 input feature + 2, # in_zero_point + torch.tensor([1, 2], dtype=dtype), # weight_zero_point + torch.tensor( + [268435456], dtype=torch.int32 + ), # out_multiplier (1.0 * 2^31) + torch.tensor([0], dtype=torch.int64), # out_shift + 1, # out_zero_point + torch.tensor([[-23, 17]], dtype=dtype), # expected_output + False, + ) + for dtype in (torch.int8, torch.uint8) + ], + # Test case 6: Non-zero out_shift (shift=1) + *[ + ( + torch.Size([1, 2]), # src_shape: 1 sample, 2 input features + torch.Size( + [2, 2] + ), # weight_shape: 2 output features, 2 input features + 2, # in_zero_point + torch.tensor([1, 1], dtype=dtype), # weight_zero_point + torch.tensor( + [268435456], dtype=torch.int32 + ), # out_multiplier (0.125 * 2^31) + torch.tensor( + [1], dtype=torch.int64 + ), # out_shift (shift=1, doubles the scale) + 1, # out_zero_point + torch.tensor([[-7, 13]], dtype=dtype), # expected_output + per_tensor, + ) + for (per_tensor, dtype) in ((False, torch.int8), (True, torch.int8)) + ], + ] + ) + def test_quantized_linear( + self, + src_shape: torch.Size, + weight_shape: torch.Size, + in_zero_point: int, + weight_zero_point: torch.Tensor, + out_multiplier: torch.Tensor, + out_shift: torch.Tensor, + out_zero_point: int, + expected_output: torch.Tensor, + per_tensor: bool, + ) -> None: + src = ( + torch.arange(np.prod(src_shape)) + .reshape(src_shape) + .to(expected_output.dtype) + ) + weight = ( + torch.arange(np.prod(weight_shape)) + .reshape(weight_shape) + .to(expected_output.dtype) + ) + bias = torch.arange(weight_shape[0]).to(torch.int32) + if per_tensor: + weight_zero_point = weight_zero_point[0] + out_multiplier = out_multiplier[0] + out_shift = out_shift[0] + + if per_tensor: + match expected_output.dtype: + case torch.int8: + linear_ops = ( + torch.ops.cadence.quantized_linear_asym8sxasym8s_asym8s.per_tensor, + torch.ops.cadence.quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor, + ) + case torch.uint8: + linear_ops = ( + torch.ops.cadence.quantized_linear_asym8uxasym8u_asym8u.per_tensor, + torch.ops.cadence.quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor, + ) + case _: + linear_ops = ( + torch.ops.cadence.quantized_linear.per_tensor, + torch.ops.cadence.quantized_fully_connected.per_tensor, + ) + else: + linear_ops = ( + torch.ops.cadence.quantized_linear, + torch.ops.cadence.quantized_fully_connected, + ) + + for linear_op in linear_ops: + output = linear_op( + src, + weight, + bias, + in_zero_point, + weight_zero_point, + out_multiplier, + out_shift, + out_zero_point, + typing.cast(torch.Tensor, None), + ) + + self.assertTrue(output.dtype == expected_output.dtype, "Dtype mismatch") + + self.assertTrue( + torch.equal(output, expected_output), + f"Values don't match: got {output}, expected {expected_output}", + ) + + @expand( + [ + # Test case 1: Simple case with int8, zero mean input + ( + torch.tensor( + [[-1, 1]], dtype=torch.int8 + ), # input: dequantized to [-0.1, 0.1] + 0.1, # X_scale + 0, # X_zero_point + [2], # normalized_shape (last dimension) + torch.tensor([1.0, 1.0]), # weight + torch.tensor([0.0, 0.0]), # bias + 1e-5, # eps + 0.1, # output_scale + 0, # output_zero_point + torch.int8, # dtype + torch.tensor([[-10, 10]], dtype=torch.int8), # expected_output + ), + # Test case 2: uint8 with zero_point offset + ( + torch.tensor( + [[127, 129]], dtype=torch.uint8 + ), # input: dequantized to [-0.05, 0.05] + 0.05, # X_scale + 128, # X_zero_point + [2], # normalized_shape (last dimension) + torch.tensor([1.0, 1.0]), # weight + torch.tensor([0.0, 0.0]), # bias + 1e-5, # eps + 0.05, # output_scale + 128, # output_zero_point + torch.uint8, # dtype + torch.tensor([[108, 148]], dtype=torch.uint8), # expected_output + ), + # Test case 3: Test with weight and bias scaling + ( + torch.tensor( + [[-2, 2]], dtype=torch.int8 + ), # input: dequantized to [-0.2, 0.2] + 0.1, # X_scale + 0, # X_zero_point + [2], # normalized_shape (last dimension) + torch.tensor( + [2.0, 0.5] + ), # weight: scale first element by 2, second by 0.5 + torch.tensor( + [0.1, -0.1] + ), # bias: add 0.1 to first, subtract 0.1 from second + 1e-5, # eps + 0.1, # output_scale + 0, # output_zero_point + torch.int8, # dtype + torch.tensor([[-19, 4]], dtype=torch.int8), # expected_output + ), + ] + ) + def test_quantized_layer_norm_per_tensor( + self, + input_tensor: torch.Tensor, + X_scale: float, + X_zero_point: int, + normalized_shape: list[int], + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + output_scale: float, + output_zero_point: int, + dtype: torch.dtype, + expected_output: torch.Tensor, + ) -> None: + output = torch.ops.cadence.quantized_layer_norm.per_tensor( + input_tensor, + X_scale, + X_zero_point, + normalized_shape, + weight, + bias, + eps, + output_scale, + output_zero_point, + ) + + # Verify output properties + self.assertEqual(output.dtype, dtype, f"Output dtype should be {dtype}") + self.assertEqual( + output.shape, input_tensor.shape, "Output shape should match input shape" + ) + + # Verify output matches expected values + self.assertTrue( + torch.equal(output, expected_output), + f"Output values don't match expected. Got {output}, expected {expected_output}", + ) + + @expand( + [ + # Test case 1: Basic 2D convolution with int8 + *[ + ( + torch.tensor( + [[[[1, 2], [3, 4]]]], dtype=torch.int8 + ), # input: 1x1x2x2 + torch.tensor( + [[[[1, 0], [0, 1]]]], dtype=torch.int8 + ), # weight: 1x1x2x2 (identity-like) + torch.tensor([0], dtype=torch.int32), # bias + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.1, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[[50]]]], dtype=torch.int8 + ), # expected_output: (1*1 + 4*1) / 0.1 = 50 + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 2: 2D convolution with stride and padding + *[ + ( + torch.tensor( + [[[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]], dtype=torch.int8 + ), # input: 1x1x3x3 + torch.tensor( + [[[[1, 1], [1, 1]]]], dtype=torch.int8 + ), # weight: 1x1x2x2 (sum filter) + torch.tensor([0], dtype=torch.int32), # bias + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.25, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[[48, 64], [96, 112]]]], dtype=torch.int8 + ), # expected_output: convolution results with output_scale=0.25 + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 3: uint8 with non-zero zero points + *[ + ( + torch.tensor( + [[[[130, 132], [134, 136]]]], dtype=torch.uint8 + ), # input: 1x1x2x2 + torch.tensor( + [[[[129, 128], [128, 129]]]], dtype=torch.uint8 + ), # weight: 1x1x2x2 (values close to zero_point) + torch.tensor([10], dtype=torch.int32), # bias + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 128, # in_zero_point + 128, # weight_zero_point + 0.1, # bias_scale + 0.1, # output_scale + 128, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.uint8, # dtype + torch.tensor( + [[[[238]]]], dtype=torch.uint8 + ), # (130 - 128) + (134 - 128) = 10 + # + bias -> 10 + 1 = 11 + # round(11 / 0.1 + 128) = 238, + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 4: 1D convolution (3D input tensor) + *[ + ( + torch.tensor( + [[[1, 2, 3, 4]]], dtype=torch.int8 + ), # input: 1x1x4 (N, C, W) + torch.tensor( + [[[1, 1]]], dtype=torch.int8 + ), # weight: 1x1x2 (OC, IC, KW) + torch.tensor([0], dtype=torch.int32), # bias + (1, 1), # stride (padding for 2D, actual stride is stride[1]) + (0, 0), # padding (padding for 2D, actual padding is padding[1]) + (1, 1), # dilation (padding for 2D, actual dilation is dilation[1]) + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.5, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[6, 10, 14]]], dtype=torch.int8 + ), # expected_output: [1+2, 2+3, 3+4] / 0.5 = [6, 10, 14] + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 5: Multiple output channels + *[ + ( + torch.tensor( + [[[[1, 2], [3, 4]]]], dtype=torch.int8 + ), # input: 1x1x2x2 + torch.tensor( + [ + [[[1, 0], [0, 1]]], # first output channel + [[[0, 1], [1, 0]]], # second output channel + ], + dtype=torch.int8, + ), # weight: 2x1x2x2 + torch.tensor( + [0, 5], dtype=torch.int32 + ), # bias for each output channel + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.2, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[[25]], [[50]]]], dtype=torch.int8 + ), # expected_output: [5/0.2, 10/0.2] = [25, 50] + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 6: Multiple input channels + *[ + ( + torch.tensor( + [ + [ + [[1, 2], [3, 4]], # first input channel + [[5, 6], [7, 8]], + ] # second input channel + ], + dtype=torch.int16, + ), # input: 1x2x2x2 + torch.tensor( + [ + [ + [[1, 0], [0, 1]], # weights for first input channel + [[0, 1], [1, 0]], + ] # weights for second input channel + ], + dtype=torch.int16, + ), # weight: 1x2x2x2 (1 output channel, 2 input channels) + torch.tensor([0], dtype=torch.int32), # bias + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.1, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int16, # dtype + torch.tensor( + [[[[180]]]], dtype=torch.int16 + ), # expected_output: (1 + 4 + 6 + 7) / 0.1 = 180 + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 7: Multiple input and output channels + *[ + ( + torch.tensor( + [ + [ + [[1, 2], [3, 4]], # first input channel + [[2, 1], [4, 3]], + ] # second input channel + ], + dtype=torch.int16, + ), # input: 1x2x2x2 + torch.tensor( + [ + [ + [ + [1, 1], + [1, 1], + ], # first output channel, first input channel + [[1, 1], [1, 1]], + ], # first output channel, second input channel + [ + [ + [1, 0], + [0, 1], + ], # second output channel, first input channel + [[0, 1], [1, 0]], + ], # second output channel, second input channel + ], + dtype=torch.int16, + ), # weight: 2x2x2x2 (2 output channels, 2 input channels) + torch.tensor( + [0, 0], dtype=torch.int32 + ), # bias for each output channel + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.05, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int16, # dtype + torch.tensor([[[[400]], [[200]]]], dtype=torch.int16), + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 8: Grouped convolution (groups=2) + *[ + ( + torch.tensor( + [ + [ + [[1, 2], [3, 4]], # first input channel (group 1) + [[5, 6], [7, 8]], + ] # second input channel (group 2) + ], + dtype=torch.int8, + ), # input: 1x2x2x2 + torch.tensor( + [ + [ + [[1, 1], [1, 1]] + ], # first output channel (processes first input channel) + [ + [[1, 0], [0, 1]] + ], # second output channel (processes second input channel) + ], + dtype=torch.int8, + ), # weight: 2x1x2x2 (2 output channels, 1 input channel each due to groups=2) + torch.tensor( + [0, 0], dtype=torch.int32 + ), # bias for each output channel + (1, 1), # stride + (0, 0), # padding + (1, 1), # dilation + 2, # groups (grouped convolution) + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.2, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[[50]], [[65]]]], dtype=torch.int8 + ), # expected_output: [(1+2+3+4)/0.2, (5+8)/0.2] = [50, 65] + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + # Test case 9: Convolution with stride=2 and padding=1 + *[ + ( + torch.tensor( + [ + [ + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16], + ] + ] + ], + dtype=torch.int8, + ), # input: 1x1x4x4 + torch.tensor( + [[[[1, 1], [1, 1]]]], dtype=torch.int8 + ), # weight: 1x1x2x2 (sum filter) + torch.tensor([0], dtype=torch.int32), # bias + (2, 2), # stride=2 + (1, 1), # padding=1 + (1, 1), # dilation + 1, # groups + 0, # in_zero_point + 0, # weight_zero_point + 1.0, # bias_scale + 0.5, # output_scale + 0, # output_zero_point + 0, # unused out_multiplier + 0, # unused out_shift + torch.int8, # dtype + torch.tensor( + [[[[2, 10, 8], [28, 68, 40], [26, 58, 32]]]], dtype=torch.int8 + ), + memory_format, + ) + for memory_format in [torch.contiguous_format, torch.channels_last] + ], + ] + ) + def test_quantized_conv_per_tensor( + self, + input_tensor: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + stride: tuple[int, int], + padding: tuple[int, int], + dilation: tuple[int, int], + groups: int, + in_zero_point: int, + weight_zero_point: int, + bias_scale: float, + output_scale: float, + output_zero_point: int, + out_multiplier: int, + out_shift: int, + dtype: torch.dtype, + expected_output: torch.Tensor, + memory_format: torch.memory_format, + ) -> None: + assert memory_format in [torch.contiguous_format, torch.channels_last] + + if memory_format == torch.channels_last: + if input_tensor.ndim == 3: + input_tensor = input_tensor.movedim(1, -1) + weight = weight.movedim(1, -1) + else: + input_tensor = input_tensor.movedim(-3, -1) + weight = weight.movedim(-3, -1) + + convs = [ + ( + torch.ops.cadence.quantized_conv_nchw.per_tensor + if memory_format == torch.contiguous_format + else torch.ops.cadence.quantized_conv_nhwc.per_tensor + ) + ] + + optimized_convs = [] + if input_tensor.dtype == torch.int8 and weight.dtype == torch.int8: + if memory_format == torch.contiguous_format: + optimized_convs = [ + torch.ops.cadence.quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor, + torch.ops.cadence.quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor, + torch.ops.cadence.quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor, + ] + + else: + optimized_convs = [ + torch.ops.cadence.quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor, + torch.ops.cadence.quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor, + torch.ops.cadence.quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor, + ] + elif input_tensor.dtype == torch.uint8 and weight.dtype == torch.uint8: + if memory_format == torch.contiguous_format: + optimized_convs = [ + torch.ops.cadence.quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor, + torch.ops.cadence.quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor, + torch.ops.cadence.quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor, + ] + + else: + optimized_convs = [ + torch.ops.cadence.quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor, + torch.ops.cadence.quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor, + torch.ops.cadence.quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor, + ] + + convs.extend(optimized_convs) + for conv in convs: + output = conv( + input_tensor, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out_multiplier, + out_shift, + ) + + if memory_format == torch.channels_last: + if input_tensor.ndim == 3: + output = output.movedim(-1, 1) + else: + output = output.movedim(-1, -3) + + # Verify output properties + self.assertEqual(output.dtype, dtype, f"Output dtype should be {dtype}") + self.assertEqual( + output.shape, + expected_output.shape, + "Output shape should match expected shape", + ) + + # Verify output matches expected values + self.assertTrue( + torch.equal(output, expected_output), + f"Output values don't match expected. Got {output}, expected {expected_output}", + ) + + @expand( + [ + # Test case 1: Basic int8 case with negative scale + *[ + ( + "basic_int8", + torch.tensor([-1, 0, 1, 3], dtype=dtype), # input + 0, # X_zero_point (scalar broadcast) + 0, # out_zero_point + 1073741824, # out_multiplier (0.5 * 2^31) + 0, # out_shift + dtype, # dtype + torch.tensor( + [0, 0, 0, -2], dtype=dtype + ), # expected: relu(-1,0,1,3) = (0,0,1,3) * (-0.5) + 0 = (0,0,-0.5,-1.5) -> (0,0,0,-2) + ) + for dtype in [torch.int8] + ], + # Test case 2: uint8 with non-zero zero point + *[ + ( + "uint8_with_zp", + torch.tensor([126, 128, 130, 132], dtype=dtype), # input + 128, # X_zero_point + 64, # out_zero_point + 536870912, # out_multiplier (0.25 * 2^31) + 0, # out_shift + dtype, # dtype + torch.tensor( + [64, 64, 64, 63], dtype=dtype + ), # expected: relu(-2,0,2,4) = (0,0,2,4) * (-0.25) + 64 = (64,64,63.5,63) -> (64,64,64,63) + ) + for dtype in [torch.uint8] + ], + # Test case 3: All negative values (should all become zero after ReLU) + *[ + ( + "all_negative_int8", + torch.tensor([-5, -3, -1], dtype=dtype), # input + 0, # X_zero_point + 10, # out_zero_point + 1073741824, # out_multiplier (0.5 * 2^31) + 0, # out_shift + dtype, # dtype + torch.tensor( + [10, 10, 10], dtype=dtype + ), # expected: relu(-5,-3,-1) = (0,0,0) * (-0.5) + 10 = (10,10,10) + ) + for dtype in [torch.int8] + ], + # Test case 4: All positive values with shift (scale becomes -0.25) + *[ + ( + "positive_with_shift", + torch.tensor([2, 4, 6, 8], dtype=dtype), # input + 1, # X_zero_point + 5, # out_zero_point + 1073741824, # out_multiplier (0.5 * 2^31) + 1, # out_shift (multiply by 2^1 = 2) + dtype, # dtype + torch.tensor( + [4, 2, 0, -2], dtype=dtype + ), # expected: relu(1,3,5,7) = (1,3,5,7) * (-1.0) + 5 = (4,2,0,-2) + ) + for dtype in [torch.int8, torch.uint8] + ], + # Test case 4: Non-per-tensor + *[ + ( + "non_per_tensor", + torch.tensor([-1, -2, -3, 1, 2, 3], dtype=dtype), # input + torch.tensor([0, 0, 0, 1, 1, 1]), # X_zero_point + 5, # out_zero_point + torch.tensor([1073741824]), # out_multiplier (0.5 * 2^31) + torch.tensor([1]), # out_shift (multiply by 2^1 = 2) + dtype, # dtype + torch.tensor([5, 5, 5, 5, 4, 3], dtype=dtype), + ) + for dtype in [torch.int8] + ], + ] + ) + def test_quantized_relu( + self, + name: str, + X: torch.Tensor, + X_zero_point: torch.Tensor | int, + out_zero_point: int, + out_multiplier: torch.Tensor | int, + out_shift: torch.Tensor | int, + dtype: torch.dtype, + expected_output: torch.Tensor, + ) -> None: + + if isinstance(X_zero_point, int): + assert isinstance(out_multiplier, int) + assert isinstance(out_shift, int) + + match dtype: + case torch.int8: + quantized_relu = ( + torch.ops.cadence.quantized_relu_asym8s_asym8s.per_tensor + ) + case torch.uint8: + quantized_relu = ( + torch.ops.cadence.quantized_relu_asym8u_asym8u.per_tensor + ) + case _: + quantized_relu = torch.ops.cadence.quantized_relu_per_tensor + + output = quantized_relu( + X, + X_zero_point, + out_zero_point, + out_multiplier, + out_shift, + ) + else: + output = torch.ops.cadence.quantized_relu( + X, X_zero_point, out_zero_point, out_multiplier, out_shift + ) + + # Verify output properties + self.assertEqual(output.dtype, dtype, f"Output dtype should be {dtype}") + self.assertEqual(output.shape, X.shape, "Output shape should match input shape") + + # Verify output matches expected values + self.assertTrue( + torch.equal(output, expected_output), + f"Output values don't match expected in {name}. Got {output}, expected {expected_output}", + ) diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index d778cd5b898..ca5168db2be 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/backends/cadence/aot/tests/test_replace_ops_passes.py @@ -17,14 +17,15 @@ ) from executorch.backends.cadence.aot.pass_utils import count_node, op_counts_match from executorch.backends.cadence.aot.replace_ops import ( - ForceChannelLastForConvPass, MakeSliceAndCatDimOutermostPass, ReplaceAdaptiveAvgPoolWithAtenAvgPoolPass, ReplaceAddMMWithLinearPass, ReplaceAtenApproxGeluWithApproxGeluPass, - ReplaceAtenConvolutionWithJarvisConvolutionPass, + ReplaceAtenConvolutionWithCadenceConvolutionPass, + ReplaceAtenLinalgSvdWithCadenceLinalgSvdPass, ReplaceConstantPadNdWithSlicePass, ReplaceConvolutionOptionalArgsWithConcreteArgsPass, + ReplaceConvWithChannelLastConvPass, ReplaceConvWithIm2RowAndLinear, ReplaceEmptyTensorsWithFullPass, ReplaceFunctionallyEquivalentOpTargets, @@ -411,7 +412,7 @@ def test_replace_transposed_conv_with_linear( builder.output([convolution]) original_gm = builder.get_graph_module() - p1 = ReplaceAtenConvolutionWithJarvisConvolutionPass() + p1 = ReplaceAtenConvolutionWithCadenceConvolutionPass() p2 = ReplaceTransposedConvWithLinearPass() graph_after_passes = cast( PassResult, p2(cast(PassResult, p1(original_gm)).graph_module) @@ -969,7 +970,7 @@ def test_replace_conv1d_with_linear(self) -> None: args=(x, weights, bias, [1], [0], [1], 1, False), ) # First, replace the aten convolution with a cadence.convolution op - p1 = ReplaceAtenConvolutionWithJarvisConvolutionPass() + p1 = ReplaceAtenConvolutionWithCadenceConvolutionPass() temp_graph = cast(PassResult, p1(original_gm)).graph_module # temp_graph = p1(original_gm).graph_module self.assertIsNotNone(temp_graph) @@ -1003,7 +1004,7 @@ def test_replace_conv2d_with_linear(self) -> None: args=(x, weights, bias, [1, 1], [0, 0], [1, 1], 1, False), ) # First, replace the aten convolution with a cadence.convolution op - p1 = ReplaceAtenConvolutionWithJarvisConvolutionPass() + p1 = ReplaceAtenConvolutionWithCadenceConvolutionPass() temp_graph = cast(PassResult, p1(original_gm)).graph_module self.assertIsNotNone(temp_graph) @@ -1454,7 +1455,7 @@ def test_replace_linear_like_conv(self) -> None: ) -class TestForceChannelLastForConvPass(unittest.TestCase): +class TestReplaceConvWithChannelLastConvPass(unittest.TestCase): def create_conv1d_graphmodule( self, channels_last: Optional[bool] = None ) -> torch.fx.GraphModule: @@ -1489,7 +1490,7 @@ def test_conv1d_default_channel_last(self) -> None: self.assertEqual(count_node(gm, exir_ops.edge.aten.transpose_copy.int), 0) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( @@ -1514,7 +1515,7 @@ def test_conv1d_no_transpose_if_already_channel_last(self) -> None: self.assertEqual(count_node(gm, exir_ops.edge.cadence.convolution.default), 1) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( @@ -1566,7 +1567,7 @@ def test_convolution_default_channel_last(self) -> None: self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( @@ -1591,7 +1592,7 @@ def test_no_transpose_if_already_channel_last(self) -> None: self.assertEqual(count_node(gm, exir_ops.edge.cadence.convolution.default), 1) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( @@ -1655,28 +1656,49 @@ def create_quantized_convolution_graph_module( out_shift, ) if channels_last is not None: - args = args + (channels_last,) - return single_op_builder( - placeholders=(x, w, b, w_zero_point, b_scale, out_multiplier, out_shift), - op=exir_ops.edge.cadence.quantized_conv.default, - args=args, - ) + return single_op_builder( + placeholders=( + x, + w, + b, + w_zero_point, + b_scale, + out_multiplier, + out_shift, + ), + op=exir_ops.edge.cadence.quantized_conv_nhwc.default, + args=args, + ) + else: + return single_op_builder( + placeholders=( + x, + w, + b, + w_zero_point, + b_scale, + out_multiplier, + out_shift, + ), + op=exir_ops.edge.cadence.quantized_conv_nchw.default, + args=args, + ) def test_quantized_convolution_default_channel_last(self) -> None: # Create a graph with a single convolution node. gm = self.create_quantized_convolution_graph_module() self.assertEqual( - count_node(gm, exir_ops.edge.cadence.quantized_conv.default), 1 + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.default), 1 ) self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( count_node( - gm_after_replacement, exir_ops.edge.cadence.quantized_conv.default + gm_after_replacement, exir_ops.edge.cadence.quantized_conv_nhwc.default ), 1, ) @@ -1685,12 +1707,6 @@ def test_quantized_convolution_default_channel_last(self) -> None: count_node(gm_after_replacement, exir_ops.edge.aten.permute_copy.default), 3, ) - for node in gm_after_replacement.graph.nodes: - if node.target != exir_ops.edge.cadence.quantized_conv.default: - continue - # Check that the channel_last argument is set to True. - self.assertEqual(len(node.args), 15, f"{node=}") - self.assertTrue(node.args[14]) def test_no_transpose_if_already_quantized_conv_channel_last(self) -> None: # Create a graph with a single im2row node. @@ -1698,26 +1714,20 @@ def test_no_transpose_if_already_quantized_conv_channel_last(self) -> None: # Check if graph module is valid by running exportpass on it. gm = ExportPass().call(gm).graph_module self.assertEqual( - count_node(gm, exir_ops.edge.cadence.quantized_conv.default), 1 + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.default), 1 ) # Apply replacement pass. - p = ForceChannelLastForConvPass() + p = ReplaceConvWithChannelLastConvPass() gm_after_replacement = p.call(gm).graph_module # Check that no replacement was made. self.assertEqual( count_node( - gm_after_replacement, exir_ops.edge.cadence.quantized_conv.default + gm_after_replacement, exir_ops.edge.cadence.quantized_conv_nhwc.default ), 1, ) self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0) - for node in gm_after_replacement.graph.nodes: - if node.target != exir_ops.edge.cadence.quantized_conv.default: - continue - # Check that the channel_last argument is set to True. - self.assertEqual(len(node.args), 15, f"{node=}") - self.assertTrue(node.args[14]) class TestMakeSliceAndCatDimOutermostPass(unittest.TestCase): @@ -2036,3 +2046,38 @@ def test_replace_adaptive_avg_pool_with_aten_avg_pool_irregular(self) -> None: len(avg_pool2d_nodes), 0, ) + + +class TestReplaceLinalgSvdPass(unittest.TestCase): + @expand( + [ + ("2x2", (2, 2)), + ("3x3", (3, 3)), + ("4x5", (4, 5)), + ("10x10", (10, 10)), + ] + ) + @torch.no_grad() + def test_replace_aten_linalg_svd_with_cadence_linalg_svd( + self, _: str, shape: Tuple[int, int] + ) -> None: + x = torch.randn(shape, dtype=torch.float32) + original_gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.aten._linalg_svd.default, + args=(x, False, True), + kwargs={"driver": None}, + ) + + p = ReplaceAtenLinalgSvdWithCadenceLinalgSvdPass() + graph_after_passes = cast(PassResult, p(original_gm)).graph_module + + # Assert that the aten linalg_svd op was replaced with cadence linalg_svd op + self.assertEqual( + count_node(graph_after_passes, exir_ops.edge.aten._linalg_svd.default), + 0, + ) + self.assertEqual( + count_node(graph_after_passes, exir_ops.edge.cadence.linalg_svd.default), + 1, + ) diff --git a/backends/cadence/aot/tests/test_type_dispatch_passes.py b/backends/cadence/aot/tests/test_type_dispatch_passes.py new file mode 100644 index 00000000000..52904aecb41 --- /dev/null +++ b/backends/cadence/aot/tests/test_type_dispatch_passes.py @@ -0,0 +1,673 @@ +# 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 unittest +from typing import cast + +import executorch.backends.cadence.aot.ops_registrations # noqa +import torch +from executorch.backends.cadence.aot.graph_builder import single_op_builder +from executorch.backends.cadence.aot.pass_utils import count_node +from executorch.backends.cadence.aot.type_dispatch import CompileTimeTypeDispatchPass +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.passes.infra.pass_base import PassResult + + +class TestTypeDispatchPasses(unittest.TestCase): + def test_int8_dispatch_quantized_fully_connected(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant""" + x = torch.randint(-128, 127, (1, 3), dtype=torch.int8) + w = torch.randint(-128, 127, (4, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_fully_connected.per_tensor, + args=(x, w, b, 0, 0, 1, 0, 0, None), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_fully_connected.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_fully_connected_asym8sxasym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_fully_connected(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant""" + x = torch.randint(0, 255, (1, 3), dtype=torch.uint8) + w = torch.randint(0, 255, (4, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_fully_connected.per_tensor, + args=(x, w, b, 0, 0, 1, 0, 0, None), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_fully_connected.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_fully_connected_asym8uxasym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_linear(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_linear""" + x = torch.randint(-128, 127, (2, 3), dtype=torch.int8) + w = torch.randint(-128, 127, (4, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_linear.per_tensor, + args=(x, w, b, 0, 0, 1, 0, 0, None), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_linear.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_linear_asym8sxasym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_quantized_linear_dispatch(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_linear""" + x = torch.randint(0, 255, (2, 3), dtype=torch.uint8) + w = torch.randint(0, 255, (4, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_linear.per_tensor, + args=(x, w, b, 0, 0, 1, 0, 0, None), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_linear.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_linear_asym8uxasym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_mixed_types_error(self) -> None: + """Test mixed int8/uint8 inputs should raise RuntimeError""" + x = torch.randint(-128, 127, (1, 3), dtype=torch.int8) + w = torch.randint(0, 255, (4, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_fully_connected.per_tensor, + args=(x, w, b, 0, 0, 1, 0, 0, None), + ) + p = CompileTimeTypeDispatchPass() + # Mixed types should raise RuntimeError + with self.assertRaises(RuntimeError) as context: + cast(PassResult, p(gm)).graph_module + self.assertIn("Unsupported input types", str(context.exception)) + + def test_int8_dispatch_quantized_relu(self) -> None: + """Test int8 input should dispatch to asym8s_asym8s variant for quantized_relu""" + x = torch.randint(-128, 127, (2, 3), dtype=torch.int8) + gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.cadence.quantized_relu.per_tensor, + args=(x, 0, 0, 1, 0), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_relu.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_relu_asym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_relu(self) -> None: + """Test uint8 input should dispatch to asym8u_asym8u variant for quantized_relu""" + x = torch.randint(0, 255, (2, 3), dtype=torch.uint8) + gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.cadence.quantized_relu.per_tensor, + args=(x, 0, 0, 1, 0), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_relu.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_relu_asym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_matmul(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_matmul""" + x = torch.randint(-128, 127, (2, 3), dtype=torch.int8) + y = torch.randint(-128, 127, (3, 4), dtype=torch.int8) + bias = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, y, bias), + op=exir_ops.edge.cadence.quantized_matmul.default, + args=(x, 0, y, 0, bias, 1, 0, 0, False), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_matmul.default), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_matmul_asym8sxasym8s_asym8s.default, + ), + 1, + ) + + def test_uint8_dispatch_quantized_matmul(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_matmul""" + x = torch.randint(0, 255, (2, 3), dtype=torch.uint8) + y = torch.randint(0, 255, (3, 4), dtype=torch.uint8) + bias = torch.randint(-2147483648, 2147483647, (4,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, y, bias), + op=exir_ops.edge.cadence.quantized_matmul.default, + args=(x, 0, y, 0, bias, 1, 0, 0, False), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_matmul.default), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_matmul_asym8uxasym8u_asym8u.default, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nchw(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_conv_nchw""" + x = torch.randint(-128, 127, (1, 3, 8, 8), dtype=torch.int8) + w = torch.randint(-128, 127, (16, 3, 3, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [1, 1], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nchw(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_conv_nchw""" + x = torch.randint(0, 255, (1, 3, 8, 8), dtype=torch.uint8) + w = torch.randint(0, 255, (16, 3, 3, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [1, 1], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nhwc(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_conv_nhwc""" + x = torch.randint(-128, 127, (1, 8, 8, 3), dtype=torch.int8) + w = torch.randint(-128, 127, (16, 3, 3, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [1, 1], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nhwc(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_conv_nhwc""" + x = torch.randint(0, 255, (1, 8, 8, 3), dtype=torch.uint8) + w = torch.randint(0, 255, (16, 3, 3, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [1, 1], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nchw_dilated(self) -> None: + """Test int8 x int8 inputs with dilation should dispatch to dilated_asym8sxasym8s_asym8s variant for quantized_conv_nchw_dilated""" + x = torch.randint(-128, 127, (1, 3, 8, 8), dtype=torch.int8) + w = torch.randint(-128, 127, (16, 3, 3, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [2, 2], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_dilated_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nchw_dilated(self) -> None: + """Test uint8 x uint8 inputs with dilation should dispatch to dilated_asym8uxasym8u_asym8u variant for quantized_conv_nchw""" + x = torch.randint(0, 255, (1, 3, 8, 8), dtype=torch.uint8) + w = torch.randint(0, 255, (16, 3, 3, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [2, 2], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_dilated_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nhwc_dilated(self) -> None: + """Test int8 x int8 inputs with dilation should dispatch to dilated_asym8sxasym8s_asym8s variant for quantized_conv_nhwc""" + x = torch.randint(-128, 127, (1, 8, 8, 3), dtype=torch.int8) + w = torch.randint(-128, 127, (16, 3, 3, 3), dtype=torch.int8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [2, 2], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nhwc_dilated(self) -> None: + """Test uint8 x uint8 inputs with dilation should dispatch to dilated_asym8uxasym8u_asym8u variant for quantized_conv_nhwc""" + x = torch.randint(0, 255, (1, 8, 8, 3), dtype=torch.uint8) + w = torch.randint(0, 255, (16, 3, 3, 3), dtype=torch.uint8) + b = torch.randint(-2147483648, 2147483647, (16,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=(x, w, b, [1, 1], [0, 0], [2, 2], 1, 0, 0, 1.0, 1.0, 0, 1, 1), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_add(self) -> None: + """Test int8 x int8 inputs should dispatch to asym8sxasym8s_asym8s variant for quantized_add""" + x = torch.randint(-128, 127, (2, 3), dtype=torch.int8) + y = torch.randint(-128, 127, (2, 3), dtype=torch.int8) + gm = single_op_builder( + placeholders=(x, y), + op=exir_ops.edge.cadence.quantized_add.per_tensor, + args=(x, 1.0, 0, y, 1.0, 0, 1.0, 0), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_add.per_tensor), + 0, + ) + # Should be replaced with int8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_add_asym8sxasym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_add(self) -> None: + """Test uint8 x uint8 inputs should dispatch to asym8uxasym8u_asym8u variant for quantized_add""" + x = torch.randint(0, 255, (2, 3), dtype=torch.uint8) + y = torch.randint(0, 255, (2, 3), dtype=torch.uint8) + gm = single_op_builder( + placeholders=(x, y), + op=exir_ops.edge.cadence.quantized_add.per_tensor, + args=(x, 1.0, 0, y, 1.0, 0, 1.0, 0), + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_add.per_tensor), + 0, + ) + # Should be replaced with uint8 specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_add_asym8uxasym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nchw_depthwise(self) -> None: + """Test int8 x int8 inputs with depthwise should dispatch to depthwise_asym8sxsym8s_asym8s variant for quantized_conv_nchw""" + # Depthwise convolution: groups == input_channels + x = torch.randint(-128, 127, (1, 3, 8, 8), dtype=torch.int8) + w = torch.randint( + -128, 127, (3, 1, 3, 3), dtype=torch.int8 + ) # groups=3, input_channels=3 + b = torch.randint(-2147483648, 2147483647, (3,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=( + x, + w, + b, + [1, 1], + [0, 0], + [1, 1], + 3, + 0, + 0, + 1.0, + 1.0, + 0, + 1, + 1, + ), # groups=3 + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with int8 depthwise specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nchw_depthwise(self) -> None: + """Test uint8 x uint8 inputs with depthwise should dispatch to depthwise_asym8uxasym8u_asym8u variant for quantized_conv_nchw""" + # Depthwise convolution: groups == input_channels + x = torch.randint(0, 255, (1, 3, 8, 8), dtype=torch.uint8) + w = torch.randint( + 0, 255, (3, 1, 3, 3), dtype=torch.uint8 + ) # groups=3, input_channels=3 + b = torch.randint(-2147483648, 2147483647, (3,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + args=( + x, + w, + b, + [1, 1], + [0, 0], + [1, 1], + 3, + 0, + 0, + 1.0, + 1.0, + 0, + 1, + 1, + ), # groups=3 + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nchw.per_tensor), + 0, + ) + # Should be replaced with uint8 depthwise specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) + + def test_int8_dispatch_quantized_conv_nhwc_depthwise(self) -> None: + """Test int8 x int8 inputs with depthwise should dispatch to depthwise_asym8sxsym8s_asym8s variant for quantized_conv_nhwc""" + # Depthwise convolution: groups == input_channels + x = torch.randint(-128, 127, (1, 8, 8, 3), dtype=torch.int8) + w = torch.randint( + -128, 127, (3, 3, 3, 1), dtype=torch.int8 + ) # groups=3, input_channels=3 + b = torch.randint(-2147483648, 2147483647, (3,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=( + x, + w, + b, + [1, 1], + [0, 0], + [1, 1], + 3, + 0, + 0, + 1.0, + 1.0, + 0, + 1, + 1, + ), # groups=3 + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with int8 depthwise specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s.per_tensor, + ), + 1, + ) + + def test_uint8_dispatch_quantized_conv_nhwc_depthwise(self) -> None: + """Test uint8 x uint8 inputs with depthwise should dispatch to depthwise_asym8uxasym8u_asym8u variant for quantized_conv_nhwc""" + # Depthwise convolution: groups == input_channels + x = torch.randint(0, 255, (1, 8, 8, 3), dtype=torch.uint8) + w = torch.randint( + 0, 255, (3, 3, 3, 1), dtype=torch.uint8 + ) # groups=3, input_channels=3 + b = torch.randint(-2147483648, 2147483647, (3,), dtype=torch.int32) + gm = single_op_builder( + placeholders=(x, w, b), + op=exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + args=( + x, + w, + b, + [1, 1], + [0, 0], + [1, 1], + 3, + 0, + 0, + 1.0, + 1.0, + 0, + 1, + 1, + ), # groups=3 + ) + p = CompileTimeTypeDispatchPass() + gm = cast(PassResult, p(gm)).graph_module + # Original op should be replaced + self.assertEqual( + count_node(gm, exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor), + 0, + ) + # Should be replaced with uint8 depthwise specific variant + self.assertEqual( + count_node( + gm, + exir_ops.edge.cadence.quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u.per_tensor, + ), + 1, + ) diff --git a/backends/cadence/aot/type_dispatch.py b/backends/cadence/aot/type_dispatch.py new file mode 100644 index 00000000000..108c4fb1a92 --- /dev/null +++ b/backends/cadence/aot/type_dispatch.py @@ -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. + +# pyre-strict + +from dataclasses import dataclass +from typing import Optional + +import torch +from executorch.backends.cadence.aot.pass_utils import ( + CadencePassAttribute, + register_cadence_pass, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from torch._ops import OpOverload +from torch.fx.node import Argument + + +@dataclass +class OpConfig: + """Configuration for type dispatch operations.""" + + base_name: str + type_dispatch_suffixes: dict[tuple[torch.dtype, ...], str] + weight_arg_idx: Optional[int] = None + variant: str = "per_tensor" + + +@register_cadence_pass(CadencePassAttribute(opt_level=4)) +class CompileTimeTypeDispatchPass(ExportPass): + """ + Replaces generic ops with ops that have explicit types. + """ + + _SUPPORTED_OPS: dict[OpOverload, OpConfig] = { + exir_ops.edge.cadence.quantized_fully_connected.per_tensor: OpConfig( + "quantized_fully_connected", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxasym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxasym8u_asym8u", + }, + weight_arg_idx=1, + ), + exir_ops.edge.cadence.quantized_linear.per_tensor: OpConfig( + "quantized_linear", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxasym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxasym8u_asym8u", + }, + weight_arg_idx=1, + ), + exir_ops.edge.cadence.quantized_matmul.default: OpConfig( + "quantized_matmul", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxasym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxasym8u_asym8u", + }, + weight_arg_idx=2, + variant="default", + ), + exir_ops.edge.cadence.quantized_conv_nchw.per_tensor: OpConfig( + "quantized_conv_nchw", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxsym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxsym8u_asym8u", + }, + weight_arg_idx=1, + ), + exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor: OpConfig( + "quantized_conv_nhwc", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxsym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxsym8u_asym8u", + }, + weight_arg_idx=1, + ), + exir_ops.edge.cadence.quantized_relu.per_tensor: OpConfig( + "quantized_relu", + type_dispatch_suffixes={ + (torch.int8,): "asym8s_asym8s", + (torch.uint8,): "asym8u_asym8u", + }, + ), + exir_ops.edge.cadence.quantized_add.per_tensor: OpConfig( + "quantized_add", + type_dispatch_suffixes={ + (torch.int8, torch.int8): "asym8sxasym8s_asym8s", + (torch.uint8, torch.uint8): "asym8uxasym8u_asym8u", + }, + weight_arg_idx=3, + ), + exir_ops.edge.aten._softmax.default: OpConfig( + "_softmax", + type_dispatch_suffixes={ + (torch.float32,): "f32_f32", + }, + variant="default", + ), + } + + def call_operator( + self, + op: OpOverload, + args: tuple[Argument, ...], + kwargs: dict[str, Argument], + meta: NodeMetadata, + ) -> ProxyValue: + if op not in self._SUPPORTED_OPS: + return super().call_operator(op, args, kwargs, meta) + + config = self._SUPPORTED_OPS[op] + + # pyre-ignore[16]: None has no attribute `to_tensor`. + input_dtype = args[0].to_tensor().dtype + + if config.weight_arg_idx is not None: + weight_dtype = args[config.weight_arg_idx].to_tensor().dtype + dtype_key = (input_dtype, weight_dtype) + else: + dtype_key = (input_dtype,) + + if dtype_key not in config.type_dispatch_suffixes: + raise RuntimeError(f"Unsupported input types for {op}: {dtype_key}") + + type_suffix = config.type_dispatch_suffixes[dtype_key] + base_name = config.base_name + + if op in [ + exir_ops.edge.cadence.quantized_conv_nchw.per_tensor, + exir_ops.edge.cadence.quantized_conv_nhwc.per_tensor, + ]: + groups = args[6] + input_channels = ( + args[0].to_tensor().shape[1] + if op == exir_ops.edge.cadence.quantized_conv_nchw.per_tensor + else args[0].to_tensor().shape[-1] + ) + is_depthwise = groups == input_channels + + dilation = args[5] + # pyre-ignore[16]: None has no attribute '__iter__'. + is_dilated = any(d > 1 for d in dilation) + + if is_dilated: + type_suffix = f"dilated_{type_suffix}" + elif is_depthwise: + type_suffix = f"depthwise_{type_suffix}" + + typed_op_name = f"{base_name}_{type_suffix}" + + typed_op = getattr( + getattr(exir_ops.edge.cadence, typed_op_name), config.variant + ) + + return super().call_operator(typed_op, args, kwargs, meta) diff --git a/backends/cadence/build_cadence_fusionG3.sh b/backends/cadence/build_cadence_fusionG3.sh index 1c84ae99364..93295bc9aa5 100644 --- a/backends/cadence/build_cadence_fusionG3.sh +++ b/backends/cadence/build_cadence_fusionG3.sh @@ -36,7 +36,7 @@ if $STEPWISE_BUILD; then -Bcmake-out . echo "Building any Cadence-specific binaries on top" - CXXFLAGS="-fno-exceptions -fno-rtti" cmake -DBUCK2="$BUCK" \ + CXXFLAGS="-fno-exceptions -fno-rtti" cmake \ -DCMAKE_TOOLCHAIN_FILE=/home/zonglinpeng/ws/zonglinpeng/executorch/backends/cadence/cadence.cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE=Release \ @@ -57,7 +57,7 @@ if $STEPWISE_BUILD; then else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" - CXXFLAGS="-fno-exceptions -fno-rtti" cmake -DBUCK2="$BUCK" \ + CXXFLAGS="-fno-exceptions -fno-rtti" cmake \ -DCMAKE_PREFIX_PATH="${cmake_prefix_path}" \ -DHAVE_SYS_STAT_H=ON \ -DCMAKE_TOOLCHAIN_FILE=./backends/cadence/cadence.cmake \ diff --git a/backends/cadence/build_cadence_hifi4.sh b/backends/cadence/build_cadence_hifi4.sh index e0a48da4074..33078b7ff2f 100644 --- a/backends/cadence/build_cadence_hifi4.sh +++ b/backends/cadence/build_cadence_hifi4.sh @@ -35,7 +35,7 @@ if $STEPWISE_BUILD; then -Bcmake-out . echo "Building any Cadence-specific binaries on top" - CXXFLAGS="-fno-exceptions -fno-rtti" cmake -DBUCK2="$BUCK" \ + CXXFLAGS="-fno-exceptions -fno-rtti" cmake \ -DCMAKE_TOOLCHAIN_FILE=./backends/cadence/cadence.cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE=Release \ @@ -56,7 +56,7 @@ if $STEPWISE_BUILD; then else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" - CXXFLAGS="-fno-exceptions -fno-rtti" cmake -DBUCK2="$BUCK" \ + CXXFLAGS="-fno-exceptions -fno-rtti" cmake \ -DCMAKE_PREFIX_PATH="${cmake_prefix_path}" \ -DCMAKE_TOOLCHAIN_FILE=./backends/cadence/cadence.cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ diff --git a/backends/cadence/cadence.cmake b/backends/cadence/cadence.cmake index 0fa55c6a65b..a0e5ea86da1 100644 --- a/backends/cadence/cadence.cmake +++ b/backends/cadence/cadence.cmake @@ -43,7 +43,7 @@ set(CMAKE_CXX_COMPILER ${TOOLCHAIN_HOME}/bin/${CROSS_COMPILE_TARGET}-clang++) set(CMAKE_C_FLAGS_INIT "-stdlib=libc++ -mtext-section-literals -mlongcalls") set(CMAKE_CXX_FLAGS_INIT "-stdlib=libc++ -mtext-section-literals -mlongcalls") -#workaround for larger compilation time +# workaround for larger compilation time set(CMAKE_CXX_FLAGS_INIT "${CMAKE_CXX_FLAGS_INIT} -fno-strict-aliasing") set(CMAKE_SYSROOT ${TOOLCHAIN_HOME}/${SYSROOT_TARGET}) diff --git a/backends/cadence/fusion_g3/operators/CMakeLists.txt b/backends/cadence/fusion_g3/operators/CMakeLists.txt index c29ffa91af9..a9501c687bb 100644 --- a/backends/cadence/fusion_g3/operators/CMakeLists.txt +++ b/backends/cadence/fusion_g3/operators/CMakeLists.txt @@ -69,16 +69,20 @@ target_link_libraries(aten_ops_cadence PUBLIC executorch) target_link_libraries(aten_ops_cadence PRIVATE xa_nnlib) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. -${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) target_include_directories( - aten_ops_cadence PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} - ${_common_include_directories} - ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/algo/common/include/ - ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/include/nnlib - ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/include - ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/algo/kernels/tables/include + aten_ops_cadence + PUBLIC + ${ROOT_DIR}/.. + ${CMAKE_BINARY_DIR} + ${_common_include_directories} + ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/algo/common/include/ + ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/include/nnlib + ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/include + ${EXECUTORCH_ROOT}/backends/cadence/fusion_g3/third-party/nnlib/nnlib-FusionG3/xa_nnlib/algo/kernels/tables/include ) # Generate C++ bindings to register kernels into both PyTorch (for AOT) and @@ -93,6 +97,4 @@ generate_bindings_for_kernels( ) message("Generated files ${gen_command_sources}") -gen_operators_lib( - LIB_NAME "cadence_ops_lib" KERNEL_LIBS DEPS aten_ops_cadence -) +gen_operators_lib(LIB_NAME "cadence_ops_lib" KERNEL_LIBS DEPS aten_ops_cadence) diff --git a/backends/cadence/fusion_g3/operators/op_clamp.cpp b/backends/cadence/fusion_g3/operators/op_clamp.cpp index 9f3f72a674f..92fb97b1260 100644 --- a/backends/cadence/fusion_g3/operators/op_clamp.cpp +++ b/backends/cadence/fusion_g3/operators/op_clamp.cpp @@ -45,6 +45,7 @@ bool is_out_of_bounds(CTYPE_VAL val) { } ET_NODISCARD bool check_bounds( + KernelRuntimeContext& ctx, const Scalar& val_scalar, const ScalarType& val_type, const ScalarType& out_type, @@ -107,14 +108,14 @@ Tensor& clamp_out( if (has_min) { ET_KERNEL_CHECK( ctx, - check_bounds(min_opt.value(), min_type, out_type, "minimum"), + check_bounds(ctx, min_opt.value(), min_type, out_type, "minimum"), InvalidArgument, out); } if (has_max) { ET_KERNEL_CHECK( ctx, - check_bounds(max_opt.value(), max_type, out_type, "maximum"), + check_bounds(ctx, max_opt.value(), max_type, out_type, "maximum"), InvalidArgument, out); } diff --git a/backends/cadence/fusion_g3/operators/op_native_layer_norm.cpp b/backends/cadence/fusion_g3/operators/op_native_layer_norm.cpp index 7e71df62d54..11c3edbb6a2 100644 --- a/backends/cadence/fusion_g3/operators/op_native_layer_norm.cpp +++ b/backends/cadence/fusion_g3/operators/op_native_layer_norm.cpp @@ -221,12 +221,15 @@ std::tuple native_layer_norm_out( num_elm *= normalized_shape[i]; } + constexpr size_t kAlignment = + 16; // 16-byte alignment for vectorized operations + float* weight_data; if (weight.has_value()) { weight_data = weight.value().mutable_data_ptr(); } else { executorch::runtime::Result temp_mem_weight = - ctx.allocate_temp(num_elm * sizeof(float)); + ctx.allocate_temp(num_elm * sizeof(float), kAlignment); weight_data = (float*)(temp_mem_weight.get()); for (int i = 0; i < num_elm; i++) { @@ -238,7 +241,7 @@ std::tuple native_layer_norm_out( bias_data = bias.value().mutable_data_ptr(); } else { executorch::runtime::Result temp_mem_bias = - ctx.allocate_temp(num_elm * sizeof(float)); + ctx.allocate_temp(num_elm * sizeof(float), kAlignment); bias_data = (float*)(temp_mem_bias.get()); for (int i = 0; i < num_elm; i++) { diff --git a/backends/cadence/hifi/kernels/CMakeLists.txt b/backends/cadence/hifi/kernels/CMakeLists.txt index 972bb4b7ab1..936e28e2241 100644 --- a/backends/cadence/hifi/kernels/CMakeLists.txt +++ b/backends/cadence/hifi/kernels/CMakeLists.txt @@ -28,8 +28,9 @@ add_library( ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_transpose_32.c ) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. -${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) target_include_directories( cadence_kernels @@ -39,7 +40,7 @@ target_include_directories( ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/nnlib-hifi4/xa_nnlib/include/nnlib ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/nnlib-hifi4/xa_nnlib/include ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/nnlib-hifi4/xa_nnlib/algo/ndsp/hifi4/include/ - ${_common_include_directories} + ${_common_include_directories} ) target_link_libraries(cadence_kernels PRIVATE xa_nnlib) diff --git a/backends/cadence/hifi/kernels/kernels.cpp b/backends/cadence/hifi/kernels/kernels.cpp index bf4a2d143fd..feabe6e1828 100644 --- a/backends/cadence/hifi/kernels/kernels.cpp +++ b/backends/cadence/hifi/kernels/kernels.cpp @@ -21,8 +21,27 @@ memcpy(void* dst, const void* src, size_t num_bytes) { } void* allocate_temp_memory(KernelRuntimeContext& ctx, size_t size) { - Result temp_mem_res = ctx.allocate_temp(size); - return temp_mem_res.ok() ? temp_mem_res.get() : nullptr; + constexpr size_t kAlignment = + 16; // 16-byte alignment for vectorized operations + ET_LOG( + Info, + "Attempting to allocate %zu bytes of temp memory (16-byte aligned)", + size); + Result temp_mem_res = ctx.allocate_temp(size, kAlignment); + if (temp_mem_res.ok()) { + void* ptr = temp_mem_res.get(); + ET_LOG( + Info, + "Successfully allocated temp memory at %p (16-byte aligned)", + ptr); + return ptr; + } else { + ET_LOG( + Error, + "Failed to allocate temp memory, error: 0x%x", + static_cast(temp_mem_res.error())); + return nullptr; + } } // Quantize a fp32 value to an int8_t/uint8_t value @@ -108,60 +127,6 @@ void dequantize( } } -// Requantize the int8_t/uint8_t in value to a uint8_t/int8_t out value. -// The scale and zero_point for requantization are in the args. -template -__attribute__((always_inline)) OT requantize( - const IT in, - float in_scale, - int32_t in_zero_point, - float inv_out_scale, - int32_t out_zero_point) { - float dequant = dequantize(in, in_scale, in_zero_point); - return quantize(dequant, inv_out_scale, out_zero_point); -} - -// Requantize the int8_t/uint8_t in array to a uint8_t/int8_t out array. -// The scale and zero_point for requantization are in the args. -template -void requantize( - OT* __restrict__ out, - const IT* __restrict__ in, - float in_scale, - int32_t in_zero_point, - float inv_out_scale, - int32_t out_zero_point, - size_t size) { - xtfloatx2 in_scale_vec = (xtfloatx2)in_scale; - xtfloatx2 in_zero_vec = XT_FLOAT_SX2(in_zero_point, 0); - xtfloatx2 inv_out_scale_vec = (xtfloatx2)inv_out_scale; - xtfloatx2 out_zero_vec = XT_FLOAT_SX2(out_zero_point, 0); - - float min_val = std::numeric_limits::min(); - float max_val = std::numeric_limits::max(); - - size_t i = 0; - // Vectorize by 2 - for (; i < (size & ~1); i += 2) { - xtfloatx2 in_vec = {(float)in[i], (float)in[i + 1]}; - xtfloatx2 t0 = XT_SUB_SX2(in_vec, in_zero_vec); - xtfloatx2 t1 = XT_MUL_SX2(t0, in_scale_vec); - - xtfloatx2 acc = out_zero_vec; - XT_MADD_SX2(acc, inv_out_scale_vec, t1); - xtfloatx2 t2 = XT_FIROUND_SX2(acc); - ae_int32x2 t3 = - XT_UTRUNC_SX2(XT_MAX_SX2(XT_MIN_SX2(t2, max_val), min_val), 0); - out[i] = AE_MOVAD32_H(t3); - out[i + 1] = AE_MOVAD32_L(t3); - } - // Handle residual iteration - if (i < size) { - out[i] = requantize( - in[i], in_scale, in_zero_point, inv_out_scale, out_zero_point); - } -} - // explicit template instantiation #define typed_quantize_val(dtype) \ @@ -210,34 +175,6 @@ typed_dequantize_vec(uint16_t); typed_dequantize_vec(int32_t); #undef typed_dequantize_vec -#define typed_requantize_val(itype, otype) \ - template __attribute__((always_inline)) otype requantize( \ - const itype in, \ - float in_scale, \ - int32_t in_zero_point, \ - float inv_out_scale, \ - int32_t out_zero_point); -typed_requantize_val(int8_t, int8_t); -typed_requantize_val(uint8_t, uint8_t); -typed_requantize_val(int8_t, uint8_t); -typed_requantize_val(uint8_t, int8_t); -#undef typed_requantize_val - -#define typed_requantize_vec(itype, otype) \ - template void requantize( \ - otype* __restrict__ out, \ - const itype* __restrict__ in, \ - float in_scale, \ - int32_t in_zero_point, \ - float inv_out_scale, \ - int32_t out_zero_point, \ - size_t size); -typed_requantize_vec(int8_t, int8_t); -typed_requantize_vec(uint8_t, uint8_t); -typed_requantize_vec(int8_t, uint8_t); -typed_requantize_vec(uint8_t, int8_t); -#undef typed_requantize_vec - }; // namespace kernels }; // namespace HiFi }; // namespace impl diff --git a/backends/cadence/hifi/kernels/targets.bzl b/backends/cadence/hifi/kernels/targets.bzl index 596acc90791..fe5feed69b5 100644 --- a/backends/cadence/hifi/kernels/targets.bzl +++ b/backends/cadence/hifi/kernels/targets.bzl @@ -16,6 +16,7 @@ def define_common_targets(): compatible_with = ["ovr_config//cpu:xtensa"], visibility = [ "//executorch/backends/cadence/...", + "@EXECUTORCH_CLIENTS", ], exported_deps = [ "fbsource//third-party/nnlib-hifi4/xa_nnlib:libxa_nnlib_common", diff --git a/backends/cadence/hifi/operators/CMakeLists.txt b/backends/cadence/hifi/operators/CMakeLists.txt index 806e2e41ff5..6bd63c6d9f6 100644 --- a/backends/cadence/hifi/operators/CMakeLists.txt +++ b/backends/cadence/hifi/operators/CMakeLists.txt @@ -72,14 +72,15 @@ set(_aten_ops__srcs "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/select_copy_util.cpp" "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/slice_util.cpp" "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/delinearize_index.cpp" - ) +) add_library(aten_ops_cadence ${_aten_ops__srcs}) target_link_libraries(aten_ops_cadence PUBLIC executorch) target_link_libraries(aten_ops_cadence PRIVATE cadence_kernels) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. -${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) target_include_directories( aten_ops_cadence PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} @@ -88,9 +89,16 @@ target_include_directories( # Custom ops that are needed to run the test model. add_library( - custom_ops "op_quantized_linear_out.cpp" "op_quantized_layer_norm.cpp" "op_quantized_matmul_out.cpp" - "op_quantize_per_tensor.cpp" "op_quantized_relu_out.cpp" "op_dequantize_per_tensor.cpp" - "op_quantized_conv_out.cpp" "op_quantized_fully_connected_out" + custom_ops + "op_quantized_linear_out.cpp" + "op_quantized_layer_norm.cpp" + "op_quantized_matmul_out.cpp" + "op_quantize_per_tensor.cpp" + "op_quantized_relu_out.cpp" + "op_dequantize_per_tensor.cpp" + "op_quantized_conv_nchw_out.cpp" + "op_quantized_conv_nhwc_out.cpp" + "op_quantized_fully_connected_out" ) target_include_directories( custom_ops PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} diff --git a/backends/cadence/hifi/operators/op_cat.cpp b/backends/cadence/hifi/operators/op_cat.cpp index 8ad52753de3..d4fd51871ce 100644 --- a/backends/cadence/hifi/operators/op_cat.cpp +++ b/backends/cadence/hifi/operators/op_cat.cpp @@ -126,29 +126,25 @@ Tensor& cat_out( const size_t outer = getLeadingDims(out, dim); const size_t dim_stride = getTrailingDims(out, dim); const size_t ninputs = tensors.size(); + const size_t element_size = out.element_size(); + char* out_ptr = static_cast(out.mutable_data_ptr()); - const auto out_type = out.scalar_type(); - ET_SWITCH_REALHB_TYPES(out_type, ctx, name, CTYPE_OUT, [&] { - CTYPE_OUT* out_ptr = out.mutable_data_ptr(); - for (size_t i = 0; i < outer; ++i) { - for (size_t j = 0; j < ninputs; ++j) { - const auto in_type = tensors[j].scalar_type(); - ET_SWITCH_REALHB_TYPES(in_type, ctx, name, CTYPE_IN, [&] { - if (tensors[j].numel() == 0) { - return; - } - size_t inner = tensors[j].size(dim) * dim_stride; - const CTYPE_IN* const in_ptr = - tensors[j].const_data_ptr() + i * inner; - - for (size_t k = 0; k < inner; ++k) { - out_ptr[k] = static_cast(in_ptr[k]); - } - out_ptr += inner; - }); + for (size_t i = 0; i < outer; ++i) { + for (size_t j = 0; j < ninputs; ++j) { + if (tensors[j].numel() == 0) { + continue; } + size_t inner_elements = tensors[j].size(dim) * dim_stride; + size_t contiguous_bytes = inner_elements * element_size; + + const char* const in_ptr = + static_cast(tensors[j].const_data_ptr()) + + i * contiguous_bytes; + + std::memcpy(out_ptr, in_ptr, contiguous_bytes); + out_ptr += contiguous_bytes; } - }); + } return out; } @@ -156,4 +152,4 @@ Tensor& cat_out( } // namespace native } // namespace HiFi } // namespace impl -} // namespace cadence \ No newline at end of file +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_dequantize_per_tensor.cpp b/backends/cadence/hifi/operators/op_dequantize_per_tensor.cpp index 243eeef2e1d..7dce0050d7f 100644 --- a/backends/cadence/hifi/operators/op_dequantize_per_tensor.cpp +++ b/backends/cadence/hifi/operators/op_dequantize_per_tensor.cpp @@ -46,9 +46,6 @@ void dequantize_per_tensor_out( input.scalar_type() == ScalarType::UInt16) { const uint16_t* input_data = input.const_data_ptr(); dequantize(out_data, input_data, scale, zero_point, numel); - } else if (input.scalar_type() == ScalarType::Int) { - const int32_t* input_data = input.const_data_ptr(); - dequantize(out_data, input_data, scale, zero_point, numel); } else { ET_CHECK_MSG( false, diff --git a/backends/cadence/hifi/operators/op_mm.cpp b/backends/cadence/hifi/operators/op_mm.cpp index abb53a7ad7c..9cf922cbf56 100644 --- a/backends/cadence/hifi/operators/op_mm.cpp +++ b/backends/cadence/hifi/operators/op_mm.cpp @@ -79,6 +79,15 @@ Tensor& mm_out( (WORD32* __restrict__)kernels::allocate_temp_memory( ctx, (n * p) * sizeof(WORD32)); + // Allocate zero-initialized bias for matmul function (it doesn't accept + // NULL) + FLOAT32* __restrict__ p_bias_zero = + (FLOAT32* __restrict__)kernels::allocate_temp_memory( + ctx, m * sizeof(FLOAT32)); + + // Initialize bias to zero since mm operation has no bias + memset(p_bias_zero, 0, m * sizeof(FLOAT32)); + WORD32 p_inp_shape[2]; p_inp_shape[0] = n; p_inp_shape[1] = p; @@ -109,11 +118,13 @@ Tensor& mm_out( const FLOAT32* __restrict__ p_vec = (const FLOAT32* __restrict__)p_o; + // mm will always be converted to addmm and to linear, and move transpose to + // graph WORD32 val = xa_nn_matmul_f32xf32_f32( p_out, p_mat1, p_vec, - NULL, + p_bias_zero, rows, cols1, row_stride1, @@ -121,7 +132,6 @@ Tensor& mm_out( vec_offset, out_offset, out_stride); - return out; } diff --git a/backends/cadence/hifi/operators/op_permute_copy.cpp b/backends/cadence/hifi/operators/op_permute_copy.cpp index 1d56d79dfd5..c5f33435733 100644 --- a/backends/cadence/hifi/operators/op_permute_copy.cpp +++ b/backends/cadence/hifi/operators/op_permute_copy.cpp @@ -70,8 +70,6 @@ Tensor& permute_copy_out( out); const auto in_type = out.scalar_type(); - - constexpr auto name = "permute_copy.out"; constexpr int kNnlibMaxDim = 16; bool optimized = false; @@ -150,18 +148,22 @@ Tensor& permute_copy_out( size_t trailing_dims_memo[kTensorDimensionLimit]; executorch::runtime::memoizeTrailingDims(in, trailing_dims_memo); - // in and out must be the same dtype - ET_SWITCH_ALL_TYPES(in_type, ctx, name, CTYPE, [&] { - const CTYPE* const in_data = in.const_data_ptr(); - CTYPE* const out_data = out.mutable_data_ptr(); + const char* const in_data = static_cast(in.const_data_ptr()); + char* const out_data = static_cast(out.mutable_data_ptr()); + const size_t element_size = out.element_size(); - for (size_t i = 0; i < out.numel(); ++i) { - out_data[i] = - in_data[executorch::runtime::coordinateToIndexWithTrailingDimsMemo( - in, in_coord, trailing_dims_memo)]; - increment_coordinate_permuted(in, in_coord, dims); - } - }); + for (size_t i = 0; i < out.numel(); ++i) { + const size_t in_index = + executorch::runtime::coordinateToIndexWithTrailingDimsMemo( + in, in_coord, trailing_dims_memo); + + std::memcpy( + out_data + i * element_size, + in_data + in_index * element_size, + element_size); + + increment_coordinate_permuted(in, in_coord, dims); + } return out; } @@ -169,4 +171,4 @@ Tensor& permute_copy_out( } // namespace native } // namespace HiFi } // namespace impl -} // namespace cadence \ No newline at end of file +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantize_per_tensor.cpp b/backends/cadence/hifi/operators/op_quantize_per_tensor.cpp index df09c7a12dd..ec649a29c5b 100644 --- a/backends/cadence/hifi/operators/op_quantize_per_tensor.cpp +++ b/backends/cadence/hifi/operators/op_quantize_per_tensor.cpp @@ -109,10 +109,6 @@ void quantize_per_tensor_out( uint16_t* out_data = out.mutable_data_ptr(); cadence::impl::HiFi::kernels::quantize( out_data, input_data, 1. / scale, zero_point, numel); - } else if (out.scalar_type() == ScalarType::Int) { - int32_t* out_data = out.mutable_data_ptr(); - cadence::impl::HiFi::kernels::quantize( - out_data, input_data, 1. / scale, zero_point, numel); } else { ET_KERNEL_CHECK_MSG( ctx, diff --git a/backends/cadence/hifi/operators/op_quantized_add_asym8sxasym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_add_asym8sxasym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..71ca15636ef --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_add_asym8sxasym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,172 @@ +/* + * 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 + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using ::impl::reference::kernels::dequantize; +using ::impl::reference::kernels::quantize; + +void quantized_add_asym8sxasym8s_asym8s_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + const int8_t* __restrict__ X_data = X.const_data_ptr(); + const int8_t* __restrict__ Y_data = Y.const_data_ptr(); + int8_t* __restrict__ out_data = out.mutable_data_ptr(); + + ssize_t Y_numel = Y.numel(); + ssize_t X_numel = X.numel(); + ssize_t out_numel = out.numel(); + + float X_scale_f = static_cast(X_scale); + float Y_scale_f = static_cast(Y_scale); + float out_scale_f = static_cast(out_scale); + int32_t X_zero_point_i32 = static_cast(X_zero_point); + int32_t Y_zero_point_i32 = static_cast(Y_zero_point); + int32_t out_zero_point_i32 = static_cast(out_zero_point); + + float inv_out_scale = 1.0f / out_scale_f; + constexpr float min_val = + static_cast(std::numeric_limits::min()); + constexpr float max_val = + static_cast(std::numeric_limits::max()); + + /* Tensor X exactly matches Y in shape, no broadcasting */ + if (X_numel == Y_numel && Y_numel == out_numel) { + for (size_t i = 0; i < X_numel; ++i) { + float x = X_scale_f * (X_data[i] - X_zero_point_i32); + float y = Y_scale_f * (Y_data[i] - Y_zero_point_i32); + float z = x + y; + float tmp = roundf(z * inv_out_scale + out_zero_point_i32); + out_data[i] = + static_cast(std::max(std::min(tmp, max_val), min_val)); + } + } /* if Y is a scalar Tensor */ + else if (Y_numel == 1) { + float y = dequantize(Y_data[0], Y_scale_f, Y_zero_point_i32); + for (size_t i = 0; i < X_numel; ++i) { + float x = dequantize(X_data[i], X_scale_f, X_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } /* if X is a scalar Tensor */ + else if (X_numel == 1) { + float x = dequantize(X_data[0], X_scale_f, X_zero_point_i32); + for (size_t i = 0; i < Y_numel; ++i) { + float y = dequantize(Y_data[i], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } /* other broadcasting cases */ + else { + /* Broadcasting implementation */ + ssize_t X_dim = X.dim(); + ssize_t Y_dim = Y.dim(); + ssize_t out_dim = out.dim(); + + /* Precompute strides for X and Y tensors */ + constexpr size_t max_dim = executorch::runtime::kTensorDimensionLimit; + size_t X_strides[max_dim] = {0}; + size_t Y_strides[max_dim] = {0}; + size_t X_stride_val = 1; + size_t Y_stride_val = 1; + + /* Calculate strides from last dimension to first */ + for (int d = out_dim - 1; d >= 0 && d >= out_dim - max_dim; --d) { + int idx = out_dim - 1 - d; /* Index into the fixed-size array */ + if (d >= out_dim - X_dim) { + size_t X_d = d - (out_dim - X_dim); + X_strides[idx] = X_stride_val; + X_stride_val *= X.size(X_d); + } + + if (d >= out_dim - Y_dim) { + size_t Y_d = d - (out_dim - Y_dim); + Y_strides[idx] = Y_stride_val; + Y_stride_val *= Y.size(Y_d); + } + } + + /* Iterate over output tensor */ + for (ssize_t i = 0; i < out_numel; ++i) { + size_t out_idx = i; + size_t X_idx = 0; + size_t Y_idx = 0; + + /* Compute corresponding indices in input tensors */ + for (int d = out_dim - 1; d >= 0; --d) { + size_t out_dim_idx = out_idx % out.size(d); + out_idx /= out.size(d); + + /* Compute X index */ + if (d >= out_dim - X_dim) { + size_t X_d = d - (out_dim - X_dim); + size_t X_dim_idx = out_dim_idx % X.size(X_d); + if (d >= out_dim - max_dim) { + int idx = out_dim - 1 - d; + X_idx += X_dim_idx * X_strides[idx]; + } else { + size_t X_stride = 1; + for (int k = out_dim - 1; k > d; --k) { + if (k >= out_dim - X_dim) { + size_t X_k = k - (out_dim - X_dim); + X_stride *= X.size(X_k); + } + } + X_idx += X_dim_idx * X_stride; + } + } + + /* Compute Y index */ + if (d >= out_dim - Y_dim) { + size_t Y_d = d - (out_dim - Y_dim); + size_t Y_dim_idx = out_dim_idx % Y.size(Y_d); + if (d >= out_dim - max_dim) { + int idx = out_dim - 1 - d; + Y_idx += Y_dim_idx * Y_strides[idx]; + } else { + size_t Y_stride = 1; + for (int k = out_dim - 1; k > d; --k) { + if (k >= out_dim - Y_dim) { + size_t Y_k = k - (out_dim - Y_dim); + Y_stride *= Y.size(Y_k); + } + } + Y_idx += Y_dim_idx * Y_stride; + } + } + } + + /* Apply the operation */ + float x = dequantize(X_data[X_idx], X_scale_f, X_zero_point_i32); + float y = dequantize(Y_data[Y_idx], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_add_asym8uxasym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_add_asym8uxasym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..60f4a98ec76 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_add_asym8uxasym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,172 @@ +/* + * 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 + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using ::impl::reference::kernels::dequantize; +using ::impl::reference::kernels::quantize; + +void quantized_add_asym8uxasym8u_asym8u_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + const uint8_t* __restrict__ X_data = X.const_data_ptr(); + const uint8_t* __restrict__ Y_data = Y.const_data_ptr(); + uint8_t* __restrict__ out_data = out.mutable_data_ptr(); + + ssize_t Y_numel = Y.numel(); + ssize_t X_numel = X.numel(); + ssize_t out_numel = out.numel(); + + float X_scale_f = static_cast(X_scale); + float Y_scale_f = static_cast(Y_scale); + float out_scale_f = static_cast(out_scale); + int32_t X_zero_point_i32 = static_cast(X_zero_point); + int32_t Y_zero_point_i32 = static_cast(Y_zero_point); + int32_t out_zero_point_i32 = static_cast(out_zero_point); + + float inv_out_scale = 1.0f / out_scale_f; + constexpr float min_val = + static_cast(std::numeric_limits::min()); + constexpr float max_val = + static_cast(std::numeric_limits::max()); + + /* Tensor X exactly matches Y in shape, no broadcasting */ + if (X_numel == Y_numel && Y_numel == out_numel) { + for (size_t i = 0; i < X_numel; ++i) { + float x = X_scale_f * (X_data[i] - X_zero_point_i32); + float y = Y_scale_f * (Y_data[i] - Y_zero_point_i32); + float z = x + y; + float tmp = roundf(z * inv_out_scale + out_zero_point_i32); + out_data[i] = + static_cast(std::max(std::min(tmp, max_val), min_val)); + } + } /* if Y is a scalar Tensor */ + else if (Y_numel == 1) { + float y = dequantize(Y_data[0], Y_scale_f, Y_zero_point_i32); + for (size_t i = 0; i < X_numel; ++i) { + float x = dequantize(X_data[i], X_scale_f, X_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } /* if X is a scalar Tensor */ + else if (X_numel == 1) { + float x = dequantize(X_data[0], X_scale_f, X_zero_point_i32); + for (size_t i = 0; i < Y_numel; ++i) { + float y = dequantize(Y_data[i], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } /* other broadcasting cases */ + else { + /* Broadcasting implementation */ + ssize_t X_dim = X.dim(); + ssize_t Y_dim = Y.dim(); + ssize_t out_dim = out.dim(); + + /* Precompute strides for X and Y tensors */ + constexpr size_t max_dim = executorch::runtime::kTensorDimensionLimit; + size_t X_strides[max_dim] = {0}; + size_t Y_strides[max_dim] = {0}; + size_t X_stride_val = 1; + size_t Y_stride_val = 1; + + /* Calculate strides from last dimension to first */ + for (int d = out_dim - 1; d >= 0 && d >= out_dim - max_dim; --d) { + int idx = out_dim - 1 - d; /* Index into the fixed-size array */ + if (d >= out_dim - X_dim) { + size_t X_d = d - (out_dim - X_dim); + X_strides[idx] = X_stride_val; + X_stride_val *= X.size(X_d); + } + + if (d >= out_dim - Y_dim) { + size_t Y_d = d - (out_dim - Y_dim); + Y_strides[idx] = Y_stride_val; + Y_stride_val *= Y.size(Y_d); + } + } + + /* Iterate over output tensor */ + for (ssize_t i = 0; i < out_numel; ++i) { + size_t out_idx = i; + size_t X_idx = 0; + size_t Y_idx = 0; + + /* Compute corresponding indices in input tensors */ + for (int d = out_dim - 1; d >= 0; --d) { + size_t out_dim_idx = out_idx % out.size(d); + out_idx /= out.size(d); + + /* Compute X index */ + if (d >= out_dim - X_dim) { + size_t X_d = d - (out_dim - X_dim); + size_t X_dim_idx = out_dim_idx % X.size(X_d); + if (d >= out_dim - max_dim) { + int idx = out_dim - 1 - d; + X_idx += X_dim_idx * X_strides[idx]; + } else { + size_t X_stride = 1; + for (int k = out_dim - 1; k > d; --k) { + if (k >= out_dim - X_dim) { + size_t X_k = k - (out_dim - X_dim); + X_stride *= X.size(X_k); + } + } + X_idx += X_dim_idx * X_stride; + } + } + + /* Compute Y index */ + if (d >= out_dim - Y_dim) { + size_t Y_d = d - (out_dim - Y_dim); + size_t Y_dim_idx = out_dim_idx % Y.size(Y_d); + if (d >= out_dim - max_dim) { + int idx = out_dim - 1 - d; + Y_idx += Y_dim_idx * Y_strides[idx]; + } else { + size_t Y_stride = 1; + for (int k = out_dim - 1; k > d; --k) { + if (k >= out_dim - Y_dim) { + size_t Y_k = k - (out_dim - Y_dim); + Y_stride *= Y.size(Y_k); + } + } + Y_idx += Y_dim_idx * Y_stride; + } + } + } + + /* Apply the operation */ + float x = dequantize(X_data[X_idx], X_scale_f, X_zero_point_i32); + float y = dequantize(Y_data[Y_idx], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..2788de589cf --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,248 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Optimized NCHW convolution for int8 x int8 -> int8 +void xa_opt_quantized_conv_nchw_asym8sxsym8s_asym8s( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + ET_CHECK_MSG(groups == 1, "Only groups=1 supported for regular convolution"); + WORD32 out_data_format = 1; + + WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * input_channels * input_height * input_width) + 8) * + sizeof(WORD8)); + + WORD8* ptr2 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((out_channels * kernel_channels * kernel_height * kernel_width) + 8) * + sizeof(WORD8)); + + WORD8* pin = (WORD8*)ALIGN_PTR(ptr1, 8); + WORD8* pkernel = (WORD8*)ALIGN_PTR(ptr2, 8); + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = input.size(0); + p_inp_shape[1] = input_channels; + p_inp_shape[2] = input_height; + p_inp_shape[3] = input_width; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = input.size(0); + p_out_shape[1] = input_height; + p_out_shape[2] = input_width; + p_out_shape[3] = input_channels; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 2, 3, 1}; + + xa_nn_transpose_8_8( + pin, + p_out_shape, + p_inp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); + + WORD32 p_inp_shape1[kNnlibMaxDim]; + p_inp_shape1[0] = out_channels; + p_inp_shape1[1] = kernel_channels; + p_inp_shape1[2] = kernel_height; + p_inp_shape1[3] = kernel_width; + + WORD32 p_out_shape1[kNnlibMaxDim]; + p_out_shape1[0] = out_channels; + p_out_shape1[1] = kernel_height; + p_out_shape1[2] = kernel_width; + p_out_shape1[3] = kernel_channels; + + xa_nn_transpose_8_8( + pkernel, + p_out_shape1, + p_kernel, + p_inp_shape1, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = pin + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + out_batch, + in_batch, + pkernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); + } +} + +void quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nchw_asym8sxsym8s_asym8s( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..9fd2d69dda9 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,248 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Optimized NCHW convolution for uint8 x uint8 -> uint8 +void xa_opt_quantized_conv_nchw_asym8uxsym8u_asym8u( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + UWORD8* __restrict__ p_out = + (UWORD8* __restrict__)out.mutable_data_ptr(); + UWORD8* __restrict__ p_inp = + (UWORD8* __restrict__)input.const_data_ptr(); + UWORD8* __restrict__ p_kernel = + (UWORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + ET_CHECK_MSG(groups == 1, "Only groups=1 supported for regular convolution"); + WORD32 out_data_format = 1; + + UWORD8* ptr1 = (UWORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * input_channels * input_height * input_width) + 8) * + sizeof(UWORD8)); + + UWORD8* ptr2 = (UWORD8*)kernels::allocate_temp_memory( + ctx, + ((out_channels * kernel_channels * kernel_height * kernel_width) + 8) * + sizeof(UWORD8)); + + UWORD8* pin = (UWORD8*)ALIGN_PTR(ptr1, 8); + UWORD8* pkernel = (UWORD8*)ALIGN_PTR(ptr2, 8); + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = input.size(0); + p_inp_shape[1] = input_channels; + p_inp_shape[2] = input_height; + p_inp_shape[3] = input_width; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = input.size(0); + p_out_shape[1] = input_height; + p_out_shape[2] = input_width; + p_out_shape[3] = input_channels; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 2, 3, 1}; + + xa_nn_transpose_8_8( + (WORD8*)pin, + p_out_shape, + (WORD8*)p_inp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); + + WORD32 p_inp_shape1[kNnlibMaxDim]; + p_inp_shape1[0] = out_channels; + p_inp_shape1[1] = kernel_channels; + p_inp_shape1[2] = kernel_height; + p_inp_shape1[3] = kernel_width; + + WORD32 p_out_shape1[kNnlibMaxDim]; + p_out_shape1[0] = out_channels; + p_out_shape1[1] = kernel_height; + p_out_shape1[2] = kernel_width; + p_out_shape1[3] = kernel_channels; + + xa_nn_transpose_8_8( + (WORD8*)pkernel, + p_out_shape1, + (WORD8*)p_kernel, + p_inp_shape1, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + UWORD8* in_batch = pin + _n * input_channels * input_height * input_width; + UWORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + (WORD8*)out_batch, + (WORD8*)in_batch, + (WORD8*)pkernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); + } +} + +void quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nchw_asym8uxsym8u_asym8u( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..3e2c9c58401 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,203 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Specialized depthwise NCHW convolution for int8 x int8 -> int8 +void xa_opt_quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + + WORD32 channels_multiplier = out_channels / input_channels; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 1); // NCHW + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + WORD32* ptr_scratch = + (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + pVOID p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * out_channels * out_height * out_width) + 8) * sizeof(WORD8)); + + WORD8* p_out_temp = (WORD8*)ALIGN_PTR(ptr1, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out_temp + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + out_batch, + p_kernel, + in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 1, // NCHW + 0, // NHWC + p_scratch); + } + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = batches; + p_inp_shape[1] = out_height; + p_inp_shape[2] = out_width; + p_inp_shape[3] = out_channels; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = batches; + p_out_shape[1] = out_channels; + p_out_shape[2] = out_height; + p_out_shape[3] = out_width; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 3, 1, 2}; + + xa_nn_transpose_8_8( + p_out, + p_out_shape, + p_out_temp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); +} + +void quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..103ce9568c5 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,203 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Specialized depthwise NCHW convolution for uint8 x uint8 -> uint8 +void xa_opt_quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + UWORD8* __restrict__ p_out = + (UWORD8* __restrict__)out.mutable_data_ptr(); + UWORD8* __restrict__ p_inp = + (UWORD8* __restrict__)input.const_data_ptr(); + UWORD8* __restrict__ p_kernel = + (UWORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + + WORD32 channels_multiplier = out_channels / input_channels; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 1); // NCHW + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + WORD32* ptr_scratch = + (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + pVOID p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + UWORD8* ptr1 = (UWORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * out_channels * out_height * out_width) + 8) * sizeof(UWORD8)); + + UWORD8* p_out_temp = (UWORD8*)ALIGN_PTR(ptr1, 8); + + for (int _n = 0; _n < batches; _n++) { + UWORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + UWORD8* out_batch = p_out_temp + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + (WORD8*)out_batch, + (WORD8*)p_kernel, + (WORD8*)in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 1, // NCHW + 0, // NHWC + p_scratch); + } + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = batches; + p_inp_shape[1] = out_height; + p_inp_shape[2] = out_width; + p_inp_shape[3] = out_channels; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = batches; + p_out_shape[1] = out_channels; + p_out_shape[2] = out_height; + p_out_shape[3] = out_width; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 3, 1, 2}; + + xa_nn_transpose_8_8( + (WORD8*)p_out, + p_out_shape, + (WORD8*)p_out_temp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, + kNnlibMaxDim); +} + +void quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..cdc1ecd8526 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,190 @@ +/* + * 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 + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Dilated fallback implementation for int8 x int8 -> int8 quantized 2d conv +// kernel for NCHW layout. This variant is optimized for asymmetric int8 inputs, +// weights, and outputs. The input is of shape [n x c x h x w] The weight is of +// shape [oc x wc x wh x ww], where wc == c The output is of shape [n x oc x oh +// x ow] The bias is of shape [oc] +template +__attribute__((noinline)) void conv2d_nchw_dilated_asym8sxsym8s_asym8s_core( + // All the arrays + const int8_t* __restrict__ p_in, + const int8_t* __restrict__ p_weight, + const int32_t* __restrict__ p_bias, + int8_t* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t c, + int32_t h, + int32_t w, + int32_t oc, + int32_t wc, + int32_t wh, + int32_t ww, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Quantization parameters + int8_t in_zero_point = 0, + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + int8_t out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const int8_t* in_batch = p_in + _n * c * h * w; + int8_t* out_batch = p_out + _n * oc * oh * ow; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + int8_t* out_plane = out_batch + _oc * oh * ow; + const int8_t* weight_batch = p_weight + _oc * wc * wh * ww; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of size + // icpg x h x w, with a stencil of size icpg x wh x ww, to compute an + // output channel of size 1 x oh x ow. + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to the + // output channel being computed) with the corresponding weight + // channel. + // General path for dilated convolutions with padding support + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + const int8_t* in_plane = in_batch + _ic * h * w; + const int8_t* weight_plane = weight_batch + (_ic - sic) * wh * ww; + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + int input_h = _h + d0 * _wh - p0; + int input_w = _w + d1 * _ww - p1; + if ((input_h >= 0) && (input_h < h) && (input_w >= 0) && + (input_w < w)) { + int ioff = input_h * w + input_w; + int woff = _wh * ww + _ww; + float lhs = static_cast(in_plane[ioff]) - + static_cast(in_zero_point); + float rhs = static_cast(weight_plane[woff]) - + static_cast(weight_zero_point); + acc += lhs * rhs; + } + } + } + } + // Quantize the accumulated result + float val = bias_scale * acc; + out_plane[_oh * ow + _ow] = + kernels::quantize(val, inv_out_scale, out_zero_point); + } + } + } + } + } +} + +void quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, c, h, w] + const int n = input.size(0); + const int c = input.size(1); + const int h = conv1d ? 1 : input.size(2); + const int w = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wc, wh, ww] + const int oc = weight.size(0); + const int wc = weight.size(1); + const int wh = conv1d ? 1 : weight.size(2); + const int ww = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oc, oh, ow] + const int oh = conv1d ? 1 : out.size(2); + const int ow = conv1d ? out.size(2) : out.size(3); + + conv2d_nchw_dilated_asym8sxsym8s_asym8s_core( + input.const_data_ptr(), + weight.const_data_ptr(), + bias.const_data_ptr(), + out.mutable_data_ptr(), + n, + c, + h, + w, + oc, + wc, + wh, + ww, + oh, + ow, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + groups, + static_cast(in_zero_point), + weight_zero_point, + bias_scale, + output_scale, + static_cast(output_zero_point)); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..9281dcea496 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,191 @@ +/* + * 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 + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Dilated fallback implementation for uint8 x uint8 -> uint8 quantized 2d conv +// kernel for NCHW layout. This variant is optimized for asymmetric uint8 +// inputs, weights, and outputs. The input is of shape [n x c x h x w] The +// weight is of shape [oc x wc x wh x ww], where wc == c The output is of shape +// [n x oc x oh x ow] The bias is of shape [oc] +template +__attribute__((noinline)) void conv2d_nchw_dilated_asym8uxsym8u_asym8u_core( + // All the arrays + const uint8_t* __restrict__ p_in, + const uint8_t* __restrict__ p_weight, + const int32_t* __restrict__ p_bias, + uint8_t* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t c, + int32_t h, + int32_t w, + int32_t oc, + int32_t wc, + int32_t wh, + int32_t ww, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Quantization parameters + uint8_t in_zero_point = 0, + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + uint8_t out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const uint8_t* in_batch = p_in + _n * c * h * w; + uint8_t* out_batch = p_out + _n * oc * oh * ow; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + uint8_t* out_plane = out_batch + _oc * oh * ow; + const uint8_t* weight_batch = p_weight + _oc * wc * wh * ww; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of size + // icpg x h x w, with a stencil of size icpg x wh x ww, to compute an + // output channel of size 1 x oh x ow. + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to the + // output channel being computed) with the corresponding weight + // channel. + // General path for dilated convolutions with padding support + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + const uint8_t* in_plane = in_batch + _ic * h * w; + const uint8_t* weight_plane = + weight_batch + (_ic - sic) * wh * ww; + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + int input_h = _h + d0 * _wh - p0; + int input_w = _w + d1 * _ww - p1; + if ((input_h >= 0) && (input_h < h) && (input_w >= 0) && + (input_w < w)) { + int ioff = input_h * w + input_w; + int woff = _wh * ww + _ww; + float lhs = static_cast(in_plane[ioff]) - + static_cast(in_zero_point); + float rhs = static_cast(weight_plane[woff]) - + static_cast(weight_zero_point); + acc += lhs * rhs; + } + } + } + } + // Quantize the accumulated result + float val = bias_scale * acc; + out_plane[_oh * ow + _ow] = + kernels::quantize(val, inv_out_scale, out_zero_point); + } + } + } + } + } +} + +void quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, c, h, w] + const int n = input.size(0); + const int c = input.size(1); + const int h = conv1d ? 1 : input.size(2); + const int w = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wc, wh, ww] + const int oc = weight.size(0); + const int wc = weight.size(1); + const int wh = conv1d ? 1 : weight.size(2); + const int ww = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oc, oh, ow] + const int oh = conv1d ? 1 : out.size(2); + const int ow = conv1d ? out.size(2) : out.size(3); + + conv2d_nchw_dilated_asym8uxsym8u_asym8u_core( + input.const_data_ptr(), + weight.const_data_ptr(), + bias.const_data_ptr(), + out.mutable_data_ptr(), + n, + c, + h, + w, + oc, + wc, + wh, + ww, + oh, + ow, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + groups, + static_cast(in_zero_point), + weight_zero_point, + bias_scale, + output_scale, + static_cast(output_zero_point)); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/reference/operators/quantized_conv_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nchw_out.cpp similarity index 56% rename from backends/cadence/reference/operators/quantized_conv_out.cpp rename to backends/cadence/hifi/operators/op_quantized_conv_nchw_out.cpp index 87ff264a258..fbc97a4c37b 100644 --- a/backends/cadence/reference/operators/quantized_conv_out.cpp +++ b/backends/cadence/hifi/operators/op_quantized_conv_nchw_out.cpp @@ -6,17 +6,22 @@ * LICENSE file in the root directory of this source tree. */ -#include -#include +#include +#include +#include +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +namespace cadence { namespace impl { -namespace reference { +namespace HiFi { namespace native { +using ::cadence::impl::HiFi::kernels::quantize; using ::executorch::aten::IntArrayRef; using ::executorch::aten::ScalarType; using ::executorch::aten::Tensor; -using ::executorch::runtime::KernelRuntimeContext; +using ::torch::executor::KernelRuntimeContext; // This implements a generic 2d conv kernel that operates on raw pointers. // The version handles both quantized and fp32 convolutions. @@ -141,8 +146,7 @@ __attribute__((noinline)) void conv2d_nchw_core_generic( if (quantized) { float val = bias_scale * acc; out_plane[_oh * ow + _ow] = - ::impl::reference::kernels::quantize( - val, inv_out_scale, out_zero_point); + quantize(val, inv_out_scale, out_zero_point); } else { out_plane[_oh * ow + _ow] = acc; } @@ -153,128 +157,286 @@ __attribute__((noinline)) void conv2d_nchw_core_generic( } } -template < - typename IT = float, - typename WT = IT, - typename BT = IT, - typename OT = IT, - bool quantized = false> -__attribute__((noinline)) void conv2d_nhwc_core_generic( - // All the arrays - const IT* __restrict__ p_in, - const WT* __restrict__ p_weight, - const BT* __restrict__ p_bias, - OT* __restrict__ p_out, - // The array sizes - int32_t n, - int32_t h, - int32_t w, - int32_t c, - int32_t oc, - int32_t wh, - int32_t ww, - int32_t wc, - int32_t oh, - int32_t ow, - // Stride - int16_t s0, - int16_t s1, - // Padding - int16_t p0, - int16_t p1, - // Dilation - int16_t d0, - int16_t d1, - // Group for depthwise conv +void xa_opt_quantized_conv_nchw( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, int16_t groups, - // Optional args that are only relevant for quantized convolution - // input zero point - IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, - float out_scale = 1, - OT out_zero_point = 0) { - float inv_out_scale = 1. / out_scale; - bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; - // Compute the number of in and out channels per group - const int ocpg = oc / groups; - const int icpg = c / groups; + if (input.scalar_type() == ScalarType::Char) { + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); - // Iterate over all the output batches (i.e., n) - for (int _n = 0; _n < n; ++_n) { - const IT* in_batch = p_in + _n * h * w * c; - OT* out_batch = p_out + _n * oh * ow * oc; - for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { - for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { - OT* out_line = out_batch + (_oh * ow + _ow) * oc; - // Compute separable convolution for each group - for (int _g = 0; _g < groups; ++_g) { - // Identify the input and output channels involved in the computation - // of this group - int sic = _g * icpg; - int soc = _g * ocpg; - // Populate all the output channels in the group - for (int _oc = soc; _oc < soc + ocpg; ++_oc) { - const WT* weight_batch = p_weight + _oc * wh * ww * wc; - // We compute one output channel at a time. The computation can be - // thought of as a stencil computation: we iterate over an input of - // size h x w x icpg, with a stencil of size wh x ww x icpg, to - // compute an output channel of size oh x ow x 1. - float acc = p_bias[_oc]; - // Below is the stencil computation that performs the hadamard - // product+accumulation of each input channel (contributing to - // the output channel being computed) with the corresponding - // weight channel. If the padding is 0, and dilation is 1, then - // we can remove the unnecessary checks, and simplify the code - // so that it can be vectorized by Tensilica compiler.x`` - if (zero_pad_unit_dilation) { - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - const IT* in_line = - in_batch + (_h + _wh) * w * c + (_w + _ww) * c; - const WT* weight_line = - weight_batch + _wh * ww * wc + _ww * wc; - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - float lhs = in_line[_ic] - in_zero_point; - float rhs = weight_line[_ic - sic] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } else { - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - if (((_h + d0 * _wh - p0) >= 0) && - ((_h + d0 * _wh - p0) < h) && - ((_w + d1 * _ww - p1) >= 0) && - ((_w + d1 * _ww - p1 < w))) { - const IT* in_line = in_batch + - (_h + d0 * _wh - p0) * w * c + (_w + d1 * _ww - p1) * c; - const WT* weight_line = - weight_batch + _wh * ww * wc + _ww * wc; - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - float lhs = in_line[_ic] - in_zero_point; - float rhs = weight_line[_ic - sic] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } - } - if (quantized) { - float val = bias_scale * acc; - out_line[_oc] = ::impl::reference::kernels::quantize( - val, inv_out_scale, out_zero_point); - } else { - out_line[_oc] = acc; - } - } - } + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + // WORD32* kernel_bias_ptr = + // (WORD32*)weight_zero_point.const_data_ptr(); + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + if (groups == 1) { + WORD32 out_data_format = 1; + + WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * input_channels * input_height * input_width) + 8) * + sizeof(WORD8)); + + WORD8* ptr2 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((out_channels * kernel_channels * kernel_height * kernel_width) + + 8) * + sizeof(WORD8)); + + WORD8* pin = (WORD8*)ALIGN_PTR(ptr1, 8); + WORD8* pkernel = (WORD8*)ALIGN_PTR(ptr2, 8); + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = input.size(0); + p_inp_shape[1] = input_channels; + p_inp_shape[2] = input_height; + p_inp_shape[3] = input_width; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = input.size(0); + p_out_shape[1] = input_height; + p_out_shape[2] = input_width; + p_out_shape[3] = input_channels; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 2, 3, 1}; + + xa_nn_transpose_8_8( + pin, + p_out_shape, + p_inp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, // input dimensions + kNnlibMaxDim); // output dimensions + + WORD32 p_inp_shape1[kNnlibMaxDim]; + p_inp_shape1[0] = out_channels; + p_inp_shape1[1] = kernel_channels; + p_inp_shape1[2] = kernel_height; + p_inp_shape1[3] = kernel_width; + + WORD32 p_out_shape1[kNnlibMaxDim]; + p_out_shape1[0] = out_channels; + p_out_shape1[1] = kernel_height; + p_out_shape1[2] = kernel_width; + p_out_shape1[3] = kernel_channels; + + xa_nn_transpose_8_8( + pkernel, + p_out_shape1, + p_kernel, + p_inp_shape1, + p_permute_vec, + kNnlibMaxDim, // input dimensions + kNnlibMaxDim); // output dimensions + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = + pin + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + out_batch, + in_batch, + pkernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); } + return; + } + + if (groups == input_channels) { + WORD32 channels_multiplier = out_channels / input_channels; + + scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 1); // NCHW + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * out_channels * out_height * out_width) + 8) * + sizeof(WORD8)); + + WORD8* p_out_temp = (WORD8*)ALIGN_PTR(ptr1, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = + p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = + p_out_temp + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + out_batch, + p_kernel, + in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 1, // NCHW + 0, // NHWC + p_scratch); + } + + WORD32 p_inp_shape[kNnlibMaxDim]; + p_inp_shape[0] = batches; + p_inp_shape[1] = out_height; + p_inp_shape[2] = out_width; + p_inp_shape[3] = out_channels; + + WORD32 p_out_shape[kNnlibMaxDim]; + p_out_shape[0] = batches; + p_out_shape[1] = out_channels; + p_out_shape[2] = out_height; + p_out_shape[3] = out_width; + + WORD32 p_permute_vec[kNnlibMaxDim] = {0, 3, 1, 2}; + + xa_nn_transpose_8_8( + p_out, + p_out_shape, + p_out_temp, + p_inp_shape, + p_permute_vec, + kNnlibMaxDim, // input dimensions + kNnlibMaxDim); // output dimensions + + return; } } } @@ -354,78 +516,7 @@ void quantized_conv_nchw( #undef typed_quantized_conv2d_nchw } -void quantized_conv_nhwc( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int16_t groups, - int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, - float output_scale, - int32_t output_zero_point, - Tensor& out) { - bool conv1d = input.dim() == 3; - // input = [n, h, w, c] - const int n = input.size(0); - const int h = conv1d ? 1 : input.size(1); - const int w = conv1d ? input.size(1) : input.size(2); - const int c = conv1d ? input.size(2) : input.size(3); - // weight = [oc, wh, ww, wc] - const int oc = weight.size(0); - const int wh = conv1d ? 1 : weight.size(1); - const int ww = conv1d ? weight.size(1) : weight.size(2); - const int wc = conv1d ? weight.size(2) : weight.size(3); - // output = [n, oh, ow, oc] - const int oh = conv1d ? 1 : out.size(1); - const int ow = conv1d ? out.size(1) : out.size(2); - -#define typed_quantized_conv2d_nhwc(ctype, dtype) \ - case ScalarType::dtype: { \ - conv2d_nhwc_core_generic( \ - input.const_data_ptr(), \ - weight.const_data_ptr(), \ - bias.const_data_ptr(), \ - out.mutable_data_ptr(), \ - n, \ - h, \ - w, \ - c, \ - oc, \ - wh, \ - ww, \ - wc, \ - oh, \ - ow, \ - stride[0], \ - stride[1], \ - padding[0], \ - padding[1], \ - dilation[0], \ - dilation[1], \ - groups, \ - in_zero_point, \ - weight_zero_point, \ - bias_scale, \ - output_scale, \ - (ctype)output_zero_point); \ - break; \ - } - ScalarType dtype = out.scalar_type(); - switch (dtype) { - ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nhwc); - default: - ET_DCHECK_MSG( - false, "Unhandled dtype %s", torch::executor::toString(dtype)); - } - -#undef typed_quantized_conv2d_nhwc -} - -void quantized_conv_out( +void quantized_conv_nchw_out( __ET_UNUSED KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -441,13 +532,23 @@ void quantized_conv_out( int64_t output_zero_point, __ET_UNUSED const Tensor& out_multiplier, __ET_UNUSED const Tensor& out_shift, - bool channel_last, Tensor& out) { const float bias_scale_float = bias_scale.const_data_ptr()[0]; const int32_t weight_zero_point_int = weight_zero_point.const_data_ptr()[0]; - if (channel_last) { - quantized_conv_nhwc( + + bool optimized = 0; + + if ((input.scalar_type() == ScalarType::Char) || + (input.scalar_type() == ScalarType::Byte)) + optimized = 1; + + if ((dilation[0] != 1) || (dilation[1] != 1)) + optimized = 0; + + if (optimized) { + xa_opt_quantized_conv_nchw( + ctx, input, weight, bias, @@ -479,7 +580,7 @@ void quantized_conv_out( } } -void quantized_conv_per_tensor_out( +void quantized_conv_nchw_per_tensor_out( __ET_UNUSED KernelRuntimeContext& ctx, const Tensor& input, const Tensor& weight, @@ -495,10 +596,19 @@ void quantized_conv_per_tensor_out( int64_t output_zero_point, __ET_UNUSED int64_t out_multiplier, __ET_UNUSED int64_t out_shift, - bool channel_last, Tensor& out) { - if (channel_last) { - quantized_conv_nhwc( + bool optimized = 0; + + if ((input.scalar_type() == ScalarType::Char) || + (input.scalar_type() == ScalarType::Byte)) + optimized = 1; + + if ((dilation[0] != 1) || (dilation[1] != 1)) + optimized = 0; + + if (optimized) { + xa_opt_quantized_conv_nchw( + ctx, input, weight, bias, @@ -531,5 +641,6 @@ void quantized_conv_per_tensor_out( } } // namespace native -} // namespace reference +} // namespace HiFi } // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..b1e023736cf --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,191 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Optimized NHWC convolution for int8 x int8 -> int8 +void xa_opt_quantized_conv_nhwc_asym8sxsym8s_asym8s( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + ET_CHECK_MSG(groups == 1, "Only groups=1 supported for regular convolution"); + WORD32 out_data_format = 1; + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + out_batch, + in_batch, + p_kernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); + } +} + +void quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nhwc_asym8sxsym8s_asym8s( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..0678cb1b821 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,191 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Optimized NHWC convolution for uint8 x uint8 -> uint8 +void xa_opt_quantized_conv_nhwc_asym8uxsym8u_asym8u( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + UWORD8* __restrict__ p_out = + (UWORD8* __restrict__)out.mutable_data_ptr(); + UWORD8* __restrict__ p_inp = + (UWORD8* __restrict__)input.const_data_ptr(); + UWORD8* __restrict__ p_kernel = + (UWORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + ET_CHECK_MSG(groups == 1, "Only groups=1 supported for regular convolution"); + WORD32 out_data_format = 1; + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + UWORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + UWORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + (WORD8*)out_batch, + (WORD8*)in_batch, + (WORD8*)p_kernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); + } +} + +void quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nhwc_asym8uxsym8u_asym8u( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..6512622f221 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,173 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Specialized depthwise NHWC convolution for int8 x int8 -> int8 +void xa_opt_quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + + WORD32 channels_multiplier = out_channels / input_channels; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 0); // NHWC + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + WORD32* ptr_scratch = + (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + pVOID p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + out_batch, + p_kernel, + in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 0, // NHWC + 0, // NHWC + p_scratch); + } +} + +void quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..d41a9c8d4b7 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,173 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Specialized depthwise NHWC convolution for uint8 x uint8 -> uint8 +void xa_opt_quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + + UWORD8* __restrict__ p_out = + (UWORD8* __restrict__)out.mutable_data_ptr(); + UWORD8* __restrict__ p_inp = + (UWORD8* __restrict__)input.const_data_ptr(); + UWORD8* __restrict__ p_kernel = + (UWORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + + WORD32 input_zero_bias = -in_zero_point; + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + + WORD32 channels_multiplier = out_channels / input_channels; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 0); // NHWC + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + WORD32* ptr_scratch = + (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + pVOID p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + UWORD8* in_batch = p_inp + _n * input_channels * input_height * input_width; + UWORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + (WORD8*)out_batch, + (WORD8*)p_kernel, + (WORD8*)in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 0, // NHWC + 0, // NHWC + p_scratch); + } +} + +void quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + xa_opt_quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..be661334acf --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,190 @@ +/* + * 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 + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Dilated fallback implementation for int8 x int8 -> int8 quantized 2d conv +// kernel for NHWC layout. This variant is optimized for asymmetric int8 inputs, +// weights, and outputs. The input is of shape [n x h x w x c] The weight is of +// shape [oc x wh x ww x wc] The output is of shape [n x oh x ow x oc] The bias +// is of shape [oc] +template +__attribute__((noinline)) void conv2d_nhwc_dilated_asym8sxsym8s_asym8s_core( + // All the arrays + const int8_t* __restrict__ p_in, + const int8_t* __restrict__ p_weight, + const int32_t* __restrict__ p_bias, + int8_t* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t h, + int32_t w, + int32_t c, + int32_t oc, + int32_t wh, + int32_t ww, + int32_t wc, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Quantization parameters + int8_t in_zero_point = 0, + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + int8_t out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const int8_t* in_batch = p_in + _n * h * w * c; + int8_t* out_batch = p_out + _n * oh * ow * oc; + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + int8_t* out_line = out_batch + (_oh * ow + _ow) * oc; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + const int8_t* weight_batch = p_weight + _oc * wh * ww * wc; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of + // size h x w x icpg, with a stencil of size wh x ww x icpg, to + // compute an output channel of size oh x ow x 1. + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to + // the output channel being computed) with the corresponding + // weight channel. + // General path for dilated convolutions with padding support + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + int input_h = _h + d0 * _wh - p0; + int input_w = _w + d1 * _ww - p1; + if ((input_h >= 0) && (input_h < h) && (input_w >= 0) && + (input_w < w)) { + const int8_t* in_line = + in_batch + input_h * w * c + input_w * c; + const int8_t* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = static_cast(in_line[_ic]) - + static_cast(in_zero_point); + float rhs = static_cast(weight_line[_ic - sic]) - + static_cast(weight_zero_point); + acc += lhs * rhs; + } + } + } + } + // Quantize the accumulated result + float val = bias_scale * acc; + out_line[_oc] = + kernels::quantize(val, inv_out_scale, out_zero_point); + } + } + } + } + } +} + +void quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, h, w, c] + const int n = input.size(0); + const int h = conv1d ? 1 : input.size(1); + const int w = conv1d ? input.size(1) : input.size(2); + const int c = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wh, ww, wc] + const int oc = weight.size(0); + const int wh = conv1d ? 1 : weight.size(1); + const int ww = conv1d ? weight.size(1) : weight.size(2); + const int wc = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oh, ow, oc] + const int oh = conv1d ? 1 : out.size(1); + const int ow = conv1d ? out.size(1) : out.size(2); + + conv2d_nhwc_dilated_asym8sxsym8s_asym8s_core( + input.const_data_ptr(), + weight.const_data_ptr(), + bias.const_data_ptr(), + out.mutable_data_ptr(), + n, + h, + w, + c, + oc, + wh, + ww, + wc, + oh, + ow, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + groups, + static_cast(in_zero_point), + weight_zero_point, + bias_scale, + output_scale, + static_cast(output_zero_point)); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..cab4897f5f0 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,190 @@ +/* + * 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 + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +// Dilated fallback implementation for uint8 x uint8 -> uint8 quantized 2d conv +// kernel for NHWC layout. This variant is optimized for asymmetric uint8 +// inputs, weights, and outputs. The input is of shape [n x h x w x c] The +// weight is of shape [oc x wh x ww x wc] The output is of shape [n x oh x ow x +// oc] The bias is of shape [oc] +template +__attribute__((noinline)) void conv2d_nhwc_dilated_asym8uxsym8u_asym8u_core( + // All the arrays + const uint8_t* __restrict__ p_in, + const uint8_t* __restrict__ p_weight, + const int32_t* __restrict__ p_bias, + uint8_t* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t h, + int32_t w, + int32_t c, + int32_t oc, + int32_t wh, + int32_t ww, + int32_t wc, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Quantization parameters + uint8_t in_zero_point = 0, + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + uint8_t out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const uint8_t* in_batch = p_in + _n * h * w * c; + uint8_t* out_batch = p_out + _n * oh * ow * oc; + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + uint8_t* out_line = out_batch + (_oh * ow + _ow) * oc; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + const uint8_t* weight_batch = p_weight + _oc * wh * ww * wc; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of + // size h x w x icpg, with a stencil of size wh x ww x icpg, to + // compute an output channel of size oh x ow x 1. + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to + // the output channel being computed) with the corresponding + // weight channel. + // General path for dilated convolutions with padding support + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + int input_h = _h + d0 * _wh - p0; + int input_w = _w + d1 * _ww - p1; + if ((input_h >= 0) && (input_h < h) && (input_w >= 0) && + (input_w < w)) { + const uint8_t* in_line = + in_batch + input_h * w * c + input_w * c; + const uint8_t* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = static_cast(in_line[_ic]) - + static_cast(in_zero_point); + float rhs = static_cast(weight_line[_ic - sic]) - + static_cast(weight_zero_point); + acc += lhs * rhs; + } + } + } + } + // Quantize the accumulated result + float val = bias_scale * acc; + out_line[_oc] = + kernels::quantize(val, inv_out_scale, out_zero_point); + } + } + } + } + } +} + +void quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, h, w, c] + const int n = input.size(0); + const int h = conv1d ? 1 : input.size(1); + const int w = conv1d ? input.size(1) : input.size(2); + const int c = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wh, ww, wc] + const int oc = weight.size(0); + const int wh = conv1d ? 1 : weight.size(1); + const int ww = conv1d ? weight.size(1) : weight.size(2); + const int wc = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oh, ow, oc] + const int oh = conv1d ? 1 : out.size(1); + const int ow = conv1d ? out.size(1) : out.size(2); + + conv2d_nhwc_dilated_asym8uxsym8u_asym8u_core( + input.const_data_ptr(), + weight.const_data_ptr(), + bias.const_data_ptr(), + out.mutable_data_ptr(), + n, + h, + w, + c, + oc, + wh, + ww, + wc, + oh, + ow, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + groups, + static_cast(in_zero_point), + weight_zero_point, + bias_scale, + output_scale, + static_cast(output_zero_point)); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_nhwc_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_out.cpp new file mode 100644 index 00000000000..8af7c0da3ef --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_conv_nhwc_out.cpp @@ -0,0 +1,552 @@ +/* + * 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 + +#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) + +using Tensor = executorch::aten::Tensor; +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; +using ScalarType = executorch::aten::ScalarType; +using ::executorch::aten::IntArrayRef; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +template < + typename IT = float, + typename WT = IT, + typename BT = IT, + typename OT = IT, + bool quantized = false> +__attribute__((noinline)) void conv2d_nhwc_core_generic( + // All the arrays + const IT* __restrict__ p_in, + const WT* __restrict__ p_weight, + const BT* __restrict__ p_bias, + OT* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t h, + int32_t w, + int32_t c, + int32_t oc, + int32_t wh, + int32_t ww, + int32_t wc, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Optional args that are only relevant for quantized convolution + // input zero point + IT in_zero_point = 0, + // weight zero point + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + OT out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const IT* in_batch = p_in + _n * h * w * c; + OT* out_batch = p_out + _n * oh * ow * oc; + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + OT* out_line = out_batch + (_oh * ow + _ow) * oc; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + const WT* weight_batch = p_weight + _oc * wh * ww * wc; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of + // size h x w x icpg, with a stencil of size wh x ww x icpg, to + // compute an output channel of size oh x ow x 1. + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to + // the output channel being computed) with the corresponding + // weight channel. If the padding is 0, and dilation is 1, then + // we can remove the unnecessary checks, and simplify the code + // so that it can be vectorized by Tensilica compiler.x`` + if (zero_pad_unit_dilation) { + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + const IT* in_line = + in_batch + (_h + _wh) * w * c + (_w + _ww) * c; + const WT* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = in_line[_ic] - in_zero_point; + float rhs = weight_line[_ic - sic] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } else { + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + if (((_h + d0 * _wh - p0) >= 0) && + ((_h + d0 * _wh - p0) < h) && + ((_w + d1 * _ww - p1) >= 0) && + ((_w + d1 * _ww - p1 < w))) { + const IT* in_line = in_batch + + (_h + d0 * _wh - p0) * w * c + (_w + d1 * _ww - p1) * c; + const WT* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = in_line[_ic] - in_zero_point; + float rhs = weight_line[_ic - sic] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } + } + if (quantized) { + float val = bias_scale * acc; + out_line[_oc] = + kernels::quantize(val, inv_out_scale, out_zero_point); + } else { + out_line[_oc] = acc; + } + } + } + } + } + } +} + +void xa_opt_quantized_conv_nhwc( + KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + constexpr int kNnlibMaxDim = 4; + + if (input.scalar_type() == ScalarType::Char) { + WORD8* __restrict__ p_out = + (WORD8* __restrict__)out.mutable_data_ptr(); + WORD8* __restrict__ p_inp = + (WORD8* __restrict__)input.const_data_ptr(); + WORD8* __restrict__ p_kernel = + (WORD8* __restrict__)weight.const_data_ptr(); + WORD32* __restrict__ p_bias = + (WORD32* __restrict__)bias.const_data_ptr(); + + WORD32 input_height = conv1d ? 1 : input.size(2); + WORD32 input_width = conv1d ? input.size(2) : input.size(3); + WORD32 input_channels = input.size(1); + WORD32 kernel_height = conv1d ? 1 : weight.size(2); + WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); + WORD32 kernel_channels = weight.size(1); + WORD32 out_channels = weight.size(0); + WORD32 out_height = conv1d ? 1 : out.size(2); + WORD32 out_width = conv1d ? out.size(2) : out.size(3); + WORD32 batches = input.size(0); + + WORD32 x_stride = stride[1]; + WORD32 y_stride = stride[0]; + WORD32 x_padding = padding[1]; + WORD32 y_padding = padding[0]; + WORD32 dilation_width = dilation[1]; + WORD32 dilation_height = dilation[0]; + + // WORD32* kernel_bias_ptr = + // (WORD32*)weight_zero_point.const_data_ptr(); + + WORD32 input_zero_bias = -in_zero_point; + WORD32 kernel_zero_bias = -weight_zero_point; + + WORD32 out_multiplier32[out_channels]; + WORD32 out_shift32[out_channels]; + + float out_scale = 1. / output_scale; + + for (int i = 0; i < out_channels; i++) { + out_multiplier32[i] = bias_scale * out_scale * 2147483648; + out_shift32[i] = 0; + } + + WORD32 out_zero_bias = output_zero_point; + WORD32 inp_precision = 8; + WORD32 kernel_precision = 8; + pVOID p_scratch = nullptr; + WORD32* ptr_scratch; + + WORD32 scratch_size = 0; + + if (groups == 1) { + WORD32 out_data_format = 1; + + scratch_size = xa_nn_conv2d_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + y_stride, + y_padding, + x_stride, + x_padding, + out_height, + out_width, + out_channels, + inp_precision, + kernel_precision, + out_data_format); + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = + p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_per_chan_sym8sxasym8s( + out_batch, + in_batch, + p_kernel, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + kernel_channels, + dilation_height, + dilation_width, + out_channels, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + out_data_format, + p_scratch); + } + return; + } + + if (groups == input_channels) { + WORD32 channels_multiplier = out_channels / input_channels; + + scratch_size = xa_nn_conv2d_depthwise_getsize( + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + inp_precision, + 0); // NHWC + + scratch_size = scratch_size < 0 ? 0 : scratch_size; + + ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); + + p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); + + WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( + ctx, + ((batches * out_channels * out_height * out_width) + 8) * + sizeof(WORD8)); + + WORD8* p_out_temp = (WORD8*)ALIGN_PTR(ptr1, 8); + + for (int _n = 0; _n < batches; _n++) { + WORD8* in_batch = + p_inp + _n * input_channels * input_height * input_width; + WORD8* out_batch = + p_out_temp + _n * out_channels * out_height * out_width; + + xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( + out_batch, + p_kernel, + in_batch, + p_bias, + input_height, + input_width, + input_channels, + kernel_height, + kernel_width, + channels_multiplier, + x_stride, + y_stride, + x_padding, + y_padding, + out_height, + out_width, + input_zero_bias, + out_multiplier32, + out_shift32, + out_zero_bias, + 0, // NHWC + 0, // NHWC + p_scratch); + } + + return; + } + } +} + +void quantized_conv_nhwc( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, h, w, c] + const int n = input.size(0); + const int h = conv1d ? 1 : input.size(1); + const int w = conv1d ? input.size(1) : input.size(2); + const int c = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wh, ww, wc] + const int oc = weight.size(0); + const int wh = conv1d ? 1 : weight.size(1); + const int ww = conv1d ? weight.size(1) : weight.size(2); + const int wc = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oh, ow, oc] + const int oh = conv1d ? 1 : out.size(1); + const int ow = conv1d ? out.size(1) : out.size(2); + +#define typed_quantized_conv2d_nhwc(ctype, dtype) \ + case ScalarType::dtype: { \ + conv2d_nhwc_core_generic( \ + input.const_data_ptr(), \ + weight.const_data_ptr(), \ + bias.const_data_ptr(), \ + out.mutable_data_ptr(), \ + n, \ + h, \ + w, \ + c, \ + oc, \ + wh, \ + ww, \ + wc, \ + oh, \ + ow, \ + stride[0], \ + stride[1], \ + padding[0], \ + padding[1], \ + dilation[0], \ + dilation[1], \ + groups, \ + in_zero_point, \ + weight_zero_point, \ + bias_scale, \ + output_scale, \ + (ctype)output_zero_point); \ + break; \ + } + ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nhwc); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } + +#undef typed_quantized_conv2d_nhwc +} + +void quantized_conv_nhwc_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED const Tensor& out_multiplier, + __ET_UNUSED const Tensor& out_shift, + Tensor& out) { + const float bias_scale_float = bias_scale.const_data_ptr()[0]; + const int32_t weight_zero_point_int = + weight_zero_point.const_data_ptr()[0]; + + bool optimized = 0; + + if ((input.scalar_type() == ScalarType::Char) || + (input.scalar_type() == ScalarType::Byte)) + optimized = 1; + + if ((dilation[0] != 1) || (dilation[1] != 1)) + optimized = 0; + + if (optimized) { + xa_opt_quantized_conv_nhwc( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point_int, + bias_scale_float, + output_scale, + output_zero_point, + out); + } else { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point_int, + bias_scale_float, + output_scale, + output_zero_point, + out); + } +} + +void quantized_conv_nhwc_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + bool optimized = 0; + + if ((input.scalar_type() == ScalarType::Char) || + (input.scalar_type() == ScalarType::Byte)) + optimized = 1; + + if ((dilation[0] != 1) || (dilation[1] != 1)) + optimized = 0; + + if (optimized) { + xa_opt_quantized_conv_nhwc( + ctx, + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); + } else { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); + } +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_conv_out.cpp b/backends/cadence/hifi/operators/op_quantized_conv_out.cpp deleted file mode 100644 index a24bad5f9a5..00000000000 --- a/backends/cadence/hifi/operators/op_quantized_conv_out.cpp +++ /dev/null @@ -1,1117 +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. - */ - -#include -#include -#include - -#define ALIGN_PTR(x, bytes) ((((unsigned)(x)) + (bytes - 1)) & (~(bytes - 1))) - -using Tensor = executorch::aten::Tensor; -using KernelRuntimeContext = torch::executor::KernelRuntimeContext; -using ScalarType = executorch::aten::ScalarType; -using ::executorch::aten::IntArrayRef; - -namespace cadence { -namespace impl { -namespace HiFi { -namespace native { - -// This implements a generic 2d conv kernel that operates on raw pointers. -// The version handles both quantized and fp32 convolutions. -// The input is of shape [n x c x h x w] -// The weight is of shape [oc x wc x wh x ww], where wc == c -// The output is of shape [n x oc x oh x ow] -// The bias is of shape [oc] -template < - typename IT = float, - typename WT = IT, - typename BT = IT, - typename OT = IT, - bool quantized = false> -__attribute__((noinline)) void conv2d_nchw_core_generic( - // All the arrays - const IT* __restrict__ p_in, - const WT* __restrict__ p_weight, - const BT* __restrict__ p_bias, - OT* __restrict__ p_out, - // The array sizes - int32_t n, - int32_t c, - int32_t h, - int32_t w, - int32_t oc, - int32_t wc, - int32_t wh, - int32_t ww, - int32_t oh, - int32_t ow, - // Stride - int16_t s0, - int16_t s1, - // Padding - int16_t p0, - int16_t p1, - // Dilation - int16_t d0, - int16_t d1, - // Group for depthwise conv - int16_t groups, - // Optional args that are only relevant for quantized convolution - // input zero point - IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, - float out_scale = 1, - OT out_zero_point = 0) { - float inv_out_scale = 1. / out_scale; - bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; - - // Compute the number of in and out channels per group - const int ocpg = oc / groups; - const int icpg = c / groups; - - // Iterate over all the output batches (i.e., n) - for (int _n = 0; _n < n; ++_n) { - const IT* in_batch = p_in + _n * c * h * w; - OT* out_batch = p_out + _n * oc * oh * ow; - // Compute separable convolution for each group - for (int _g = 0; _g < groups; ++_g) { - // Identify the input and output channels involved in the computation - // of this group - int sic = _g * icpg; - int soc = _g * ocpg; - // Populate all the output channels in the group - for (int _oc = soc; _oc < soc + ocpg; ++_oc) { - OT* out_plane = out_batch + _oc * oh * ow; - const WT* weight_batch = p_weight + _oc * wc * wh * ww; - // We compute one output channel at a time. The computation can be - // thought of as a stencil computation: we iterate over an input of size - // icpg x h x w, with a stencil of size icpg x wh x ww, to compute an - // output channel of size 1 x oh x ow. - for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { - for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { - float acc = p_bias[_oc]; - // Below is the stencil computation that performs the hadamard - // product+accumulation of each input channel (contributing to the - // output channel being computed) with the corresponding weight - // channel. - // If the padding is 0, and dilation is 1, then we can remove the - // unnecessary checks, and simplify the code so that it can be - // vectorized by Tensilica compiler. - if (zero_pad_unit_dilation) { - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - const IT* in_plane = in_batch + _ic * h * w; - const WT* weight_plane = weight_batch + (_ic - sic) * wh * ww; - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - int ioff = (_h + _wh) * w + (_w + _ww); - int woff = _wh * ww + _ww; - float lhs = in_plane[ioff] - in_zero_point; - float rhs = weight_plane[woff] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } else { - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - const IT* in_plane = in_batch + _ic * h * w; - const WT* weight_plane = weight_batch + (_ic - sic) * wh * ww; - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - if (((_h + d0 * _wh - p0) >= 0) && - ((_h + d0 * _wh - p0) < h) && - ((_w + d1 * _ww - p1) >= 0) && - ((_w + d1 * _ww - p1) < w)) { - int ioff = - (_h + d0 * _wh - p0) * w + (_w + d1 * _ww - p1); - int woff = _wh * ww + _ww; - float lhs = in_plane[ioff] - in_zero_point; - float rhs = weight_plane[woff] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } - } - if (quantized) { - float val = bias_scale * acc; - out_plane[_oh * ow + _ow] = - kernels::quantize(val, inv_out_scale, out_zero_point); - } else { - out_plane[_oh * ow + _ow] = acc; - } - } - } - } - } - } -} - -template < - typename IT = float, - typename WT = IT, - typename BT = IT, - typename OT = IT, - bool quantized = false> -__attribute__((noinline)) void conv2d_nhwc_core_generic( - // All the arrays - const IT* __restrict__ p_in, - const WT* __restrict__ p_weight, - const BT* __restrict__ p_bias, - OT* __restrict__ p_out, - // The array sizes - int32_t n, - int32_t h, - int32_t w, - int32_t c, - int32_t oc, - int32_t wh, - int32_t ww, - int32_t wc, - int32_t oh, - int32_t ow, - // Stride - int16_t s0, - int16_t s1, - // Padding - int16_t p0, - int16_t p1, - // Dilation - int16_t d0, - int16_t d1, - // Group for depthwise conv - int16_t groups, - // Optional args that are only relevant for quantized convolution - // input zero point - IT in_zero_point = 0, - // weight zero point - int32_t weight_zero_point = 0, - float bias_scale = 1, - float out_scale = 1, - OT out_zero_point = 0) { - float inv_out_scale = 1. / out_scale; - bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; - - // Compute the number of in and out channels per group - const int ocpg = oc / groups; - const int icpg = c / groups; - - // Iterate over all the output batches (i.e., n) - for (int _n = 0; _n < n; ++_n) { - const IT* in_batch = p_in + _n * h * w * c; - OT* out_batch = p_out + _n * oh * ow * oc; - for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { - for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { - OT* out_line = out_batch + (_oh * ow + _ow) * oc; - // Compute separable convolution for each group - for (int _g = 0; _g < groups; ++_g) { - // Identify the input and output channels involved in the computation - // of this group - int sic = _g * icpg; - int soc = _g * ocpg; - // Populate all the output channels in the group - for (int _oc = soc; _oc < soc + ocpg; ++_oc) { - const WT* weight_batch = p_weight + _oc * wh * ww * wc; - // We compute one output channel at a time. The computation can be - // thought of as a stencil computation: we iterate over an input of - // size h x w x icpg, with a stencil of size wh x ww x icpg, to - // compute an output channel of size oh x ow x 1. - float acc = p_bias[_oc]; - // Below is the stencil computation that performs the hadamard - // product+accumulation of each input channel (contributing to - // the output channel being computed) with the corresponding - // weight channel. If the padding is 0, and dilation is 1, then - // we can remove the unnecessary checks, and simplify the code - // so that it can be vectorized by Tensilica compiler.x`` - if (zero_pad_unit_dilation) { - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - const IT* in_line = - in_batch + (_h + _wh) * w * c + (_w + _ww) * c; - const WT* weight_line = - weight_batch + _wh * ww * wc + _ww * wc; - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - float lhs = in_line[_ic] - in_zero_point; - float rhs = weight_line[_ic - sic] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } else { - for (int _wh = 0; _wh < wh; ++_wh) { - for (int _ww = 0; _ww < ww; ++_ww) { - if (((_h + d0 * _wh - p0) >= 0) && - ((_h + d0 * _wh - p0) < h) && - ((_w + d1 * _ww - p1) >= 0) && - ((_w + d1 * _ww - p1 < w))) { - const IT* in_line = in_batch + - (_h + d0 * _wh - p0) * w * c + (_w + d1 * _ww - p1) * c; - const WT* weight_line = - weight_batch + _wh * ww * wc + _ww * wc; - for (int _ic = sic; _ic < sic + icpg; ++_ic) { - float lhs = in_line[_ic] - in_zero_point; - float rhs = weight_line[_ic - sic] - - (quantized ? weight_zero_point : 0); - acc += lhs * rhs; - } - } - } - } - } - if (quantized) { - float val = bias_scale * acc; - out_line[_oc] = - kernels::quantize(val, inv_out_scale, out_zero_point); - } else { - out_line[_oc] = acc; - } - } - } - } - } - } -} - -void xa_opt_quantized_conv_nhwc( - KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int16_t groups, - int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, - float output_scale, - int32_t output_zero_point, - Tensor& out) { - bool conv1d = input.dim() == 3; - constexpr int kNnlibMaxDim = 4; - - if (input.scalar_type() == ScalarType::Char) { - WORD8* __restrict__ p_out = - (WORD8* __restrict__)out.mutable_data_ptr(); - WORD8* __restrict__ p_inp = - (WORD8* __restrict__)input.const_data_ptr(); - WORD8* __restrict__ p_kernel = - (WORD8* __restrict__)weight.const_data_ptr(); - WORD32* __restrict__ p_bias = - (WORD32* __restrict__)bias.const_data_ptr(); - - WORD32 input_height = conv1d ? 1 : input.size(2); - WORD32 input_width = conv1d ? input.size(2) : input.size(3); - WORD32 input_channels = input.size(1); - WORD32 kernel_height = conv1d ? 1 : weight.size(2); - WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); - WORD32 kernel_channels = weight.size(1); - WORD32 out_channels = weight.size(0); - WORD32 out_height = conv1d ? 1 : out.size(2); - WORD32 out_width = conv1d ? out.size(2) : out.size(3); - WORD32 batches = input.size(0); - - WORD32 x_stride = stride[1]; - WORD32 y_stride = stride[0]; - WORD32 x_padding = padding[1]; - WORD32 y_padding = padding[0]; - WORD32 dilation_width = dilation[1]; - WORD32 dilation_height = dilation[0]; - - // WORD32* kernel_bias_ptr = - // (WORD32*)weight_zero_point.const_data_ptr(); - - WORD32 input_zero_bias = -in_zero_point; - WORD32 kernel_zero_bias = -weight_zero_point; - - WORD32 out_multiplier32[out_channels]; - WORD32 out_shift32[out_channels]; - - float out_scale = 1. / output_scale; - - for (int i = 0; i < out_channels; i++) { - out_multiplier32[i] = bias_scale * out_scale * 2147483648; - out_shift32[i] = 0; - } - - WORD32 out_zero_bias = output_zero_point; - WORD32 inp_precision = 8; - WORD32 kernel_precision = 8; - pVOID p_scratch = nullptr; - WORD32* ptr_scratch; - - WORD32 scratch_size = 0; - - if (groups == 1) { - WORD32 out_data_format = 1; - - scratch_size = xa_nn_conv2d_getsize( - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - kernel_channels, - dilation_height, - dilation_width, - y_stride, - y_padding, - x_stride, - x_padding, - out_height, - out_width, - out_channels, - inp_precision, - kernel_precision, - out_data_format); - - scratch_size = scratch_size < 0 ? 0 : scratch_size; - - ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); - - p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); - - for (int _n = 0; _n < batches; _n++) { - WORD8* in_batch = - p_inp + _n * input_channels * input_height * input_width; - WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; - - xa_nn_conv2d_per_chan_sym8sxasym8s( - out_batch, - in_batch, - p_kernel, - p_bias, - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - kernel_channels, - dilation_height, - dilation_width, - out_channels, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - input_zero_bias, - out_multiplier32, - out_shift32, - out_zero_bias, - out_data_format, - p_scratch); - } - return; - } - - if (groups == input_channels) { - WORD32 channels_multiplier = out_channels / input_channels; - - scratch_size = xa_nn_conv2d_depthwise_getsize( - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - channels_multiplier, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - inp_precision, - 0); // NHWC - - scratch_size = scratch_size < 0 ? 0 : scratch_size; - - ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); - - p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); - - WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( - ctx, - ((batches * out_channels * out_height * out_width) + 8) * - sizeof(WORD8)); - - WORD8* p_out_temp = (WORD8*)ALIGN_PTR(ptr1, 8); - - for (int _n = 0; _n < batches; _n++) { - WORD8* in_batch = - p_inp + _n * input_channels * input_height * input_width; - WORD8* out_batch = - p_out_temp + _n * out_channels * out_height * out_width; - - xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( - out_batch, - p_kernel, - in_batch, - p_bias, - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - channels_multiplier, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - input_zero_bias, - out_multiplier32, - out_shift32, - out_zero_bias, - 0, // NHWC - 0, // NHWC - p_scratch); - } - - return; - } - } -} - -void xa_opt_quantized_conv_nchw( - KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int16_t groups, - int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, - float output_scale, - int32_t output_zero_point, - Tensor& out) { - bool conv1d = input.dim() == 3; - constexpr int kNnlibMaxDim = 4; - - if (input.scalar_type() == ScalarType::Char) { - WORD8* __restrict__ p_out = - (WORD8* __restrict__)out.mutable_data_ptr(); - WORD8* __restrict__ p_inp = - (WORD8* __restrict__)input.const_data_ptr(); - WORD8* __restrict__ p_kernel = - (WORD8* __restrict__)weight.const_data_ptr(); - WORD32* __restrict__ p_bias = - (WORD32* __restrict__)bias.const_data_ptr(); - - WORD32 input_height = conv1d ? 1 : input.size(2); - WORD32 input_width = conv1d ? input.size(2) : input.size(3); - WORD32 input_channels = input.size(1); - WORD32 kernel_height = conv1d ? 1 : weight.size(2); - WORD32 kernel_width = conv1d ? weight.size(2) : weight.size(3); - WORD32 kernel_channels = weight.size(1); - WORD32 out_channels = weight.size(0); - WORD32 out_height = conv1d ? 1 : out.size(2); - WORD32 out_width = conv1d ? out.size(2) : out.size(3); - WORD32 batches = input.size(0); - - WORD32 x_stride = stride[1]; - WORD32 y_stride = stride[0]; - WORD32 x_padding = padding[1]; - WORD32 y_padding = padding[0]; - WORD32 dilation_width = dilation[1]; - WORD32 dilation_height = dilation[0]; - - // WORD32* kernel_bias_ptr = - // (WORD32*)weight_zero_point.const_data_ptr(); - - WORD32 input_zero_bias = -in_zero_point; - WORD32 kernel_zero_bias = -weight_zero_point; - - WORD32 out_multiplier32[out_channels]; - WORD32 out_shift32[out_channels]; - - float out_scale = 1. / output_scale; - - for (int i = 0; i < out_channels; i++) { - out_multiplier32[i] = bias_scale * out_scale * 2147483648; - out_shift32[i] = 0; - } - - WORD32 out_zero_bias = output_zero_point; - WORD32 inp_precision = 8; - WORD32 kernel_precision = 8; - pVOID p_scratch = nullptr; - WORD32* ptr_scratch; - - WORD32 scratch_size = 0; - - if (groups == 1) { - WORD32 out_data_format = 1; - - WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( - ctx, - ((batches * input_channels * input_height * input_width) + 8) * - sizeof(WORD8)); - - WORD8* ptr2 = (WORD8*)kernels::allocate_temp_memory( - ctx, - ((out_channels * kernel_channels * kernel_height * kernel_width) + - 8) * - sizeof(WORD8)); - - WORD8* pin = (WORD8*)ALIGN_PTR(ptr1, 8); - WORD8* pkernel = (WORD8*)ALIGN_PTR(ptr2, 8); - - WORD32 p_inp_shape[kNnlibMaxDim]; - p_inp_shape[0] = input.size(0); - p_inp_shape[1] = input_channels; - p_inp_shape[2] = input_height; - p_inp_shape[3] = input_width; - - WORD32 p_out_shape[kNnlibMaxDim]; - p_out_shape[0] = input.size(0); - p_out_shape[1] = input_height; - p_out_shape[2] = input_width; - p_out_shape[3] = input_channels; - - WORD32 p_permute_vec[kNnlibMaxDim] = {0, 2, 3, 1}; - - xa_nn_transpose_8_8( - pin, - p_out_shape, - p_inp, - p_inp_shape, - p_permute_vec, - kNnlibMaxDim, // input dimensions - kNnlibMaxDim); // output dimensions - - WORD32 p_inp_shape1[kNnlibMaxDim]; - p_inp_shape1[0] = out_channels; - p_inp_shape1[1] = kernel_channels; - p_inp_shape1[2] = kernel_height; - p_inp_shape1[3] = kernel_width; - - WORD32 p_out_shape1[kNnlibMaxDim]; - p_out_shape1[0] = out_channels; - p_out_shape1[1] = kernel_height; - p_out_shape1[2] = kernel_width; - p_out_shape1[3] = kernel_channels; - - xa_nn_transpose_8_8( - pkernel, - p_out_shape1, - p_kernel, - p_inp_shape1, - p_permute_vec, - kNnlibMaxDim, // input dimensions - kNnlibMaxDim); // output dimensions - - scratch_size = xa_nn_conv2d_getsize( - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - kernel_channels, - dilation_height, - dilation_width, - y_stride, - y_padding, - x_stride, - x_padding, - out_height, - out_width, - out_channels, - inp_precision, - kernel_precision, - out_data_format); - - scratch_size = scratch_size < 0 ? 0 : scratch_size; - - ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); - - p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); - - for (int _n = 0; _n < batches; _n++) { - WORD8* in_batch = - pin + _n * input_channels * input_height * input_width; - WORD8* out_batch = p_out + _n * out_channels * out_height * out_width; - - xa_nn_conv2d_per_chan_sym8sxasym8s( - out_batch, - in_batch, - pkernel, - p_bias, - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - kernel_channels, - dilation_height, - dilation_width, - out_channels, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - input_zero_bias, - out_multiplier32, - out_shift32, - out_zero_bias, - out_data_format, - p_scratch); - } - return; - } - - if (groups == input_channels) { - WORD32 channels_multiplier = out_channels / input_channels; - - scratch_size = xa_nn_conv2d_depthwise_getsize( - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - channels_multiplier, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - inp_precision, - 1); // NCHW - - scratch_size = scratch_size < 0 ? 0 : scratch_size; - - ptr_scratch = (WORD32*)kernels::allocate_temp_memory(ctx, scratch_size); - - p_scratch = (pVOID)ALIGN_PTR(ptr_scratch, 8); - - WORD8* ptr1 = (WORD8*)kernels::allocate_temp_memory( - ctx, - ((batches * out_channels * out_height * out_width) + 8) * - sizeof(WORD8)); - - WORD8* p_out_temp = (WORD8*)ALIGN_PTR(ptr1, 8); - - for (int _n = 0; _n < batches; _n++) { - WORD8* in_batch = - p_inp + _n * input_channels * input_height * input_width; - WORD8* out_batch = - p_out_temp + _n * out_channels * out_height * out_width; - - xa_nn_conv2d_depthwise_per_chan_sym8sxasym8s( - out_batch, - p_kernel, - in_batch, - p_bias, - input_height, - input_width, - input_channels, - kernel_height, - kernel_width, - channels_multiplier, - x_stride, - y_stride, - x_padding, - y_padding, - out_height, - out_width, - input_zero_bias, - out_multiplier32, - out_shift32, - out_zero_bias, - 1, // NCHW - 0, // NHWC - p_scratch); - } - - WORD32 p_inp_shape[kNnlibMaxDim]; - p_inp_shape[0] = batches; - p_inp_shape[1] = out_height; - p_inp_shape[2] = out_width; - p_inp_shape[3] = out_channels; - - WORD32 p_out_shape[kNnlibMaxDim]; - p_out_shape[0] = batches; - p_out_shape[1] = out_channels; - p_out_shape[2] = out_height; - p_out_shape[3] = out_width; - - WORD32 p_permute_vec[kNnlibMaxDim] = {0, 3, 1, 2}; - - xa_nn_transpose_8_8( - p_out, - p_out_shape, - p_out_temp, - p_inp_shape, - p_permute_vec, - kNnlibMaxDim, // input dimensions - kNnlibMaxDim); // output dimensions - - return; - } - } -} - -// The quantized convolution kernel. in_scale and weight_scale are implicit in -// bias_scale, since it is a product of the two. The kernel will branch to -// quantized::conv1d or quantized::conv2d based on the dimensionality of -// activation tensor. -void quantized_conv_nchw( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int16_t groups, - int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, - float output_scale, - int32_t output_zero_point, - Tensor& out) { - bool conv1d = input.dim() == 3; - // input = [n, c, h, w] - const int n = input.size(0); - const int c = input.size(1); - const int h = conv1d ? 1 : input.size(2); - const int w = conv1d ? input.size(2) : input.size(3); - // weight = [oc, wc, wh, ww] - const int oc = weight.size(0); - const int wc = weight.size(1); - const int wh = conv1d ? 1 : weight.size(2); - const int ww = conv1d ? weight.size(2) : weight.size(3); - // output = [n, oc, oh, ow] - const int oh = conv1d ? 1 : out.size(2); - const int ow = conv1d ? out.size(2) : out.size(3); - -#define typed_quantized_conv2d_nchw(ctype, dtype) \ - case ScalarType::dtype: { \ - conv2d_nchw_core_generic( \ - input.const_data_ptr(), \ - weight.const_data_ptr(), \ - bias.const_data_ptr(), \ - out.mutable_data_ptr(), \ - n, \ - c, \ - h, \ - w, \ - oc, \ - wc, \ - wh, \ - ww, \ - oh, \ - ow, \ - stride[0], \ - stride[1], \ - padding[0], \ - padding[1], \ - dilation[0], \ - dilation[1], \ - groups, \ - in_zero_point, \ - weight_zero_point, \ - bias_scale, \ - output_scale, \ - (ctype)output_zero_point); \ - break; \ - } - ScalarType dtype = out.scalar_type(); - switch (dtype) { - ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nchw); - default: - ET_DCHECK_MSG( - false, "Unhandled dtype %s", torch::executor::toString(dtype)); - } - -#undef typed_quantized_conv2d_nchw -} - -void quantized_conv_nhwc( - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int16_t groups, - int32_t in_zero_point, - int32_t weight_zero_point, - float bias_scale, - float output_scale, - int32_t output_zero_point, - Tensor& out) { - bool conv1d = input.dim() == 3; - // input = [n, h, w, c] - const int n = input.size(0); - const int h = conv1d ? 1 : input.size(1); - const int w = conv1d ? input.size(1) : input.size(2); - const int c = conv1d ? input.size(2) : input.size(3); - // weight = [oc, wh, ww, wc] - const int oc = weight.size(0); - const int wh = conv1d ? 1 : weight.size(1); - const int ww = conv1d ? weight.size(1) : weight.size(2); - const int wc = conv1d ? weight.size(2) : weight.size(3); - // output = [n, oh, ow, oc] - const int oh = conv1d ? 1 : out.size(1); - const int ow = conv1d ? out.size(1) : out.size(2); - -#define typed_quantized_conv2d_nhwc(ctype, dtype) \ - case ScalarType::dtype: { \ - conv2d_nhwc_core_generic( \ - input.const_data_ptr(), \ - weight.const_data_ptr(), \ - bias.const_data_ptr(), \ - out.mutable_data_ptr(), \ - n, \ - h, \ - w, \ - c, \ - oc, \ - wh, \ - ww, \ - wc, \ - oh, \ - ow, \ - stride[0], \ - stride[1], \ - padding[0], \ - padding[1], \ - dilation[0], \ - dilation[1], \ - groups, \ - in_zero_point, \ - weight_zero_point, \ - bias_scale, \ - output_scale, \ - (ctype)output_zero_point); \ - break; \ - } - ScalarType dtype = out.scalar_type(); - switch (dtype) { - ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nhwc); - default: - ET_DCHECK_MSG( - false, "Unhandled dtype %s", torch::executor::toString(dtype)); - } - -#undef typed_quantized_conv2d_nhwc -} - -void quantized_conv_out( - __ET_UNUSED KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - const Tensor& weight_zero_point, - const Tensor& bias_scale, - double output_scale, - int64_t output_zero_point, - __ET_UNUSED const Tensor& out_multiplier, - __ET_UNUSED const Tensor& out_shift, - bool channel_last, - Tensor& out) { - const float bias_scale_float = bias_scale.const_data_ptr()[0]; - const int32_t weight_zero_point_int = - weight_zero_point.const_data_ptr()[0]; - - bool optimized = 0; - - if ((input.scalar_type() == ScalarType::Char) || - (input.scalar_type() == ScalarType::Byte)) - optimized = 1; - - if ((dilation[0] != 1) || (dilation[1] != 1)) - optimized = 0; - - if (channel_last) { - if (optimized) { - xa_opt_quantized_conv_nhwc( - ctx, - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point_int, - bias_scale_float, - output_scale, - output_zero_point, - out); - } else { - quantized_conv_nhwc( - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point_int, - bias_scale_float, - output_scale, - output_zero_point, - out); - } - } else { - if (optimized) { - xa_opt_quantized_conv_nchw( - ctx, - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point_int, - bias_scale_float, - output_scale, - output_zero_point, - out); - } else { - quantized_conv_nchw( - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point_int, - bias_scale_float, - output_scale, - output_zero_point, - out); - } - } -} - -void quantized_conv_per_tensor_out( - __ET_UNUSED KernelRuntimeContext& ctx, - const Tensor& input, - const Tensor& weight, - const Tensor& bias, - IntArrayRef stride, - IntArrayRef padding, - IntArrayRef dilation, - int64_t groups, - int64_t in_zero_point, - int64_t weight_zero_point, - double bias_scale, - double output_scale, - int64_t output_zero_point, - __ET_UNUSED int64_t out_multiplier, - __ET_UNUSED int64_t out_shift, - bool channel_last, - Tensor& out) { - bool optimized = 0; - - if ((input.scalar_type() == ScalarType::Char) || - (input.scalar_type() == ScalarType::Byte)) - optimized = 1; - - if ((dilation[0] != 1) || (dilation[1] != 1)) - optimized = 0; - - if (channel_last) { - if (optimized) { - xa_opt_quantized_conv_nhwc( - ctx, - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point, - bias_scale, - output_scale, - output_zero_point, - out); - } else { - quantized_conv_nhwc( - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point, - bias_scale, - output_scale, - output_zero_point, - out); - } - } else { - if (optimized) { - xa_opt_quantized_conv_nchw( - ctx, - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point, - bias_scale, - output_scale, - output_zero_point, - out); - } else { - quantized_conv_nchw( - input, - weight, - bias, - stride, - padding, - dilation, - groups, - in_zero_point, - weight_zero_point, - bias_scale, - output_scale, - output_zero_point, - out); - } - } -} - -} // namespace native -} // namespace HiFi -} // namespace impl -} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..5e3a5173f32 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,64 @@ +/* + * 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 + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using std::optional; + +void quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { + // input comes in shape [leading_dims, in_dim] + // weight comes in shape [out_dim, in_dim] + // output comes in empty with shape [leading_dims, out_dim] + // Perform matrix multiply (M x N) x (N x P)' => M x P + int64_t leading_dims = 1; + int64_t out_dim = weight.size(0); // = out_dim + int64_t in_dim = weight.size(1); // = in_dim + + const int8_t* __restrict__ in_data = in.const_data_ptr(); + const int8_t* __restrict__ weight_data = weight.const_data_ptr(); + const int32_t* __restrict__ bias_data = bias.const_data_ptr(); + int8_t* __restrict__ out_data = out.mutable_data_ptr(); + + int32_t ret = xa_nn_fully_connected_asym8sxasym8s_asym8s( + out_data, + weight_data, + in_data, + bias_data, + in_dim, // weight_depth, number of columns in weight + out_dim, // out_depth, number of rows in weight + -in_zero_point, + -static_cast(weight_zero_point), + static_cast(out_multiplier), + static_cast(out_shift), + out_zero_point); + ET_DCHECK_MSG(ret == 0, "HiFi quantized::fully_connected failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..80509fdd5db --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,64 @@ +/* + * 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 + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using std::optional; + +void quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { + // input comes in shape [leading_dims, in_dim] + // weight comes in shape [out_dim, in_dim] + // output comes in empty with shape [leading_dims, out_dim] + // Perform matrix multiply (M x N) x (N x P)' => M x P + int64_t leading_dims = 1; + int64_t out_dim = weight.size(0); // = out_dim + int64_t in_dim = weight.size(1); // = in_dim + + const uint8_t* __restrict__ in_data = in.const_data_ptr(); + const uint8_t* __restrict__ weight_data = weight.const_data_ptr(); + const int32_t* __restrict__ bias_data = bias.const_data_ptr(); + uint8_t* __restrict__ out_data = out.mutable_data_ptr(); + + int32_t ret = xa_nn_fully_connected_asym8uxasym8u_asym8u( + out_data, + weight_data, + in_data, + bias_data, + in_dim, // weight_depth, number of columns in weight + out_dim, // out_depth, number of rows in weight + -in_zero_point, + -static_cast(weight_zero_point), + static_cast(out_multiplier), + static_cast(out_shift), + out_zero_point); + ET_DCHECK_MSG(ret == 0, "HiFi quantized::fully_connected failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp b/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp index 0606fb28523..7906f245b03 100644 --- a/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp +++ b/backends/cadence/hifi/operators/op_quantized_layer_norm.cpp @@ -13,6 +13,8 @@ #include #include +using ::cadence::impl::HiFi::kernels::dequantize; +using ::cadence::impl::HiFi::kernels::quantize; using ::executorch::aten::IntArrayRef; using ::executorch::aten::ScalarType; using ::executorch::aten::Tensor; @@ -80,11 +82,9 @@ void quantized_layer_norm_per_tensor_( for (size_t j = 0; j < last_dim; ++j) { // Since X is quantized, we dequantize it, compute fp32 result, and // quantize the result to an int8/uint8 value. - float val = ::cadence::impl::HiFi::kernels::dequantize( - x[j], input_scale, input_zero_point); + float val = dequantize(x[j], input_scale, input_zero_point); val = (val - mean) * inv_std * weight_data[j] + bias_data[j]; - y[j] = ::cadence::impl::HiFi::kernels::quantize( - val, output_inv_scale, output_zero_point); + y[j] = quantize(val, output_inv_scale, output_zero_point); } } } diff --git a/backends/cadence/hifi/operators/op_quantized_linear_asym8sxasym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_linear_asym8sxasym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..7b8ab8e91b9 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_linear_asym8sxasym8s_asym8s_per_tensor_out.cpp @@ -0,0 +1,75 @@ +/* + * 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 cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::getLeadingDims; +using ::executorch::runtime::KernelRuntimeContext; +using std::optional; + +void quantized_linear_asym8sxasym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { + // input comes in shape [leading_dims, in_dim] + // weight comes in shape [out_dim, in_dim] + // output comes in empty with shape [leading_dims, out_dim] + // Perform matrix multiply (M x N) x (N x P)' => M x P + const int64_t leading_dims = getLeadingDims(in, in.dim() - 1); + const int64_t out_dim = weight.size(0); // = out_dim + const int64_t in_dim = weight.size(1); // = in_dim + + const int8_t* __restrict__ in_data = in.const_data_ptr(); + const int8_t* __restrict__ weight_data = weight.const_data_ptr(); + const int32_t* __restrict__ bias_data = bias.const_data_ptr(); + int8_t* __restrict__ out_data = out.mutable_data_ptr(); + + const int32_t out_multipler_int32 = static_cast(out_multiplier); + const int32_t out_shift_int32 = static_cast(out_shift); + + // The nnlib kernel to compute quantized linear via matmul. + const int32_t ret = xa_nn_matmul_asym8sxasym8s_asym8s( + out_data, // p_out + weight_data, // p_mat1, + in_data, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dims, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -weight_zero_point, // mat1_zero_bias + -in_zero_point, // mat2_zero_bias + out_multipler_int32, // out_multiplier + out_shift_int32, // out_shift + out_zero_point); // out_zero_bias + ET_DCHECK_MSG(ret == 0, "HiFi quantized::linear_per_tensor failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_linear_asym8uxasym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_linear_asym8uxasym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..e9632e77eeb --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_linear_asym8uxasym8u_asym8u_per_tensor_out.cpp @@ -0,0 +1,75 @@ +/* + * 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 cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::getLeadingDims; +using ::executorch::runtime::KernelRuntimeContext; +using std::optional; + +void quantized_linear_asym8uxasym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { + // input comes in shape [leading_dims, in_dim] + // weight comes in shape [out_dim, in_dim] + // output comes in empty with shape [leading_dims, out_dim] + // Perform matrix multiply (M x N) x (N x P)' => M x P + const int64_t leading_dims = getLeadingDims(in, in.dim() - 1); + const int64_t out_dim = weight.size(0); // = out_dim + const int64_t in_dim = weight.size(1); // = in_dim + + const uint8_t* __restrict__ in_data = in.const_data_ptr(); + const uint8_t* __restrict__ weight_data = weight.const_data_ptr(); + const int32_t* __restrict__ bias_data = bias.const_data_ptr(); + uint8_t* __restrict__ out_data = out.mutable_data_ptr(); + + const int32_t out_multipler_int32 = static_cast(out_multiplier); + const int32_t out_shift_int32 = static_cast(out_shift); + + // The nnlib kernel to compute quantized linear via matmul. + const int32_t ret = xa_nn_matmul_asym8uxasym8u_asym8u( + out_data, // p_out + weight_data, // p_mat1, + in_data, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dims, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -weight_zero_point, // mat1_zero_bias + -in_zero_point, // mat2_zero_bias + out_multipler_int32, // out_multiplier + out_shift_int32, // out_shift + out_zero_point); // out_zero_bias + ET_DCHECK_MSG(ret == 0, "HiFi quantized::linear_per_tensor failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_matmul_asym8sxasym8s_asym8s_out.cpp b/backends/cadence/hifi/operators/op_quantized_matmul_asym8sxasym8s_asym8s_out.cpp new file mode 100644 index 00000000000..0e7b3f1a2aa --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_matmul_asym8sxasym8s_asym8s_out.cpp @@ -0,0 +1,135 @@ +/* + * 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 + +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::runtime::getLeadingDims; +using torch::executor::RuntimeContext; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +void quantized_matmul_asym8sxasym8s_asym8s_out( + RuntimeContext& ctx, + const Tensor& X, + int64_t X_zero_point, + const Tensor& Y, + int64_t Y_zero_point, + const exec_aten::optional& bias, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + bool transposed, + Tensor& out) { + int8_t* __restrict__ out_data = out.mutable_data_ptr(); + const int8_t* __restrict__ X_data = X.const_data_ptr(); + const int8_t* __restrict__ Y_data = Y.const_data_ptr(); + size_t batch_size = getLeadingDims(X, X.dim() - 2); + size_t leading_dim = X.size(X.dim() - 2); + size_t out_dim = Y.size(Y.dim() - 1 - transposed); + size_t in_dim = X.size(X.dim() - 1); + + const int32_t* __restrict__ bias_data = + (WORD32* __restrict__)kernels::allocate_temp_memory( + ctx, (leading_dim * in_dim) * sizeof(int32_t)); + + ET_CHECK_MSG(bias_data != nullptr, "MemoryAllocationFailed"); + + std::memset((void*)bias_data, 0, (leading_dim * in_dim) * sizeof(int32_t)); + + int8_t* y_data_temp = NULL; + + if (!transposed) { + y_data_temp = + (int8_t*)kernels::allocate_temp_memory(ctx, (leading_dim * in_dim)); + + ET_CHECK_MSG(y_data_temp != nullptr, "MemoryAllocationFailed"); + } + + for (size_t i = 0; i < batch_size; ++i) { + const int8_t* x = X_data + i * leading_dim * in_dim; + const int8_t* y = Y_data + i * in_dim * out_dim; + int8_t* z = out_data + i * leading_dim * out_dim; + if (transposed) { + WORD32 ret_val = xa_nn_matmul_asym8sxasym8s_asym8s( + z, // p_out + y, // p_mat1, + x, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dim, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -(static_cast(Y_zero_point)), // mat1_zero_bias + -(static_cast(X_zero_point)), // mat2_zero_bias + static_cast(out_multiplier), // out_multiplier + static_cast(out_shift), // out_shift + static_cast(out_zero_point)); // out_zero_bias + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + } else { + /* Assuming matmul is 2D always */ + WORD32 num_inp_dims = 2; + WORD32 num_out_dims = 2; + + WORD32 p_inp_shape[2]; + WORD32 p_out_shape[2]; + WORD32 p_permute_vec[2] = {1, 0}; + + p_inp_shape[0] = leading_dim; + p_inp_shape[1] = in_dim; + p_out_shape[0] = in_dim; + p_out_shape[1] = leading_dim; + + WORD32 ret_val = xa_nn_transpose_8_8( + y_data_temp, + p_out_shape, + y, + p_inp_shape, + p_permute_vec, + num_out_dims, + num_inp_dims); + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + + ret_val = xa_nn_matmul_asym8sxasym8s_asym8s( + z, // p_out + y_data_temp, // p_mat1, + x, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dim, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -(static_cast(Y_zero_point)), // mat1_zero_bias + -(static_cast(X_zero_point)), // mat2_zero_bias + static_cast(out_multiplier), // out_multiplier + static_cast(out_shift), // out_shift + static_cast(out_zero_point)); // out_zero_bias + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + } + } +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_matmul_asym8uxasym8u_asym8u_out.cpp b/backends/cadence/hifi/operators/op_quantized_matmul_asym8uxasym8u_asym8u_out.cpp new file mode 100644 index 00000000000..7016e6635dc --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_matmul_asym8uxasym8u_asym8u_out.cpp @@ -0,0 +1,135 @@ +/* + * 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 + +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::runtime::getLeadingDims; +using torch::executor::RuntimeContext; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +void quantized_matmul_asym8uxasym8u_asym8u_out( + RuntimeContext& ctx, + const Tensor& X, + int64_t X_zero_point, + const Tensor& Y, + int64_t Y_zero_point, + const exec_aten::optional& bias, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + bool transposed, + Tensor& out) { + uint8_t* __restrict__ out_data = out.mutable_data_ptr(); + const uint8_t* __restrict__ X_data = X.const_data_ptr(); + const uint8_t* __restrict__ Y_data = Y.const_data_ptr(); + size_t batch_size = getLeadingDims(X, X.dim() - 2); + size_t leading_dim = X.size(X.dim() - 2); + size_t out_dim = Y.size(Y.dim() - 1 - transposed); + size_t in_dim = X.size(X.dim() - 1); + + const int32_t* __restrict__ bias_data = + (WORD32* __restrict__)kernels::allocate_temp_memory( + ctx, (leading_dim * in_dim) * sizeof(int32_t)); + + ET_CHECK_MSG(bias_data != nullptr, "MemoryAllocationFailed"); + + std::memset((void*)bias_data, 0, (leading_dim * in_dim) * sizeof(int32_t)); + + uint8_t* y_data_temp = NULL; + + if (!transposed) { + y_data_temp = + (uint8_t*)kernels::allocate_temp_memory(ctx, (leading_dim * in_dim)); + + ET_CHECK_MSG(y_data_temp != nullptr, "MemoryAllocationFailed"); + } + + for (size_t i = 0; i < batch_size; ++i) { + const uint8_t* x = X_data + i * leading_dim * in_dim; + const uint8_t* y = Y_data + i * in_dim * out_dim; + uint8_t* z = out_data + i * leading_dim * out_dim; + if (transposed) { + WORD32 ret_val = xa_nn_matmul_asym8uxasym8u_asym8u( + z, // p_out + y, // p_mat1, + x, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dim, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -(static_cast(Y_zero_point)), // mat1_zero_bias + -(static_cast(X_zero_point)), // mat2_zero_bias + static_cast(out_multiplier), // out_multiplier + static_cast(out_shift), // out_shift + static_cast(out_zero_point)); // out_zero_bias + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + } else { + /* Assuming matmul is 2D always */ + WORD32 num_inp_dims = 2; + WORD32 num_out_dims = 2; + + WORD32 p_inp_shape[2]; + WORD32 p_out_shape[2]; + WORD32 p_permute_vec[2] = {1, 0}; + + p_inp_shape[0] = leading_dim; + p_inp_shape[1] = in_dim; + p_out_shape[0] = in_dim; + p_out_shape[1] = leading_dim; + + WORD32 ret_val = xa_nn_transpose_8_8( + (int8_t*)y_data_temp, + p_out_shape, + (int8_t*)y, + p_inp_shape, + p_permute_vec, + num_out_dims, + num_inp_dims); + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + + ret_val = xa_nn_matmul_asym8uxasym8u_asym8u( + z, // p_out + y_data_temp, // p_mat1, + x, // p_mat2, + bias_data, // p_bias + out_dim, // rows of p_mat1 + in_dim, // cols of p_mat1 + in_dim, // row_stride of p_mat1 + leading_dim, // vec_count, i.e., rows of p_mat2 + in_dim, // vec_offset of p_mat2. + out_dim, // out_offset, i.e., offset of next output element written + 1, // out_stride, i.e., stride to go to next output row + -(static_cast(Y_zero_point)), // mat1_zero_bias + -(static_cast(X_zero_point)), // mat2_zero_bias + static_cast(out_multiplier), // out_multiplier + static_cast(out_shift), // out_shift + static_cast(out_zero_point)); // out_zero_bias + + ET_CHECK_MSG(ret_val == 0, "An internal error occured"); + } + } +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_relu_asym8s_asym8s_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_relu_asym8s_asym8s_per_tensor_out.cpp new file mode 100644 index 00000000000..deae48d4411 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_relu_asym8s_asym8s_per_tensor_out.cpp @@ -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. + */ + +#include +#include +#include + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; + +void quantized_relu_asym8s_asym8s_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const int64_t in_zero_point, + const int64_t out_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + Tensor& output) { + const int8_t* __restrict__ input_data = input.const_data_ptr(); + int8_t* __restrict__ output_data = output.mutable_data_ptr(); + + const int32_t out_multipler_int32 = static_cast(out_multiplier); + const int32_t out_shift_int32 = static_cast(out_shift); + + const int32_t ret = xa_nn_vec_relu_asym8s_asym8s( + output_data, + input_data, + in_zero_point, + out_multipler_int32, + out_shift_int32, + out_zero_point, + -128, + 127, + input.numel()); + ET_DCHECK_MSG( + ret == 0, "HiFi quantized_relu_asym8s_asym8s_per_tensor failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_quantized_relu_asym8u_asym8u_per_tensor_out.cpp b/backends/cadence/hifi/operators/op_quantized_relu_asym8u_asym8u_per_tensor_out.cpp new file mode 100644 index 00000000000..8aaca463cf9 --- /dev/null +++ b/backends/cadence/hifi/operators/op_quantized_relu_asym8u_asym8u_per_tensor_out.cpp @@ -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. + */ + +#include +#include +#include + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; + +void quantized_relu_asym8u_asym8u_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const int64_t in_zero_point, + const int64_t out_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + Tensor& output) { + const uint8_t* __restrict__ input_data = input.const_data_ptr(); + uint8_t* __restrict__ output_data = output.mutable_data_ptr(); + + const int32_t out_multipler_int32 = static_cast(out_multiplier); + const int32_t out_shift_int32 = static_cast(out_shift); + + const int32_t ret = xa_nn_vec_relu_asym8u_asym8u( + output_data, + input_data, + in_zero_point, + out_multipler_int32, + out_shift_int32, + out_zero_point, + 0, + 255, + input.numel()); + ET_DCHECK_MSG( + ret == 0, "HiFi quantized_relu_asym8u_asym8u_per_tensor failed"); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/op_softmax.cpp b/backends/cadence/hifi/operators/op_softmax.cpp index 645b9febef0..be496813ce8 100644 --- a/backends/cadence/hifi/operators/op_softmax.cpp +++ b/backends/cadence/hifi/operators/op_softmax.cpp @@ -72,7 +72,6 @@ Tensor& _softmax_out( if (optimized) { int* p_inp = (int*)in.const_data_ptr(); int* out_data = (int*)out.mutable_data_ptr(); - int num_inp_dims = in.dim(); int num_out_dims = num_inp_dims; @@ -99,6 +98,37 @@ Tensor& _softmax_out( outer_stride = size; + WORD32 ret_val = 0; + + // Check if the input is permuted. If not, then we don't need to transpose + bool is_permuted = false; + for (int i = 0; i < num_inp_dims; i++) { + if (p_permute_vec[i] != i) { + is_permuted = true; + break; + } + } + + if (!is_permuted) { + const float* p_inpf = in.const_data_ptr(); + float* out_dataf = out.mutable_data_ptr(); + + for (size_t outer_idx = 0; outer_idx < outer_size; ++outer_idx) { + size_t outer = outer_idx * outer_stride; + for (size_t inner_idx = 0; inner_idx < stride; ++inner_idx) { + size_t base = outer + inner_idx; + + float* p_in_data = (float*)&p_inpf[base]; + float* p_out_data = (float*)&out_dataf[base]; + + ret_val = xa_nn_vec_softmax_f32_f32(p_out_data, p_in_data, size); + + ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); + } + } + return out; + } + int* p_out = (int*)kernels::allocate_temp_memory(ctx, out.numel() * sizeof(int)); @@ -109,7 +139,7 @@ Tensor& _softmax_out( ET_KERNEL_CHECK(ctx, p_out1 != nullptr, MemoryAllocationFailed, out); - WORD32 ret_val = xa_nn_transpose_32_32( + ret_val = xa_nn_transpose_32_32( p_out, p_out_shape, p_inp, @@ -142,9 +172,7 @@ Tensor& _softmax_out( p_permute_vec, num_out_dims, num_inp_dims); - ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); - return out; } diff --git a/backends/cadence/hifi/operators/op_softmax_f32_f32.cpp b/backends/cadence/hifi/operators/op_softmax_f32_f32.cpp new file mode 100644 index 00000000000..bbcf2c66c3d --- /dev/null +++ b/backends/cadence/hifi/operators/op_softmax_f32_f32.cpp @@ -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. + */ + +#include +#include + +using executorch::aten::ScalarType; +using executorch::aten::Tensor; +using executorch::runtime::KernelRuntimeContext; +using torch::executor::Error; + +namespace cadence { +namespace impl { +namespace HiFi { +namespace native { + +inline Tensor& _softmax_f32_f32_out( + KernelRuntimeContext& ctx, + const Tensor& in, + int64_t dim, + ::executorch::aten::optional half_to_float, + Tensor& out) { + constexpr int kNnlibMaxDim = 16; + + const std::optional& dim_t = dim; + const size_t d = ET_NORMALIZE_IX(dim_t.value(), in.dim()); + const size_t size = in.size(d); + + size_t stride = 1, outer_size = 1; + + size_t outer_stride = 1; + + int* p_inp = (int*)in.const_data_ptr(); + int* out_data = (int*)out.mutable_data_ptr(); + + int num_inp_dims = in.dim(); + int num_out_dims = num_inp_dims; + + int p_inp_shape[kNnlibMaxDim]; + int p_out_shape[kNnlibMaxDim]; + int p_permute_vec[kNnlibMaxDim]; + + for (int i = 0; i < num_inp_dims; i++) + p_inp_shape[i] = in.size(i); + for (int i = 0; i < num_inp_dims; i++) { + if (i == d) + p_permute_vec[i] = num_inp_dims - 1; + else if (i == (num_inp_dims - 1)) + p_permute_vec[num_inp_dims - 1] = d; + else + p_permute_vec[i] = i; + + p_out_shape[i] = p_inp_shape[p_permute_vec[i]]; + + if (i != d) + outer_size = outer_size * p_inp_shape[i]; + } + + outer_stride = size; + + WORD32 ret_val = 0; + + // Check if the input is permuted. If not, then we don't need to transpose + bool is_permuted = false; + for (int i = 0; i < num_inp_dims; i++) { + if (p_permute_vec[i] != i) { + is_permuted = true; + break; + } + } + + if (!is_permuted) { + const float* p_inpf = in.const_data_ptr(); + float* out_dataf = out.mutable_data_ptr(); + + for (size_t outer_idx = 0; outer_idx < outer_size; ++outer_idx) { + size_t outer = outer_idx * outer_stride; + for (size_t inner_idx = 0; inner_idx < stride; ++inner_idx) { + size_t base = outer + inner_idx; + + float* p_in_data = (float*)&p_inpf[base]; + float* p_out_data = (float*)&out_dataf[base]; + + ret_val = xa_nn_vec_softmax_f32_f32(p_out_data, p_in_data, size); + + ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); + } + } + return out; + } + + int* p_out = + (int*)kernels::allocate_temp_memory(ctx, out.numel() * sizeof(int)); + + ET_KERNEL_CHECK(ctx, p_out != nullptr, MemoryAllocationFailed, out); + + int* p_out1 = + (int*)kernels::allocate_temp_memory(ctx, out.numel() * sizeof(int)); + + ET_KERNEL_CHECK(ctx, p_out1 != nullptr, MemoryAllocationFailed, out); + + ret_val = xa_nn_transpose_32_32( + p_out, + p_out_shape, + p_inp, + p_inp_shape, + p_permute_vec, + num_out_dims, + num_inp_dims); + + ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); + + for (size_t outer_idx = 0; outer_idx < outer_size; ++outer_idx) { + size_t outer = outer_idx * outer_stride; + for (size_t inner_idx = 0; inner_idx < stride; ++inner_idx) { + size_t base = outer + inner_idx; + + float* p_in_data = (float*)&p_out[base]; + float* p_out_data = (float*)&p_out1[base]; + + ret_val = xa_nn_vec_softmax_f32_f32(p_out_data, p_in_data, size); + + ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); + } + } + + ret_val = xa_nn_transpose_32_32( + out_data, + p_inp_shape, + p_out1, + p_out_shape, + p_permute_vec, + num_out_dims, + num_inp_dims); + + ET_KERNEL_CHECK(ctx, ret_val == 0, Internal, out); + + return out; +} + +Tensor& softmax_f32_f32_out( + KernelRuntimeContext& ctx, + const Tensor& in, + int64_t dim, + ::executorch::aten::optional half_to_float, + Tensor& out) { + return _softmax_f32_f32_out(ctx, in, dim, half_to_float, out); +} + +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/operators.h b/backends/cadence/hifi/operators/operators.h index 85a71dd5092..5b8a1e253c1 100644 --- a/backends/cadence/hifi/operators/operators.h +++ b/backends/cadence/hifi/operators/operators.h @@ -84,7 +84,7 @@ void quantized_linear_per_tensor_out( const ::executorch::aten::optional<::executorch::aten::Tensor>& offset, ::executorch::aten::Tensor& out); -void quantized_conv_out( +void quantized_conv_nhwc_out( ::executorch::runtime::KernelRuntimeContext& ctx, const ::executorch::aten::Tensor& input, const ::executorch::aten::Tensor& weight, @@ -100,10 +100,27 @@ void quantized_conv_out( int64_t output_zero_point, const ::executorch::aten::Tensor& out_multiplier, const ::executorch::aten::Tensor& out_shift, - bool channel_last, ::executorch::aten::Tensor& out); -void quantized_conv_per_tensor_out( +void quantized_conv_nchw_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + const ::executorch::aten::Tensor& input, + const ::executorch::aten::Tensor& weight, + const ::executorch::aten::Tensor& bias, + ::executorch::aten::IntArrayRef stride, + ::executorch::aten::IntArrayRef padding, + ::executorch::aten::IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const ::executorch::aten::Tensor& weight_zero_point, + const ::executorch::aten::Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + const ::executorch::aten::Tensor& out_multiplier, + const ::executorch::aten::Tensor& out_shift, + ::executorch::aten::Tensor& out); + +void quantized_conv_nchw_per_tensor_out( ::executorch::runtime::KernelRuntimeContext& ctx, const ::executorch::aten::Tensor& input, const ::executorch::aten::Tensor& weight, @@ -119,7 +136,60 @@ void quantized_conv_per_tensor_out( int64_t output_zero_point, int64_t out_multiplier, int64_t out_shift, - bool channel_last, + ::executorch::aten::Tensor& out); + +void quantized_conv_nhwc_per_tensor_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + const ::executorch::aten::Tensor& input, + const ::executorch::aten::Tensor& weight, + const ::executorch::aten::Tensor& bias, + ::executorch::aten::IntArrayRef stride, + ::executorch::aten::IntArrayRef padding, + ::executorch::aten::IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + int64_t out_multiplier, + int64_t out_shift, + ::executorch::aten::Tensor& out); + +::executorch::aten::Tensor& cat_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + ::executorch::aten::ArrayRef<::executorch::aten::Tensor> tensors, + int64_t dim, + ::executorch::aten::Tensor& out); + +::executorch::aten::Tensor& permute_copy_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + const ::executorch::aten::Tensor& in, + ::executorch::aten::IntArrayRef dims, + ::executorch::aten::Tensor& out); + +void quantized_add_asym8sxasym8s_asym8s_per_tensor_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + const ::executorch::aten::Tensor& X, + double X_scale, + int64_t X_zero_point, + const ::executorch::aten::Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + ::executorch::aten::Tensor& out); + +void quantized_add_asym8uxasym8u_asym8u_per_tensor_out( + ::executorch::runtime::KernelRuntimeContext& ctx, + const ::executorch::aten::Tensor& X, + double X_scale, + int64_t X_zero_point, + const ::executorch::aten::Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, ::executorch::aten::Tensor& out); } // namespace native diff --git a/backends/cadence/hifi/operators/targets.bzl b/backends/cadence/hifi/operators/targets.bzl index bd9658cc2f9..d310396c262 100644 --- a/backends/cadence/hifi/operators/targets.bzl +++ b/backends/cadence/hifi/operators/targets.bzl @@ -63,12 +63,33 @@ OPERATORS = [ "ne", "permute_copy", "pow", - "quantized_conv_out", + "quantized_conv_nchw_out", + "quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out", + "quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out", + "quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out", + "quantized_conv_nhwc_out", + "quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out", + "quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out", + "quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out", + "quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out", "quantized_fully_connected_out", + "quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out", + "quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out", "quantized_layer_norm", "quantized_linear_out", + "quantized_linear_asym8sxasym8s_asym8s_per_tensor_out", + "quantized_linear_asym8uxasym8u_asym8u_per_tensor_out", "quantized_matmul_out", + "quantized_matmul_asym8sxasym8s_asym8s_out", + "quantized_matmul_asym8uxasym8u_asym8u_out", "quantized_relu_out", + "quantized_relu_asym8s_asym8s_per_tensor_out", + "quantized_relu_asym8u_asym8u_per_tensor_out", "quantize_per_tensor", "remainder", "rsqrt", @@ -76,6 +97,7 @@ OPERATORS = [ "sigmoid", "slice_copy", "softmax", + "softmax_f32_f32", "split_with_sizes_copy", "sub", "tanh", diff --git a/backends/cadence/hifi/operators/tests/test_op_cat.cpp b/backends/cadence/hifi/operators/tests/test_op_cat.cpp new file mode 100644 index 00000000000..2f012ed6c81 --- /dev/null +++ b/backends/cadence/hifi/operators/tests/test_op_cat.cpp @@ -0,0 +1,136 @@ +/* + * 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 cadence { +namespace impl { +namespace HiFi { +namespace native { +namespace { + +using ::executorch::aten::ArrayRef; +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::aten::TensorImpl; +using ::executorch::runtime::Error; +using ::executorch::runtime::KernelRuntimeContext; +using ::executorch::runtime::runtime_init; +using ::executorch::runtime::testing::TensorFactory; + +class HiFiCatTest : public OperatorTest { + public: + protected: + Tensor& cat_out(ArrayRef tensors, int64_t dim, Tensor& out) { + return ::cadence::impl::HiFi::native::cat_out(context_, tensors, dim, out); + } +}; + +TEST_F(HiFiCatTest, FloatCatDim0Test) { + TensorFactory tf; + Tensor a = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + Tensor b = tf.make({1, 3}, {7.0, 8.0, 9.0}); + Tensor c = tf.make({2, 3}, {10.0, 11.0, 12.0, 13.0, 14.0, 15.0}); + + Tensor expected = tf.make( + {5, 3}, + {1.0, + 2.0, + 3.0, + 4.0, + 5.0, + 6.0, + 7.0, + 8.0, + 9.0, + 10.0, + 11.0, + 12.0, + 13.0, + 14.0, + 15.0}); + + Tensor out = tf.zeros({5, 3}); + std::vector tensors = {a, b, c}; + + cat_out(ArrayRef(tensors.data(), tensors.size()), 0, out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiCatTest, FloatCatDim1Test) { + TensorFactory tf; + Tensor a = tf.make({2, 2}, {1.0, 2.0, 3.0, 4.0}); + Tensor b = tf.make({2, 1}, {5.0, 6.0}); + Tensor c = tf.make({2, 3}, {7.0, 8.0, 9.0, 10.0, 11.0, 12.0}); + + Tensor expected = tf.make( + {2, 6}, {1.0, 2.0, 5.0, 7.0, 8.0, 9.0, 3.0, 4.0, 6.0, 10.0, 11.0, 12.0}); + + Tensor out = tf.zeros({2, 6}); + std::vector tensors = {a, b, c}; + + cat_out(ArrayRef(tensors.data(), tensors.size()), 1, out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiCatTest, IntCatDim0Test) { + TensorFactory tf; + Tensor a = tf.make({2, 3}, {1, 2, 3, 4, 5, 6}); + Tensor b = tf.make({1, 3}, {7, 8, 9}); + + Tensor expected = tf.make({3, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + + Tensor out = tf.zeros({3, 3}); + std::vector tensors = {a, b}; + cat_out(ArrayRef(tensors.data(), tensors.size()), 0, out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiCatTest, SingleTensorTest) { + TensorFactory tf; + Tensor a = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + Tensor expected = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + + Tensor out = tf.zeros({2, 3}); + std::vector tensors = {a}; + cat_out(ArrayRef(tensors.data(), tensors.size()), 0, out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiCatTest, ThreeDimensionalCatTest) { + TensorFactory tf; + Tensor a = tf.make({2, 2, 2}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); + Tensor b = tf.make({2, 2, 1}, {9.0, 10.0, 11.0, 12.0}); + + Tensor expected = tf.make( + {2, 2, 3}, + {1.0, 2.0, 9.0, 3.0, 4.0, 10.0, 5.0, 6.0, 11.0, 7.0, 8.0, 12.0}); + + Tensor out = tf.zeros({2, 2, 3}); + std::vector tensors = {a, b}; + + cat_out(ArrayRef(tensors.data(), tensors.size()), 2, out); + EXPECT_TENSOR_EQ(out, expected); +} + +} // namespace +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/tests/test_op_permute_copy.cpp b/backends/cadence/hifi/operators/tests/test_op_permute_copy.cpp new file mode 100644 index 00000000000..a549fac786e --- /dev/null +++ b/backends/cadence/hifi/operators/tests/test_op_permute_copy.cpp @@ -0,0 +1,232 @@ +/* + * 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 cadence { +namespace impl { +namespace HiFi { +namespace native { +namespace { + +using ::executorch::aten::IntArrayRef; +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::aten::TensorImpl; +using ::executorch::runtime::Error; +using ::executorch::runtime::KernelRuntimeContext; +using ::executorch::runtime::runtime_init; +using ::executorch::runtime::testing::TensorFactory; + +class HiFiPermuteCopyTest : public OperatorTest { + public: + protected: + Tensor& permute_copy_out(const Tensor& in, IntArrayRef dims, Tensor& out) { + return ::cadence::impl::HiFi::native::permute_copy_out( + context_, in, dims, out); + } +}; + +TEST_F(HiFiPermuteCopyTest, FloatPermute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + Tensor expected = tf.make({3, 2}, {1.0, 4.0, 2.0, 5.0, 3.0, 6.0}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, IntPermute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6}); + Tensor expected = tf.make({3, 2}, {1, 4, 2, 5, 3, 6}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, Int8Permute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6}); + Tensor expected = tf.make({3, 2}, {1, 4, 2, 5, 3, 6}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, UInt8Permute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6}); + Tensor expected = tf.make({3, 2}, {1, 4, 2, 5, 3, 6}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, DoublePermute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + Tensor expected = tf.make({3, 2}, {1.0, 4.0, 2.0, 5.0, 3.0, 6.0}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, Long8Permute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1, 2, 3, 4, 5, 6}); + Tensor expected = tf.make({3, 2}, {1, 4, 2, 5, 3, 6}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, BoolPermute2DTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {true, false, true, false, true, false}); + Tensor expected = tf.make({3, 2}, {true, false, false, true, true, false}); + + Tensor out = tf.zeros({3, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, Float3DPermuteTest) { + TensorFactory tf; + Tensor in = tf.make({2, 2, 2}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); + Tensor expected = + tf.make({2, 2, 2}, {1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0}); + + Tensor out = tf.zeros({2, 2, 2}); + std::vector dims = {2, 0, 1}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, Float4DPermuteTest) { + TensorFactory tf; + Tensor in = tf.make({1, 2, 2, 2}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); + Tensor expected = + tf.make({2, 1, 2, 2}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}); + + Tensor out = tf.zeros({2, 1, 2, 2}); + std::vector dims = {1, 0, 2, 3}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, IdentityPermuteTest) { + TensorFactory tf; + Tensor in = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + Tensor expected = tf.make({2, 3}, {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}); + + Tensor out = tf.zeros({2, 3}); + std::vector dims = {0, 1}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, LargeTensorPermuteTest) { + TensorFactory tf; + std::vector input_data; + for (int i = 0; i < 60; ++i) { + input_data.push_back(static_cast(i + 1)); + } + Tensor in = tf.make({3, 4, 5}, input_data); + + // Permute: [3, 4, 5] -> [5, 3, 4] with dims [2, 0, 1] + std::vector expected_data(60); + for (int i = 0; i < 3; ++i) { + for (int j = 0; j < 4; ++j) { + for (int k = 0; k < 5; ++k) { + int old_idx = i * 20 + j * 5 + k; + int new_idx = k * 12 + i * 4 + j; + expected_data[new_idx] = static_cast(old_idx + 1); + } + } + } + + Tensor expected = tf.make({5, 3, 4}, expected_data); + Tensor out = tf.zeros({5, 3, 4}); + std::vector dims = {2, 0, 1}; + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_TENSOR_EQ(out, expected); +} + +TEST_F(HiFiPermuteCopyTest, HighDimPermuteTest) { + TensorFactory tf; + std::vector shape = {2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2}; + std::vector input_data = {1.0, 2.0, 3.0, 4.0}; + Tensor in = tf.make(shape, input_data); + + // Simple transpose: swap first and last dimension + std::vector dims(16); + for (int i = 0; i < 16; ++i) { + dims[i] = i; + } + std::swap(dims[0], dims[15]); + Tensor out = tf.zeros(shape); + + permute_copy_out(in, IntArrayRef(dims.data(), dims.size()), out); + EXPECT_DOUBLE_EQ(out.const_data_ptr()[0], 1.0); + EXPECT_DOUBLE_EQ(out.const_data_ptr()[1], 3.0); + EXPECT_DOUBLE_EQ(out.const_data_ptr()[2], 2.0); + EXPECT_DOUBLE_EQ(out.const_data_ptr()[3], 4.0); +} + +TEST_F(HiFiPermuteCopyTest, MixedDataTypesTest) { + TensorFactory tf_short; + Tensor in_short = tf_short.make({2, 2}, {1, 2, 3, 4}); + Tensor expected_short = tf_short.make({2, 2}, {1, 3, 2, 4}); + Tensor out_short = tf_short.zeros({2, 2}); + std::vector dims = {1, 0}; + + permute_copy_out(in_short, IntArrayRef(dims.data(), dims.size()), out_short); + EXPECT_TENSOR_EQ(out_short, expected_short); +} + +} // namespace +} // namespace native +} // namespace HiFi +} // namespace impl +} // namespace cadence diff --git a/backends/cadence/hifi/operators/tests/test_op_quantize_per_tensor.cpp b/backends/cadence/hifi/operators/tests/test_op_quantize_per_tensor.cpp index c8d5b03ce75..6f910cb76a8 100644 --- a/backends/cadence/hifi/operators/tests/test_op_quantize_per_tensor.cpp +++ b/backends/cadence/hifi/operators/tests/test_op_quantize_per_tensor.cpp @@ -118,8 +118,8 @@ TEST_F(HiFiQuantizePerTensorTest, CheckSingleElementIntQuantize) { constexpr int64_t kQuantMin = std::numeric_limits::min(); constexpr int64_t kQuantMax = std::numeric_limits::max(); constexpr float kInputValue = 100.0f; - constexpr int32_t kExpectedOutputValue = - static_cast(kInputValue / kScale + kZeroPoint); + constexpr int32_t kExpectedOutputValue = static_cast( + static_cast(kInputValue) / kScale + kZeroPoint); quantize_per_tensor_out( tf.make(sizes, {kInputValue}), @@ -144,8 +144,8 @@ TEST_F(HiFiQuantizePerTensorTest, CheckSingleElementUInt16Quantize) { constexpr int64_t kQuantMin = std::numeric_limits::min(); constexpr int64_t kQuantMax = std::numeric_limits::max(); constexpr float kInputValue = 100.0f; - constexpr uint16_t kExpectedOutputValue = - static_cast(kInputValue / kScale + kZeroPoint); + constexpr uint16_t kExpectedOutputValue = static_cast( + static_cast(kInputValue) / kScale + kZeroPoint); quantize_per_tensor_out( tf.make(sizes, {kInputValue}), diff --git a/backends/cadence/reference/kernels/CMakeLists.txt b/backends/cadence/reference/kernels/CMakeLists.txt index 3fe0fe2101f..5af049418ce 100644 --- a/backends/cadence/reference/kernels/CMakeLists.txt +++ b/backends/cadence/reference/kernels/CMakeLists.txt @@ -8,9 +8,10 @@ add_library(cadence_kernels kernels.cpp) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. -${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) -target_include_directories(cadence_kernels PUBLIC . - ${_common_include_directories} +target_include_directories( + cadence_kernels PUBLIC . ${_common_include_directories} ) diff --git a/backends/cadence/reference/kernels/kernels.cpp b/backends/cadence/reference/kernels/kernels.cpp index 0961b1ac658..ad8746f51eb 100644 --- a/backends/cadence/reference/kernels/kernels.cpp +++ b/backends/cadence/reference/kernels/kernels.cpp @@ -7,11 +7,10 @@ */ #include -#include #include +#include #include #include -#include namespace impl { namespace reference { @@ -20,10 +19,18 @@ namespace kernels { // Quantize a fp32 value to an int8_t/uint8_t value template T quantize(const float x, float scale, int32_t zero_point) { - constexpr float min_val = std::numeric_limits::min(); - constexpr float max_val = std::numeric_limits::max(); - float tmp = roundf(x * scale + zero_point); - return std::max(std::min(tmp, max_val), min_val); + // constexpr float min_val = std::numeric_limits::min(); + // constexpr float max_val = std::numeric_limits::max(); + // float tmp = roundf(x * scale + zero_point); + // return std::max(std::min(tmp, max_val), min_val); + // Match Executorch CPU kernel implementation at + // https://fburl.com/code/fxizw6u6 + int64_t qvalue; + qvalue = static_cast(zero_point + std::nearbyint(scale * x)); + + qvalue = std::max(qvalue, std::numeric_limits::min()); + qvalue = std::min(qvalue, std::numeric_limits::max()); + return static_cast(qvalue); } // Quantize an fp32 array to an int8_t/uint8_t array @@ -58,36 +65,6 @@ void dequantize( } } -// Requantize the int8_t/uint8_t in value to a uint8_t/int8_t out value. -// The scale and zero_point for requantization are in the args. -template -OT requantize( - const IT in, - float in_scale, - int32_t in_zero_point, - float inv_out_scale, - int32_t out_zero_point) { - float dequant = dequantize(in, in_scale, in_zero_point); - return quantize(dequant, inv_out_scale, out_zero_point); -} - -// Requantize the int8_t/uint8_t in array to a uint8_t/int8_t out array. -// The scale and zero_point for requantization are in the args. -template -void requantize( - OT* __restrict__ out, - const IT* __restrict__ in, - float in_scale, - int32_t in_zero_point, - float inv_out_scale, - int32_t out_zero_point, - size_t size) { - for (size_t i = 0; i < size; ++i) { - out[i] = requantize( - in[i], in_scale, in_zero_point, inv_out_scale, out_zero_point); - } -} - // explicit template instantiation #define typed_quantize_val(dtype) \ @@ -96,7 +73,6 @@ typed_quantize_val(int8_t); typed_quantize_val(uint8_t); typed_quantize_val(int16_t); typed_quantize_val(uint16_t); -typed_quantize_val(int32_t); #undef typed_quantize_val #define typed_quantize_vec(dtype) \ @@ -110,7 +86,6 @@ typed_quantize_vec(int8_t); typed_quantize_vec(uint8_t); typed_quantize_vec(int16_t); typed_quantize_vec(uint16_t); -typed_quantize_vec(int32_t); #undef typed_quantize_vec #define typed_dequantize_val(dtype) \ @@ -119,7 +94,6 @@ typed_dequantize_val(int8_t); typed_dequantize_val(uint8_t); typed_dequantize_val(int16_t); typed_dequantize_val(uint16_t); -typed_dequantize_val(int32_t); #undef typed_dequantize_val #define typed_dequantize_vec(dtype) \ @@ -133,61 +107,8 @@ typed_dequantize_vec(int8_t); typed_dequantize_vec(uint8_t); typed_dequantize_vec(int16_t); typed_dequantize_vec(uint16_t); -typed_dequantize_vec(int32_t); #undef typed_dequantize_vec -#define typed_requantize_val(itype, otype) \ - template otype requantize( \ - const itype in, \ - float in_scale, \ - int32_t in_zero_point, \ - float inv_out_scale, \ - int32_t out_zero_point); -typed_requantize_val(int8_t, int8_t); -typed_requantize_val(int8_t, uint8_t); -typed_requantize_val(int8_t, int16_t); -typed_requantize_val(int8_t, uint16_t); -typed_requantize_val(uint8_t, int8_t); -typed_requantize_val(uint8_t, uint8_t); -typed_requantize_val(uint8_t, int16_t); -typed_requantize_val(uint8_t, uint16_t); -typed_requantize_val(int16_t, int8_t); -typed_requantize_val(int16_t, uint8_t); -typed_requantize_val(int16_t, int16_t); -typed_requantize_val(int16_t, uint16_t); -typed_requantize_val(uint16_t, int8_t); -typed_requantize_val(uint16_t, uint8_t); -typed_requantize_val(uint16_t, int16_t); -typed_requantize_val(uint16_t, uint16_t); -#undef typed_requantize_val - -#define typed_requantize_vec(itype, otype) \ - template void requantize( \ - otype* __restrict__ out, \ - const itype* __restrict__ in, \ - float in_scale, \ - int32_t in_zero_point, \ - float inv_out_scale, \ - int32_t out_zero_point, \ - size_t size); -typed_requantize_vec(int8_t, int8_t); -typed_requantize_vec(int8_t, uint8_t); -typed_requantize_vec(int8_t, int16_t); -typed_requantize_vec(int8_t, uint16_t); -typed_requantize_vec(uint8_t, int8_t); -typed_requantize_vec(uint8_t, uint8_t); -typed_requantize_vec(uint8_t, int16_t); -typed_requantize_vec(uint8_t, uint16_t); -typed_requantize_vec(int16_t, int8_t); -typed_requantize_vec(int16_t, uint8_t); -typed_requantize_vec(int16_t, int16_t); -typed_requantize_vec(int16_t, uint16_t); -typed_requantize_vec(uint16_t, int8_t); -typed_requantize_vec(uint16_t, uint8_t); -typed_requantize_vec(uint16_t, int16_t); -typed_requantize_vec(uint16_t, uint16_t); -#undef typed_requantize_vec - }; // namespace kernels }; // namespace reference }; // namespace impl diff --git a/backends/cadence/reference/kernels/targets.bzl b/backends/cadence/reference/kernels/targets.bzl index d3fe3fa39db..d50cfe8f130 100644 --- a/backends/cadence/reference/kernels/targets.bzl +++ b/backends/cadence/reference/kernels/targets.bzl @@ -10,6 +10,7 @@ def define_common_targets(): ], visibility = [ "//executorch/backends/cadence/...", + "@EXECUTORCH_CLIENTS", ], platforms = CXX, ) diff --git a/backends/cadence/reference/operators/CMakeLists.txt b/backends/cadence/reference/operators/CMakeLists.txt index 6a71af012e4..ea5b699f441 100644 --- a/backends/cadence/reference/operators/CMakeLists.txt +++ b/backends/cadence/reference/operators/CMakeLists.txt @@ -67,8 +67,9 @@ target_link_libraries(aten_ops_cadence PUBLIC executorch) target_link_libraries(aten_ops_cadence PRIVATE cadence_kernels) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. -${EXECUTORCH_ROOT}/runtime/core/portable_type/c10) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 +) target_include_directories( aten_ops_cadence PUBLIC ${ROOT_DIR}/.. ${CMAKE_BINARY_DIR} @@ -79,7 +80,8 @@ target_include_directories( add_library( custom_ops "quantized_linear_out.cpp" - "quantized_conv_out.cpp" + "quantized_conv_nchw_out.cpp" + "quantized_conv_nhwc_out.cpp" "quantized_relu_out.cpp" "quantized_layer_norm.cpp" "quantize_per_tensor.cpp" diff --git a/backends/cadence/reference/operators/dequantize_per_tensor.cpp b/backends/cadence/reference/operators/dequantize_per_tensor.cpp index 1f48fc78201..f53292e312d 100644 --- a/backends/cadence/reference/operators/dequantize_per_tensor.cpp +++ b/backends/cadence/reference/operators/dequantize_per_tensor.cpp @@ -13,9 +13,10 @@ namespace impl { namespace reference { namespace native { -using executorch::aten::ScalarType; -using executorch::aten::Tensor; -using executorch::runtime::KernelRuntimeContext; +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using ::impl::reference::kernels::dequantize; void dequantize_per_tensor_out( KernelRuntimeContext& context, @@ -31,26 +32,18 @@ void dequantize_per_tensor_out( if (input.scalar_type() == ScalarType::Byte) { const uint8_t* input_data = input.const_data_ptr(); - impl::reference::kernels::dequantize( - out_data, input_data, scale, zero_point, numel); + dequantize(out_data, input_data, scale, zero_point, numel); } else if (input.scalar_type() == ScalarType::Char) { const int8_t* input_data = input.const_data_ptr(); - impl::reference::kernels::dequantize( - out_data, input_data, scale, zero_point, numel); + dequantize(out_data, input_data, scale, zero_point, numel); } else if ( input.scalar_type() == ScalarType::Bits16 || input.scalar_type() == ScalarType::UInt16) { const uint16_t* input_data = input.const_data_ptr(); - impl::reference::kernels::dequantize( - out_data, input_data, scale, zero_point, numel); + dequantize(out_data, input_data, scale, zero_point, numel); } else if (input.scalar_type() == ScalarType::Short) { const int16_t* input_data = input.const_data_ptr(); - impl::reference::kernels::dequantize( - out_data, input_data, scale, zero_point, numel); - } else if (input.scalar_type() == ScalarType::Int) { - const int32_t* input_data = input.const_data_ptr(); - impl::reference::kernels::dequantize( - out_data, input_data, scale, zero_point, numel); + dequantize(out_data, input_data, scale, zero_point, numel); } else { ET_CHECK_MSG( false, diff --git a/backends/cadence/reference/operators/requantize_out.cpp b/backends/cadence/reference/operators/op_requantize_out.cpp similarity index 86% rename from backends/cadence/reference/operators/requantize_out.cpp rename to backends/cadence/reference/operators/op_requantize_out.cpp index e57a6e1614e..5cb9ee3943d 100644 --- a/backends/cadence/reference/operators/requantize_out.cpp +++ b/backends/cadence/reference/operators/op_requantize_out.cpp @@ -86,18 +86,15 @@ Tensor& requantize_out( torch::executor::toString(out.scalar_type()), torch::executor::toString(out_dtype)); -#define typed_requantize(ctype, dtype) \ - const ctype* input_data = input.const_data_ptr(); \ - dtype* out_data = out.mutable_data_ptr(); \ - kernels::requantize( \ - out_data, \ - input_data, \ - in_scale, \ - in_zero_point, \ - 1.0 / out_scale, \ - out_zero_point, \ - numel); - +#define typed_requantize(ctype, dtype) \ + const ctype* input_data = input.const_data_ptr(); \ + dtype* out_data = out.mutable_data_ptr(); \ + for (size_t i = 0; i < numel; ++i) { \ + float dequant = \ + kernels::dequantize(input_data[i], in_scale, in_zero_point); \ + out_data[i] = \ + kernels::quantize(dequant, 1 / out_scale, out_zero_point); \ + }; #define typed_requantize_in(ctype) \ switch (out_dtype) { \ case ScalarType::Byte: { \ @@ -190,17 +187,15 @@ Tensor& requantize_per_tensor_out( torch::executor::toString(out.scalar_type()), torch::executor::toString(out_dtype)); -#define typed_requantize(ctype, dtype) \ - const ctype* input_data = input.const_data_ptr(); \ - dtype* out_data = out.mutable_data_ptr(); \ - kernels::requantize( \ - out_data, \ - input_data, \ - static_cast(in_scale), \ - static_cast(in_zero_point), \ - 1.0 / static_cast(out_scale), \ - static_cast(out_zero_point), \ - numel); +#define typed_requantize(ctype, dtype) \ + const ctype* input_data = input.const_data_ptr(); \ + dtype* out_data = out.mutable_data_ptr(); \ + for (size_t i = 0; i < numel; ++i) { \ + float dequant = \ + kernels::dequantize(input_data[i], in_scale, in_zero_point); \ + out_data[i] = \ + kernels::quantize(dequant, 1 / out_scale, out_zero_point); \ + }; #define typed_requantize_in(ctype) \ switch (out_dtype) { \ diff --git a/backends/cadence/reference/operators/quantize_per_tensor.cpp b/backends/cadence/reference/operators/quantize_per_tensor.cpp index 9debf37364e..8f8cc961dd8 100644 --- a/backends/cadence/reference/operators/quantize_per_tensor.cpp +++ b/backends/cadence/reference/operators/quantize_per_tensor.cpp @@ -49,10 +49,6 @@ void quantize_per_tensor_out( int16_t* out_data = out.mutable_data_ptr(); impl::reference::kernels::quantize( out_data, input_data, 1. / scale, zero_point, numel); - } else if (out.scalar_type() == ScalarType::Int) { - int32_t* out_data = out.mutable_data_ptr(); - impl::reference::kernels::quantize( - out_data, input_data, 1. / scale, zero_point, numel); } else { ET_CHECK_MSG( false, diff --git a/backends/cadence/reference/operators/quantized_add_out.cpp b/backends/cadence/reference/operators/quantized_add_out.cpp new file mode 100644 index 00000000000..7e5834de7bf --- /dev/null +++ b/backends/cadence/reference/operators/quantized_add_out.cpp @@ -0,0 +1,192 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +namespace impl { +namespace reference { +namespace native { + +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; +using ::impl::reference::kernels::dequantize; +using ::impl::reference::kernels::quantize; + +template +void quantized_add_per_tensor_impl( + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + const T* __restrict__ X_data = X.const_data_ptr(); + const T* __restrict__ Y_data = Y.const_data_ptr(); + T* __restrict__ out_data = out.mutable_data_ptr(); + + ssize_t Y_numel = Y.numel(); + ssize_t X_numel = X.numel(); + ssize_t out_numel = out.numel(); + + float X_scale_f = static_cast(X_scale); + float Y_scale_f = static_cast(Y_scale); + float out_scale_f = static_cast(out_scale); + int32_t X_zero_point_i32 = static_cast(X_zero_point); + int32_t Y_zero_point_i32 = static_cast(Y_zero_point); + int32_t out_zero_point_i32 = static_cast(out_zero_point); + + float inv_out_scale = 1.0f / out_scale_f; + + // Simple case: tensors have the same shape, no broadcasting + if (X_numel == Y_numel && Y_numel == out_numel) { + for (size_t i = 0; i < X_numel; ++i) { + float x = dequantize(X_data[i], X_scale_f, X_zero_point_i32); + float y = dequantize(Y_data[i], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } + // Y is a scalar tensor + else if (Y_numel == 1) { + float y = dequantize(Y_data[0], Y_scale_f, Y_zero_point_i32); + for (size_t i = 0; i < X_numel; ++i) { + float x = dequantize(X_data[i], X_scale_f, X_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } + // X is a scalar tensor + else if (X_numel == 1) { + float x = dequantize(X_data[0], X_scale_f, X_zero_point_i32); + for (size_t i = 0; i < Y_numel; ++i) { + float y = dequantize(Y_data[i], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } + // General broadcasting case - simplified implementation + else { + for (ssize_t i = 0; i < out_numel; ++i) { + // Simple broadcasting: repeat elements as needed + size_t x_idx = (X_numel == 1) ? 0 : i % X_numel; + size_t y_idx = (Y_numel == 1) ? 0 : i % Y_numel; + + float x = dequantize(X_data[x_idx], X_scale_f, X_zero_point_i32); + float y = dequantize(Y_data[y_idx], Y_scale_f, Y_zero_point_i32); + float z = x + y; + out_data[i] = quantize(z, inv_out_scale, out_zero_point_i32); + } + } +} + +// Generic quantized add with type dispatch +void quantized_add_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + (void)ctx; + + executorch::aten::ScalarType dtype = X.scalar_type(); + switch (dtype) { + case executorch::aten::ScalarType::Byte: + quantized_add_per_tensor_impl( + X, + X_scale, + X_zero_point, + Y, + Y_scale, + Y_zero_point, + out_scale, + out_zero_point, + out); + break; + case executorch::aten::ScalarType::Char: + quantized_add_per_tensor_impl( + X, + X_scale, + X_zero_point, + Y, + Y_scale, + Y_zero_point, + out_scale, + out_zero_point, + out); + break; + default: + ET_CHECK_MSG( + false, "Unhandled input dtype %hhd", static_cast(dtype)); + } +} + +// int8-specific quantized add +void quantized_add_asym8sxasym8s_asym8s_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + (void)ctx; + + quantized_add_per_tensor_impl( + X, + X_scale, + X_zero_point, + Y, + Y_scale, + Y_zero_point, + out_scale, + out_zero_point, + out); +} + +// uint8-specific quantized add +void quantized_add_asym8uxasym8u_asym8u_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& X, + double X_scale, + int64_t X_zero_point, + const Tensor& Y, + double Y_scale, + int64_t Y_zero_point, + double out_scale, + int64_t out_zero_point, + Tensor& out) { + (void)ctx; + + quantized_add_per_tensor_impl( + X, + X_scale, + X_zero_point, + Y, + Y_scale, + Y_zero_point, + out_scale, + out_zero_point, + out); +} + +} // namespace native +} // namespace reference +} // namespace impl diff --git a/backends/cadence/reference/operators/quantized_conv_nchw_out.cpp b/backends/cadence/reference/operators/quantized_conv_nchw_out.cpp new file mode 100644 index 00000000000..aefa75d7047 --- /dev/null +++ b/backends/cadence/reference/operators/quantized_conv_nchw_out.cpp @@ -0,0 +1,501 @@ +/* + * 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 + +namespace impl { +namespace reference { +namespace native { + +using ::executorch::aten::IntArrayRef; +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; + +// This implements a generic 2d conv kernel that operates on raw pointers. +// The version handles both quantized and fp32 convolutions. +// The input is of shape [n x c x h x w] +// The weight is of shape [oc x wc x wh x ww], where wc == c +// The output is of shape [n x oc x oh x ow] +// The bias is of shape [oc] +template < + typename IT = float, + typename WT = IT, + typename BT = IT, + typename OT = IT, + bool quantized = false> +__attribute__((noinline)) void conv2d_nchw_core_generic( + // All the arrays + const IT* __restrict__ p_in, + const WT* __restrict__ p_weight, + const BT* __restrict__ p_bias, + OT* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t c, + int32_t h, + int32_t w, + int32_t oc, + int32_t wc, + int32_t wh, + int32_t ww, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Optional args that are only relevant for quantized convolution + // input zero point + IT in_zero_point = 0, + // weight zero point + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + OT out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const IT* in_batch = p_in + _n * c * h * w; + OT* out_batch = p_out + _n * oc * oh * ow; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + OT* out_plane = out_batch + _oc * oh * ow; + const WT* weight_batch = p_weight + _oc * wc * wh * ww; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of size + // icpg x h x w, with a stencil of size icpg x wh x ww, to compute an + // output channel of size 1 x oh x ow. + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to the + // output channel being computed) with the corresponding weight + // channel. + // If the padding is 0, and dilation is 1, then we can remove the + // unnecessary checks, and simplify the code so that it can be + // vectorized by Tensilica compiler. + if (zero_pad_unit_dilation) { + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + const IT* in_plane = in_batch + _ic * h * w; + const WT* weight_plane = weight_batch + (_ic - sic) * wh * ww; + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + int ioff = (_h + _wh) * w + (_w + _ww); + int woff = _wh * ww + _ww; + float lhs = in_plane[ioff] - in_zero_point; + float rhs = weight_plane[woff] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } else { + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + const IT* in_plane = in_batch + _ic * h * w; + const WT* weight_plane = weight_batch + (_ic - sic) * wh * ww; + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + if (((_h + d0 * _wh - p0) >= 0) && + ((_h + d0 * _wh - p0) < h) && + ((_w + d1 * _ww - p1) >= 0) && + ((_w + d1 * _ww - p1) < w)) { + int ioff = + (_h + d0 * _wh - p0) * w + (_w + d1 * _ww - p1); + int woff = _wh * ww + _ww; + float lhs = in_plane[ioff] - in_zero_point; + float rhs = weight_plane[woff] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } + } + if (quantized) { + float val = bias_scale * acc; + out_plane[_oh * ow + _ow] = + ::impl::reference::kernels::quantize( + val, inv_out_scale, out_zero_point); + } else { + out_plane[_oh * ow + _ow] = acc; + } + } + } + } + } + } +} + +// The quantized convolution kernel. in_scale and weight_scale are implicit in +// bias_scale, since it is a product of the two. The kernel will branch to +// quantized::conv1d or quantized::conv2d based on the dimensionality of +// activation tensor. +void quantized_conv_nchw( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, c, h, w] + const int n = input.size(0); + const int c = input.size(1); + const int h = conv1d ? 1 : input.size(2); + const int w = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wc, wh, ww] + const int oc = weight.size(0); + const int wc = weight.size(1); + const int wh = conv1d ? 1 : weight.size(2); + const int ww = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oc, oh, ow] + const int oh = conv1d ? 1 : out.size(2); + const int ow = conv1d ? out.size(2) : out.size(3); + +#define typed_quantized_conv2d_nchw(ctype, dtype) \ + case ScalarType::dtype: { \ + conv2d_nchw_core_generic( \ + input.const_data_ptr(), \ + weight.const_data_ptr(), \ + bias.const_data_ptr(), \ + out.mutable_data_ptr(), \ + n, \ + c, \ + h, \ + w, \ + oc, \ + wc, \ + wh, \ + ww, \ + oh, \ + ow, \ + stride[0], \ + stride[1], \ + padding[0], \ + padding[1], \ + dilation[0], \ + dilation[1], \ + groups, \ + in_zero_point, \ + weight_zero_point, \ + bias_scale, \ + output_scale, \ + (ctype)output_zero_point); \ + break; \ + } + ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nchw); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } + +#undef typed_quantized_conv2d_nchw +} + +void quantized_conv_nchw_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED const Tensor& out_multiplier, + __ET_UNUSED const Tensor& out_shift, + Tensor& out) { + const float bias_scale_float = bias_scale.const_data_ptr()[0]; + const int32_t weight_zero_point_int = + weight_zero_point.const_data_ptr()[0]; + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point_int, + bias_scale_float, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + bool channel_last, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_dilated_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_dilated_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_depthwise_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nchw_depthwise_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nchw( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace reference +} // namespace impl diff --git a/backends/cadence/reference/operators/quantized_conv_nhwc_out.cpp b/backends/cadence/reference/operators/quantized_conv_nhwc_out.cpp new file mode 100644 index 00000000000..26fbc86d5b0 --- /dev/null +++ b/backends/cadence/reference/operators/quantized_conv_nhwc_out.cpp @@ -0,0 +1,488 @@ +/* + * 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 + +namespace impl { +namespace reference { +namespace native { + +using ::executorch::aten::IntArrayRef; +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::runtime::KernelRuntimeContext; + +template < + typename IT = float, + typename WT = IT, + typename BT = IT, + typename OT = IT, + bool quantized = false> +__attribute__((noinline)) void conv2d_nhwc_core_generic( + // All the arrays + const IT* __restrict__ p_in, + const WT* __restrict__ p_weight, + const BT* __restrict__ p_bias, + OT* __restrict__ p_out, + // The array sizes + int32_t n, + int32_t h, + int32_t w, + int32_t c, + int32_t oc, + int32_t wh, + int32_t ww, + int32_t wc, + int32_t oh, + int32_t ow, + // Stride + int16_t s0, + int16_t s1, + // Padding + int16_t p0, + int16_t p1, + // Dilation + int16_t d0, + int16_t d1, + // Group for depthwise conv + int16_t groups, + // Optional args that are only relevant for quantized convolution + // input zero point + IT in_zero_point = 0, + // weight zero point + int32_t weight_zero_point = 0, + float bias_scale = 1, + float out_scale = 1, + OT out_zero_point = 0) { + float inv_out_scale = 1. / out_scale; + bool zero_pad_unit_dilation = d0 == 1 && d1 == 1 && p0 == 0 && p1 == 0; + + // Compute the number of in and out channels per group + const int ocpg = oc / groups; + const int icpg = c / groups; + + // Iterate over all the output batches (i.e., n) + for (int _n = 0; _n < n; ++_n) { + const IT* in_batch = p_in + _n * h * w * c; + OT* out_batch = p_out + _n * oh * ow * oc; + for (int _h = 0, _oh = 0; _oh < oh; _h += s0, ++_oh) { + for (int _w = 0, _ow = 0; _ow < ow; _w += s1, ++_ow) { + OT* out_line = out_batch + (_oh * ow + _ow) * oc; + // Compute separable convolution for each group + for (int _g = 0; _g < groups; ++_g) { + // Identify the input and output channels involved in the computation + // of this group + int sic = _g * icpg; + int soc = _g * ocpg; + // Populate all the output channels in the group + for (int _oc = soc; _oc < soc + ocpg; ++_oc) { + const WT* weight_batch = p_weight + _oc * wh * ww * wc; + // We compute one output channel at a time. The computation can be + // thought of as a stencil computation: we iterate over an input of + // size h x w x icpg, with a stencil of size wh x ww x icpg, to + // compute an output channel of size oh x ow x 1. + float acc = p_bias[_oc]; + // Below is the stencil computation that performs the hadamard + // product+accumulation of each input channel (contributing to + // the output channel being computed) with the corresponding + // weight channel. If the padding is 0, and dilation is 1, then + // we can remove the unnecessary checks, and simplify the code + // so that it can be vectorized by Tensilica compiler.x`` + if (zero_pad_unit_dilation) { + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + const IT* in_line = + in_batch + (_h + _wh) * w * c + (_w + _ww) * c; + const WT* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = in_line[_ic] - in_zero_point; + float rhs = weight_line[_ic - sic] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } else { + for (int _wh = 0; _wh < wh; ++_wh) { + for (int _ww = 0; _ww < ww; ++_ww) { + if (((_h + d0 * _wh - p0) >= 0) && + ((_h + d0 * _wh - p0) < h) && + ((_w + d1 * _ww - p1) >= 0) && + ((_w + d1 * _ww - p1 < w))) { + const IT* in_line = in_batch + + (_h + d0 * _wh - p0) * w * c + (_w + d1 * _ww - p1) * c; + const WT* weight_line = + weight_batch + _wh * ww * wc + _ww * wc; + for (int _ic = sic; _ic < sic + icpg; ++_ic) { + float lhs = in_line[_ic] - in_zero_point; + float rhs = weight_line[_ic - sic] - + (quantized ? weight_zero_point : 0); + acc += lhs * rhs; + } + } + } + } + } + if (quantized) { + float val = bias_scale * acc; + out_line[_oc] = ::impl::reference::kernels::quantize( + val, inv_out_scale, out_zero_point); + } else { + out_line[_oc] = acc; + } + } + } + } + } + } +} + +void quantized_conv_nhwc( + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int16_t groups, + int32_t in_zero_point, + int32_t weight_zero_point, + float bias_scale, + float output_scale, + int32_t output_zero_point, + Tensor& out) { + bool conv1d = input.dim() == 3; + // input = [n, h, w, c] + const int n = input.size(0); + const int h = conv1d ? 1 : input.size(1); + const int w = conv1d ? input.size(1) : input.size(2); + const int c = conv1d ? input.size(2) : input.size(3); + // weight = [oc, wh, ww, wc] + const int oc = weight.size(0); + const int wh = conv1d ? 1 : weight.size(1); + const int ww = conv1d ? weight.size(1) : weight.size(2); + const int wc = conv1d ? weight.size(2) : weight.size(3); + // output = [n, oh, ow, oc] + const int oh = conv1d ? 1 : out.size(1); + const int ow = conv1d ? out.size(1) : out.size(2); + +#define typed_quantized_conv2d_nhwc(ctype, dtype) \ + case ScalarType::dtype: { \ + conv2d_nhwc_core_generic( \ + input.const_data_ptr(), \ + weight.const_data_ptr(), \ + bias.const_data_ptr(), \ + out.mutable_data_ptr(), \ + n, \ + h, \ + w, \ + c, \ + oc, \ + wh, \ + ww, \ + wc, \ + oh, \ + ow, \ + stride[0], \ + stride[1], \ + padding[0], \ + padding[1], \ + dilation[0], \ + dilation[1], \ + groups, \ + in_zero_point, \ + weight_zero_point, \ + bias_scale, \ + output_scale, \ + (ctype)output_zero_point); \ + break; \ + } + ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_conv2d_nhwc); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } + +#undef typed_quantized_conv2d_nhwc +} + +void quantized_conv_nhwc_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + const Tensor& weight_zero_point, + const Tensor& bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED const Tensor& out_multiplier, + __ET_UNUSED const Tensor& out_shift, + Tensor& out) { + const float bias_scale_float = bias_scale.const_data_ptr()[0]; + const int32_t weight_zero_point_int = + weight_zero_point.const_data_ptr()[0]; + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point_int, + bias_scale_float, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + bool channel_last, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_dilated_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_dilated_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_depthwise_asym8sxsym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +void quantized_conv_nhwc_depthwise_asym8uxsym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& input, + const Tensor& weight, + const Tensor& bias, + IntArrayRef stride, + IntArrayRef padding, + IntArrayRef dilation, + int64_t groups, + int64_t in_zero_point, + int64_t weight_zero_point, + double bias_scale, + double output_scale, + int64_t output_zero_point, + __ET_UNUSED int64_t out_multiplier, + __ET_UNUSED int64_t out_shift, + Tensor& out) { + quantized_conv_nhwc( + input, + weight, + bias, + stride, + padding, + dilation, + groups, + in_zero_point, + weight_zero_point, + bias_scale, + output_scale, + output_zero_point, + out); +} + +} // namespace native +} // namespace reference +} // namespace impl diff --git a/backends/cadence/reference/operators/quantized_fully_connected_out.cpp b/backends/cadence/reference/operators/quantized_fully_connected_out.cpp index fe41c2d7e77..136055de70a 100644 --- a/backends/cadence/reference/operators/quantized_fully_connected_out.cpp +++ b/backends/cadence/reference/operators/quantized_fully_connected_out.cpp @@ -92,6 +92,80 @@ void quantized_fully_connected_per_tensor_out( #undef typed_quantized_linear } +void quantized_fully_connected_asym8sxasym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { +#define typed_quantized_linear(ctype, dtype) \ + case ScalarType::dtype: { \ + quantized_linear_per_tensor_( \ + in, \ + weight, \ + bias, \ + in_zero_point, \ + weight_zero_point, \ + out_multiplier, \ + out_shift, \ + out_zero_point, \ + out); \ + break; \ + } + + ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_linear); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } +#undef typed_quantized_linear +} + +void quantized_fully_connected_asym8uxasym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& in, + const Tensor& weight, + const Tensor& bias, + int64_t in_zero_point, + int64_t weight_zero_point, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + __ET_UNUSED const optional& offset, + Tensor& out) { +#define typed_quantized_linear(ctype, dtype) \ + case ScalarType::dtype: { \ + quantized_linear_per_tensor_( \ + in, \ + weight, \ + bias, \ + in_zero_point, \ + weight_zero_point, \ + out_multiplier, \ + out_shift, \ + out_zero_point, \ + out); \ + break; \ + } + + ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_linear); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } +#undef typed_quantized_linear +} + }; // namespace native }; // namespace reference }; // namespace impl diff --git a/backends/cadence/reference/operators/quantized_layer_norm.cpp b/backends/cadence/reference/operators/quantized_layer_norm.cpp index efddecc3f23..64dcb000cc1 100644 --- a/backends/cadence/reference/operators/quantized_layer_norm.cpp +++ b/backends/cadence/reference/operators/quantized_layer_norm.cpp @@ -16,6 +16,8 @@ using ::executorch::aten::ScalarType; using ::executorch::aten::Tensor; using ::executorch::runtime::getLeadingDims; using ::executorch::runtime::KernelRuntimeContext; +using ::impl::reference::kernels::dequantize; +using ::impl::reference::kernels::quantize; namespace impl { namespace reference { @@ -74,10 +76,10 @@ void quantized_layer_norm_per_tensor_( // y[j] = (x[j] - mean) / std * kGamma + kBeta; // Since X is quantized, we dequantize it, compute fp32 result, and // quantize the result to an int8/uint8 value. - float val = kernels::dequantize(x[j], input_scale, input_zero_point); + float val = dequantize(x[j], input_scale, input_zero_point); val = (val - mean) * inv_std * weight_data[j] + bias_data[j]; - y[j] = kernels::quantize(val, output_inv_scale, output_zero_point); + y[j] = quantize(val, output_inv_scale, output_zero_point); } } } diff --git a/backends/cadence/reference/operators/quantized_linear_out.cpp b/backends/cadence/reference/operators/quantized_linear_out.cpp index edd8634d56e..f60c98e5875 100644 --- a/backends/cadence/reference/operators/quantized_linear_out.cpp +++ b/backends/cadence/reference/operators/quantized_linear_out.cpp @@ -154,6 +154,80 @@ void quantized_linear_per_tensor_out( #undef typed_quantized_linear_per_tensor } +void quantized_linear_asym8sxasym8s_asym8s_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& src, + const Tensor& weight, + const Tensor& bias, + const int64_t src_zero_point, + const int64_t weight_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + const int64_t out_zero_point, + __ET_UNUSED const std::optional& offset, + Tensor& out) { +#define typed_quantized_linear_per_tensor(ctype, dtype) \ + case executorch::aten::ScalarType::dtype: { \ + quantized_linear_per_tensor_( \ + src, \ + weight, \ + bias, \ + src_zero_point, \ + weight_zero_point, \ + out_multiplier, \ + out_shift, \ + out_zero_point, \ + out); \ + break; \ + } + + executorch::aten::ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_linear_per_tensor); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", executorch::runtime::toString(dtype)); + } +#undef typed_quantized_linear_per_tensor +} + +void quantized_linear_asym8uxasym8u_asym8u_per_tensor_out( + __ET_UNUSED KernelRuntimeContext& ctx, + const Tensor& src, + const Tensor& weight, + const Tensor& bias, + const int64_t src_zero_point, + const int64_t weight_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + const int64_t out_zero_point, + __ET_UNUSED const std::optional& offset, + Tensor& out) { +#define typed_quantized_linear_per_tensor(ctype, dtype) \ + case executorch::aten::ScalarType::dtype: { \ + quantized_linear_per_tensor_( \ + src, \ + weight, \ + bias, \ + src_zero_point, \ + weight_zero_point, \ + out_multiplier, \ + out_shift, \ + out_zero_point, \ + out); \ + break; \ + } + + executorch::aten::ScalarType dtype = out.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_linear_per_tensor); + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", executorch::runtime::toString(dtype)); + } +#undef typed_quantized_linear_per_tensor +} + }; // namespace native }; // namespace reference }; // namespace impl diff --git a/backends/cadence/reference/operators/quantized_matmul_out.cpp b/backends/cadence/reference/operators/quantized_matmul_out.cpp index cc0fa05351c..3c2070c70dc 100644 --- a/backends/cadence/reference/operators/quantized_matmul_out.cpp +++ b/backends/cadence/reference/operators/quantized_matmul_out.cpp @@ -152,6 +152,56 @@ void quantized_matmul_out( } } +void quantized_matmul_asym8sxasym8s_asym8s_out( + KernelRuntimeContext& ctx, + const Tensor& X, + int64_t X_zero_point, + const Tensor& Y, + int64_t Y_zero_point, + const std::optional& bias, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + bool transposed, + Tensor& out) { + _typed_quantized_matmul( + X, + X_zero_point, + Y, + Y_zero_point, + bias, + out_multiplier, + out_shift, + out_zero_point, + transposed, + out); +} + +void quantized_matmul_asym8uxasym8u_asym8u_out( + KernelRuntimeContext& ctx, + const Tensor& X, + int64_t X_zero_point, + const Tensor& Y, + int64_t Y_zero_point, + const std::optional& bias, + int64_t out_multiplier, + int64_t out_shift, + int64_t out_zero_point, + bool transposed, + Tensor& out) { + _typed_quantized_matmul( + X, + X_zero_point, + Y, + Y_zero_point, + bias, + out_multiplier, + out_shift, + out_zero_point, + transposed, + out); +} + }; // namespace native }; // namespace reference }; // namespace impl diff --git a/backends/cadence/reference/operators/quantized_relu_out.cpp b/backends/cadence/reference/operators/quantized_relu_out.cpp index 7a385849aee..8dab01cf982 100644 --- a/backends/cadence/reference/operators/quantized_relu_out.cpp +++ b/backends/cadence/reference/operators/quantized_relu_out.cpp @@ -129,6 +129,70 @@ void quantized_relu_per_tensor_out( #undef typed_quantized_relu } +void quantized_relu_asym8s_asym8s_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const int64_t in_zero_point, + const int64_t out_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + Tensor& output) { +#define typed_quantized_relu(ctype, dtype) \ + case executorch::aten::ScalarType::dtype: { \ + quantized_relu_per_tensor_out_( \ + ctx, \ + input, \ + in_zero_point, \ + out_zero_point, \ + out_multiplier, \ + out_shift, \ + output); \ + break; \ + } + + executorch::aten::ScalarType dtype = input.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_relu) + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } + +#undef typed_quantized_relu +} + +void quantized_relu_asym8u_asym8u_per_tensor_out( + KernelRuntimeContext& ctx, + const Tensor& input, + const int64_t in_zero_point, + const int64_t out_zero_point, + const int64_t out_multiplier, + const int64_t out_shift, + Tensor& output) { +#define typed_quantized_relu(ctype, dtype) \ + case executorch::aten::ScalarType::dtype: { \ + quantized_relu_per_tensor_out_( \ + ctx, \ + input, \ + in_zero_point, \ + out_zero_point, \ + out_multiplier, \ + out_shift, \ + output); \ + break; \ + } + + executorch::aten::ScalarType dtype = input.scalar_type(); + switch (dtype) { + ET_FORALL_CADENCE_QUANTIZED_TYPES(typed_quantized_relu) + default: + ET_DCHECK_MSG( + false, "Unhandled dtype %s", torch::executor::toString(dtype)); + } + +#undef typed_quantized_relu +} + }; // namespace native }; // namespace reference }; // namespace impl diff --git a/backends/cadence/runtime/TARGETS b/backends/cadence/runtime/TARGETS index 4055f1922a1..9c65c469280 100644 --- a/backends/cadence/runtime/TARGETS +++ b/backends/cadence/runtime/TARGETS @@ -1,9 +1,9 @@ load(":targets.bzl", "define_common_targets") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("odai_jarvis") -python_library( +runtime.python_library( name = "runtime", srcs = [ "__init__.py", diff --git a/backends/cadence/utils/facto_util.py b/backends/cadence/utils/facto_util.py index b896f8a8e89..fd056cd08cc 100644 --- a/backends/cadence/utils/facto_util.py +++ b/backends/cadence/utils/facto_util.py @@ -10,6 +10,8 @@ from functools import lru_cache from typing import List, OrderedDict, Tuple +import facto.specdb.function as fn + import torch from facto.inputgen.argtuple.gen import ArgumentTupleGenerator from facto.inputgen.specs.model import ConstraintProducer as cp @@ -21,22 +23,55 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]: - tensor_constraints = [ - cp.Dtype.In(lambda deps: [torch.int, torch.float]), - cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]), - cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), - cp.Value.Le(lambda deps, dtype, struct: 2**4), - cp.Rank.Ge(lambda deps: 1), - cp.Size.Ge(lambda deps, r, d: 1), - cp.Size.Le(lambda deps, r, d: 2**9), - ] + tensor_constraints = ( + [ + cp.Dtype.In( + lambda deps: [ + torch.int8, + torch.int16, + torch.uint8, + torch.uint16, + torch.int32, + torch.float32, + ] + ), + cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + cp.Rank.Ge(lambda deps: 1), + cp.Size.Ge(lambda deps, r, d: 1), + cp.Size.Le(lambda deps, r, d: 2**9), + cp.Rank.Le(lambda deps: 2**3), + ] + if op_name + not in ( + "slice_copy.Tensor", + "add.Scalar", + "sub.Scalar", + "mul.Scalar", + "div.Tensor", + "neg.default", + ) + else [ + cp.Dtype.In( + lambda deps: [ + torch.int32, + torch.float32, + ] + ), + cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + cp.Rank.Ge(lambda deps: 1), + cp.Size.Ge(lambda deps, r, d: 1), + cp.Size.Le(lambda deps, r, d: 2**9), + cp.Rank.Le(lambda deps: 2**3), + ] + ) match op_name: case "where.self": if index == 0: # condition tensor_constraints = [ cp.Dtype.In(lambda deps: [torch.bool]), - cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]), cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), cp.Value.Le(lambda deps, dtype, struct: 2**4), cp.Rank.Ge(lambda deps: 1), @@ -45,19 +80,36 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]: ] else: tensor_constraints = [ - cp.Dtype.In(lambda deps: [torch.float, torch.int]), - cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]), + cp.Dtype.In( + lambda deps: [ + torch.int8, + torch.int16, + torch.uint8, + torch.uint16, + torch.int32, + torch.float32, + ] + ), cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), cp.Value.Le(lambda deps, dtype, struct: 2**4), cp.Rank.Ge(lambda deps: 1), cp.Size.Ge(lambda deps, r, d: 1), cp.Size.Le(lambda deps, r, d: 2**9), ] + case "embedding.default": + tensor_constraints = [ + cp.Dtype.In(lambda deps: [torch.float, torch.int]), + cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]), + cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + cp.Rank.Ge(lambda deps: 1), + cp.Size.Ge(lambda deps, r, d: 1), + cp.Size.Le(lambda deps, r, d: 2**9), + ] case "sigmoid.default": tensor_constraints.extend( [ - cp.Dtype.In(lambda deps: [torch.float]), - cp.Rank.Le(lambda deps: 2**2), + cp.Dtype.In(lambda deps: [torch.float32]), cp.Value.Ge(lambda deps, dtype, struct: -2), cp.Value.Le(lambda deps, dtype, struct: 2), ] @@ -65,8 +117,7 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]: case "rsqrt.default": tensor_constraints.extend( [ - cp.Dtype.In(lambda deps: [torch.float]), - cp.Rank.Le(lambda deps: 2**2), + cp.Dtype.In(lambda deps: [torch.float32]), cp.Value.Gt( lambda deps, dtype, struct: 0 ), # only generate real numbers @@ -76,14 +127,12 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]: case "mean.dim": tensor_constraints.extend( [ - cp.Dtype.In(lambda deps: [torch.float]), - cp.Rank.Le(lambda deps: 2**2), + cp.Dtype.In(lambda deps: [torch.float32]), ] ) case "exp.default": tensor_constraints.extend( [ - cp.Rank.Le(lambda deps: 2**3), cp.Value.Ge(lambda deps, dtype, struct: -(2**2)), cp.Value.Le(lambda deps, dtype, struct: 2**2), ] @@ -96,26 +145,99 @@ def apply_tensor_contraints(op_name: str, index: int) -> list[object]: cp.Value.Le(lambda deps, dtype, struct: 2), ] ) - case _: + case "constant_pad_nd.default": + tensor_constraints.extend( + [ + cp.Dtype.In(lambda deps: [torch.float32]), + cp.Size.Le(lambda deps, r, d: 2**2), + ] + ) + case "avg_pool2d.default": + tensor_constraints.extend( + [ + cp.Rank.Eq(lambda deps: 4), + ] + ) + case "bmm.default" | "addmm.default" | "mm.default": tensor_constraints.extend( [ + cp.Dtype.Eq(lambda deps: torch.float), + cp.Size.Le(lambda deps, r, d: 2**2), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + ] + ) + case "div.Tensor": + tensor_constraints.extend( + [ + cp.Value.Ne(lambda deps, dtype, struct: 0), + cp.Value.Le(lambda deps, dtype, struct: 2**3), + cp.Size.Le(lambda deps, r, d: 2**3), cp.Rank.Le(lambda deps: 2**2), ] ) + case "div.Tensor_mode" | "minimum.default": + if index == 0: + tensor_constraints = [ + cp.Dtype.In(lambda deps: [torch.int64, torch.int32, torch.float32]), + cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + cp.Rank.Ge(lambda deps: 1), + cp.Size.Ge(lambda deps, r, d: 1), + cp.Size.Le(lambda deps, r, d: 2**2), + ] + else: + tensor_constraints = [ + cp.Dtype.In(lambda deps: [torch.int64, torch.int32, torch.float32]), + cp.Value.Ge(lambda deps, dtype, struct: -(2**4)), + cp.Value.Le(lambda deps, dtype, struct: 2**4), + cp.Rank.Ge(lambda deps: 1), + cp.Rank.Eq(lambda deps: deps[0].dim()), + cp.Size.Eq(lambda deps, r, d: fn.safe_size(deps[0], d)), + ] + case "_native_batch_norm_legit_no_training.default": + tensor_constraints.extend( + [ + cp.Rank.Le(lambda deps: 3), + ], + ) + case "reciprocal.default": + tensor_constraints = [ + cp.Value.Ge(lambda deps, dtype, struct: -(2**2)), + cp.Value.Le(lambda deps, dtype, struct: 2**2), + cp.Size.Le(lambda deps, r, d: 2**3), + ] + case "_softmax.default": + tensor_constraints.extend( + [ + cp.Dtype.Eq(lambda deps: torch.float32), + cp.Size.Le(lambda deps, r, d: 2**2), + ] + ) + case _: + pass return tensor_constraints def apply_scalar_contraints(op_name: str) -> list[ScalarDtype]: match op_name: - case "add.Scalar" | "sub.Scalar" | "mul.Scalar" | "div.Scalar": + case ( + "add.Scalar" + | "sub.Scalar" + | "mul.Scalar" + | "div.Scalar" + | "constant_pad_nd.default" + ): + return [ScalarDtype.int] + case "full.default": return [ScalarDtype.int] - case _: return [ScalarDtype.float, ScalarDtype.int] @lru_cache(maxsize=None) -def facto_testcase_gen(op_name: str) -> List[Tuple[List[str], OrderedDict[str, str]]]: +def facto_testcase_gen( # noqa: C901 + op_name: str, +) -> List[Tuple[List[str], OrderedDict[str, str]]]: # minimal example to test add.Tensor using FACTO spec = SpecDictDB[op_name] @@ -149,6 +271,12 @@ def facto_testcase_gen(op_name: str) -> List[Tuple[List[str], OrderedDict[str, s cp.Dtype.In(lambda deps: apply_scalar_contraints(op_name)), ] ) + if in_spec.name == "dtype": # full.default + spec.inspec[index].constraints.extend( + [ + cp.Dtype.In(lambda deps: [torch.long, torch.float]), + ] + ) elif in_spec.type.is_tensor(): spec.inspec[index].constraints.extend( apply_tensor_contraints(op_name, index) @@ -166,6 +294,29 @@ def facto_testcase_gen(op_name: str) -> List[Tuple[List[str], OrderedDict[str, s cp.Dtype.In(lambda deps: [torch.bool]), ] ) + elif in_spec.type.is_length_list(): + spec.inspec[index].constraints.extend( + [ + cp.Value.Ge(lambda deps, dtype, struct: 0), + ] + ) + if op_name == "avg_pool2d.default": + spec.inspec[index].constraints.extend( + [ + cp.Length.Eq(lambda deps: 2), + ] + ) + elif in_spec.type.is_shape(): + spec.inspec[index].constraints.extend( + [ + cp.Rank.Ge(lambda deps: 1), + cp.Rank.Le(lambda deps: 2**2), + cp.Value.Gt(lambda deps, dtype, struct: 0), + cp.Value.Le(lambda deps, dtype, struct: 2**2), + cp.Size.Ge(lambda deps, r, d: 1), + cp.Size.Le(lambda deps, r, d: 2**2), + ] + ) return [ (posargs, inkwargs) diff --git a/backends/cadence/utils/targets.bzl b/backends/cadence/utils/targets.bzl index 5a8d18ec8a2..8b14f7684da 100644 --- a/backends/cadence/utils/targets.bzl +++ b/backends/cadence/utils/targets.bzl @@ -3,11 +3,11 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") def define_common_targets(): - python_library( + runtime.python_library( name = "facto_util", srcs = [ "facto_util.py", diff --git a/backends/cortex_m/CMakeLists.txt b/backends/cortex_m/CMakeLists.txt index 5c353389d94..1567b8b5e1c 100644 --- a/backends/cortex_m/CMakeLists.txt +++ b/backends/cortex_m/CMakeLists.txt @@ -1,13 +1,10 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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. -# Kernel library for Cortex-M operators. Please keep this file formatted by running: -# ~~~ -# cmake-format -i CMakeLists.txt -# ~~~ cmake_minimum_required(VERSION 3.19) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) @@ -22,29 +19,65 @@ endif() include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) +include(FetchContent) + +# CMSIS-NN version to download +set(CMSIS_NN_VERSION + "v4.1.0" + CACHE STRING "CMSIS-NN version to download" +) + +# Declare CMSIS-NN as a FetchContent project +FetchContent_Declare( + cmsis_nn + GIT_REPOSITORY https://github.com/ARM-software/CMSIS-NN.git + GIT_TAG ${CMSIS_NN_VERSION} +) + +# Download and make CMSIS-NN available +FetchContent_MakeAvailable(cmsis_nn) + +# Print paths for debugging +message(STATUS "CMSIS-NN source dir: ${cmsis_nn_SOURCE_DIR}") +message(STATUS "CMSIS-NN binary dir: ${cmsis_nn_BINARY_DIR}") # Cortex-M ops kernel sources set(_cortex_m_kernels__srcs ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantize_per_tensor.cpp ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_add.cpp ) -# Generate C++ bindings to register kernels into Executorch (for runtime). -# Here select all ops in operators.yaml +# Generate C++ bindings to register kernels into Executorch (for runtime) set(_yaml_file ${CMAKE_CURRENT_LIST_DIR}/ops/operators.yaml) gen_selected_ops(LIB_NAME "cortex_m_ops_lib" OPS_SCHEMA_YAML "${_yaml_file}") -# Generate bindings for the kernels generate_bindings_for_kernels( LIB_NAME "cortex_m_ops_lib" CUSTOM_OPS_YAML "${_yaml_file}" ) message("Generated files ${gen_command_sources}") -# Build a library for _cortex_m_kernels_srcs +# Build a library for cortex_m_kernels add_library(cortex_m_kernels ${_cortex_m_kernels__srcs}) -target_link_libraries(cortex_m_kernels PRIVATE executorch) target_compile_options(cortex_m_kernels PUBLIC ${_common_compile_options}) +# Include directories for cortex_m_kernels +target_include_directories( + cortex_m_kernels + PRIVATE ${EXECUTORCH_ROOT}/.. + ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 + ${cmsis_nn_SOURCE_DIR}/Include +) + +# Link directly to the CMSIS-NN static library file +target_link_libraries( + cortex_m_kernels PUBLIC ${cmsis_nn_BINARY_DIR}/libcmsis-nn.a executorch +) + +# Add dependency to ensure CMSIS-NN builds before we try to link. Use the actual +# CMSIS-NN target name (usually 'cmsis-nn') +add_dependencies(cortex_m_kernels cmsis-nn) + # cortex_m_ops_lib: Register Cortex-M ops kernels into Executorch runtime gen_operators_lib( LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS executorch @@ -52,6 +85,7 @@ gen_operators_lib( install( TARGETS cortex_m_kernels cortex_m_ops_lib + EXPORT ExecuTorchTargets DESTINATION lib PUBLIC_HEADER DESTINATION include/executorch/backends/cortex_m/ops/ ) diff --git a/backends/cortex_m/ops/TARGETS b/backends/cortex_m/ops/TARGETS index e02f096fd83..12044266ca1 100644 --- a/backends/cortex_m/ops/TARGETS +++ b/backends/cortex_m/ops/TARGETS @@ -4,18 +4,19 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("targets.bzl", "define_common_targets") oncall("executorch") -python_library( +runtime.python_library( name = "ops", srcs = [ "operators.py", ], deps = [ "fbcode//caffe2:torch", + "//executorch/backends/cortex_m/passes:passes_utils", ], ) diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h new file mode 100644 index 00000000000..5ef2d9d4bf9 --- /dev/null +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -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. + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +using Tensor = torch::executor::Tensor; +using ScalarType = executorch::aten::ScalarType; +using Scalar = torch::executor::Scalar; +using Error = executorch::runtime::Error; + +// Basic tensor type / layout validation and dimension order checking +inline void validate_cmsis_nn_tensor_requirements( + const Tensor& input1, + const Tensor& input2, + Tensor& output, + ScalarType expected_dtype = ScalarType::Char, + bool require_channels_last = false) { + // Basic dtype validation + ET_CHECK_MSG( + input1.scalar_type() == expected_dtype, + "Input1 dtype must be %hhd", + expected_dtype); + ET_CHECK_MSG( + input2.scalar_type() == expected_dtype, + "Input2 dtype must be %hhd", + expected_dtype); + ET_CHECK_MSG( + output.scalar_type() == expected_dtype, + "Output dtype must be %hhd", + expected_dtype); + + // Dim order consistency + ET_CHECK_MSG( + executorch::runtime::tensors_have_same_dim_order(input1, input2, output), + "Tensors must have same dimension order"); + + // TBD: Validate memory alignment (CMSIS-NN requirement) +} + +inline void validate_single_quant_params( + const Scalar& zero_point, + const Scalar& multiplier, + const Scalar& shift, + const char* param_name) { + int64_t zp_val = zero_point.to(); + int64_t mult_val = multiplier.to(); + int64_t shift_val = shift.to(); + + ET_CHECK_MSG( + zp_val >= std::numeric_limits::min() && + zp_val <= std::numeric_limits::max(), + "%s zero point must be in int8 range [Value: %d]", + param_name, + zp_val); + + ET_CHECK_MSG( + mult_val >= std::numeric_limits::min() && + mult_val <= std::numeric_limits::max(), + "%s multiplier must be in int32 range [Value: %d]", + param_name, + mult_val); + + ET_CHECK_MSG( + shift_val >= -31 && shift_val <= 31, + "%s shift must be in range [-31, 31] [Value: %d]", + param_name, + shift_val); +} + +/** + * Validate quantization parameters for inputs and output. + * + * Checks that zero points fit in int8 range, multipliers fit in int32 range, + * and shifts are within a valid bit-shift range (0-31). + * + * Ensures parameters comply with Ahead-Of-Time (AOT) quantization requirements + * and CMSIS-NN kernel expectations. + * + * Raises errors via ET_KERNEL_CHECK if any check fails. + */ +inline void validate_quantization_params( + const Scalar& zero_point1, + const Scalar& multiplier1, + const Scalar& shift1, + const Scalar& zero_point2, + const Scalar& multiplier2, + const Scalar& shift2, + const Scalar& output_zero_point, + const Scalar& output_multiplier, + const Scalar& output_shift, + Tensor& output) { + validate_single_quant_params( + zero_point1, multiplier1, shift1, "Single quant Input1"); + validate_single_quant_params( + zero_point2, multiplier2, shift2, "Single quant Input2"); + validate_single_quant_params( + output_zero_point, + output_multiplier, + output_shift, + "Single quant Output"); +} + +inline Error resize_to_broadcast_target_size( + const Tensor& input1, + const Tensor& input2, + Tensor& output) { + static constexpr int kTensorDimensionLimit = 5; + Tensor::SizesType expected_output_size[kTensorDimensionLimit]; + size_t expected_output_dim = 0; + auto err = torch::executor::get_broadcast_target_size( + input1, + input2, + expected_output_size, + kTensorDimensionLimit, + &expected_output_dim); + + if (err != Error::Ok) + return err; + + return executorch::runtime::resize_tensor( + output, {expected_output_size, expected_output_dim}); +} + +/** + * Convert Scalar to CMSIS-NN int32 format + * For multipliers, zero_points, etc. from quantize_multiplier_aot + */ +inline int32_t extractScalarToInt32(const Scalar& scalar_value) { + return static_cast(scalar_value.to()); +} + +/** + * Convert Scalar to CMSIS-NN int format + * For shift values from quantize_multiplier_aot + */ +inline int extractScalarToInt(const Scalar& scalar_value) { + return static_cast(scalar_value.to()); +} diff --git a/backends/cortex_m/ops/op_quantized_add.cpp b/backends/cortex_m/ops/op_quantized_add.cpp new file mode 100644 index 00000000000..044c2bd92d5 --- /dev/null +++ b/backends/cortex_m/ops/op_quantized_add.cpp @@ -0,0 +1,150 @@ +/* + * 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 "cortex_m_ops_common.h" + +// Include CMSIS-NN headers with C linkage +extern "C" { +#include "arm_nnfunctions.h" +} + +namespace cortex_m { +namespace native { +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; + +Tensor& quantized_add_out( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const Scalar& input1_zero_point, + const Scalar& input1_multiplier, + const Scalar& input1_shift, + const Tensor& input2_int8, + const Scalar& input2_zero_point, + const Scalar& input2_multiplier, + const Scalar& input2_shift, + const Scalar& output_zero_point, + const Scalar& output_multiplier, + const Scalar& output_shift, + Tensor& out) { + // Validate tensor types and dim order + validate_cmsis_nn_tensor_requirements(input1_int8, input2_int8, out); + + // Validate quantization parameters + validate_quantization_params( + input1_zero_point, + input1_multiplier, + input1_shift, + input2_zero_point, + input2_multiplier, + input2_shift, + output_zero_point, + output_multiplier, + output_shift, + out); + + // Broadcast if needed + auto result = resize_to_broadcast_target_size(input1_int8, input2_int8, out); + ET_CHECK_MSG( + (result == Error::Ok), + "Failed to resize output tensor. Status: [%d]", + result); + + ET_LOG( + Info, + "quantized_add_out: input1_int8.sizes() = %zu", + input1_int8.sizes().size()); + + int32_t zp1 = extractScalarToInt32(input1_zero_point); + int32_t input1_mult = extractScalarToInt32(input1_multiplier); + int input1_shift_val = extractScalarToInt(input1_shift); + int32_t zp2 = extractScalarToInt32(input2_zero_point); + int32_t input2_mult = extractScalarToInt32(input2_multiplier); + int input2_shift_val = extractScalarToInt(input2_shift); + int32_t out_zp = extractScalarToInt32(output_zero_point); + int32_t output_mult = extractScalarToInt32(output_multiplier); + int output_shift_val = extractScalarToInt(output_shift); + + // Left shift to maximize precision (tune as needed) + const int32_t left_shift = 20; + const int32_t activation_min = std::numeric_limits::min(); + const int32_t activation_max = std::numeric_limits::max(); + + ET_LOG( + Info, + "Using AoT-computed parameters: input1[mult=%d, shift=%d], input2[mult=%d, shift=%d], output[mult=%d, shift=%d]", + input1_mult, + input1_shift_val, + input2_mult, + input2_shift_val, + output_mult, + output_shift_val); + + // Call CMSIS-NN kernel with precomputed parameters + arm_cmsis_nn_status status = arm_elementwise_add_s8( + input1_int8.const_data_ptr(), + input2_int8.const_data_ptr(), + static_cast(zp1), + input1_mult, + input1_shift_val, + static_cast(zp2), + input2_mult, + input2_shift_val, + left_shift, + out.mutable_data_ptr(), + static_cast(out_zp), + output_mult, + output_shift_val, + static_cast(out.numel()), + activation_min, + activation_max); + + if (status != ARM_CMSIS_NN_SUCCESS) { + ET_LOG( + Error, + "quantized_add_out: arm_elementwise_add_s8 failed with status [%d]", + status); + + context.fail(Error::Internal); // Fail the execution context + return out; + } + ET_LOG( + Info, + "quantized_add_out: Successfully completed with AoT-computed parameters!"); + + return out; +} + +// Stub Implementation: Non-out variant for compatibility (functional variant) +// EXIR/ExecuTorch runs an out-variant pass that converts +// .default operations to .out variants before memory planning. +// In the pass we are calling quantized_add's default variant +// but ExecuTorch's kernel dispatch mechanism will end up calling the out +// variant. This stub is to make sure that compiler doesn't complain. +Tensor quantized_add( + KernelRuntimeContext& context, + const Tensor& input1_int8, + const Scalar& input1_zero_point, + const Scalar& input1_multiplier, + const Scalar& input1_shift, + const Tensor& input2_int8, + const Scalar& input2_zero_point, + const Scalar& input2_multiplier, + const Scalar& input2_shift, + const Scalar& output_zero_point, + const Scalar& output_multiplier, + const Scalar& output_shift) { + ET_LOG(Info, "quantized_add: input1_int8.sizes() = %zu", input1_int8.sizes()); + + // Crash on Debug builds if invoked + assert(False); + // This is to make sure compiler doesn't complain. + return const_cast(input1_int8); +} + +} // namespace native +} // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index b5ba3c6ddc8..926dcd85e4b 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -5,11 +5,16 @@ # LICENSE file in the root directory of this source tree. import torch -from executorch.exir.dialects._ops import ( - ops as exir_ops, -) # To provide the implementation of the operators +from executorch.backends.cortex_m.passes.passes_utils import ( + dequantize_per_tensor_cmsis, + quantize_per_tensor_cmsis, +) +from executorch.exir.dialects._ops import ops as exir_ops + +# To provide the implementation of the operators from torch.library import impl, Library, register_fake + # New operator library with a custom namespace to allow fusion etc. lib = Library("cortex_m", "DEF") @@ -96,3 +101,125 @@ def dequantize_per_tensor_impl( return exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default( input, scale, zero_point, quant_min, quant_max, dtype ) + + +# Define the operator schema with multipliers and shifts (11 args) +lib.define( + "quantized_add(" + "Tensor self, Scalar self_zero_point, Scalar self_multiplier, Scalar self_shift, " + "Tensor other, Scalar other_zero_point, Scalar other_multiplier, Scalar other_shift, " + "Scalar output_zero_point, Scalar output_multiplier, Scalar output_shift) -> Tensor" +) + + +@register_fake("cortex_m::quantized_add") +def quantized_add_meta( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + broadcasted_shape = torch.broadcast_shapes(self.shape, other.shape) + return torch.empty(broadcasted_shape, dtype=torch.int8, device=self.device) + + +@impl(lib, "quantized_add", "CompositeExplicitAutograd") +def quantized_add_impl( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, +) -> torch.Tensor: + self_fp = dequantize_per_tensor_cmsis( + self, self_zero_point, self_multiplier, self_shift + ) + other_fp = dequantize_per_tensor_cmsis( + other, other_zero_point, other_multiplier, other_shift + ) + result_fp = self_fp + other_fp + result_quantized = quantize_per_tensor_cmsis( + result_fp, output_zero_point, output_multiplier, output_shift + ) + return result_quantized + + +# Define the operator schema with multipliers and shifts (11 args + out tensor) +lib.define( + "quantized_add.out(" + "Tensor self, Scalar self_zero_point, Scalar self_multiplier, Scalar self_shift, " + "Tensor other, Scalar other_zero_point, Scalar other_multiplier, Scalar other_shift, " + "Scalar output_zero_point, Scalar output_multiplier, Scalar output_shift, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +# Fake meta function for shape and dtype inference during compilation +@register_fake("cortex_m::quantized_add.out") +def quantized_add_out_meta( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + out: torch.Tensor, +) -> torch.Tensor: + # Validate against correct broadcasted shape + expected_shape = torch.broadcast_shapes(self.shape, other.shape) + assert ( + out.shape == expected_shape + ), f"Output shape {out.shape} must match broadcasted shape {expected_shape}" + return out + + +# Actual implementation delegating to backend or custom kernel +@impl(lib, "quantized_add.out", "CompositeExplicitAutograd") +def quantized_add_out_impl( + self: torch.Tensor, + self_zero_point: int, + self_multiplier: int, + self_shift: int, + other: torch.Tensor, + other_zero_point: int, + other_multiplier: int, + other_shift: int, + output_zero_point: int, + output_multiplier: int, + output_shift: int, + *, + out: torch.Tensor, +) -> torch.Tensor: + self_fp = dequantize_per_tensor_cmsis( + self, self_zero_point, self_multiplier, self_shift + ) + other_fp = dequantize_per_tensor_cmsis( + other, other_zero_point, other_multiplier, other_shift + ) + result_fp = self_fp + other_fp + result_quantized = quantize_per_tensor_cmsis( + result_fp, output_zero_point, output_multiplier, output_shift + ) + + # Write into the provided output tensor + out.copy_(result_quantized) + + return out diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 0cc248effaa..f2615a1f525 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -15,3 +15,15 @@ kernels: - arg_meta: null kernel_name: cortex_m::dequantize_per_tensor_out + +- func: cortex_m::quantized_add(Tensor self, Scalar self_zero_point, Scalar self_multiplier, Scalar self_shift, Tensor other, Scalar other_zero_point, Scalar other_multiplier, Scalar other_shift, Scalar output_zero_point, Scalar output_multiplier, Scalar output_shift) -> Tensor + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_add + +- func: cortex_m::quantized_add.out(Tensor self, Scalar self_zero_point, Scalar self_multiplier, Scalar self_shift, Tensor other, Scalar other_zero_point, Scalar other_multiplier, Scalar other_shift, Scalar output_zero_point, Scalar output_multiplier, Scalar output_shift, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_add_out diff --git a/backends/cortex_m/passes/TARGETS b/backends/cortex_m/passes/TARGETS index e5d1fc6149d..687aafe57d4 100644 --- a/backends/cortex_m/passes/TARGETS +++ b/backends/cortex_m/passes/TARGETS @@ -4,18 +4,32 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( - name = "replace_quant_nodes_pass", - srcs = ["replace_quant_nodes_pass.py"], - deps = [ +runtime.python_library( + name="replace_quant_nodes_pass", + srcs=[ + "replace_quant_nodes_pass.py", + "quantized_op_fusion_pass.py", + ], + deps=[ "//caffe2:torch", "//executorch/exir:lib", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", "//executorch/backends/cortex_m/ops:ops", - ] + "//executorch/backends/cortex_m/passes:passes_utils", + ], +) + +runtime.python_library( + name="passes_utils", + srcs=[ + "passes_utils.py", + ], + deps=[ + "fbcode//caffe2:torch", + ], ) diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py new file mode 100644 index 00000000000..3f6e05fc4de --- /dev/null +++ b/backends/cortex_m/passes/passes_utils.py @@ -0,0 +1,94 @@ +# 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 + + +def dequantize_per_tensor_cmsis( + qtensor: torch.Tensor, zero_point: int, multiplier: int, shift: int +) -> torch.Tensor: + """ + Simulate CMSIS-NN fixed-point dequantization: + result = (qtensor - zero_point) * multiplier * 2^shift / 2^31 + """ + scale = multiplier * (2**shift) / (1 << 31) + return (qtensor.float() - zero_point) * scale + + +def quantize_per_tensor_cmsis( + tensor: torch.Tensor, + zero_point: int, + multiplier: int, + shift: int, + qmin=-128, + qmax=127, +) -> torch.Tensor: + """ + Simulate CMSIS-NN fixed-point quantization: + result = round(tensor / scale) + zero_point, clamped to [qmin, qmax] + """ + scale = multiplier * (2**shift) / (1 << 31) + quantized = torch.round(tensor / scale) + zero_point + return quantized.clamp(qmin, qmax).to(torch.int8) + + +def extract_scalar_value(node_arg) -> float: + """ + Extract scalar value from various PyTorch scalar representations. + """ + if hasattr(node_arg, "op") and node_arg.op == "get_attr": + # Handle case where scalar is a graph attribute + return float(node_arg.target) + elif isinstance(node_arg, (int, float)): + return float(node_arg) + elif hasattr(node_arg, "item"): + return float(node_arg.item()) + else: + # Try to extract from meta if available + if hasattr(node_arg, "meta") and "val" in node_arg.meta: + val = node_arg.meta["val"] + if hasattr(val, "item"): + return float(val.item()) + return float(val) + raise ValueError( + f"Cannot extract scalar value from {type(node_arg)}: {node_arg}" + ) + + +def is_qualified_int8_node(args) -> bool: + try: + if len(args) < 6: + return False + qmin = int(args[3]) + qmax = int(args[4]) + dtype_str = str(args[5]) + is_int8_range = ( + qmin >= torch.iinfo(torch.int8).min and qmax <= torch.iinfo(torch.int8).max + ) + is_int8_dtype = "int8" in dtype_str.lower() + return is_int8_range and is_int8_dtype + except (IndexError, ValueError, TypeError): + return False + + +def quantize_multiplier_aot(scale: float) -> tuple[int, int]: + if scale == 0.0: + return 0, 0 + mantissa, exponent = math.frexp(scale) + shift = -exponent + q_fixed = int(round(mantissa * (1 << 31))) + if q_fixed == (1 << 31): + q_fixed //= 2 + shift -= 1 + multiplier = max(-2147483648, min(2147483647, q_fixed)) + return multiplier, shift + + +def cleanup_erased_nodes(graph_module: torch.fx.GraphModule): + # Placeholder for any additional cleanup if needed + pass diff --git a/backends/cortex_m/passes/quantized_op_fusion_pass.py b/backends/cortex_m/passes/quantized_op_fusion_pass.py new file mode 100644 index 00000000000..ca6d8b97795 --- /dev/null +++ b/backends/cortex_m/passes/quantized_op_fusion_pass.py @@ -0,0 +1,255 @@ +# 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 logging +from typing import Set + +import executorch.backends.cortex_m.ops.operators # noqa +import torch + +from executorch.backends.cortex_m.passes.passes_utils import ( + extract_scalar_value, + quantize_multiplier_aot, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from torch.fx.passes.infra.pass_manager import PassResult + +logger = logging.getLogger("quant_op_fusion_pass") +logger.setLevel(logging.INFO) + + +class QuantizedOpFusionPass(ExportPass): + """ + Generic ExportPass that: + 1. Replaces certain ops with cortex_m variants based on qualifiers. + 2. Fuses patterns: dequantize_per_tensor -> [binary_op] -> quantize_per_tensor + into cortex_m.quantized_[op].default with AoT computed multipliers/shifts. + + + Supports multiple binary operations with backward compatibility for add. + """ + + # Generic operation mapping + SUPPORTED_OPS_MAPPING = { + exir_ops.edge.aten.add.Tensor: exir_ops.edge.cortex_m.quantized_add.default, + # Future ops to be added here: + } + + def __init__(self): + super().__init__() + + def _get_dequant_targets(self) -> Set: + """Support both decomposed and cortex_m dequant targets for flexible pass ordering.""" + return { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.cortex_m.dequantize_per_tensor.default, + } + + def _get_quant_targets(self) -> Set: + """Support both decomposed and cortex_m quant targets for flexible pass ordering.""" + return { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.cortex_m.quantize_per_tensor.default, + } + + def _is_supported_binary_op(self, node: torch.fx.Node) -> bool: + """Check if node is a supported binary operation.""" + return node.op == "call_function" and node.target in self.SUPPORTED_OPS_MAPPING + + def _is_dequant_node(self, node: torch.fx.Node) -> bool: + """Check if node is a dequantize operation.""" + return ( + hasattr(node, "op") + and node.op == "call_function" + and node.target in self._get_dequant_targets() + ) + + def _is_quant_node(self, node: torch.fx.Node) -> bool: + """Check if node is a quantize operation.""" + return ( + hasattr(node, "op") + and node.op == "call_function" + and node.target in self._get_quant_targets() + ) + + def _transfer_metadata( + self, + new_node: torch.fx.Node, + source_node: torch.fx.Node, + pass_name: str = "QuantizedOpFusionPass", + ) -> None: + """Metadata transfer with proper provenance tracking.""" + if hasattr(source_node, "meta") and source_node.meta: + new_node.meta = source_node.meta.copy() + + if "from_node" in new_node.meta: + from_node_list = new_node.meta.get("from_node", []).copy() + from_node_list.append( + {"source": source_node.name, "pass": pass_name, "op": "fuse"} + ) + new_node.meta["from_node"] = from_node_list + + # Copy essential fields + for field in ["tensor_meta", "stack_trace"]: + if field in source_node.meta: + new_node.meta[field] = source_node.meta[field] + + def _normalize_to_cortex_m_targets(self, graph_module: torch.fx.GraphModule) -> int: + """Convert decomposed targets to cortex_m equivalents for consistent handling.""" + target_mapping = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: exir_ops.edge.cortex_m.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default: exir_ops.edge.cortex_m.quantize_per_tensor.default, + } + + normalization_count = 0 + for node in list(graph_module.graph.nodes): + if node.op == "call_function" and node.target in target_mapping: + logger.info(f"Normalizing {node.target} to cortex_m equivalent") + node.target = target_mapping[node.target] + normalization_count += 1 + + return normalization_count + + def _fuse_quantized_binary_patterns( + self, graph_module: torch.fx.GraphModule + ) -> int: + """Generic fusion for quantized binary operation patterns.""" + fusion_count = 0 + nodes_to_erase = [] + + for node in list(graph_module.graph.nodes): + if not self._is_quant_node(node): + continue + + quantize_node = node + if not quantize_node.args: + continue + + binary_op_node = quantize_node.args[0] + if not self._is_supported_binary_op(binary_op_node): + continue + + if len(binary_op_node.args) < 2: + continue + + dequant_node1, dequant_node2 = binary_op_node.args[:2] + if not ( + self._is_dequant_node(dequant_node1) + and self._is_dequant_node(dequant_node2) + ): + continue + + # Get the target quantized operation + quantized_target = self.SUPPORTED_OPS_MAPPING[binary_op_node.target] + # Extract op name (e.g., 'Tensor' -> 'add') + op_name = str(binary_op_node.target).split(".")[-1] + logger.info(f"✅ Found complete cortex_m Q/DQ + {op_name} pattern!") + + try: + # Extract values + int8_tensor1, scale1, zero_point1 = dequant_node1.args[:3] + int8_tensor2, scale2, zero_point2 = dequant_node2.args[:3] + output_scale, output_zero_point = quantize_node.args[1:3] + + # Convert to Python floats + scale1_val = extract_scalar_value(scale1) + scale2_val = extract_scalar_value(scale2) + output_scale_val = extract_scalar_value(output_scale) + zp1_val = int(extract_scalar_value(zero_point1)) + zp2_val = int(extract_scalar_value(zero_point2)) + output_zp_val = int(extract_scalar_value(output_zero_point)) + + # AoT COMPUTATION: Calculate multipliers and shifts + input1_mult, input1_shift = quantize_multiplier_aot( + scale1_val / output_scale_val + ) + input2_mult, input2_shift = quantize_multiplier_aot( + scale2_val / output_scale_val + ) + output_mult, output_shift = quantize_multiplier_aot( + 1.0 + ) # Output multiplier is 1 + + logger.info("AoT computed parameters:") + logger.info(f" Input1: mult={input1_mult}, shift={input1_shift}") + logger.info(f" Input2: mult={input2_mult}, shift={input2_shift}") + logger.info(f" Output: mult={output_mult}, shift={output_shift}") + + with graph_module.graph.inserting_after(quantize_node): + fused = graph_module.graph.create_node( + "call_function", + target=quantized_target, + args=( + int8_tensor1, + zp1_val, + input1_mult, + input1_shift, + int8_tensor2, + zp2_val, + input2_mult, + input2_shift, + output_zp_val, + output_mult, + output_shift, + ), + kwargs={}, + ) + + # metadata transfer + self._transfer_metadata(fused, quantize_node) + + logger.info(f"✅ Created fused quantized_{op_name} node: {fused}") + + # Replace all uses + quantize_node.replace_all_uses_with(fused) + binary_op_node.replace_all_uses_with(fused) + dequant_node1.replace_all_uses_with(fused) + dequant_node2.replace_all_uses_with(fused) + + nodes_to_erase.extend( + [quantize_node, binary_op_node, dequant_node1, dequant_node2] + ) + fusion_count += 1 + logger.info(f"Pattern fused, total so far: {fusion_count}") + + except Exception as e: + logger.info(f"❌ Error during AoT computation: {e}") + logger.info(" Skipping fusion for this pattern") + continue + + for old_node in reversed(nodes_to_erase): + if old_node in graph_module.graph.nodes and len(old_node.users) == 0: + logger.info(f"🗑️ Erasing node: {old_node}") + graph_module.graph.erase_node(old_node) + + return fusion_count + + def call(self, graph_module: torch.fx.GraphModule): + logger.info("QuantizedOpFusionPass.call() started") + + # Normalize targets for flexible pass ordering + normalization_count = self._normalize_to_cortex_m_targets(graph_module) + + # Generic fusion for supported binary operations + fusion_count = self._fuse_quantized_binary_patterns(graph_module) + + total_changes = normalization_count + fusion_count + logger.info(f"Total changes: {total_changes}") + + if total_changes > 0: + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + + logger.debug("=== AFTER FUSION: All nodes in the graph ===") + for i, node in enumerate(graph_module.graph.nodes): + logger.debug(f"Node {i}: op={node.op}, target={node.target}") + if "quantized_" in str(node.target) and "add" in str(node.target): + logger.debug(" ⭐ FOUND QUANTIZED BINARY OP NODE! ⭐") + logger.debug("=== END DEBUG ===") + + return PassResult(graph_module, total_changes > 0) diff --git a/backends/cortex_m/passes/replace_quant_nodes_pass.py b/backends/cortex_m/passes/replace_quant_nodes_pass.py index d9d7506a146..a8153136db9 100644 --- a/backends/cortex_m/passes/replace_quant_nodes_pass.py +++ b/backends/cortex_m/passes/replace_quant_nodes_pass.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -10,7 +11,6 @@ import torch from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue @@ -40,6 +40,10 @@ def __init__(self): "qualifier": self._is_qualified_int8_node, }, } + self.disallowed_targets = { + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.default, + } def call_operator( self, @@ -48,9 +52,10 @@ def call_operator( kwargs: Dict[str, object], meta: NodeMetadata, ) -> ProxyValue: - assert isinstance( - op, EdgeOpOverload - ), "Op must be an EdgeOpOverload. Run this pass after to_edge()." + if op in self.disallowed_targets: + raise RuntimeError( + f"Found unexpected aten op '{op}'. Make sure you run this pass after lowering to edge." + ) if op in self.op_replacements and self.op_replacements[op]["qualifier"](args): return super().call_operator( diff --git a/backends/cortex_m/test/TARGETS b/backends/cortex_m/test/TARGETS index d381011b648..b7a04f3efab 100644 --- a/backends/cortex_m/test/TARGETS +++ b/backends/cortex_m/test/TARGETS @@ -10,14 +10,18 @@ load("targets.bzl", "define_common_targets") oncall("executorch") python_unittest( - name = "test_replace_quant_nodes", - srcs = ["test_replace_quant_nodes.py"], - deps = [ + name="test_replace_quant_nodes", + srcs=[ + "test_helpers_passes_utils.py", + "test_replace_quant_nodes.py", + "test_quantize_op_fusion_pass.py", + ], + deps=[ "//pytorch/ao:torchao", # @manual - "//caffe2:torch", + "//caffe2:torch", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/backends/cortex_m/ops:ops", ], -) +) define_common_targets() diff --git a/backends/cortex_m/test/test_helpers_passes_utils.py b/backends/cortex_m/test/test_helpers_passes_utils.py new file mode 100644 index 00000000000..ccaf8476867 --- /dev/null +++ b/backends/cortex_m/test/test_helpers_passes_utils.py @@ -0,0 +1,103 @@ +# 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 dataclasses import dataclass +from typing import Optional + +import torch +from torch.fx import GraphModule +from torchao.quantization.pt2e.observer import HistogramObserver +from torchao.quantization.pt2e.quantizer import ( + QuantizationAnnotation, + QuantizationSpec, + Quantizer, +) +from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY + + +@dataclass(eq=True, frozen=True) +class QuantizationConfig: + input_activation: Optional[QuantizationSpec] + output_activation: Optional[QuantizationSpec] + + +class AddQuantizer(Quantizer): + def __init__(self): + super().__init__() + + @staticmethod + def _get_qspec(): + return QuantizationSpec( + dtype=torch.int8, + quant_min=-128, + quant_max=127, + qscheme=torch.per_tensor_symmetric, + is_dynamic=False, + observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12), + ) + + @staticmethod + def _get_qconfig(): + qspec = AddQuantizer._get_qspec() + return QuantizationConfig( + input_activation=qspec, + output_activation=qspec, + ) + + def annotate(self, model: GraphModule): + config = self._get_qconfig() + annotated_partitions = [] + + for node in model.graph.nodes: + if node.op != "call_function" or node.target not in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.add_.Tensor, + ]: + continue + + if Q_ANNOTATION_KEY in node.meta and node.meta[Q_ANNOTATION_KEY]._annotated: + continue + + input_qspec_map = { + node.args[0]: config.input_activation, + node.args[1]: config.input_activation, + } + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=config.output_activation, + _annotated=True, + ) + annotated_partitions.append([node]) + + return annotated_partitions + + def validate(self, model: GraphModule) -> None: + pass + + +def check_count( + graph_module: GraphModule, op: torch.fx.node.Target, expected_count: int +): + actual_count = sum( + 1 + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target == op + ) + + assert ( + actual_count == expected_count + ), f"Expected {expected_count} {op} nodes, got {actual_count}" + + +def get_node_args(graph_module: GraphModule, op: torch.fx.node.Target): + """Helper to get arguments of specific operator nodes""" + nodes = [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target == op + ] + return [node.args for node in nodes] diff --git a/backends/cortex_m/test/test_quantize_op_fusion_pass.py b/backends/cortex_m/test/test_quantize_op_fusion_pass.py new file mode 100644 index 00000000000..1595b0cfbc3 --- /dev/null +++ b/backends/cortex_m/test/test_quantize_op_fusion_pass.py @@ -0,0 +1,369 @@ +# 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 executorch +import executorch.backends.cortex_m.ops.operators # noqa + +import torch + +from executorch.backends.cortex_m.passes.quantized_op_fusion_pass import ( + QuantizedOpFusionPass, +) +from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import ( + ReplaceQuantNodesPass, +) +from executorch.backends.cortex_m.test.test_helpers_passes_utils import ( + AddQuantizer, + check_count, + get_node_args, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +class TestQuantizedOpFusionPass(unittest.TestCase): + """ + Test suite for the QuantizedOpFusionPass which fuses dequantize->add->quantize patterns + into a single quantized_add operation with AoT-computed parameters. + """ + + def setUp(self): + """Set up common test fixtures""" + self.example_inputs = (torch.randn(4, 8), torch.randn(4, 8)) + + def _prepare_quantized_model(self, model_class): + """Helper to prepare a quantized model for testing""" + model = model_class() + + # Export and quantize + exported_model = export(model.eval(), self.example_inputs, strict=True).module() + prepared_model = prepare_pt2e(exported_model, AddQuantizer()) + quantized_model = convert_pt2e(prepared_model) + + # Export to EXIR Edge + exported = export(quantized_model, self.example_inputs, strict=True) + edge_program = executorch.exir.to_edge( + exported, + compile_config=executorch.exir.EdgeCompileConfig(_check_ir_validity=False), + ) + return edge_program + + def _apply_passes(self, edge_program): + """Apply both ReplaceQuantNodesPass and QuantizedOpFusionPass""" + passes = [QuantizedOpFusionPass(), ReplaceQuantNodesPass()] + final_program = edge_program.transform(passes) + return final_program + + def test_single_add_fusion(self): + """Single add with full Q/DQ pattern should fuse into one quantized_add node""" + + class SingleAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # Prepare model + edge_program = self._prepare_quantized_model(SingleAddModel) + edge_graph = edge_program.exported_program().graph_module + + # Get reference output + reference_output = edge_graph(*self.example_inputs) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Verify fusion occurred + check_count( + transformed_graph, + exir_ops.edge.cortex_m.quantized_add.default, + 1, # Should have exactly 1 fused quantized_add + ) + + # Verify the following + # Before fusion: + # x --> quantize_per_tensor --> dequantize_per_tensor --> add --> quantize_per_tensor --> + # dequantize_per_tensor --> output y --> quantize_per_tensor --> dequantize_per_tensor --^ + # After fusion: + # x --> quantize_per_tensor --> quantized_add --> dequantize_per_tensor --> output + # y --> quantize_per_tensor --^ + check_count( + transformed_graph, exir_ops.edge.cortex_m.quantize_per_tensor.default, 2 + ) + check_count( + transformed_graph, exir_ops.edge.cortex_m.dequantize_per_tensor.default, 1 + ) + check_count(transformed_graph, exir_ops.edge.cortex_m.quantized_add.default, 1) + + # Verify numerical equivalence + fused_output = transformed_graph(*self.example_inputs) + torch.testing.assert_close(reference_output, fused_output, rtol=1e-3, atol=1e-3) + + def test_multiple_add_fusion(self): + """Multiple independent adds should create multiple quantized_add nodes""" + + class MultipleAddModel(torch.nn.Module): + def forward(self, x, y): + z1 = x + y # First add + z2 = x + z1 # Second add + return z2 + + # Prepare model + edge_program = self._prepare_quantized_model(MultipleAddModel) + edge_graph = edge_program.exported_program().graph_module + + # Get reference output + reference_output = edge_graph(*self.example_inputs) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Verify multiple fusions occurred + check_count( + transformed_graph, + exir_ops.edge.cortex_m.quantized_add.default, + 2, # Should have 2 fused quantized_add nodes + ) + + # Verify numerical equivalence + fused_output = transformed_graph(*self.example_inputs) + torch.testing.assert_close(reference_output, fused_output, rtol=1e-3, atol=1e-3) + + def test_no_fusion_without_pattern(self): + """Add without proper Q/DQ pattern should not be fused""" + + class NonQuantizedAddModel(torch.nn.Module): + def forward(self, x, y): + # This will have add but not the full Q/DQ pattern after quantization + return torch.relu(x + y) # ReLU breaks the pattern + + # For this test, we'll create a model that doesn't have the complete pattern + # We need to manually construct a graph that has add without full Q/DQ + + model = NonQuantizedAddModel() + exported = export(model, self.example_inputs, strict=True) + edge_program = executorch.exir.to_edge( + exported, + compile_config=executorch.exir.EdgeCompileConfig(_check_ir_validity=False), + ) + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Verify no fusion occurred + check_count( + transformed_graph, + exir_ops.edge.cortex_m.quantized_add.default, + 0, # Should have no fused quantized_add nodes + ) + + def test_precomputed_parameters(self): + """Fused node should have precomputed multipliers/shifts instead of scales""" + + class SingleAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # Prepare model + edge_program = self._prepare_quantized_model(SingleAddModel) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Get arguments of the fused quantized_add node + quantized_add_args = get_node_args( + transformed_graph, exir_ops.edge.cortex_m.quantized_add.default + ) + + # Should have exactly one quantized_add node + self.assertEqual(len(quantized_add_args), 1) + args = quantized_add_args[0] + + # Verify argument structure: (tensor1, zp1, mult1, shift1, tensor2, zp2, mult2, shift2, out_zp, out_mult, out_shift) + self.assertEqual(len(args), 11, "quantized_add should have 11 arguments") + + # Check that multipliers and shifts are integers (not floats/scales) + # args[2], args[3] = input1 multiplier, shift + # args[6], args[7] = input2 multiplier, shift + # args[9], args[10] = output multiplier, shift + for i in [2, 3, 6, 7, 9, 10]: # multiplier and shift positions + self.assertIsInstance( + args[i], int, f"Argument {i} should be an integer (precomputed)" + ) + + def test_mixed_fusion_pattern(self): + """Mixed pattern (some fusable, some not) should partially fuse""" + + class MixedModel(torch.nn.Module): + def forward(self, x, y): + z1 = x + y # This should fuse + z2 = torch.relu(z1) # ReLU breaks next fusion + z3 = z2 + x # This won't have full Q/DQ pattern + return z3 + + # Prepare model + edge_program = self._prepare_quantized_model(MixedModel) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Should have partial fusion (at least 1, but not necessarily all adds) + quantized_add_count = sum( + 1 + for node in transformed_graph.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.edge.cortex_m.quantized_add.default + ) + + self.assertGreaterEqual( + quantized_add_count, 1, "Should have at least 1 fused operation" + ) + + def test_different_tensor_shapes(self): + """Different tensor shapes should still fuse correctly""" + + class SingleAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # Test with different input shapes + for shape in [(2, 3), (10, 20, 30), (1,)]: + with self.subTest(shape=shape): + inputs = (torch.randn(shape), torch.randn(shape)) + + model = SingleAddModel() + exported_model = export(model.eval(), inputs, strict=True).module() + prepared_model = prepare_pt2e(exported_model, AddQuantizer()) + quantized_model = convert_pt2e(prepared_model) + + exported = export(quantized_model, inputs, strict=True) + edge_program = executorch.exir.to_edge( + exported, + compile_config=executorch.exir.EdgeCompileConfig( + _check_ir_validity=False + ), + ) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Verify fusion occurred regardless of shape + check_count( + transformed_graph, exir_ops.edge.cortex_m.quantized_add.default, 1 + ) + + def test_aot_parameter_computation_accuracy(self): + """Verify that AoT-computed parameters match runtime computation""" + + class SingleAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # Prepare model + edge_program = self._prepare_quantized_model(SingleAddModel) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + transformed_graph = transformed_program.exported_program().graph_module + + # Get the fused node arguments + quantized_add_args = get_node_args( + transformed_graph, exir_ops.edge.cortex_m.quantized_add.default + )[0] + + # Extract the computed multipliers and shifts + input1_mult, input1_shift = quantized_add_args[2], quantized_add_args[3] + input2_mult, input2_shift = quantized_add_args[6], quantized_add_args[7] + output_mult, output_shift = quantized_add_args[9], quantized_add_args[10] + + # Verify they are reasonable values + # Multipliers should be in int32 range + self.assertTrue(-(2**31) <= input1_mult < 2**31) + self.assertTrue(-(2**31) <= input2_mult < 2**31) + self.assertTrue(-(2**31) <= output_mult < 2**31) + + # Shifts should be reasonable (typically -31 to 31) + self.assertTrue(-50 <= input1_shift <= 50) + self.assertTrue(-50 <= input2_shift <= 50) + self.assertTrue(-50 <= output_shift <= 50) + + # Output multiplier should be close to 2^30 (for 1.0 scale) + self.assertAlmostEqual(output_mult, 2**30, delta=1000) + self.assertEqual(output_shift, -1) + + def test_executorch_program_generation(self): + """Verify ExecuTorch program generation with fused ops""" + + class SingleAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # Prepare model + edge_program = self._prepare_quantized_model(SingleAddModel) + + # Apply passes + transformed_program = self._apply_passes(edge_program) + + # Generate ExecutorTorch program + executorch_program = transformed_program.to_executorch() + + # Verify the program contains the expected fused operator + operator_names = [ + op.name + for op in executorch_program.executorch_program.execution_plan[0].operators + ] + + self.assertIn("cortex_m::quantized_add", operator_names) + self.assertIn("cortex_m::quantize_per_tensor", operator_names) + self.assertIn("cortex_m::dequantize_per_tensor", operator_names) + # quantize_per_tensor --> dequantize_per_tensor --> add --> quantize_per_tensor --> dequantize_per_tensor + # (input quant) (dequant) (fp32 add) (re-quant) (dequant) + # ↓ + # Fusion Pass detects pattern: + # dequantize_per_tensor --> quantized_add (Fused node) --> quantize_per_tensor + + def test_broadcastable_shapes(self): + """Verify that broadcastable shapes are supported""" + + class BroadcastAddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + # input broadcastable shapes + inputs = (torch.randn(4, 1), torch.randn(4, 8)) + print(inputs) + + # Prepare quantized model + edge_program = self._prepare_quantized_model(BroadcastAddModel) + + # Get unfused output + unfused_graph = edge_program.exported_program().graph_module + unfused_output = unfused_graph(*inputs) + if isinstance(unfused_output, tuple): + unfused_output = unfused_output[0] + + # Apply fusion pass + fused_program = self._apply_passes(edge_program) + fused_graph = fused_program.exported_program().graph_module + fused_output = fused_graph(*inputs) + if isinstance(fused_output, tuple): + fused_output = fused_output[0] + + # Check fusion occurred + check_count(fused_graph, exir_ops.edge.cortex_m.quantized_add.default, 1) + + # Compare fused vs unfused (both quantized) + torch.testing.assert_close(fused_output, unfused_output, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cortex_m/test/test_replace_quant_nodes.py b/backends/cortex_m/test/test_replace_quant_nodes.py index 08e75b17d9d..94c7762251d 100644 --- a/backends/cortex_m/test/test_replace_quant_nodes.py +++ b/backends/cortex_m/test/test_replace_quant_nodes.py @@ -1,106 +1,30 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 -from dataclasses import dataclass -from typing import Optional -import executorch import executorch.backends.cortex_m.ops.operators # noqa +import executorch.exir + import torch from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import ( ReplaceQuantNodesPass, ) -from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import export, export_for_training -from torch.fx import GraphModule -from torchao.quantization.pt2e.observer import HistogramObserver -from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e -from torchao.quantization.pt2e.quantizer import ( - QuantizationAnnotation, - QuantizationSpec, - Quantizer, +from executorch.backends.cortex_m.test.test_helpers_passes_utils import ( + AddQuantizer, + check_count, ) -from torchao.quantization.pt2e.quantizer.quantizer import Q_ANNOTATION_KEY - - -@dataclass(eq=True, frozen=True) -class QuantizationConfig: - input_activation: Optional[QuantizationSpec] - output_activation: Optional[QuantizationSpec] - - -class AddQuantizer(Quantizer): - def __init__(self): - super().__init__() - - @staticmethod - def _get_qspec(): - return QuantizationSpec( - dtype=torch.int8, - quant_min=-128, - quant_max=127, - qscheme=torch.per_tensor_symmetric, - is_dynamic=False, - observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12), - ) - - @staticmethod - def _get_qconfig(): - qspec = AddQuantizer._get_qspec() - return QuantizationConfig( - input_activation=qspec, - output_activation=qspec, - ) - - def annotate(self, model: GraphModule): - config = self._get_qconfig() - annotated_partitions = [] - - for node in model.graph.nodes: - if node.op != "call_function" or node.target not in [ - torch.ops.aten.add.Tensor, - torch.ops.aten.add_.Tensor, - ]: - continue - - if Q_ANNOTATION_KEY in node.meta and node.meta[Q_ANNOTATION_KEY]._annotated: - continue - - input_qspec_map = { - node.args[0]: config.input_activation, - node.args[1]: config.input_activation, - } - - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=config.output_activation, - _annotated=True, - ) - annotated_partitions.append([node]) - - return annotated_partitions - - def validate(self, model: GraphModule) -> None: - pass - +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.program._program import _transform -def check_count( - graph_module: GraphModule, op: torch.fx.node.Target, expected_count: int -): - actual_count = sum( - 1 - for node in graph_module.graph.nodes - if node.op == "call_function" and node.target == op - ) +from torch.export import export - assert ( - actual_count == expected_count - ), f"Expected {expected_count} {op} nodes, got {actual_count}" +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e class TestReplaceQuantOps(unittest.TestCase): @@ -125,15 +49,20 @@ def forward(self, x): example_inputs = (torch.randn(10, 11, 12),) # Step 1: Export and quantize the model - exported_model = export_for_training( - model.eval(), example_inputs, strict=True - ).module() + exported_model = export(model.eval(), example_inputs, strict=True).module() prepared_model = prepare_pt2e(exported_model, AddQuantizer()) + prepared_model(*example_inputs) quantized_model = convert_pt2e(prepared_model) # Step 2: Export to EXIR exported = export(quantized_model, example_inputs, strict=True) + # The pass should raise an Exception if ran before to_edge. + with self.assertRaisesRegex( + Exception, "An error occurred when running the 'ReplaceQuantNodesPass' pass" + ): + _transform(exported, ReplaceQuantNodesPass()) + # Step 3: Convert to Edge edge_program = executorch.exir.to_edge( exported, @@ -199,3 +128,7 @@ def forward(self, x): "cortex_m::quantize_per_tensor", "cortex_m::dequantize_per_tensor", ], f"Unexpected op {op.name}" + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/example/TARGETS b/backends/example/TARGETS index e99a408cdb7..48fb971428b 100644 --- a/backends/example/TARGETS +++ b/backends/example/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") -python_library( +runtime.python_library( name = "example_quantizer", srcs = [ "example_quantizer.py", @@ -15,7 +15,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "example_backend", srcs = [ "example_backend.py", @@ -27,7 +27,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "example_partitioner", srcs = [ "example_partitioner.py", diff --git a/backends/example/example_backend_delegate_passes/TARGETS b/backends/example/example_backend_delegate_passes/TARGETS index bc546ce410a..9e3b06486a8 100644 --- a/backends/example/example_backend_delegate_passes/TARGETS +++ b/backends/example/example_backend_delegate_passes/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "merge_to_dim_pass.py", diff --git a/backends/example/example_operators/TARGETS b/backends/example/example_operators/TARGETS index a4d85c8fb71..70f8f3cc0c5 100644 --- a/backends/example/example_operators/TARGETS +++ b/backends/example/example_operators/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "example_operators_lib", srcs = [ "adaptive_avg_pool2d.py", diff --git a/backends/example/test_example_delegate.py b/backends/example/test_example_delegate.py index bc6ad4d7e4c..fd4c6652787 100644 --- a/backends/example/test_example_delegate.py +++ b/backends/example/test_example_delegate.py @@ -46,9 +46,7 @@ def get_example_inputs(): ) m = model.eval() - m = torch.export.export_for_training( - m, copy.deepcopy(example_inputs), strict=True - ).module() + m = torch.export.export(m, copy.deepcopy(example_inputs), strict=True).module() # print("original model:", m) quantizer = ExampleQuantizer() # quantizer = XNNPACKQuantizer() @@ -84,9 +82,7 @@ def test_delegate_mobilenet_v2(self): ) m = model.eval() - m = torch.export.export_for_training( - m, copy.deepcopy(example_inputs), strict=True - ).module() + m = torch.export.export(m, copy.deepcopy(example_inputs), strict=True).module() quantizer = ExampleQuantizer() m = prepare_pt2e(m, quantizer) diff --git a/backends/mediatek/CMakeLists.txt b/backends/mediatek/CMakeLists.txt index 23e50e8cd8a..ed9b37e1998 100644 --- a/backends/mediatek/CMakeLists.txt +++ b/backends/mediatek/CMakeLists.txt @@ -30,12 +30,13 @@ target_link_libraries( ) target_sources( neuron_backend - INTERFACE $ - $ - $ - $ - $ - $ + INTERFACE + $ + $ + $ + $ + $ + $ PRIVATE ${CMAKE_CURRENT_LIST_DIR}/runtime/NeuronBackend.cpp ${CMAKE_CURRENT_LIST_DIR}/runtime/NeuronExecutor.cpp ${CMAKE_CURRENT_LIST_DIR}/runtime/NeuronBufferAllocator.cpp diff --git a/backends/mediatek/preprocess.py b/backends/mediatek/preprocess.py index 6ffaebb7589..0e2b68335e0 100644 --- a/backends/mediatek/preprocess.py +++ b/backends/mediatek/preprocess.py @@ -4,14 +4,16 @@ # except in compliance with the License. See the license file in the root # directory of this source tree for more details. +import collections import contextlib import struct -from typing import final, List +from typing import Dict, final, List import mtk_converter import mtk_neuron import torch +from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import ( BackendDetails, ExportedProgram, @@ -20,8 +22,11 @@ from executorch.exir.backend.compile_spec_schema import CompileSpec SKIP_COMPILE_SPEC_KEYS = {"ImportForever"} +EXTRACT_SHARED_BLOB_KEY = "ExtractSharedBlobKey" +HEADER_SIZE = 13 +HEADER_VERSION = 1 REQUIRED_COMPILE_SPEC_KEYS = {"platform-config"} -SUPPORTED_PLATFORM_CONFIGS = {"mt6989", "mt6991"} +SUPPORTED_PLATFORM_CONFIGS = {"mt6989", "mt6991", "mt6993"} def assert_default_dim_order(edge_graph_module: torch.fx.GraphModule) -> None: @@ -41,6 +46,21 @@ def assert_default_dim_order(edge_graph_module: torch.fx.GraphModule) -> None: ) +def _pack_header(num_inputs, num_outputs, model_bytes_size): + header_bytes = struct.pack( + " Dict[str, list[PreprocessResult]]: + + # Follow the default behavior of `preprocess_multimethod` + preprocess_results = {} + for method_name, programs in edge_programs.items(): + assert ( + method_name in compile_specs + ), f"Error: missing compile specs for {method_name}" + compile_specs_for_method = compile_specs[method_name] + assert len(compile_specs_for_method) == len( + programs + ), f"Error: method {method_name} has {len(programs)} partitions but only {len(compile_specs_for_method)}" + results_for_method = [] + for program, compile_spec_for_program in zip( + programs, compile_specs_for_method + ): + preprocess_result = cls.preprocess(program, compile_spec_for_program) + results_for_method.append(preprocess_result) + + preprocess_results[method_name] = results_for_method + + # Try extract shared data blob if necessary + infos_dict = collections.defaultdict(list) + models_dict = collections.defaultdict(list) + result_dict = collections.defaultdict(list) + for method_name, method_results in preprocess_results.items(): + for idx, result in enumerate(method_results): + shared_blob_key = None + for spec in compile_specs[method_name][idx]: + if spec.key == EXTRACT_SHARED_BLOB_KEY: + shared_blob_key = spec.value.decode("utf-8") + + if shared_blob_key is None: + continue + + header_bytes = result.processed_bytes[:HEADER_SIZE] + model_bytes = result.processed_bytes[HEADER_SIZE:] + num_inputs, num_outputs, model_bytes_size = _unpack_header(header_bytes) + assert len(model_bytes) == model_bytes_size + infos_dict[shared_blob_key].append((num_inputs, num_outputs)) + models_dict[shared_blob_key].append(model_bytes) + result_dict[shared_blob_key].append(result) + + data_store_output_dict = {} + for key, models in models_dict.items(): + ndm = NamedDataStore() + blob, new_models = mtk_neuron.extract_shared_data( + models, options="-e union" + ) + ndm.add_named_data(key, bytes(blob)) + data_store_output_dict[key] = ndm.get_named_data_store_output() + models.clear() + models.extend(new_models) + + for key, data_store_output in data_store_output_dict.items(): + for idx, (model_info, model_bytes) in enumerate( + zip(infos_dict[key], models_dict[key]) + ): + num_inputs, num_outputs = model_info + header_bytes = _pack_header(num_inputs, num_outputs, len(model_bytes)) + result_dict[key][idx].data_store_output = data_store_output + result_dict[key][idx].processed_bytes = bytes( + header_bytes + model_bytes + ) + + return preprocess_results diff --git a/backends/mediatek/quantizer/annotator.py b/backends/mediatek/quantizer/annotator.py index 8c0e42627e0..6fd4a1d23c3 100644 --- a/backends/mediatek/quantizer/annotator.py +++ b/backends/mediatek/quantizer/annotator.py @@ -10,7 +10,7 @@ from torch._ops import OpOverload from torch._subclasses import FakeTensor -from torch.export import export_for_training +from torch.export import export from torch.fx import Graph, Node from torch.fx.passes.utils.matcher_with_name_node_map_utils import ( SubgraphMatcherWithNameNodeMap, @@ -158,9 +158,7 @@ def forward(self, x): return norm, {} for pattern_cls in (ExecuTorchPattern, MTKPattern): - pattern_gm = export_for_training( - pattern_cls(), (torch.randn(3, 3),), strict=True - ).module() + pattern_gm = export(pattern_cls(), (torch.randn(3, 3),), strict=True).module() matcher = SubgraphMatcherWithNameNodeMap( pattern_gm, ignore_literals=True, remove_overlapping_matches=False ) diff --git a/backends/mediatek/runtime/NeuronBackend.cpp b/backends/mediatek/runtime/NeuronBackend.cpp index 15b82e04129..7b4084f66b8 100644 --- a/backends/mediatek/runtime/NeuronBackend.cpp +++ b/backends/mediatek/runtime/NeuronBackend.cpp @@ -12,8 +12,8 @@ #include "NeuronPayloadHeader.h" #include "api/NeuronAdapter.h" +#include #include "executorch/runtime/core/error.h" -#include "executorch/runtime/core/exec_aten/util/dim_order_util.h" #include #include @@ -24,6 +24,7 @@ namespace executorch { namespace backends { namespace neuron { +using executorch::ET_RUNTIME_NAMESPACE::NamedDataMap; using executorch::runtime::ArrayRef; using executorch::runtime::BackendExecutionContext; using executorch::runtime::BackendInitContext; @@ -34,15 +35,26 @@ using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::MemoryAllocator; using executorch::runtime::Result; +using executorch::runtime::Span; const char kHighAddrKey[] = "HighAddr"; const char kImportForeverKey[] = "ImportForever"; +const char kSharedWeightsKey[] = "ExtractSharedBlobKey"; Result NeuronBackend::init( BackendInitContext& context, FreeableBuffer* processed, ArrayRef compile_specs) const { NeuronDelegateSetting setting; + MemoryAllocator* runtime_allocator = context.get_runtime_allocator(); + NeuronExecuTorchDelegate* delegate = + runtime_allocator->allocateInstance(); + if (delegate == nullptr) { + return Error::MemoryAllocationFailed; + } + + new (delegate) NeuronExecuTorchDelegate(); + for (auto& compile_spec : compile_specs) { if (std::strcmp(compile_spec.key, kHighAddrKey) == 0) { setting.mHighAddr = *static_cast(compile_spec.value.buffer); @@ -53,11 +65,62 @@ Result NeuronBackend::init( "NeuronBackend", "IsImportForever Enable : %d", setting.mImportForever); + } else if (std::strcmp(compile_spec.key, kSharedWeightsKey) == 0) { + setting.mSharedWeights = true; + std::string shared_weights_key( + static_cast(compile_spec.value.buffer), + compile_spec.value.nbytes); + LogInfo( + "NeuronBackend", + "SharedWeights Enabled for %s", + shared_weights_key.c_str()); + std::shared_ptr neuron_shared_weights; + if (neuron_shared_weights_cache_.find(shared_weights_key) != + neuron_shared_weights_cache_.end()) { + neuron_shared_weights = + neuron_shared_weights_cache_.at(shared_weights_key).lock(); + if (neuron_shared_weights) { + LogInfo( + "NeuronBackend", + "Reusing cached shared weights with key %s", + shared_weights_key.c_str()); + delegate->SetSharedWeights(neuron_shared_weights); + continue; + } else { + LogInfo( + "NeuronBackend", + "Shared weights cache expired: %s", + shared_weights_key.c_str()); + neuron_shared_weights_cache_.erase(shared_weights_key); // Expired + } + } + const NamedDataMap* named_data_map = context.get_named_data_map(); + Result shared_weights = + named_data_map->get_data(shared_weights_key.c_str()); + + if (shared_weights.ok()) { + LogInfo( + "NeuronBackend", + "Loaded shared weights from named_data_map. Size: %zu", + shared_weights.get().size()); + FreeableBuffer& buffer = shared_weights.get(); + neuron_shared_weights = + std::make_shared(std::move(buffer)); + delegate->SetSharedWeights(neuron_shared_weights); + neuron_shared_weights_cache_[shared_weights_key] = + neuron_shared_weights; + } else { + LogError( + "NeuronBackend", + "Failed to load shared weights from named_data_map."); + return Error::Internal; + } } else { LogWarn("NeuronBackend", "unknown compile spec: %s", compile_spec.key); } } auto Payload = NeuronPayload(processed->data(), processed->size()); + LogInfo( "NeuronBackend", "version %u, input %u, output %u, length %u, payload size: %zu", @@ -67,26 +130,14 @@ Result NeuronBackend::init( Payload.Header.DataLen, processed->size()); - MemoryAllocator* runtime_allocator = context.get_runtime_allocator(); - NeuronExecuTorchDelegate* delegate = - runtime_allocator->allocateInstance(); - if (delegate == nullptr) { - return Error::MemoryAllocationFailed; - } - - new (delegate) NeuronExecuTorchDelegate(); - - if (delegate == nullptr) { - return nullptr; - } - auto res = delegate->LoadCompiledNetwork(Payload, setting); + int res = delegate->LoadCompiledNetwork(Payload, setting); return res == NEURON_NO_ERROR ? delegate : nullptr; } Error NeuronBackend::execute( ET_UNUSED BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const { + Span args) const { NeuronExecuTorchDelegate* delegate = reinterpret_cast(handle); return delegate->execute(context, args); @@ -106,26 +157,27 @@ bool NeuronBackend::is_available() const { Error NeuronExecuTorchDelegate::execute( BackendExecutionContext& context, - EValue** args) const { + Span args) const { if (HintNeuronBackend(args) != NEURON_NO_ERROR) { return Error::InvalidState; }; + ET_CHECK_OR_RETURN_ERROR( + CheckDimOrder(args) == NEURON_NO_ERROR, + Internal, + "Expecting default dim_order but got a non default dim_order tensor input"); + + PrepareInputsOuputs(args); + auto allocator = dynamic_cast(context.get_temp_allocator()); - size_t inputCount = mInputSizes.size(), outputCount = mOutputSizes.size(); - - for (int i = 0; i < inputCount; i++) { - auto tensor_in = args[i]->toTensor(); - ET_CHECK_OR_RETURN_ERROR( - runtime::is_contiguous_dim_order( - tensor_in.dim_order().data(), tensor_in.dim()), - Internal, - "Expecting default dim_order but got a non default dim_order tensor for external input %u", - i); - - auto data_ptr = args[i]->toTensor().data_ptr(); - auto data_size = args[i]->toTensor().nbytes(); + + size_t inputCount = mInputSizes.size() + neuron_shared_weights_.size(); + size_t outputCount = mOutputSizes.size(); + + for (size_t i = 0; i < inputCount; i++) { + auto data_ptr = mPreparedInputs[i].data_ptr; + auto data_size = mPreparedInputs[i].size; if (IsCached(i, data_ptr)) { continue; }; @@ -140,22 +192,20 @@ Error NeuronExecuTorchDelegate::execute( } } - for (int o = inputCount; o < inputCount + outputCount; o++) { - auto data_ptr = args[o]->toTensor().data_ptr(); - auto data_size = args[o]->toTensor().nbytes(); - auto output_index = o - inputCount; - if (IsCached(output_index, data_ptr)) { + for (size_t o = 0; o < outputCount; o++) { + auto data_ptr = mPreparedOutputs[o].data_ptr; + auto data_size = mPreparedOutputs[o].size; + if (IsCached(o, data_ptr)) { continue; }; auto unit = allocator != nullptr ? allocator->Find(data_ptr) : nullptr; if (unit) { - UpdateCache(output_index, data_ptr); + UpdateCache(o, data_ptr); size_t offset = (char*)data_ptr - (char*)unit->GetAddress(); mExecutor.SetInputOutputFromMemory( - output_index, unit->GetNeuronMemory(), offset, data_size); + o, unit->GetNeuronMemory(), offset, data_size); } else { - mExecutor.SetInputOutput( - output_index, data_ptr, data_size); + mExecutor.SetInputOutput(o, data_ptr, data_size); } } @@ -163,8 +213,8 @@ Error NeuronExecuTorchDelegate::execute( : Error::InvalidState; }; -int NeuronExecuTorchDelegate::HintNeuronBackend(EValue** args) const { - auto HintImportForever = [this](EValue** args) -> int { +int NeuronExecuTorchDelegate::HintNeuronBackend(Span args) const { + auto HintImportForever = [this](Span args) -> int { auto& allocator = GET_NEURON_ALLOCATOR; size_t inputCount = mInputSizes.size(), outputCount = mOutputSizes.size(); for (int i = 0; i < inputCount; i++) { diff --git a/backends/mediatek/runtime/include/NeuronBackend.h b/backends/mediatek/runtime/include/NeuronBackend.h index 570cc5dca59..1d2e8563ab3 100644 --- a/backends/mediatek/runtime/include/NeuronBackend.h +++ b/backends/mediatek/runtime/include/NeuronBackend.h @@ -18,6 +18,7 @@ #include #include #include +#include "executorch/runtime/core/exec_aten/util/dim_order_util.h" #include #include @@ -27,6 +28,50 @@ namespace executorch { namespace backends { namespace neuron { +using executorch::runtime::EValue; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; +using executorch::runtime::Span; + +class NeuronSharedWeights { + public: + explicit NeuronSharedWeights(const FreeableBuffer& shared_weights_buffer) { + auto& buffer_allocator = GET_NEURON_ALLOCATOR; + nbytes_ = shared_weights_buffer.size(); + data_ = buffer_allocator.Allocate(nbytes_); + ET_CHECK_MSG( + data_ != nullptr, + "Error: Failed to allocate memory for shared weights of size %zu", + nbytes_); + std::memcpy(data_, shared_weights_buffer.data(), nbytes_); + } + + explicit NeuronSharedWeights(FreeableBuffer&& shared_weights_buffer) + : NeuronSharedWeights(shared_weights_buffer) { + shared_weights_buffer.Free(); + } + + ~NeuronSharedWeights() { + if (data_ == nullptr || nbytes_ == 0) { + return; + } + auto& buffer_allocator = GET_NEURON_ALLOCATOR; + buffer_allocator.RemoveBuffer(data_); + } + + void* data() const { + return data_; + } + + size_t size() const { + return nbytes_; + } + + private: + void* data_ = nullptr; + size_t nbytes_ = 0; +}; + class NeuronBackend final : public ::executorch::runtime::BackendInterface { public: ::executorch::runtime::Result<::executorch::runtime::DelegateHandle*> init( @@ -38,11 +83,16 @@ class NeuronBackend final : public ::executorch::runtime::BackendInterface { ::executorch::runtime::Error execute( ET_UNUSED ::executorch::runtime::BackendExecutionContext& context, ::executorch::runtime::DelegateHandle* handle, - ::executorch::runtime::EValue** args) const override; + ::executorch::runtime::Span<::executorch::runtime::EValue*> args) + const override; void destroy(::executorch::runtime::DelegateHandle* handle) const override; bool is_available() const override; + + private: + mutable std::unordered_map> + neuron_shared_weights_cache_; }; extern const char kHighAddrKey[]; @@ -53,6 +103,8 @@ struct NeuronDelegateSetting { bool mImportForever = false; + bool mSharedWeights = false; + std::string ToRuntimeOption() { if (mHighAddr && mImportForever) { return "--apusys-config \"{ \\\"high_addr\\\": true, \\\"import_forever\\\": true }\""; @@ -68,6 +120,13 @@ struct NeuronDelegateSetting { class NeuronExecuTorchDelegate { public: + struct InputOutputInfo { + void* data_ptr; + size_t size; + + InputOutputInfo(void* ptr, size_t sz) : data_ptr(ptr), size(sz) {} + }; + class MemoryCache { public: template @@ -103,19 +162,25 @@ class NeuronExecuTorchDelegate { auto res = mExecutor.LoadFromCompiledNetwork( payload.CompiledNetwork, payload.Header.DataLen, - payload.Header.InputCount, + mSettings.mSharedWeights ? payload.Header.InputCount + 1 + : payload.Header.InputCount, payload.Header.OutputCount, runtimeOption); CHECK_NO_ERROR(res); CHECK_TRUE(mExecutor.IsValid()); - SummaryIoCounts(); + SummarizeIoSizes(payload.Header.InputCount, payload.Header.OutputCount); mPLock = std::unique_ptr(new ScopePerformancer); return NEURON_NO_ERROR; } + int SetSharedWeights(std::shared_ptr sharedWeights) { + neuron_shared_weights_.push_back(sharedWeights); + return NEURON_NO_ERROR; + } + ::executorch::runtime::Error execute( ET_UNUSED ::executorch::runtime::BackendExecutionContext& context, - ::executorch::runtime::EValue** args) const; + ::executorch::runtime::Span<::executorch::runtime::EValue*> args) const; private: template @@ -128,19 +193,19 @@ class NeuronExecuTorchDelegate { mCache.UpdateCache(index, ptr); } - int SummaryIoCounts() { - for (int i = 0;; i++) { + int SummarizeIoSizes(uint32_t input_count, uint32_t output_count) { + for (int i = 0; i < input_count; i++) { size_t size = mExecutor.GetInputOutputPaddedSize(i); if (size == 0) { - break; + LogWarn("NeuronBackend", "Model input:%d got size: %lu", i, size); } LogInfo("NeuronBackend", "Model input:%d size: %lu", i, size); mInputSizes.push_back(size); } - for (int o = 0;; o++) { + for (int o = 0; o < output_count; o++) { size_t size = mExecutor.GetInputOutputPaddedSize(o); if (size == 0) { - break; + LogWarn("NeuronBackend", "Model output:%d got size: %lu", o, size); } LogInfo("NeuronBackend", "Model output:%d size: %lu", o, size); mOutputSizes.push_back(size); @@ -148,13 +213,70 @@ class NeuronExecuTorchDelegate { return NEURON_NO_ERROR; } - int HintNeuronBackend(::executorch::runtime::EValue** args) const; + int CheckDimOrder(Span args) const { + size_t data_input_count = mInputSizes.size(); + for (int i = 0; i < data_input_count; i++) { + auto tensor_in = args[i]->toTensor(); + LogInfo("NeuronBackend", "Checking dim order for input %d", i); + if (!runtime::is_contiguous_dim_order( + tensor_in.dim_order().data(), tensor_in.dim())) { + return NEURON_BAD_DATA; + } + } + + return NEURON_NO_ERROR; + } + + int PrepareInputsOuputs(Span args) const { + bool has_shared_weights_input = neuron_shared_weights_.size() > 0; + + size_t data_input_count = mInputSizes.size(); + size_t data_output_count = mOutputSizes.size(); + + mPreparedInputs.clear(); + mPreparedOutputs.clear(); + mPreparedInputs.reserve(data_input_count); + mPreparedOutputs.reserve(data_output_count); + + // Prepare input data + for (int i = 0; i < data_input_count; i++) { + auto tensor_in = args[i]->toTensor(); + auto data_ptr = tensor_in.data_ptr(); + auto data_size = tensor_in.nbytes(); + mPreparedInputs.push_back(InputOutputInfo{data_ptr, data_size}); + } + + // Prepare shared weights if any as the last model inputs + if (has_shared_weights_input) { + for (const auto& shared_weights : neuron_shared_weights_) { + mPreparedInputs.push_back( + InputOutputInfo{shared_weights->data(), shared_weights->size()}); + } + } + + // Prepare output data + for (int o = data_input_count; o < data_input_count + data_output_count; + o++) { + auto tensor_out = args[o]->toTensor(); + auto data_ptr = tensor_out.data_ptr(); + auto data_size = tensor_out.nbytes(); + mPreparedOutputs.push_back(InputOutputInfo{data_ptr, data_size}); + } + + return NEURON_NO_ERROR; + } + + int HintNeuronBackend(Span args) const; private: std::vector mInputSizes; std::vector mOutputSizes; + mutable std::vector mPreparedInputs; + + mutable std::vector mPreparedOutputs; + mutable MemoryCache mCache; std::unique_ptr mPLock; @@ -165,6 +287,9 @@ class NeuronExecuTorchDelegate { mutable std::unordered_set mHasImported; + mutable std::vector> + neuron_shared_weights_; + private: NeuronExecuTorchDelegate(const NeuronExecuTorchDelegate&); diff --git a/backends/nxp/CMakeLists.txt b/backends/nxp/CMakeLists.txt index 0a9d72d3555..43fcaa24d19 100644 --- a/backends/nxp/CMakeLists.txt +++ b/backends/nxp/CMakeLists.txt @@ -3,16 +3,19 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -set( - _common_include_directories - ${CMAKE_CURRENT_SOURCE_DIR}/../../.. - ${CMAKE_CURRENT_SOURCE_DIR}/../../runtime/core/portable_type/c10 -) +set(_common_include_directories ${CMAKE_CURRENT_SOURCE_DIR}/../../..) add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) -set(_neutron_sources ${CMAKE_CURRENT_SOURCE_DIR}/runtime/NeutronBackend.cpp ) +set(_neutron_sources ${CMAKE_CURRENT_SOURCE_DIR}/runtime/NeutronBackend.cpp) add_library(executorch_delegate_neutron STATIC ${_neutron_sources}) target_include_directories( - executorch_delegate_neutron PUBLIC ${_common_include_directories} + executorch_delegate_neutron + PRIVATE ${_common_include_directories} + ${PROJECT_SOURCE_DIR}/runtime/core/portable_type/c10 +) +install( + TARGETS executorch_delegate_neutron + EXPORT ExecuTorchTargets + DESTINATION lib ) diff --git a/backends/nxp/TARGETS b/backends/nxp/TARGETS index 08c250c5c20..d56ac60242c 100644 --- a/backends/nxp/TARGETS +++ b/backends/nxp/TARGETS @@ -1,4 +1,4 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load( "@fbsource//arvr/tools/build_defs:oxx_python.bzl", @@ -8,7 +8,7 @@ load( oncall("executorch") -python_library( +runtime.python_library( name = "aten_passes", srcs = glob([ "aten_passes/*.py", @@ -19,7 +19,20 @@ python_library( ], ) -python_library( +runtime.python_library( + name = "edge_passes", + srcs = glob([ + "edge_passes/*.py", + ]), + deps = [ + ":neutron_backend", + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/exir:pass_manager", + ], +) + +runtime.python_library( name = "quantizer", srcs = [ "quantizer/neutron_quantizer.py", @@ -33,15 +46,15 @@ python_library( ], ) -python_library( +runtime.python_library( name = "neutron_sdk", srcs = glob(["backend/**/*.py"]), deps = [ - "fbsource//third-party/pypi/neutron_convertor_SDK_25_03:neutron_convertor_SDK_25_03", + "fbsource//third-party/pypi/neutron_converter:neutron_converter", ], ) -python_library( +runtime.python_library( name = "neutron_backend", srcs = [ "nxp_backend.py", diff --git a/backends/nxp/aten_passes/neutron_aten_pass_manager.py b/backends/nxp/aten_passes/neutron_aten_pass_manager.py index 6c989b6c136..f6e3c374b19 100644 --- a/backends/nxp/aten_passes/neutron_aten_pass_manager.py +++ b/backends/nxp/aten_passes/neutron_aten_pass_manager.py @@ -13,11 +13,20 @@ from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import ( FuseBatchNormWithLinearPass, ) +from executorch.backends.nxp.aten_passes.remove_nodes_with_known_outputs import ( + RemoveNodesWithKnownOutputs, +) +from executorch.backends.nxp.aten_passes.split_group_convolution import ( + SplitGroupConvolution, +) +from executorch.backends.nxp.aten_passes.split_gru_based_on_num_layers import ( + SplitGRUBasedOnNumLayers, +) from executorch.exir.pass_manager import PassManager from torch import nn from torch.fx.passes.infra.pass_base import PassResult -PassType = list[type[Callable[[torch.fx.GraphModule], PassResult]]] +PassType = type[Callable[[torch.fx.GraphModule], PassResult]] class NeutronAtenPassManager(PassManager): @@ -26,6 +35,9 @@ def __init__(self, passes: list[PassType] = None): passes: list[PassType] = passes or [ FuseBatchNormWithConvPass(), FuseBatchNormWithLinearPass(), + SplitGroupConvolution(), + SplitGRUBasedOnNumLayers(), + RemoveNodesWithKnownOutputs(), ] super().__init__(passes) diff --git a/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py b/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py new file mode 100644 index 00000000000..3c08ac6c3fb --- /dev/null +++ b/backends/nxp/aten_passes/remove_nodes_with_known_outputs.py @@ -0,0 +1,196 @@ +# Copyright 2025 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 Collection + +import torch + +from executorch.backends.nxp.backend.edge_helper import ( + try_get_tensor_constant_from_node, +) +from torch._subclasses import FakeTensor, FakeTensorMode +from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix +from torch.export.unflatten import _assign_attr, _AttrKind +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.nn import Parameter + + +class RemoveNodesWithKnownOutputs(PassBase): + """In some situations, a node will always produce the same output data at runtime. If these cases are identified, + the nodes can simply be removed and replaced by a static parameter node, which holds the data the original + node would produce. + This pass identifies some of these cases and performs the replacement. + """ + + # Nodes which don't have the `.meta['val']` attribute. The datatype and shape of their inferred output data will + # therefore not be checked against the expected values in the `.meta['val']`. + nodes_without_val_meta = [ + torch.ops.aten.empty.memory_format, + ] + + module: GraphModule + + def replace_nodes_in_list_with_their_data(self, list_of_args: list) -> list | None: + """Replace the nodes in `list_of_args` by their static data. If not all data is available, return `None`. + + :param list_of_args: List of arguments of an aten operator. Can include nodes, generic arguments, lists... + :return:`list_of_args` but with tensors replaced by their static data, or `None` if not all data is available. + """ + args_with_data = [] + for arg in list_of_args: + match arg: + case Node(): + # `arg` is either another operator, a model input, or a static parameter. + data = try_get_tensor_constant_from_node(self.module, arg) + if data is None: + # No static data is available. + return None + + args_with_data.append(data) + case list(): + nested = self.replace_nodes_in_list_with_their_data(arg) + if nested is None: + return None + args_with_data.append(nested) + + case _: + # Generic argument. Not an input from a previous node. + args_with_data.append(arg) + + return args_with_data + + @staticmethod + def node_is_followed_only_by_getitem_nodes(node: Node) -> bool: + def _is_getitem(node_: Node) -> bool: + return node_.op == "call_function" and node_.target.__name__ == "getitem" + + users = list(node.users.keys()) + return all(_is_getitem(user) for user in users) + + def replace_node_with_static_data(self, node: Node, static_data: Parameter): + """Remove the given `node` from the graph and replace it with a parameter node containing the `static_data`.""" + # Generate a unique name for the new static parameter. + new_name = get_new_attr_name_with_prefix(node.name)(self.module) + + # Create the node for the parameter. + param = torch.nn.Parameter(static_data, False) + _assign_attr(param, self.module, str(new_name), _AttrKind.PARAMETER) + with self.module.graph.inserting_before(node): + static_parameter_node = self.module.graph.get_attr(new_name) + + with FakeTensorMode() as mode: + # Assign the parameter node its shape and data type. + static_parameter_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(static_data.shape, dtype=static_data.dtype), mode + ) + + # Replace the old node with the new static parameter. + node.replace_all_uses_with(static_parameter_node) + self.module.graph.erase_node(node) + + def replace_following_getitem_nodes_with_static_data( + self, root_node: Node, static_data: Collection[Parameter] + ) -> bool: + """Remove the `root_node` and all `GetItem` nodes that consume its output from the graph, and replace their + uses with parameter nodes containing the provided `static_data`. + If something other than just `GetItem` nodes follow after the `root_node`, nothing is done and `False` is + returned. + + :param root_node: The main compute node which is followed only by `GetItem` nodes. + :param static_data: A tuple of static tensors with the data that will be used to replace the `GetItem` nodes + after the `root_node`. + :return: `True` if the replacement was successfully executed. `False` otherwise. + """ + + if not self.node_is_followed_only_by_getitem_nodes(root_node): + return False # Unexpected case. + + users = list(root_node.users.keys()) + if len(users) != len(static_data): + return False # Unexpected missmatch. + + # Replace the individual `GetItem` nodes. + for get_item_node in users: + idx = get_item_node.args[1] + self.replace_node_with_static_data(get_item_node, static_data[idx]) + + # Finally remove the root node from the graph. + self.module.graph.erase_node(root_node) + + return True + + def data_matches_node_meta(self, node: Node, data: Parameter) -> bool: + """Verify that the provided `data` tensor has the same shape and datatype as the `node`.""" + if node.target not in self.nodes_without_val_meta: + if node.meta["val"].shape != data.shape: + return False # The inferred data has a different shape than expected. + + if node.meta["val"].dtype != data.dtype: + return ( + False # The inferred data has a different data type than expected. + ) + + return True + + def data_matches_meta_of_following_getitem_nodes( + self, root_node: Node, data: Collection[Parameter] + ) -> bool: + """Verify that the provided `data` tensor has the same shape and datatype as the `GetItem` nodes which consume + the output of the `root_node`. + """ + if not self.node_is_followed_only_by_getitem_nodes(root_node): + return False # Unexpected case + + users = list(root_node.users.keys()) + return all( + self.data_matches_node_meta(get_item, data[get_item.args[1]]) + for get_item in users + ) + + def call(self, module: GraphModule) -> bool: + self.module = module + made_changes = False + + for node in module.graph.nodes: + if node.op != "call_function": + continue # Not a compute operator. + + # Try to access the static data for the inputs of the node. + args_with_data = self.replace_nodes_in_list_with_their_data(node.args) + + if args_with_data is None: + # Output data inference is not possible. + continue + + # All input data is static. Run the operator to compute the input it would produce at runtime. + # noinspection PyBroadException + try: + output = node.target(*args_with_data, **node.kwargs) + + if isinstance(output, tuple) or isinstance(output, list): + if not self.data_matches_meta_of_following_getitem_nodes( + node, output + ): + continue # The inferred data does not have the expected type/shape. + else: + if not self.data_matches_node_meta(node, output): + continue # The inferred data does not have the expected type/shape. + + except Exception: + continue # Failed to infer the data. Continue with the other nodes. + # The output data appears to have been correctly inferred. Create a static parameter node for it. + + if isinstance(output, tuple) or isinstance(output, list): + # The node produces multiple outputs (e.g. `split`). If the node is followed only by `GetItem` nodes + # which extract the individual outputs, replace them by the static data. + if self.replace_following_getitem_nodes_with_static_data(node, output): + made_changes = True + + else: + self.replace_node_with_static_data(node, output) + made_changes = True # Indicate that changes were made. + + return PassResult(module, made_changes) diff --git a/backends/nxp/aten_passes/split_group_convolution.py b/backends/nxp/aten_passes/split_group_convolution.py new file mode 100644 index 00000000000..58c87730c84 --- /dev/null +++ b/backends/nxp/aten_passes/split_group_convolution.py @@ -0,0 +1,281 @@ +# Copyright 2025 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 operator + +import torch + +from executorch.backends.nxp.backend.ir.converter.node_converters.shared.conv_utils import ( + group_conv_convertible_into_multiple_convolutions, +) +from torch._subclasses import FakeTensor, FakeTensorMode +from torch.ao.quantization.fx.utils import get_new_attr_name_with_prefix +from torch.export.unflatten import _assign_attr, _AttrKind +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.nn.parameter import Parameter + + +class SplitGroupConvolution(PassBase): + """The eIQ Neutron NPU supports only regular and depthwise convolutions. Group convolutions must be decomposed into + multiple parallel single group convolutions. + Replace the nodes in the following pattern. The square brackets indicate the tensor shapes. + + + │[N, Ic, ...] + ┌───▼───┐ + │ split │ + └┬─────┬┘ + ┌──────────────────┘ ... └────────────────┐ + │[N, Ic, ...] │[N, Ic/G, ...] │[N, Ic/G, ...] + ┌──────▼──────┐ ┌──────▼──────┐ ┌──────▼──────┐ + │ convolution ◄──W [Oc, Ic/G, ...] replace │ convolution ◄──W [Oc/G, Ic/G, ...] │ convolution ◄──W [Oc/G, Ic/G, ...] + │ group=G ◄──B [Oc] ────────► │ group=1 ◄──B [Oc/G] ... │ group=1 ◄──B [Oc/G] + └──────┬──────┘ with └──────┬──────┘ └──────┬──────┘ + ▼[N, Oc, ...] │ [N, Oc/G, ...] │[N, Oc/G, ...] + └──────────────────┐ ... ┌────────────────┘ + ┌▼─────▼┐ + │ cat │ + └───┬───┘ + ▼[N, Oc, ...] + """ + + module: GraphModule + + def _get_tensor_constant_from_node(self, node) -> Parameter | None: + """Get the static data from a given node. If it doesn't have any data, return `None`.""" + if node is None or node.op != "get_attr": + return None + + target_atoms = node.target.split(".") + attr_itr = self.module + for atom in target_atoms: + if not hasattr(attr_itr, atom): + return None + attr_itr = getattr(attr_itr, atom) + return attr_itr + + def _create_and_insert_get_item_node(self, input_node: Node, idx: int) -> Node: + """Create a `GetItem` node which extracts the output of `input_node` on index `idx`. + The `GetItem` is also added to the graph right after the `input_node`. + """ + with self.module.graph.inserting_after(input_node): + get_item_node = self.module.graph.create_node( + "call_function", + operator.getitem, + (input_node, idx), + {}, + ) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + get_item_node.meta["source_fn_stack"] = [ + (get_item_node.name, input_node.meta["source_fn_stack"]) + ] + get_item_node.meta["val"] = input_node.meta["val"][idx] + + return get_item_node + + def _create_split_node(self, *split_args) -> Node: + split_target = torch.ops.aten.split.default + split_node = self.module.graph.call_function(split_target, split_args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + split_node.meta["source_fn_stack"] = [(split_node.name, torch.split)] + + # Compute the output shapes for the `split`, and assign the `val` meta. + x_val = split_args[0].meta["val"] + with FakeTensorMode() as mode: + fake_input = FakeTensor.from_tensor( + torch.empty(x_val.shape, dtype=x_val.dtype), mode + ) + output_shapes = [t.shape for t in split_target(fake_input, *split_args[1:])] + split_node.meta["val"] = tuple( + [ + FakeTensor.from_tensor(torch.empty(shape, dtype=x_val.dtype), mode) + for shape in output_shapes + ] + ) + + return split_node + + def _create_convolution_node(self, conv_target, args: tuple) -> Node: + convolution_node = self.module.graph.call_function(conv_target, args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + convolution_node.meta["source_fn_stack"] = [ + (convolution_node.name, torch.convolution) + ] + + # Compute the output shapes for the `convolution`, and assign the `val` meta. + with FakeTensorMode() as mode: + input_shapes = [ + input_.meta["val"].shape if hasattr(input_, "meta") else input_.shape + for input_ in args[:3] + ] + input_dtypes = [ + input_.meta["val"].dtype if hasattr(input_, "meta") else input_.dtype + for input_ in args[:3] + ] + fake_inputs = [ + FakeTensor.from_tensor(torch.empty(shape, dtype=dtype), mode) + for shape, dtype in zip(input_shapes, input_dtypes) + ] + output = conv_target(*fake_inputs, *args[3:]) + convolution_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(output.shape, dtype=output.dtype), mode + ) + + return convolution_node + + def _create_concat_node(self, *cat_args) -> Node: + cat_target = torch.ops.aten.cat.default + concat_node = self.module.graph.call_function(cat_target, cat_args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + concat_node.meta["source_fn_stack"] = [(concat_node.name, torch.cat)] + + # Compute the output shape for the `concat`, and assign the `val` meta. + with FakeTensorMode() as mode: + fake_inputs = [ + FakeTensor.from_tensor( + torch.empty( + input_.meta["val"].shape, dtype=input_.meta["val"].dtype + ), + mode, + ) + for input_ in cat_args[0] + ] + output = cat_target(fake_inputs, *cat_args[1:]) + concat_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(output.shape, dtype=output.dtype), mode + ) + + return concat_node + + def _get_topologically_last_node(self, nodes: list[Node]) -> Node: + """Return the node from `nodes` which appears last in the graph.""" + for node in reversed(self.module.graph.nodes): + if node in nodes: + return node + + raise RuntimeError(f"None of the nodes `{nodes}` are in the graph.") + + def _create_parameter_node_for_data( + self, data: torch.Tensor, name: str, insert_after_node: torch.Node + ) -> torch.Node: + """Create a parameter node in the graph, which contains the provided `data`.""" + new_name = get_new_attr_name_with_prefix(name)(self.module) + + # Create the node for the parameter. + param = torch.nn.Parameter(data, False) + _assign_attr(param, self.module, str(new_name), _AttrKind.PARAMETER) + with self.module.graph.inserting_after(insert_after_node): + static_parameter_node = self.module.graph.get_attr(new_name) + + with FakeTensorMode() as mode: + static_parameter_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(data.shape, dtype=data.dtype), mode + ) + + return static_parameter_node + + def call(self, module: GraphModule): + self.module = module + + def _is_conv(node_: Node): + return node_.op == "call_function" and node_.target in ( + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + ) + + made_changes = False + + for node in self.module.graph.nodes: + if not _is_conv(conv_node := node): + continue + + if len(conv_node.args) < 7: + # The `aten.conv` can have fewer args if the others use default values. + # So in this case, `groups == 1`. + continue + x, w, b, stride, padding, dilation, groups = conv_node.args + + if not group_conv_convertible_into_multiple_convolutions(conv_node, groups): + continue + + if len(x.meta["val"].shape) not in [3, 4]: + # Only 1D and 2D convolutions are supported by the Neutron backend. Don't decompose anything else. + continue + + w_data = self._get_tensor_constant_from_node(w) + b_data = self._get_tensor_constant_from_node(b) + if w_data is None or b_data is None: + continue # Only the standard case with static weights and bias is supported. + + # Create a `split` node to split the main input. + # Split across dimension `1` (channels), `groups` slices of size `input_split_size`. + num_input_channels = x.meta["val"].shape[1] + input_split_sizes = [num_input_channels // groups] * groups + with self.module.graph.inserting_before(conv_node): + split_node = self._create_split_node(x, input_split_sizes, 1) + + # Add `GetItem` nodes to extract the outputs of the `split_node`. + split_getitem_nodes = [ + self._create_and_insert_get_item_node(split_node, i) + for i in range(groups) + ] + + # Split the weights and bias, across dimension `0`, slices of size `weight_split_size`. + weight_split_size = w.meta["val"].shape[0] // groups + split_weights_data = torch.split(w_data, weight_split_size, 0) + split_bias_data = torch.split(b_data, weight_split_size, 0) + + # Turn the weights and biases into parameter nodes containing the data. + # Use a different name for every parameter. The function internally ensures the name's uniqueness, but + # relying on it sometimes causes strange failures when `groups > 5` for some weird reason. + split_weight_nodes = [ + self._create_parameter_node_for_data( + weight_data, w.name + f"_{i}_", split_node + ) + for i, weight_data in enumerate(split_weights_data) + ] + split_bias_nodes = [ + self._create_parameter_node_for_data( + bias_data, b.name + f"_{i}_", split_node + ) + for i, bias_data in enumerate(split_bias_data) + ] + + # Create the `conv` nodes. + with self.module.graph.inserting_after( + self._get_topologically_last_node( + split_getitem_nodes + split_weight_nodes + split_bias_nodes + ) + ): + split_conv_nodes = [ + self._create_convolution_node( + conv_node.target, # Use the same target as the original convolution (1d/2d/3d/...). + (input_getitem, weight, bias, stride, padding, dilation, 1), + ) + for input_getitem, weight, bias in zip( + split_getitem_nodes, split_weight_nodes, split_bias_nodes + ) + ] + + # Create the `cat` node. + with self.module.graph.inserting_after( + self._get_topologically_last_node(split_conv_nodes) + ): + concat_node = self._create_concat_node( + split_conv_nodes, 1 + ) # Concatenate along the channels. + + # Replace the uses of the original convolution with the `concat_node`. + conv_node.replace_all_uses_with(concat_node) + self.module.graph.erase_node(conv_node) + + made_changes = True + + return PassResult(self.module, made_changes) diff --git a/backends/nxp/aten_passes/split_gru_based_on_num_layers.py b/backends/nxp/aten_passes/split_gru_based_on_num_layers.py new file mode 100644 index 00000000000..0c20f71d4e3 --- /dev/null +++ b/backends/nxp/aten_passes/split_gru_based_on_num_layers.py @@ -0,0 +1,307 @@ +# Copyright 2025 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 operator + +import torch +from torch._subclasses import FakeTensor, FakeTensorMode +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase, PassResult + + +class SplitGRUBasedOnNumLayers(PassBase): + """Replace an `aten.gru.input` operator with `num_layers > 1` with a subgraph consisting of multiple chained + `aten.gru.input` operators with `num_layers == 1`, according to the following schematic. + + + X H_h + │ │ + │ ┌───▼───┐ + │ │ Split │ + │ └─┬───┬─┘ + │ ┌───────┘ └───────┐ + │ ┌─────▼──────┐ ┌──────▼──────────────────┐ + │ │ GetItem[0] │ ... │ GetItem[-1] │ + │ └─────┬──────┘ └──────────────┬──────────┘ + X X_h └────┐ │ │ + │ │ ┌─────▼──▼─────┐ │ + ┌──────────▼───▼──────────┐ W11─► GRU ◄─B11 │ + W11─► ◄─B11 W12─► num_layers=1 ◄─B12 │ + W12─► GRU ◄─B12 └───────┬──────┘ │ + ... │ num_layers= │ ... ┌───────┴───────┐ │ + W2─► ◄─B2 ┌─────▼──────┐ ┌──────▼─────┐ │ + └────────────┬────────────┘ │ GetItem[0] │ │ GetItem[1] │ │ + ┌────────┴────────┐ └─────┬──────┘ └──────┬─────┘ │ + ┌─────▼──────┐ ┌──────▼─────┐ replace with │ ┌──────┘ │ + │ GetItem[0] │ │ GetItem[1] │ ─────────► └────────┼─────── ... ────────────┐ │ + └─────┬──────┘ └──────┬─────┘ │ ┌─────▼──▼─────┐ + ▼ ▼ │ W1─► GRU ◄─B1 + Y Y_h │ W2─► num_layers=1 ◄─B2 + │ └──────┬───────┘ + │ ┌───────┴───────┐ + │ ┌─────▼──────┐ ┌──────▼─────┐ + │ │ GetItem[0] │ │ GetItem[1] │ + │ └─────┬──────┘ └──────┬─────┘ + │ ... ┌──────────┼───────────────┘ + ┌▼──────▼┐ │ + │ Concat │ │ + └───┬────┘ │ + ▼ ▼ + Y_h Y + + The `aten.gru.input` has the following schema: + aten::gru.input( + Tensor input, Tensor hx, Tensor[] params, + bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first + ) -> (Tensor, Tensor) + """ + + module: GraphModule + + def _get_topologically_last_node(self, nodes: list[Node]) -> Node: + """Return the node from `nodes` which appears last in the graph.""" + for node in reversed(self.module.graph.nodes): + if node in nodes: + return node + + raise RuntimeError(f"None of the nodes `{nodes}` are in the graph.") + + def _create_and_insert_get_item_node(self, input_node: Node, idx: int) -> Node: + """Create a `GetItem` node which extracts the output of `input_node` on index `idx`. + The `GetItem` is also added to the graph right after the `input_node`. + """ + with self.module.graph.inserting_after(input_node): + get_item_node = self.module.graph.create_node( + "call_function", + operator.getitem, + (input_node, idx), + {}, + ) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + get_item_node.meta["source_fn_stack"] = [ + (get_item_node.name, input_node.meta["source_fn_stack"]) + ] + get_item_node.meta["val"] = input_node.meta["val"][idx] + + return get_item_node + + def _create_gru_node(self, *gru_args) -> Node: + """Create an `aten.gru.input` node with the provided arguments. The node will NOT be added to the graph + automatically. + + :param gru_args: Arguments for the `aten.gru.input` operation. + :return: The created GRU Node. + """ + gru_target = torch.ops.aten.gru.input + gru_node = self.module.graph.call_function(gru_target, gru_args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + gru_node.meta["source_fn_stack"] = [(gru_node.name, torch.nn.modules.rnn.GRU)] + + # Compute the shapes of the GRU outputs, and assign the `val` meta. + x_val, h_val = gru_args[0].meta["val"], gru_args[1].meta["val"] + with FakeTensorMode() as mode: + fake_x = FakeTensor.from_tensor( + torch.empty(x_val.shape, dtype=x_val.dtype), mode + ) + fake_h = FakeTensor.from_tensor( + torch.empty(h_val.shape, dtype=h_val.dtype), mode + ) + fake_weights = [ + FakeTensor.from_tensor( + torch.empty(w.meta["val"].shape, dtype=x_val.dtype), mode + ) + for w in gru_args[2] + ] + output_shapes = [ + t.shape for t in gru_target(fake_x, fake_h, fake_weights, *gru_args[3:]) + ] + gru_node.meta["val"] = tuple( + [ + FakeTensor.from_tensor(torch.empty(shape, dtype=h_val.dtype), mode) + for shape in output_shapes + ] + ) + + return gru_node + + def _create_split_node(self, *split_args) -> Node: + """Create an `aten.split.default` node with the provided arguments. The node will NOT be added to the graph + automatically. + + :param split_args: Arguments for the `aten.split.default` operation. + :return: The created Split Node. + """ + split_target = torch.ops.aten.split.default + split_node = self.module.graph.call_function(split_target, split_args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + split_node.meta["source_fn_stack"] = [(split_node.name, torch.split)] + + # Compute the output shapes for the `split`, and assign the `val` meta. + x_val = split_args[0].meta["val"] + with FakeTensorMode() as mode: + fake_input = FakeTensor.from_tensor( + torch.empty(x_val.shape, dtype=x_val.dtype), mode + ) + output_shapes = [t.shape for t in split_target(fake_input, *split_args[1:])] + split_node.meta["val"] = tuple( + [ + FakeTensor.from_tensor(torch.empty(shape, dtype=x_val.dtype), mode) + for shape in output_shapes + ] + ) + + return split_node + + def create_concat_node(self, *cat_args) -> Node: + """Create an `aten.cat.default` node with the provided arguments. The node will NOT be added to the graph + automatically. + + :param cat_args: Arguments for the `aten.cat.default` operation. + :return: The created Cat Node. + """ + cat_target = torch.ops.aten.cat.default + cat_node = self.module.graph.call_function(cat_target, cat_args) + + # Assign the `source_fn_stack` and `val` meta fields as they are required for quantization. + cat_node.meta["source_fn_stack"] = [(cat_node.name, torch.cat)] + + # Compute the output shape for the `concat`, and assign the `val` meta. + with FakeTensorMode() as mode: + fake_inputs = [ + FakeTensor.from_tensor( + torch.empty( + input_.meta["val"].shape, dtype=input_.meta["val"].dtype + ), + mode, + ) + for input_ in cat_args[0] + ] + output = cat_target(fake_inputs, *cat_args[1:]) + cat_node.meta["val"] = FakeTensor.from_tensor( + torch.empty(output.shape, dtype=output.dtype), mode + ) + + return cat_node + + def call(self, module: GraphModule) -> PassResult: + self.module = module + made_changes = False + + def _is_gru(node_: Node) -> bool: + return ( + node_.op == "call_function" and node_.target == torch.ops.aten.gru.input + ) + + if not any(map(_is_gru, module.graph.nodes)): + return PassResult(module, False) # No GRU nodes in the model. + + for node in module.graph.nodes: + if not _is_gru(node): + continue # Not GRU. + + original_gru_node = node + if (num_layers := original_gru_node.args[4]) == 1: + # Basic 1-layer GRU. + continue + + if (dropout := node.args[5]) != 0.0 or (train := node.args[6]): + # Conversion for these cases is not supported, so the pre-processing should not be applied. + continue + + # The `hx` (initial hidden state) has shape: + # [D * num_layers, hidden_size] or [D * num_layers, batch_size, hidden_size] + # where D = 2 if bidirectional else 1. + # Split the `hx` into `num_layers` different tensors. + h_x = original_gru_node.args[1] + bidirectional = original_gru_node.args[7] + d = 2 if bidirectional else 1 + with module.graph.inserting_before(original_gru_node): + # Split across the dimension `0`. Slices of size `d`. + num_slices = h_x.meta["val"].shape[0] // d + split_node = self._create_split_node(h_x, [d] * num_slices, 0) + + # Add `GetItem` nodes to extract the outputs of the `split_node`. + h_0_get_item_nodes = [ + self._create_and_insert_get_item_node(split_node, i) + for i in range(num_layers) + ] + + # ---- Create new GRU nodes ---- + + all_weights = original_gru_node.args[2] + has_biases = original_gru_node.args[3] + # The `all_weights` list contains + # [w11, w12, b11, b12, w21, w22, b21, b22, ...] if `has_biases` else [w11, w12, w21, w22, ...]. + step = 4 if has_biases else 2 + if bidirectional: + # Every other set of weights and biases (2 or 4) represents the reverse connections for the layer. + step *= 2 + + gru_nodes = [] + batch_first = original_gru_node.args[-1] + + # The `GetItem` node which extracts the main output (y) of the previous GRU. (Or the main input for the + # first GRU). + prev_gru_main_output_get_item = original_gru_node.args[0] + output_h_get_item_nodes = ( + [] + ) # `GetItem` nodes which extract the output hidden states of the GRU nodes. + for i in range(num_layers): + current_gru_weights = tuple(all_weights[step * i : step * (i + 1)]) + + # Select the node to insert the new `GRU` after. + prev_node = ( + self._get_topologically_last_node(h_0_get_item_nodes) + if i == 0 + else prev_gru_main_output_get_item + ) + + # Create the new `GRU`. + with module.graph.inserting_after(prev_node): + gru = self._create_gru_node( + prev_gru_main_output_get_item, + h_0_get_item_nodes[i], + current_gru_weights, + has_biases, + 1, # New `num_layers`. + dropout, + train, + bidirectional, + batch_first, + ) + gru_nodes.append(gru) + + # Create the `GetItem` nodes to extract the outputs of the `GRU`. + prev_gru_main_output_get_item = self._create_and_insert_get_item_node( + gru_nodes[i], 0 + ) + output_h_get_item_nodes.append( + self._create_and_insert_get_item_node(gru_nodes[i], 1) + ) + + # Add a `Concat` to collect all the output hidden states. + with module.graph.inserting_after(prev_gru_main_output_get_item): + concat_node = self.create_concat_node( + output_h_get_item_nodes, 0 # Concatenate along the dimension `0`. + ) + + # Replace the uses of the original `GRU` outputs with the new corresponding outputs. + original_y_get_item, original_yh_get_item = list( + original_gru_node.users.keys() + ) + original_y_get_item.replace_all_uses_with(prev_gru_main_output_get_item) + original_yh_get_item.replace_all_uses_with(concat_node) + + # Remove the old nodes. + module.graph.erase_node(original_y_get_item) + module.graph.erase_node(original_yh_get_item) + module.graph.erase_node(original_gru_node) + + made_changes = True + + return PassResult(module, made_changes) diff --git a/backends/nxp/backend/custom_delegation_options.py b/backends/nxp/backend/custom_delegation_options.py new file mode 100644 index 00000000000..a5552c6ec89 --- /dev/null +++ b/backends/nxp/backend/custom_delegation_options.py @@ -0,0 +1,19 @@ +# Copyright 2025 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 dataclasses import dataclass + + +@dataclass +class CustomDelegationOptions: + """The class allows the user to specify details which affect which nodes will be delegated.""" + + # Neutron requires the channel dimension to be multiple of `num_macs` for concatenation (cat op). + # Due to different dim ordering in torch (channel_first) and Neutron IR (channel last), dim of the channel is + # ambiguous. Cat converter will defensively require both possible dimension index for the channels to be multiple + # of `num_macs`. The `force_delegate_cat` allows the user to turn off the defensive check if from the model design + # it is known this constraint will be satisfied. + force_delegate_cat: bool = False diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index 3b74d86f599..061295ead79 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree. import torch -from torch.fx import Node +from torch.fx import GraphModule, Node from torch.nn import Parameter @@ -71,3 +71,19 @@ def _is_dequantize(node_: Node) -> bool: return _is_dequantize(node) and node_is_static_tensor( node.args[0], parameters_mapping ) + + +def try_get_tensor_constant_from_node( + graph_module: GraphModule, node: Node +) -> Parameter | None: + """Get the static data from a given node. If it doesn't have any data, return `None`.""" + if node is None or node.op != "get_attr": + return None + + target_atoms = node.target.split(".") + attr_itr = graph_module + for atom in target_atoms: + if not hasattr(attr_itr, atom): + return None + attr_itr = getattr(attr_itr, atom) + return attr_itr diff --git a/backends/nxp/backend/edge_program_converter.py b/backends/nxp/backend/edge_program_converter.py index 1e930d37a6a..ddbbf5b2e3a 100644 --- a/backends/nxp/backend/edge_program_converter.py +++ b/backends/nxp/backend/edge_program_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -10,6 +10,9 @@ from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import ( AtenModelBuilderDirector, ) +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, +) from torch.export import ExportedProgram from torch.export.graph_signature import InputKind from torch.fx import Node @@ -28,6 +31,7 @@ exir_ops.edge.aten.addmm.default: AddMMConverter, # noqa F405 exir_ops.edge.aten.add.Tensor: AddTensorConverter, # noqa F405 exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter, # noqa F405 + exir_ops.edge.aten.cat.default: CatConverter, # noqa F405 exir_ops.edge.aten.clone.default: CloneConverter, # noqa F405 exir_ops.edge.aten.constant_pad_nd.default: ConstantPadNDConverter, # noqa F405 exir_ops.edge.aten.convolution.default: ConvolutionConverter, # noqa F405 @@ -38,6 +42,7 @@ exir_ops.edge.aten.permute_copy.default: PermuteCopyConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 + exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 exir_ops.edge.aten.view_copy.default: ViewCopyConverter, # noqa F405 exir_ops.edge.aten.sigmoid.default: SigmoidConverter, # noqa F405 } @@ -49,24 +54,30 @@ class EdgeProgramToIRConverter: """ _default_conversion_config = ConversionConfig() + _default_delegation_options = CustomDelegationOptions() def convert_program( self, edge_program: ExportedProgram, conversion_config=_default_conversion_config, + custom_delegation_options: CustomDelegationOptions = _default_delegation_options, ) -> (bytes, dict): """ Convert ExportedProgram in Edge dialect to IR (TFLite flatbuffers) as bytes. :param edge_program: Converter ExportedProgram. :param conversion_config: ConversionConfig instance. + :param custom_delegation_options: Custom user options which affect node delegation. :return: TFLite flatbuffers as bytes. """ node_formats = NodeFormatInference(edge_program).identify_node_formats() parameters_mapping = self.map_inputs_to_parameters(edge_program) cc = self.build_conversion_context( - parameters_mapping, node_formats, conversion_config + parameters_mapping, + node_formats, + conversion_config, + custom_delegation_options, ) # Program conversion @@ -162,6 +173,7 @@ def build_conversion_context( parameters_mapping: dict, node_formats: dict[Node, NodeFormat], conversion_config: ConversionConfig = _default_conversion_config, + custom_delegation_options: CustomDelegationOptions = _default_delegation_options, ) -> ConversionContext: tflite_builder = AtenModelBuilderDirector( 3, "TFLite from EdgeProgram", conversion_config @@ -171,7 +183,11 @@ def build_conversion_context( tflite_builder.build_empty_buffer() context = ConversionContext( - tflite_builder, conversion_config, parameters_mapping, node_formats + tflite_builder, + conversion_config, + parameters_mapping, + node_formats, + custom_delegation_options, ) return context diff --git a/backends/nxp/backend/ir/conversion_context.py b/backends/nxp/backend/ir/conversion_context.py index 6ec80f02a66..6fb7e98424e 100644 --- a/backends/nxp/backend/ir/conversion_context.py +++ b/backends/nxp/backend/ir/conversion_context.py @@ -1,8 +1,11 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 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.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import ( AtenModelBuilderDirector, @@ -17,6 +20,7 @@ class ConversionContext: conversion_config: ConversionConfig parameters_mapping: dict[str, Parameter] node_formats: dict[Node, NodeFormat] + custom_delegation_options: CustomDelegationOptions def __init__( self, @@ -24,6 +28,7 @@ def __init__( conversion_config: ConversionConfig, parameters_mapping: dict, node_formats: dict[Node, NodeFormat], + custom_delegation_options: CustomDelegationOptions, ): """ Context with data related to current conversion. @@ -35,3 +40,4 @@ def __init__( self.conversion_config = conversion_config self.parameters_mapping = parameters_mapping self.node_formats = node_formats + self.custom_delegation_options = custom_delegation_options diff --git a/backends/nxp/backend/ir/converter/conversion/common.py b/backends/nxp/backend/ir/converter/conversion/common.py index 0f69b152ec7..8230e39a7fa 100755 --- a/backends/nxp/backend/ir/converter/conversion/common.py +++ b/backends/nxp/backend/ir/converter/conversion/common.py @@ -70,29 +70,22 @@ def try_get_input(t_op: tflite_model.Operator, idx: int) -> tflite_model.Tensor return tensor -def extend_1d_pads_to_2d(onnx_1d_pads: MutableSequence): - """Extend the onnx 'pads' operator attribute that represents padding for a 1D kernel to 2D, by adding '0's.""" - if onnx_1d_pads is not None: - onnx_1d_pads.insert(1, 0) - onnx_1d_pads.append(0) +def extend_1d_padding_to_2d(tflite_1d_padding: MutableSequence): + """Extend the PyTorch 'padding' operator attribute that represents padding for a 1D kernel to 2D, by adding '0's.""" + if tflite_1d_padding is not None: + tflite_1d_padding.append(0) -def extend_1d_strides_to_2d(onnx_1d_strides: MutableSequence): - """Extend the onnx 'strides' operator attribute that represents strides for a 1D kernel to 2D, by adding '1'.""" - if onnx_1d_strides is not None: - onnx_1d_strides.append(1) +def extend_1d_stride_to_2d(tflite_1d_stride: MutableSequence): + """Extend the PyTorch 'stride' operator attribute that represents stride for a 1D kernel to 2D, by adding '1'.""" + if tflite_1d_stride is not None: + tflite_1d_stride.append(1) -def extend_1d_dilations_to_2d(onnx_1d_dilations: MutableSequence): - """Extend the onnx 'dilations' operator attribute that represents dilations for a 1D kernel to 2D, by adding '1'.""" - if onnx_1d_dilations is not None: - onnx_1d_dilations.append(1) - - -def extend_1d_kernel_shape_to_2d(onnx_1d_kernel_shape: MutableSequence): - """Extend the onnx 1D 'kernel_shape' operator attribute to 2D, by adding '1'.""" - if onnx_1d_kernel_shape is not None: - onnx_1d_kernel_shape.append(1) +def extend_1d_dilation_to_2d(tflite_1d_dilation: MutableSequence): + """Extend the PyTorch 'dilation' operator attribute that represents dilation for a 1D kernel to 2D, by adding '1'.""" + if tflite_1d_dilation is not None: + tflite_1d_dilation.append(1) StridedOptions = ( diff --git a/backends/nxp/backend/ir/converter/node_converter.py b/backends/nxp/backend/ir/converter/node_converter.py index a9f94d543f2..ed624aaa411 100755 --- a/backends/nxp/backend/ir/converter/node_converter.py +++ b/backends/nxp/backend/ir/converter/node_converter.py @@ -1,14 +1,16 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 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 abc import ABC, abstractmethod from enum import Enum -from typing import Collection import torch +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) from executorch.backends.nxp.backend.ir.conversion_context import ConversionContext from executorch.backends.nxp.backend.ir.converter.builder.aten_model_builder_director import ( AtenModelBuilderDirector, @@ -53,7 +55,6 @@ class NodeConverter(ABC): """ context: ConversionContext - supported_targets: Collection def __init__(self, context: ConversionContext): self.context = context @@ -72,45 +73,57 @@ def convert(self, node: Node): @staticmethod @abstractmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: """Check if the `node` can be converted to the intermediate representation. Classes which implement conversion for individual operators must overwrite this method. :param node: torch.Node to check. + :param parameters_mapping: Dictionary mapping tensor names to their static data (if they have it). + :param custom_delegation_options: Custom options which affect delegation. """ pass - @classmethod - def _is_supported_on_target(cls, target: Target) -> bool: - """Check if the node is supported on the target platform. It uses the 'supported_platform' attribute, which is - a list of supported target platforms, and it must be defined by the specific `NodeConverter`. + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + """Check if the node is supported on the target platform. + Child classes should overwrite this method to implement specific target checks. The default implementation + can be used by operators with no target specific requirements. + :param node: The node (edge operator) to check. :param target: Value of the `Target` enum representing the target platform to check for. + :param parameters_mapping: Dictionary mapping tensor names to their static data (if they have it). + :param custom_delegation_options: Custom options which affect delegation. """ - if not ( - hasattr(cls, "supported_targets") - and isinstance(cls.supported_targets, Collection) - ): - raise NotImplementedError( - f"The NodeConverter `{cls}` does not define its `supported_targets` collection." - ) - - return target == Target.IGNORE or target in cls.supported_targets + return target == Target.RT700 @classmethod def is_supported( - cls, node: Node, target: Target, parameters_mapping: dict[str, Parameter] + cls, + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: """Check if the given `node` is supported in the IR and on the given `target` platform. :param node: torch.Node to check. :param target: Value of the `Target` enum representing the target platform to check for. :param parameters_mapping: Dict mapping tensor names to their data. + :param custom_delegation_options: Custom user options which affect node delegation. """ return cls._is_supported_in_IR( - node, parameters_mapping - ) and cls._is_supported_on_target(target) + node, parameters_mapping, custom_delegation_options + ) and cls._is_supported_on_target( + node, target, parameters_mapping, custom_delegation_options + ) @staticmethod def _has_shared_q_params_if_quantized(node: Node) -> bool: @@ -119,13 +132,8 @@ def _has_shared_q_params_if_quantized(node: Node) -> bool: # Some exotic operator (only consumer or only produces) return True - pre_node = node.prev - post_node = node.next - - if pre_node.name == node.all_input_nodes[0] and post_node.name == node.users[0]: - raise RuntimeError( - "Prev & next nodes are not the same as inputs and outputs." - ) + pre_node = node.all_input_nodes[0] + post_node = list(node.users)[0] if _is_dequant_node(pre_node) and _is_quant_node(post_node): # Node is quantized @@ -149,7 +157,11 @@ def assert_convertible(self, node): """Assert that the call `_is_supported_in_IR()` returns `True`. Otherwise, raise an exception and print an error message. """ - assert self._is_supported_in_IR(node, self.context.parameters_mapping), ( + assert self._is_supported_in_IR( + node, + self.context.parameters_mapping, + self.context.custom_delegation_options, + ), ( f"Node `{node}` is not convertible to the intermediate representation. " "There is an error in the partitioner." ) @@ -173,7 +185,15 @@ def _create_tflite_op_with_io_tensors(self, node: Node) -> tflite_model.Operator # Initialize node's inputs t_operator.inputs = tflite_model.OperatorInputs() - input_nodes = [arg for arg in node.args if isinstance(arg, Node)] + + input_nodes = [] + for arg in node.args: + match arg: + case Node(): + input_nodes.append(arg) + case list() if all(isinstance(node_, Node) for node_ in arg): + input_nodes.extend(arg) + for ancestor_node in input_nodes: assert self.context.tflite_builder.tensor_exists(ancestor_node.name) t_operator.tmp_inputs.append( diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py index 8a0498810ce..d1674e16a9f 100755 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/__init__.py @@ -13,6 +13,9 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.avg_pool_2d_converter import ( AvgPool2dConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.cat_converter import ( + CatConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.clone_converter import ( CloneConverter, ) @@ -52,12 +55,16 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.softmax_converter import ( SoftmaxConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.tanh_converter import ( + TanhConverter, +) from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.view_copy_converter import ( ViewCopyConverter, ) __all__ = [ "AddMMConverter", + "CatConverter", "ConvolutionConverter", "MMConverter", "PermuteCopyConverter", @@ -76,4 +83,5 @@ "AdaptiveAvgPool2dConverter", "HardTanhConverter", "SigmoidConverter", + "TanhConverter", ] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py index 2dbb903c8f9..f2b26d6512e 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/abs_converter.py @@ -1,13 +1,12 @@ -# Copyright (c) 2025 NXP -# All rights reserved. +# Copyright 2025 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.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( abs_options, @@ -17,11 +16,12 @@ class AbsConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: return True 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 f0eab0ccbf6..4b9ff6fe85a 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 @@ -1,5 +1,4 @@ -# Copyright (c) 2025 NXP -# All rights reserved. +# Copyright 2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -7,8 +6,8 @@ import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding from executorch.backends.nxp.backend.ir.converter.conversion import common from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( @@ -20,11 +19,12 @@ class AdaptiveAvgPool2dConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: 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/add_tensor_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/add_tensor_converter.py index c4ce2e44bd0..c74baa61f67 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/add_tensor_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/add_tensor_converter.py @@ -1,5 +1,4 @@ -# Copyright (c) 2025 NXP -# All rights reserved. +# Copyright 2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -8,6 +7,7 @@ node_uses_shape_broadcasting, ) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, Target, ) @@ -19,11 +19,29 @@ class AddTensorConverter(NodeConverter): - supported_targets = [Target.RT700] + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + match target: + case Target.RT700: + if node_uses_shape_broadcasting(node): + # Shape broadcasting may require the addition of `Transpose` ops during conversion. + return False + + return True + + case _: + return False @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: if len(node.args) != 2: return False @@ -31,10 +49,6 @@ def _is_supported_in_IR( if hasattr(node.kwargs, "alpha"): return False - # Don't convert if broadcasting input tensors - if node_uses_shape_broadcasting(node): - return False - return True # add.Tensor Node format: (Tensor self, Tensor other, *, Scalar alpha=1) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/addmm_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/addmm_converter.py index 820d1414f3b..0df41526da2 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/addmm_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/addmm_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -6,8 +6,8 @@ from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( fully_connected_options, @@ -21,7 +21,9 @@ class AddMMConverter(NodeConverter): @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: if len(node.all_input_nodes) != 3: return False @@ -32,8 +34,6 @@ def _is_supported_in_IR( return True - supported_targets = [Target.RT700] - def convert(self, node: Node): self.assert_convertible(node) 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 41150f52d98..99ae0a30dbb 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 @@ -1,17 +1,21 @@ -# Copyright (c) 2025 NXP -# All rights reserved. +# Copyright 2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import numpy as np + from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, ) from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList +from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( + tf_lite_type_to_numpy, +) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( @@ -22,11 +26,12 @@ class AvgPool2dConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: n_args = len(node.args) @@ -62,9 +67,20 @@ def _convert_2d_avg_pool( ) if explicit_padding is not None: - # Need to prepend a 'Pad' operator, which adds 0s. But these will be included in the computation! + # Need to prepend a 'Pad' operator, which adds 0s (or `zero_point` for the quantized case). But these will + # be included in the computation! + input_quantization = t_op.tmp_inputs[0].quantization + pad_value = ( + None + if input_quantization is None + else np.array(input_quantization.zero_point[0]).astype( + tf_lite_type_to_numpy(t_op.tmp_inputs[0].type) + ) + ) ops.add_pre( - self.builder.create_pad_operator_before(t_op, 0, explicit_padding) + self.builder.create_pad_operator_before( + t_op, 0, explicit_padding, pad_value + ) ) return ops.flatten() diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/cat_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/cat_converter.py new file mode 100644 index 00000000000..4f7f00fe5ba --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/cat_converter.py @@ -0,0 +1,148 @@ +# Copyright 2025 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.ir.converter.conversion import translator +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + _is_dequant_node, + _is_quant_node, + NodeConverter, + Target, +) +from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.concatenation_options import ( + Concatenation, +) +from torch.fx import Node +from torch.nn import Parameter + + +def _get_shape(node: torch.fx.Node) -> list[int]: + return node.meta["val"].shape + + +class CatConverter(NodeConverter): + + @staticmethod + def _get_normalized_dim(node: torch.fx.Node) -> int: + dim = node.args[1] if len(node.args) >= 2 else 0 # Default `dim` value. + rank = len(_get_shape(node)) + if dim < 0: + dim += rank + + if not (0 <= dim < rank): + raise RuntimeError("`Cat` operator has invalid `dim`.") + + return dim + + @staticmethod + def _all_io_shares_quantization_parameters(node: Node) -> bool: + post_node = list(node.users.keys())[0] + if not _is_quant_node(post_node): + return False + output_zp, output_scale, output_type = ( + post_node.args[1], + post_node.args[2], + post_node.args[5], + ) + + for input_node in node.args[0]: + if not _is_dequant_node(input_node): + return False + + input_zp, input_scale, input_type = ( + input_node.args[1], + input_node.args[2], + input_node.args[5], + ) + if (input_zp, input_scale, input_type) != ( + output_zp, + output_scale, + output_type, + ): + return False + + return True + + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if custom_delegation_options.force_delegate_cat: + return True + + match target: + case Target.RT700: + dim = CatConverter._get_normalized_dim(node) + + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1491 + if dim == 0: + return False + + # Neutron requires the channels to be a multiple of `8`. The channels could either be the second or the + # last dimension, depending on the formats of the node. The format, however, cannot be determined + # during conversion, as it depends on what other nodes are delegated. + input_channels = [ + # The second dimension is the channels in PyTorch. If the inputs/output are not channels first, it + # will still be the channels in the IR. + _get_shape(input_)[1] + for input_ in node.all_input_nodes + ] + [ + # If the inputs/outputs are channels first, the last dimension will be the channels. + _get_shape(input_)[-1] + for input_ in node.all_input_nodes + ] + if any((input_channel % 8) != 0 for input_channel in input_channels): + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1492 + return False + + output_channels = [_get_shape(node)[1], _get_shape(node)[-1]] + # neutron-library/src/utils/NeutronLibraryInterrogation.cpp#1493 + if any((out_c % 8) != 0 for out_c in output_channels): + return False + + if len(node.all_input_nodes) < 2: # Not supported on Neutron + # TODO Try to skip the operator if this case is realistic. + return False + + return True + + case _: + return False + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + if not CatConverter._all_io_shares_quantization_parameters(node): + # The IR requires all inputs to have the same quantization parameters as the output. + # The quantizer should quantize the operator so that this case does not happen. + return False + + return True + + def convert(self, node: Node): + """Convert the 'aten.cat' operator to TFLite 'Concatenation'.""" + self.assert_convertible(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + + dim = self._get_normalized_dim(node) # Also checks the validity of `dim`. + + if t_op.tmp_inputs[0].tensor_format.is_channels_last(): + dim = translator.create_channels_last_to_channels_first_permutation( + t_op.tmp_inputs[0].rank + )[dim] + + t_op.builtin_options = Concatenation(dim) + self.builder.append_operators([t_op]) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/clone_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/clone_converter.py index 5b51fc72dc1..1d370ab8c48 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/clone_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/clone_converter.py @@ -5,8 +5,8 @@ import torch from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from torch.fx import Node from torch.nn import Parameter @@ -20,11 +20,12 @@ def _has_supported_memory_format(node: Node) -> bool: class CloneConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: return _has_supported_memory_format(node) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py index 7b749818f5e..f58df1a88d9 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/constant_pad_nd_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -15,6 +15,7 @@ tf_lite_type_to_numpy, ) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, Target, ) @@ -31,11 +32,31 @@ class ConstantPadNDConverter(NodeConverter): - supported_targets = [Target.RT700] + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + match target: + case Target.RT700: + # TODO: Consider different tensor formats (dim-order) + paddings = node.args[1] + if len(paddings) > 4 and paddings[4:6] != [0, 0]: + # Attempt to Pad channels dimension, which is not supported on Neutron. + return False + + return True + + case _: + return False @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: paddings = node.args[1] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py index 6aac32649d3..0f3a4b9bb5a 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py @@ -1,8 +1,10 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 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 numpy as np import torch @@ -14,9 +16,14 @@ from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, + translator, ) from executorch.backends.nxp.backend.ir.converter.conversion.common import try_get_input +from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( + tf_lite_type_to_numpy, +) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, Target, ) @@ -36,51 +43,71 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( conv_2d_options, depthwise_conv_2d_options, + reshape_options, ) from torch.fx import Node from torch.nn import Parameter class ConvolutionConverter(NodeConverter): - supported_targets = [Target.RT700] + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + match target: + case Target.RT700: + activations = node.args[0] + weights = node.args[1] + groups = node.args[8] + + if activations.meta["val"].shape[0] != 1: + # Only batch size 1 is supported on neutron. + return False + + if groups == 1: # Regular convolution. + pass + elif conv_utils.group_conv_convertible_as_depthwise( + node, groups + ): # Depthwise convolution. + # Only supported if the weights are static, because TFLite `DepthwiseConv2D` uses permuted + # weights. In case the weights are dynamic, a Transpose operator would have to be added, which + # is not supported on Neutron. + if not node_is_effectively_static_tensor( + weights, parameters_mapping + ): + return False + elif conv_utils.group_conv_convertible_into_multiple_convolutions( + node, groups + ): # Separable conv. This should never be reached, as the node should have been decomposed into + # multiple parallel convolutions by the `SplitGroupConvolution` pre-processing pass. + logging.warning("Group convolution was not decomposed.") + return False + else: # Unexpected case (should never happen). + return False + + return True + + case _: + return False @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: + input_tensor_rank = len(node.meta["val"].shape) + dimensions = input_tensor_rank - 2 is_transposed = node.args[6] output_padding = node.args[7] - groups = node.args[8] if is_transposed: return False - if output_padding != [0, 0]: - return False - - if groups == 1: - # Regular (pointwise) convolution. - pass - - elif conv_utils.group_conv_convertible_as_depthwise( - node, groups - ) and node_is_effectively_static_tensor(node.args[1], parameters_mapping): - # Depthwise convolution. - # Only supported if the weights are static, because TFLite `DepthwiseConv2D` uses permuted weights. In case - # the weights are dynamic, a Transpose operator would have to be added, which is not supported on Neutron. - pass - - elif conv_utils.group_conv_convertible_into_multiple_convolutions(node, groups): - # Group Separable convolution. - # Not supported natively by the eIQ Neutron so Group Separable Convolution. - # In practice it can be computed by splitting the Group Separable Convolution into multiple Pointwise - # Convo it will use the Split and Concat operation. The Concat operation in Neutron Converter - # SDK 25.03 requires the # of channels to be multipy of # of MAC units in the eIQ Neutron. - # For this reason Group Separable Convolution is not delegated by default at this moment. - return False - - else: - # All conversion options related to the `group` attribute have been checked and none of them can be used. + if output_padding != [0] * dimensions: return False if input_tensor_safe(node, 2) is None: @@ -89,10 +116,6 @@ def _is_supported_in_IR( if weight_tensor.dtype not in [torch.float32, torch.int8, torch.uint8]: return False - if node.args[0].meta["val"].shape[0] != 1: - # Only batch size 1 is supported on neutron. - return False - return True Stride = Padding = Dilation = OutPadding = list[int] @@ -109,7 +132,107 @@ def _get_convolution_arguments( _, _, _, stride, padding, dilation, transposed, out_padding, groups = ( conv_node.args ) - return stride, padding, dilation, transposed, out_padding, groups + return ( + list(stride), + list(padding), + list(dilation), + transposed, + out_padding, + groups, + ) + + def _convert_1d_conv( + self, t_op: tflite_model.Operator, conv_params: ConvParameters + ) -> list[tflite_model.Operator]: + """Convert the 'Conv' operator with a 1D kernel to TFLite 'Conv2D'. + TFLite doesn't support 1D convolution, but this behaviour can be represented using + Reshape -> Conv2D -> Reshape. + The first reshape introduces a 4th dimension with size 1. The second Reshape removes the temporary dimension. + """ + # -- Calculate the shapes for equivalent 2D convolution -- + conv_2d_input_shape = translator.nhc_dimensions_to_nhwc( + t_op.tmp_inputs[0].shape.vector + ) + conv_2d_weight_shape = translator.nhc_dimensions_to_nhwc( + t_op.tmp_inputs[1].shape.vector + ) + conv_2d_output_shape = translator.nhc_dimensions_to_nhwc( + t_op.tmp_outputs[0].shape.vector + ) + + # -- Generate tensors taking part in the conversion -- + reshape1_input = t_op.tmp_inputs[0] + + reshape1_output = self.builder.duplicate_tensor( + reshape1_input, name_suffix="_4D_" + ) + reshape1_output.shape = tflite_model.Shape(conv_2d_input_shape) + + reshape2_input = self.builder.duplicate_tensor( + t_op.tmp_outputs[0], name_suffix="_4D_" + ) + reshape2_input.shape = tflite_model.Shape(conv_2d_output_shape) + + reshape2_output = t_op.tmp_outputs[0] + + pre_reshapes = [] + + # Extend the weights tensor to 4D + weights_tensor = t_op.tmp_inputs[1] + if tensor_has_data(weights_tensor): + # Do it statically + weights_tensor.shape = tflite_model.Shape(conv_2d_weight_shape) + weights_tensor.tmp_buffer.data = weights_tensor.tmp_buffer.data.reshape( + conv_2d_weight_shape + ) + + else: + # Add a Reshape before the weights tensor + new_weights_tensor = self.builder.duplicate_tensor( + weights_tensor, name_suffix="_4D_" + ) + new_weights_tensor.shape = tflite_model.Shape(conv_2d_weight_shape) + + weight_reshape = tflite_model.Operator( + builtin_options=reshape_options.Reshape(conv_2d_weight_shape) + ) + weight_reshape.tmp_inputs = [weights_tensor] + weight_reshape.tmp_outputs = [new_weights_tensor] + + pre_reshapes.append(weight_reshape) + + # Save the new weights tensor, to assign it later. + weights_tensor = new_weights_tensor + + # -- Create the new operators -- + reshape1 = tflite_model.Operator( + builtin_options=reshape_options.Reshape(conv_2d_input_shape) + ) + reshape1.tmp_inputs = [reshape1_input] + reshape1.tmp_outputs = [reshape1_output] + pre_reshapes.append(reshape1) + + reshape2 = tflite_model.Operator( + builtin_options=reshape_options.Reshape(reshape2_output.shape.vector) + ) + reshape2.tmp_inputs = [reshape2_input] + reshape2.tmp_outputs = [reshape2_output] + + # Assign the new input and output of the Conv2D + t_op.tmp_inputs = [reshape1_output, weights_tensor] + t_op.tmp_inputs[ + 2: + ] # Add bias as well, if present + t_op.tmp_outputs = [reshape2_input] + + # Extend all Conv attributes to 2D + common.extend_1d_stride_to_2d(conv_params.stride) + common.extend_1d_dilation_to_2d(conv_params.dilation) + common.extend_1d_padding_to_2d(conv_params.padding) + + # Convert the now 2D Conv + converted_conv_ops = self._convert_2d_conv(t_op, conv_params) + + return pre_reshapes + converted_conv_ops + [reshape2] # noinspection PyPep8Naming def _convert_unpadded_2D( @@ -140,7 +263,7 @@ def _convert_unpadded_2D( ) b = self.builder.create_zeros_tensor( - [output_channels], "zero_bias", bias_type, True + [output_channels], "zero_bias", bias_type, False ) # Compute scale and zero point for bias tensor @@ -175,9 +298,19 @@ def _convert_2d_conv( aten_translator.convert_padding(conv_params.padding) ) if explicit_padding is not None: - # Need to prepend a 'Pad' operator, which adds 0s. + # Need to prepend a 'Pad' operator, which adds 0s (or `zero_point` for the quantized case). + input_quantization = t_op.tmp_inputs[0].quantization + pad_value = ( + None + if input_quantization is None + else np.array(input_quantization.zero_point[0]).astype( + tf_lite_type_to_numpy(t_op.tmp_inputs[0].type) + ) + ) conversion_result.ops_list.add_pre( - self.builder.create_pad_operator_before(t_op, 0, explicit_padding) + self.builder.create_pad_operator_before( + t_op, 0, explicit_padding, constant_value=pad_value + ) ) # DepthwiseConv2D expects weights in format [kernel_channels, kernel_height, kernel_width, output_channels] @@ -194,17 +327,8 @@ def _convert_2d_conv( elif conv_utils.group_conv_convertible_into_multiple_convolutions( t_op, conv_params.groups ): - # Note: by default the Group Separable Convolution is rejected by the Neutron Partitioner, see the - # ConvolutionConveter._is_supported_in_IR() - t_op.builtin_options = conv_2d_options.Conv2D() - - return conv_utils.create_separated_convolutions_based_on_group( - t_op, - conv_params, - self.builder, - self._convert_unpadded_2D, - conv_utils.conv_op_factory, - ) + # This case should have been rejected in the `is_supported_on_target()` method. + raise RuntimeError("Group convolution was not decomposed.") else: # Convert to regular `Conv2D`. @@ -214,9 +338,19 @@ def _convert_2d_conv( aten_translator.convert_padding(conv_params.padding) ) if explicit_padding is not None: - # Need to prepend a 'Pad' operator, which adds 0s. + # Need to prepend a 'Pad' operator, which adds 0s (or `zero_point` for the quantized case). + input_quantization = t_op.tmp_inputs[0].quantization + pad_value = ( + None + if input_quantization is None + else np.array(input_quantization.zero_point[0]).astype( + tf_lite_type_to_numpy(t_op.tmp_inputs[0].type) + ) + ) conversion_result.ops_list.add_pre( - self.builder.create_pad_operator_before(t_op, 0, explicit_padding) + self.builder.create_pad_operator_before( + t_op, 0, explicit_padding, constant_value=pad_value + ) ) return conversion_result.ops_list.flatten() @@ -230,7 +364,9 @@ def convert(self, node: Node): conv_params = ConvParameters(stride, padding, dilation, groups) rank = t_op.tmp_inputs[1].shape.len() - if rank == 4: # Conv2D + if rank == 3: # Conv1D + ops_to_add = self._convert_1d_conv(t_op, conv_params) + elif rank == 4: # Conv2D ops_to_add = self._convert_2d_conv(t_op, conv_params) else: raise NotImplementedError( diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/hardtanh_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/hardtanh_converter.py index 53f493f4ed9..14d69ed42fb 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/hardtanh_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/hardtanh_converter.py @@ -1,12 +1,11 @@ -# Copyright (c) 2025 NXP -# All rights reserved. +# Copyright 2025 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.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator, @@ -16,7 +15,6 @@ class HardTanhConverter(NodeConverter): - supported_targets = [Target.RT700] # Maps possible input parameters of HardTanh to equivalent ReLU-based operators supported by TFLite. supported_modes_map = { @@ -28,7 +26,9 @@ class HardTanhConverter(NodeConverter): @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: _, min_value, max_value = node.args return (min_value, max_value) in HardTanhConverter.supported_modes_map.keys() diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool_2d_converter.py index cd917e9d217..ce9a3697318 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool_2d_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -10,8 +10,8 @@ ) from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.lib.tflite.TensorType import TensorType from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model @@ -27,11 +27,11 @@ class MaxPool2dConverter(NodeConverter): NOTE: max_pool2d_with_indices is a different operator and is unsupported. """ - supported_targets = [Target.RT700] - @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: n_args = len(node.args) 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 659efa24fb7..f03c403876f 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 @@ -10,6 +10,7 @@ create_channels_last_to_channels_first_permutation, ) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, Target, ) @@ -24,29 +25,44 @@ class MeanDimConverter(NodeConverter): - supported_targets = [Target.RT700] - @staticmethod - def _to_neg_dim(d, rank): - return d - rank if d > 0 else d + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + match target: + case Target.RT700: + # TODO: Consider different tensor formats (dim-order) + dim = node.args[1] + keepdim = node.args[2] if len(node.args) >= 3 else False + rank = len(node.args[0].meta["val"].shape) + dim = [MeanDimConverter._to_neg_dim(d, rank) for d in dim] + + # Only last 2 dimensions (H, W) and keepdim=True with rank=4 are supported on Neutron. + if rank != 4 or dim not in [[-1, -2], [-2, -1]] or not keepdim: + return False + + return True + + case _: + return False @staticmethod def _to_pos_dim(d, rank): return d + rank if d < 0 else d + @staticmethod + def _to_neg_dim(d, rank): + return d - rank if d > 0 else d + @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: - dim = node.args[1] - keepdim = node.args[2] if len(node.args) >= 3 else False - rank = len(node.args[0].meta["val"].shape) - dim = [MeanDimConverter._to_neg_dim(d, rank) for d in dim] - - # Only last 2 dimensions (H, W) and keepdim=True with rank=4 are supported on Neutron. - if rank != 4 or dim not in [[-1, -2], [-2, -1]] or not keepdim: - return False - if hasattr(node.kwargs, "dtype") and node.kwargs["dtype"] not in [ torch.float32, torch.uint32, diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mm_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mm_converter.py index fc513240c44..dd9e3e2da54 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mm_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mm_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -6,8 +6,8 @@ from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( fully_connected_options, @@ -17,11 +17,12 @@ class MMConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: if len(node.all_input_nodes) != 2: 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 e24ed4f6863..f0150b4bc1f 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 @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -8,8 +8,8 @@ from executorch.backends.nxp.backend.ir.converter import quantization_utils from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( transpose_options, @@ -19,11 +19,12 @@ class PermuteCopyConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: return True diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_dequantize_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_dequantize_converter.py index 8731b3f6ed2..c6ea7f90042 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_dequantize_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_dequantize_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -9,8 +9,8 @@ torch_type_to_numpy_type, ) from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.converter.quantization_utils import ( set_quantization_parameters_to_tensor, @@ -20,11 +20,12 @@ class QDQDequantizeConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: zero_point_type = torch_type_to_numpy_type(node.args[5]) if "cluster" not in node.meta or zero_point_type not in [np.int8, np.int32]: diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_quantize_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_quantize_converter.py index b0680e9b949..32bcd9445d3 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_quantize_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/qdq_quantize_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -7,8 +7,8 @@ import torch from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.converter.quantization_utils import ( set_quantization_parameters_to_tensor, @@ -18,11 +18,12 @@ class QDQQuantizeConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: if "cluster" not in node.meta or node.args[5] != torch.int8: return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/relu_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/relu_converter.py index d1af0ec2de5..eb9d62287c0 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/relu_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/relu_converter.py @@ -1,12 +1,11 @@ -# Copyright (c) 2024 NXP -# All rights reserved. +# Copyright 2024-2025 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.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator, @@ -16,11 +15,12 @@ class ReLUConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: return True diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py index a9af12f60dd..96e4655d011 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/sigmoid_converter.py @@ -1,12 +1,11 @@ # Copyright 2025 NXP -# 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.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator, @@ -16,18 +15,12 @@ class SigmoidConverter(NodeConverter): - @staticmethod - def _is_supported_on_target(target: Target) -> bool: - match target: - case Target.RT700: - return True - - case _: - return False @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: return True diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/softmax_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/softmax_converter.py index 99932602c2f..aa74c78ca24 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/softmax_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/softmax_converter.py @@ -1,10 +1,16 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 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.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) from executorch.backends.nxp.backend.edge_helper import input_rank -from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + NodeConverter, + Target, +) from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( softmax_options, ) @@ -13,11 +19,27 @@ class SoftmaxConverter(NodeConverter): - supported_targets = [] + @staticmethod + def _is_supported_on_target( + node: Node, + target: Target, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + match target: + case Target.RT700: + # The eIQ Neutron NPU runtime software has a known issue with the SoftMax operation. + # As long as the issue is present, return False for the i.MX RT700 target also. + return False + + case _: + return False @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: # The IR only supports the `dim` as the last dimension. But that depends on the format of the input tensor, # which is only known after the `Partitioner` has divided the model. So if the input shape can be channels diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/tanh_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/tanh_converter.py new file mode 100644 index 00000000000..427865f8ee7 --- /dev/null +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/tanh_converter.py @@ -0,0 +1,33 @@ +# Copyright 2025 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.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter +from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( + BuiltinOperator, +) +from torch.fx import Node +from torch.nn import Parameter + + +class TanhConverter(NodeConverter): + + @staticmethod + def _is_supported_in_IR( + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, + ) -> bool: + return True + + def convert(self, node: Node): + self.assert_convertible(node) + + t_op = self._create_tflite_op_with_io_tensors(node) + t_op.opcode_index = self.builder.op_code_index_for_op_type(BuiltinOperator.TANH) + + self.builder.append_operators([t_op]) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/view_copy_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/view_copy_converter.py index 2eceeba9b24..95a42d5d078 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/view_copy_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/view_copy_converter.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -13,8 +13,8 @@ from executorch.backends.nxp.backend.ir.converter import quantization_utils from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList from executorch.backends.nxp.backend.ir.converter.node_converter import ( + CustomDelegationOptions, NodeConverter, - Target, ) from executorch.backends.nxp.backend.ir.converter.node_converters.shared.reshape_transposition import ( ensure_reshape_transposition, @@ -27,11 +27,12 @@ class ViewCopyConverter(NodeConverter): - supported_targets = [Target.RT700] @staticmethod def _is_supported_in_IR( - node: Node, parameters_mapping: dict[str, Parameter] + node: Node, + parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ) -> bool: x = input_tensor(node, 0) y = output_tensor(node) diff --git a/backends/nxp/backend/ir/converter/node_converters/shared/conv_utils.py b/backends/nxp/backend/ir/converter/node_converters/shared/conv_utils.py index 73bf76a830d..5817fd127b3 100755 --- a/backends/nxp/backend/ir/converter/node_converters/shared/conv_utils.py +++ b/backends/nxp/backend/ir/converter/node_converters/shared/conv_utils.py @@ -1,28 +1,13 @@ # Copyright 2023-2025 NXP # -# License: LA_OPT_NXP_Software_License -# See the LICENSE_LA_OPT_NXP_Software_License for more details. -# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. -from copy import copy from dataclasses import dataclass -from typing import Callable, cast - -import numpy as np -from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( - ModelBuilder, -) -from executorch.backends.nxp.backend.ir.converter.conversion import aten_translator from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList -from executorch.backends.nxp.backend.ir.converter.tensor_utils import tensor_has_data -from executorch.backends.nxp.backend.ir.lib.tflite.Padding import Padding from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model -from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import ( - concatenation_options, - conv_2d_options, - split_options, -) + from torch.fx import Node @@ -62,6 +47,9 @@ def group_conv_convertible_into_multiple_convolutions( if group == 1: return False + if group_conv_convertible_as_depthwise(node, group): + return False + _, output_channels = _get_IO_channels(node) if output_channels % group != 0: return False # Unable to split group Conv into separated convolutions because out_channels % group != 0. @@ -88,313 +76,3 @@ def __init__( self.conv_bias_tensor = bias_tensor self.conv_output_tensor = output_tensor self.ops_list = OpsList() - - -ConvBuiltinOptions = conv_2d_options.Conv2D -ConvOpFactory = Callable[ - [ - ConvParameters, - tflite_model.Tensor, - tflite_model.Tensor, - tflite_model.Tensor, - tflite_model.Tensor, - ModelBuilder, - ConvBuiltinOptions, - ], - OpsList, -] -ConvConversionFn = Callable[ - [tflite_model.Operator, ConvParameters], ConvConversionResult -] - - -class _InputTensorsSplitter: - """Splits the tensors of a `Conv2D` operator. Static tensors are split statically, and for dynamic tensors, a - TFLite `Split` operator is added. - """ - - input_tensors: list[tflite_model.Tensor] - weight_tensors: list[tflite_model.Tensor] - bias_tensors: list[tflite_model.Tensor] - split_ops: list[tflite_model.Operator] - - def __init__( - self, - input_tensor: tflite_model.Tensor, - weight_tensor: tflite_model.Tensor, - bias_tensor: tflite_model.Tensor, - groups: int, - builder: ModelBuilder, - ): - self.input_tensors = [] - self.weight_tensors = [] - self.bias_tensors = [] - self.split_ops = [] - - inputs = [ - # input tensor, split by axis, output tensors container - (input_tensor, -1, self.input_tensors), - (weight_tensor, 0, self.weight_tensors), - (bias_tensor, 0, self.bias_tensors), - ] - - for i in inputs: - if tensor_has_data(i[0]): - self._generate_static_tensors(builder, groups, i[0], i[1], i[2]) - else: - self._generate_dynamic_tensors(builder, groups, i[0], i[1], i[2]) - - def _generate_dynamic_tensors( - self, builder, groups, split_tensor, axis, target_list - ): - quantization = None - if split_tensor.quantization is not None: - if split_tensor.quantization.is_per_channel(): - scale = np.split( - np.array(split_tensor.quantization.scale.vector, "float32"), groups - ) - zero_point = np.split( - np.array(split_tensor.quantization.zero_point.vector, "int32"), - groups, - ) - quantization = [ - tflite_model.Quantization( - scale=tflite_model.Scale(s), - zero_point=tflite_model.ZeroPoint(zp), - ) - for s, zp in zip(scale, zero_point) - ] - else: - quantization = [split_tensor.quantization] * groups - - split_op = self._create_split_op(builder, groups, split_tensor, axis) - - new_tensor_shape = split_tensor.shape.vector.copy() - new_tensor_shape[axis] = new_tensor_shape[axis] // groups - - for i in range(groups): - conv_split_tensor = builder.duplicate_tensor( - split_tensor, name_suffix="_group_" + str(i) - ) - conv_split_tensor.shape = tflite_model.Shape(new_tensor_shape) - if quantization is not None: - conv_split_tensor.quantization = copy(quantization[i]) - - split_op.tmp_outputs.append(conv_split_tensor) - target_list.append(conv_split_tensor) - self.split_ops.append(split_op) - - # noinspection PyMethodMayBeStatic - def _generate_static_tensors( - self, builder, groups, split_tensor, axis, target_list - ): - quantization = None - if split_tensor.quantization is not None: - if split_tensor.quantization.is_per_channel(): - scale = np.split( - np.array(split_tensor.quantization.scale.vector, "float32"), groups - ) - zero_point = np.split( - np.array(split_tensor.quantization.zero_point.vector, "int32"), - groups, - ) - quantization = [ - tflite_model.Quantization( - scale=tflite_model.Scale(s), - zero_point=tflite_model.ZeroPoint(zp), - ) - for s, zp in zip(scale, zero_point) - ] - else: - quantization = [split_tensor.quantization] * groups - - input_data = np.split(split_tensor.tmp_buffer.data, groups, axis) - - for i in range(len(input_data)): - tensor_name = split_tensor.name + "_group_" + str(i) - conv_input_tensor = builder.create_tensor_for_data( - input_data[i], tensor_name - ) - if quantization is not None: - conv_input_tensor.quantization = copy(quantization[i]) - - target_list.append(conv_input_tensor) - - # noinspection PyMethodMayBeStatic - def _create_split_op(self, builder, groups, input_tensor, axis): - axis_tensor = builder.create_tensor_for_data( - np.asarray([axis], np.int32), "split_dim_" - ) - input_split_op = tflite_model.Operator( - builtin_options=split_options.Split(groups) - ) - input_split_op.tmp_inputs = [axis_tensor, input_tensor] - - return input_split_op - - def get_input_tensor(self, idx) -> tflite_model.Tensor: - return self.input_tensors[idx] - - def get_weight_tensor(self, idx) -> tflite_model.Tensor: - return self.weight_tensors[idx] - - def get_bias_tensor(self, idx) -> tflite_model.Tensor: - return self.bias_tensors[idx] - - def get_ops(self) -> list[tflite_model.Operator]: - return self.split_ops - - -class _OutputTensorsCombiner: - """Handles creation and aggregation of the TFLite Conv2D output tensors. - Aggregation is done with `Concatenation` op. - """ - - output_tensors: list[tflite_model.Tensor] - concat_op: tflite_model.Operator - - def __init__(self, output_tensor, groups, builder): - self.output_tensors = [] - combine_axis = -1 - - new_conv_output_shape = output_tensor.shape.vector.copy() - new_conv_output_shape[combine_axis] = ( - new_conv_output_shape[combine_axis] // groups - ) - conv_output_shape = tflite_model.Shape(new_conv_output_shape) - - self.concat_op = tflite_model.Operator( - builtin_options=concatenation_options.Concatenation(combine_axis) - ) - self.concat_op.tmp_outputs = [output_tensor] - - for i in range(groups): - tensor_name = output_tensor.name + "_group_" + str(i) - output_tensor = builder.duplicate_tensor(output_tensor, tensor_name) - output_tensor.shape = conv_output_shape - - self.output_tensors.append(output_tensor) - self.concat_op.tmp_inputs.append(output_tensor) - - def get_output_tensor(self, idx): - return self.output_tensors[idx] - - def get_ops(self): - return [self.concat_op] - - -def build_input_tensor_padding( - t_op, conv_params: ConvParameters, builder, input_idx=0 -) -> (Padding, tflite_model.Operator | None): - """Build padding for input tensor of Conv2D op 't_op'.""" - - tfl_padding, explicit_padding = aten_translator.convert_padding(conv_params.padding) - if explicit_padding is not None: - # Must add extra 'Pad' operator - return tfl_padding, builder.create_pad_operator_before( - t_op, input_idx, explicit_padding - ) - - return tfl_padding, None - - -def conv_op_factory( - conv_params: ConvParameters, - input_tensor: tflite_model.Tensor, - weight_tensor: tflite_model.Tensor, - bias_tensor: tflite_model.Tensor, - output_tensor: tflite_model.Tensor, - builder, - builtin_options, -) -> OpsList: - """Build padded 'Conv2D' TFLite operator. Padding is realized by 'builtin_options.padding' definition and by - optional prepended 'Pad' operator. - """ - - conv_op = tflite_model.Operator(builtin_options=copy(builtin_options)) - conv_op.tmp_inputs = [input_tensor, weight_tensor, bias_tensor] - conv_op.tmp_outputs = [output_tensor] - - padding, pad_op = build_input_tensor_padding(conv_op, conv_params, builder) - conv_op.builtin_options.padding = padding - - if pad_op is not None: - return OpsList(pre_ops=[pad_op], middle_op=conv_op) - else: - return OpsList(middle_op=conv_op) - - -# noinspection GrazieInspection -def create_separated_convolutions_based_on_group( - t_op: tflite_model.Operator, - conv_params: ConvParameters, - builder: ModelBuilder, - conv_conversion_fn: ConvConversionFn, - conv_op_factory_fn: ConvOpFactory, -) -> list[tflite_model.Operator]: - """Build a subgraph with multiple TFLite Conv2D operators that replace an `aten.convolution` operator with 'group' - attribute higher than one. The number of new Conv2D operators corresponds to the number of groups. Input - tensors of the Aten operator are split and distributed into related convolution operators. Outputs are then - concatenated back together. - - Example: 'aten.convolution' operator with group=2 converted into TFLite subgraph will have - the following structure (tensor dimensions are just for illustrative purposes): - - │ (1,4,4,48) - ┌───▼──┐ - │Split │ - └┬────┬┘ - (1,4,4,24) │ │ (1,4,4,24) - ┌─────▼┐ ┌▼─────┐ - │Conv2D│ │Conv2D│ - └────┬─┘ └─┬────┘ - (1,4,4,18)│ │(1,4,4,18) - ┌─▼──────▼──┐ - │Concatenate│ - └─────┬─────┘ - │ (1,4,4,36) - ▼ - """ - - conversion_result = conv_conversion_fn(t_op, conv_params) - - splitter = _InputTensorsSplitter( - conversion_result.conv_input_tensor, - conversion_result.conv_weight_tensor, - conversion_result.conv_bias_tensor, - conv_params.groups, - builder, - ) - combiner = _OutputTensorsCombiner( - conversion_result.conv_output_tensor, conv_params.groups, builder - ) - - conv_ops = [] - for i in range(conv_params.groups): - input_tensor = splitter.get_input_tensor(i) - weight_tensor = splitter.get_weight_tensor(i) - bias_tensor = splitter.get_bias_tensor(i) - output_tensor = combiner.get_output_tensor(i) - - conv_builtin_options = cast( - ConvBuiltinOptions, conversion_result.ops_list.middle_op.builtin_options - ) - conv_ops_list = conv_op_factory_fn( - conv_params, - input_tensor, - weight_tensor, - bias_tensor, - output_tensor, - builder, - conv_builtin_options, - ) - - conv_ops.extend(conv_ops_list.flatten()) - - return ( - conversion_result.ops_list.pre_ops # `Pad` operator - + splitter.get_ops() - + conv_ops - + combiner.get_ops() # Split, Conv2D, Concatenate ops - + conversion_result.ops_list.post_ops - ) # Currently not used diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/add_n_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/add_n_options.py index 2646f326852..744d2b332b3 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/add_n_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/add_n_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/add_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/add_options.py index 37c04a84588..48c82a9974f 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/add_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/add_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/average_pool_2d_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/average_pool_2d_options.py index d3f59b3844d..1bafc61cb60 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/average_pool_2d_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/average_pool_2d_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/leaky_relu_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/leaky_relu_options.py index 6ba7bb65d72..848faa6c34b 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/leaky_relu_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/leaky_relu_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/log_softmax_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/log_softmax_options.py index 163cbfb7cf9..a700c524562 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/log_softmax_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/log_softmax_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/max_pool_2d_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/max_pool_2d_options.py index b87a2f46de2..13d827d98f3 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/max_pool_2d_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/max_pool_2d_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/reshape_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/reshape_options.py index 800bd645b8a..66e1e836c38 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/reshape_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/reshape_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/softmax_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/softmax_options.py index 3001f659d40..ce828c0e1fe 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/softmax_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/softmax_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/sub_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/sub_options.py index 16dcd1e64ab..226b5bb498d 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/sub_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/sub_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_generator/builtin_options/transpose_options.py b/backends/nxp/backend/ir/tflite_generator/builtin_options/transpose_options.py index 5869b1ed315..48052690b18 100755 --- a/backends/nxp/backend/ir/tflite_generator/builtin_options/transpose_options.py +++ b/backends/nxp/backend/ir/tflite_generator/builtin_options/transpose_options.py @@ -1,5 +1,6 @@ # # Copyright 2023 Martin Pavella +# Copyright 2024 NXP # # License: MIT # See the LICENSE_MIT for more details. diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/eliminate_dead_branches.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/eliminate_dead_branches.py deleted file mode 100755 index cea179dfb09..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/eliminate_dead_branches.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright 2024 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.nxp.backend.ir import logger -from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) - - -class EliminateDeadBranches(BaseOptimization): - - def __call__(self) -> bool: - _, output_to_ops = self._create_tensor_to_operator_dictionaries() - - output_names = [ - tensor.name for tensor in self._builder.get_sub_graph().outputs.tmp_outputs - ] - - tensor_names_to_process = set(output_names) - tensors_to_keep = set() - ops_to_keep = set() - processed_ops = set() - - # Iterate from output tensors to inputs and mark all visited nodes & tensors - while len(tensor_names_to_process) != 0: - tensor = tensor_names_to_process.pop() - tensors_to_keep.add(tensor) - - if tensor not in output_to_ops: - # Input tensor or already processed - continue - - op: tflite_model.Operator = output_to_ops[tensor] - - if op in processed_ops: - continue - - # Append all inputs and outputs to next processing. Outputs of nodes aren't - # necessarily outputs of the model but must be preserved. - for tensor in op.tmp_inputs + op.tmp_outputs: - tensor_names_to_process.add(tensor.name) - - ops_to_keep.add(op) - processed_ops.add(op) - - if not self._conversion_config.allow_inputs_stripping: - # Keep all inputs (even if they are not used) when prohibited by user - tensors_to_keep.update( - [ - tensor.name - for tensor in self._builder.get_sub_graph().inputs.tmp_inputs - ] - ) - - # Remove unused ops - ops = self._builder.get_operators().vector - i, removed_ops_count = 0, 0 - while i < len(ops): - if ops[i] in ops_to_keep: - i += 1 - else: - removed_ops_count += 1 - del ops[i] - - # Remove unused tensors - tensors = self._builder.get_tensors().vector - i = 0 - while i < len(tensors): - if tensors[i].name in tensors_to_keep: - i += 1 - else: - del tensors[i] - - if removed_ops_count != 0: - logger.i( - f"Dead branch elimination optimization removed {removed_ops_count} unused ops from the graph." - ) - - return removed_ops_count != 0 diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/fuse_quanitze_into_preceding_ops.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/fuse_quanitze_into_preceding_ops.py deleted file mode 100755 index 6b3bd70cc01..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/fuse_quanitze_into_preceding_ops.py +++ /dev/null @@ -1,94 +0,0 @@ -# Copyright 2024 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.nxp.backend.ir.tflite_optimizer.operator_rules import ( - WasNotInTheOriginalONNXModel, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.pattern_matcher import ( - Op, - PatternMatcher, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.tensor_rules import ( - TensorHasOneConsumer, - TensorsArePerTensorQuantized, - TensorsHaveSameType, -) - - -class FuseQuantizeIntoPrecedingOps(BaseOptimization): - """Remove some `Quantize` operators in the following pattern. - - │ - ┌─▼──┐ - │ Op │ │ - └─┬──┘ ┌─▼──┐ - │ 'x' (same type, quantization params `A`) ─────► │ Op │ - ┌────▼─────┐ └─┬──┘ - │ Quantize │ │ (same type, quantization params `B`) - └────┬─────┘ - │ 'y' (same type, quantization params `B`) - """ - - ops_that_can_have_any_output_quantization = [ - # List of operators which don't have restrictions placed on their output quantization and are currently - # supported by `onnx2quant`. - "Add", - "BatchMatMul", - "FullyConnected", - "HardSwish", - "LeakyRelu", - "Mean", - "Mul", - "PRelu", - "ReduceProd", - "Relu", - "Sub", - "Sum", - ] - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [ - Op( - self.ops_that_can_have_any_output_quantization, - outputs=[..., "x", ...], - ), - Op( - ["Quantize"], - ["x"], - ["y"], - [ - # Restrict this optimization to extra `Quantize` operators which were added during conversion. - # Sometimes the `Quantize` operators which are present in the ONNX model can be essential and - # shouldn't be removed. They can for example perform clipping. - WasNotInTheOriginalONNXModel() - ], - ), - ], - [ - TensorHasOneConsumer("x"), - # Make sure the `Quantize` is just changing quantization parameters. Otherwise, it couldn't be fused. - TensorsHaveSameType(["x", "y"]), - TensorsArePerTensorQuantized(["x", "y"]), - ], - ) - - to_remove = [] - for [leading_op, quantize], tensor_map, _, _ in matcher.match_patterns(): - x, y = tensor_map["x"], tensor_map["y"] - - x_idx = leading_op.tmp_outputs.index(x) - leading_op.tmp_outputs[x_idx] = y - - to_remove.append(quantize) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_cast_operators.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_cast_operators.py deleted file mode 100755 index 8cce0bb61e8..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_cast_operators.py +++ /dev/null @@ -1,117 +0,0 @@ -# Copyright 2024 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.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.pattern_matcher import ( - MultipleSameOps, - Op, - PatternMatcher, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.tensor_rules import ( - RuleOr, - TensorIsNotModelOutput, - TensorIsNotQuantized, - TensorsAreNotQuantized, - TensorsHaveSameType, -) - - -class FuseCastOperators(BaseOptimization): - """Remove some `Cast` operators in the following pattern. - - │ 'x' - ┌──▼───┐ - │ Cast │ - └──┬───┘ │ 'x' - ┌─┴─── ... ──────┐ 'y' ─────► ┌──┴── ... ─────┐ ('y' is not in the model anymore) - ┌──▼───┐ ┌──▼───┐ ┌──▼───┐ ┌──▼───┐ - │ Cast │ ... │ Cast │ │ Cast │ ... │ Cast │ - └──┬───┘ └──┬───┘ └──┬───┘ └──┬───┘ - │ │ 'z' │ │ 'z' - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [ - Op(["Cast"], outputs=["y"]), - MultipleSameOps(["Cast"], ["y", ...]), # Only `Cast` ops can use `y`. - ], - [TensorIsNotModelOutput("y"), TensorIsNotQuantized("y")], - ) - - to_remove = [] - for [leading_cast, following_cast_ops], _, _, _ in matcher.match_patterns(): - # Remove the leading cast. - for cast in following_cast_ops: - cast.tmp_inputs[0] = leading_cast.tmp_inputs[0] - - to_remove.append(leading_cast) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 - - -class RemoveCastOperatorsWithNoEffect(BaseOptimization): - """Remove operators that match the following pattern. - - │ 'x' - ┌──▼───┐ - │ Cast │ - └──┬───┘ - │ 'y' (same type as 'x') - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [Op(["Cast"], ["x", ...], ["y"])], - [ - TensorsHaveSameType(["x", "y"]), - TensorsAreNotQuantized(["x", "y"]), - RuleOr( - TensorIsNotModelOutput("x"), - TensorIsNotModelOutput("y"), - # If both 'x' and 'y' are model outputs, the `Cast` cannot be removed. If the op was removed, its - # input and output would be combined into 1 tensor, which would have to represent 2 model outputs - # with 2 different names, which is not possible. - ), - ], - ) - - to_remove = [] - for [cast], tensor_map, input_to_ops, _ in matcher.match_patterns(): - if not self._builder.operator_can_be_skipped(cast): - continue - - x = tensor_map["x"] - y = tensor_map["y"] - model_outputs = self._builder.get_sub_graph().outputs.tmp_outputs - - # Replace `y` with `x` in the inputs of all following operators. - following_ops = input_to_ops.get(y.name, []) - for op in following_ops: - while y in op.tmp_inputs: - input_idx = op.tmp_inputs.index(y) - op.tmp_inputs[input_idx] = x - - if y in model_outputs: - # Replace the output as well. - while y in model_outputs: - idx = model_outputs.index(y) - model_outputs[idx] = x - - self._builder.swap_tensor_names(x, y) - - to_remove.append(cast) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_quantize_operators.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_quantize_operators.py deleted file mode 100755 index 317654cde9a..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_quantize_operators.py +++ /dev/null @@ -1,304 +0,0 @@ -# Copyright 2024 NXP -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import numpy as np - -from executorch.backends.nxp.backend.ir.lib.tflite.TensorType import TensorType -from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.pattern_matcher import ( - MultipleSameOps, - Op, - PatternMatcher, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.tensor_rules import ( - TensorIsNotModelOutput, - TensorsHaveSameQuantization, -) - - -class FuseParallelQuantizeOperators(BaseOptimization): - """Fuse some `Quantize` operators in the following pattern. - - │ 'x' │ 'x' - ┌───┴──── ... ───────┐ ┌────▼─────┐ - ┌────▼─────┐ ┌────▼─────┐ ─────► │ Quantize │ - │ Quantize │ ... │ Quantize │ └────┬─────┘ - └────┬─────┘ └────┬─────┘ ┌───┴─ ... ─┐ - │ │ 'y' (same quantization) │ │ 'y' - - - The pattern below only has 2 `Quantize` operators. But the `PatternMatcher` will gradually match all parallel - `Quantize` operators which fit the pattern above, and remove the unnecessary ones. - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [Op(["Quantize"], ["x"], ["y1"]), Op(["Quantize"], ["x"], ["y2"])], - [ - TensorsHaveSameQuantization(["y1", "y2"]), - # 'y2' will be removed from the model, so it cannot be a model output. But thanks to the nature of the - # `PatternMatcher`, it doesn't matter which `Quantize` produces the model output. The `PatternMatcher` - # will first match 1 `Quantize` as the first `Op` and try to optimize. If it doesn't work, it will then - # match the second `Quantize` operator with the first `Op` and try to optimize that way. This will - # result in a perfectly optimized pattern every time. - TensorIsNotModelOutput("y2"), - ], - ) - - to_remove = [] - for ( - [_, quant_to_remove], - tensor_map, - input_to_ops, - _, - ) in matcher.match_patterns(): - to_remove.append(quant_to_remove) - - y1 = tensor_map["y1"] - y2 = tensor_map["y2"] - next_ops = input_to_ops.get(y2.name, []) - for next_op in next_ops: - while y2 in next_op.tmp_inputs: - idx = next_op.tmp_inputs.index(y2) - next_op.tmp_inputs[idx] = y1 - - quant_to_remove.tmp_inputs = ( - [] - ) # To prevent future matches of this operator. - - ops = self._builder.get_operators() - for op in to_remove: - ops.remove(op) - - return ops.len() != 0 - - -# noinspection PyMethodMayBeStatic -class PruneQuantizeOperators(BaseOptimization): - """Remove some `Quantize` operators in the following pattern. - - │ 'x' - ┌────▼─────┐ - │ Quantize │ - └────┬─────┘ - ┌───┴──── ... ───────┐ 'y' - ┌────▼─────┐ ┌────▼─────┐ - │ Quantize │ ... │ Quantize │ - └────┬─────┘ └────┬─────┘ - │ │ 'z' - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [ - Op(["Quantize"], ["x"], ["y"]), - MultipleSameOps( - ["Quantize"], ["y"] - ), # Nothing other than `Quantize` ops can use `y`. - ], - [TensorIsNotModelOutput("y")], - ) - - to_remove = [] - for ( - [leading_quantize, following_quantize_ops], - tensor_map, - input_to_ops, - _, - ) in matcher.match_patterns(): - x = tensor_map["x"] - - if self._is_quantization_recasting_from_float(x, following_quantize_ops): - # First Quantize can be skipped because it does only recasting - to_remove.append(leading_quantize) - - for next_quantize in following_quantize_ops: - next_quantize.tmp_inputs[0] = x - - elif self._is_quantization_recasting_from_integer( - x, following_quantize_ops - ): - # The Quantize ops negate each other -> remove them both - to_remove.append(leading_quantize) - - graph_outputs = self._builder.get_sub_graph().outputs.tmp_outputs - for next_quantize in following_quantize_ops: - to_remove.append(next_quantize) - - # Replace the output of the next Quantize with the input of the first Quantize - next_quantize_output = next_quantize.tmp_outputs[0] - self._bypass_to_next_quantize_ops( - input_to_ops, next_quantize_output, x - ) - - # If the output of the first Quantize is also the graph output -> replace the graph output too - if next_quantize_output in graph_outputs: - graph_outputs.remove(next_quantize_output) - if x not in graph_outputs: - graph_outputs.append(x) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 - - def _is_quantization_recasting_from_float( - self, quantize_input: tflite_model.Tensor, next_ops: list[tflite_model.Operator] - ): - """ - Check if 'next_ops' just recast from one type to another. Scale + recalculated zp - must be the same for all nodes. Input of first Quantize op has to be float to match - criteria. - - float int8 uint8 - ----> [quant] ---> [quant] -----> - zp zp-128 - OR - - float uint8 int8 - ----> [quant] ----> [quant] -----> - zp zp+128 - - OR (forked variant with similar restrictions as mentioned above) - - u / int u/int - float | ----> [quant] ----> - ----> [quant] --| - | ----> [quant] ----> - u/int u/int - - :param quantize_input: Input tensor of first QuantizeLinear node. - :param next_ops: QuantizeLinear ops that consume output of 'quantize_input'. - :return: True if pattern with recasting is found. - """ - - if not quantize_input.type == TensorType.FLOAT32: - return False - - # All 'next_ops' has the same output type and q-params - next_op_output_match_first = [ - self._same_type_and_quantization( - next_ops[0].tmp_outputs[0], next_op.tmp_outputs[0] - ) - for next_op in next_ops - ] - if not all(next_op_output_match_first): - return False - - # All 'next_ops' are the same, do some additional checks on the first one - - next_op_input = next_ops[0].tmp_inputs[0] - next_op_output = next_ops[0].tmp_outputs[0] - - input_zp = next_op_input.quantization.zero_point.vector - output_zp = next_op_output.quantization.zero_point.vector - - if next_op_input.quantization.scale != next_op_output.quantization.scale: - return False - - if ( - next_op_input.type == TensorType.INT8 - and next_op_output.type == TensorType.UINT8 - ): - return np.equal(input_zp, np.array(output_zp) - 128) - elif ( - next_op_input.type == TensorType.UINT8 - and next_op_output.type == TensorType.INT8 - ): - return np.equal(input_zp, np.array(output_zp) + 128) - - return False - - def _is_quantization_recasting_from_integer( - self, quantize_input: tflite_model.Tensor, next_ops: list[tflite_model.Operator] - ): - """ - Check if 'next_ops' just recast from one type to another. Scale + recalculated zp - must be the same for all nodes. Input of first Quantize op has to be (u)int8 to - match criteria. - - uint8 int8 uint8 - ----> [quant] -----> [quant] ----> - zp zp+128 zp - - OR - - int8 uint8 int8 - ---> [quant] -----> [quant] ---> - zp zp-128 zp - - OR (forked variant with similar restrictions as mentioned above) - - u/int u/int - u/int | ----> [quant] ----> - ----> [quant] --| - | ----> [quant] ----> - u/int u/int - - :param quantize_input: Input tensor of first QuantizeLinear node. - :param next_ops: QuantizeLinear ops that consume output of 'quantize_input'. - :return: True if pattern with recasting is found. - """ - - if quantize_input.type not in [TensorType.INT8, TensorType.UINT8]: - return False - - # All 'next_ops' has the same output type and q-params as input of first Quantize - next_op_output_match_first = [ - self._same_type_and_quantization(quantize_input, next_op.tmp_outputs[0]) - for next_op in next_ops - ] - if not all(next_op_output_match_first): - return False - - # All 'next_ops' are the same, do some additional checks on the first one - - next_op_input = next_ops[0].tmp_inputs[0] - next_op_output = next_ops[0].tmp_outputs[0] - - input_zp = next_op_input.quantization.zero_point.vector - output_zp = next_op_output.quantization.zero_point.vector - - if quantize_input.quantization.scale != next_op_input.quantization.scale: - return False - - if next_op_input.quantization.scale != next_op_output.quantization.scale: - return False - - if ( - next_op_input.type == TensorType.INT8 - and next_op_output.type == TensorType.UINT8 - ): - return np.equal(input_zp, np.array(output_zp) - 128) - elif ( - next_op_input.type == TensorType.UINT8 - and next_op_output.type == TensorType.INT8 - ): - return np.equal(input_zp, np.array(output_zp) + 128) - - return False - - def _same_type_and_quantization( - self, a: tflite_model.Tensor, b: tflite_model.Tensor - ): - same_type = a.type == b.type - same_quantization = a.quantization == b.quantization - - return same_type and same_quantization - - def _bypass_to_next_quantize_ops( - self, input_to_ops, next_quantize_output, quantize_input - ): - ops_after_next_quantize = input_to_ops.get(next_quantize_output.name, []) - for op_after_next_quantize in ops_after_next_quantize: - for index, input_tensor in enumerate(op_after_next_quantize.tmp_inputs): - if input_tensor == next_quantize_output: - # Replace the input - op_after_next_quantize.tmp_inputs[index] = quantize_input diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_reshape_operators.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_reshape_operators.py deleted file mode 100755 index 229d4747a7c..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/prune_reshape_operators.py +++ /dev/null @@ -1,116 +0,0 @@ -# Copyright 2024 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.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.pattern_matcher import ( - MultipleSameOps, - Op, - PatternMatcher, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.tensor_rules import ( - RuleOr, - TensorIsNotModelOutput, - TensorsHaveSameShape, -) - - -class FuseReshapeOperators(BaseOptimization): - """Remove some `Reshape` operator in the following pattern. - - │ 'x' - ┌────▼────┐ - │ Reshape │ - └────┬────┘ │ 'x' - ┌───┴─── ... ───────┐ 'y' ─────► ┌───┴─── ... ───────┐ ('y' is not in the model anymore) - ┌────▼────┐ ┌────▼────┐ ┌────▼────┐ ┌────▼────┐ - │ Reshape │ ... │ Reshape │ │ Reshape │ ... │ Reshape │ - └────┬────┘ └────┬────┘ └────┬────┘ └────┬────┘ - │ │ 'z' │ │ 'z' - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [ - Op(["Reshape"], outputs=["y"]), - MultipleSameOps( - ["Reshape"], ["y", ...] - ), # Nothing other than `Reshape` ops can use `y`. - ], - [TensorIsNotModelOutput("y")], - ) - - to_remove = [] - for [leading_reshape, following_reshapes], _, _, _ in matcher.match_patterns(): - # Remove the leading reshape. - for r in following_reshapes: - r.tmp_inputs[0] = leading_reshape.tmp_inputs[0] - - to_remove.append(leading_reshape) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 - - -class RemoveReshapeOperatorsWithNoEffect(BaseOptimization): - """Remove operators that match the following pattern. - - │ 'x' - ┌────▼────┐ - │ Reshape │ - └────┬────┘ - │ 'y' (same shape as 'x') - """ - - def __call__(self) -> bool: - matcher = PatternMatcher( - self._builder, - [Op(["Reshape"], ["x", ...], ["y"])], - [ - TensorsHaveSameShape(["x", "y"]), - RuleOr( - TensorIsNotModelOutput("x"), - TensorIsNotModelOutput("y"), - # If both 'x' and 'y' are model outputs, the `Reshape` cannot be removed. If the op was removed, its - # input and output would be combined into 1 tensor, which would have to represent 2 model outputs - # with 2 different names, which is not possible. - ), - ], - ) - - to_remove = [] - for [reshape], tensor_map, input_to_ops, _ in matcher.match_patterns(): - if not self._builder.operator_can_be_skipped(reshape): - continue - - x = tensor_map["x"] - y = tensor_map["y"] - model_outputs = self._builder.get_sub_graph().outputs.tmp_outputs - - # Replace `y` with `x` in the inputs of all following operators. - following_ops = input_to_ops.get(y.name, []) - for op in following_ops: - while y in op.tmp_inputs: - input_idx = op.tmp_inputs.index(y) - op.tmp_inputs[input_idx] = x - - if y in model_outputs: - # Replace the output as well. - while y in model_outputs: - idx = model_outputs.index(y) - model_outputs[idx] = x - - self._builder.swap_tensor_names(x, y) - - to_remove.append(reshape) - - for op in to_remove: - self._builder.get_operators().remove(op) - - return len(to_remove) != 0 diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizations/remove_unused_tensors_and_buffers.py b/backends/nxp/backend/ir/tflite_optimizer/optimizations/remove_unused_tensors_and_buffers.py deleted file mode 100755 index 105c06c3709..00000000000 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizations/remove_unused_tensors_and_buffers.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright 2024 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.nxp.backend.ir.tflite_generator import tflite_model -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.base_optimization import ( - BaseOptimization, -) - - -class RemoveUnusedTensorsAndBuffers(BaseOptimization): - - def _get_used_tensors_and_buffers( - self, - ) -> (set[tflite_model.Tensor], set[tflite_model.Buffer]): - """Get a set of all tensors used by the operators in the model, and a set of all buffers used by these tensors.""" - used_tensors = set() - used_buffers = set() - - for op in self._builder.get_operators(): - for tensor in op.tmp_inputs + op.tmp_outputs: - used_tensors.add(tensor) - if tensor.tmp_buffer is not None: - used_buffers.add(tensor.tmp_buffer) - - return used_tensors, used_buffers - - def __call__(self) -> bool: - """Remove all tensors and buffers from the model, that are not used. - :return: True, if any tensors/buffers were removed. Otherwise, False. - """ - - used_tensors, used_buffers = self._get_used_tensors_and_buffers() - - made_changes = False - model_inputs = self._builder.get_sub_graph().inputs.tmp_inputs - to_remove = [] - for tensor in self._builder.get_tensors(): - if tensor not in used_tensors: - if tensor in model_inputs: - # It is possible that an input tensor ended up not being used by any operators. But removing it from - # the model would cause errors at runtime, so it must stay. - pass - - else: - to_remove.append(tensor) - - for tensor in to_remove: - made_changes = True - self._builder.get_tensors().remove(tensor) - - to_remove = [] - for buffer in self._builder.get_buffers().vector: - if buffer not in used_buffers: - to_remove.append(buffer) - - for buffer in to_remove: - made_changes = True - self._builder.get_buffers().remove(buffer) - - return made_changes diff --git a/backends/nxp/backend/ir/tflite_optimizer/optimizer.py b/backends/nxp/backend/ir/tflite_optimizer/optimizer.py index 0d075c2cdaa..eb4ce6a5992 100755 --- a/backends/nxp/backend/ir/tflite_optimizer/optimizer.py +++ b/backends/nxp/backend/ir/tflite_optimizer/optimizer.py @@ -1,6 +1,6 @@ # # Copyright 2023 Martin Pavella -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # License: MIT # See the LICENSE_MIT for more details. @@ -14,18 +14,12 @@ from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.combine_hard_sigmoid_and_mul_to_hard_swish import ( CombineHardSigmoidAndMulIntoHardSwish, ) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.eliminate_dead_branches import ( - EliminateDeadBranches, -) from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.fuse_activation_functions import ( FuseActivationFunctions, ) from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.fuse_fully_connected_and_add_operators import ( FuseFullyConnectedAndAddOperators, ) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.fuse_quanitze_into_preceding_ops import ( - FuseQuantizeIntoPrecedingOps, -) from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.keep_one_empty_buffer import ( KeepOneEmptyBuffer, ) @@ -35,25 +29,10 @@ from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.permute_fully_connected_weights_after_reshape import ( PermuteFullyConnectedWeightsAfterReshape, ) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.prune_cast_operators import ( - FuseCastOperators, - RemoveCastOperatorsWithNoEffect, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.prune_quantize_operators import ( - FuseParallelQuantizeOperators, - PruneQuantizeOperators, -) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.prune_reshape_operators import ( - FuseReshapeOperators, - RemoveReshapeOperatorsWithNoEffect, -) from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.prune_transpose_operators import ( FuseTransposeOperators, RemoveIdentityTransposeOperators, ) -from executorch.backends.nxp.backend.ir.tflite_optimizer.optimizations.remove_unused_tensors_and_buffers import ( - RemoveUnusedTensorsAndBuffers, -) class Optimization(Enum): @@ -61,23 +40,11 @@ class Optimization(Enum): FUSE_ACTIVATION_FUNCTIONS = 1 FUSE_FULLY_CONNECTED_AND_ADD = 2 - FUSE_RESHAPE_OPERATORS = 3 - REMOVE_RESHAPE_OPERATORS_WITH_NO_EFFECT = 4 - FUSE_TRANSPOSE_OPERATORS = 5 REMOVE_IDENTITY_TRANSPOSE_OPERATORS = 6 - PRUNE_QUANTIZE_OPERATORS = 7 - FUSE_PARALLEL_QUANTIZE_OPERATORS = 8 - FUSE_QUANTIZE_INTO_PRECEDING_OPS = 9 - - REMOVE_UNUSED_TENSORS = 10 - ELIMINATE_DEAD_BRANCHES = 11 PERMUTE_FULLY_CONNECTED_WEIGHTS_AFTER_RESHAPE = 12 - FUSE_CAST_OPERATORS = 13 - REMOVE_CAST_OPERATORS_WITH_NO_EFFECT = 14 - MOVE_ACTIVATION_BEFORE_CONCAT = 15 COMBINE_HARD_SIGMOID_AND_MUL_INTO_HARD_SWISH = 16 @@ -118,42 +85,15 @@ def __init__( Optimization.FUSE_FULLY_CONNECTED_AND_ADD: FuseFullyConnectedAndAddOperators( builder, conversion_config ), - Optimization.FUSE_RESHAPE_OPERATORS: FuseReshapeOperators( - builder, conversion_config - ), - Optimization.REMOVE_RESHAPE_OPERATORS_WITH_NO_EFFECT: RemoveReshapeOperatorsWithNoEffect( - builder, conversion_config - ), Optimization.FUSE_TRANSPOSE_OPERATORS: FuseTransposeOperators( builder, conversion_config ), Optimization.REMOVE_IDENTITY_TRANSPOSE_OPERATORS: RemoveIdentityTransposeOperators( builder, conversion_config ), - Optimization.PRUNE_QUANTIZE_OPERATORS: PruneQuantizeOperators( - builder, conversion_config - ), - Optimization.FUSE_PARALLEL_QUANTIZE_OPERATORS: FuseParallelQuantizeOperators( - builder, conversion_config - ), - Optimization.FUSE_QUANTIZE_INTO_PRECEDING_OPS: FuseQuantizeIntoPrecedingOps( - builder, conversion_config - ), - Optimization.REMOVE_UNUSED_TENSORS: RemoveUnusedTensorsAndBuffers( - builder, conversion_config - ), - Optimization.ELIMINATE_DEAD_BRANCHES: EliminateDeadBranches( - builder, conversion_config - ), Optimization.PERMUTE_FULLY_CONNECTED_WEIGHTS_AFTER_RESHAPE: PermuteFullyConnectedWeightsAfterReshape( builder, conversion_config ), - Optimization.FUSE_CAST_OPERATORS: FuseCastOperators( - builder, conversion_config - ), - Optimization.REMOVE_CAST_OPERATORS_WITH_NO_EFFECT: RemoveCastOperatorsWithNoEffect( - builder, conversion_config - ), Optimization.MOVE_ACTIVATION_BEFORE_CONCAT: MoveActivationBeforeConcatenation( builder, conversion_config ), diff --git a/backends/nxp/backend/neutron_converter_manager.py b/backends/nxp/backend/neutron_converter_manager.py index 8826c60dc3b..99d7715ebf0 100644 --- a/backends/nxp/backend/neutron_converter_manager.py +++ b/backends/nxp/backend/neutron_converter_manager.py @@ -50,6 +50,8 @@ def convert( cctx = neutron_converter.CompilationContext() cctx.targetOpts = neutron_converter.getNeutronTarget(target) + # New switch since Neutron Converter SDK_25.06 + cctx.compilationOpts.minNumOpsPerGraph = 1 model_converted = neutron_converter.convertModel(list(tflite_model), cctx) return bytes(model_converted) 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 new file mode 100644 index 00000000000..d88684b86f0 --- /dev/null +++ b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py @@ -0,0 +1,231 @@ +# Copyright 2025 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.nxp.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 import Node +from torch.fx.passes.infra.pass_base import PassResult + +# Operator aliases for better readability. +AddMM = exir_ops.edge.aten.addmm.default +ViewCopy = exir_ops.edge.aten.view_copy.default +MM = exir_ops.edge.aten.mm.default + + +def insert_qdq_pair_after_node( + graph: torch.fx.Graph, anchor: torch.fx.Node, q_params: tuple +): + # Insert a Quantize node. + with graph.inserting_after(anchor): + quantize_op = graph.create_node( + op="call_function", + target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + args=(), # Will be added later. + ) + quantize_op.meta = anchor.meta + + # Insert a Dequantize 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, + args=(quantize_op,) + q_params, + ) + dequantize_op.meta = quantize_op.meta + anchor.replace_all_uses_with(dequantize_op) + + # Add this at the end, so the `anchor.replace_all_uses_with(dequantize_op)` does not replace the first use of the + # `quantize_op`. + quantize_op.args = (anchor,) + q_params + + +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 + ) + + +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 + ) + + +class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): + """ + │ + ┌─────▼──────┐ + │ │ dequantize │ + ┌─────▼──────┐ └─────┬──────┘ + │ dequantize │ ┌─────▼──────┐ + └─────┬──────┘ │ │ + ┌─────▼──────┐ └─────┬──────┘ + │ │ ┌────▼─────┐ ┐ + └─────┬──────┘ │ quantize │ │ + ┌──────────▼──────────┐ replaced with └────┬─────┘ │ + ⋯┤ ├⋯ ──────────────► │ │ newly added nodes + └──────────┬──────────┘ ┌─────▼──────┐ │ + ▼ │ dequantize │ │ + ⋮ └─────┬──────┘ ┘ + ┌────▼─────┐ ┌──────────▼──────────┐ + │ quantize │ ⋯┤ ├⋯ + └────┬─────┘ └──────────┬──────────┘ + ▼ ▼ + ⋮ + ┌────▼─────┐ + │ quantize │ + └────┬─────┘ + ▼ + """ + + # Dictionary mapping main cluster nodes to auxiliary nodes, for which this optimization will be applied. + main_cluster_node_to_auxiliary_nodes = { + AddMM: [ + ViewCopy, + ], + MM: [ + ViewCopy, + ], + } + + def run(self, graph_module: torch.fx.GraphModule) -> PassResult: + for aux_node in graph_module.graph.nodes: + if aux_node.op != "call_function": + continue + + dequantize_node = aux_node.args[0] + if not _is_dequantize(dequantize_node): + # Not the intended use case. + continue + + users = list(aux_node.users.keys()) + if len(users) != 1: + # Not the intended use case. + continue + + main_cluster_node = users[0] + if main_cluster_node.op != "call_function": + continue + + if aux_node.target not in self.main_cluster_node_to_auxiliary_nodes.get( + main_cluster_node.target, [] + ): + # Unsupported main cluster node and auxiliary node pair. + continue + + # Make sure the nodes are part of the same QDQ cluster. + cluster = QDQClusterRecognizer().get_qdq_cluster(main_cluster_node) + if any( + node_ not in cluster + for node_ in [dequantize_node, aux_node, main_cluster_node] + ): + continue + + # ---- The nodes follow the pattern described in the header. ---- + + q_params = dequantize_node.args[1:] + insert_qdq_pair_after_node(graph_module.graph, aux_node, q_params) + + # The graph has now changed, and we shouldn't keep iterating through it. Return the new graph and the parent + # class will call this pass again. + return PassResult(graph_module, True) + + # Nothing was changed. + return PassResult(graph_module, False) + + +class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): + """ + │ + ┌─────▼──────┐ + │ │ dequantize │ + ┌─────▼──────┐ └─────┬──────┘ + │ dequantize │ ⋮ + └─────┬──────┘ ┌──────────▼──────────┐ + ▼ ⋯┤ ├⋯ + ⋮ └──────────┬──────────┘ + ┌──────────▼──────────┐ replaced with ┌────▼─────┐ ┐ + ⋯┤ ├⋯ ──────────────► │ quantize │ │ + └──────────┬──────────┘ └────┬─────┘ │ + ┌─────▼──────┐ │ │ newly added nodes + │ │ ┌─────▼──────┐ │ + └─────┬──────┘ │ dequantize │ │ + ┌────▼─────┐ └─────┬──────┘ ┘ + │ quantize │ ┌─────▼──────┐ + └────┬─────┘ │ │ + ▼ └─────┬──────┘ + ┌────▼─────┐ + │ quantize │ + └────┬─────┘ + ▼ + """ + + # Dictionary mapping main cluster nodes to auxiliary nodes, for which this optimization will be applied. + main_cluster_node_to_auxiliary_nodes = { + AddMM: [ + ViewCopy, + ], + MM: [ + ViewCopy, + ], + } + + def run(self, graph_module: torch.fx.GraphModule) -> PassResult: + + for aux_node in graph_module.graph.nodes: + if aux_node.op != "call_function": + continue + + main_cluster_node = aux_node.args[0] + if not ( + hasattr(main_cluster_node, "op") + and main_cluster_node.op == "call_function" + ): + continue + + if aux_node.target not in self.main_cluster_node_to_auxiliary_nodes.get( + main_cluster_node.target, [] + ): + # Unsupported main cluster node and auxiliary node pair. + continue + + users = list(aux_node.users.keys()) + if len(users) != 1: + # Not the intended use case. + continue + + quantize_node = users[0] + if not _is_quantize(quantize_node): + # Not the intended use case. + continue + + # Make sure the nodes are part of the same QDQ cluster. + cluster = QDQClusterRecognizer().get_qdq_cluster(main_cluster_node) + if any( + node_ not in cluster + for node_ in [quantize_node, aux_node, main_cluster_node] + ): + continue + + # ---- The nodes follow the pattern described in the header. ---- + + q_params = quantize_node.args[1:] + insert_qdq_pair_after_node(graph_module.graph, main_cluster_node, q_params) + + # The graph has now changed, and we shouldn't keep iterating through it. Return the new graph and the parent + # class will call this pass again. + return PassResult(graph_module, True) + + # Nothing was changed. + return PassResult(graph_module, False) diff --git a/backends/nxp/edge_passes/neutron_edge_pass.py b/backends/nxp/edge_passes/neutron_edge_pass.py new file mode 100644 index 00000000000..8f77ce022fc --- /dev/null +++ b/backends/nxp/edge_passes/neutron_edge_pass.py @@ -0,0 +1,55 @@ +# Copyright 2025 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 abc import abstractmethod + +import torch + +from executorch.exir.pass_base import ExportPass +from torch.fx.passes.infra.pass_base import PassResult + + +class NeutronEdgePass(ExportPass): + """Abstract parent class for pre-processing passes on the edge dialect level.""" + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + """Call `self.run()` as long as changes are being made. After a pass modifies the graph, it cannot keep on + iterating through its nodes, and must return. This method allows the pass to go through the whole model. + """ + + # Every pass will return once it makes a change to the graph, to avoid traversing and modifying a graph at the + # same time. Therefore, it must be called multiple times (at most `iteration_limit` times). + iteration_limit = len(graph_module.graph.nodes) + modified = False + for _ in range(iteration_limit): + res = self.run(graph_module) + if res.modified: + modified = True + graph_module = res.graph_module + + else: + # No more changes have been made. + graph_module = self.recompile_module(graph_module) + return PassResult(graph_module, modified) + + # Iteration limit was reached. + logging.warning( + f"The NeutronEdgePass `{self.__class__.__name__}` reached the iteration limit." + ) + graph_module = self.recompile_module(graph_module) + return PassResult(graph_module, modified) + + @abstractmethod + def run(self, graph_module: torch.fx.GraphModule) -> PassResult: + """Child classes should implement their graph modification here.""" + pass + + def recompile_module( + self, graph_module: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + """Recompile the graph and re-trace the metadata. This should ensure that the datatypes and shapes are correct.""" + graph_module.recompile() + return super().call(graph_module).graph_module diff --git a/backends/nxp/edge_passes/neutron_edge_pass_manager.py b/backends/nxp/edge_passes/neutron_edge_pass_manager.py new file mode 100644 index 00000000000..ec46070ac31 --- /dev/null +++ b/backends/nxp/edge_passes/neutron_edge_pass_manager.py @@ -0,0 +1,89 @@ +# Copyright 2025 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 + +from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( + MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass, + MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass, +) +from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass +from executorch.exir import EdgeProgramManager +from executorch.exir.program._program import ( + _get_updated_graph_signature, + _get_updated_range_constraints, +) + +from torch import nn +from torch.export import ExportedProgram +from torch.fx.passes.infra.pass_base import PassResult +from torch.fx.passes.infra.pass_manager import PassManager + + +class NeutronEdgePassManager(PassManager): + + def __init__(self, passes: list[NeutronEdgePass] = None): + passes: list[NeutronEdgePass] = passes or [ + MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(), + MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(), + ] + + super().__init__( + passes, + steps=10, # Empirical value. At most 10 cycles of passes will be run. + ) + + def _transform_graph_module(self, module: nn.Module) -> PassResult: + """Apply the passes to a single graph module.""" + pass_result: PassResult = super().__call__(module) + + graph_module = pass_result.graph_module + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + + return pass_result + + def __call__(self, epm: EdgeProgramManager) -> EdgeProgramManager: + """Apply the passes to all graph modules in the edge program.""" + new_programs: dict[str, ExportedProgram] = {} + + for name, program in epm._edge_programs.items(): + pass_result = self._transform_graph_module(program.graph_module) + + if pass_result.modified: + # Create a new exported program. + new_program = ExportedProgram( + root=pass_result.graph_module, + graph=pass_result.graph_module.graph, + graph_signature=_get_updated_graph_signature( + program.graph_signature, pass_result.graph_module + ), + state_dict=program.state_dict, + range_constraints=_get_updated_range_constraints( + pass_result.graph_module + ), + module_call_graph=copy.deepcopy(program._module_call_graph), + example_inputs=program.example_inputs, + constants=program.constants, + verifiers=[program.verifier], + ) + new_program.graph_module.meta.update(program.graph_module.meta) + new_program.graph_module.meta.update(pass_result.graph_module.meta) + + else: + # Keep the old exported program. + new_program = program + + new_programs[name] = new_program + + if len(new_programs) == 0: + # No passes were run, return the old EdgeProgramManager. + return epm + + else: + # Return a new EdgeProgramManager with the updated programs. + return EdgeProgramManager( + new_programs, copy.deepcopy(epm._config_methods), epm.compile_config + ) diff --git a/backends/nxp/neutron_node_extraction.py b/backends/nxp/neutron_node_extraction.py index 2eb0f2d18c0..9d2431d29ed 100644 --- a/backends/nxp/neutron_node_extraction.py +++ b/backends/nxp/neutron_node_extraction.py @@ -45,7 +45,8 @@ def extract_artifacts_from_neutron_node( if sub_graph.OperatorsLength() == 0: raise RuntimeError( - "Model converted with neutron-converter has `0` operators instead of `1`." + "Model converted with neutron-converter has `0` operators instead of `1`.", + sub_graph.OperatorsLength(), ) elif sub_graph.OperatorsLength() > 1: builtin_operators_map: dict[int, str] = { @@ -61,7 +62,8 @@ def extract_artifacts_from_neutron_node( raise RuntimeError( f"Model converted with neutron-converter has `{sub_graph.OperatorsLength()}` operators " - f'instead of `1`. Operators found: {", ".join(ops_found)}.' + f'instead of `1`. Operators found: {", ".join(ops_found)}.', + sub_graph.OperatorsLength(), ) neutron_node = None diff --git a/backends/nxp/neutron_partitioner.py b/backends/nxp/neutron_partitioner.py index 952946ae26d..5bcdee0f8b6 100644 --- a/backends/nxp/neutron_partitioner.py +++ b/backends/nxp/neutron_partitioner.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024-2025 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -12,6 +12,9 @@ import torch +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) @@ -192,6 +195,7 @@ def tag_qdq_clusters(self, nodes: List[torch.fx.Node]): exir_ops.edge.aten.addmm.default: AddMMConverter, # noqa F405 exir_ops.edge.aten.add.Tensor: AddTensorConverter, # noqa F405 exir_ops.edge.aten.avg_pool2d.default: AvgPool2dConverter, # noqa F405 + exir_ops.edge.aten.cat.default: CatConverter, # noqa F405 exir_ops.edge.aten.clone.default: CloneConverter, # noqa F405 exir_ops.edge.aten.constant_pad_nd.default: ConstantPadNDConverter, # noqa F405 exir_ops.edge.aten.convolution.default: ConvolutionConverter, # noqa F405 @@ -202,6 +206,7 @@ def tag_qdq_clusters(self, nodes: List[torch.fx.Node]): exir_ops.edge.aten.mm.default: MMConverter, # noqa F405 exir_ops.edge.aten.relu.default: ReLUConverter, # noqa F405 exir_ops.edge.aten._softmax.default: SoftmaxConverter, # noqa F405 + exir_ops.edge.aten.tanh.default: TanhConverter, # noqa F405 exir_ops.edge.aten.view_copy.default: ViewCopyConverter, # noqa F405 exir_ops.edge.aten.sigmoid.default: SigmoidConverter, # noqa F405 } @@ -215,11 +220,13 @@ def __init__( target: Target, operators_not_to_delegate: List[str], parameters_mapping: dict[str, Parameter], + custom_delegation_options: CustomDelegationOptions, ): self.qdq_clusters = qdq_clusters self.target = target self.operators_not_to_delegate = operators_not_to_delegate self.parameters_mapping = parameters_mapping + self.custom_delegation_options = custom_delegation_options def _is_node_quantized(self, node: torch.fx.node.Node): return "cluster" in node.meta @@ -251,7 +258,12 @@ def _is_node_supported_compute(self, node: torch.fx.node.Node) -> bool: and self._is_node_quantized(node) and # TODO: `view_copy` node should be delegated only if it's not the only operator in the cluster. - node_converter.is_supported(node, self.target, self.parameters_mapping) + node_converter.is_supported( + node, + self.target, + self.parameters_mapping, + self.custom_delegation_options, + ) ) def _is_node_supported_non_compute(self, node: torch.fx.node.Node) -> bool: @@ -282,8 +294,15 @@ def is_node_supported( @final class NeutronPartitioner(Partitioner): - def __init__(self, compile_spec: List[CompileSpec]) -> None: + def __init__( + self, + compile_spec: List[CompileSpec], + custom_delegation_options: CustomDelegationOptions | None = None, + ) -> None: self.delegation_spec = DelegationSpec(NeutronBackend.__name__, compile_spec) + self.custom_delegation_options = ( + custom_delegation_options or CustomDelegationOptions() + ) def partition(self, exported_program: ExportedProgram) -> PartitionResult: # Run the CapabilityBasedPartitioner to return the largest possible @@ -318,6 +337,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: target, operators_not_to_delegate, parameters_mapping, + self.custom_delegation_options, ), allows_single_node_partition=True, ) diff --git a/backends/nxp/nxp_backend.py b/backends/nxp/nxp_backend.py index dd7d64227e3..c801eefec81 100644 --- a/backends/nxp/nxp_backend.py +++ b/backends/nxp/nxp_backend.py @@ -64,7 +64,7 @@ def neutron_compile_spec( Args: config: Neutron accelerator configuration, e.g. "imxrt700" neutron_converter_flavor: Flavor of the neutron-converter module to use. Neutron-converter module named " - "'neutron_converter_SDK_25_03' has flavor 'SDK_25_03'. + "'neutron_converter_SDK_25_06' has flavor 'SDK_25_06'. extra_flags: Extra flags for the Neutron compiler operators_not_to_delegate: List of operators that should not be delegated """ @@ -186,7 +186,7 @@ def preprocess( # noqa C901 # Convert the edge program to TFLite. tflite_model, io_formats = EdgeProgramToIRConverter().convert_program( - edge_program + edge_program, ) neutron_model = NeutronConverterManager().convert( diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 2279c177f59..d3f84144aa3 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -7,6 +7,7 @@ from typing import List, Optional, Tuple, Union import torch + from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( NeutronAtenPassManager, ) @@ -16,6 +17,7 @@ AddmmPattern, AddTensorPattern, AvgPoolPattern, + CatPattern, Conv1dPattern, Conv2dPattern, DropoutPattern, @@ -34,6 +36,8 @@ SharedSpecPattern, SigmoidPattern, SoftMaxPattern, + TanhInPlacePattern, + TanhPattern, ViewPattern, ) from executorch.backends.nxp.quantizer.utils import ( @@ -205,6 +209,7 @@ def __init__(self): NeutronAtenQuantizer(AddTensorPattern(), static_qconfig), NeutronAtenQuantizer(AddmmPattern(), static_fc_qconfig), NeutronAtenQuantizer(AvgPoolPattern(), static_qconfig), + NeutronAtenQuantizer(CatPattern(), static_qconfig), NeutronAtenQuantizer(Conv1dPattern(), static_qconfig), NeutronAtenQuantizer(Conv2dPattern(), static_qconfig), NeutronAtenQuantizer(DropoutPattern(), static_qconfig), @@ -221,6 +226,8 @@ def __init__(self): NeutronAtenQuantizer(ReshapePattern(), static_qconfig), NeutronAtenQuantizer(SigmoidPattern(), static_qconfig), NeutronAtenQuantizer(SoftMaxPattern(), static_qconfig), + NeutronAtenQuantizer(TanhPattern(), static_qconfig), + NeutronAtenQuantizer(TanhInPlacePattern(), static_qconfig), NeutronAtenQuantizer(ViewPattern(), static_qconfig), ] ) @@ -236,8 +243,13 @@ def __init__(self): def transform_for_annotation( self, model: torch.fx.GraphModule ) -> torch.fx.GraphModule: - pass_runner = NeutronAtenPassManager() - return pass_runner(model).graph_module + model.graph.eliminate_dead_code() # Remove dead code to simplify the graph for the passes. + + model = NeutronAtenPassManager()(model).graph_module + + model.graph.eliminate_dead_code() # Remove dead code again, in case it was created by the passes. + + return model def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule: self._annotate_inputs(model) diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 35649f0c0fc..651f995d570 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -106,6 +106,35 @@ def get_anchors( ) +def get_anchors_for_fixed_quant_specs( + fused_partition: list[fx.GraphModule], + scale: float, + zero_point: int, + quant_min: int = -128, + quant_max: int = 127, +) -> PartitionAnchors: + node = fused_partition[0].nodes[-1] + assert len(fused_partition[0].input_nodes) == 1 + + qspec = FixedQParamsQuantizationSpec( + dtype=torch.int8, + scale=scale, + zero_point=zero_point, + quant_min=quant_min, + quant_max=quant_max, + qscheme=torch.per_tensor_affine, + ) + + return PartitionAnchors( + inputs=[(node, 0)], + weights=[], + biases=[], + output=[ + (node, qspec), + ], + ) + + class AbsPattern(SharedSpecPattern): """ Quantizer for Abs operator. @@ -189,6 +218,47 @@ def partition_types(self): return [torch.ops.aten.avg_pool2d.default] +class CatPattern(QuantizationPattern): + """ + Quantizer for the Cat operator. The pattern is designed for the `NeutronAtenQuantizer`. + + The node can have an arbitrary number of inputs, which are all quantized. + """ + + def partition_types(self) -> list[OpOverload]: + return [torch.ops.aten.cat.default] + + def get_anchors( + self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] + ) -> PartitionAnchors | None: + node = fused_partition[0].nodes[-1] + + quantized_input = None + for prev_node in node.args[0]: + if "quantization_annotation" in prev_node.meta: + quantized_input = prev_node + break + + if quantized_input is not None: + inputs = [] + for idx, _ in enumerate(node.args[0]): + inputs.append((node, (0, idx), SharedQuantizationSpec(quantized_input))) + outputs = [(node, SharedQuantizationSpec(quantized_input))] + + else: + # No previous node was quantized => we are not able to share q-params. The conversion to IR will have to + # re-quantize the inputs if necessary. + inputs = [(node, (0, idx)) for idx in range(len(node.args[0]))] + outputs = [(node,)] + + return PartitionAnchors( + inputs=inputs, + weights=[], + biases=[], + output=outputs, + ) + + class Conv1dPattern(QuantizationPattern): def partition_types(self) -> List[OpOverload]: return [torch.ops.aten.conv1d.default] @@ -279,7 +349,7 @@ def partition_types(self): return [torch.ops.aten.flatten.using_ints] -class HardTanhPattern(SharedSpecPattern): +class HardTanhPattern(QuantizationPattern): """ Quantizer for HardTanh operator. Shared quantization spec is selected, as activation functions usually follows computation layer. @@ -288,8 +358,23 @@ class HardTanhPattern(SharedSpecPattern): def partition_types(self): return [torch.ops.aten.hardtanh.default] + def get_anchors( + self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule] + ) -> PartitionAnchors | None: + node = fused_partition[0].nodes[-1] -class HardTanhInPlacePattern(SharedSpecPattern): + return PartitionAnchors( + inputs=[(node, 0)], + weights=[], + biases=[], + output=[(node,)], + ) + + def replacement_op(self): + raise AssertionError() + + +class HardTanhInPlacePattern(QuantizationPattern): """ Quantizer for HardTanh operator with param inplace=True. Shared quantization spec is selected, as activation functions usually follows computation layer. @@ -298,6 +383,21 @@ class HardTanhInPlacePattern(SharedSpecPattern): def partition_types(self): return [torch.ops.aten.hardtanh_.default] + def get_anchors( + self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule] + ) -> PartitionAnchors | None: + node = fused_partition[0].nodes[-1] + + return PartitionAnchors( + inputs=[(node, 0)], + weights=[], + biases=[], + output=[(node,)], + ) + + def replacement_op(self): + raise AssertionError() + class LinearPattern(QuantizationPattern): def partition_types(self) -> List[OpOverload]: @@ -408,31 +508,6 @@ def partition_types(self): return [torch.ops.aten.view.default] -def get_anchors_for_softmax_like_operators( - fused_partition: List[fx.GraphModule], -) -> PartitionAnchors: - node = fused_partition[0].nodes[-1] - assert len(fused_partition[0].input_nodes) == 1 - - qspec = FixedQParamsQuantizationSpec( - dtype=torch.int8, - scale=1.0 / 256.0, - zero_point=-128, - quant_min=-128, - quant_max=127, - qscheme=torch.per_tensor_affine, - ) - - return PartitionAnchors( - inputs=[(node, 0)], - weights=[], - biases=[], - output=[ - (node, qspec), - ], - ) - - class SoftMaxPattern(QuantizationPattern): """ Quantizer for Softmax operator. @@ -444,9 +519,47 @@ def partition_types(self) -> List[OpOverload]: return [torch.ops.aten.softmax.int] def get_anchors( - self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule] + self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] ) -> PartitionAnchors: - return get_anchors_for_softmax_like_operators(fused_partition) + return get_anchors_for_fixed_quant_specs( + fused_partition, scale=1.0 / 256.0, zero_point=-128 + ) + + +class TanhPattern(QuantizationPattern): + """ + Quantizer for Tanh operator. + + The quantization of Tanh output is fixed to scale 1/128, zero point 0, dtype int8. + """ + + def partition_types(self): + return [torch.ops.aten.tanh.default] + + def get_anchors( + self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] + ) -> PartitionAnchors: + return get_anchors_for_fixed_quant_specs( + fused_partition, scale=1.0 / 128.0, zero_point=0 + ) + + +class TanhInPlacePattern(QuantizationPattern): + """ + Quantizer for inplace version of Tanh operator (torch.tanh_). + + The quantization of Tanh output is fixed to scale 1/128, zero point 0, dtype int8. + """ + + def partition_types(self): + return [torch.ops.aten.tanh_.default] + + def get_anchors( + self, gm: fx.GraphModule, fused_partition: list[fx.GraphModule] + ) -> PartitionAnchors: + return get_anchors_for_fixed_quant_specs( + fused_partition, scale=1.0 / 128.0, zero_point=0 + ) class SigmoidPattern(QuantizationPattern): @@ -462,4 +575,6 @@ def partition_types(self) -> List[OpOverload]: def get_anchors( self, gm: fx.GraphModule, fused_partition: List[fx.GraphModule] ) -> PartitionAnchors: - return get_anchors_for_softmax_like_operators(fused_partition) + return get_anchors_for_fixed_quant_specs( + fused_partition, scale=1.0 / 256.0, zero_point=-128 + ) diff --git a/backends/nxp/requirements-tests-eiq.txt b/backends/nxp/requirements-tests-eiq.txt new file mode 100644 index 00000000000..896d2b8c07e --- /dev/null +++ b/backends/nxp/requirements-tests-eiq.txt @@ -0,0 +1,2 @@ +--index-url https://eiq.nxp.com/repository +neutron_converter_SDK_25_06 diff --git a/backends/nxp/requirements-tests-pypi.txt b/backends/nxp/requirements-tests-pypi.txt new file mode 100644 index 00000000000..f31cf445420 --- /dev/null +++ b/backends/nxp/requirements-tests-pypi.txt @@ -0,0 +1,4 @@ +tensorflow==2.18.0 +pytest-mock +tflite +GvGen \ No newline at end of file diff --git a/backends/nxp/requirements-tests.txt b/backends/nxp/requirements-tests.txt deleted file mode 100644 index ea6d56a43ec..00000000000 --- a/backends/nxp/requirements-tests.txt +++ /dev/null @@ -1,6 +0,0 @@ ---extra-index-url https://eiq.nxp.com/repository -tensorflow==2.18.0 -pytest-mock -tflite -GvGen -neutron_converter_SDK_25_03 diff --git a/backends/nxp/run_unittests.sh b/backends/nxp/run_unittests.sh index f0a91e2a65d..78e35d2617a 100755 --- a/backends/nxp/run_unittests.sh +++ b/backends/nxp/run_unittests.sh @@ -12,3 +12,5 @@ cd $EXECUTORCH_DIR # '-c /dev/null' is used to ignore root level pytest.ini. pytest -c /dev/null backends/nxp/tests/ + +python -m unittest discover -s backends/nxp/tests/ -v diff --git a/backends/nxp/runtime/NeutronBackend.cpp b/backends/nxp/runtime/NeutronBackend.cpp index ef31054e933..3568ab72580 100644 --- a/backends/nxp/runtime/NeutronBackend.cpp +++ b/backends/nxp/runtime/NeutronBackend.cpp @@ -330,7 +330,7 @@ class NeutronBackend final : public PyTorchBackendInterface { Error execute( BackendExecutionContext& context, DelegateHandle* input_handle, - EValue** args) const override { + Span args) const override { NeutronConfig* cfg = static_cast(input_handle); // Allocate place for input and output pointers. diff --git a/backends/nxp/tests/TARGETS b/backends/nxp/tests/TARGETS index 1846423ffe9..f492111aff2 100644 --- a/backends/nxp/tests/TARGETS +++ b/backends/nxp/tests/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") oncall("executorch") -python_library( +runtime.python_library( name = "models", srcs = [ "models.py", @@ -13,7 +13,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "executorch_pipeline", srcs = [ "executorch_pipeline.py", @@ -21,9 +21,11 @@ python_library( ], deps = [ "//executorch/exir:lib", + "//executorch/extension/export_util:export_util", "//pytorch/ao:torchao", "//executorch/backends/nxp:quantizer", "//executorch/backends/nxp:neutron_backend", + "//executorch/backends/nxp:edge_passes", ] ) diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index 36ef76f8a2c..881bdeeec7b 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -3,12 +3,21 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from dataclasses import dataclass +from typing import Callable + import torch from executorch import exir +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import ( RemoveIOQuantOpsPass, ) +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) 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 @@ -17,13 +26,19 @@ EdgeProgramManager, ExecutorchBackendConfig, ExecutorchProgramManager, - to_edge_transform_and_lower, ) +from executorch.extension.export_util.utils import export_to_edge from torch import nn from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e -def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]): +@dataclass +class ModelInputSpec: + shape: tuple[int, ...] + dtype: torch.dtype = torch.float32 + + +def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor, ...]]): quantizer = NeutronQuantizer() m = prepare_pt2e(model, quantizer) @@ -34,52 +49,78 @@ def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]): return m +def get_random_calibration_inputs( + input_spec: tuple[ModelInputSpec, ...] +) -> list[tuple[torch.Tensor, ...]]: + return [ + tuple([torch.randn(spec.shape, dtype=spec.dtype) for spec in input_spec]) + for _ in range(4) + ] + + +def to_model_input_spec( + input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]] +) -> tuple[ModelInputSpec, ...]: + + if isinstance(input_spec, tuple) and all( + isinstance(spec, ModelInputSpec) for spec in input_spec + ): + return input_spec + + elif isinstance(input_spec, tuple) and all( + isinstance(spec, int) for spec in input_spec + ): + return (ModelInputSpec(input_spec),) + + elif isinstance(input_spec, list) and all( + isinstance(input_shape, tuple) for input_shape in input_spec + ): + return tuple([ModelInputSpec(spec) for spec in input_spec]) + else: + raise TypeError(f"Unsupported type {type(input_spec)}") + + def to_quantized_edge_program( model: torch.nn.Module, - input_shapes: tuple[int] | list[tuple[int]], + input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]], operators_not_to_delegate: list[str] = None, + get_calibration_inputs_fn: Callable[ + [tuple[ModelInputSpec, ...]], list[tuple[torch.Tensor, ...]] + ] = get_random_calibration_inputs, target="imxrt700", - neutron_converter_flavor="SDK_25_03", + neutron_converter_flavor="SDK_25_06", remove_quant_io_ops=False, + custom_delegation_options=CustomDelegationOptions(), # noqa B008 ) -> EdgeProgramManager: - if isinstance(input_shapes, list): - assert all(isinstance(input_shape, tuple) for input_shape in input_shapes), ( - "For multiple inputs, provide" " list[tuple[int]]." - ) + calibration_inputs = get_calibration_inputs_fn(to_model_input_spec(input_spec)) - random_tensors = ( - (torch.randn(input_shapes),) - if type(input_shapes) is tuple - else tuple(torch.randn(input_shape) for input_shape in input_shapes) - ) - calibration_inputs = [random_tensors, random_tensors] - example_input = ( - (torch.ones(input_shapes),) - if type(input_shapes) is tuple - else tuple(torch.ones(input_shape) for input_shape in input_shapes) - ) + example_input = calibration_inputs[0] - exir_program_aten = torch.export.export_for_training( - model, example_input, strict=True - ) + # Make sure the model is in the evaluation mode. + model.eval() + + exir_program_aten = torch.export.export(model, example_input, strict=True) exir_program_aten__module_quant = _quantize_model( exir_program_aten.module(), calibration_inputs ) + edge_compile_config = EdgeCompileConfig(_check_ir_validity=False) + edge_program_manager = export_to_edge( + exir_program_aten__module_quant, + example_input, + edge_compile_config=edge_compile_config, + ) + + edge_program_manager = NeutronEdgePassManager()(edge_program_manager) + compile_spec = generate_neutron_compile_spec( target, operators_not_to_delegate=operators_not_to_delegate, neutron_converter_flavor=neutron_converter_flavor, ) - partitioner = NeutronPartitioner(compile_spec) - edge_program_manager = to_edge_transform_and_lower( - torch.export.export( - exir_program_aten__module_quant, example_input, strict=True - ), - partitioner=[partitioner], - compile_config=EdgeCompileConfig(_check_ir_validity=False), - ) + partitioner = NeutronPartitioner(compile_spec, custom_delegation_options) + edge_program_manager = edge_program_manager.to_backend(partitioner) if remove_quant_io_ops: edge_program_manager = edge_program_manager.transform( @@ -90,9 +131,10 @@ def to_quantized_edge_program( def to_quantized_executorch_program( - model: torch.nn.Module, input_shapes: tuple[int] | list[tuple[int]] + model: torch.nn.Module, + input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]], ) -> ExecutorchProgramManager: - edge_program_manager = to_quantized_edge_program(model, input_shapes) + edge_program_manager = to_quantized_edge_program(model, input_spec) return edge_program_manager.to_executorch( config=ExecutorchBackendConfig(extract_delegate_segments=False) @@ -100,17 +142,15 @@ def to_quantized_executorch_program( def to_edge_program( - model: nn.Module, input_shapes: tuple[int] | list[tuple[int]] + model: nn.Module, + input_spec: tuple[ModelInputSpec, ...] | tuple[int, ...] | list[tuple[int, ...]], ) -> EdgeProgramManager: - if isinstance(input_shapes, list): - assert all(isinstance(input_shape, tuple) for input_shape in input_shapes), ( - "For multiple inputs, provide" " list[tuple[int]]." - ) + calibration_inputs = get_random_calibration_inputs(to_model_input_spec(input_spec)) + + example_input = calibration_inputs[0] + + # Make sure the model is in the evaluation mode. + model.eval() - example_input = ( - (torch.ones(input_shapes),) - if type(input_shapes) is tuple - else tuple(torch.ones(input_shape) for input_shape in input_shapes) - ) exir_program = torch.export.export(model, example_input) return exir.to_edge(exir_program) diff --git a/backends/nxp/tests/executors.py b/backends/nxp/tests/executors.py index fb1c1c4b4cb..afdb15af106 100644 --- a/backends/nxp/tests/executors.py +++ b/backends/nxp/tests/executors.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 warnings -from typing import Dict, Union +from typing import Callable, Dict, Union import numpy import numpy as np @@ -18,7 +18,12 @@ create_channels_first_to_channels_last_permutation, create_channels_last_to_channels_first_permutation, ) +from executorch.backends.nxp.backend.ir.converter.node_converter import ( + NodeConverter, + Target, +) from torch.export import ExportedProgram +from torch.fx import Node from torch.fx.graph import Graph @@ -52,6 +57,13 @@ def inference( return output.detach().numpy() elif isinstance(output, tuple) and len(output) == 1: return output[0].detach().numpy() + elif isinstance(output, tuple): + output_names = self.edge_program.graph_signature.user_outputs + + return { + name: tensor.detach().numpy() + for (name, tensor) in zip(output_names, output) + } raise RuntimeError( "Edge program inference with multiple outputs not implemented" @@ -356,16 +368,23 @@ def graph_contains_any_of_ops(graph: Graph, ops: list) -> bool: return any(node.target in ops for node in graph.nodes) -class OverrideSupportedTargets: +target_support_check_function = Callable[[Node, Target], bool] - def __init__(self, converter_class, *, new_targets): - self._converter_class = converter_class - self._new_targets = new_targets - self._old_targets = self._converter_class.supported_targets +class OverrideTargetSupportCheck: + + def __init__( + self, + converter_class: type[NodeConverter], + *, + new_target_support_check: target_support_check_function, + ): + self._converter_class = converter_class + self.new_target_support_check = new_target_support_check + self.old_target_support_check = converter_class._is_supported_on_target def __enter__(self): - self._converter_class.supported_targets = self._new_targets + self._converter_class._is_supported_on_target = self.new_target_support_check def __exit__(self, exc_type, exc_val, exc_tb): - self._converter_class.supported_targets = self._old_targets + self._converter_class._is_supported_on_target = self.old_target_support_check diff --git a/backends/nxp/tests/ir/__init__.py b/backends/nxp/tests/ir/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/nxp/tests/ir/converter/__init__.py b/backends/nxp/tests/ir/converter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/nxp/tests/ir/converter/node_converter/__init__.py b/backends/nxp/tests/ir/converter/node_converter/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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 8b6b63bb53f..bcdbd955c71 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 @@ -10,6 +10,12 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( + ModelBuilder, +) +from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( + BuiltinOperator, +) from executorch.backends.nxp.tests.executorch_pipeline import ( to_edge_program, to_quantized_edge_program, @@ -156,3 +162,49 @@ def test_avg_pool_2d_quant_conversion(mocker, input_shape, padding, count_includ tflite_output_preprocess=ToNCHWPreprocess(), input_data=input_data, ) + + +def test_avg_pool_2d_quant_conversion__padded(mocker): + input_shape = (1, 8, 8, 8) + model = AvgPool2dModule(True, 1) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + ops_spy = mocker.spy(ModelBuilder, "finish") + + # Run conversion + _ = to_quantized_edge_program(model, input_shape) + + # Capture the converter operators. + ops = ops_spy.spy_return.sub_graphs[0].operators.vector + + # Capture generated model + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + + # Capture converted program + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8) + + convert_run_compare( + exported_program, + tflite_input_preprocess=ToNHWCPreprocess(), + tfl_model=tflite_flatbuffers_model, + tflite_output_preprocess=ToNCHWPreprocess(), + input_data=input_data, + ) + + assert len(ops) == 2 + assert ops[0].builtin_options.operator_type == BuiltinOperator.PADV2 + assert ops[1].builtin_options.operator_type == BuiltinOperator.AVERAGE_POOL_2D + + # Make sure the padding used the `zero-point`. + pad_value = ops[0].tmp_inputs[2].tmp_buffer.data.item() + assert ( + pad_value == ops[0].tmp_inputs[0].quantization.zero_point[0] + ) # `Pad` input zp. + assert ( + pad_value == ops[0].tmp_outputs[0].quantization.zero_point[0] + ) # `Pad` output zp. + assert ( + pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0] + ) # `AvgPool` input zp. 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 new file mode 100644 index 00000000000..3df703f5bba --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py @@ -0,0 +1,292 @@ +# Copyright 2025 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import pytest +import torch + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.edge_program_converter import ( + EdgeProgramToIRConverter, +) +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import ( + convert_run_compare, + graph_contains_any_of_ops, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram + + +def _normalized_dim(dim, rank): + return dim if dim >= 0 else dim + rank + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +class CatModule(torch.nn.Module): + + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, *inputs: torch.Tensor): + return torch.cat(list(inputs), self.dim) + + +class CatConvModule(torch.nn.Module): + + def __init__(self, dim: int, channels: int = 4): + super().__init__() + self.dim = dim + self.conv = torch.nn.Conv2d(channels, channels, 2) + + def forward(self, *inputs: torch.Tensor): + x = torch.cat(list(inputs), self.dim) + return self.conv(x) + + +@pytest.mark.parametrize( + "rank, num_inputs, dim", + [ + pytest.param(2, 2, 1, id="2D, 2 inputs, dim=1"), + pytest.param(2, 2, -1, id="2D, 2 inputs, dim=-1"), + pytest.param(2, 3, 1, id="2D, 3 inputs, dim=1"), + pytest.param(2, 3, -1, id="2D, 3 inputs, dim=-1"), + pytest.param(2, 4, -1, id="2D, 4 inputs, dim=-1"), + pytest.param(3, 2, 1, id="3D, 2 inputs, dim=1"), + pytest.param(3, 2, -1, id="3D, 2 inputs, dim=-1"), + pytest.param(3, 5, -1, id="3D, 5 inputs, dim=-2"), + pytest.param(4, 2, -1, id="4D, 2 inputs, dim=-1"), + pytest.param(4, 3, 2, id="4D, 3 inputs, dim=2"), + pytest.param(4, 5, -3, id="4D, 5 inputs, dim=-3"), + ], +) +def test_cat__same_shapes(dim, num_inputs, rank, mocker): + input_shape = tuple([2, 8, 8, 8, 8][-rank:]) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + quantized_program = to_quantized_edge_program( + CatModule(dim), [input_shape] * num_inputs + ).exported_program() + + # Make sure the `Cat` was delegated. + assert not graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert any("lowered_module" in node.name for node in quantized_program.graph.nodes) + + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + exported_program: ExportedProgram = converter_spy.call_args.args[1] + input_data = { + i: (np.random.random(input_shape) * 50).astype(np.int8) + for i in range(num_inputs) + } + convert_run_compare( + exported_program, + tfl_model=tflite_flatbuffers_model, + input_data=input_data, + atol=1, + ) + + +@pytest.mark.parametrize("dim", [3, -2, -3]) +@pytest.mark.parametrize("num_inputs", [2, 5]) +def test_cat__channels_first__same_shapes(dim, num_inputs, mocker): + input_shape = (2, 8, 6, 8) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + channels = input_shape[1] if dim not in {1, -3} else input_shape[1] * num_inputs + quantized_program = to_quantized_edge_program( + CatConvModule(dim, channels), [input_shape] * num_inputs + ).exported_program() + + # Make sure the `Cat` was delegated. + assert not graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert any("lowered_module" in node.name for node in quantized_program.graph.nodes) + + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + exported_program: ExportedProgram = converter_spy.call_args.args[1] + input_data = { + i: (np.random.random(input_shape) * 50).astype(np.int8) + for i in range(num_inputs) + } + convert_run_compare( + exported_program, + tfl_model=tflite_flatbuffers_model, + input_data=input_data, + atol=1, + ) + + +@pytest.mark.parametrize("dim", [0, -4]) +@pytest.mark.parametrize("num_inputs", [2]) +def test_cat__unsupported_dim__imxrt700(dim, num_inputs): + input_shape = (2, 8, 6, 8) + + quantized_program = to_quantized_edge_program( + CatModule(dim), [input_shape] * num_inputs, target="imxrt700" + ).exported_program() + + # Make sure the `Cat` was NOT delegated. + assert graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert not any( + "lowered_module" in node.name for node in quantized_program.graph.nodes + ) + + +@pytest.mark.parametrize( + "rank, num_inputs, dim", + [ + pytest.param(2, 2, 1, id="2D, 2 inputs, dim=1"), + pytest.param(2, 2, -1, id="2D, 2 inputs, dim=-1"), + pytest.param(2, 3, 1, id="2D, 3 inputs, dim=1"), + pytest.param(2, 3, -1, id="2D, 3 inputs, dim=-1"), + pytest.param(2, 4, -1, id="2D, 4 inputs, dim=-1"), + pytest.param(3, 2, 1, id="3D, 2 inputs, dim=1"), + pytest.param(3, 2, -1, id="3D, 2 inputs, dim=-1"), + pytest.param(3, 5, -1, id="3D, 5 inputs, dim=-2"), + pytest.param(4, 2, -1, id="4D, 2 inputs, dim=-1"), + pytest.param(4, 3, 2, id="4D, 3 inputs, dim=2"), + pytest.param(4, 5, -3, id="4D, 5 inputs, dim=-3"), + ], +) +def test_cat__different_shapes(dim, num_inputs, rank, mocker): + input_shape = tuple([2, 8, 8, 8, 8][-rank:]) + + # The shape of every input will be different along the concatenated dimension. + input_shapes = [] + for i in range(num_inputs): + tmp_shape = list(input_shape) + tmp_shape[dim] = 8 * (i + 1) # RT700 requires multiples of 8 for the channels. + input_shapes.append(tuple(tmp_shape)) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + quantized_program = to_quantized_edge_program( + CatModule(dim), input_shapes + ).exported_program() + + # Make sure the `Cat` was delegated. + assert not graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert any("lowered_module" in node.name for node in quantized_program.graph.nodes) + + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + exported_program: ExportedProgram = converter_spy.call_args.args[1] + input_data = { + i: (np.random.random(shape) * 50).astype(np.int8) + for i, shape in enumerate(input_shapes) + } + convert_run_compare( + exported_program, + tfl_model=tflite_flatbuffers_model, + input_data=input_data, + atol=1, + ) + + +@pytest.mark.parametrize("dim", [1, -1, -2], ids=lambda dim: f"dim = {dim}") +@pytest.mark.parametrize( + "num_inputs", [2, 5], ids=lambda num_inputs: f"num_inputs = {num_inputs}" +) +def test_cat__channels_first__different_shapes(dim, num_inputs, mocker): + input_shape = (2, 8, 6, 8) + + # The shape of every input will be different along the concatenated dimension. + input_shapes = [] + for i in range(num_inputs): + tmp_shape = list(input_shape) + tmp_shape[dim] = 8 * ( + i + 1 + ) # Neutron only supports channels that are multiples of 8 (on RT700). + input_shapes.append(tuple(tmp_shape)) + + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + + channels = ( + sum(shape[1] for shape in input_shapes) if dim in [1, -3] else input_shape[1] + ) + quantized_program = to_quantized_edge_program( + CatConvModule(dim, channels), input_shapes + ).exported_program() + + # Make sure the `Cat` was delegated. + assert not graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert any("lowered_module" in node.name for node in quantized_program.graph.nodes) + + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + exported_program: ExportedProgram = converter_spy.call_args.args[1] + input_data = { + i: (np.random.random(shape) * 50).astype(np.int8) + for i, shape in enumerate(input_shapes) + } + convert_run_compare( + exported_program, + tfl_model=tflite_flatbuffers_model, + input_data=input_data, + atol=1, + ) + + +def test_cat__different_shapes__unsupported_channels__imxrt700(): + input_shape = (2, 4, 6, 7) # (channels % 8) != 0 + + num_inputs = 2 + dim = -1 + + # The shape of every input will be different along the concatenated dimension. + input_shapes = [] + for i in range(num_inputs): + tmp_shape = list(input_shape) + tmp_shape[dim] = i + 2 + input_shapes.append(tuple(tmp_shape)) + + quantized_program = to_quantized_edge_program( + CatModule(dim), input_shapes, target="imxrt700" + ).exported_program() + + # Make sure the `Cat` was NOT delegated. + assert graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert not any( + "lowered_module" in node.name for node in quantized_program.graph.nodes + ) + + +def test_cat__force_delegate(): + target = "imxrt700" + + # The Partitioner doesn't know if the `8` or the `1` will become the channels in the IR. Therefore, it would + # normally not delegate the `cat`. But we know that the `8` will be the channels, so we can force the delegation. + input_shape = (8, 1, 8) + + quantized_program = to_quantized_edge_program( + CatModule(1), + [input_shape, input_shape], + target=target, + custom_delegation_options=CustomDelegationOptions(force_delegate_cat=True), + ).exported_program() + + # Make sure the `Cat` was delegated. + assert not graph_contains_any_of_ops( + graph=quantized_program.graph, ops=[exir_ops.edge.aten.cat.default] + ) + assert any("lowered_module" in node.name for node in quantized_program.graph.nodes) 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 8194bf3cb8c..745b26ef8ff 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 @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -22,10 +22,10 @@ ) from executorch.backends.nxp.tests.executors import ( convert_run_compare, - ToNCHWPreprocess, - ToNHWCPreprocess, + ToChannelFirstPreprocess, + ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.models import Conv2dModule +from executorch.backends.nxp.tests.models import Conv1dModule, Conv2dModule from torch.export import ExportedProgram @@ -35,48 +35,209 @@ def reseed_model_per_test_run(): np.random.seed(23) -@pytest.mark.parametrize( - "input_shape, padding", - [ - pytest.param((1, 4, 32, 32), (0, 0), id="No padding."), - pytest.param( - (1, 4, 32, 32), - (1, 1), - id="Padding, keep the same output tensor size as input.", - ), - pytest.param( - (1, 4, 32, 32), (1, 0), id="Padding, change the output tensor size." - ), - pytest.param( - (1, 4, 31, 31), (1, 0), id="Padding, change the output tensor size." - ), - pytest.param( - (1, 4, 31, 31), (0, 1), id="Padding, change the output tensor size." - ), - ], -) -@pytest.mark.parametrize( - "dilation", - [ - pytest.param(1, id="No dilation."), - pytest.param(2, id="2 dilation."), - pytest.param((1, 3), id="Side-different dilation."), - ], -) -def test_conv2d_conversion(input_shape, padding, dilation: int): - edge_program = to_edge_program( - Conv2dModule(padding=padding, dilation=dilation), input_shape - ).exported_program() +@pytest.mark.parametrize("stride", [1, 2]) +@pytest.mark.parametrize("dilation", [2, 1]) +@pytest.mark.parametrize("kernel_size", [(1,), (3,)]) +def test_conv1d_quant_conversion(stride, dilation, kernel_size, mocker): + input_shape = (1, 4, 16) + model = Conv1dModule(stride=stride, dilation=dilation, kernel_size=kernel_size) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + ops_spy = mocker.spy(ModelBuilder, "finish") - input_data = np.random.random(input_shape).astype(np.float32) + # Run conversion + _ = to_quantized_edge_program(model, input_shape) + + # Capture generated model + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + + # Capture converted program + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8) convert_run_compare( - edge_program, - input_data, - tflite_input_preprocess=ToNHWCPreprocess(), - tflite_output_preprocess=ToNCHWPreprocess(), - atol=4e-7, + exported_program, + tflite_input_preprocess=ToChannelLastPreprocess(), + tfl_model=tflite_flatbuffers_model, + tflite_output_preprocess=ToChannelFirstPreprocess(), + input_data=input_data, + atol=1.0, + ) + + # Capture IR model ops + conversion_result = ops_spy.spy_return + ops = conversion_result.sub_graphs[0].operators.vector + + assert len(ops) == 3 + assert ops[0].builtin_options.operator_type == BuiltinOperator.RESHAPE + assert ops[1].builtin_options.operator_type == BuiltinOperator.CONV_2D + assert ops[2].builtin_options.operator_type == BuiltinOperator.RESHAPE + + +@pytest.mark.parametrize("stride", [1, 2]) +@pytest.mark.parametrize("dilation", [2, 1]) +@pytest.mark.parametrize("kernel_size", [(1,), (3,)]) +@pytest.mark.parametrize("padding", [(1,), 2]) +def test_conv1d_quant_conversion__padded( + stride, dilation, kernel_size, padding, mocker +): + input_shape = (1, 4, 16) + model = Conv1dModule( + stride=stride, dilation=dilation, kernel_size=kernel_size, padding=padding ) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + ops_spy = mocker.spy(ModelBuilder, "finish") + + # Run conversion + _ = to_quantized_edge_program(model, input_shape) + + # Capture generated model + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + + # Capture converted program + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8) + + convert_run_compare( + exported_program, + tflite_input_preprocess=ToChannelLastPreprocess(), + tfl_model=tflite_flatbuffers_model, + tflite_output_preprocess=ToChannelFirstPreprocess(), + input_data=input_data, + atol=1.0, + ) + + # Capture IR model ops + conversion_result = ops_spy.spy_return + ops = conversion_result.sub_graphs[0].operators.vector + + assert len(ops) == 4 + assert ops[0].builtin_options.operator_type == BuiltinOperator.RESHAPE + assert ops[1].builtin_options.operator_type == BuiltinOperator.PADV2 + assert ops[2].builtin_options.operator_type == BuiltinOperator.CONV_2D + assert ops[3].builtin_options.operator_type == BuiltinOperator.RESHAPE + + # Make sure the padding used the `zero-point`. + pad_value = ops[1].tmp_inputs[2].tmp_buffer.data.item() + assert ( + pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0] + ) # `Pad` input zp. + assert ( + pad_value == ops[1].tmp_outputs[0].quantization.zero_point[0] + ) # `Pad` output zp. + assert ( + pad_value == ops[2].tmp_inputs[0].quantization.zero_point[0] + ) # `Conv` input zp. + + +@pytest.mark.parametrize("stride", [1, 2]) +@pytest.mark.parametrize("dilation", [2, 1]) +@pytest.mark.parametrize("kernel_size", [(1,), (3,)]) +def test_conv1d_quant_conversion__depthwise(stride, dilation, kernel_size, mocker): + input_shape = (1, 4, 16) + group = input_shape[1] + model = Conv1dModule( + group=group, + in_channels=group, + out_channels=group, + stride=stride, + dilation=dilation, + kernel_size=kernel_size, + ) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + ops_spy = mocker.spy(ModelBuilder, "finish") + + # Run conversion + _ = to_quantized_edge_program(model, input_shape) + + # Capture generated model + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + + # Capture converted program + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8) + + convert_run_compare( + exported_program, + tflite_input_preprocess=ToChannelLastPreprocess(), + tfl_model=tflite_flatbuffers_model, + tflite_output_preprocess=ToChannelFirstPreprocess(), + input_data=input_data, + atol=1.0, + ) + + # Capture IR model ops + ops = ops_spy.spy_return.sub_graphs[0].operators.vector + + assert len(ops) == 3 + assert ops[0].builtin_options.operator_type == BuiltinOperator.RESHAPE + assert ops[1].builtin_options.operator_type == BuiltinOperator.DEPTHWISE_CONV_2D + assert ops[2].builtin_options.operator_type == BuiltinOperator.RESHAPE + + +@pytest.mark.parametrize("stride", [1, 2]) +@pytest.mark.parametrize("dilation", [2, 1]) +@pytest.mark.parametrize("kernel_size", [(1,), (3,)]) +@pytest.mark.parametrize("padding", [(1,), 2]) +def test_conv1d_quant_conversion__depthwise__padded( + stride, dilation, kernel_size, padding, mocker +): + input_shape = (1, 4, 16) + group = input_shape[1] + model = Conv1dModule( + group=group, + in_channels=group, + out_channels=group, + stride=stride, + dilation=dilation, + kernel_size=kernel_size, + padding=padding, + ) + converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") + ops_spy = mocker.spy(ModelBuilder, "finish") + + # Run conversion + _ = to_quantized_edge_program(model, input_shape) + + # Capture generated model + tflite_flatbuffers_model, io_formats = converter_spy.spy_return + + # Capture converted program + exported_program: ExportedProgram = converter_spy.call_args.args[1] + + input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8) + + convert_run_compare( + exported_program, + tflite_input_preprocess=ToChannelLastPreprocess(), + tfl_model=tflite_flatbuffers_model, + tflite_output_preprocess=ToChannelFirstPreprocess(), + input_data=input_data, + atol=1.0, + ) + + # Capture IR model ops + ops = ops_spy.spy_return.sub_graphs[0].operators.vector + + assert len(ops) == 4 + assert ops[0].builtin_options.operator_type == BuiltinOperator.RESHAPE + assert ops[1].builtin_options.operator_type == BuiltinOperator.PADV2 + assert ops[2].builtin_options.operator_type == BuiltinOperator.DEPTHWISE_CONV_2D + assert ops[3].builtin_options.operator_type == BuiltinOperator.RESHAPE + + # Make sure the padding used the `zero-point`. + pad_value = ops[1].tmp_inputs[2].tmp_buffer.data.item() + assert ( + pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0] + ) # `Pad` input zp. + assert ( + pad_value == ops[1].tmp_outputs[0].quantization.zero_point[0] + ) # `Pad` output zp. + assert ( + pad_value == ops[2].tmp_inputs[0].quantization.zero_point[0] + ) # `Conv` input zp. @pytest.mark.parametrize( @@ -204,9 +365,9 @@ def test_conv2d_quant_conversion(mocker, model: torch.nn.Module, input_shape): convert_run_compare( exported_program, - tflite_input_preprocess=ToNHWCPreprocess(), + tflite_input_preprocess=ToChannelLastPreprocess(), tfl_model=tflite_flatbuffers_model, - tflite_output_preprocess=ToNCHWPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), input_data=input_data, atol=1.0, ) @@ -237,8 +398,8 @@ def test_conv2d_conversion__depthwise(stride, dilation, kernel_shape, mocker): convert_run_compare( edge_program, input_data, - tflite_input_preprocess=ToNHWCPreprocess(), - tflite_output_preprocess=ToNCHWPreprocess(), + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), atol=4e-7, ) conversion_result = spy.spy_return @@ -299,8 +460,8 @@ def test_conv2d_conversion__depthwise__padded(padding, mocker): convert_run_compare( edge_program, input_data, - tflite_input_preprocess=ToNHWCPreprocess(), - tflite_output_preprocess=ToNCHWPreprocess(), + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), atol=4e-7, ) conversion_result = spy.spy_return @@ -326,7 +487,7 @@ def test_conv2d_conversion__depthwise__padded__quantized(padding, mocker): ops = spy.spy_return.sub_graphs[0].operators.vector assert len(ops) == 2 - assert ops[0].builtin_options.operator_type == BuiltinOperator.PAD + assert ops[0].builtin_options.operator_type == BuiltinOperator.PADV2 assert ops[1].builtin_options.operator_type == BuiltinOperator.DEPTHWISE_CONV_2D nodes = list(edge_program.graph.nodes) @@ -334,184 +495,3 @@ def test_conv2d_conversion__depthwise__padded__quantized(padding, mocker): len(nodes) == 7 ) # input, Quant, lowered_module, delegate_call, getitem, Deq, output assert nodes[2].target == "lowered_module_0" - - -@pytest.mark.parametrize("stride", [1, 2]) -@pytest.mark.parametrize("dilation", [1, 2]) -@pytest.mark.parametrize( - "input_shape, group, out_channels", - [((1, 4, 12, 12), 2, 2), ((2, 3, 8, 15), 3, 6), ((11, 16, 9, 8), 4, 16)], -) -def test_conv2d_conversion__separated( - input_shape, group, out_channels, stride, dilation -): - edge_program = to_edge_program( - Conv2dModule( - group=group, - in_channels=input_shape[1], - out_channels=out_channels, - stride=stride, - dilation=dilation, - ), - input_shape, - ).exported_program() - - input_data = np.random.random(input_shape).astype(np.float32) - - # Note: The generic group convolution is not yet supported by Neutron Converter. Once supported, the - # commented out code allows usual testing flow for this test-case. - - # spy = mocker.spy(ModelBuilder, 'finish') - - # The convert_run_compare skips the partitioner call, hence conversion failure indicated by exception - # is expected behavior now. - with pytest.raises(AssertionError) as e: - convert_run_compare( - edge_program, - input_data, - tflite_input_preprocess=ToNHWCPreprocess(), - tflite_output_preprocess=ToNCHWPreprocess(), - atol=3.0e-7, - ) - assert ( - "`aten_convolution_default` is not convertible to the intermediate representation" - in str(e) - ) - - # ops = spy.spy_return.sub_graphs[0].operators.vector - # assert len(ops) == 1 + group + 1 # Split -> Conv (group times) -> Concat - # assert ops[0].builtin_options.operator_type == BuiltinOperator.SPLIT - # for op in ops[1:-1]: - # assert op.builtin_options.operator_type == BuiltinOperator.CONV_2D - # assert ops[-1].builtin_options.operator_type == BuiltinOperator.CONCATENATION - - -@pytest.mark.parametrize("stride", [1, 2]) -@pytest.mark.parametrize("dilation", [1, 2]) -@pytest.mark.parametrize( - "input_shape, group, out_channels", - [((1, 4, 12, 12), 2, 2), ((2, 3, 17, 9), 3, 6), ((11, 16, 9, 8), 4, 16)], -) -def test_conv2d_conversion__separated__quantized( - input_shape, group, out_channels, stride, dilation -): - - # Note: The generic group convolution is not yet supported by Neutron Converter. Once supported, the - # commented out code allows usuall testing flow for this test-case. - # spy = mocker.spy(ModelBuilder, 'finish') - - # The convert_run_compare skips the partitioner call, hence conversion failure indicated by exception - # is expected behavior now. - edge_program = to_quantized_edge_program( - Conv2dModule( - group=group, - in_channels=input_shape[1], - out_channels=out_channels, - stride=stride, - dilation=dilation, - ), - tuple(input_shape), - ).exported_program() - - # ops = spy.spy_return.sub_graphs[0].operators.vector - # assert len(ops) == 1 + group + 1 # Split -> Conv (group times) -> Concat - # assert ops[0].builtin_options.operator_type == BuiltinOperator.SPLIT - # for op in ops[1:-1]: - # assert op.builtin_options.operator_type == BuiltinOperator.CONV_2D - # assert ops[-1].builtin_options.operator_type == BuiltinOperator.CONCATENATION - - nodes = list(edge_program.graph.nodes) - assert len(nodes) == 11 - assert ( - nodes[7].target.__name__ == "aten.convolution.default" - ) # Convolution not delegated. - - -@pytest.mark.parametrize("padding", [1, 2]) -@pytest.mark.parametrize( - "input_shape, group, out_channels", - [((1, 4, 12, 12), 2, 2), ((2, 3, 4, 5), 3, 6), ((11, 16, 9, 8), 4, 16)], -) -def test_conv2d_conversion__separated__padded( - input_shape, group, out_channels, padding -): - edge_program = to_edge_program( - Conv2dModule( - group=group, - in_channels=input_shape[1], - out_channels=out_channels, - padding=padding, - ), - input_shape, - ).exported_program() - - input_data = np.random.random(input_shape).astype(np.float32) - - # Note: The generic group convolution is not yet supported by Neutron Converter. Once supported, the - # commented out code allows usuall testing flow for this test-case. - - # spy = mocker.spy(ModelBuilder, 'finish') - - # The convert_run_compare skips the partitioner call, hence conversion failure indicated by exception - # is expected behavior now. - with pytest.raises(AssertionError) as e: - convert_run_compare( - edge_program, - input_data, - tflite_input_preprocess=ToNHWCPreprocess(), - tflite_output_preprocess=ToNCHWPreprocess(), - atol=3.0e-7, - ) - assert ( - "`aten_convolution_default` is not convertible to the intermediate representation" - in str(e) - ) - - # conversion_result = spy.spy_return - # ops = conversion_result.sub_graphs[0].operators.vector - # assert len(ops) == 1 + 2 * group + 1 # Split -> Pad + Conv (group times) -> Concat - # assert ops[0].builtin_options.operator_type == BuiltinOperator.SPLIT - # for op in ops[1:-2:2]: - # assert op.builtin_options.operator_type == BuiltinOperator.PAD - # for op in ops[2:-1:2]: - # assert op.builtin_options.operator_type == BuiltinOperator.CONV_2D - # assert ops[-1].builtin_options.operator_type == BuiltinOperator.CONCATENATION - - -@pytest.mark.parametrize("padding", [1, 2]) -@pytest.mark.parametrize( - "input_shape, group, out_channels", - [((1, 4, 12, 12), 2, 2), ((2, 3, 4, 5), 3, 6), ((11, 16, 9, 8), 4, 16)], -) -def test_conv2d_conversion__separated__padded__quantized( - input_shape, group, out_channels, padding -): - - # Note: The generic group convolution is not yet supported by Neutron Converter. Once supported, the - # commented out code allows usuall testing flow for this test-case. - # spy = mocker.spy(ModelBuilder, 'finish') - - edge_program = to_quantized_edge_program( - Conv2dModule( - group=group, - in_channels=input_shape[1], - out_channels=out_channels, - padding=padding, - ), - tuple(input_shape), - ).exported_program() - - # ops = spy.spy_return.sub_graphs[0].operators.vector - # assert len(ops) == 1 + 2 * group + 1 # Split -> Pad + Conv (group times) -> Concat - # assert ops[0].builtin_options.operator_type == BuiltinOperator.SPLIT - # for op in ops[1:-2:2]: - # assert op.builtin_options.operator_type == BuiltinOperator.PAD - # for op in ops[2:-1:2]: - # assert op.builtin_options.operator_type == BuiltinOperator.CONV_2D - # assert ops[-1].builtin_options.operator_type == BuiltinOperator.CONCATENATION - - nodes = list(edge_program.graph.nodes) - assert len(nodes) == 11 - assert ( - nodes[7].target.__name__ == "aten.convolution.default" - ) # Convolution not delegated. 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 f90118f4bed..e17868d16e2 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 @@ -1,3 +1,8 @@ +# Copyright 2025 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + import numpy as np import pytest import torch @@ -15,6 +20,7 @@ ToNCHWPreprocess, ToNHWCPreprocess, ) +from executorch.backends.nxp.tests.models import Conv2dWithActivation from executorch.exir.dialects._ops import ops as exir_ops from torch.export import ExportedProgram @@ -25,48 +31,14 @@ def reseed_model_per_test_run(): np.random.seed(23) -class Relu6ConvBlock(torch.nn.Module): - def __init__(self, conv_in_channels: int = 3, inplace: bool = False): - super().__init__() - self.block = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4) - ), - torch.nn.ReLU6(inplace=inplace), - ) - - def forward(self, x): - return self.block(x) - - -class CustomHardTanhBlock(torch.nn.Module): - def __init__( - self, - conv_in_channels: int = 3, - min_act_val: float = -1.0, - max_act_val: float = 1.0, - inplace: bool = False, - ): - super().__init__() - self.block = torch.nn.Sequential( - torch.nn.Conv2d( - in_channels=conv_in_channels, out_channels=64, kernel_size=(4, 4) - ), - torch.nn.Hardtanh( - min_val=min_act_val, max_val=max_act_val, inplace=inplace - ), - ) - - def forward(self, x): - return self.block(x) - - -@pytest.mark.parametrize("input_shape", [(1, 3, 128, 128), (1, 3, 256, 256)]) +@pytest.mark.parametrize("input_shape", [(1, 3, 128, 128)]) @pytest.mark.parametrize("inplace", [True, False]) def test_relu6_quant(mocker, input_shape: tuple[int], inplace: bool): # The torch.nn.Relu6 inherits from torch.nn.Hardtanh, and hence represented as HardTanh in ATen. # Testing the hardtanh originated from torch.nn.Relu6 op. - model = Relu6ConvBlock(conv_in_channels=input_shape[1], inplace=inplace) + model = Conv2dWithActivation( + activation=torch.nn.ReLU6(inplace=inplace), in_channels=input_shape[1] + ) converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") @@ -89,7 +61,7 @@ def test_relu6_quant(mocker, input_shape: tuple[int], inplace: bool): ) -@pytest.mark.parametrize("input_shape", [(1, 3, 128, 128), (1, 3, 256, 256)]) +@pytest.mark.parametrize("input_shape", [(1, 3, 16, 16), (1, 3, 32, 32)]) @pytest.mark.parametrize( "activation_range", list(HardTanhConverter.supported_modes_map.keys()) ) @@ -97,12 +69,12 @@ def test_relu6_quant(mocker, input_shape: tuple[int], inplace: bool): def test_custom_hardtanh_quant( mocker, input_shape: tuple[int], activation_range: tuple[int, int], inplace: bool ): + # TODO(13063): This test suffers from non-ideal testing random quantization, because we always use range <0,1>. + # We should update (decrease atol) when the Conv/Linear + Activation fuse at quantization is in place. min_val, max_val = activation_range - model = CustomHardTanhBlock( - conv_in_channels=input_shape[1], - min_act_val=min_val, - max_act_val=max_val, - inplace=inplace, + model = Conv2dWithActivation( + activation=torch.nn.Hardtanh(min_val=min_val, max_val=max_val, inplace=inplace), + in_channels=input_shape[1], ) converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program") @@ -122,5 +94,5 @@ def test_custom_hardtanh_quant( tflite_input_preprocess=ToNHWCPreprocess(), tflite_output_preprocess=ToNCHWPreprocess(), input_data=input_data, - atol=1.0, + atol=2.0, ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_linear_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_linear_converter.py index 4a19c3d8c4b..858724522cd 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_linear_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_linear_converter.py @@ -10,6 +10,7 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_edge_program from executorch.backends.nxp.tests.executors import convert_run_compare from executorch.backends.nxp.tests.models import LinearModule +from executorch.exir.dialects._ops import ops as exir_ops @pytest.fixture(autouse=True) @@ -26,15 +27,23 @@ def test_linear_conversion__with_bias(): input_data = np.random.random(input_shape).astype(np.float32) - convert_run_compare(edge_program, input_data=input_data, atol=1.0e-6) + nodes = list(edge_program.graph.nodes) + assert nodes[4].target == exir_ops.edge.aten.addmm.default + assert len(nodes[4].args) == 3 # Has bias. + + convert_run_compare(edge_program, input_data=input_data) def test_linear_conversion__without_bias(): input_shape = (10, 32) edge_program = to_edge_program( - LinearModule(bias=True), input_shape + LinearModule(bias=False), input_shape ).exported_program() input_data = np.random.random(input_shape).astype(np.float32) - convert_run_compare(edge_program, input_data=input_data, atol=1.0e-6) + nodes = list(edge_program.graph.nodes) + assert nodes[3].target == exir_ops.edge.aten.mm.default + assert len(nodes[3].args) == 2 # No bias. + + convert_run_compare(edge_program, input_data=input_data) 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 9139dd97f9a..c5d7d4d6a38 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py @@ -1,5 +1,4 @@ # Copyright 2025 NXP -# 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. 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 new file mode 100644 index 00000000000..40857d18eb8 --- /dev/null +++ b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py @@ -0,0 +1,85 @@ +# Copyright 2025 NXP +# 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 kgb +import numpy as np +import torch + +from executorch.backends.nxp.nxp_backend import EdgeProgramToIRConverter +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import ( + convert_run_compare, + graph_contains_any_of_ops, + ToChannelFirstPreprocess, + ToChannelLastPreprocess, +) +from executorch.backends.nxp.tests.models import Conv2dWithActivation +from executorch.exir.dialects._ops import ops as exir_ops +from parameterized import parameterized +from torch.export import ExportedProgram + + +class TestTanhConverter(unittest.TestCase): + __test__ = False # Prevent interfering with PyTest tests + + @parameterized.expand( + input=[ + ( + "inplace", + True, + ), + ( + "not_inplace", + False, + ), + ] + ) + def test_conv_tanh( + self, _: str, inplace: bool, input_shape: tuple[int] = (1, 3, 112, 112) + ): + with kgb.spy_on( + EdgeProgramToIRConverter.convert_program, call_original=True + ) as converter_spy: + if inplace: + model = Conv2dWithActivation( + activation=torch.tanh_, in_channels=input_shape[1] + ) + else: + model = Conv2dWithActivation( + activation=torch.tanh, in_channels=input_shape[1] + ) + + quantized_program = to_quantized_edge_program( + model, input_shape + ).exported_program() + tflite_flatbuffers_model, io_formats = converter_spy.calls[-1].return_value + exported_program: ExportedProgram = converter_spy.calls[-1].args[0] + + lowered_module_graph = ( + quantized_program.graph_module.lowered_module_0.original_module.graph + ) + tanh_ops = [ + exir_ops.edge.aten.tanh.default, + exir_ops.edge.aten.tanh_.default, + ] + assert graph_contains_any_of_ops(graph=lowered_module_graph, ops=tanh_ops) + + input_data = (np.random.random(input_shape) * 50).astype(np.int8) + convert_run_compare( + exported_program, + tfl_model=tflite_flatbuffers_model, + tflite_input_preprocess=ToChannelLastPreprocess(), + tflite_output_preprocess=ToChannelFirstPreprocess(), + input_data=input_data, + atol=1.0, + ) + + @classmethod + def setUpClass(cls): + torch.manual_seed(23) + np.random.seed(23) diff --git a/backends/nxp/tests/ir/edge_passes/__init__.py b/backends/nxp/tests/ir/edge_passes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d 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 d7920aa55d8..17b040fbc3d 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 @@ -1,9 +1,11 @@ -# Copyright 2025 NXP +# Copyright 2024-2025 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 itertools +import executorch.extension.pybindings.portable_lib import executorch.kernels.quantized # noqa F401 import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -57,12 +59,12 @@ def test_remove_io_quant_ops_pass__cifarnet(): ) nodes = list(exec_prog.exported_program().graph.nodes) - assert len(nodes) == 17 + assert len(nodes) == 11 assert ( nodes[0].meta["val"].dtype == torch.int8 ), "Input tensor doesn't have type INT8." assert ( - nodes[16].meta["val"][0].dtype == torch.int8 + nodes[10].meta["val"][0].dtype == torch.int8 ), "Output tensor doesn't have type INT8." assert ( diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/models.py index 3aafab36a95..bdad9ddc4b4 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/models.py @@ -4,11 +4,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 Collection, Union +from typing import Callable, Collection, Union import torch +class Conv1dModule(torch.nn.Module): + def __init__( + self, + bias: bool = True, + dilation: Union[int, tuple[int, int]] = 1, + in_channels: int = 4, + kernel_size: Union[int, tuple[int, int]] = 3, + out_channels: int = 8, + padding: Union[str, int, Collection[int]] = 0, + stride: Union[int, tuple[int, int]] = 2, + group: int = 1, + ): + super().__init__() + + self.conv = torch.nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + groups=group, + ) + + def forward(self, x): + return self.conv(x) + + class Conv2dModule(torch.nn.Module): def __init__( self, @@ -38,6 +67,35 @@ def forward(self, x): return self.conv(x) +class Conv3dModule(torch.nn.Module): + def __init__( + self, + bias: bool = True, + dilation: Union[int, tuple[int, int]] = 1, + in_channels: int = 4, + kernel_size: Union[int, tuple[int, int]] = 3, + out_channels: int = 8, + padding: Union[str, int, Collection[int]] = 0, + stride: Union[int, tuple[int, int]] = 2, + group: int = 1, + ): + super().__init__() + + self.conv = torch.nn.Conv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + groups=group, + ) + + def forward(self, x): + return self.conv(x) + + class Conv2dAndMaxPool2DModule(torch.nn.Module): def __init__(self): super().__init__() @@ -142,6 +200,24 @@ def forward(self, x): return x +class ConvFCFCSoftmaxModuleWithoutReshape(torch.nn.Module): + def __init__(self): + super().__init__() + + self.conv = torch.nn.Conv2d(4, 5, 2, bias=False) + self.fc1 = torch.nn.Linear(32, 16) + self.fc2 = torch.nn.Linear(16, 8) + self.softmax = torch.nn.Softmax(1) + + def forward(self, x): + x = self.conv(x) + x = self.fc1(x) + x = self.fc2(x) + x = self.softmax(x) + + return x + + class ConstantPadNDModule(torch.nn.Module): def __init__(self, paddings: Collection[int], constant: float | int | None = None): super().__init__() @@ -271,6 +347,20 @@ def forward(self, x): return self.relu(x) +class Conv2dWithActivation(torch.nn.Module): + def __init__(self, activation: torch.nn.Module | Callable, in_channels: int = 3): + super().__init__() + + self.conv = torch.nn.Conv2d( + in_channels=in_channels, out_channels=64, kernel_size=(3, 3) + ) + self.activation = activation + + def forward(self, x): + x = self.conv(x) + return self.activation(x) + + class Conv2dReLUModule(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/nxp/tests/test_batch_norm_fusion.py b/backends/nxp/tests/test_batch_norm_fusion.py index c058543be2d..3f1106c6d24 100644 --- a/backends/nxp/tests/test_batch_norm_fusion.py +++ b/backends/nxp/tests/test_batch_norm_fusion.py @@ -15,8 +15,11 @@ AddMMConverter, MMConverter, ) +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters.view_copy_converter import ( + ViewCopyConverter, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import OverrideSupportedTargets +from executorch.backends.nxp.tests.executors import OverrideTargetSupportCheck from torch import nn @@ -92,7 +95,7 @@ def test_batch_norm_conv_fusing(bias: bool, input_shape: list[int]): example_input = (torch.ones(*input_shape),) module = ConvBatchNormModule(bias, len(input_shape), 4) - program = torch.export.export_for_training(module, example_input, strict=True) + program = torch.export.export(module, example_input, strict=True) og_module = program.module() pm = NeutronAtenPassManager() @@ -126,7 +129,7 @@ def test_batch_norm_linear_fusing(bias: bool): example_input = (torch.ones(*input_shape),) module = LinearBatchNormModule(bias, 4, input_shape[-1], input_shape[1]) - program = torch.export.export_for_training(module, example_input, strict=True) + program = torch.export.export(module, example_input, strict=True) og_module = program.module() pm = NeutronAtenPassManager() @@ -201,14 +204,24 @@ def test_batch_norm_linear_fusing__full_pipeline(bias: bool): # Don't delegate the Linear node, because there seems to be a bug with the NeutronConverter/NeutronPartitioner. # But that doesn't affect the validity of this test. - with OverrideSupportedTargets(AddMMConverter, new_targets=[]): - with OverrideSupportedTargets(MMConverter, new_targets=[]): - edge_program = to_quantized_edge_program( - module, tuple(input_shape) - ).exported_program() - nodes = list(edge_program.graph.nodes) - - assert len(nodes) == 14 + def unsupported_target(*_): # Accept all input arguments and return `False`. + return False + + with OverrideTargetSupportCheck( + AddMMConverter, new_target_support_check=unsupported_target + ): + with OverrideTargetSupportCheck( + MMConverter, new_target_support_check=unsupported_target + ): + with OverrideTargetSupportCheck( + ViewCopyConverter, new_target_support_check=unsupported_target + ): + edge_program = to_quantized_edge_program( + module, tuple(input_shape) + ).exported_program() + nodes = list(edge_program.graph.nodes) + + assert len(nodes) == 18 assert not any( node.op == "call_function" and "batch_norm" in node.target.__name__ for node in nodes diff --git a/backends/nxp/tests/test_edge_passes.py b/backends/nxp/tests/test_edge_passes.py new file mode 100644 index 00000000000..a189299be52 --- /dev/null +++ b/backends/nxp/tests/test_edge_passes.py @@ -0,0 +1,88 @@ +import numpy as np +from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import ( + ViewCopyConverter, +) +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import ( + EdgeProgramExecutor, + OverrideTargetSupportCheck, +) +from executorch.backends.nxp.tests.models import ConvFCFCSoftmaxModuleWithoutReshape +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph, Node + + +def _is_view_copy(node_: Node) -> bool: + return ( + node_.op == "call_function" + and node_.target == exir_ops.edge.aten.view_copy.default + ) + + +def _is_dequantize(node_: Node) -> bool: + return ( + node_.op == "call_function" + and node_.target.__name__ + == "quantized_decomposed.dequantize_per_tensor.default" + ) + + +def _is_quantize(node_: Node) -> bool: + return ( + node_.op == "call_function" + and node_.target.__name__ == "quantized_decomposed.quantize_per_tensor.default" + ) + + +def _find_view_copy_node_indices(graph_nodes: list[Node]) -> list[int]: + view_copy_nodes_indices = [] + + for idx, node in enumerate(graph_nodes): + if _is_view_copy(node): + view_copy_nodes_indices.append(idx) + + return view_copy_nodes_indices + + +def _assert_nodes_form_a_view_copy_qdq_cluster(graph: Graph, node_indices: list[int]): + assert len(node_indices) == 3 + + nodes = list(graph.nodes) + assert _is_dequantize(dequantize := nodes[node_indices[0]]) + assert _is_view_copy(view_copy := nodes[node_indices[1]]) + assert _is_quantize(quantize := nodes[node_indices[2]]) + + # Make sure the nodes are properly connected. + assert view_copy.args[0] == dequantize + assert quantize.args[0] == view_copy + + +def test_moving_view_copy_into_separate_qdq_clusters(): + model = ConvFCFCSoftmaxModuleWithoutReshape() + input_shape = (1, 4, 3, 33) + + # Prohibit `view_copy` conversion for the testing purposes. + def unsupported_target(*_): + return False + + with OverrideTargetSupportCheck( + ViewCopyConverter, new_target_support_check=unsupported_target + ): + epm = to_quantized_edge_program(model, input_shape, target="imxrt700") + exported_program = epm.exported_program() + + nodes = list(exported_program.graph_module.graph.nodes) + assert len(nodes) == 28 + + view_copy_indices = _find_view_copy_node_indices(nodes) + + assert len(view_copy_indices) == 4 + for idx in view_copy_indices: + _assert_nodes_form_a_view_copy_qdq_cluster( + exported_program.graph, node_indices=[idx - 1, idx, idx + 1] + ) + + # Make sure the program is runnable. + input_data = np.random.random(input_shape).astype("float32") + program_executor = EdgeProgramExecutor(exported_program) + program_executor.inference(input_data) diff --git a/backends/nxp/tests/test_gru_splitting.py b/backends/nxp/tests/test_gru_splitting.py new file mode 100644 index 00000000000..a2e9d324f69 --- /dev/null +++ b/backends/nxp/tests/test_gru_splitting.py @@ -0,0 +1,258 @@ +# Copyright 2025 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import pytest +import torch + +from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( + NeutronAtenPassManager, +) +from executorch.backends.nxp.aten_passes.split_gru_based_on_num_layers import ( + SplitGRUBasedOnNumLayers, +) + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(42) + np.random.seed(23) + + +def get_gru_input_shapes( + input_size=8, + hidden_size=8, + sequence_length=8, + batch_size=8, + num_layers=1, + D=1, # Unidirectional. +): + input_shapes = [ + (batch_size, sequence_length, input_size), + (D * num_layers, batch_size, hidden_size), + ] + + return input_shapes, input_size, hidden_size, sequence_length + + +class GruModule(torch.nn.Module): + def __init__( + self, + input_size, + hidden_size, + num_layers=1, + bias=True, + batch_first=False, + dropout=0.0, + bidirectional=False, + ): + super().__init__() + self.gru = torch.nn.GRU( + input_size, + hidden_size, + num_layers, + bias, + batch_first, + dropout, + bidirectional, + ) + + def forward(self, input_, h_0): + # `input_` has shape + # [sequence_length, batch_size] or [sequence_length, batch_size, input_size] if `batch_first` is False + # [batch_size, sequence_length, input_size] if `batch_first` is True + # `h_0` has shape [D * num_layers, hidden_size] or [D * num_layers, batch_size, hidden_size] + # where `D` is equal to 2 if bidirectional==True otherwise 1 + return self.gru(input_, h_0) + + +@pytest.mark.parametrize( + "num_layers", [2, 3, 8], ids=lambda num_layers: f"num_layers = {num_layers}" +) +def test_gru_splitting__with_bias(num_layers): + input_shapes, input_size, hidden_size, _ = get_gru_input_shapes( + num_layers=num_layers + ) + model = GruModule(input_size, hidden_size, num_layers=num_layers).eval() + + example_input = tuple(torch.ones(input_shape) for input_shape in input_shapes) + exir_program_aten = torch.export.export(model, example_input).module() + + pre_pass_output = [t.detach() for t in exir_program_aten(*example_input)] + assert len(exir_program_aten.graph.nodes) == 6 + (num_layers) * 4 + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == 1 + ) # Just 1 `GRU` in the model. + + # Run pre-processing passes of the float32 aten dialect program. + pytorch_pass_manager = NeutronAtenPassManager([SplitGRUBasedOnNumLayers()]) + pytorch_pass_manager(exir_program_aten) + + post_pass_output = [t.detach() for t in exir_program_aten(*example_input)] + nodes = list(exir_program_aten.graph.nodes) + assert len(nodes) == 5 + (num_layers) * 8 + assert nodes[2 + num_layers * 4].target == torch.ops.aten.split.default + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == num_layers + ) # Many `GRU` nodes. + assert nodes[-2].target == torch.ops.aten.cat.default + + assert np.allclose(pre_pass_output[0], post_pass_output[0]), "Main outputs differ" + assert np.allclose(pre_pass_output[1], post_pass_output[1]), "Hidden outputs differ" + + +@pytest.mark.parametrize( + "num_layers", [2, 5], ids=lambda num_layers: f"num_layers = {num_layers}" +) +def test_gru_splitting__no_bias(num_layers): + input_shapes, input_size, hidden_size, _ = get_gru_input_shapes( + num_layers=num_layers + ) + model = GruModule(input_size, hidden_size, num_layers=num_layers, bias=False).eval() + + example_input = tuple(torch.ones(input_shape) for input_shape in input_shapes) + exir_program_aten = torch.export.export(model, example_input).module() + + pre_pass_output = [t.detach() for t in exir_program_aten(*example_input)] + assert len(exir_program_aten.graph.nodes) == 6 + (num_layers) * 2 + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == 1 + ) # Just 1 `GRU` in the model. + + # Run pre-processing passes of the float32 aten dialect program. + pytorch_pass_manager = NeutronAtenPassManager([SplitGRUBasedOnNumLayers()]) + pytorch_pass_manager(exir_program_aten) + + post_pass_output = [t.detach() for t in exir_program_aten(*example_input)] + nodes = list(exir_program_aten.graph.nodes) + assert len(nodes) == 5 + (num_layers) * 6 + assert nodes[2 + num_layers * 2].target == torch.ops.aten.split.default + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == num_layers + ) # Many `GRU` nodes. + assert nodes[-2].target == torch.ops.aten.cat.default + + assert np.allclose(pre_pass_output[0], post_pass_output[0]), "Main outputs differ" + assert np.allclose(pre_pass_output[1], post_pass_output[1]), "Hidden outputs differ" + + +@pytest.mark.parametrize( + "num_layers", [2, 3, 5], ids=lambda num_layers: f"num_layers = {num_layers}" +) +def test_gru_splitting__bidirectional__no_bias(num_layers): + input_shapes, input_size, hidden_size, _ = get_gru_input_shapes( + num_layers=num_layers, D=2 + ) + model = GruModule( + input_size, hidden_size, num_layers=num_layers, bidirectional=True, bias=False + ).eval() + + example_input = tuple(torch.ones(input_shape) for input_shape in input_shapes) + exir_program_aten = torch.export.export(model, example_input).module() + + assert len(exir_program_aten.graph.nodes) == 6 + (num_layers) * 4 + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == 1 + ) # Just 1 `GRU` in the model. + + # Run pre-processing passes of the float32 aten dialect program. + pytorch_pass_manager = NeutronAtenPassManager([SplitGRUBasedOnNumLayers()]) + pytorch_pass_manager(exir_program_aten) + + nodes = list(exir_program_aten.graph.nodes) + assert len(nodes) == 5 + (num_layers) * 8 + assert nodes[2 + num_layers * 4].target == torch.ops.aten.split.default + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == num_layers + ) # Many `GRU` in the model. + assert nodes[-2].target == torch.ops.aten.cat.default + + +@pytest.mark.parametrize( + "num_layers", [2, 3, 7], ids=lambda num_layers: f"num_layers = {num_layers}" +) +def test_gru_splitting__bidirectional__with_bias(num_layers): + input_shapes, input_size, hidden_size, _ = get_gru_input_shapes( + num_layers=num_layers, D=2 + ) + model = GruModule( + input_size, hidden_size, num_layers=num_layers, bidirectional=True, bias=True + ).eval() + + example_input = tuple(torch.ones(input_shape) for input_shape in input_shapes) + exir_program_aten = torch.export.export(model, example_input).module() + + assert len(exir_program_aten.graph.nodes) == 6 + (num_layers) * 8 + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == 1 + ) # Just 1 `GRU` in the model. + + # Run pre-processing passes of the float32 aten dialect program. + pytorch_pass_manager = NeutronAtenPassManager([SplitGRUBasedOnNumLayers()]) + pytorch_pass_manager(exir_program_aten) + + nodes = list(exir_program_aten.graph.nodes) + assert len(nodes) == 5 + (num_layers) * 12 + assert nodes[2 + num_layers * 8].target == torch.ops.aten.split.default + assert ( + len( + [ + n + for n in exir_program_aten.graph.nodes + if n.target == torch.ops.aten.gru.input + ] + ) + == num_layers + ) # Many `GRU` in the model. + assert nodes[-2].target == torch.ops.aten.cat.default diff --git a/backends/nxp/tests/test_integration.py b/backends/nxp/tests/test_integration.py index 6c143df79b3..d31b22c9ce9 100644 --- a/backends/nxp/tests/test_integration.py +++ b/backends/nxp/tests/test_integration.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -43,8 +43,8 @@ def test_cifarnet(): delegation_info = get_delegation_info(exec_prog.exported_program().graph_module) assert delegation_info.num_delegated_subgraphs == 1 - assert delegation_info.num_non_delegated_nodes == 17 - assert delegation_info.num_delegated_nodes == 42 + assert delegation_info.num_non_delegated_nodes == 11 + assert delegation_info.num_delegated_nodes == 45 nodes = list(exec_prog.exported_program().graph.nodes) assert nodes[2].name == "quantized_decomposed_quantize_per_tensor_default" diff --git a/backends/nxp/tests/test_neutron_converter_manager.py b/backends/nxp/tests/test_neutron_converter_manager.py index fb816ef199f..af723ec9c7a 100644 --- a/backends/nxp/tests/test_neutron_converter_manager.py +++ b/backends/nxp/tests/test_neutron_converter_manager.py @@ -30,7 +30,7 @@ def test_conv2d_neutron_conversion__default_flavor(): neutron_converter_manager = NeutronConverterManager() neutron_model = neutron_converter_manager.convert( - tflite_model, "imxrt700", "SDK_25_03" + tflite_model, "imxrt700", "SDK_25_06" ) assert len( diff --git a/backends/nxp/tests/test_quantizer.py b/backends/nxp/tests/test_quantizer.py index e97889e09a2..ef5fbb0cbca 100644 --- a/backends/nxp/tests/test_quantizer.py +++ b/backends/nxp/tests/test_quantizer.py @@ -23,9 +23,7 @@ def test_quantizer_conv2d(): example_input = (torch.ones(1, 4, 32, 32),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -64,9 +62,7 @@ def test_quantizer_linear(): example_input = (torch.ones(10, 32),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -105,9 +101,7 @@ def test_quantizer_maxpool2d(): example_input = (torch.ones(1, 8, 32, 32),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -143,9 +137,7 @@ def test_quantizer_softmax(): example_input = (torch.ones(1, 10),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -182,9 +174,7 @@ def test_quantizer_single_maxpool2d(): example_input = (torch.ones(1, 4, 32, 32),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -206,9 +196,7 @@ def test_quantizer_conv2d_relu(): example_input = (torch.ones(1, 4, 32, 32),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -231,9 +219,7 @@ def test_quantizer_conv2d_avg_pool2d(): example_input = (torch.ones(1, 4, 16, 16),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -256,9 +242,7 @@ def test_quantizer_conv2d_permute(): example_input = (torch.ones(1, 4, 16, 16),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -285,9 +269,7 @@ def test_multiple_shared_spec_ops_in_row(): example_input = (torch.ones(1, 3, 64, 64),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() # noinspection PyTypeChecker m = prepare_pt2e(graph_module, quantizer) @@ -321,9 +303,7 @@ def test_quantizers_order_invariance(): example_input = (torch.ones(1, 4, 64, 64),) quantizer = NeutronQuantizer() - graph_module = torch.export.export_for_training( - model, example_input, strict=True - ).module() + graph_module = torch.export.export(model, example_input, strict=True).module() m = prepare_pt2e(deepcopy(graph_module), quantizer) m(*example_input) diff --git a/backends/nxp/tests/test_removing_dead_code.py b/backends/nxp/tests/test_removing_dead_code.py new file mode 100644 index 00000000000..7b8641fb247 --- /dev/null +++ b/backends/nxp/tests/test_removing_dead_code.py @@ -0,0 +1,60 @@ +# Copyright 2025 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 unittest + +import numpy as np +import pytest +import torch + +from executorch.backends.nxp.tests.executorch_pipeline import _quantize_model +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(42) + np.random.seed(23) + + +class DeadCodeModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.eval() + + def forward(self, x): + _ = torch.add(x, x) # Dead code + return torch.mul(x, x) + + +class TestRemovingDeadCode(unittest.TestCase): + __test__ = False # Prevent interfering with PyTest tests + + def test_removing_dead_code(self): + input_shape = (42,) + example_inputs = (torch.ones(input_shape),) + model = DeadCodeModule() + + exir_program_aten = torch.export.export(model, example_inputs, strict=True) + + # Make sure the model contains the dead code. + assert graph_contains_any_of_ops( + exir_program_aten.module().graph, [torch.ops.aten.add.Tensor] + ) + + # The `NeutronQuantizer` should remove the dead code in the `transform_for_annotation()` method. + exir_program_aten_quant = _quantize_model( + exir_program_aten.module(), [example_inputs] + ) + + # Make sure the is no `add` operation in the graph anymore. + assert not any( + "add" in str(node.target) for node in exir_program_aten_quant.graph.nodes + ) + + @classmethod + def setUpClass(cls): + torch.manual_seed(23) + np.random.seed(23) diff --git a/backends/nxp/tests/test_removing_nodes_with_known_outputs.py b/backends/nxp/tests/test_removing_nodes_with_known_outputs.py new file mode 100644 index 00000000000..8f5549c8526 --- /dev/null +++ b/backends/nxp/tests/test_removing_nodes_with_known_outputs.py @@ -0,0 +1,109 @@ +# Copyright 2025 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 unittest + +import numpy as np +import torch + +from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( + NeutronAtenPassManager, +) +from executorch.backends.nxp.aten_passes.remove_nodes_with_known_outputs import ( + RemoveNodesWithKnownOutputs, +) +from executorch.backends.nxp.aten_passes.split_gru_based_on_num_layers import ( + SplitGRUBasedOnNumLayers, +) +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from parameterized import parameterized +from torch import nn + + +class GRUModel(nn.Module): + def __init__(self, num_layers=1): + super().__init__() + self.gru = torch.nn.GRU(8, 8, num_layers=num_layers) + + def forward(self, input_): + # `input_` has shape [sequence_length, batch_size, input_size] ([8, 1, 8]) + return self.gru( + input_, None + ) # The initial hidden is `None`, which will result in a `Zeros` node being added. + + +class TestRemovingNodesWithKnownOutputs(unittest.TestCase): + __test__ = False # Prevent interfering with PyTest tests. + + @classmethod + def setUpClass(cls): + torch.manual_seed(23) + np.random.seed(42) + + def test_removing_nodes__zeros(self): + model = GRUModel() + + input_shape = (8, 1, 8) + example_input = (torch.ones(input_shape),) + + exir_program_aten = torch.export.export(model, example_input).module() + + # Make sure the `aten.zeros` is in the model. + assert graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.zeros.default] + ) + outputs_before = [o.detach().numpy() for o in exir_program_aten(*example_input)] + + # Apply the optimization. + NeutronAtenPassManager([RemoveNodesWithKnownOutputs()])(exir_program_aten) + + # Make sure the `aten.zeros` is no longer in the model. + assert not graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.zeros.default] + ) + outputs_after = [o.detach().numpy() for o in exir_program_aten(*example_input)] + + # Make sure the model still produces the exact same output. + assert np.allclose(outputs_before[0], outputs_after[0]) + assert np.allclose(outputs_before[1], outputs_after[1]) + + @parameterized.expand([2, 8]) + def test_removing_nodes__split(self, num_layers): + # `num_layers > 1` will result in a `split` operator being added. It's input will be a `zeros` operator, which + # provides the static 0s input data. + model = GRUModel(num_layers).eval() + + input_shape = (8, 1, 8) + example_input = (torch.ones(input_shape),) + + exir_program_aten = torch.export.export(model, example_input).module() + + # Apply the pass to split the `aten.gru.input` into multiple instances, and add a `split` node. + NeutronAtenPassManager([SplitGRUBasedOnNumLayers()])(exir_program_aten) + + # Make sure the `aten.zeros` and `torch.split` are in the model. + assert graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.zeros.default] + ) + assert graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.split.default] + ) + outputs_before = [o.detach().numpy() for o in exir_program_aten(*example_input)] + + # Apply the optimization. + NeutronAtenPassManager([RemoveNodesWithKnownOutputs()])(exir_program_aten) + + # Make sure the `aten.zeros` and `torch.split` are no longer in the model. + assert not graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.zeros.default] + ) + assert not graph_contains_any_of_ops( + exir_program_aten.graph, [torch.ops.aten.split.default] + ) + outputs_after = [o.detach().numpy() for o in exir_program_aten(*example_input)] + + # Make sure the model still produces the exact same output. + assert np.allclose(outputs_before[0], outputs_after[0]) + assert np.allclose(outputs_before[1], outputs_after[1]) diff --git a/backends/nxp/tests/test_split_group_convolution.py b/backends/nxp/tests/test_split_group_convolution.py new file mode 100644 index 00000000000..1da53af794d --- /dev/null +++ b/backends/nxp/tests/test_split_group_convolution.py @@ -0,0 +1,265 @@ +# Copyright 2025 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 unittest +from copy import deepcopy + +import numpy as np +import torch + +from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( + NeutronAtenPassManager, +) +from executorch.backends.nxp.aten_passes.split_group_convolution import ( + SplitGroupConvolution, +) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec +from executorch.backends.nxp.tests.executorch_pipeline import ( + _quantize_model, + get_random_calibration_inputs, + to_model_input_spec, +) +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.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 + + +def _quantize_and_lower_module( + module: GraphModule, input_shape: tuple[int, ...], target="imxrt700" +) -> EdgeProgramManager: + calibration_inputs = get_random_calibration_inputs(to_model_input_spec(input_shape)) + + exir_program_aten__module_quant = _quantize_model(module, calibration_inputs) + + edge_compile_config = EdgeCompileConfig(_check_ir_validity=False) + edge_program_manager = export_to_edge( + exir_program_aten__module_quant, + calibration_inputs[0], + edge_compile_config=edge_compile_config, + ) + + compile_spec = generate_neutron_compile_spec(target, "SDK_25_06") + partitioner = NeutronPartitioner(compile_spec) + return edge_program_manager.to_backend(partitioner) + + +class TestSplitGroupConvolution(unittest.TestCase): + __test__ = False # Prevent interfering with PyTest tests. + + @classmethod + def setUp(cls): + torch.manual_seed(23) + np.random.seed(42) + + @parameterized.expand( + [ + ["group = 2", [1, 16, 10, 10], 2], + ["group = 3", [1, 24, 10, 10], 3], + ["group = 8", [1, 8, 10, 10], 8], + ] + ) + def test_split_group_convolution__2d(self, _, input_shape: list[int], group: int): + example_input = (torch.ones(input_shape),) + + module = Conv2dModule( + bias=True, + in_channels=input_shape[1], + out_channels=8 + * group, # Make sure the output channels are multiple of 8, so the `cat` can be delegated. + group=group, + stride=1, + ) + graph_module = torch.export.export(module, example_input, strict=True).module() + original_module = deepcopy(graph_module) + + modified_module = NeutronAtenPassManager([SplitGroupConvolution()])( + graph_module + ).graph_module + + # Make sure the fusion worked. + original_nodes = list(original_module.graph.nodes) + modified_nodes = list(modified_module.graph.nodes) + + assert len(original_nodes) == 5 + assert original_nodes[3].target == torch.ops.aten.conv2d.default + assert original_nodes[3].args[-1] == group + + assert len(modified_nodes) == 4 + group * 4 + assert modified_nodes[1].target == torch.ops.aten.split.default + for node in modified_nodes[2 + 3 * group : 4 + 3 * group]: + assert node.target == torch.ops.aten.conv2d.default + assert node.args[-1] == 1 # Groups. + assert modified_nodes[-2].target == torch.ops.aten.cat.default + + # Verify that the behavior has not changed. + input_data = torch.randn(input_shape, dtype=torch.float32) + out1 = original_module(input_data).detach().numpy() + out2 = modified_module(input_data).detach().numpy() + assert np.allclose(out1, out2, atol=2.0e-7) + + # Make sure the graph can be correctly quantized and lowered to edge. + ep = _quantize_and_lower_module( + modified_module, tuple(input_shape) + ).exported_program() + nodes = list(ep.graph.nodes) + 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], + ) + + @parameterized.expand( + [ + ["group = 2", [1, 16, 10], 2], + ["group = 3", [1, 24, 10], 3], + ["group = 6", [1, 24, 10], 6], + ] + ) + def test_split_group_convolution__1d(self, _, input_shape: list[int], group: int): + example_input = (torch.ones(input_shape),) + + module = Conv1dModule( + bias=True, + in_channels=input_shape[1], + out_channels=8 + * group, # Make sure the output channels are multiple of 8, so the `cat` can be delegated. + group=group, + stride=1, + ) + graph_module = torch.export.export(module, example_input).module() + original_module = deepcopy(graph_module) + + modified_module = NeutronAtenPassManager([SplitGroupConvolution()])( + graph_module + ).graph_module + + # Make sure the fusion worked. + original_nodes = list(original_module.graph.nodes) + modified_nodes = list(modified_module.graph.nodes) + + assert len(original_nodes) == 5 + assert original_nodes[3].target == torch.ops.aten.conv1d.default + assert original_nodes[3].args[-1] == group + + assert len(modified_nodes) == 4 + group * 4 + assert modified_nodes[1].target == torch.ops.aten.split.default + for node in modified_nodes[2 + 3 * group : 4 + 3 * group]: + assert node.target == torch.ops.aten.conv1d.default + assert node.args[-1] == 1 # Groups. + assert modified_nodes[-2].target == torch.ops.aten.cat.default + + # Verify that the behavior has not changed. + input_data = torch.randn(input_shape, dtype=torch.float32) + out1 = original_module(input_data).detach().numpy() + out2 = modified_module(input_data).detach().numpy() + assert np.allclose(out1, out2, atol=2.0e-7) + + # Make sure the graph can be correctly quantized and lowered to edge. + ep = _quantize_and_lower_module( + modified_module, tuple(input_shape) + ).exported_program() + nodes = list(ep.graph.nodes) + 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], + ) + + @parameterized.expand( + [ + ["group = 2", [1, 16, 10, 10, 10], 2], + ] + ) + def test_split_group_convolution__3d(self, _, input_shape: list[int], group: int): + example_input = (torch.ones(input_shape),) + + module = Conv3dModule( + bias=True, + in_channels=input_shape[1], + out_channels=8 + * group, # Make sure the output channels are multiple of 8, so the `cat` can be delegated. + group=group, + ) + graph_module = torch.export.export(module, example_input).module() + original_module = deepcopy(graph_module) + + modified_module = NeutronAtenPassManager([SplitGroupConvolution()])( + graph_module + ).graph_module + + # Verify that the pass has NOT made any changes, as it is disabled for 3D convolution. + original_nodes = list(original_module.graph.nodes) + modified_nodes = list(modified_module.graph.nodes) + + assert len(original_nodes) == len(modified_nodes) + for original_node, modified_node in zip(original_nodes, modified_nodes): + assert original_node.name == modified_node.name + assert original_node.target == modified_node.target + + # Verify that the behavior has not changed. + input_data = torch.randn(input_shape, dtype=torch.float32) + out1 = original_module(input_data).detach().numpy() + out2 = modified_module(input_data).detach().numpy() + assert np.allclose(out1, out2) + + def test_split_group_convolution__applied_by_default(self): + input_shape = [1, 16, 10, 10] + group = 2 + example_input = (torch.ones(input_shape),) + + module = Conv2dModule( + in_channels=input_shape[1], + out_channels=8 + * group, # Make sure the output channels are multiple of 8, so the `cat` can be delegated. + group=group, + stride=1, + ) + graph_module = torch.export.export(module, example_input).module() + original_module = deepcopy(graph_module) + + modified_module = NeutronAtenPassManager()( + graph_module + ).graph_module # Default passes. + + # Make sure the fusion worked. + original_nodes = list(original_module.graph.nodes) + modified_nodes = list(modified_module.graph.nodes) + + assert len(original_nodes) == 5 + assert original_nodes[3].target == torch.ops.aten.conv2d.default + assert original_nodes[3].args[-1] == group + + assert len(modified_nodes) == 4 + group * 4 + assert modified_nodes[1].target == torch.ops.aten.split.default + for node in modified_nodes[2 + 3 * group : 4 + 3 * group]: + assert node.target == torch.ops.aten.conv2d.default + assert node.args[-1] == 1 # Groups. + assert modified_nodes[-2].target == torch.ops.aten.cat.default + + # Verify that the behavior has not changed. + input_data = torch.randn(input_shape, dtype=torch.float32) + out1 = original_module(input_data).detach().numpy() + out2 = modified_module(input_data).detach().numpy() + assert np.allclose(out1, out2, atol=5.0e-7) + + # Make sure the graph can be correctly quantized and lowered to edge. + ep = _quantize_and_lower_module( + modified_module, tuple(input_shape) + ).exported_program() + nodes = list(ep.graph.nodes) + 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], + ) diff --git a/backends/openvino/CMakeLists.txt b/backends/openvino/CMakeLists.txt index bf5fc3b217e..4d32d8932c2 100644 --- a/backends/openvino/CMakeLists.txt +++ b/backends/openvino/CMakeLists.txt @@ -38,38 +38,50 @@ add_library(openvino_backend STATIC .) target_compile_options(openvino_backend PRIVATE -frtti -fexceptions) # Include Executorch directories -target_include_directories(openvino_backend PUBLIC ${COMMON_INCLUDE_DIRS}) +target_include_directories(openvino_backend PRIVATE ${COMMON_INCLUDE_DIRS}) # Link OpenVINO and ExecuteTorch core libraries -target_link_libraries(openvino_backend PRIVATE openvino::runtime executorch_core) +target_link_libraries( + openvino_backend PRIVATE openvino::runtime executorch_core +) # Add source files for OpenVINO backend -target_sources(openvino_backend PRIVATE ${CMAKE_CURRENT_LIST_DIR}/runtime/OpenvinoBackend.cpp) +target_sources( + openvino_backend + PRIVATE ${CMAKE_CURRENT_LIST_DIR}/runtime/OpenvinoBackend.cpp +) executorch_target_link_options_shared_lib(openvino_backend) if(EXECUTORCH_BUILD_OPENVINO_EXECUTOR_RUNNER) - # Build executor runner binary for openvino backend - list(APPEND openvino_executor_runner_libs openvino_backend executorch) - - set(_openvino_executor_runner__srcs - ${EXECUTORCH_ROOT}/examples/portable/executor_runner/executor_runner.cpp - ${EXECUTORCH_ROOT}/extension/data_loader/file_data_loader.cpp - ${EXECUTORCH_ROOT}/extension/evalue_util/print_evalue.cpp - ${EXECUTORCH_ROOT}/extension/runner_util/inputs.cpp - ${EXECUTORCH_ROOT}/extension/runner_util/inputs_portable.cpp - ) - add_executable(openvino_executor_runner ${_openvino_executor_runner__srcs}) - - list(APPEND openvino_executor_runner_libs) - - target_link_libraries( - openvino_executor_runner gflags portable_ops_lib ${openvino_executor_runner_libs} - ) - target_compile_options(openvino_executor_runner PUBLIC ${_common_compile_options}) + # Build executor runner binary for openvino backend + list(APPEND openvino_executor_runner_libs openvino_backend executorch) + + set(_openvino_executor_runner__srcs + ${EXECUTORCH_ROOT}/examples/portable/executor_runner/executor_runner.cpp + ${EXECUTORCH_ROOT}/extension/data_loader/file_data_loader.cpp + ${EXECUTORCH_ROOT}/extension/evalue_util/print_evalue.cpp + ${EXECUTORCH_ROOT}/extension/runner_util/inputs.cpp + ${EXECUTORCH_ROOT}/extension/runner_util/inputs_portable.cpp + ) + add_executable(openvino_executor_runner ${_openvino_executor_runner__srcs}) + + list(APPEND openvino_executor_runner_libs) + + target_link_libraries( + openvino_executor_runner gflags portable_ops_lib + ${openvino_executor_runner_libs} + ) + target_compile_options( + openvino_executor_runner PUBLIC ${_common_compile_options} + ) endif() - - # Install OpenVINO backend library to the lib directory -install(TARGETS openvino_backend DESTINATION lib) +install( + TARGETS openvino_backend + EXPORT ExecuTorchTargets + DESTINATION lib + INCLUDES + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} +) diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index a3134f72b4b..8ec40d7f7c6 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -93,7 +93,7 @@ exr::Result OpenvinoBackend::init( exr::Error OpenvinoBackend::execute( exr::BackendExecutionContext& context, exr::DelegateHandle* input_handle, - exr::EValue** args) const { + exr::Span args) const { ExecutionHandle* execution_handle = (ExecutionHandle*)input_handle; auto infer_request = execution_handle->infer_request; diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index 069e4659d37..d84e3ba1f86 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -45,7 +45,7 @@ class OpenvinoBackend final : public ::exr::BackendInterface { exr::Error execute( exr::BackendExecutionContext& context, exr::DelegateHandle* input_handle, - exr::EValue** args) const override; + exr::Span args) const override; void destroy(exr::DelegateHandle* handle) const override; private: diff --git a/backends/openvino/scripts/openvino_build.sh b/backends/openvino/scripts/openvino_build.sh index bc85d6b8410..5a26f0b6dae 100755 --- a/backends/openvino/scripts/openvino_build.sh +++ b/backends/openvino/scripts/openvino_build.sh @@ -52,7 +52,7 @@ main() { export CMAKE_BUILD_ARGS="--target openvino_backend" # Build the package - ./install_executorch.sh + ./install_executorch.sh --minimal # Install torchao pip install third-party/ao diff --git a/backends/qualcomm/CMakeLists.txt b/backends/qualcomm/CMakeLists.txt index 33f150413a3..32105597260 100644 --- a/backends/qualcomm/CMakeLists.txt +++ b/backends/qualcomm/CMakeLists.txt @@ -39,17 +39,6 @@ if(${ANDROID}) find_library(android_log log) endif() -set(qcir_schema_include_dir ${CMAKE_CURRENT_LIST_DIR}/aot/ir) -set(qcir_schema_output ${qcir_schema_include_dir}/qcir_generated.h) -add_custom_command( - OUTPUT qcir_schema_output - COMMAND flatc --cpp --cpp-std c++11 --scoped-enums -o - ${qcir_schema_include_dir} ${qcir_schema_include_dir}/qcir.fbs - DEPENDS flatc - COMMENT "Generating qualcomm ir schema headers" - VERBATIM -) - add_compile_options("-Wall" "-Werror" "-Wno-sign-compare") add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) @@ -69,11 +58,8 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") endif() include_directories( - BEFORE - ${_common_include_directories} - ${QNN_SDK_ROOT}/include/QNN + BEFORE ${_common_include_directories} ${QNN_SDK_ROOT}/include/QNN ${QNN_SDK_ROOT}/share/QNN/converter/jni - ${EXECUTORCH_SOURCE_DIR}/third-party/flatbuffers/include ${EXECUTORCH_SOURCE_DIR}/runtime/core/portable_type/c10 ) @@ -112,10 +98,9 @@ include_directories( # declare targets # add_library(executorch_backend INTERFACE) -add_library(qcir INTERFACE qcir_schema_output) -add_library(qcir_utils STATIC) add_library(qnn_backend STATIC) add_library(qnn_backend_cache STATIC) +add_library(qnn_backend_options STATIC) add_library(qnn_context STATIC) add_library(qnn_custom_protocol STATIC) add_library(qnn_dlc_manager STATIC) @@ -142,7 +127,6 @@ add_library(utils STATIC) # # declare dependency # -target_link_libraries(qcir_utils PRIVATE qcir) target_link_libraries(wrappers PRIVATE qnn_executorch_logging) target_link_libraries( qnn_implementation PRIVATE qnn_function_interface qnn_executorch_logging @@ -159,6 +143,7 @@ target_link_libraries( qnn_backend PRIVATE qnn_implementation qnn_logger qnn_op_package_manager ) target_link_libraries(qnn_custom_protocol PRIVATE qnn_logger) +target_link_libraries(qnn_backend_options PRIVATE qnn_schema) target_link_libraries( qnn_device PRIVATE qnn_executorch_logging qnn_implementation qnn_logger ) @@ -196,8 +181,9 @@ target_link_libraries( qnn_dlc_manager ) target_link_libraries( - qnn_executorch_backend PRIVATE qnn_executorch_header qnn_schema qnn_manager - executorch_core extension_tensor + qnn_executorch_backend + PRIVATE qnn_executorch_header qnn_schema qnn_manager executorch_core + extension_tensor qnn_backend_options ) set_target_properties( qnn_executorch_backend PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" @@ -225,10 +211,6 @@ add_subdirectory( ${QNN_EXECUTORCH_ROOT_DIR}/aot/wrappers ${CMAKE_CURRENT_BINARY_DIR}/qnn_executorch/wrappers ) -add_subdirectory( - ${QNN_EXECUTORCH_ROOT_DIR}/aot/ir - ${CMAKE_CURRENT_BINARY_DIR}/qnn_executorch/ir -) install( TARGETS qnn_executorch_backend EXPORT ExecuTorchTargets @@ -261,6 +243,7 @@ if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64") qnn_executorch_header executorch extension_tensor + qnn_backend_options ) target_link_libraries( PyQnnWrapperAdaptor PRIVATE pybind11::module pybind11::lto wrappers diff --git a/backends/qualcomm/__init__.py b/backends/qualcomm/__init__.py new file mode 100644 index 00000000000..04ba5fcf24b --- /dev/null +++ b/backends/qualcomm/__init__.py @@ -0,0 +1,23 @@ +import os + +from .scripts.download_qnn_sdk import ( + check_glibc_exist_and_validate, + install_qnn_sdk, + is_linux_x86, +) + + +env_flag = os.getenv("EXECUTORCH_BUILDING_WHEEL", "0").lower() +# If users have preinstalled QNN_SDK_ROOT, we will use it. +qnn_sdk_root_flag = os.getenv("QNN_SDK_ROOT", None) + +if ( + env_flag not in ("1", "true", "yes") + and not qnn_sdk_root_flag + and is_linux_x86() + and check_glibc_exist_and_validate() +): + ok = install_qnn_sdk() + + if not ok: + raise RuntimeError("Failed to install QNN SDK. Please check the logs above.") diff --git a/backends/qualcomm/_passes/TARGETS b/backends/qualcomm/_passes/TARGETS index a824ca9f6e5..62a0fc43a78 100644 --- a/backends/qualcomm/_passes/TARGETS +++ b/backends/qualcomm/_passes/TARGETS @@ -12,6 +12,7 @@ runtime.python_library( ], deps = [ "//executorch/backends/transforms:addmm_mm_to_linear", + "//executorch/backends/transforms:decompose_sdpa", "//executorch/exir/backend:backend_details", "//executorch/exir/backend:compile_spec_schema", ], diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index aaf65afd279..15fce79ea12 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -8,8 +8,9 @@ from .annotate_quant_attrs import AnnotateQuantAttrs from .annotate_stack import AnnotateStack from .annotate_unbind import AnnotateUnbind +from .canonicalize_conv import CanonicalizeConv from .convert_bmm_to_matmul import ConvertBmmToMatmul -from .convert_conv1d_to_conv2d import ConvertConv1dToConv2d +from .convert_linear_to_conv2d import ConvertLinearToConv2d from .convert_square_to_pow import ConvertSquareToPow from .decompose_any import DecomposeAny from .decompose_cdist import DecomposeCDist @@ -38,6 +39,7 @@ from .remove_redundancy import RemoveRedundancy from .replace_arange_args import ReplaceArangeArgs from .replace_inf_values import ReplaceInfValues +from .seq_mse import SeqMSE from .tag_quant_io import TagQuantIO @@ -46,8 +48,9 @@ AnnotateQuantAttrs, AnnotateStack, AnnotateUnbind, + CanonicalizeConv, ConvertBmmToMatmul, - ConvertConv1dToConv2d, + ConvertLinearToConv2d, ConvertSquareToPow, DecomposeAny, DecomposeCDist, @@ -76,5 +79,6 @@ RemoveRedundancy, ReplaceArangeArgs, ReplaceInfValues, + SeqMSE, TagQuantIO, ] diff --git a/backends/qualcomm/_passes/annotate_quant_attrs.py b/backends/qualcomm/_passes/annotate_quant_attrs.py index 64496b71f1c..610e88e6d3b 100644 --- a/backends/qualcomm/_passes/annotate_quant_attrs.py +++ b/backends/qualcomm/_passes/annotate_quant_attrs.py @@ -30,9 +30,14 @@ class AnnotateQuantAttrs(ExportPass): generated after quantization process. """ - def __init__(self, edge_program: torch.export.ExportedProgram): + def __init__( + self, + edge_program: torch.export.ExportedProgram, + skip_advanced_requant: bool = False, + ): super(AnnotateQuantAttrs, self).__init__() self.edge_program = edge_program + self.skip_advanced_requant = skip_advanced_requant def _annotate_source_nodes( self, quant_node: torch.fx.Node, quant_attrs: Dict[str, Any] @@ -82,16 +87,29 @@ def _annotate_requant(self, n): # TODO: Store multiple pairs of requantize attributes when we have an op builder # that has multiple outputs that requires quant attributes. - if any( - q_attrs[attr] != dq_attrs[attr] - for attr in [ - QCOM_SCALE, - QCOM_ZERO_POINT, - QCOM_QUANT_MIN, - QCOM_QUANT_MAX, - QCOM_DTYPE, - ] - ): + # Determine if requantization is needed based on configuration and attribute mismatch. + is_requant_needed = False + if self.skip_advanced_requant: + # In skip_advanced_requant mode, only consider requant if dtypes differ. + if q_attrs[QCOM_DTYPE] != dq_attrs[QCOM_DTYPE]: + is_requant_needed = True + else: + # In full requant mode, consider requant if any key attribute differs. + # This aims to improve accuracy by adjusting scale, zero_point, etc. + # Users can disable this if it causes regressions. + if any( + q_attrs[attr] != dq_attrs[attr] + for attr in [ + QCOM_SCALE, + QCOM_ZERO_POINT, + QCOM_QUANT_MIN, + QCOM_QUANT_MAX, + QCOM_DTYPE, + ] + ): + is_requant_needed = True + + if is_requant_needed: dq_attrs[QCOM_ENCODING] = q_attrs[QCOM_ENCODING] user_node = list(dq_node.users)[0] n.args[0].meta.setdefault(QCOM_REQUANTIZE, {}) diff --git a/backends/qualcomm/_passes/build_quant_io.py b/backends/qualcomm/_passes/build_quant_io.py index b34d50c4e24..d43842e84a5 100644 --- a/backends/qualcomm/_passes/build_quant_io.py +++ b/backends/qualcomm/_passes/build_quant_io.py @@ -5,9 +5,11 @@ # LICENSE file in the root directory of this source tree. import torch from executorch.backends.qualcomm.utils.constants import QCOM_QUANTIZED_IO +from executorch.exir.delegate import executorch_call_delegate -from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_base import ExportPass, ProxyValue from executorch.exir.tensor import TensorSpec +from torch.utils import _pytree as pytree class BuildQuantIo(ExportPass): @@ -26,21 +28,22 @@ def _make_spec(self, x): else: return None - def _build(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: - # Forcedly update delegate node's meta['spec'] to get correct output - # tensor size in runtime - call_delegate = [ - node - for node in graph_module.graph.nodes - if node.op == "call_function" and node.name == "executorch_call_delegate" - ] - assert len(call_delegate) == 1 - for n in graph_module.graph.nodes: - if QCOM_QUANTIZED_IO in n.meta: - n.meta["val"] = n.meta["val"].to(dtype=n.meta[QCOM_QUANTIZED_IO]) - - def call(self, graph_module: torch.fx.GraphModule): - self._build(graph_module) - graph_module.graph.eliminate_dead_code() - graph_module.recompile() - return PassResult(graph_module, True) + def placeholder(self, name: str, arg, meta): + if quantized_dtype := meta.data.get(QCOM_QUANTIZED_IO, None): + arg = arg.to(dtype=quantized_dtype) + meta["spec"] = self._make_spec(arg) + return super().placeholder(name, arg, meta) + + def call_getitem(self, value, key: int, meta): + meta["spec"] = value.node.meta["spec"][key] + return super().call_getitem(value, key, meta) + + def call_delegate(self, lowered_module, args, kwargs, meta): + args_data, _ = pytree.tree_map_only( + ProxyValue, lambda x: x.data, (args, kwargs) + ) + meta["spec"] = pytree.tree_map( + self._make_spec, + executorch_call_delegate(lowered_module, *args_data), + ) + return super().call_delegate(lowered_module, args, kwargs, meta) diff --git a/backends/qualcomm/_passes/convert_conv1d_to_conv2d.py b/backends/qualcomm/_passes/canonicalize_conv.py similarity index 66% rename from backends/qualcomm/_passes/convert_conv1d_to_conv2d.py rename to backends/qualcomm/_passes/canonicalize_conv.py index 1ee71d42bd4..3804fb05da0 100644 --- a/backends/qualcomm/_passes/convert_conv1d_to_conv2d.py +++ b/backends/qualcomm/_passes/canonicalize_conv.py @@ -4,57 +4,96 @@ # 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 cast, Tuple + import torch + from executorch.backends.qualcomm.builders.utils import get_parameter, set_parameter from executorch.backends.qualcomm.utils.constants import QCOM_REQUANTIZE from executorch.exir.pass_base import ExportPass, PassResult +from torch._guards import detect_fake_mode -from .utils import copy_meta +from .utils import append_qdq, copy_meta -class ConvertConv1dToConv2d(ExportPass): +class CanonicalizeConv(ExportPass): """ - Conv1d is not supported by QNN. - Change it to input -> unsqueeze -> conv2d -> squeeze -> output + 1. QNN does not support dilation on TransposeConvND + Dilate the kernel manually for math-equivalent operation + 2. Conv1d is not supported by QNN. + Change it to input -> unsqueeze -> conv2d -> squeeze -> output """ def __init__(self, edge_program: torch.export.ExportedProgram): - super(ConvertConv1dToConv2d, self).__init__() + super(CanonicalizeConv, self).__init__() self.edge_program = edge_program - self.conv_op_map = { + self.conv1d_op_map = { torch.ops.aten.conv1d.default: torch.ops.aten.conv2d.default, torch.ops.aten.conv_transpose1d.default: torch.ops.aten.conv_transpose2d.input, } + self.transpose_conv_set = { + torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.conv_transpose2d.input, + } - def append_qdq( - self, - graph_module: torch.fx.GraphModule, - node: torch.fx.Node, - qdq_node: torch.fx.Node, - ): - q_op = torch.ops.quantized_decomposed.quantize_per_tensor.default - dq_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default - if qdq_node.target not in {q_op, dq_op}: - return node - - with graph_module.graph.inserting_after(node): - q_args = (node, *qdq_node.args[1:]) - q_node = graph_module.graph.create_node("call_function", q_op, q_args) - q_node.meta = copy_meta(node.meta) - q_node.meta["val"] = q_node.meta["val"].to(q_args[-1]) - with graph_module.graph.inserting_after(q_node): - dq_args = (q_node, *qdq_node.args[1:]) - dq_node = graph_module.graph.create_node( - "call_function", dq_op, dq_args - ) - dq_node.meta = copy_meta(node.meta) - - return dq_node + def dilate(self, tensor, dilation): + # e.g. + # for 3x3 kernel with dilation == (2, 3) + # 1, 0, 0, 2, 0, 0, 3 + # 1, 2, 3 0, 0, 0, 0, 0, 0, 0 + # 4, 5, 6 --> 4, 0, 0, 5, 0, 0, 6 + # 7, 8, 9 0, 0, 0, 0, 0, 0, 0 + # 7, 0, 0, 8, 0, 0, 9 + i, o, *k = tensor.shape + new_k = [dim + (dim - 1) * (s - 1) for s, dim in zip(dilation, k)] + new_tensor = torch.zeros((i, o, *new_k), dtype=tensor.dtype) + indexing = (...,) + tuple([slice(None, None, d) for d in dilation]) + new_tensor[indexing] = tensor + return new_tensor def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph + # condition 1 + for node in graph.nodes: + # arg order (https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.conv_transpose2d.html) + # > input, weight, bias, stride, padding, output_padding, groups, dilation + if node.target in self.transpose_conv_set and len(node.args) > 7: + dilation = cast(Tuple[int], node.args[7]) + # dilate kernel in advance + filter_arg = node.args[1] + filter_node = ( + # fp graph + filter_arg + if filter_arg.op == "placeholder" + # qdq graph + else node.args[1].args[0] + ) + filter_tensor = self.dilate( + get_parameter(filter_node, self.edge_program), + dilation, + ) + # update tensor meta for kernel node + fake_mode = detect_fake_mode(filter_node.meta["val"]) + converter = fake_mode.fake_tensor_converter + filter_node.meta["val"] = converter.from_real_tensor( + fake_mode, filter_tensor + ) + # update kernel + set_parameter( + ( + torch.nn.Parameter(filter_tensor) + if filter_tensor.dtype == torch.float + else filter_tensor + ), + filter_node, + self.edge_program, + ) + # pop dilation for graph in cpu + node.args = node.args[0:-1] + + # condition 2 for node in graph.nodes: - if node.target in self.conv_op_map: + if node.target in self.conv1d_op_map: input_node = node.args[0] with graph_module.graph.inserting_after(input_node): unsqueeze_op = torch.ops.aten.unsqueeze_copy.default @@ -69,7 +108,7 @@ def call(self, graph_module: torch.fx.GraphModule): unsqueeze_node.meta = copy_meta( input_node.meta, lambda m: {**m, "val": m["val"].unsqueeze(2)} ) - qdq_node_after_unsqueeze = self.append_qdq( + qdq_node_after_unsqueeze = append_qdq( graph_module=graph_module, node=unsqueeze_node, qdq_node=input_node, @@ -105,7 +144,7 @@ def call(self, graph_module: torch.fx.GraphModule): padding = [0] + node.args[4] if num_args > 4 else [0, 0] if node.target == torch.ops.aten.conv1d.default: dilation = [1] + node.args[5] if num_args > 5 else [1, 1] - groups = node.args[6] if num_args > 5 else 1 + groups = node.args[6] if num_args > 6 else 1 conv_args = ( qdq_node_after_unsqueeze, node.args[1], @@ -133,13 +172,13 @@ def call(self, graph_module: torch.fx.GraphModule): ) conv2d_node = graph.create_node( "call_function", - self.conv_op_map[node.target], + self.conv1d_op_map[node.target], conv_args, ) conv2d_node.meta = copy_meta( node.meta, lambda m: {**m, "val": m["val"].unsqueeze(2)} ) - qdq_node_after_conv2d = self.append_qdq( + qdq_node_after_conv2d = append_qdq( graph_module=graph_module, node=conv2d_node, qdq_node=list(node.users)[0], diff --git a/backends/qualcomm/_passes/convert_bmm_to_matmul.py b/backends/qualcomm/_passes/convert_bmm_to_matmul.py index 84e1ff26aa1..3d4e44dfa42 100644 --- a/backends/qualcomm/_passes/convert_bmm_to_matmul.py +++ b/backends/qualcomm/_passes/convert_bmm_to_matmul.py @@ -21,7 +21,7 @@ class ConvertBmmToMatmul(ExportPass): view_copy = exir_ops.edge.aten.view_copy.default expand_copy = exir_ops.edge.aten.expand_copy.default - clone = exir_ops.edge.aten.clone.default + clone = exir_ops.edge.dim_order_ops._clone_dim_order.default bmm = exir_ops.edge.aten.bmm.default matmul = exir_ops.edge.aten.matmul.default patterns = [ diff --git a/backends/qualcomm/_passes/convert_linear_to_conv2d.py b/backends/qualcomm/_passes/convert_linear_to_conv2d.py new file mode 100644 index 00000000000..19caa7017d5 --- /dev/null +++ b/backends/qualcomm/_passes/convert_linear_to_conv2d.py @@ -0,0 +1,232 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.backends.qualcomm._passes.utils import append_qdq, copy_meta +from executorch.backends.qualcomm.builders.utils import get_parameter, set_parameter +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule +from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix + + +def _pad_list_to_4(lst): + return lst + [1] * (4 - len(lst)) if len(lst) < 4 else lst[:4] + + +class ConvertLinearToConv2d(ExportPass): + """ + Replace aten.linear.default with equivalent 1x1 conv2d using call_function nodes. + """ + + def __init__(self, edge_program: torch.export.ExportedProgram): + super().__init__() + self.edge_program = edge_program + self.per_block_dq = torch.ops.torchao.dequantize_affine.default + + def _register_tensor( + self, + gm: torch.fx.GraphModule, + node: torch.fx.Node, + tensor_constant: torch.Tensor, + ) -> torch.fx.Node: + new_node_name = get_new_attr_name_with_prefix(node.name)(gm) + gm.register_buffer(new_node_name, tensor_constant) + + with gm.graph.inserting_before(node): + get_attr_node = gm.graph.get_attr(new_node_name) + get_attr_node.meta["val"] = tensor_constant + return get_attr_node + + def _append_dq( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + qdq_node: torch.fx.Node, + ): + q_op = torch.ops.quantized_decomposed.quantize_per_tensor.default + dq_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default + + if qdq_node.target not in {q_op, dq_op}: + return node + + with graph_module.graph.inserting_after(node): + dq_args = (node, *qdq_node.args[1:]) + dq_node = graph_module.graph.create_node("call_function", dq_op, dq_args) + dq_node.meta = copy_meta(node.meta) + return dq_node + + def _create_node( + self, graph_module, target, args, meta_node, new_meta_val, qdq_node + ): + new_node = graph_module.graph.call_function(target, args) + new_node.meta = copy_meta( + meta_node.meta, + lambda m, new_meta_val=new_meta_val: { + **m, + "val": new_meta_val, + }, + ) + dq_node = append_qdq( + graph_module=graph_module, + node=new_node, + qdq_node=qdq_node, + ) + return dq_node + + def _reshape_weight(self, graph_module, weight_node, dq_node): + # After export, constant node will be placeholder from edge_program + weight_val = get_parameter(weight_node, self.edge_program) + assert weight_val is not None, "Cannot get the weight in linear node." + + weight_val = weight_val.reshape(*weight_val.shape, 1, 1) + # Create the new weight node when several node share the same weight + # such as embedding and lm_head in LLM. + if len(list(weight_node.users)) > 1: + weight_node = self._register_tensor(graph_module, weight_node, weight_val) + dq_node = self._append_dq(graph_module, weight_node, dq_node) + else: + set_parameter( + ( + torch.nn.Parameter(weight_val) + if weight_val.dtype == torch.float + else weight_val + ), + weight_node, + self.edge_program, + ) + + # Update node meta val + weight_node.meta["val"] = weight_node.meta["val"].reshape(weight_val.shape) + dq_node.meta["val"] = dq_node.meta["val"].reshape(weight_val.shape) + # Update block size for per-block quant + if dq_node.target == self.per_block_dq: + new_args = list(dq_node.args) + # pad block size + new_args[1] = _pad_list_to_4(list(new_args[1])) + dq_node.args = tuple(new_args) + + return dq_node + + def call(self, graph_module: GraphModule): + graph = graph_module.graph + + for node in list(graph.nodes): + if node.target == torch.ops.aten.linear.default: + input_node = node.args[0] + # In quantization flow, weight_arg will be dq node. + weight_arg = node.args[1] + weight_node = ( + weight_arg if weight_arg.op == "placeholder" else weight_arg.args[0] + ) + bias_arg = node.args[2] if len(node.args) > 2 else None + + input_meta_val = input_node.meta["val"] + output_meta_val = node.meta["val"] + if bias_arg: + bias_meta_val = bias_arg.meta["val"] + + rank = input_meta_val.ndim + with graph.inserting_before(node): + # Step 1: reshape input + # rank = 2: (dim, C) -> (1, C, 1, dim) + # rank = 3: (N, dim, C) -> (N, C, 1, dim) + # rank = 4: (N, H, W, C) -> (N, C, H, W) + order = (0, 3, 1, 2) + if rank <= 3: + # (dim, C) -> (1, C, 1, dim) + # (N, dim, C) -> (N, C, 1, dim) + shape = ( + (1, *input_meta_val.shape, 1) + if rank == 2 + else (*input_meta_val.shape, 1) + ) + x_meta_val = input_meta_val.reshape(shape) + input_node = self._create_node( + graph_module, + torch.ops.aten.reshape.default, + (input_node, shape), + node, + x_meta_val, + input_node, + ) + order = (0, 2, 3, 1) + + x_meta_val = x_meta_val.permute(order) + x = self._create_node( + graph_module, + torch.ops.aten.permute.default, + (input_node, order), + node, + x_meta_val, + input_node, + ) + + # Step 2: reshape weight + weight_arg = self._reshape_weight( + graph_module, weight_node, weight_arg + ) + weight_meta_val = weight_arg.meta["val"] + + conv_args = [x, weight_arg] + conv_args_meta_val = [x_meta_val, weight_meta_val] + if bias_arg: + conv_args.append(bias_arg) + conv_args_meta_val.append(bias_meta_val) + else: + conv_args.append(None) + conv_args_meta_val.append(None) + + conv_args.extend( + [[1, 1], [0, 0], [1, 1], 1] + ) # stride, padding, dilation, groups + conv_node_val = torch.nn.functional.conv2d( + *conv_args_meta_val, + stride=(1, 1), + padding=(0, 0), + dilation=(1, 1), + groups=1, + ) + conv_node = self._create_node( + graph_module, + torch.ops.aten.conv2d.default, + tuple(conv_args), + node, + conv_node_val, + list(node.users)[0], + ) + + # Step 3: restore shape + # rank = 2: (1, C, 1, dim) -> (dim, C) + # rank = 3: (N, C, 1, dim) -> (N, dim C) + # rank = 4: (N, C, H, W) -> (N, H, W, C) + order = (0, 2, 3, 1) if rank == 4 else (0, 3, 1, 2) + y_meta_val = conv_node_val.permute(order) + y = self._create_node( + graph_module, + torch.ops.aten.permute.default, + (conv_node, order), + node, + y_meta_val, + list(node.users)[0], + ) + if rank <= 3: + target_shape = output_meta_val.shape + y_meta_val = y_meta_val.reshape(target_shape) + y = self._create_node( + graph_module, + torch.ops.aten.reshape.default, + (y, target_shape), + node, + y_meta_val, + list(node.users)[0], + ) + + node.replace_all_uses_with(y) + graph.erase_node(node) + + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/fixed_linear_keep_dim.py b/backends/qualcomm/_passes/fixed_linear_keep_dim.py index 4f625b96f0e..19f5c631921 100644 --- a/backends/qualcomm/_passes/fixed_linear_keep_dim.py +++ b/backends/qualcomm/_passes/fixed_linear_keep_dim.py @@ -9,8 +9,6 @@ from executorch.exir.pass_base import ExportPass, PassResult from executorch.exir.passes import dead_code_elimination_pass -from torch.fx.passes.utils.source_matcher_utils import get_source_partitions - class FixedLinearKeepDim(ExportPass): """ @@ -24,61 +22,58 @@ def __init__(self): super(FixedLinearKeepDim, self).__init__() def _fixed_keep_dim(self, graph_module: torch.fx.GraphModule): - partitions = get_source_partitions( - graph_module.graph, [torch.nn.Linear, torch.ops.aten.linear.default] - ) - for _, src_partitions in partitions.items(): - for src_partition in src_partitions: - linear_node = [ - n for n in src_partition.nodes if n.target == self.linear - ][0] - input_node = linear_node.args[0] - # Since QNN has no keep dims for linear op, we will need to add squeeze and unsqueeze around linear node - # TODO: Find a more general conditional statement. - linear_output = linear_node.meta["val"] - if linear_output.dim() >= 3: - with graph_module.graph.inserting_after(input_node): - input_users = list(input_node.users.keys()) - input_tensor = input_node.meta["val"] - squeeze_dim = (-1, input_tensor.shape[-1]) - squeeze_node = graph_module.graph.create_node( - "call_function", - self.view_copy, - ( - input_node, - squeeze_dim, - ), - ) - # meta needs to be copied elementwisely for fake-tensor - # to be updated correctly and not affect meta of input_node - for k, v in input_node.meta.items(): - squeeze_node.meta[k] = v - squeeze_node.meta["val"] = input_tensor.reshape(squeeze_dim) - for user in input_users: - if user == linear_node: - user.replace_input_with(input_node, squeeze_node) + for node in graph_module.graph.nodes: + if node.target != self.linear: + continue + + linear_node = node + input_node = linear_node.args[0] + # Since QNN has no keep dims for linear op, we will need to add squeeze and unsqueeze around linear node + # TODO: Find a more general conditional statement. + linear_output = linear_node.meta["val"] + if linear_output.dim() >= 3: + with graph_module.graph.inserting_after(input_node): + input_users = list(input_node.users.keys()) + input_tensor = input_node.meta["val"] + squeeze_dim = (-1, input_tensor.shape[-1]) + squeeze_node = graph_module.graph.create_node( + "call_function", + self.view_copy, + ( + input_node, + squeeze_dim, + ), + ) + # meta needs to be copied elementwisely for fake-tensor + # to be updated correctly and not affect meta of input_node + for k, v in input_node.meta.items(): + squeeze_node.meta[k] = v + squeeze_node.meta["val"] = input_tensor.reshape(squeeze_dim) + for user in input_users: + if user == linear_node: + user.replace_input_with(input_node, squeeze_node) - with graph_module.graph.inserting_after(linear_node): - output_users = list(linear_node.users.keys()) - unsqueeze_dim = linear_output.shape - unsqueeze_node = graph_module.graph.create_node( - "call_function", - self.view_copy, - ( - linear_node, - unsqueeze_dim, - ), - ) - # meta needs to be copied elementwisely for fake-tensor - # to be updated correctly and not affect meta of unsqueeze_node - for k, v in linear_node.meta.items(): - unsqueeze_node.meta[k] = v - # update linear node's shape - linear_node.meta["val"] = linear_output.reshape( - (squeeze_node.meta["val"].shape[0], linear_output.shape[-1]) - ) - for user in output_users: - user.replace_input_with(linear_node, unsqueeze_node) + with graph_module.graph.inserting_after(linear_node): + output_users = list(linear_node.users.keys()) + unsqueeze_dim = linear_output.shape + unsqueeze_node = graph_module.graph.create_node( + "call_function", + self.view_copy, + ( + linear_node, + unsqueeze_dim, + ), + ) + # meta needs to be copied elementwisely for fake-tensor + # to be updated correctly and not affect meta of unsqueeze_node + for k, v in linear_node.meta.items(): + unsqueeze_node.meta[k] = v + # update linear node's shape + linear_node.meta["val"] = linear_output.reshape( + (squeeze_node.meta["val"].shape[0], linear_output.shape[-1]) + ) + for user in output_users: + user.replace_input_with(linear_node, unsqueeze_node) def call(self, graph_module: torch.fx.GraphModule): self._fixed_keep_dim(graph_module) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 13175fe41bd..f285eb79bfb 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -64,8 +64,10 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.amax.default, exir_ops.edge.aten.amin.default, + exir_ops.edge.aten.asin.default, exir_ops.edge.aten.atan.default, exir_ops.edge.aten.bitwise_or.Tensor, + exir_ops.edge.aten.bitwise_xor.Tensor, exir_ops.edge.aten.bmm.default, exir_ops.edge.aten.bitwise_and.Tensor, exir_ops.edge.aten.cat.default, @@ -77,7 +79,9 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.elu.default, exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.exp.default, + exir_ops.edge.aten.flip.default, exir_ops.edge.aten.floor.default, + exir_ops.edge.aten.floor_divide.default, exir_ops.edge.aten.full.default, exir_ops.edge.aten.full_like.default, exir_ops.edge.aten.ge.Tensor, @@ -107,6 +111,8 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.relu.default, exir_ops.edge.aten.round.default, exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.sign.default, + exir_ops.edge.aten.slice_copy.Tensor, exir_ops.edge.aten.split_with_sizes.default, exir_ops.edge.aten.split_with_sizes_copy.default, exir_ops.edge.aten.sqrt.default, @@ -175,6 +181,7 @@ def is_layout_agnostic(self, node: torch.fx.Node) -> bool: exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.min.dim, exir_ops.edge.aten.sum.dim_IntList, + exir_ops.edge.aten.amax.default, }: # if dimemsion is not kept, we'll have no clue how to do layout transform if len(node.args) < 3 or not node.args[2]: diff --git a/backends/qualcomm/_passes/lift_constant_scalar_operands.py b/backends/qualcomm/_passes/lift_constant_scalar_operands.py index dc9592e415b..f5c5915cab2 100644 --- a/backends/qualcomm/_passes/lift_constant_scalar_operands.py +++ b/backends/qualcomm/_passes/lift_constant_scalar_operands.py @@ -54,6 +54,7 @@ class TensorOpInfo: aten.where.ScalarOther: TensorOpInfo(aten.where.self, False, True), aten.where.Scalar: TensorOpInfo(aten.where.self, False, True), aten.masked_fill.Scalar: TensorOpInfo(aten.masked_fill.Tensor, False, False), + aten.bitwise_xor.Scalar: TensorOpInfo(aten.bitwise_xor.Tensor, False, False), } @@ -64,6 +65,7 @@ class TensorOpInfo: aten.arange.default, aten.scalar_tensor.default, aten.elu.default, + aten.hardtanh.default, } diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 152433195cd..ffb9f3221df 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -13,8 +13,9 @@ AnnotateQuantAttrs, AnnotateStack, AnnotateUnbind, + CanonicalizeConv, ConvertBmmToMatmul, - ConvertConv1dToConv2d, + ConvertLinearToConv2d, ConvertSquareToPow, DecomposeAny, DecomposeCDist, @@ -81,8 +82,8 @@ def get_capture_program_passes(): (AnnotateQuantAttrs, True), (AnnotateStack, True), (AnnotateUnbind, True), + (CanonicalizeConv, True), (ConvertBmmToMatmul, False), - (ConvertConv1dToConv2d, True), (DecomposeAny, True), (DecomposeColIm, True), (DecomposeMinMaxDim, True), @@ -92,7 +93,7 @@ def get_capture_program_passes(): (I64toI32, True), (LayoutTransform, True), (RecomposePixelUnshuffle, True), - (RecomposeRmsNorm, False), + (RecomposeRmsNorm, True), (Remove0DTensor, True), (RemoveRedundancy, True), (TagQuantIO, False), @@ -190,6 +191,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): self.add_pass(RemoveRedundancy(quantization_capture=True)) self.add_pass(ReduceDynamicRange()) self.add_pass(RecomposePixelUnshuffle(quantization_capture=True)) + self.add_pass(RecomposeRmsNorm(quantization_capture=True)) self.add_pass(ReplaceArangeArgs()) self.add_pass(DecomposeCDist()) self.add_pass(DecomposeScaledDotProductAttention()) @@ -203,7 +205,9 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): self.add_pass(LiftConstantScalarOperands()) return self._transform(graph_module) - def transform_for_export_pipeline(self, exported_program: ExportedProgram): + def transform_for_export_pipeline( + self, exported_program: ExportedProgram, convert_linear_to_conv2d: bool = False + ): self.add_pass(DecomposeCDist()) self.add_pass(DecomposeScaledDotProductAttention()) self.add_pass(DecomposeRoll()) @@ -212,7 +216,9 @@ def transform_for_export_pipeline(self, exported_program: ExportedProgram): self.add_pass(DecomposeWrapWithAutocast()) # this pass will rewrite state_dict, it needs to be accomplished before # to_edge_transform_and_lower - self.add_pass(ConvertConv1dToConv2d(exported_program)) + self.add_pass(CanonicalizeConv(exported_program)) + if convert_linear_to_conv2d: + self.add_pass(ConvertLinearToConv2d(exported_program)) self.add_pass(ConvertSquareToPow()) self.add_pass(LiftConstantScalarOperands()) self._transform(exported_program.graph_module) diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index 2e2063cdf6e..7c5b864377e 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -4,70 +4,139 @@ # 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.qualcomm._passes.utils import find_patterns from executorch.backends.qualcomm.builders.node_visitor import dq_ops -from executorch.backends.qualcomm.builders.utils import get_parameter, is_parameter from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult -from torch.fx.passes.utils.source_matcher_utils import get_source_partitions + + +def _is_node(node): + return isinstance(node, torch.fx.Node) + + +def _is_call(node): + return _is_node(node) and node.op == "call_function" + + +def _is_placeholder(node): + return _is_node(node) and node.op == "placeholder" + + +def _is_get_attr(node): + return _is_node(node) and node.op == "get_attr" + + +def _is_add(node): + return _is_call(node) and node.target in [ + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.add.Scalar, + torch.ops.aten.add.Tensor, + torch.ops.aten.add.Scalar, + ] + + +def _is_dq(node): + return _is_call(node) and node.target in dq_ops + + +def _is_mean(node): + return _is_call(node) and node.target in [ + exir_ops.edge.aten.mean.dim, + torch.ops.aten.mean.dim, + ] + + +def _is_mul(node): + return _is_call(node) and node.target in [ + exir_ops.edge.aten.mul.Tensor, + torch.ops.aten.mul.Tensor, + ] + + +def _is_pow(node): + return _is_call(node) and node.target in [ + exir_ops.edge.aten.pow.Tensor_Tensor, + exir_ops.edge.aten.pow.Tensor_Scalar, + torch.ops.aten.pow.Tensor_Scalar, + ] + + +def _is_rsqrt(node): + return _is_call(node) and node.target in [ + exir_ops.edge.aten.rsqrt.default, + torch.ops.aten.rsqrt.default, + ] class RecomposeRmsNorm(ExportPass): """ Merge decomposed operators back to one super node. - TODO: After replacing export_to_edge with to_edge_transform_and_lowering - in examples/models/llama/export_llama_lib.py, this pass can be removed """ - def __init__(self, edge_program: torch.export.ExportedProgram): + def __init__(self, quantization_capture=False): super(RecomposeRmsNorm, self).__init__() - self.edge_program = edge_program + self.rms_norm_target = exir_ops.edge.aten.rms_norm.default + self.skip_targets = [ + exir_ops.edge.aten.to.dtype, + ] + self.quantization_capture = quantization_capture + if quantization_capture: + self.rms_norm_target = torch.ops.aten.rms_norm.default + self.skip_targets = [ + torch.ops.aten.to.dtype, + ] - def _get_eps_node(self, nodes): - # eps: one of inputs of add node - add_node = [n for n in nodes if hasattr(n, "name") and "add" in n.name][0] - for a in add_node.args: - if isinstance(a, float) or a.op != "call_function": - return a - - def _get_gamma_node(self, output_node): - # gamma: one of inputs of output node - for a in output_node.args: - if a.op != "call_function" or a.target in dq_ops: - return a + def _get_input_node(self, node): + input_node = node.args[0] + while input_node.target in self.skip_targets: + input_node = input_node.args[0] + return input_node def call(self, graph_module: torch.fx.GraphModule): graph = graph_module.graph - partitions = get_source_partitions( - graph, [torch.nn.RMSNorm, torch.ops.aten.rms_norm.default] - ) - for _, src_partitions in partitions.items(): - for src_partition in src_partitions: - input_len = len(src_partition.input_nodes) - if input_len == 1: - input_node = src_partition.input_nodes[0] - elif input_len == 2: - inp_0, inp_1 = src_partition.input_nodes - input_node = inp_0 if len(inp_0.users) == 2 else inp_1 - else: - raise RuntimeError( - f"Found a edge case of rms_node partition {src_partition}, which has {input_len} inputs" - ) - output_node = src_partition.output_nodes[0] - eps = self._get_eps_node(src_partition.nodes) - if isinstance(eps, torch.fx.Node) and is_parameter( - eps, self.edge_program - ): - eps = get_parameter(eps, self.edge_program).item() - gamma_node = self._get_gamma_node(output_node) + # Root Mean Square normalization math equivalent implementation + patterns = [ + # transformers.models.qwen2.modeling_qwen2.Qwen2RMSNorm + [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_pow], + # executorch.examples.models.llama.norm.RMSNorm + [_is_mul, "*", _is_mul, _is_rsqrt, _is_add, _is_mean, _is_mul], + ] + + for node in graph.nodes: + if not _is_mul(node): + continue + + rms_norm_patterns = [ + pattern + for pattern in find_patterns(node, patterns) + if pattern is not None + ] + + if len(rms_norm_patterns) > 0: + # Use first matched pattern + rms_norm_pattern = rms_norm_patterns[0][0] + last_mul_node = rms_norm_pattern[0] + gamma_node = None + # weight should be a constant + for arg in last_mul_node.args: + if _is_get_attr(arg) or _is_placeholder(arg) or _is_dq(arg): + gamma_node = arg + + if gamma_node is None: + continue + + eps = rms_norm_pattern[4].args[1] + if isinstance(eps, torch.fx.Node): + eps = eps.meta["val"].constant.item() + input_node = self._get_input_node(rms_norm_pattern[6]) - with graph.inserting_before(output_node): + with graph.inserting_before(last_mul_node): # args schema # (Tensor input, int[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor rms_node = graph.create_node( "call_function", - exir_ops.edge.aten.rms_norm.default, + self.rms_norm_target, ( input_node, list(gamma_node.meta["val"].shape), @@ -75,11 +144,11 @@ def call(self, graph_module: torch.fx.GraphModule): eps, ), ) - users = output_node.users.copy() + users = last_mul_node.users.copy() for user in users: - user.replace_input_with(output_node, rms_node) + user.replace_input_with(last_mul_node, rms_node) # copy metadata - rms_node.meta = output_node.meta + rms_node.meta = last_mul_node.meta graph.eliminate_dead_code() graph_module.recompile() diff --git a/backends/qualcomm/_passes/remove_redundancy.py b/backends/qualcomm/_passes/remove_redundancy.py index 2ec8161613b..d75637beaf5 100644 --- a/backends/qualcomm/_passes/remove_redundancy.py +++ b/backends/qualcomm/_passes/remove_redundancy.py @@ -19,7 +19,7 @@ def __init__(self, quantization_capture=False): self.redundant_ops_general = { torch.clone: self._default_condition, torch.ops.aten.clone.default: self._default_condition, - exir_ops.edge.aten.clone.default: self._default_condition, + exir_ops.edge.dim_order_ops._clone_dim_order.default: self._default_condition, torch.ops.aten.alias.default: self._default_condition, exir_ops.edge.aten.alias.default: self._default_condition, exir_ops.edge.aten.alias_copy.default: self._default_condition, @@ -41,12 +41,7 @@ def __init__(self, quantization_capture=False): ) def _dim_order_op_condition(self, node): - dim_order = node.kwargs.get("dim_order") - # skip if there contains layout hint - # e.g. (0, 2, 3, 1) != (0, 1, 2, 3) - if node.meta["val"].dtype != node.args[0].meta["val"].dtype: - return False - return dim_order != list(range(len(dim_order))) + return node.meta["val"].dtype == node.args[0].meta["val"].dtype def _to_copy_op_condition(self, node): return "memory_format" in node.kwargs diff --git a/backends/qualcomm/_passes/seq_mse.py b/backends/qualcomm/_passes/seq_mse.py new file mode 100644 index 00000000000..a1d8343c928 --- /dev/null +++ b/backends/qualcomm/_passes/seq_mse.py @@ -0,0 +1,223 @@ +# 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 types +from contextlib import contextmanager + +import torch +import torchao +from executorch.backends.qualcomm.quantizer.observers.per_block_param_observer import ( + PerBlockParamObserver, +) +from executorch.exir.pass_base import ExportPass, PassResult +from torchao.quantization.pt2e import PerChannelMinMaxObserver + + +class SeqMseModule(torch.nn.Module): + """ + Args: + nominal_weight: Tensor + nominal parameters from operator + nominal_bias: Tensor + nominal parameters from operator + operator: fx.Node + operator to be executed + observer: UniformQuantizationObserverBase + parameter observer (specific for weight) + num_candidates: int + grids to search minimal mse loss + """ + + def __init__( + self, + nominal_weight, + nominal_bias, + operator, + observer, + num_candidates, + ): + super().__init__() + self.nominal_weight = nominal_weight + self.nominal_bias = nominal_bias + self.observer = observer + self.steps = torch.linspace( + 1 / num_candidates, 1, steps=num_candidates + ).tolist() + self.operator = self._make_operator(operator) + self.best_candidate_step = 1.0 + + def _make_operator(self, aten_op): + if aten_op.target == torch.ops.aten.conv2d.default: + stride = [1, 1] if len(aten_op.args) < 4 else aten_op.args[3] + padding = [0, 0] if len(aten_op.args) < 5 else aten_op.args[4] + dilation = [1, 1] if len(aten_op.args) < 6 else aten_op.args[5] + groups = 1 if len(aten_op.args) < 7 else aten_op.args[6] + has_bias = self.nominal_bias is not None + module = torch.nn.Conv2d( + in_channels=self.nominal_weight.shape[1], + out_channels=self.nominal_weight.shape[0], + kernel_size=self.nominal_weight.shape[-2:], + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=has_bias, + ) + module.weight.data = self.nominal_weight + if has_bias: + module.bias.data = self.nominal_bias + return module + else: + raise NotImplementedError(f"target of {aten_op.target} is not implemented") + + def _per_block_qdq(self, scale, zero_point): + return torchao.quantization.quant_primitives._fake_quantize_affine( + input=self.nominal_weight, + block_size=self.observer.block_size, + scale=scale, + zero_point=zero_point, + quant_dtype=self.observer.dtype, + quant_min=self.observer.quant_min, + quant_max=self.observer.quant_max, + ) + + def _per_channel_qdq(self, scale, zero_point): + return torch.fake_quantize_per_channel_affine( + input=self.nominal_weight, + scale=scale, + zero_point=zero_point, + axis=0, + quant_min=self.observer.quant_min, + quant_max=self.observer.quant_max, + ) + + def _fake_quant(self, scale, zero_point): + dispatcher = { + PerChannelMinMaxObserver: self._per_channel_qdq, + PerBlockParamObserver: self._per_block_qdq, + } + return dispatcher[type(self.observer)](scale, zero_point) + + def _find_best_candidate(self, nominal_input, nominal_output): + # calculate current baseline + scale, zero_point = self.observer.calculate_qparams() + zero_point = zero_point.to(torch.int32) + self.operator.weight.data = self._fake_quant(scale, zero_point) + candidate, current_loss = ( + 1, + torch.nn.functional.mse_loss( + self.operator(nominal_input), nominal_output + ).item(), + ) + for step in self.steps: + self.operator.weight.data = self._fake_quant(scale * step, zero_point) + loss = torch.nn.functional.mse_loss( + self.operator(nominal_input), nominal_output + ).item() + if loss < current_loss: + candidate, current_loss = step, loss + return candidate + + def forward(self, nominal_input, nominal_output): + self.best_candidate_step = self._find_best_candidate( + nominal_input=nominal_input, nominal_output=nominal_output + ) + + +class InsertSeqMse(ExportPass): + """ + Insert Seq Mse Observer to find the best quant config for certain node's weight. + """ + + seq_mse_ops = {torch.ops.aten.conv2d.default} + + def __init__(self, num_candidates=1000): + super(InsertSeqMse, self).__init__() + self.num_candidates = num_candidates + + def _insert_seq_mse( + self, graph_module: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + count = 0 + for node in graph_module.graph.nodes: + if node.target in self.seq_mse_ops: + # extract observer + weight_node_obs = node.args[1] + observer = getattr(graph_module, weight_node_obs.name) + # extract parameters + weight_node = weight_node_obs.args[0] + weight_tensor = graph_module.get_parameter(weight_node.target).detach() + bias_tensor = None + if len(node.args) > 2 and node.args[2] is not None: + bias_tensor = graph_module.get_parameter( + node.args[2].args[0].target + ).detach() + + with graph_module.graph.inserting_after(node): + seq_mse_mod = SeqMseModule( + nominal_weight=weight_tensor, + nominal_bias=bias_tensor, + operator=node, + observer=observer, + num_candidates=self.num_candidates, + ) + module_name = f"seq_mse_{count}" + count += 1 + setattr(graph_module, module_name, seq_mse_mod) + input_nodes = (node.args[0], node) + graph_module.graph.create_node( + "call_module", module_name, input_nodes, {} + ) + + def call(self, graph_module: torch.fx.GraphModule): + self._insert_seq_mse(graph_module) + graph_module.recompile() + return PassResult(graph_module, True) + + +class RemoveSeqMse(ExportPass): + """ + Remove Seq Mse before invoking convert_pt2e and update final quantization encoding. + """ + + def __init__(self): + super(RemoveSeqMse, self).__init__() + + def _remove_seq_mse( + self, graph_module: torch.fx.GraphModule + ) -> torch.fx.GraphModule: + node_to_erase = [] + for node in graph_module.graph.nodes: + if node.op == "call_module": + # try extracting SeqMse module + module = getattr(graph_module, node.target) + if isinstance(module, SeqMseModule): + # rewrite observer method for pre-calculated scale + scale, zero_point = module.observer.calculate_qparams() + module.observer.updated_encoding = ( + scale * module.best_candidate_step, + zero_point, + ) + module.observer.calculate_qparams = types.MethodType( + lambda s: s.updated_encoding, module.observer + ) + node_to_erase.append(node) + + for node in node_to_erase: + graph_module.graph.erase_node(node) + + def call(self, graph_module: torch.fx.GraphModule): + self._remove_seq_mse(graph_module) + graph_module.recompile() + return PassResult(graph_module, True) + + +@contextmanager +def SeqMSE(prepared_gm, num_candidates): + prepared_gm = InsertSeqMse(num_candidates)(prepared_gm).graph_module + try: + yield + finally: + prepared_gm = RemoveSeqMse()(prepared_gm).graph_module diff --git a/backends/qualcomm/_passes/utils.py b/backends/qualcomm/_passes/utils.py index ae11ba7b325..6d908707892 100755 --- a/backends/qualcomm/_passes/utils.py +++ b/backends/qualcomm/_passes/utils.py @@ -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. -from typing import Dict +from typing import Callable, Dict, List import torch from executorch.backends.qualcomm.builders.utils import get_parameter @@ -64,8 +64,8 @@ def get_passes_dependency_for_capture_program(): AnnotateQuantAttrs, AnnotateStack, AnnotateUnbind, + CanonicalizeConv, ConvertBmmToMatmul, - ConvertConv1dToConv2d, DecomposeAny, DecomposeColIm, DecomposeLinalgVectorNorm, @@ -99,7 +99,7 @@ def get_passes_dependency_for_capture_program(): I64toI32: [RemoveRedundancy], LayoutTransform: [ AnnotateQuantAttrs, - ConvertConv1dToConv2d, + CanonicalizeConv, ExpandBroadcastTensorShape, FixedLinearKeepDim, ], @@ -121,3 +121,118 @@ def is_float_tensor(node: torch.fx.Node) -> bool: if "val" not in node.meta or not isinstance(node.meta["val"], FakeTensor): return False return node.meta["val"].dtype == torch.float32 + + +def _is_node(node): + return isinstance(node, torch.fx.Node) + + +def _pred(node, pat): + return isinstance(pat, Callable) and pat(node) + + +def _next(node, from_args=True): + if from_args: + yield from [i for i in node.args if _is_node(i)] + else: + yield from list(node.users) + + +def _find_pattern( + node: torch.fx.Node, + pattern: List[Callable[[torch.fx.Node], bool] | str], + from_args: bool = True, + max_wildcard_life: int = 3, + verbose: bool = False, +): + """ + Implement wildcard pattern matching + - node: fx.Node + - pattern: predicate list, can contain followings + Callable(fx.node): predicate + '*': wildcard + - from_args: if True find from node.args, otherwise from node.users + - max_wildcard_life: max number of skips for wildcard + + If not matched, return None. + Otherwise, return list of matched node list, which is the same length as pattern + """ + + asterisk = "*" + + def _probe( + cur, hist, pat_idx, asterisk_life_count=max_wildcard_life, verbose=verbose + ): + if pat_idx == len(pattern): + # Expected len(hist) is equal to pat_idx + assert len(hist) == len(pattern) + if list(hist) not in matched: + matched.append(list(hist)) + return + if verbose: + print( + f"cur:{cur}, idx:{pat_idx}, life={asterisk_life_count}, pattern:{pattern[pat_idx]} hist={hist}" + ) + if _pred(cur, pattern[pat_idx]): + hist.append(cur) + for child in _next(cur, from_args): + _probe(child, hist, pat_idx + 1) + hist.pop() + elif pattern[pat_idx] == asterisk and asterisk_life_count > 0: + # 3 cases: ignore/consume/keep asterisk + # 1, Ignore asterisk + hist.append(None) + _probe(cur, hist, pat_idx + 1) + hist.pop() + + # 2. Consume asterisk + hist.append(None) + for child in _next(cur, from_args): + _probe(child, hist, pat_idx + 1) + hist.pop() + + # 3. keep asterisk and skip to next node + for child in _next(cur, from_args): + _probe(child, hist, pat_idx, asterisk_life_count - 1) + + # Check if pattern is valid + assert all( + isinstance(i, Callable) or (isinstance(i, str) and i == "*") for i in pattern + ), f"Invalid pattern: {pattern}" + + # Start probing + matched = [] + _probe(node, [], 0) + return matched if matched else None + + +def find_patterns(node, patterns, **kwargs): + assert isinstance(patterns, list) and isinstance(patterns[0], list) + results = [] + for pattern in patterns: + result = _find_pattern(node, pattern, **kwargs) + results.append(result) + return results + + +def append_qdq( + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + qdq_node: torch.fx.Node, +): + q_op = torch.ops.quantized_decomposed.quantize_per_tensor.default + dq_op = torch.ops.quantized_decomposed.dequantize_per_tensor.default + + if qdq_node.target not in {q_op, dq_op}: + return node + + with graph_module.graph.inserting_after(node): + q_args = (node, *qdq_node.args[1:]) + q_node = graph_module.graph.create_node("call_function", q_op, q_args) + q_node.meta = copy_meta(node.meta) + q_node.meta["val"] = q_node.meta["val"].to(q_args[-1]) + with graph_module.graph.inserting_after(q_node): + dq_args = (q_node, *qdq_node.args[1:]) + dq_node = graph_module.graph.create_node("call_function", dq_op, dq_args) + dq_node.meta = copy_meta(node.meta) + return dq_node diff --git a/backends/qualcomm/aot/ir/CMakeLists.txt b/backends/qualcomm/aot/ir/CMakeLists.txt deleted file mode 100755 index 48cb07c5dd2..00000000000 --- a/backends/qualcomm/aot/ir/CMakeLists.txt +++ /dev/null @@ -1,11 +0,0 @@ -# Copyright (c) Qualcomm Innovation Center, Inc. -# All rights reserved -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -# QCIR -target_sources( - qcir_utils PRIVATE ${CMAKE_CURRENT_LIST_DIR}/qcir_utils.h - ${CMAKE_CURRENT_LIST_DIR}/qcir_utils.cpp -) diff --git a/backends/qualcomm/aot/ir/qcir.fbs b/backends/qualcomm/aot/ir/qcir.fbs deleted file mode 100755 index 82e56c405cc..00000000000 --- a/backends/qualcomm/aot/ir/qcir.fbs +++ /dev/null @@ -1,119 +0,0 @@ -// -// Copyright (c) Qualcomm Innovation Center, Inc. -// All rights reserved. -// -// This source code is licensed under the BSD-style license found in the -// LICENSE file in the root directory of this source tree. -// - -namespace qcir; - -enum TensorType : byte { - WRITE = 0, - READ, - READWRITE, - NATIVE, - STATIC, - OPTIONAL, - UNDEFINED, -} - -enum DataType : byte { - INT8 = 0, - INT16, - INT32, - INT64, - UINT8, - UINT16, - UINT32, - UINT64, - FLOAT16, - FLOAT32, - FLOAT64, - SFIXED4, - SFIXED8, - SFIXED16, - SFIXED32, - UFIXED4, - UFIXED8, - UFIXED16, - UFIXED32, - BOOL, - STRING, - UNDEFINED, -} - -enum QuantizeDef : byte { - IMPL_GENERATED = 0, - DEFINED, - UNDEFINED, -} - -enum QuantizeType : byte { - SCALE_OFFSET = 0, - AXIS_SCALE_OFFSET, - BW_SCALE_OFFSET, - BW_AXIS_SCALE_OFFSET, - BLOCKWISE_EXPANSION, - UNDEFINED, -} - -enum BlockScaleStorageType: byte { - BITWIDTH_SCALE_STORAGE_8 = 0, - BITWIDTH_SCALE_STORAGE_16, - UNDEFINED, -} - -struct ScaleOffset { - scale: float; - offset: int; -} - -table QuantizeParam { - def: QuantizeDef; - type: QuantizeType; - bitwidth: uint; - axis: int; - // used by bitwidth quantization - scales: [float]; - offsets: [int]; - // used by general quantization - data: [ScaleOffset]; - // used by block quantization - num_blocks_per_axis: uint; - block_scale_storage_type: BlockScaleStorageType; - block_scale: [ubyte]; -} - -table Tensor { - name: string; - shape: [uint]; - dynamic_dims: [ubyte]; - type: TensorType; - dtype: DataType; - qparam: QuantizeParam; - size: uint; - offset: ulong; -} - -table Operator { - name: string; - package_name: string; - type_name: string; - // keep only tensor indexes - inputs: [uint]; - outputs: [uint]; - params: [uint]; -} - -table Graph { - name: string; - nodes: [Operator]; - tensors: [Tensor]; -} - -table Context { - graphs: [Graph]; -} - -root_type Context; diff --git a/backends/qualcomm/aot/ir/qcir_utils.cpp b/backends/qualcomm/aot/ir/qcir_utils.cpp deleted file mode 100755 index de9e349abe7..00000000000 --- a/backends/qualcomm/aot/ir/qcir_utils.cpp +++ /dev/null @@ -1,345 +0,0 @@ -/* - * Copyright (c) Qualcomm Innovation Center, Inc. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#include -#include - -#include - -namespace executorch { -namespace backends { -namespace qnn { - -qcir::TensorType ToTensorType(Qnn_TensorType_t type) { - static const std::unordered_map type_map{ - {QNN_TENSOR_TYPE_APP_WRITE, qcir::TensorType::WRITE}, - {QNN_TENSOR_TYPE_APP_READ, qcir::TensorType::READ}, - {QNN_TENSOR_TYPE_APP_READWRITE, qcir::TensorType::READWRITE}, - {QNN_TENSOR_TYPE_NATIVE, qcir::TensorType::NATIVE}, - {QNN_TENSOR_TYPE_STATIC, qcir::TensorType::STATIC}, - {QNN_TENSOR_TYPE_NULL, qcir::TensorType::OPTIONAL}, - {QNN_TENSOR_TYPE_UNDEFINED, qcir::TensorType::UNDEFINED}, - }; - return type_map.at(type); -} - -Qnn_TensorType_t ToTensorType(qcir::TensorType type) { - static const std::unordered_map type_map{ - {qcir::TensorType::WRITE, QNN_TENSOR_TYPE_APP_WRITE}, - {qcir::TensorType::READ, QNN_TENSOR_TYPE_APP_READ}, - {qcir::TensorType::READWRITE, QNN_TENSOR_TYPE_APP_READWRITE}, - {qcir::TensorType::NATIVE, QNN_TENSOR_TYPE_NATIVE}, - {qcir::TensorType::STATIC, QNN_TENSOR_TYPE_STATIC}, - {qcir::TensorType::OPTIONAL, QNN_TENSOR_TYPE_NULL}, - {qcir::TensorType::UNDEFINED, QNN_TENSOR_TYPE_UNDEFINED}, - }; - return type_map.at(type); -} - -// TODO: enable commented type by QNN version control -qcir::DataType ToDataType(Qnn_DataType_t type) { - static const std::unordered_map type_map{ - {QNN_DATATYPE_INT_8, qcir::DataType::INT8}, - {QNN_DATATYPE_INT_16, qcir::DataType::INT16}, - {QNN_DATATYPE_INT_32, qcir::DataType::INT32}, - {QNN_DATATYPE_INT_64, qcir::DataType::INT64}, - {QNN_DATATYPE_UINT_8, qcir::DataType::UINT8}, - {QNN_DATATYPE_UINT_16, qcir::DataType::UINT16}, - {QNN_DATATYPE_UINT_32, qcir::DataType::UINT32}, - {QNN_DATATYPE_UINT_64, qcir::DataType::UINT64}, - {QNN_DATATYPE_FLOAT_16, qcir::DataType::FLOAT16}, - {QNN_DATATYPE_FLOAT_32, qcir::DataType::FLOAT32}, - // {QNN_DATATYPE_FLOAT_64, qcir::DataType::FLOAT64}, - {QNN_DATATYPE_SFIXED_POINT_4, qcir::DataType::SFIXED4}, - {QNN_DATATYPE_SFIXED_POINT_8, qcir::DataType::SFIXED8}, - {QNN_DATATYPE_SFIXED_POINT_16, qcir::DataType::SFIXED16}, - {QNN_DATATYPE_SFIXED_POINT_32, qcir::DataType::SFIXED32}, - {QNN_DATATYPE_UFIXED_POINT_4, qcir::DataType::UFIXED4}, - {QNN_DATATYPE_UFIXED_POINT_8, qcir::DataType::UFIXED8}, - {QNN_DATATYPE_UFIXED_POINT_16, qcir::DataType::UFIXED16}, - {QNN_DATATYPE_UFIXED_POINT_32, qcir::DataType::UFIXED32}, - {QNN_DATATYPE_BOOL_8, qcir::DataType::BOOL}, - // {QNN_DATATYPE_STRING, qcir::DataType::STRING}, - {QNN_DATATYPE_UNDEFINED, qcir::DataType::UNDEFINED}, - }; - return type_map.at(type); -} - -// TODO: enable commented type by QNN version control -Qnn_DataType_t ToDataType(qcir::DataType type) { - static const std::unordered_map type_map{ - {qcir::DataType::INT8, QNN_DATATYPE_INT_8}, - {qcir::DataType::INT16, QNN_DATATYPE_INT_16}, - {qcir::DataType::INT32, QNN_DATATYPE_INT_32}, - {qcir::DataType::INT64, QNN_DATATYPE_INT_64}, - {qcir::DataType::UINT8, QNN_DATATYPE_UINT_8}, - {qcir::DataType::UINT16, QNN_DATATYPE_UINT_16}, - {qcir::DataType::UINT32, QNN_DATATYPE_UINT_32}, - {qcir::DataType::UINT64, QNN_DATATYPE_UINT_64}, - {qcir::DataType::FLOAT16, QNN_DATATYPE_FLOAT_16}, - {qcir::DataType::FLOAT32, QNN_DATATYPE_FLOAT_32}, - // {qcir::DataType::FLOAT64, QNN_DATATYPE_FLOAT_64}, - {qcir::DataType::SFIXED4, QNN_DATATYPE_SFIXED_POINT_4}, - {qcir::DataType::SFIXED8, QNN_DATATYPE_SFIXED_POINT_8}, - {qcir::DataType::SFIXED16, QNN_DATATYPE_SFIXED_POINT_16}, - {qcir::DataType::SFIXED32, QNN_DATATYPE_SFIXED_POINT_32}, - {qcir::DataType::UFIXED4, QNN_DATATYPE_UFIXED_POINT_4}, - {qcir::DataType::UFIXED8, QNN_DATATYPE_UFIXED_POINT_8}, - {qcir::DataType::UFIXED16, QNN_DATATYPE_UFIXED_POINT_16}, - {qcir::DataType::UFIXED32, QNN_DATATYPE_UFIXED_POINT_32}, - {qcir::DataType::BOOL, QNN_DATATYPE_BOOL_8}, - // {qcir::DataType::STRING, QNN_DATATYPE_STRING}, - {qcir::DataType::UNDEFINED, QNN_DATATYPE_UNDEFINED}, - }; - return type_map.at(type); -} - -flatbuffers::Offset ToQuantizeParam( - const Qnn_Tensor_t& tensor, - flatbuffers::FlatBufferBuilder* builder) { - static const std::unordered_map def_map{ - {QNN_DEFINITION_IMPL_GENERATED, qcir::QuantizeDef::IMPL_GENERATED}, - {QNN_DEFINITION_DEFINED, qcir::QuantizeDef::DEFINED}, - {QNN_DEFINITION_UNDEFINED, qcir::QuantizeDef::UNDEFINED}, - }; - static const std:: - unordered_map - type_map{ - {QNN_QUANTIZATION_ENCODING_SCALE_OFFSET, - qcir::QuantizeType::SCALE_OFFSET}, - {QNN_QUANTIZATION_ENCODING_AXIS_SCALE_OFFSET, - qcir::QuantizeType::AXIS_SCALE_OFFSET}, - {QNN_QUANTIZATION_ENCODING_BW_SCALE_OFFSET, - qcir::QuantizeType::BW_SCALE_OFFSET}, - {QNN_QUANTIZATION_ENCODING_BW_AXIS_SCALE_OFFSET, - qcir::QuantizeType::BW_AXIS_SCALE_OFFSET}, - {QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION, - qcir::QuantizeType::BLOCKWISE_EXPANSION}, - {QNN_QUANTIZATION_ENCODING_UNDEFINED, - qcir::QuantizeType::UNDEFINED}, - }; - - int32_t axis = 0; - uint32_t bitwidth = 0, num_blocks_per_axis = 0; - auto param = QNN_TENSOR_VER_PTR(tensor)->quantizeParams; - auto quant_type = type_map.at(param.quantizationEncoding); - std::vector data; - std::vector block_scale; - std::vector scales; - std::vector offsets; - qcir::BlockScaleStorageType block_scale_storage_type = - qcir::BlockScaleStorageType::BITWIDTH_SCALE_STORAGE_8; - switch (quant_type) { - case qcir::QuantizeType::SCALE_OFFSET: { - data.emplace_back(qcir::ScaleOffset( - param.scaleOffsetEncoding.scale, param.scaleOffsetEncoding.offset)); - } break; - case qcir::QuantizeType::AXIS_SCALE_OFFSET: { - size_t len = param.axisScaleOffsetEncoding.numScaleOffsets; - axis = param.axisScaleOffsetEncoding.axis; - data.reserve(len); - for (uint i = 0; i < len; ++i) { - data.emplace_back(qcir::ScaleOffset( - param.axisScaleOffsetEncoding.scaleOffset[i].scale, - param.axisScaleOffsetEncoding.scaleOffset[i].offset)); - } - } break; - case qcir::QuantizeType::BW_SCALE_OFFSET: { - bitwidth = param.bwScaleOffsetEncoding.bitwidth; - scales.push_back(param.bwScaleOffsetEncoding.scale); - offsets.push_back(param.bwScaleOffsetEncoding.offset); - } break; - case qcir::QuantizeType::BW_AXIS_SCALE_OFFSET: { - bitwidth = param.bwAxisScaleOffsetEncoding.bitwidth; - axis = param.bwAxisScaleOffsetEncoding.axis; - size_t len = param.bwAxisScaleOffsetEncoding.numElements; - scales.reserve(len); - offsets.reserve(len); - for (size_t i = 0; i < len; ++i) { - scales.push_back(param.bwAxisScaleOffsetEncoding.scales[i]); - offsets.push_back(param.bwAxisScaleOffsetEncoding.offsets[i]); - } - } break; - case qcir::QuantizeType::BLOCKWISE_EXPANSION: { - bitwidth = param.blockwiseExpansion->blockScaleBitwidth; - axis = param.blockwiseExpansion->axis; - uint num_channels = QNN_TENSOR_VER_PTR(tensor)->dimensions[axis]; - for (uint i = 0; i < num_channels; ++i) { - data.emplace_back(qcir::ScaleOffset( - param.blockwiseExpansion->scaleOffsets[i].scale, - param.blockwiseExpansion->scaleOffsets[i].offset)); - } - num_blocks_per_axis = param.blockwiseExpansion->numBlocksPerAxis; - uint multiplier = 1; - if (param.blockwiseExpansion->blockScaleStorageType == - QNN_BLOCKWISE_EXPANSION_BITWIDTH_SCALE_STORAGE_16) { - multiplier = 2; - block_scale_storage_type = - qcir::BlockScaleStorageType::BITWIDTH_SCALE_STORAGE_16; - } - uint total_bytes = num_channels * num_blocks_per_axis * multiplier; - block_scale = std::vector( - param.blockwiseExpansion->blocksScale8, - param.blockwiseExpansion->blocksScale8 + total_bytes); - } break; - default: - // encodings are not required if lowering with floating point precision - break; - } - return CreateQuantizeParamDirect( - *builder, - def_map.at(param.encodingDefinition), - quant_type, - bitwidth, - axis, - &scales, - &offsets, - &data, - num_blocks_per_axis, - block_scale_storage_type, - &block_scale); -} - -Qnn_QuantizeParams_t ToQuantizeParam(const tensor_type& tensor) { - static const std::unordered_map def_map{ - {qcir::QuantizeDef::IMPL_GENERATED, QNN_DEFINITION_IMPL_GENERATED}, - {qcir::QuantizeDef::DEFINED, QNN_DEFINITION_DEFINED}, - {qcir::QuantizeDef::UNDEFINED, QNN_DEFINITION_UNDEFINED}, - }; - static const std:: - unordered_map - type_map{ - {qcir::QuantizeType::SCALE_OFFSET, - QNN_QUANTIZATION_ENCODING_SCALE_OFFSET}, - {qcir::QuantizeType::AXIS_SCALE_OFFSET, - QNN_QUANTIZATION_ENCODING_AXIS_SCALE_OFFSET}, - {qcir::QuantizeType::BW_SCALE_OFFSET, - QNN_QUANTIZATION_ENCODING_BW_SCALE_OFFSET}, - {qcir::QuantizeType::BW_AXIS_SCALE_OFFSET, - QNN_QUANTIZATION_ENCODING_BW_AXIS_SCALE_OFFSET}, - {qcir::QuantizeType::BLOCKWISE_EXPANSION, - QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION}, - {qcir::QuantizeType::UNDEFINED, - QNN_QUANTIZATION_ENCODING_UNDEFINED}, - }; - // Qnn_BlockwiseExpansion_t is a pointer type in Qnn_QuantizeParams_t - // need a bookkeeper for guarding life cycle - static std::vector> block_param; - - Qnn_QuantizeParams_t p = QNN_QUANTIZE_PARAMS_INIT; - auto param = tensor->qparam(); - p.encodingDefinition = def_map.at(param->def()); - p.quantizationEncoding = type_map.at(param->type()); - switch (p.quantizationEncoding) { - case QNN_QUANTIZATION_ENCODING_SCALE_OFFSET: { - p.scaleOffsetEncoding.scale = param->data()->Get(0)->scale(); - p.scaleOffsetEncoding.offset = param->data()->Get(0)->offset(); - } break; - case QNN_QUANTIZATION_ENCODING_AXIS_SCALE_OFFSET: { - p.axisScaleOffsetEncoding.axis = param->axis(); - p.axisScaleOffsetEncoding.numScaleOffsets = param->data()->size(); - p.axisScaleOffsetEncoding.scaleOffset = - reinterpret_cast( - const_cast(param->data()->Data())); - } break; - case QNN_QUANTIZATION_ENCODING_BW_SCALE_OFFSET: { - p.bwAxisScaleOffsetEncoding.bitwidth = param->bitwidth(); - p.bwScaleOffsetEncoding.scale = param->scales()->Get(0); - p.bwScaleOffsetEncoding.offset = param->offsets()->Get(0); - } break; - case QNN_QUANTIZATION_ENCODING_BW_AXIS_SCALE_OFFSET: { - p.bwAxisScaleOffsetEncoding.bitwidth = param->bitwidth(); - p.bwAxisScaleOffsetEncoding.axis = param->axis(); - p.bwAxisScaleOffsetEncoding.numElements = param->scales()->size(); - p.bwAxisScaleOffsetEncoding.scales = - const_cast(param->scales()->data()); - p.bwAxisScaleOffsetEncoding.offsets = - const_cast(param->offsets()->data()); - } break; - case QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION: { - block_param.emplace_back(std::make_unique()); - p.blockwiseExpansion = block_param.back().get(); - p.blockwiseExpansion->axis = param->axis(); - p.blockwiseExpansion->scaleOffsets = reinterpret_cast( - const_cast(param->data()->Data())); - p.blockwiseExpansion->numBlocksPerAxis = param->num_blocks_per_axis(); - switch (param->block_scale_storage_type()) { - case qcir::BlockScaleStorageType::BITWIDTH_SCALE_STORAGE_8: - p.blockwiseExpansion->blockScaleStorageType = - QNN_BLOCKWISE_EXPANSION_BITWIDTH_SCALE_STORAGE_8; - break; - case qcir::BlockScaleStorageType::BITWIDTH_SCALE_STORAGE_16: - p.blockwiseExpansion->blockScaleStorageType = - QNN_BLOCKWISE_EXPANSION_BITWIDTH_SCALE_STORAGE_16; - break; - default: - p.blockwiseExpansion->blockScaleStorageType = - QNN_BLOCKWISE_EXPANSION_BITWIDTH_SCALE_STORAGE_UNDEFINED; - break; - } - p.blockwiseExpansion->blocksScale8 = - const_cast(param->block_scale()->Data()); - } break; - default: - // encodings are not required if lowering with floating point precision - break; - } - return p; -} - -flatbuffers::Offset ToTensor( - const Qnn_Tensor_t& tensor, - const uint64_t data_offset, - flatbuffers::FlatBufferBuilder* builder) { - std::vector shape( - QNN_TENSOR_VER_PTR(tensor)->dimensions, - QNN_TENSOR_VER_PTR(tensor)->dimensions + - QNN_TENSOR_VER_PTR(tensor)->rank); - std::vector dynamic_dims( - QNN_TENSOR_VER_PTR(tensor)->isDynamicDimensions, - QNN_TENSOR_VER_PTR(tensor)->isDynamicDimensions + - QNN_TENSOR_VER_PTR(tensor)->rank); - - return qcir::CreateTensorDirect( - *builder, - QNN_TENSOR_VER_PTR(tensor)->name, - &shape, - &dynamic_dims, - ToTensorType(QNN_TENSOR_VER_PTR(tensor)->type), - ToDataType(QNN_TENSOR_VER_PTR(tensor)->dataType), - ToQuantizeParam(tensor, builder), - QNN_TENSOR_VER_PTR(tensor)->clientBuf.dataSize, - data_offset); -} - -Qnn_Tensor_t ToTensor(const tensor_type& tensor, const uint8_t* data_ptr) { - auto is_io_tensor = [](Qnn_TensorType_t type) { - return type < QNN_TENSOR_TYPE_STATIC; - }; - - Qnn_Tensor_t t({.version = QNN_TENSOR_VERSION_2, .v2 = QNN_TENSOR_V2_INIT}); - QNN_TENSOR_VER_PTR(t)->name = tensor->name()->c_str(); - QNN_TENSOR_VER_PTR(t)->type = ToTensorType(tensor->type()); - QNN_TENSOR_VER_PTR(t)->dataType = ToDataType(tensor->dtype()); - QNN_TENSOR_VER_PTR(t)->quantizeParams = ToQuantizeParam(tensor); - QNN_TENSOR_VER_PTR(t)->rank = tensor->shape()->size(); - QNN_TENSOR_VER_PTR(t)->dimensions = - const_cast(tensor->shape()->data()); - QNN_TENSOR_VER_PTR(t)->isDynamicDimensions = - const_cast(tensor->dynamic_dims()->data()); - QNN_TENSOR_VER_PTR(t)->clientBuf.dataSize = tensor->size(); - QNN_TENSOR_VER_PTR(t)->clientBuf.data = - is_io_tensor(QNN_TENSOR_VER_PTR(t)->type) - ? nullptr - : static_cast(const_cast(data_ptr)); - return t; -} - -} // namespace qnn -} // namespace backends -} // namespace executorch diff --git a/backends/qualcomm/aot/ir/qcir_utils.h b/backends/qualcomm/aot/ir/qcir_utils.h deleted file mode 100755 index 085f09bf145..00000000000 --- a/backends/qualcomm/aot/ir/qcir_utils.h +++ /dev/null @@ -1,41 +0,0 @@ -/* - * Copyright (c) Qualcomm Innovation Center, Inc. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -#include -#include "QnnTypes.h" - -namespace executorch { -namespace backends { -namespace qnn { - -typedef flatbuffers::Vector<::flatbuffers::Offset>::return_type - tensor_type; -typedef flatbuffers::Vector< - ::flatbuffers::Offset>::return_type qparam_type; - -qcir::TensorType ToTensorType(Qnn_TensorType_t type); -Qnn_TensorType_t ToTensorType(qcir::TensorType type); -qcir::DataType ToDataType(Qnn_DataType_t type); -Qnn_DataType_t ToDataType(qcir::DataType type); - -flatbuffers::Offset ToQuantizeParam( - const Qnn_Tensor_t& tensor, - flatbuffers::FlatBufferBuilder* builder); -Qnn_QuantizeParams_t ToQuantizeParam(const tensor_type& tensor); - -flatbuffers::Offset ToTensor( - const Qnn_Tensor_t& tensor, - const uint64_t data_offset, - flatbuffers::FlatBufferBuilder* builder); -Qnn_Tensor_t ToTensor(const tensor_type& tensor, const uint8_t* data_ptr); - -} // namespace qnn -} // namespace backends -} // namespace executorch diff --git a/backends/qualcomm/aot/ir/targets.bzl b/backends/qualcomm/aot/ir/targets.bzl deleted file mode 100644 index 2405af35d6c..00000000000 --- a/backends/qualcomm/aot/ir/targets.bzl +++ /dev/null @@ -1,68 +0,0 @@ -load( - "@fbsource//tools/build_defs:default_platform_defs.bzl", - "ANDROID", -) -load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -load("@fbsource//xplat/executorch/backends/qualcomm:targets.bzl", "generate_schema_header") -load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") - -QCIR_NAME = "qcir" -INPUT_QCIR = QCIR_NAME + ".fbs" -OUTPUT_QCIR_HEADER = QCIR_NAME + "_generated.h" -QCIR_GEN_RULE_NAME = "qcir_generated" - -def define_common_targets(): - """Defines targets that should be shared between fbcode and xplat. - The directory containing this targets.bzl file should also contain both - TARGETS and BUCK files that call this function. - """ - - generate_schema_header( - QCIR_GEN_RULE_NAME, - [INPUT_QCIR], - [OUTPUT_QCIR_HEADER], - OUTPUT_QCIR_HEADER, - ) - - # Header-only library target with the generate executorch program schema header. - runtime.cxx_library( - name = "qcir_schema", - srcs = [], - exported_headers = { - OUTPUT_QCIR_HEADER: ":{}[{}]".format(QCIR_GEN_RULE_NAME, OUTPUT_QCIR_HEADER), - }, - visibility = [ - # Lock this down as tightly as possible to ensure that flatbuffers - # are an implementation detail. Ideally this list would only include - # //executorch/runtime/executor/... - "//executorch/backends/qualcomm/...", - "//executorch/backends/qualcomm/aot/ir/...", - ], - exported_external_deps = ["flatbuffers-api"], - define_static_target = True, - platforms = [ANDROID], - ) - - - runtime.cxx_library( - name = "qcir_utils", - srcs = [ - "qcir_utils.cpp", - ], - exported_headers = [ - "qcir_utils.h", - ], - define_static_target = True, - platforms = [ANDROID], - visibility = ["@EXECUTORCH_CLIENTS"], - deps = [ - "fbsource//third-party/qualcomm/qnn/qnn-{0}:api".format(get_qnn_library_version()), - "fbsource//third-party/qualcomm/qnn/qnn-{0}:app_sources".format(get_qnn_library_version()), - "//executorch/runtime/backend:interface", - "//executorch/runtime/core:core", - "//executorch/backends/qualcomm/aot/wrappers:wrappers", - ], - exported_deps = [ - ":qcir_schema", - ], - ) diff --git a/backends/qualcomm/aot/python/PyQnnManagerAdaptor.h b/backends/qualcomm/aot/python/PyQnnManagerAdaptor.h index 409ec1a4294..c8044e5db0e 100644 --- a/backends/qualcomm/aot/python/PyQnnManagerAdaptor.h +++ b/backends/qualcomm/aot/python/PyQnnManagerAdaptor.h @@ -132,16 +132,6 @@ class PyQnnManager { return qnn_manager_->GetSpillFillBufferSize(); } - QnnExecuTorchContextBinary MakeQcirCustomBinaryInfo( - const QnnExecuTorchContextBinary& ctx_bin, - const std::vector& tensor_data) { - custom_qcir_protocol_buffer_ = - QnnQcirCustomProtocol(ctx_bin.nbytes, tensor_data.size()); - custom_qcir_protocol_buffer_.BuildQcirCustomBuffer(ctx_bin, tensor_data); - auto [ptr, size] = custom_qcir_protocol_buffer_.GetCustomProtocolBuffer(); - return {ptr, size}; - } - py::array_t MakeBinaryInfo(const py::bytes& ctx_bin) { py::buffer_info info(py::buffer(ctx_bin).request()); QnnExecuTorchContextBinary binary( @@ -171,22 +161,10 @@ class PyQnnManager { buf_size = ctx_size; buf_ptr = ctx_bin; } else { - // check if it's a qcir flatbuffers, return fbs if matched - auto - [status, - qcir_fbs_size, - qcir_tensor_size, - qcir_fbs_ptr, - qcir_tensor_ptr] = - QnnQcirCustomProtocol().DeserializeQcirCustomBuffer(info.ptr); - if (status == Error::Ok) { - buf_size = qcir_fbs_size; - buf_ptr = qcir_fbs_ptr; - } else { - // the format should be DLC, return nothing here - return py::array_t(0); - } + // the format should be DLC, return nothing here + return py::array_t(0); } + auto result = py::array_t(buf_size); auto result_buffer = result.request(); std::memcpy(result_buffer.ptr, buf_ptr, buf_size); @@ -199,7 +177,6 @@ class PyQnnManager { const py::bytes qnn_executorch_option_ptr_; QnnExecuTorchContextBinary qnn_executorch_context_binary_; std::shared_ptr qnn_manager_; - QnnQcirCustomProtocol custom_qcir_protocol_buffer_; QnnContextCustomProtocol custom_context_custom_buffer_; flatbuffers::FlatBufferBuilder builder_; }; diff --git a/backends/qualcomm/aot/python/targets.bzl b/backends/qualcomm/aot/python/targets.bzl index da27997808b..74fbd1da511 100644 --- a/backends/qualcomm/aot/python/targets.bzl +++ b/backends/qualcomm/aot/python/targets.bzl @@ -31,7 +31,6 @@ def define_common_targets(): "//executorch/backends/qualcomm/aot/wrappers:wrappers", "//executorch/backends/qualcomm/runtime:logging", "//executorch/backends/qualcomm:schema", - "//executorch/backends/qualcomm/aot/ir:qcir_utils", "//executorch/backends/qualcomm/runtime:runtime", "fbsource//third-party/pybind11:pybind11", "fbsource//third-party/qualcomm/qnn/qnn-{0}:api".format(get_qnn_library_version()), @@ -65,7 +64,6 @@ def define_common_targets(): "//executorch/backends/qualcomm/aot/wrappers:wrappers", "//executorch/backends/qualcomm/runtime:logging", "//executorch/backends/qualcomm:schema", - "//executorch/backends/qualcomm/aot/ir:qcir_utils", "//executorch/backends/qualcomm/runtime:runtime", "fbsource//third-party/pybind11:pybind11", "fbsource//third-party/qualcomm/qnn/qnn-{0}:api".format(get_qnn_library_version()), @@ -94,7 +92,6 @@ def define_common_targets(): "//executorch/backends/qualcomm/aot/wrappers:wrappers", "//executorch/backends/qualcomm/runtime:logging", "//executorch/backends/qualcomm:schema", - "//executorch/backends/qualcomm/aot/ir:qcir_utils", "//executorch/backends/qualcomm/runtime:runtime", "fbsource//third-party/pybind11:pybind11", "fbsource//third-party/qualcomm/qnn/qnn-{0}:api".format(get_qnn_library_version()), diff --git a/backends/qualcomm/bc/test_qnn_static_llama_bc.sh b/backends/qualcomm/bc/test_qnn_static_llama_bc.sh index c76485a664c..478e6118641 100644 --- a/backends/qualcomm/bc/test_qnn_static_llama_bc.sh +++ b/backends/qualcomm/bc/test_qnn_static_llama_bc.sh @@ -13,15 +13,18 @@ fi which "${PYTHON_EXECUTABLE}" -llama_artifacts="." +llama_artifacts="260k_stories" PTE_ARTIFACT="examples/qualcomm/oss_scripts/llama/artifacts" +mkdir ${llama_artifacts} # Download stories260K.pt and tokenizer from Github -curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/stories260K.pt" --output stories260K.pt -curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/tok512.model" --output tokenizer.model +curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/stories260K.pt" --output ${llama_artifacts}/stories260K.pt +curl -Ls "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K/tok512.model" --output ${llama_artifacts}/tokenizer.model + +$PYTHON_EXECUTABLE -m pytorch_tokenizers.tools.llama2c.convert -t ${llama_artifacts}/tokenizer.model -o ${llama_artifacts}/tokenizer.bin # Create params.json file -touch params.json -echo '{"dim": 64, "n_layers": 5, "n_heads": 8, "n_kv_heads": 4, "vocab_size": 512, "multiple_of": 4, "max_seq_len": 512}' > params.json +touch ${llama_artifacts}/params.json +echo '{"dim": 64, "n_layers": 5, "n_heads": 8, "n_kv_heads": 4, "vocab_size": 512, "multiple_of": 4, "max_seq_len": 512}' > ${llama_artifacts}/params.json # Checks e2e accuracy expected=$($PYTHON_EXECUTABLE backends/qualcomm/tests/test_qnn_delegate.py -k TestExampleLLMScript.test_llama_stories_260k --model SM8650 --build_folder build-x86/ --executorch_root . --artifact_dir . --llama_artifacts $llama_artifacts --enable_x86_64 | grep "Model CI result:") diff --git a/backends/qualcomm/builders/README.md b/backends/qualcomm/builders/README.md index 9c62e1080fe..6ba4eafb01f 100644 --- a/backends/qualcomm/builders/README.md +++ b/backends/qualcomm/builders/README.md @@ -360,7 +360,12 @@ The operator now should be functional for Qualcomm backends. For operator to wor ## Operator Support Status Please help update following table if you are contributing new operators: -| Operators | HTP - 82/116 Enabled | ++ ✓ = Supported ++ ✗ = Not Supported ++ 🚫 = Deprecated, supported with other QNN Ops + + +| Operators | HTP - 90/116 Enabled | |-----------|---------| | Argmax | ✓ | | Argmin | ✓ | @@ -381,16 +386,16 @@ Please help update following table if you are contributing new operators: | ElementWiseAbs | ✓ | | ElementWiseAdd | ✓ | | ElementWiseAnd | ✓ | -| ElementWiseAsin | ✗ | +| ElementWiseAsin | ✓ | | ElementWiseAtan | ✓ | -| ElementWiseBinary | ✗ | +| ElementWiseBinary | ✓ | | ElementWiseCeil | ✓ | | ElementWiseCos | ✓ | | ElementWiseDivide | ✓ | | ElementWiseEqual | ✓ | | ElementWiseExp | ✓ | | ElementWiseFloor | ✓ | -| ElementWiseFloorDiv | ✗ | +| ElementWiseFloorDiv | ✓ | | ElementWiseGreater | ✓ | | ElementWiseGreaterEqual | ✓ | | ElementWiseLess | ✓ | @@ -408,13 +413,13 @@ Please help update following table if you are contributing new operators: | ElementWiseRound | ✓ | | ElementWiseRsqrt | ✓ | | ElementWiseSelect | ✓ | -| ElementWiseSign | ✗ | +| ElementWiseSign | ✓ | | ElementWiseSin | ✓ | | ElementWiseSquaredDifference | ✗ | | ElementWiseSquareRoot | ✓ | | ElementWiseSubtract | ✓ | | ElementWiseUnary | ✗ | -| ElementWiseXor | ✗ | +| ElementWiseXor | ✓ | | Elu | ✓ | | ExpandDims | ✓ | | ExtractGlimpse | ✗ | @@ -452,11 +457,11 @@ Please help update following table if you are contributing new operators: | ReduceMin | ✓ | | ReduceSum | ✓ | | Relu | ✓ | -| Relu1 | ✗ | -| Relu6 | ✗ | +| Relu1 | 🚫 | +| Relu6 | 🚫 | | ReluMinMax | ✓ | | Reshape | ✓ | -| Resize | ✗ | +| Resize | ✓ | | ResizeBilinear | ✓ | | ResizeNearestNeighbor | ✓ | | RoiAlign | ✗ | diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index 62e8e476257..9800fb7bdab 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -15,14 +15,17 @@ op_arange, op_argmax, op_argmin, + op_asin, op_atan, op_avg_pool2d, op_batch_norm, + op_binary, op_bmm, op_cat, op_ceil, op_clamp, op_conv2d, + op_copy, op_cos, op_cum_sum, op_depth_to_space, @@ -33,6 +36,7 @@ op_eq, op_exp, op_expand, + op_flip, op_floor, op_full, op_full_like, @@ -46,6 +50,7 @@ op_hardtanh, op_index, op_index_put, + op_index_select, op_instance_norm, op_layer_norm, op_le, @@ -79,9 +84,11 @@ op_scalar_tensor, op_select_copy, op_sigmoid, + op_sign, op_sin, op_skip_ops, op_slice_copy, + op_slice_scatter, op_softmax, op_space_to_depth, op_split_with_sizes, @@ -99,6 +106,7 @@ op_upsample_bilinear2d, op_upsample_nearest2d, op_where, + op_xor, ) __all__ = [ @@ -112,14 +120,17 @@ op_arange, op_argmax, op_argmin, + op_asin, op_atan, op_avg_pool2d, op_batch_norm, + op_binary, op_bmm, op_cat, op_ceil, op_clamp, op_conv2d, + op_copy, op_cos, op_cum_sum, op_depth_to_space, @@ -130,6 +141,7 @@ op_eq, op_exp, op_expand, + op_flip, op_floor, op_full, op_full_like, @@ -143,6 +155,7 @@ op_hardsigmoid, op_index, op_index_put, + op_index_select, op_instance_norm, op_layer_norm, op_le, @@ -176,9 +189,11 @@ op_scalar_tensor, op_select_copy, op_sigmoid, + op_sign, op_sin, op_skip_ops, op_slice_copy, + op_slice_scatter, op_softmax, op_space_to_depth, op_split_with_sizes, @@ -196,4 +211,5 @@ op_upsample_bilinear2d, op_upsample_nearest2d, op_where, + op_xor, ] diff --git a/backends/qualcomm/builders/node_visitor.py b/backends/qualcomm/builders/node_visitor.py index ae3c99ff523..bc2b62c8c0b 100644 --- a/backends/qualcomm/builders/node_visitor.py +++ b/backends/qualcomm/builders/node_visitor.py @@ -59,6 +59,8 @@ QNN_TENSOR_TYPE_MAP = { torch.bool: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8, torch.float32: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, + # Note that there is no float64 tensor data type in Qnn. + torch.float64: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32, torch.int8: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_8, torch.int16: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_16, torch.int32: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_32, @@ -162,7 +164,7 @@ def make_qnn_per_block_config(self, node: torch.fx.Node, quant_attrs: Dict): for ch in range(num_channels): max_scale = scales[ch].reshape(1, -1).amax(dim=-1) / num_steps q_scales = torch.clamp( - input=scales[ch] / max_scale, + input=torch.round(input=scales[ch] / max_scale), min=1, max=2**bitwidth_of_scale, ).to(quant_scales_dtype) @@ -176,6 +178,10 @@ def make_qnn_per_block_config(self, node: torch.fx.Node, quant_attrs: Dict): # OIHW (pytorch) -> HWIO (QNN) quant_config[QCOM_AXIS] = 3 quant_config[QCOM_AXIS_ORDER] = (2, 3, 1, 0) + elif "linear" in user_0.target.__name__: + # OI (pytorch) -> OI (QNN) + quant_config[QCOM_AXIS] = 0 + quant_config[QCOM_AXIS_ORDER] = (0, 1) else: raise AttributeError("undetermined axis for block quantization") diff --git a/backends/qualcomm/builders/op_amax.py b/backends/qualcomm/builders/op_amax.py index 051355a8b6b..d0335f95463 100644 --- a/backends/qualcomm/builders/op_amax.py +++ b/backends/qualcomm/builders/op_amax.py @@ -40,14 +40,19 @@ def define_node( ) # mean dims and keep dims - mean_dims = cast(List[int], node.args[1]) - mean_dims = [ - mean_dim % len(input_node.meta["val"].shape) for mean_dim in mean_dims - ] - if QCOM_AXIS_ORDER in node.meta: + if len(node.args) > 1: + mean_dims = cast(List[int], node.args[1]) mean_dims = [ - node.meta[QCOM_AXIS_ORDER].index(mean_dim) for mean_dim in mean_dims + mean_dim % len(input_node.meta["val"].shape) for mean_dim in mean_dims ] + if QCOM_AXIS_ORDER in node.meta: + mean_dims = [ + node.meta[QCOM_AXIS_ORDER].index(mean_dim) for mean_dim in mean_dims + ] + else: + # reduce all dimensions + mean_dims = list(range(input_node.meta["val"].dim())) + mean_dims_shape = [len(mean_dims)] output_tensor = self.get_tensor(node, node) diff --git a/backends/qualcomm/builders/op_amin.py b/backends/qualcomm/builders/op_amin.py index 9f8f17b4e37..142340dbae0 100644 --- a/backends/qualcomm/builders/op_amin.py +++ b/backends/qualcomm/builders/op_amin.py @@ -40,14 +40,19 @@ def define_node( ) # mean dims and keep dims - mean_dims = cast(List[int], node.args[1]) - mean_dims = [ - mean_dim % len(input_node.meta["val"].shape) for mean_dim in mean_dims - ] - if QCOM_AXIS_ORDER in node.meta: + if len(node.args) > 1: + mean_dims = cast(List[int], node.args[1]) mean_dims = [ - node.meta[QCOM_AXIS_ORDER].index(mean_dim) for mean_dim in mean_dims + mean_dim % len(input_node.meta["val"].shape) for mean_dim in mean_dims ] + if QCOM_AXIS_ORDER in node.meta: + mean_dims = [ + node.meta[QCOM_AXIS_ORDER].index(mean_dim) for mean_dim in mean_dims + ] + else: + # reduce all dimensions + mean_dims = list(range(input_node.meta["val"].dim())) + mean_dims_shape = [len(mean_dims)] output_tensor = self.get_tensor(node, node) diff --git a/backends/qualcomm/builders/op_asin.py b/backends/qualcomm/builders/op_asin.py new file mode 100644 index 00000000000..ff50380e62c --- /dev/null +++ b/backends/qualcomm/builders/op_asin.py @@ -0,0 +1,56 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree.from typing import cast, Dict +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper +import torch + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor + +from .qnn_constants import OpElementWiseAsin, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class asin(NodeVisitor): + target = ["aten.asin.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + asin_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseAsin.op_name, + ) + asin_op.AddInputTensors([input_tensor_wrapper]) + asin_op.AddOutputTensors([output_tensor_wrapper]) + + return asin_op diff --git a/backends/qualcomm/builders/op_binary.py b/backends/qualcomm/builders/op_binary.py new file mode 100644 index 00000000000..4f4d8b9b560 --- /dev/null +++ b/backends/qualcomm/builders/op_binary.py @@ -0,0 +1,84 @@ +# 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 warnings +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_DATA +from executorch.exir.dialects._ops import ops as exir_ops + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpElementWiseBinary, QNN_OP_PACKAGE_NAME_QTI_AISW + + +# Refer to QnnOpDef.h for the value. +QNN_BINARY_OPERATOR = { + exir_ops.edge.aten.floor_divide.default: 4, +} + + +@register_node_visitor +class Binary(NodeVisitor): + target = ["aten.floor_divide.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + out_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + binary_output_tensors = [output_tensor_wrapper] + + binary_input_tensors = [] + for index in range(2): + input_node = self.get_node(node.args[index]) + input_tensor = self.get_tensor(input_node, node) + tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + tensor_type, + nodes_to_wrappers, + ) + binary_input_tensors.append(input_tensor_wrapper) + + binary_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseBinary.op_name, + ) + binary_op.AddInputTensors(binary_input_tensors) + binary_op.AddOutputTensors(binary_output_tensors) + + if node.target not in QNN_BINARY_OPERATOR: + warnings.warn( + "[QNN Delegate Op Builder]: This binary operator is not yet supported.", + stacklevel=1, + ) + return None + + binary_op.AddScalarParam( + OpElementWiseBinary.param_operation, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: np.uint32(QNN_BINARY_OPERATOR[node.target])}, + ) + + return binary_op diff --git a/backends/qualcomm/builders/op_conv2d.py b/backends/qualcomm/builders/op_conv2d.py index 0456cd53524..1cfc1e45c9b 100644 --- a/backends/qualcomm/builders/op_conv2d.py +++ b/backends/qualcomm/builders/op_conv2d.py @@ -109,7 +109,7 @@ def define_node( input_tensor = self.get_tensor(input_node, node) assert ( input_tensor.dim() == 4 - ), "All Conv should be converted to Conv2D in ConvertConv1dToConv2d" + ), "All Conv1D should be converted to Conv2D in CanonicalizeConv," input_tensor_wrapper = self.define_tensor( input_node, node, diff --git a/backends/qualcomm/builders/op_copy.py b/backends/qualcomm/builders/op_copy.py new file mode 100644 index 00000000000..164c910835e --- /dev/null +++ b/backends/qualcomm/builders/op_copy.py @@ -0,0 +1,64 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_QUANT_ATTRS + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpReshape, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Copy(NodeVisitor): + target = ["aten.copy.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[1]) + input_tensor = self.get_tensor(input_node, node) + copy_inp_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + copy_input_tensors = [copy_inp_tensor_wrapper] + + if quant_attrs := input_node.meta.get(QCOM_QUANT_ATTRS): + quant_attrs = quant_attrs.copy() + # Because there is no output after convert_pt2e, the QCOM_QUANT_ATTRS of node is none + node.meta[QCOM_QUANT_ATTRS] = quant_attrs + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + copy_output_tensors = [output_tensor_wrapper] + + copy_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpReshape.op_name, + ) + copy_op.AddInputTensors(copy_input_tensors) + copy_op.AddOutputTensors(copy_output_tensors) + + return copy_op diff --git a/backends/qualcomm/builders/op_flip.py b/backends/qualcomm/builders/op_flip.py new file mode 100644 index 00000000000..16d68942d31 --- /dev/null +++ b/backends/qualcomm/builders/op_flip.py @@ -0,0 +1,81 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch + +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpStridedSlice, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Flip(NodeVisitor): + target = ["aten.flip.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + tensor_type, + nodes_to_wrappers, + ) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + ranges = [] + + dims = node.args[1] + if QCOM_AXIS_ORDER in node.meta: + dims = [node.meta[QCOM_AXIS_ORDER].index(dim) for dim in dims] + + for dim, size in enumerate(output_tensor.shape): + if dim in dims: + ranges.extend([size - 1, -1, -1]) + else: + ranges.extend([0, size, 1]) + + range_shape = [input_tensor.dim(), 3] + stride_slice_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpStridedSlice.op_name, + ) + stride_slice_op.AddInputTensors([input_tensor_wrapper]) + stride_slice_op.AddOutputTensors([output_tensor_wrapper]) + stride_slice_op.AddTensorParam( + OpStridedSlice.param_ranges, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_32, + len(range_shape), + range_shape, + np.array(ranges, dtype=np.int32), + True, + ) + + return stride_slice_op diff --git a/backends/qualcomm/builders/op_index_select.py b/backends/qualcomm/builders/op_index_select.py new file mode 100644 index 00000000000..22733e45397 --- /dev/null +++ b/backends/qualcomm/builders/op_index_select.py @@ -0,0 +1,81 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_DATA + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpGather, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class IndexSelect(NodeVisitor): + target = ["aten.index_select.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + axis = node.args[1] + indices_node = node.args[2] + indices_tensor = self.get_tensor(indices_node, node).to(torch.int32) + assert indices_tensor.size(0) != 0, "Not support empty indices list" + + indices_tensor_wrapper = self.define_tensor( + indices_node, + node, + indices_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + gather_input_tensors = [input_tensor_wrapper, indices_tensor_wrapper] + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + gather_output_tensors = [output_tensor_wrapper] + + gather_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpGather.op_name, + ) + gather_op.AddInputTensors(gather_input_tensors) + gather_op.AddOutputTensors(gather_output_tensors) + + # If support tuple of tensor, need to refine it based on len + gather_op.AddScalarParam( + OpGather.param_axis, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_32, + {QCOM_DATA: np.int32(axis)}, + ) + + return gather_op diff --git a/backends/qualcomm/builders/op_ne.py b/backends/qualcomm/builders/op_ne.py index 1c7f87d4f4f..660c78e3e14 100644 --- a/backends/qualcomm/builders/op_ne.py +++ b/backends/qualcomm/builders/op_ne.py @@ -16,7 +16,7 @@ @register_node_visitor class NotEqual(NodeVisitor): - target = ["aten.ne.Tensor", "aten.ne.Scalar"] + target = ["aten.ne.Tensor"] def __init__(self, *args) -> None: super().__init__(*args) diff --git a/backends/qualcomm/builders/op_sign.py b/backends/qualcomm/builders/op_sign.py new file mode 100644 index 00000000000..faf2f2e0066 --- /dev/null +++ b/backends/qualcomm/builders/op_sign.py @@ -0,0 +1,56 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpElementWiseSign, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class Sign(NodeVisitor): + target = ["aten.sign.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + sign_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseSign.op_name, + ) + sign_op.AddInputTensors([input_tensor_wrapper]) + sign_op.AddOutputTensors([output_tensor_wrapper]) + + return sign_op diff --git a/backends/qualcomm/builders/op_slice_copy.py b/backends/qualcomm/builders/op_slice_copy.py index 62688a10036..5923d438252 100644 --- a/backends/qualcomm/builders/op_slice_copy.py +++ b/backends/qualcomm/builders/op_slice_copy.py @@ -6,9 +6,9 @@ from typing import cast, Dict import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper - import numpy as np import torch +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER from .node_visitor import NodeVisitor from .node_visitor_manager import register_node_visitor @@ -47,8 +47,9 @@ def define_node( PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, nodes_to_wrappers, ) - dim = cast(int, node.args[1]) + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) if dim < 0: dim = dim % len(input_tensor.shape) @@ -62,7 +63,6 @@ def define_node( end = end % input_tensor.shape[dim] else: end = input_tensor.shape[dim] - input_tensor_rank = len(input_tensor.shape) ranges = [] for i in range(input_tensor_rank): diff --git a/backends/qualcomm/builders/op_slice_scatter.py b/backends/qualcomm/builders/op_slice_scatter.py new file mode 100644 index 00000000000..9fa162d6653 --- /dev/null +++ b/backends/qualcomm/builders/op_slice_scatter.py @@ -0,0 +1,122 @@ +from typing import cast, Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper +import torch + +from executorch.exir.dialects._ops import ops as exir_ops + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpScatterNd, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class SliceScatterVisitor(NodeVisitor): + target = ["aten.slice_scatter.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = self.get_node(node.args[0]) + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + value_node = self.get_node(node.args[1]) + value_tensor = self.get_tensor(value_node, node) + value_tensor_wrapper = self.define_tensor( + value_node, + node, + value_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + dim = cast(int, node.args[2]) if len(node.args) > 2 else 0 + if dim < 0: + dim = dim % len(input_tensor.shape) + + start = ( + cast(int, node.args[3]) + if len(node.args) > 3 and node.args[3] is not None + else 0 + ) + if start < 0: + start = start % input_tensor.shape[dim] + + if len(node.args) > 4: + end = min(cast(int, node.args[4]), input_tensor.shape[dim]) + if end < 0: + end = end % input_tensor.shape[dim] + else: + end = input_tensor.shape[dim] + + step = node.args[5] if len(node.args) > 5 else 1 + + target_index_shape = [] + ranges = [] + # Collect the index + for i in range(dim + 1): + if i == dim: + target_range = torch.tensor(range(start, end, step), dtype=torch.int32) + target_index_shape.append(target_range.size(-1)) + ranges.append(target_range) + else: + size = input_tensor.size(i) + target_index_shape.append(size) + ranges.append(torch.arange(size, dtype=torch.int32)) + # last dim means x-tuple index + target_index_shape.append(dim + 1) + target_index_tensor = ( + torch.cartesian_prod(*ranges).reshape(target_index_shape).contiguous() + ) + + target_index_node = torch.fx.Node( + node.graph, + node.name + "_target_index", + "call_function", + exir_ops.edge.aten.tensor.default, + (), # args + {}, # kwargs + ) + target_index_tensor_wrapper = self.define_tensor( + target_index_node, + node, + target_index_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC, + nodes_to_wrappers, + ) + + slice_scatter_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpScatterNd.op_name, + ) + slice_scatter_op.AddInputTensors( + [ + input_tensor_wrapper, + target_index_tensor_wrapper, + value_tensor_wrapper, + ] + ) + slice_scatter_op.AddOutputTensors([output_tensor_wrapper]) + + return slice_scatter_op diff --git a/backends/qualcomm/builders/op_xor.py b/backends/qualcomm/builders/op_xor.py new file mode 100644 index 00000000000..d4462d9c707 --- /dev/null +++ b/backends/qualcomm/builders/op_xor.py @@ -0,0 +1,60 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import torch + +from .node_visitor import NodeVisitor +from .node_visitor_manager import register_node_visitor +from .qnn_constants import OpElementWiseXor, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class OpXor(NodeVisitor): + target = ["aten.bitwise_xor.Tensor"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + out_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + out_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + xor_output_tensors = [output_tensor_wrapper] + + xor_input_tensors = [] + for index in range(2): + input_node = self.get_node(node.args[index]) + input_tensor = self.get_tensor(input_node, node) + tensor_type = PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE + + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + tensor_type, + nodes_to_wrappers, + ) + xor_input_tensors.append(input_tensor_wrapper) + xor_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpElementWiseXor.op_name, + ) + xor_op.AddInputTensors(xor_input_tensors) + xor_op.AddOutputTensors(xor_output_tensors) + return xor_op diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index 74ffe24e3c4..b0c44dcae80 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -112,11 +112,22 @@ class OpElementWiseAnd: op_name: str = "ElementWiseAnd" +@dataclass(init=False, frozen=True) +class OpElementWiseAsin: + op_name: str = "ElementWiseAsin" + + @dataclass(init=False, frozen=True) class OpElementWiseAtan: op_name: str = "ElementWiseAtan" +@dataclass(init=False, frozen=True) +class OpElementWiseBinary: + op_name: str = "ElementWiseBinary" + param_operation: str = "operation" + + @dataclass(init=False, frozen=True) class OpElementWiseCeil: op_name = "ElementWiseCeil" @@ -240,6 +251,11 @@ class OpElementWiseSelect: op_name = "ElementWiseSelect" +@dataclass(init=False, frozen=True) +class OpElementWiseSign: + op_name: str = "ElementWiseSign" + + @dataclass(init=False, frozen=True) class OpElementWiseSquareRoot: op_name = "ElementWiseSquareRoot" @@ -250,6 +266,11 @@ class OpElementWiseSubtract: op_name = "ElementWiseSubtract" +@dataclass(init=False, frozen=True) +class OpElementWiseXor: + op_name: str = "ElementWiseXor" + + @dataclass(init=False, frozen=True) class OpElu: op_name: str = "Elu" diff --git a/backends/qualcomm/debugger/TARGETS b/backends/qualcomm/debugger/TARGETS index bfdc6967ed7..85f204f9718 100644 --- a/backends/qualcomm/debugger/TARGETS +++ b/backends/qualcomm/debugger/TARGETS @@ -1,6 +1,6 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "utils", srcs = ["utils.py"], deps = [ diff --git a/backends/qualcomm/debugger/utils.py b/backends/qualcomm/debugger/utils.py index 2c7be66fb68..b1d3ea84900 100644 --- a/backends/qualcomm/debugger/utils.py +++ b/backends/qualcomm/debugger/utils.py @@ -267,11 +267,6 @@ def qnn_context_binary_generator( assert os.path.isfile(f"{self.tmp_dir}/{binary_name}.bin"), result.stderr def qnn_net_run(self, graph_name="forward.serialized"): - input_list = "" - for idx, _ in enumerate(self.sample_input): - input_name = f"input_{idx}_0.raw" - input_list += input_name + " " - input_list = input_list.strip() + "\n" self.config["backend_extension_config"]["backend_extensions"][ "shared_library_path" @@ -304,7 +299,6 @@ def qnn_net_run(self, graph_name="forward.serialized"): ] self.adb.push( inputs=self.sample_input, - input_list=input_list, files=files, ) self.adb.execute(custom_runner_cmd=" ".join(cmds)) diff --git a/backends/qualcomm/partition/common_defs.py b/backends/qualcomm/partition/common_defs.py index b427c59ce07..7a2924fe756 100644 --- a/backends/qualcomm/partition/common_defs.py +++ b/backends/qualcomm/partition/common_defs.py @@ -10,13 +10,25 @@ from executorch.exir.dialects._ops import ops as exir_ops not_supported_operator = [ - exir_ops.edge.aten.clone.default, - exir_ops.edge.aten.slice_scatter.default, - exir_ops.edge.aten.copy.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, exir_ops.edge.quantized_decomposed.embedding_4bit.dtype, ] -to_be_implemented_operator = [] +to_be_implemented_operator = [ + exir_ops.edge.aten._adaptive_avg_pool3d.default, + exir_ops.edge.aten.adaptive_max_pool2d.default, + exir_ops.edge.aten.avg_pool3d.default, + exir_ops.edge.aten.div.Tensor_mode, + exir_ops.edge.aten.log10.default, + exir_ops.edge.aten.log1p.default, + exir_ops.edge.aten.log2.default, + exir_ops.edge.aten.max_pool3d_with_indices.default, + exir_ops.edge.aten.median.default, + exir_ops.edge.aten.median.dim, + exir_ops.edge.aten.round.decimals, + exir_ops.edge.aten.le.Scalar, + exir_ops.edge.aten.trunc.default, +] constant_operator = [ exir_ops.edge.aten.arange.start_step, diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 4eee818efe5..97e0b4bd109 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -44,6 +44,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: torch.ops.aten.adaptive_avg_pool2d.default, torch.ops.aten.col2im.default, torch.ops.aten.elu.default, + torch.ops.aten.floor_divide.default, torch.ops.aten.hardsigmoid.default, torch.ops.aten.hardswish.default, torch.ops.aten.im2col.default, diff --git a/backends/qualcomm/quantizer/annotators.py b/backends/qualcomm/quantizer/annotators.py index 38a8bc6ebe6..a27ad2b4a5c 100644 --- a/backends/qualcomm/quantizer/annotators.py +++ b/backends/qualcomm/quantizer/annotators.py @@ -162,6 +162,11 @@ def annotate_single_in_single_out( ) +@register_annotator([torch.ops.aten.to.dtype]) +def annotate_to_dtype(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.atan.default]) def annotate_atan(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) @@ -209,7 +214,7 @@ def annotate_add(node: Node, quantization_config: QuantizationConfig) -> None: @register_annotator([torch.ops.aten.amax.default]) def annotate_amax(node: Node, quantization_config: QuantizationConfig) -> None: - annotate_binary(node, quantization_config) + annotate_single_in_single_out(node, quantization_config) @register_annotator([torch.ops.aten.argmax.default]) @@ -219,7 +224,7 @@ def annotate_argmax(node: Node, quantization_config: QuantizationConfig) -> None @register_annotator([torch.ops.aten.amin.default]) def annotate_amin(node: Node, quantization_config: QuantizationConfig) -> None: - annotate_binary(node, quantization_config) + annotate_single_in_single_out(node, quantization_config) @register_annotator([torch.ops.aten.argmin.default]) @@ -428,6 +433,17 @@ def annotate_clamp(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.index_select.default]) +def annotate_index_select(node: Node, quantization_config: QuantizationConfig) -> None: + # args[2] = indices, which should be int + annotate_single_in_single_out(node, quantization_config) + + +@register_annotator([torch.ops.aten.flip.default]) +def annotate_flip(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.floor.default]) def annotate_floor(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) @@ -463,6 +479,11 @@ def annotate_sin(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.floor_divide.default]) +def annotate_floor_divide(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_binary(node, quantization_config) + + @register_annotator([torch.ops.aten.scalar_tensor.default]) def annotate_scalar_tensor(node: Node, quantization_config: QuantizationConfig) -> None: if _is_annotated([node]): @@ -626,6 +647,11 @@ def annotate_softmax(node: Node, quantization_config: QuantizationConfig) -> Non annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.asin.default]) +def annotate_asin(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.linalg_vector_norm.default]) def annotate_linalg_vector_norm( node: Node, quantization_config: QuantizationConfig @@ -658,11 +684,32 @@ def annotate_select(node: Node, quantization_config: QuantizationConfig) -> None annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.sign.default]) +def annotate_sign(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.slice.Tensor]) def annotate_slice(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) +@register_annotator([torch.ops.aten.slice_scatter.default]) +def annotate_slice_scatter(node: Node, quantization_config: QuantizationConfig) -> None: + input = node.args[0] + value = node.args[1] + + input_qspec_map = {} + input_qspec_map[input] = quantization_config.input_activation + input_qspec_map[value] = SharedQuantizationSpec((input, node)) + + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=SharedQuantizationSpec((input, node)), + _annotated=True, + ) + + @register_annotator([torch.ops.aten.sqrt.default]) def annotate_sqrt(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) @@ -802,6 +849,11 @@ def annotate_bitwise_or(node: Node, quantization_config: QuantizationConfig) -> annotate_binary(node, quantization_config) +@register_annotator([torch.ops.aten.bitwise_xor.Tensor, torch.ops.aten.__xor__.Tensor]) +def annotate_xor(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_binary(node, quantization_config) + + @register_annotator([torch.ops.aten.pow.Tensor_Tensor]) def annotate_pow(node: Node, quantization_config: QuantizationConfig) -> None: annotate_single_in_single_out(node, quantization_config) @@ -1047,6 +1099,7 @@ def annotate_cdist(node: Node, quantization_config: QuantizationConfig) -> None: torch.ops.aten.conv1d.default, torch.ops.aten.conv_transpose2d.input, torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.convolution.default, ] ) def annotate_conv(node: Node, quantization_config: QuantizationConfig) -> None: @@ -1101,6 +1154,12 @@ def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None if _is_annotated([node]): return + # block quantization + if quantization_config.block_size is not None: + quantization_config.weight.observer_or_fake_quant_ctr.p.keywords.update( + {"block_size": quantization_config.block_size} + ) + annotate_input_qspec_map( node, act_node, diff --git a/backends/qualcomm/quantizer/custom_annotation.py b/backends/qualcomm/quantizer/custom_annotation.py index 331622ee71b..3f10dbaa3fc 100644 --- a/backends/qualcomm/quantizer/custom_annotation.py +++ b/backends/qualcomm/quantizer/custom_annotation.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. +from enum import Enum, unique from typing import Sequence import torch @@ -12,8 +13,11 @@ ) from executorch.backends.qualcomm.quantizer.quantizer import ( get_16a8w_qnn_ptq_config, + get_16a8w_qnn_qat_config, get_8a8w_qnn_ptq_config, + get_8a8w_qnn_qat_config, get_ptq_per_channel_quant_config, + get_qat_per_channel_quant_config, QuantizationConfig, ) from executorch.exir.dialects._ops import ops as exir_ops @@ -28,6 +32,36 @@ ) +def annotate_down_proj( + gm: torch.fx.GraphModule, quantization_config: QuantizationConfig +): + for node in gm.graph.nodes: + if ( + node.target == torch.ops.aten.conv2d.default + and any(s in node.meta["stack_trace"] for s in ["forward_feedfoward_conv"]) + and node.args[0].target == torch.ops.aten.mul.Tensor + ): + input_qspec_map = {} + input_qspec_map[node.args[0]] = quantization_config.input_activation + input_qspec_map[node.args[1]] = quantization_config.weight + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + +@unique +class StaticLLMQuantConfig(Enum): + """ + Layer namespace configuration for Qualcomm's static LLaMA quantization. + """ + + wq_sha = "wq_sha" # Query weight (single head) + wk_sha = "wk_sha" # Key weight (single head) + wv_sha = "wv_sha" # Value weight (single head) + + def annotate_eurobert(gm: torch.fx.GraphModule): """ QNN does not support int32 -> signed 16bit quant @@ -89,7 +123,12 @@ def annotate_mimi_decoder(gm: torch.fx.GraphModule): break -def annotate_linear_16a8w_in_affine_layer(gm: torch.fx.GraphModule) -> None: +def annotate_output_16a8w(gm: torch.fx.GraphModule, is_qat: bool = False) -> None: + """ + This function is for static LLM models. + This function will annotate the last conv(linear), which is the lm_head, as 16a8w. + """ + def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None: input_qspec_map = {} input_act = node.args[0] @@ -105,9 +144,14 @@ def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None _annotated=True, ) - quantization_config_16a8w_per_channel = get_ptq_per_channel_quant_config( - torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver - ) + if is_qat: + quantization_config_16a8w_per_channel = get_qat_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) + else: + quantization_config_16a8w_per_channel = get_ptq_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) for node in gm.graph.nodes: if node.op == "call_function" and node.target == torch.ops.aten.conv2d.default: if "nn_module_stack" in node.meta: @@ -153,10 +197,54 @@ def annotate_prefill_kv_output(gm: torch.fx.GraphModule, kv_quant_attrs: dict): ) -def annotate_matmul_16a8w( # noqa: C901 - gm: torch.fx.GraphModule, annotate_conv=True +def annotate_qkv_proj_sha( + gm: torch.fx.GraphModule, + quantization_config: QuantizationConfig, + qkv_tags: set[StaticLLMQuantConfig], +): + """ + Annotates QKV projection layers in a GraphModule for quantization, + specifically layers defined in StaticLLMQuantConfig. + + Args: + qkv_tags (set[StaticLLMQuantConfig]): A set of enum tags indicating which QKV layers + (e.g., wq, wk, wv) should be annotated for quantization. Only tags defined in + StaticLLMQuantConfig are allowed. + + Raises: + ValueError: If any tag in `qkv_tags` is not among the allowed enum members. + """ + + # Get all valid tags from the StaticLLMQuantConfig enum + allowed_tags = set(StaticLLMQuantConfig) + invalid_tags = qkv_tags - allowed_tags + if invalid_tags: + raise ValueError( + f"Invalid qkv tags: {invalid_tags}. Allowed tags are: {allowed_tags}" + ) + + for node in gm.graph.nodes: + if node.target == torch.ops.aten.conv2d.default and any( + tag.value in node.meta["stack_trace"] for tag in qkv_tags + ): + input_qspec_map = {} + input_qspec_map[node.args[0]] = quantization_config.input_activation + input_qspec_map[node.args[1]] = quantization_config.weight + if len(node.args) > 2 and isinstance(node.args[2], Node): + input_qspec_map[node.args[2]] = quantization_config.bias(node) + node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + +def annotate_kv_8bit( # noqa: C901 + gm: torch.fx.GraphModule, + is_qat=False, ) -> None: """ + This function is for static LLM models. This function is specific for matmul op 16a8w. For k, we will tag such as the below, and for v, we will tag 8a until conv op. @@ -202,21 +290,6 @@ def annotate_cat(node: Node, quantization_config: QuantizationConfig): _annotated=True, ) - def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None: - input_qspec_map = {} - input_act = node.args[0] - input_spec = quantization_config.input_activation - input_qspec_map[input_act] = input_spec - - weight = node.args[1] - input_qspec_map[weight] = quantization_config.weight - - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=quantization_config.output_activation, - _annotated=True, - ) - def annotate_rms_norm(node: Node, quantization_config: QuantizationConfig) -> None: act_node = node.args[0] weight_node = node.args[2] @@ -238,7 +311,6 @@ def annotate_rms_norm(node: Node, quantization_config: QuantizationConfig) -> No def annotate_single_in_single_out( node: Node, quantization_config: QuantizationConfig ) -> None: - input_qspec_map = {} input_act = node.args[0] input_qspec_map[input_act] = quantization_config.input_activation @@ -252,7 +324,6 @@ def annotate_single_in_single_out( def annotate_single_in_share_out( node: Node, quantization_config: QuantizationConfig ) -> None: - input_qspec_map = {} input_act = node.args[0] input_qspec_map[input_act] = quantization_config.input_activation @@ -283,16 +354,15 @@ def annotate_stack(node: Node, quantization_config: QuantizationConfig) -> None: _annotated=True, ) - def annotate_matmul_input1(node: Node): - quantization_config_8a8w = get_8a8w_qnn_ptq_config( - act_symmetric=True, act_observer=MinMaxObserver - ) - quantization_config_8a4w_per_channel = get_ptq_per_channel_quant_config( - act_dtype=torch.uint8, - weight_dtype=torch.int4, - act_observer=MinMaxObserver, - act_symmetric=True, - ) + def annotate_matmul_input1(node: Node, is_qat: str): + if is_qat: + quantization_config_8a8w = get_8a8w_qnn_qat_config( + act_symmetric=True, act_observer=MinMaxObserver + ) + else: + quantization_config_8a8w = get_8a8w_qnn_ptq_config( + act_symmetric=True, act_observer=MinMaxObserver + ) while isinstance(node, Node) and node.op == "call_function": if node.target in [ torch.ops.aten.permute.default, @@ -300,6 +370,7 @@ def annotate_matmul_input1(node: Node): torch.ops.aten.transpose.int, torch.ops.aten.view.default, torch.ops.aten.reshape.default, + torch.ops.aten.slice.Tensor, ]: annotate_single_in_single_out(node, quantization_config_8a8w) node = node.args[0] @@ -318,24 +389,35 @@ def annotate_matmul_input1(node: Node): # For k, we tag 8a until add or sub op (rotatary embedding). # The arguments of cat op: (the past kv cache, the new kv cache) node = node.args[0][1] - elif node.target == torch.ops.aten.conv2d.default: - if annotate_conv: - annotate_conv2d( - node, quantization_config=quantization_config_8a4w_per_channel - ) - break - elif node.target in [torch.ops.aten.add.Tensor, torch.ops.aten.sub.Tensor]: + elif node.target in [ + torch.ops.aten.add.Tensor, + torch.ops.aten.sub.Tensor, + torch.ops.aten.matmul.default, + torch.ops.aten.conv2d.default, + ]: break else: print(f"The node ({node}) is not expected in the input1 of the matmul") node = node.args[0] - quantization_config_16a8w = get_16a8w_qnn_ptq_config(act_observer=MinMaxObserver) + if is_qat: + quantization_config_16a8w = get_16a8w_qnn_qat_config( + act_observer=MinMaxObserver + ) + else: + quantization_config_16a8w = get_16a8w_qnn_ptq_config( + act_observer=MinMaxObserver + ) for node in gm.graph.nodes: - if node.op == "call_function" and node.target == torch.ops.aten.matmul.default: + if ( + node.op == "call_function" + and node.target == torch.ops.aten.matmul.default + and all(arg.op == "call_function" for arg in node.args) + ): + # Only apply custom annotation on Q @ K^T @ V annotate_matmul(node, quantization_config_16a8w) - annotate_matmul_input1(node.args[1]) + annotate_matmul_input1(node.args[1], is_qat=is_qat) def custom_annotate_llama_matmul_16a8w(gm: torch.fx.GraphModule) -> None: # noqa: C901 diff --git a/backends/qualcomm/quantizer/observers/per_block_param_observer.py b/backends/qualcomm/quantizer/observers/per_block_param_observer.py index 7d605b12cf8..b3f854db527 100644 --- a/backends/qualcomm/quantizer/observers/per_block_param_observer.py +++ b/backends/qualcomm/quantizer/observers/per_block_param_observer.py @@ -34,7 +34,11 @@ def __init__( eps=eps, **kwargs, ) + self.dtype = dtype self.block_size = block_size + # TODO: expand this when QNN starts to support more configurations + self.bitwidth_of_scale = 4 + self.num_steps = 2**self.bitwidth_of_scale self.calibrated = False def forward(self, input: torch.Tensor): diff --git a/backends/qualcomm/quantizer/qconfig.py b/backends/qualcomm/quantizer/qconfig.py index 748128ceafd..2f26cd27d31 100644 --- a/backends/qualcomm/quantizer/qconfig.py +++ b/backends/qualcomm/quantizer/qconfig.py @@ -1,3 +1,9 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple @@ -49,7 +55,10 @@ def _derive_bias_qparams_fn( derived_zero = torch.zeros(derived_scale.size()).to(torch.int32) if isinstance(weight_obs_or_fq, PerBlockParamObserver): # keep maximum scale of each channel for bias - derived_scale = derived_scale.view(derived_scale.size(0), -1).amax(dim=-1) + derived_scale = ( + derived_scale.view(derived_scale.size(0), -1).amax(dim=-1) + / weight_obs_or_fq.num_steps + ) derived_zero = derived_zero.view(derived_zero.size(0), -1).amax(dim=-1) return (derived_scale, derived_zero) @@ -187,6 +196,65 @@ def get_16a8w_qnn_ptq_config( return quantization_config +def get_16a8w_qnn_qat_config( + act_observer=MovingAverageMinMaxObserver, +) -> QuantizationConfig: + extra_args: Dict[str, Any] = {"eps": 2**-20} + act_fake_quant_ctr = FakeQuantize.with_args( + dtype=torch.int32, + quant_min=torch.iinfo(torch.uint16).min, + quant_max=torch.iinfo(torch.uint16).max, + qscheme=torch.per_tensor_affine, + reduce_range=True, + observer=act_observer.with_args(**extra_args), + ) + act_quantization_spec = QuantizationSpec( + dtype=torch.int32, + quant_min=torch.iinfo(torch.uint16).min, + quant_max=torch.iinfo(torch.uint16).max, + qscheme=torch.per_tensor_affine, + observer_or_fake_quant_ctr=act_fake_quant_ctr, + ) + weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( + dtype=torch.int8, + quant_min=torch.iinfo(torch.int8).min + 1, + quant_max=torch.iinfo(torch.int8).max, + qscheme=torch.per_tensor_symmetric, + reduce_range=True, + observer=MovingAverageMinMaxObserver, + ) + weight_quantization_spec = QuantizationSpec( + dtype=torch.int8, + quant_min=torch.iinfo(torch.int8).min + 1, + quant_max=torch.iinfo(torch.int8).max, + qscheme=torch.per_tensor_symmetric, + ch_axis=0, + observer_or_fake_quant_ctr=weight_fake_quant_ctr, + ) + bias_fake_quant_ctr = FakeQuantize.with_args( + dtype=torch.int32, + quant_min=torch.iinfo(torch.int32).min, + quant_max=torch.iinfo(torch.int32).max, + qscheme=torch.per_tensor_symmetric, + observer=MovingAverageMinMaxObserver, + ) + bias_quantization_spec = QuantizationSpec( + dtype=torch.int32, + quant_min=torch.iinfo(torch.int32).min, + quant_max=torch.iinfo(torch.int32).max, + qscheme=torch.per_tensor_symmetric, + observer_or_fake_quant_ctr=bias_fake_quant_ctr, + ) + quantization_config = QuantizationConfig( + input_activation=act_quantization_spec, + output_activation=act_quantization_spec, + weight=weight_quantization_spec, + bias=bias_quantization_spec, + ) + + return quantization_config + + def get_16a16w_qnn_ptq_config( act_observer=MovingAverageMinMaxObserver, ) -> QuantizationConfig: @@ -328,7 +396,7 @@ def get_ptq_per_block_quant_config( ) -# TODO merge qat and ptq to a fucntion, and use a bool flag to control it +# TODO merge qat and ptq to a function, and use a bool flag to control it def get_8a8w_qnn_qat_config( act_symmetric: bool = False, act_observer=MovingAverageMinMaxObserver ) -> QuantizationConfig: @@ -337,7 +405,6 @@ def get_8a8w_qnn_qat_config( qscheme=( torch.per_tensor_symmetric if act_symmetric else torch.per_tensor_affine ), - reduce_range=True, observer=act_observer, ) act_quantization_spec = QuantizationSpec( @@ -459,6 +526,7 @@ def get_qat_per_channel_quant_config( act_dtype=torch.uint8, weight_dtype=torch.int8, act_observer=MovingAverageMinMaxObserver, + act_symmetric=False, ) -> QuantizationConfig: supported_act_types = { torch.uint8, @@ -476,21 +544,38 @@ def get_qat_per_channel_quant_config( ), f"weight_dtype, {weight_dtype} is not one of supported types, {supported_weight_dtypes}" # torch does not support uint16 quantization, use int32 to bypass - act_fake_quant_ctr = FakeQuantize.with_args( - dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, - quant_min=torch.iinfo(act_dtype).min, - quant_max=torch.iinfo(act_dtype).max, - qscheme=torch.per_tensor_affine, - reduce_range=True, - observer=act_observer, - ) - act_quantization_spec = QuantizationSpec( - dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, - quant_min=torch.iinfo(act_dtype).min, - quant_max=torch.iinfo(act_dtype).max, - qscheme=torch.per_tensor_affine, - observer_or_fake_quant_ctr=act_fake_quant_ctr, - ) + if act_symmetric: + # If zero_point is 128, htp can do optimizations. + # If we keep quant_min and quant_max none, observer will default use 128 as zero_point. + # If we provide uint8 quant_min/max, it will use 127 as zero_point, which is undesired. + act_fake_quant_ctr = FakeQuantize.with_args( + dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, + qscheme=torch.per_tensor_symmetric, + reduce_range=True, + observer=act_observer, + ) + act_quantization_spec = QuantizationSpec( + dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, + qscheme=torch.per_tensor_symmetric, + ch_axis=0, + observer_or_fake_quant_ctr=act_fake_quant_ctr, + ) + else: + act_fake_quant_ctr = FakeQuantize.with_args( + dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, + quant_min=torch.iinfo(act_dtype).min, + quant_max=torch.iinfo(act_dtype).max, + qscheme=torch.per_tensor_affine, + reduce_range=True, + observer=act_observer, + ) + act_quantization_spec = QuantizationSpec( + dtype=torch.int32 if act_dtype == torch.uint16 else act_dtype, + quant_min=torch.iinfo(act_dtype).min, + quant_max=torch.iinfo(act_dtype).max, + qscheme=torch.per_tensor_affine, + observer_or_fake_quant_ctr=act_fake_quant_ctr, + ) weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( dtype=torch.int8 if weight_dtype == torch.int4 else weight_dtype, diff --git a/backends/qualcomm/quantizer/quantizer.py b/backends/qualcomm/quantizer/quantizer.py index e14d73f521d..5943b54d968 100644 --- a/backends/qualcomm/quantizer/quantizer.py +++ b/backends/qualcomm/quantizer/quantizer.py @@ -23,6 +23,7 @@ get_16a4w_qnn_ptq_config, get_16a4w_qnn_qat_config, get_16a8w_qnn_ptq_config, + get_16a8w_qnn_qat_config, get_8a8w_qnn_ptq_config, get_8a8w_qnn_qat_config, get_ptq_per_block_quant_config, @@ -39,6 +40,7 @@ "QuantDtype", "get_16a4w_qnn_ptq_config", "get_16a8w_qnn_ptq_config", + "get_16a8w_qnn_qat_config", "get_16a16w_qnn_ptq_config", "get_8a8w_qnn_ptq_config", "get_8a8w_qnn_qat_config", diff --git a/backends/qualcomm/recipes/TARGETS b/backends/qualcomm/recipes/TARGETS new file mode 100644 index 00000000000..12d1bac6f12 --- /dev/null +++ b/backends/qualcomm/recipes/TARGETS @@ -0,0 +1,53 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_library( + name = "qnn_recipes", + srcs = [ + "__init__.py", + ], + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + deps = [ + "//executorch/export:recipe_registry", + ":qnn_recipe_provider", + ":qnn_recipe_types", + ], +) + +runtime.python_library( + name = "qnn_recipe_provider", + srcs = [ + "qnn_recipe_provider.py", + ], + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + deps = [ + "//caffe2:torch", + "//executorch/export:lib", + "//executorch/backends/qualcomm/partition:partition", + "//executorch/backends/qualcomm/serialization:serialization", + "//executorch/backends/qualcomm/utils:utils", + "//executorch/backends/qualcomm/_passes:passes", + ":qnn_recipe_types", + ], +) + +runtime.python_library( + name = "qnn_recipe_types", + srcs = [ + "qnn_recipe_types.py", + ], + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + deps = [ + "//executorch/export:lib", + ], +) diff --git a/backends/qualcomm/recipes/__init__.py b/backends/qualcomm/recipes/__init__.py new file mode 100644 index 00000000000..ee0985584d6 --- /dev/null +++ b/backends/qualcomm/recipes/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""QNN Recipe module for ExecuTorch""" +from executorch.export import recipe_registry + +from .qnn_recipe_provider import QNNRecipeProvider +from .qnn_recipe_types import QNNRecipeType + +# Auto-register XNNPACK recipe provider +recipe_registry.register_backend_recipe_provider(QNNRecipeProvider()) + +__all__ = ["QNNRecipeProvider", "QNNRecipeType"] diff --git a/backends/qualcomm/recipes/qnn_recipe_provider.py b/backends/qualcomm/recipes/qnn_recipe_provider.py new file mode 100644 index 00000000000..fcfab0c3bd1 --- /dev/null +++ b/backends/qualcomm/recipes/qnn_recipe_provider.py @@ -0,0 +1,179 @@ +# 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. + +# pyre-strict + +import logging +from typing import Any, Optional, Sequence + +from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager +from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner +from executorch.backends.qualcomm.recipes.qnn_recipe_types import ( + QNN_BACKEND, + QNNRecipeType, +) +from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_soc_to_chipset_map, + qnn_edge_config, +) +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + RecipeType, +) + + +class QNNRecipeProvider(BackendRecipeProvider): + @property + def backend_name(self) -> str: + return QNN_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(QNNRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + """Create QNN recipe for different precisions and SoC targets""" + + if recipe_type not in self.get_supported_recipes(): + return None + + self._validate_recipe_kwargs(recipe_type, kwargs) + + if recipe_type == QNNRecipeType.FP16: + return self._build_fp16_recipe(recipe_type, kwargs) + + return None + + def _validate_recipe_kwargs(self, recipe_type: RecipeType, kwargs: Any) -> None: + """Validate kwargs for each recipe type""" + expected_keys = self._get_expected_keys(recipe_type) + + unexpected = set(kwargs.keys()) - expected_keys + if unexpected: + logging.warning( + f"QNN Recipe '{recipe_type.value}' received unexpected parameters: {list(unexpected)}, ignoring them" + ) + + self._validate_soc_parameter(kwargs) + self._validate_partitioner_parameters(kwargs) + + def _get_expected_keys(self, recipe_type: RecipeType) -> set: + """Get expected parameter keys for a recipe type""" + _ = recipe_type + common_keys = { + "soc_model", + "skip_node_id_set", + "skip_node_op_set", + "skip_mutable_buffer", + } + return common_keys + + def _validate_soc_parameter(self, kwargs: Any) -> None: + """Validate soc_model parameter""" + if "soc_model" in kwargs: + soc_model = kwargs["soc_model"] + if isinstance(soc_model, str): + try: + soc_model = get_soc_to_chipset_map()[soc_model] + kwargs["soc_model"] = soc_model + except KeyError: + raise ValueError( + f"Invalid SoC model '{soc_model}'. Supported models: {[e.name for e in get_soc_to_chipset_map()]}" + ) + elif not isinstance(soc_model, QcomChipset): + raise ValueError( + f"Parameter 'soc_model' must be a QcomChipset enum or string, got {type(soc_model)}" + ) + else: + raise ValueError("Parameter 'soc_model' is required") + + def _validate_partitioner_parameters(self, kwargs: Any) -> None: + """Validate partitioner parameters""" + if "skip_node_id_set" in kwargs: + skip_node_id_set = kwargs["skip_node_id_set"] + if skip_node_id_set is not None and not isinstance(skip_node_id_set, set): + raise ValueError( + f"Parameter 'skip_node_id_set' must be a set or None, got {type(skip_node_id_set)}" + ) + + if "skip_node_op_set" in kwargs: + skip_node_op_set = kwargs["skip_node_op_set"] + if skip_node_op_set is not None and not isinstance(skip_node_op_set, set): + raise ValueError( + f"Parameter 'skip_node_op_set' must be a set or None, got {type(skip_node_op_set)}" + ) + + if "skip_mutable_buffer" in kwargs: + skip_mutable_buffer = kwargs["skip_mutable_buffer"] + if not isinstance(skip_mutable_buffer, bool): + raise ValueError( + f"Parameter 'skip_mutable_buffer' must be a boolean, got {type(skip_mutable_buffer)}" + ) + + def _build_fp16_recipe( + self, + recipe_type: RecipeType, + kwargs: Any, + ) -> ExportRecipe: + soc_model = kwargs["soc_model"] + skip_node_id_set = kwargs.get("skip_node_id_set", None) + skip_node_op_set = kwargs.get("skip_node_op_set", None) + skip_mutable_buffer = kwargs.get("skip_mutable_buffer", False) + + lowering_recipe = self._get_qnn_lowering_recipe( + use_fp16=True, + soc_model=soc_model, + skip_node_id_set=skip_node_id_set, + skip_node_op_set=skip_node_op_set, + skip_mutable_buffer=skip_mutable_buffer, + ) + + return ExportRecipe( + name=recipe_type.value, + aten_transform_passes=[ + lambda method_, ep: QnnPassManager().transform_for_export_pipeline(ep) + ], + lowering_recipe=lowering_recipe, + ) + + def _get_qnn_lowering_recipe( + self, + use_fp16: bool, + soc_model: QcomChipset, + skip_node_id_set: Optional[set] = None, + skip_node_op_set: Optional[set] = None, + skip_mutable_buffer: bool = False, + ) -> LoweringRecipe: + """Get QNN lowering recipe with optional precision and SoC target""" + backend_options = generate_htp_compiler_spec(use_fp16=use_fp16) + + compile_specs = generate_qnn_executorch_compiler_spec( + soc_model=soc_model, + backend_options=backend_options, + ) + + partitioner = QnnPartitioner( + compiler_specs=compile_specs, + skip_node_id_set=skip_node_id_set, + skip_node_op_set=skip_node_op_set, + skip_mutable_buffer=skip_mutable_buffer, + ) + + edge_compile_config = qnn_edge_config() + + return LoweringRecipe( + partitioners=[partitioner], + edge_transform_passes=[ + lambda method_, ep: QnnPassManager().get_to_edge_transform_passes(ep) + ], + edge_compile_config=edge_compile_config, + ) diff --git a/backends/qualcomm/recipes/qnn_recipe_types.py b/backends/qualcomm/recipes/qnn_recipe_types.py new file mode 100644 index 00000000000..f41f494fce6 --- /dev/null +++ b/backends/qualcomm/recipes/qnn_recipe_types.py @@ -0,0 +1,28 @@ +# 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. + +# pyre-strict + +from executorch.export import RecipeType + + +QNN_BACKEND: str = "qnn" + + +class QNNRecipeType(RecipeType): + """QNN-specific recipe types""" + + # FP16 precision recipe, accepts kwargs: + # 1. soc_model + # 2. skip_node_id_set + # 3. skip_node_op_set + # 4. skip_mutable_buffer + + FP16 = "qnn_fp16" + + @classmethod + def get_backend_name(cls) -> str: + return QNN_BACKEND diff --git a/backends/qualcomm/runtime/CMakeLists.txt b/backends/qualcomm/runtime/CMakeLists.txt index eb31bee7a53..1a35ec8366f 100644 --- a/backends/qualcomm/runtime/CMakeLists.txt +++ b/backends/qualcomm/runtime/CMakeLists.txt @@ -28,6 +28,13 @@ target_sources( PRIVATE ${CMAKE_CURRENT_LIST_DIR}/QnnManager.cpp ) +# qnn_backend_options +target_sources( + qnn_backend_options + INTERFACE ${CMAKE_CURRENT_LIST_DIR}/QnnBackendOptions.h + PRIVATE ${CMAKE_CURRENT_LIST_DIR}/QnnBackendOptions.cpp +) + # logging target_sources( qnn_executorch_logging diff --git a/backends/qualcomm/runtime/QnnBackendOptions.cpp b/backends/qualcomm/runtime/QnnBackendOptions.cpp new file mode 100644 index 00000000000..1ce48cfcd61 --- /dev/null +++ b/backends/qualcomm/runtime/QnnBackendOptions.cpp @@ -0,0 +1,66 @@ +/* + * 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. + */ +#include +#include +#include + +namespace executorch { +namespace backends { +namespace qnn { + +using namespace qnn_delegate; + +template +T get_option(T aot_option) { + executorch::runtime::Error status; + executorch::runtime::BackendOption backend_option; + + if constexpr (std::is_same_v) { + std::strncpy( + backend_option.key, + QNN_RUNTIME_LOG_LEVEL, + runtime::kMaxOptionKeyLength); + backend_option.key[runtime::kMaxOptionKeyLength - 1] = '\0'; + backend_option.value = -1; + } else if constexpr (std::is_same_v) { + std::strncpy( + backend_option.key, + QNN_RUNTIME_HTP_PERFORMANCE_MODE, + runtime::kMaxOptionKeyLength); + backend_option.key[runtime::kMaxOptionKeyLength - 1] = '\0'; + backend_option.value = -1; + } else if constexpr (std::is_same_v) { + std::strncpy( + backend_option.key, + QNN_RUNTIME_PROFILE_LEVEL, + runtime::kMaxOptionKeyLength); + backend_option.key[runtime::kMaxOptionKeyLength - 1] = '\0'; + backend_option.value = -1; + } + + // This will call get_option under runtime backend interface + status = get_option(QNN_BACKEND, backend_option); + + if (status != executorch::runtime::Error::Ok) { + return aot_option; + } else { + return static_cast(std::get(backend_option.value)); + } +} + +// Explicit instantiations +template QnnExecuTorchLogLevel get_option( + QnnExecuTorchLogLevel); +template QnnExecuTorchHtpPerformanceMode get_option< + QnnExecuTorchHtpPerformanceMode>(QnnExecuTorchHtpPerformanceMode); +template QnnExecuTorchProfileLevel get_option( + QnnExecuTorchProfileLevel); + +} // namespace qnn +} // namespace backends +} // namespace executorch diff --git a/backends/qualcomm/runtime/QnnBackendOptions.h b/backends/qualcomm/runtime/QnnBackendOptions.h new file mode 100644 index 00000000000..a601a4202c0 --- /dev/null +++ b/backends/qualcomm/runtime/QnnBackendOptions.h @@ -0,0 +1,41 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace qnn { + +/** + * @brief Storing runtime option value. + * @param is_set True when user calls set_option api to set option, else False. + */ +struct RuntimeOption { + bool is_set; + executorch::runtime::OptionValue value; +}; + +/** + * @brief + * Get the backend option. + * This method checks both AOT option and runtime option. + * If runtime option is provided, it will have a higher priority. + * + * @param aot_option The flatbuffer option under qc_compiler_spec.fbs. + */ + +template +T get_option(T aot_option); + +} // namespace qnn +} // namespace backends +} // namespace executorch diff --git a/backends/qualcomm/runtime/QnnExecuTorch.h b/backends/qualcomm/runtime/QnnExecuTorch.h index 2ca0cd61cd5..d8fbade3b3b 100644 --- a/backends/qualcomm/runtime/QnnExecuTorch.h +++ b/backends/qualcomm/runtime/QnnExecuTorch.h @@ -16,14 +16,18 @@ #include #endif +#define QNN_BACKEND "QnnBackend" +#define QNN_RUNTIME_LOG_LEVEL "qnn_runtime_log_level" +#define QNN_RUNTIME_HTP_PERFORMANCE_MODE "qnn_runtime_htp_performance_mode" +#define QNN_RUNTIME_PROFILE_LEVEL "qnn_runtime_profile_level" + #ifdef __cplusplus extern "C" { #endif // __cplusplus // This could be: // 1. qnn_context_binary -// 2. QnnQcirCustomProtocol -// 3. QnnContextCustomProtocol +// 2. QnnContextCustomProtocol // To check if it is custom protocol, users can deserialize the binary using // QnnCustomProtocol and check the status typedef struct { diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp index 01bf13603d6..988c4b84a68 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp @@ -8,10 +8,12 @@ #include #include +#include #include #include #include - +#include +#include namespace executorch { namespace backends { namespace qnn { @@ -26,6 +28,7 @@ using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::MemoryAllocator; using executorch::runtime::Result; +using executorch::runtime::Span; // ========== Public method implementations ========================= constexpr const char* QNN_COMPILE_SPEC = "qnn_compile_spec"; @@ -48,8 +51,7 @@ Result QnnExecuTorchBackend::init( qnn_context_blob.buffer = ctx_bin; } else { // This buffer will be verified again in QnnBackendCache. - QNN_EXECUTORCH_LOG_INFO( - "Deserializing processed data using QnnQcirCustomProtocol"); + QNN_EXECUTORCH_LOG_INFO("Deserializing processed data using Dlc"); qnn_context_blob.buffer = const_cast(processed->data()); qnn_context_blob.nbytes = processed->size(); } @@ -114,7 +116,7 @@ Result QnnExecuTorchBackend::init( Error QnnExecuTorchBackend::execute( BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const { + Span args) const { ET_CHECK_OR_RETURN_ERROR( delegate_map_rev_.count(handle) != 0, Internal, @@ -189,6 +191,77 @@ void QnnExecuTorchBackend::destroy(DelegateHandle* handle) const { } } +executorch::runtime::Error QnnExecuTorchBackend::set_option( + executorch::runtime::BackendOptionContext& context, + const executorch::runtime::Span& + backend_options) { + std::lock_guard guard(runtime_option_mutex_); + size_t matches = backend_options.size(); + for (const auto& option : backend_options) { + if (strcmp(option.key, QNN_RUNTIME_LOG_LEVEL) == 0) { + if (auto* val = std::get_if(&option.value)) { + qnn_runtime_log_level_.value = *val; + qnn_runtime_log_level_.is_set = true; + } + } else if (strcmp(option.key, QNN_RUNTIME_HTP_PERFORMANCE_MODE) == 0) { + if (auto* val = std::get_if(&option.value)) { + qnn_runtime_performance_mode_.value = *val; + qnn_runtime_performance_mode_.is_set = true; + } + } else if (strcmp(option.key, QNN_RUNTIME_PROFILE_LEVEL) == 0) { + if (auto* val = std::get_if(&option.value)) { + qnn_runtime_profile_level_.value = *val; + qnn_runtime_profile_level_.is_set = true; + } + } else { + ET_LOG( + Error, + "Unable to set the following runtime option for QnnExecuTorchBackend: %s.", + option.key); + matches--; + } + } + + ET_CHECK_OR_RETURN_ERROR( + matches == backend_options.size(), + Internal, + "Some set options are not supported by QnnExecuTorchBackend. %zu options provided but only %zu is supported.", + backend_options.size(), + matches); + + return Error::Ok; +} + +executorch::runtime::Error QnnExecuTorchBackend::get_option( + executorch::runtime::BackendOptionContext& context, + executorch::runtime::Span& + backend_options) { + size_t matches = backend_options.size(); + for (size_t i = 0; i < backend_options.size(); ++i) { + // Set the value to what was stored by set_option + if (strcmp(backend_options[i].key, QNN_RUNTIME_LOG_LEVEL) == 0 && + qnn_runtime_log_level_.is_set) { + backend_options[i].value = qnn_runtime_log_level_.value; + } else if ( + strcmp(backend_options[i].key, QNN_RUNTIME_HTP_PERFORMANCE_MODE) == 0 && + qnn_runtime_performance_mode_.is_set) { + backend_options[i].value = qnn_runtime_performance_mode_.value; + } else if ( + strcmp(backend_options[i].key, QNN_RUNTIME_PROFILE_LEVEL) == 0 && + qnn_runtime_profile_level_.is_set) { + backend_options[i].value = qnn_runtime_profile_level_.value; + } else { + // either runtime never called set_option or key does not exist + matches--; + } + } + + if (matches != backend_options.size()) { + return Error::Internal; + } + return Error::Ok; +} + bool QnnExecuTorchBackend::is_available() const { return true; } @@ -214,7 +287,7 @@ void QnnExecuTorchBackend::erase_cached_delegate( namespace { auto cls = QnnExecuTorchBackend(); -executorch::runtime::Backend backend{"QnnBackend", &cls}; +executorch::runtime::Backend backend{QNN_BACKEND, &cls}; static auto success_with_compiler = register_backend(backend); } // namespace } // namespace qnn diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.h b/backends/qualcomm/runtime/QnnExecuTorchBackend.h index e83ec6b13b0..5cca7669b20 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.h +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.h @@ -7,6 +7,7 @@ */ #pragma once +#include #include #include #include @@ -32,7 +33,18 @@ class QnnExecuTorchBackend final executorch::runtime::Error execute( ET_UNUSED executorch::runtime::BackendExecutionContext& context, executorch::runtime::DelegateHandle* handle, - executorch::runtime::EValue** args) const override; + executorch::runtime::Span args) + const override; + + ET_NODISCARD executorch::runtime::Error set_option( + executorch::runtime::BackendOptionContext& context, + const executorch::runtime::Span& + backend_options) override; + + executorch::runtime::Error get_option( + executorch::runtime::BackendOptionContext& context, + executorch::runtime::Span& + backend_options) override; void destroy(executorch::runtime::DelegateHandle* handle) const override; @@ -45,10 +57,15 @@ class QnnExecuTorchBackend final void erase_cached_delegate(executorch::runtime::DelegateHandle* handle) const; mutable std::mutex mutex_; + mutable std::mutex runtime_option_mutex_; mutable std::unordered_map delegate_map_; mutable std::unordered_map delegate_map_rev_; + + RuntimeOption qnn_runtime_log_level_{false, 0}; + RuntimeOption qnn_runtime_performance_mode_{false, 0}; + RuntimeOption qnn_runtime_profile_level_{false, 0}; }; } // namespace qnn diff --git a/backends/qualcomm/runtime/QnnManager.cpp b/backends/qualcomm/runtime/QnnManager.cpp index 0dd0470a2b0..5e3220f25d9 100644 --- a/backends/qualcomm/runtime/QnnManager.cpp +++ b/backends/qualcomm/runtime/QnnManager.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include @@ -63,7 +64,8 @@ QnnManager::QnnManager( options->backend_options()->backend_type(); std::string library_path = options->library_path()->str(); - if (options->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo) { + if (get_option(options_->log_level()) >= + QnnExecuTorchLogLevel::kLogLevelInfo) { QNN_EXECUTORCH_LOG_INFO( "soc_model in soc_info: %s", EnumNameQcomChipset(options_->soc_info()->soc_model())); @@ -75,10 +77,12 @@ QnnManager::QnnManager( QNN_EXECUTORCH_LOG_INFO("library_path: %s", library_path.c_str()); QNN_EXECUTORCH_LOG_INFO("dump intermediate outputs: %s", IsTensorDump()); QNN_EXECUTORCH_LOG_INFO( - "log_level: %s", EnumNameQnnExecuTorchLogLevel(options_->log_level())); + "log_level: %s", + EnumNameQnnExecuTorchLogLevel(get_option(options_->log_level()))); QNN_EXECUTORCH_LOG_INFO( "profile_level: %s", - EnumNameQnnExecuTorchProfileLevel(options_->profile_level())); + EnumNameQnnExecuTorchProfileLevel( + get_option(options_->profile_level()))); QNN_EXECUTORCH_LOG_INFO( "the size of qnn context binary: %d", qnn_executorch_context_binary.nbytes); @@ -202,7 +206,8 @@ Error QnnManager::RegisterIonMem( return Error::Internal; } else if (backend_params_ptr_->qnn_mem_manager_ptr_->IsRegistered( tensor_wrapper->GetMemHandle(), data_ptr)) { - if (options_->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo) + if (get_option(options_->log_level()) >= + QnnExecuTorchLogLevel::kLogLevelInfo) QNN_EXECUTORCH_LOG_INFO( "Tensor name %s has been registered shared memory.", tensor_wrapper->GetName().c_str()); @@ -231,7 +236,8 @@ Error QnnManager::RegisterCustomMem( const std::shared_ptr& tensor_wrapper) { if (backend_params_ptr_->qnn_mem_manager_ptr_->IsRegistered( tensor_wrapper->GetMemHandle(), data_ptr)) { - if (options_->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo) + if (get_option(options_->log_level()) >= + QnnExecuTorchLogLevel::kLogLevelInfo) QNN_EXECUTORCH_LOG_INFO( "Tensor name %s has been registered shared memory.", tensor_wrapper->GetName().c_str()); @@ -251,7 +257,8 @@ Error QnnManager::RegisterCustomMem( Qnn_MemHandle_t pre_registered_handle = backend_params_ptr_->qnn_mem_manager_ptr_->GetPreRegisteredHandle(info); if (pre_registered_handle != nullptr) { - if (options_->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo) { + if (get_option(options_->log_level()) >= + QnnExecuTorchLogLevel::kLogLevelInfo) { QNN_EXECUTORCH_LOG_INFO( "Tensor name %s found a pre-registered memHandle.", tensor_wrapper->GetName().c_str()); @@ -284,7 +291,8 @@ Error QnnManager::RegisterCustomMem( data_ptr, unaligned_custom_mem_base, total_custom_mem_size, - tensor_offset) == Error::Ok, + tensor_offset, + info) == Error::Ok, Internal, "Fail to register to shared memory."); @@ -295,7 +303,7 @@ Error QnnManager::Init() { ET_CHECK_OR_RETURN_ERROR( LoadQnnLibrary() == Error::Ok, Internal, "Fail to load Qnn library"); logger_ = std::make_unique( - qnn_loaded_backend_, LoggingCallback, options_->log_level()); + qnn_loaded_backend_, LoggingCallback, get_option(options_->log_level())); std::vector graph_names; for (auto name : *options_->graph_name()) { graph_names.emplace_back(name->str()); @@ -492,7 +500,8 @@ Error QnnManager::ProfileExecuteData( const std::string& graph_name, executorch::runtime::EventTracer* event_tracer) { Qnn_ErrorHandle_t error = QNN_SUCCESS; - if (options_->profile_level() != QnnExecuTorchProfileLevel::kProfileOff) { + if (get_option(options_->profile_level()) != + QnnExecuTorchProfileLevel::kProfileOff) { error = backend_params_ptr_->qnn_graph_ptr_->ProfileExecuteData( graph_name, event_tracer); if (error != QNN_SUCCESS) { diff --git a/backends/qualcomm/runtime/backends/CMakeLists.txt b/backends/qualcomm/runtime/backends/CMakeLists.txt index 2497aa48340..6a44f3234c5 100644 --- a/backends/qualcomm/runtime/backends/CMakeLists.txt +++ b/backends/qualcomm/runtime/backends/CMakeLists.txt @@ -68,11 +68,12 @@ target_sources( PUBLIC ${CMAKE_CURRENT_LIST_DIR}/QnnContextCommon.h ${CMAKE_CURRENT_LIST_DIR}/htpbackend/HtpContext.h ${CMAKE_CURRENT_LIST_DIR}/irbackend/IrContext.h - PRIVATE ${CMAKE_CURRENT_LIST_DIR}/QnnContextCommon.cpp - ${CMAKE_CURRENT_LIST_DIR}/htpbackend/HtpContext.cpp - ${CMAKE_CURRENT_LIST_DIR}/htpbackend/HtpContextCustomConfig.h - ${HOST_ARCHITECTURE}/HtpContextCustomConfig.cpp - ${CMAKE_CURRENT_LIST_DIR}/irbackend/${CMAKE_SYSTEM_PROCESSOR}/IrContext.cpp + PRIVATE + ${CMAKE_CURRENT_LIST_DIR}/QnnContextCommon.cpp + ${CMAKE_CURRENT_LIST_DIR}/htpbackend/HtpContext.cpp + ${CMAKE_CURRENT_LIST_DIR}/htpbackend/HtpContextCustomConfig.h + ${HOST_ARCHITECTURE}/HtpContextCustomConfig.cpp + ${CMAKE_CURRENT_LIST_DIR}/irbackend/${CMAKE_SYSTEM_PROCESSOR}/IrContext.cpp ) # qnn_backend_cache @@ -137,5 +138,6 @@ target_sources( target_sources( qnn_dlc_manager PUBLIC ${CMAKE_CURRENT_LIST_DIR}/QnnDlcManager.h - PRIVATE ${CMAKE_CURRENT_LIST_DIR}/irbackend/${CMAKE_SYSTEM_PROCESSOR}/QnnDlcManager.cpp + PRIVATE + ${CMAKE_CURRENT_LIST_DIR}/irbackend/${CMAKE_SYSTEM_PROCESSOR}/QnnDlcManager.cpp ) diff --git a/backends/qualcomm/runtime/backends/QnnBackendCache.cpp b/backends/qualcomm/runtime/backends/QnnBackendCache.cpp index 4387d61ab7c..3dd1738d33b 100644 --- a/backends/qualcomm/runtime/backends/QnnBackendCache.cpp +++ b/backends/qualcomm/runtime/backends/QnnBackendCache.cpp @@ -113,7 +113,6 @@ Error QnnBackendCache::Configure(const std::vector& graph_names) { // DO DESERIALIZE state_ = DESERIALIZE; QNN_EXECUTORCH_LOG_INFO("Caching: Caching is in RESTORE MODE."); - auto [status, _, context_size, context_ptr] = QnnContextCustomProtocol().DeserializeContextCustomBuffer( qnn_context_blob_.buffer); diff --git a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp index 2fbb2243d8d..e7e9db6fed8 100644 --- a/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp +++ b/backends/qualcomm/runtime/backends/QnnBackendFactory.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ #include +#include #include #include namespace executorch { @@ -30,7 +31,8 @@ std::unique_ptr QnnBackendFactory::Create( if (!skel_library_dir.empty()) { setenv("ADSP_LIBRARY_PATH", skel_library_dir.c_str(), /*overwrite=*/1); } - if (options->log_level() >= QnnExecuTorchLogLevel::kLogLevelInfo) { + if (get_option(options->log_level()) >= + QnnExecuTorchLogLevel::kLogLevelInfo) { QNN_EXECUTORCH_LOG_INFO( "skel_library_dir: %s", skel_library_dir.c_str()); QNN_EXECUTORCH_LOG_INFO( @@ -42,7 +44,7 @@ std::unique_ptr QnnBackendFactory::Create( QNN_EXECUTORCH_LOG_INFO( "performance_mode in htp_options: %s", EnumNameQnnExecuTorchHtpPerformanceMode( - htp_options->performance_mode())); + get_option(htp_options->performance_mode()))); QNN_EXECUTORCH_LOG_INFO( "precision in htp_options: %s", EnumNameQnnExecuTorchHtpPrecision(htp_options->precision())); @@ -75,13 +77,13 @@ std::unique_ptr QnnBackendFactory::Create( implementation, backend_params->qnn_backend_ptr_.get(), backend_params->qnn_context_ptr_.get(), - options->profile_level(), + get_option(options->profile_level()), options->soc_info(), htp_options); backend_params->qnn_mem_manager_ptr_ = std::make_unique( implementation, backend_params->qnn_context_ptr_.get(), - options->log_level()); + get_option(options->log_level())); backend_params->backend_init_state_ = BackendInitializeState::INITIALIZED; } break; case QnnExecuTorchBackendType::kGpuBackend: diff --git a/backends/qualcomm/runtime/backends/QnnCustomProtocol.cpp b/backends/qualcomm/runtime/backends/QnnCustomProtocol.cpp index 12de1b3e705..b01d7ab6d80 100644 --- a/backends/qualcomm/runtime/backends/QnnCustomProtocol.cpp +++ b/backends/qualcomm/runtime/backends/QnnCustomProtocol.cpp @@ -12,87 +12,6 @@ namespace executorch { namespace backends { namespace qnn { -// we still need this for on-device op validation of other backends -void QnnQcirCustomProtocol::BuildQcirCustomBuffer( - const QnnExecuTorchContextBinary& qcir_binary, - const std::vector& tensor_data) { - if (qnn_custom_buffer_.size() == 0) { - uint8_t magic_number_proto_size = sizeof(magic_number_); - uint8_t qcir_fbs_proto_size = sizeof(qcir_fbs_size_); - uint8_t tensor_proto_size = sizeof(tensor_size_); - - uint64_t buffer_size = magic_number_proto_size + qcir_fbs_proto_size + - tensor_proto_size + qcir_fbs_size_ + tensor_size_; - qnn_custom_buffer_.resize(buffer_size, 0); - - size_t pos = 0; - // magic number itself - std::memcpy( - qnn_custom_buffer_.data(), &magic_number_, magic_number_proto_size); - pos += magic_number_proto_size; - - // size of qcir_fbs, should be 4 bytes - std::memcpy( - qnn_custom_buffer_.data() + pos, &qcir_fbs_size_, qcir_fbs_proto_size); - pos += qcir_fbs_proto_size; - - // size of tensor, should be 8 bytes - std::memcpy( - qnn_custom_buffer_.data() + pos, &tensor_size_, tensor_proto_size); - pos += tensor_proto_size; - - // qcir.fbs buffer - uint8_t* qcir_ptr = static_cast(qcir_binary.buffer); - - std::memcpy(qnn_custom_buffer_.data() + pos, qcir_ptr, qcir_fbs_size_); - pos += qcir_fbs_size_; - - // tensor data - std::memcpy( - qnn_custom_buffer_.data() + pos, tensor_data.data(), tensor_size_); - } -} - -std::tuple -QnnQcirCustomProtocol::DeserializeQcirCustomBuffer(void* processed_data) { - Error status = Error::Ok; - uint8_t* ptr = static_cast(processed_data); - size_t magic_number_proto_size = sizeof(magic_number_); - uint8_t qcir_fbs_proto_size = sizeof(qcir_fbs_size_); - uint8_t tensor_proto_size = sizeof(tensor_size_); - - uint32_t magic_number; - std::memcpy(&magic_number, ptr, magic_number_proto_size); - ptr += magic_number_proto_size; - - if (magic_number != magic_number_) { - QNN_EXECUTORCH_LOG_INFO( - "QnnQcirCustomProtocol expected magic number: 0x%x but get: 0x%x", - magic_number_, - magic_number); - status = Error::Internal; - } - - // Retrieve size of qcir.fbs - uint32_t qcir_fbs_size; - std::memcpy(&qcir_fbs_size, ptr, qcir_fbs_proto_size); - ptr += qcir_fbs_proto_size; - - // Retrieve size of tensor - uint64_t tensor_size; - std::memcpy(&tensor_size, ptr, tensor_proto_size); - ptr += tensor_proto_size; - - // Retrieve qcir.fbs pointer - void* qcir_fbs_ptr = static_cast(ptr); - ptr += qcir_fbs_size; - - // Retrieve tensor - void* tensor_ptr = static_cast(ptr); - - return {status, qcir_fbs_size, tensor_size, qcir_fbs_ptr, tensor_ptr}; -} - void QnnContextCustomProtocol::BuildContextCustomBuffer() { if (qnn_custom_buffer_.size() == 0) { signature_ = diff --git a/backends/qualcomm/runtime/backends/QnnCustomProtocol.h b/backends/qualcomm/runtime/backends/QnnCustomProtocol.h index 6ea556899f5..3cc6a6e25dc 100644 --- a/backends/qualcomm/runtime/backends/QnnCustomProtocol.h +++ b/backends/qualcomm/runtime/backends/QnnCustomProtocol.h @@ -24,13 +24,8 @@ namespace qnn { using executorch::runtime::Error; -// We have 2 kinds of protocol here: custom_qcir_protocol, -// custom_context_protocol. We need this class due to limitation of 32bits -// flatbuffer. Since larger models can exceed the maximum size for 32bits -// flatbuffer, we need to define our own protocol and store some information -// outside of the flatbuffer. The magic number helps determine if we are getting -// the correct custom protocol buffer and differentiate custom_qcir_protocol -// from custom_context_protocol. +// Required for multi-graph support to retrieve qnn manager handle via unique +// signature. class QnnCustomProtocol { public: QnnCustomProtocol() {} @@ -47,48 +42,6 @@ class QnnCustomProtocol { std::vector qnn_custom_buffer_; }; -// For custom_qcir_protocol, we expect the following format: -// -// ------------------------------ -// | qcir magic number (4 bytes)| -// ------------------------------ -// | qcir.fbs size (4 bytes) | -// ------------------------------ -// | tensor size (8 bytes) | -// ------------------------------ -// | qcir.fbs (flatbuffer) | -// ------------------------------ -// | tensor.data | -// ------------------------------ -class QnnQcirCustomProtocol : public QnnCustomProtocol { - public: - // Constructor for Serialize - QnnQcirCustomProtocol(uint32_t qcir_fbs_size, uint64_t tensor_size) - : QnnCustomProtocol(), - qcir_fbs_size_(qcir_fbs_size), - tensor_size_(tensor_size) {} - - // Constructor for Deserialize - QnnQcirCustomProtocol() : QnnCustomProtocol() {} - - void BuildQcirCustomBuffer( - const QnnExecuTorchContextBinary& qcir_binary, - const std::vector& tensor_data); - // Return a tuple with 5 elements: - // 1) Error: Status of whether deserializing is successful. - // 2) uint32_t: Size of qcir fbs - // 3) uint64_t: Size of tensor - // 4) void*: Pointer pointing to the start of qcir fbs - // 5) void*: Pointer pointing to the start of tensor - std::tuple - DeserializeQcirCustomBuffer(void* processed_data); - - private: - static constexpr uint32_t magic_number_ = 0x1234ABCD; - uint32_t qcir_fbs_size_{0}; - uint64_t tensor_size_{0}; -}; - // For custom context binary protocol, we expect the following format: // // --------------------------------- diff --git a/backends/qualcomm/runtime/backends/QnnImplementation.cpp b/backends/qualcomm/runtime/backends/QnnImplementation.cpp index a9136a83c9c..7083f2bef30 100644 --- a/backends/qualcomm/runtime/backends/QnnImplementation.cpp +++ b/backends/qualcomm/runtime/backends/QnnImplementation.cpp @@ -6,7 +6,6 @@ * LICENSE file in the root directory of this source tree. */ #include - #include "QnnInterface.h" namespace executorch { namespace backends { @@ -51,16 +50,12 @@ Error QnnImplementation::StartBackend( const std::string& lib_path, const QnnSaver_Config_t** saver_config) { Qnn_ErrorHandle_t error = QNN_SUCCESS; - // RTLD_GLOBAL is needed on x86 as HTP op package has a requirement for the - // symbols in backend to be visible. Using RTLD_LOCAL on Android to allow full - // unloading of HTP backend shared library on dlclose() as RTLD_GLOBAL isn't - // letting it happen. void* lib_handle = nullptr; -#if defined(__ANDROID__) - lib_handle = dlopen(lib_path.c_str(), RTLD_NOW | RTLD_LOCAL); -#else - lib_handle = dlopen(lib_path.c_str(), RTLD_NOW | RTLD_GLOBAL); -#endif + // If the library is already loaded, return the handle. + lib_handle = dlopen(lib_path.c_str(), RTLD_NOW | RTLD_NOLOAD); + if (!lib_handle) { + lib_handle = dlopen(lib_path.c_str(), RTLD_NOW | RTLD_GLOBAL); + } if (lib_handle == nullptr) { QNN_EXECUTORCH_LOG_ERROR( "Cannot Open QNN library %s, with error: %s", diff --git a/backends/qualcomm/runtime/backends/QnnMemManager.cpp b/backends/qualcomm/runtime/backends/QnnMemManager.cpp index e09d071075b..3b99dd10868 100644 --- a/backends/qualcomm/runtime/backends/QnnMemManager.cpp +++ b/backends/qualcomm/runtime/backends/QnnMemManager.cpp @@ -56,13 +56,16 @@ Error QnnMemManager::RegisterIonMem( return Error::Ok; } +// TODO: Find a better way to unify RegisterCustomMem and +// PreRegisterCustomMemHandle Error QnnMemManager::RegisterCustomMem( const std::shared_ptr& tensor_wrapper, int32_t mem_fd, void* mem_ptr, void* unaligned_custom_mem_base, size_t total_custom_mem_size, - size_t tensor_offset) { + size_t tensor_offset, + const CustomMemTensorInfo& info) { const QnnInterface& qnn_interface = implementation_.GetQnnInterface(); Qnn_MemDescriptor_t descriptor = { {tensor_wrapper->GetRank(), tensor_wrapper->GetDims(), nullptr}, @@ -94,6 +97,7 @@ Error QnnMemManager::RegisterCustomMem( return Error::Internal; } tensor_wrapper->SetMemHandle(handle); + pre_registered_handles_.insert({info, handle}); registered_map_.insert({handle, mem_ptr}); if (log_level_ >= QnnExecuTorchLogLevel::kLogLevelInfo) { QNN_EXECUTORCH_LOG_INFO( diff --git a/backends/qualcomm/runtime/backends/QnnMemManager.h b/backends/qualcomm/runtime/backends/QnnMemManager.h index 30bb64d78ad..6a7f00b016a 100644 --- a/backends/qualcomm/runtime/backends/QnnMemManager.h +++ b/backends/qualcomm/runtime/backends/QnnMemManager.h @@ -41,7 +41,8 @@ class QnnMemManager { void* mem_ptr, void* unaligned_custom_mem_base, size_t total_custom_mem_size, - size_t tensor_offset); + size_t tensor_offset, + const CustomMemTensorInfo& info); // Pre-register custom mem handle from SharedBuffer. Bring forward the // memHandle creating time from execution to initialization. @@ -67,7 +68,9 @@ class QnnMemManager { const QnnImplementation& implementation_; QnnContext* context_; QnnExecuTorchLogLevel log_level_; + // Store the registered Qnn_MemHandle_t for de-registration std::unordered_map registered_map_; + // Store the pre-registered custom mem handles std::unordered_map pre_registered_handles_; std::unordered_map scalar_type_to_qnn_dtype_ = { diff --git a/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.cpp b/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.cpp index 46ba3117269..35a20048fc5 100644 --- a/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.cpp +++ b/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.cpp @@ -396,11 +396,10 @@ Error HtpDevice::AfterCreateDevice() { QNN_GET_ERROR_CODE(error)); return Error::Internal; } - // Set vector of PowerConfigs and map it to a vector of pointers. perf_power_configs_ = SetVotePowerConfig( powerconfig_client_id_, - htp_options_->performance_mode(), + get_option(htp_options_->performance_mode()), PerformanceModeVoteType::kUpVote); perf_power_configs_ptr_ = ObtainNullTermPtrVector(perf_power_configs_); @@ -416,7 +415,7 @@ Error HtpDevice::AfterCreateDevice() { // Set Rpc polling mode rpc_power_configs_ = - SetRpcPollingPowerConfig(htp_options_->performance_mode()); + SetRpcPollingPowerConfig(get_option(htp_options_->performance_mode())); rpc_power_configs_ptr_ = ObtainNullTermPtrVector(rpc_power_configs_); htp_perf_infra_->setPowerConfig( diff --git a/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.h b/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.h index f75e15fc77c..9052deb6b52 100644 --- a/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.h +++ b/backends/qualcomm/runtime/backends/htpbackend/HtpDevice.h @@ -7,6 +7,7 @@ */ #pragma once +#include #include #include #include @@ -55,7 +56,7 @@ class HtpDevice : public QnnDevice { void ReleasePerformanceVote(); inline bool IsPerfModeEnabled() { - return htp_options_->performance_mode() != + return get_option(htp_options_->performance_mode()) != QnnExecuTorchHtpPerformanceMode::kHtpDefault; } diff --git a/backends/qualcomm/runtime/backends/irbackend/x86_64/QnnDlcManager.cpp b/backends/qualcomm/runtime/backends/irbackend/x86_64/QnnDlcManager.cpp index 050a679e62a..280751cf160 100644 --- a/backends/qualcomm/runtime/backends/irbackend/x86_64/QnnDlcManager.cpp +++ b/backends/qualcomm/runtime/backends/irbackend/x86_64/QnnDlcManager.cpp @@ -5,6 +5,7 @@ * 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 @@ -51,7 +52,7 @@ Error QnnDlcManager::Create() { qnn_loaded_backend_, backend_params_ptr_->qnn_backend_ptr_.get(), backend_params_ptr_->qnn_context_ptr_.get(), - options_->profile_level()); + get_option(options_->profile_level())); backend_params_ptr_->backend_init_state_ = BackendInitializeState::INITIALIZED; return backend_params_ptr_->qnn_backend_ptr_->VerifyQNNSDKVersion(); @@ -105,7 +106,7 @@ Error QnnDlcManager::SetUpDlcEnvironment(const Qnn_Version_t& coreApiVersion) { "Fail to Load Qnn IR library."); logger_ = std::make_unique( - qnn_loaded_backend_, LoggingCallback, options_->log_level()); + qnn_loaded_backend_, LoggingCallback, get_option(options_->log_level())); ET_CHECK_OR_RETURN_ERROR( Create() == Error::Ok, Internal, "Failed to load Qnn IR backend."); diff --git a/backends/qualcomm/runtime/targets.bzl b/backends/qualcomm/runtime/targets.bzl index 1bd82f8f913..db3706ba221 100644 --- a/backends/qualcomm/runtime/targets.bzl +++ b/backends/qualcomm/runtime/targets.bzl @@ -73,13 +73,12 @@ def define_common_targets(): "fbsource//third-party/qualcomm/qnn/qnn-{0}:app_sources".format(get_qnn_library_version()), ":logging", "//executorch/backends/qualcomm:schema", - "//executorch/backends/qualcomm/aot/ir:qcir_utils", "//executorch/backends/qualcomm/aot/wrappers:wrappers", - "//executorch/runtime/backend:interface", "//executorch/runtime/core:core", "//executorch/extension/tensor:tensor", ], exported_deps = [ + "//executorch/runtime/backend:interface", "//executorch/runtime/core/exec_aten/util:scalar_type_util", "//executorch/runtime/core:event_tracer", ], diff --git a/backends/qualcomm/scripts/build.sh b/backends/qualcomm/scripts/build.sh index 43d968813a9..c84911cf851 100755 --- a/backends/qualcomm/scripts/build.sh +++ b/backends/qualcomm/scripts/build.sh @@ -1,3 +1,4 @@ +#!/usr/bin/env bash # Copyright (c) Qualcomm Innovation Center, Inc. # All rights reserved # @@ -116,6 +117,21 @@ if [ "$BUILD_AARCH64" = true ]; then -B$EXAMPLE_ROOT cmake --build $EXAMPLE_ROOT -j$BUILD_JOB_NUMBER + + LLAMA_EXAMPLE_ROOT=examples/models/llama + cmake $PRJ_ROOT/$LLAMA_EXAMPLE_ROOT \ + -DBUILD_TESTING=OFF \ + -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK_ROOT/build/cmake/android.toolchain.cmake \ + -DCMAKE_BUILD_TYPE=$BUILD_TYPE \ + -DANDROID_ABI='arm64-v8a' \ + -DANDROID_PLATFORM=android-30 \ + -DCMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ + -DPYTHON_EXECUTABLE=$PYTHON_EXECUTABLE \ + -B$LLAMA_EXAMPLE_ROOT + + cmake --build $LLAMA_EXAMPLE_ROOT -j$BUILD_JOB_NUMBER fi if [ "$BUILD_X86_64" = true ]; then @@ -172,4 +188,18 @@ if [ "$BUILD_X86_64" = true ]; then -B$EXAMPLE_ROOT cmake --build $EXAMPLE_ROOT -j$BUILD_JOB_NUMBER + + LLAMA_EXAMPLE_ROOT=examples/models/llama + cmake $PRJ_ROOT/$LLAMA_EXAMPLE_ROOT \ + -DBUILD_TESTING=OFF \ + -DCMAKE_BUILD_TYPE=$BUILD_TYPE \ + -DANDROID_ABI='arm64-v8a' \ + -DANDROID_PLATFORM=android-30 \ + -DCMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ + -DPYTHON_EXECUTABLE=$PYTHON_EXECUTABLE \ + -B$LLAMA_EXAMPLE_ROOT + + cmake --build $LLAMA_EXAMPLE_ROOT -j$BUILD_JOB_NUMBER fi diff --git a/backends/qualcomm/scripts/download_qnn_sdk.py b/backends/qualcomm/scripts/download_qnn_sdk.py new file mode 100644 index 00000000000..35006a41433 --- /dev/null +++ b/backends/qualcomm/scripts/download_qnn_sdk.py @@ -0,0 +1,444 @@ +# Add these imports for additional logging +import ctypes +import logging +import os +import pathlib +import platform +import re +import shutil +import tarfile +import tempfile +import urllib.request +import zipfile +from typing import Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) +logger.addHandler(logging.NullHandler()) + +PKG_ROOT = pathlib.Path(__file__).parent.parent +SDK_DIR = PKG_ROOT / "sdk" / "qnn" + + +def is_linux_x86() -> bool: + """ + Check if the current platform is Linux x86_64. + + Returns: + bool: True if the system is Linux x86_64, False otherwise. + """ + return platform.system().lower() == "linux" and platform.machine().lower() in ( + "x86_64", + "amd64", + "i386", + "i686", + ) + + +import subprocess + +MINIMUM_LIBC_VERSION = 2.29 + +REQUIRED_LIBC_LIBS = [ + "/lib/x86_64-linux-gnu/libc.so.6", + "/lib64/libc.so.6", + "/lib/libc.so.6", +] + + +def check_glibc_exist_and_validate() -> bool: + """ + Check if users have glibc installed. + """ + exists = False + for path in REQUIRED_LIBC_LIBS: + try: + output = subprocess.check_output( + [path, "--version"], stderr=subprocess.STDOUT + ) + output = output.decode().split("\n")[0] + logger.debug(f"[QNN] glibc version for path {path} is: {output}") + match = re.search(r"version (\d+\.\d+)", output) + if match: + version = match.group(1) + if float(version) >= MINIMUM_LIBC_VERSION: + logger.debug(f"[QNN] glibc version is {version}.") + exists = True + return True + else: + logger.error( + f"[QNN] glibc version is too low. The minimum libc version is {MINIMUM_LIBC_VERSION} Please install glibc following the commands below." + ) + else: + logger.error("[QNN] glibc version not found.") + + except Exception: + continue + + if not exists: + logger.error( + r"""" + [QNN] glibc not found or the version is too low. Please install glibc following the commands below. + Ubuntu/Debian: + sudo apt update + sudo apt install libc6 + + Fedora/Red Hat: + sudo dnf install glibc + + Arch Linux: + sudo pacman -S glibc + + Also please make sure the glibc version is >= MINIMUM_LIBC_VERSION. You can verify the glibc version by running the following command: + Option 1: + ldd --version + Option 2: + /path/to/libc.so.6 --version + """ + ) + return exists + + +def _download_archive(url: str, archive_path: pathlib.Path) -> bool: + """Download archive from URL with progress reporting.""" + logger.debug("Archive will be saved to: %s", archive_path) + + try: + urllib.request.urlretrieve(url, archive_path, _make_report_progress()) + logger.info("Download completed!") + except Exception as e: + logger.exception("Error during download: %s", e) + return False + + if archive_path.exists() and archive_path.stat().st_size == 0: + logger.warning("Downloaded file is empty!") + return False + elif not archive_path.exists(): + logger.error("File was not downloaded!") + return False + return True + + +def _make_report_progress(): + """Return a callback to report download progress.""" + last_reported = 0 + + def report_progress(block_num, block_size, total_size): + nonlocal last_reported + try: + downloaded = block_num * block_size + percent = downloaded / total_size * 100 if total_size else 100.0 + except Exception: + percent, downloaded, total_size = 0.0, block_num * block_size, 0 + if percent - last_reported >= 20 or percent >= 100: + logger.info( + "Downloaded: %d/%d bytes (%.2f%%)", downloaded, total_size, percent + ) + last_reported = percent + + return report_progress + + +def _extract_archive( + url: str, archive_path: pathlib.Path, content_dir: str, dst_folder: pathlib.Path +): + """Extract archive based on type (zip or tar).""" + if url.endswith(".zip"): + logger.info("Extracting ZIP archive...") + _extract_zip(archive_path, content_dir, dst_folder) + elif url.endswith((".tar.gz", ".tgz")): + logger.info("Extracting TAR archive...") + _extract_tar(archive_path, content_dir, dst_folder) + else: + raise ValueError(f"Unsupported archive format: {url}") + + +def _verify_extraction(dst_folder: pathlib.Path): + """Check if extraction succeeded and log contents.""" + logger.info("Verifying extraction to %s", dst_folder) + if dst_folder.exists(): + logger.debug("SDK directory exists. Contents:") + for item in dst_folder.iterdir(): + logger.debug(" %s", item.name) + else: + logger.error("SDK directory was not created!") + + +def _download_qnn_sdk(dst_folder=SDK_DIR) -> Optional[pathlib.Path]: + """ + Download and extract the Qualcomm SDK into dst_folder. + Only runs on Linux x86 platforms. + """ + QNN_VERSION = "2.37.0.250724" + logger.info("Downloading Qualcomm SDK...") + QAIRT_URL = ( + f"https://softwarecenter.qualcomm.com/api/download/software/sdks/" + f"Qualcomm_AI_Runtime_Community/All/{QNN_VERSION}/v{QNN_VERSION}.zip" + ) + QAIRT_CONTENT_DIR = f"qairt/{QNN_VERSION}" + if not is_linux_x86(): + logger.info("[QNN] Skipping Qualcomm SDK (only supported on Linux x86).") + return None + elif not check_glibc_exist_and_validate(): + logger.info("[QNN] Skipping Qualcomm SDK (glibc not found or version too old).") + return None + else: + logger.info("[QNN] Downloading Qualcomm SDK for Linux x86") + + dst_folder.mkdir(parents=True, exist_ok=True) + + with tempfile.TemporaryDirectory() as tmpdir: + archive_path = pathlib.Path(tmpdir) / pathlib.Path(QAIRT_URL).name + if not _download_archive(QAIRT_URL, archive_path): + return None + + _extract_archive(QAIRT_URL, archive_path, QAIRT_CONTENT_DIR, dst_folder) + _verify_extraction(dst_folder) + + return dst_folder + + +def _extract_zip(archive_path, content_dir, target_dir): + logger.debug("Extracting %s to %s", archive_path, target_dir) + logger.debug("Looking for content in subdirectory: %s", content_dir) + + target_dir.mkdir(parents=True, exist_ok=True) + + with zipfile.ZipFile(archive_path, "r") as zip_ref: + files_to_extract = [f for f in zip_ref.namelist() if f.startswith(content_dir)] + + for file in files_to_extract: + relative_path = pathlib.Path(file).relative_to(content_dir) + if relative_path == pathlib.Path("."): + continue + + out_path = target_dir / relative_path + if file.endswith("/"): + out_path.mkdir(parents=True, exist_ok=True) + else: + out_path.parent.mkdir(parents=True, exist_ok=True) + with zip_ref.open(file) as src, open(out_path, "wb") as dst: + shutil.copyfileobj(src, dst) + + +def _extract_tar(archive_path: pathlib.Path, prefix: str, target_dir: pathlib.Path): + with tarfile.open(archive_path, "r:gz") as tf: + for m in tf.getmembers(): + if not m.name.startswith(prefix + "/"): + continue + relpath = pathlib.Path(m.name).relative_to(prefix) + if not relpath.parts or relpath.parts[0] == "..": + continue + + out_path = target_dir / relpath + if m.isdir(): + out_path.mkdir(parents=True, exist_ok=True) + else: + out_path.parent.mkdir(parents=True, exist_ok=True) + src = tf.extractfile(m) + if src is None: + continue + with src, open(out_path, "wb") as dst: + dst.write(src.read()) + + +LLVM_VERSION = "14.0.0" +LIBCXX_BASE_NAME = f"clang+llvm-{LLVM_VERSION}-x86_64-linux-gnu-ubuntu-18.04" +LLVM_URL = f"https://github.com/llvm/llvm-project/releases/download/llvmorg-{LLVM_VERSION}/{LIBCXX_BASE_NAME}.tar.xz" +REQUIRED_LIBCXX_LIBS = [ + "libc++.so.1.0", + "libc++abi.so.1.0", + "libunwind.so.1", +] + + +def _stage_libcxx(target_dir: pathlib.Path): + target_dir.mkdir(parents=True, exist_ok=True) + + if all((target_dir / libname).exists() for libname in REQUIRED_LIBCXX_LIBS): + logger.info("[libcxx] Already staged at %s, skipping download", target_dir) + return + + temp_tar = pathlib.Path("/tmp") / f"{LIBCXX_BASE_NAME}.tar.xz" + temp_extract = pathlib.Path("/tmp") / LIBCXX_BASE_NAME + + if not temp_tar.exists(): + logger.info("[libcxx] Downloading %s", LLVM_URL) + urllib.request.urlretrieve(LLVM_URL, temp_tar) + + logger.info("[libcxx] Extracting %s", temp_tar) + with tarfile.open(temp_tar, "r:xz") as tar: + tar.extractall(temp_extract.parent) + + lib_src = temp_extract / "lib" / "x86_64-unknown-linux-gnu" + for fname in REQUIRED_LIBCXX_LIBS: + src_path = lib_src / fname + if not src_path.exists(): + logger.warning( + "[libcxx] %s not found in extracted LLVM src_path %s", fname, src_path + ) + continue + shutil.copy(src_path, target_dir / fname) + + logger.info("[libcxx] Staged libc++ to %s", target_dir) + + +REQUIRED_QNN_LIBS: List[str] = [ + "libQnnHtp.so", +] + + +def _ld_library_paths() -> List[pathlib.Path]: + """Split LD_LIBRARY_PATH into ordered directories (skip empties).""" + raw = os.environ.get("LD_LIBRARY_PATH", "") + return [pathlib.Path(p) for p in raw.split(":") if p.strip()] + + +def _find_lib_in_ld_paths( + libname: str, ld_dirs: Optional[List[pathlib.Path]] = None +) -> Optional[pathlib.Path]: + """Return first matching path to `libname` in LD_LIBRARY_PATH, or None.""" + if ld_dirs is None: + ld_dirs = _ld_library_paths() + for d in ld_dirs: + candidate = d / libname + try: + if candidate.exists(): + return candidate.resolve() + except Exception: + # Ignore unreadable / permission issues, keep looking. + pass + return None + + +def _check_libs_in_ld( + libnames: List[str], +) -> Tuple[bool, Dict[str, Optional[pathlib.Path]]]: + """ + Check if each lib in `libnames` exists in LD_LIBRARY_PATH directories. + + Returns: + all_present: True iff every lib was found + locations: mapping lib -> path (or None if missing) + """ + ld_dirs = _ld_library_paths() + locations: Dict[str, Optional[pathlib.Path]] = {} + for lib in libnames: + locations[lib] = _find_lib_in_ld_paths(lib, ld_dirs) + all_present = all(locations[lib] is not None for lib in libnames) + return all_present, locations + + +# ----------------------- +# Ensure QNN SDK library +# ----------------------- +def _ensure_qnn_sdk_lib() -> bool: + """ + Ensure libQnnHtp.so is available. + - If found in LD_LIBRARY_PATH: do nothing, return True. + - Otherwise: ensure packaged SDK is present, then load libQnnHtp.so from it. + """ + all_present, locs = _check_libs_in_ld(REQUIRED_QNN_LIBS) + if all_present: + logger.info( + "[QNN] libQnnHtp.so found in LD_LIBRARY_PATH; skipping SDK install." + ) + for lib, p in locs.items(): + logger.info(" - %s: %s", lib, p) + return True + + # Not found → use packaged SDK + qnn_sdk_dir = SDK_DIR + logger.info("[QNN] libQnnHtp.so not found in LD_LIBRARY_PATH.") + if not qnn_sdk_dir.exists(): + logger.info("[QNN] SDK dir missing; downloading...") + _download_qnn_sdk() + else: + logger.info("[QNN] Using existing SDK at %s", qnn_sdk_dir) + + os.environ["QNN_SDK_ROOT"] = str(qnn_sdk_dir) + + qnn_lib = qnn_sdk_dir / "lib" / "x86_64-linux-clang" / "libQnnHtp.so" + logger.info("[QNN] Loading %s", qnn_lib) + lib_loaded = False + try: + ctypes.CDLL(str(qnn_lib), mode=ctypes.RTLD_GLOBAL) + logger.info("[QNN] Loaded libQnnHtp.so from packaged SDK.") + lib_loaded = True + except OSError as e: + logger.error("[QNN][ERROR] Failed to load %s: %s", qnn_lib, e) + return lib_loaded + + +def _load_libcxx_libs(lib_path): + logger.debug("running _load_libcxx_libs") + candidates = list(lib_path.glob("*.so*")) + priority = ["libc++abi", "libc++"] + sorted_candidates = [ + f for name in priority for f in candidates if f.name.startswith(name) + ] + sorted_candidates += [f for f in candidates if f not in sorted_candidates] + logger.debug("sorted_candidates: %s", sorted_candidates) + for sofile in sorted_candidates: + try: + ctypes.CDLL(str(sofile), mode=ctypes.RTLD_GLOBAL) + logger.info("Loaded %s", sofile.name) + except OSError as e: + logger.warning("[WARN] Failed to load %s: %s", sofile.name, e) + + +# --------------------- +# Ensure libc++ family +# --------------------- +def _ensure_libcxx_stack() -> bool: + """ + Ensure libc++ stack is available. + - If all required libc++ libs are found in LD_LIBRARY_PATH: do nothing. + - Otherwise: stage and load the packaged libc++ bundle. + """ + all_present, locs = _check_libs_in_ld(REQUIRED_LIBCXX_LIBS) + if all_present: + logger.info( + "[libcxx] All libc++ libs present in LD_LIBRARY_PATH; skipping staging." + ) + for lib, p in locs.items(): + logger.info(" - %s: %s", lib, p) + return True + + logger.info( + "[libcxx] Some libc++ libs missing in LD_LIBRARY_PATH; staging packaged libc++..." + ) + lib_loaded = False + try: + libcxx_dir = PKG_ROOT / "sdk" / f"libcxx-{LLVM_VERSION}" + _stage_libcxx(libcxx_dir) + _load_libcxx_libs(libcxx_dir) + logger.info("[libcxx] Staged and loaded libc++ from %s", libcxx_dir) + lib_loaded = True + except Exception as e: + logger.exception("[libcxx][ERROR] Failed to stage/load libc++: %s", e) + return lib_loaded + + +# --------------- +# Public entrypoint +# --------------- +def install_qnn_sdk() -> bool: + """ + Initialize Qualcomm backend with separated logic: + + QNN SDK: + - If libQnnHtp.so exists in LD_LIBRARY_PATH: do nothing. + - Else: ensure packaged SDK, load libQnnHtp.so. + + libc++ stack: + - If required libc++ libs exist in LD_LIBRARY_PATH: do nothing. + - Else: stage and load packaged libc++. + + Returns: + True if both steps succeeded (or were already satisfied), else False. + """ + if check_glibc_exist_and_validate(): + if _ensure_libcxx_stack(): + if _ensure_qnn_sdk_lib(): + return True + return False diff --git a/backends/qualcomm/scripts/install_qnn_sdk.sh b/backends/qualcomm/scripts/install_qnn_sdk.sh index 913ce34b711..a97d4258770 100644 --- a/backends/qualcomm/scripts/install_qnn_sdk.sh +++ b/backends/qualcomm/scripts/install_qnn_sdk.sh @@ -9,7 +9,7 @@ source "${SCRIPT_DIR}/qnn_config.sh" # Function to install Android NDK (only if not already set) setup_android_ndk() { # Check if ANDROID_NDK_ROOT is already set and valid - if [ -n "${ANDROID_NDK_ROOT}" ] && [ -d "${ANDROID_NDK_ROOT}" ]; then + if [ -n "${ANDROID_NDK_ROOT:-}" ] && [ -d "${ANDROID_NDK_ROOT:-}" ]; then echo "Android NDK already set to ${ANDROID_NDK_ROOT} - skipping installation" return fi @@ -41,13 +41,19 @@ verify_pkg_installed() { install_qnn() { # Check if QNN_SDK_ROOT is already set and valid - if [ -n "${QNN_SDK_ROOT}" ] && [ -d "${QNN_SDK_ROOT}" ]; then + if [ -n "${QNN_SDK_ROOT:-}" ] && [ -d "${QNN_SDK_ROOT:-}" ]; then echo "QNN SDK already set to ${QNN_SDK_ROOT} - skipping installation" return fi echo "Start installing qnn v${QNN_VERSION}" QNN_INSTALLATION_DIR="/tmp/qnn" + + if [ -d "${QNN_INSTALLATION_DIR}/${QNN_VERSION}" ]; then + echo "QNN SDK already installed at ${QNN_INSTALLATION_DIR}/${QNN_VERSION}" + export QNN_SDK_ROOT="${QNN_INSTALLATION_DIR}/${QNN_VERSION}" + return + fi # Clean up any previous installation if [ -d "${QNN_INSTALLATION_DIR}" ]; then @@ -79,7 +85,7 @@ install_qnn() { ls -lah "${QNN_INSTALLATION_DIR}" # Set QNN_SDK_ROOT environment variable - export QNN_SDK_ROOT="${QNN_INSTALLATION_DIR}" + export QNN_SDK_ROOT="${QNN_INSTALLATION_DIR}/${QNN_VERSION}" echo "Set QNN_SDK_ROOT=${QNN_SDK_ROOT}" } @@ -141,13 +147,9 @@ setup_libcpp() { popd >/dev/null # Set environment variables - export CPLUS_INCLUDE_PATH="${INSTALL_DIR}/include:$CPLUS_INCLUDE_PATH" - export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:$LD_LIBRARY_PATH" - export LIBRARY_PATH="${INSTALL_DIR}/lib:$LIBRARY_PATH" + export CPLUS_INCLUDE_PATH="${INSTALL_DIR}/include:${CPLUS_INCLUDE_PATH:-}" + export LD_LIBRARY_PATH="${INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" + export LIBRARY_PATH="${INSTALL_DIR}/lib:${LIBRARY_PATH:-}" echo "libc++ installed to ${INSTALL_DIR}" } - -setup_libcpp 12 -setup_android_ndk -install_qnn diff --git a/backends/qualcomm/scripts/qnn_config.sh b/backends/qualcomm/scripts/qnn_config.sh index fe2d82e939e..938eb0d3007 100644 --- a/backends/qualcomm/scripts/qnn_config.sh +++ b/backends/qualcomm/scripts/qnn_config.sh @@ -6,5 +6,5 @@ # LICENSE file in the root directory of this source tree. # QNN SDK Configuration -QNN_VERSION="2.28.0.241029" -QNN_ZIP_URL="https://softwarecenter.qualcomm.com/api/download/software/qualcomm_neural_processing_sdk/v${QNN_VERSION}.zip" +QNN_VERSION="2.37.0.250724" +QNN_ZIP_URL="https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/${QNN_VERSION}/v${QNN_VERSION}.zip" diff --git a/backends/qualcomm/serialization/qc_schema_serialize.py b/backends/qualcomm/serialization/qc_schema_serialize.py index 59610d7b996..328695869d6 100644 --- a/backends/qualcomm/serialization/qc_schema_serialize.py +++ b/backends/qualcomm/serialization/qc_schema_serialize.py @@ -1,14 +1,16 @@ # Copyright (c) Qualcomm Innovation Center, Inc. +# Copyright 2025 Arm Limited and/or its 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 importlib.resources as _resources import json import os import tempfile -import pkg_resources +import executorch.backends.qualcomm.serialization as serialization_package from executorch.backends.qualcomm.serialization.qc_schema import QnnExecuTorchOptions from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile @@ -19,7 +21,9 @@ def _convert_to_flatbuffer(obj, schema: str): with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, f"{schema}.fbs") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, f"{schema}.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, f"{schema}.fbs") + ) json_path = os.path.join(d, f"{schema}.json") with open(json_path, "wb") as json_file: json_file.write(obj_json.encode("ascii")) @@ -36,7 +40,9 @@ def _convert_to_object(flatbuffers: bytes, obj_type, schema: str): schema_path = os.path.join(d, f"{schema}.fbs") bin_path = os.path.join(d, f"{schema}.bin") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, f"{schema}.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, f"{schema}.fbs") + ) with open(bin_path, "wb") as bin_file: bin_file.write(flatbuffers) diff --git a/backends/qualcomm/tests/TARGETS b/backends/qualcomm/tests/TARGETS index 8078ca611f8..639303c7eb8 100644 --- a/backends/qualcomm/tests/TARGETS +++ b/backends/qualcomm/tests/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") -python_library( +runtime.python_library( name = "models", srcs = ["models.py"], deps = [ @@ -10,7 +10,7 @@ python_library( ] ) -python_library( +runtime.python_library( name = "test_qnn_delegate", srcs = [ "test_qnn_delegate.py", @@ -37,3 +37,13 @@ python_library( "//executorch/backends/qualcomm/debugger:utils", ], ) + +runtime.python_library( + name = "tester", + srcs = [ + "tester.py", + ], + deps = [ + ":test_qnn_delegate" + ] +) diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 01ed37f80a3..0e290575beb 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -102,6 +102,24 @@ def forward(self, x): return torch.amax(x, dim=self.dim, keepdim=self.keepdim) +class AMaxFollowingConv2D(torch.nn.Module): + def __init__( + self, in_channels, out_channels, kernel_size=3, dim=None, keepdim=False + ): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels, out_channels, kernel_size, padding=kernel_size // 2 + ) + self.dim = dim + self.keepdim = keepdim + + def forward(self, x): + x = self.conv( + x + ) # Apply convolution (output shape: [batch, out_channels, H, W]) + return torch.amax(x, dim=self.dim, keepdim=self.keepdim) + + class AMin(torch.nn.Module): def __init__(self, dim=None, keepdim=False): super().__init__() @@ -165,6 +183,14 @@ def forward(self, x, y): return squeeze_out, conv_out +class Asin(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.asin(x) + + class Atan(torch.nn.Module): def __init__(self): super().__init__() @@ -248,6 +274,19 @@ def forward(self, x, y): return torch.cat((y, y, x, x), axis=2) +class CausalMask(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("causal_mask", torch.zeros((1, 1, 1, 128))) + self.mask_length = 128 + + def forward(self, padding_mask): + self.causal_mask[:, :, :, : self.mask_length] = self.causal_mask[ + :, :, :, : self.mask_length + ].masked_fill(padding_mask, 1) + return self.causal_mask + 1 + + class CDist(torch.nn.Module): def __init__(self): super().__init__() @@ -550,10 +589,16 @@ def forward(self, x): class ConvTranspose1dSingle(torch.nn.Module): - def __init__(self, bias=True): + def __init__(self, bias=True, dilation=1): super().__init__() self.conv_transpose = torch.nn.ConvTranspose1d( - in_channels=1, out_channels=3, kernel_size=3, stride=2, padding=1, bias=bias + in_channels=1, + out_channels=3, + kernel_size=3, + stride=2, + padding=1, + dilation=dilation, + bias=bias, ) def forward(self, x): @@ -561,7 +606,7 @@ def forward(self, x): class ConvTranspose2dSingle(torch.nn.Module): - def __init__(self, bias=True): + def __init__(self, bias=True, dilation=1): super().__init__() self.conv_transpose = torch.nn.ConvTranspose2d( in_channels=1, @@ -569,6 +614,7 @@ def __init__(self, bias=True): kernel_size=3, stride=2, padding=1, + dilation=dilation, bias=bias, ) @@ -600,6 +646,40 @@ def forward(self, x): return self.conv_transpose(self.conv(x)) +class Conv2dFlip(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels=16, + out_channels=16, + kernel_size=3, + stride=2, + padding=1, + bias=False, + ) + self.dims = [1, 3] + + def forward(self, x): + x = self.conv(x) + return torch.flip(x, self.dims) + + +class Conv2dSliceCopy(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels=1, + out_channels=4, + kernel_size=(3, 3), + padding=1, + bias=True, + ) + + def forward(self, x): + x = self.conv(x) + return x[:, 2:, :, :] + + class Conv2dSumReduceDim(torch.nn.Module): def __init__(self): super().__init__() @@ -768,6 +848,15 @@ def forward(self, x): return torch.special.expm1(x) +class Flip(torch.nn.Module): + def __init__(self): + super().__init__() + self.dims = [0, 2] + + def forward(self, x): + return torch.flip(x, self.dims) + + class Floor(torch.nn.Module): def __init__(self): super().__init__() @@ -776,6 +865,23 @@ def forward(self, x): return torch.floor(x) +class FloorDiv(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, y): + return torch.floor_divide(x, y) + + +class FloorDivConstantFloat(torch.nn.Module): + def __init__(self, constant=2.0): + super().__init__() + self.constant = constant + + def forward(self, x): + return torch.floor(x / self.constant) + + class Fold(torch.nn.Module): def __init__(self): super().__init__() @@ -976,6 +1082,15 @@ def forward(self, input_pos, k_val): return k_out + 0 +class IndexSelect(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x, indices): + return torch.index_select(x, self.dim, indices) + + class InstanceNorm2d(torch.nn.Module): def __init__(self, n_features, affine=True): super().__init__() @@ -1081,7 +1196,7 @@ def forward(self, x): class Linear(torch.nn.Module): def __init__(self, use_bias: bool = True): super().__init__() - self.linear = torch.nn.Linear(4, 5, use_bias).eval() + self.linear = torch.nn.Linear(512, 32, use_bias).eval() def forward(self, x): return self.linear(x) @@ -1380,6 +1495,15 @@ def forward(self, x): return self.relu(x) +class Relu6(torch.nn.Module): + def __init__(self): + super().__init__() + self.relu6 = torch.nn.ReLU6() + + def forward(self, x): + return self.relu6(x) + + class Repeat(torch.nn.Module): def __init__(self): super().__init__() @@ -1564,6 +1688,14 @@ def forward(self, x): return torch.sigmoid(x) +class Sign(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.sign(x) + + class Sin(torch.nn.Module): def __init__(self): super().__init__() @@ -1642,6 +1774,20 @@ def forward(self, x, y): ) +class SliceScatter(torch.nn.Module): + def __init__(self, dim, start, end, step): + super().__init__() + self.dim = dim + self.start = start + self.end = end + self.step = step + + def forward(self, x, y): + return x.slice_scatter( + y, dim=self.dim, start=self.start, end=self.end, step=self.step + ) + + class Softmax(torch.nn.Module): def __init__(self, dim): super().__init__() @@ -1853,6 +1999,28 @@ def forward(self, x): ) +class XorBitWise(torch.nn.Module): + def __init__(self, pos, neg): + super().__init__() + self.pos = pos + self.neg = neg + + def forward(self, x, y): + bitwise_xor = torch.bitwise_xor(x, y).bool() + return torch.where(bitwise_xor, self.pos, self.neg) + + +class XorOperator(torch.nn.Module): + def __init__(self, pos, neg): + super().__init__() + self.pos = pos + self.neg = neg + + def forward(self, x, y): + operator_xor = x.to(torch.bool) ^ y.to(torch.bool) + return torch.where(operator_xor, self.pos, self.neg) + + # Mimi Decoder has 0D tensor which QNN cannot handle. class ZeroDimTensor(torch.nn.Module): def __init__(self): diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index d4eb3e4eac3..6ef4fa8fe13 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -9,6 +9,7 @@ import sys import tempfile import unittest +from functools import partial from multiprocessing.connection import Listener from pathlib import Path @@ -127,14 +128,29 @@ def test_qnn_backend_alias(self): self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_amax(self): - modules = [AMax(dim=1, keepdim=False), AMax(dim=1, keepdim=True)] # noqa: F405 + modules = [ + AMax(dim=1, keepdim=False), # noqa: F405 + AMax(dim=1, keepdim=True), # noqa: F405 + AMax(), # noqa: F405 + ] sample_input = (torch.randn(4, 4),) for i, module in enumerate(modules): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_amax_conv(self): + sample_input = (torch.randn(2, 3, 64, 64),) # [batch, channels, height, width] + module = AMaxFollowingConv2D( # noqa: F405 + in_channels=3, out_channels=16, kernel_size=3, dim=-1, keepdim=False + ) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_amin(self): - modules = [AMin(dim=1, keepdim=False), AMin(dim=1, keepdim=True)] # noqa: F405 + modules = [ + AMin(dim=1, keepdim=False), # noqa: F405 + AMin(dim=1, keepdim=True), # noqa: F405 + AMin(), # noqa: F405 + ] sample_input = (torch.randn(4, 4),) for i, module in enumerate(modules): with self.subTest(i=i): @@ -163,7 +179,13 @@ def test_qnn_backend_argmax(self): def test_qnn_backend_argmin(self): module = Argmin() # noqa: F405 - sample_input = (torch.randn(16, 3, 4, 4),) + sample_input = (torch.rand(3, 4),) + self.lower_module_and_test_output(module, sample_input) + + @unittest.expectedFailure + def test_qnn_backend_asin(self): + sample_input = (torch.rand(3, 4) * 2 - 1,) + module = Asin() # noqa: F405 self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_atan(self): @@ -264,8 +286,9 @@ def test_qnn_backend_conv_transpose1d(self): modules = [ ConvTranspose1dSingle(), # noqa: F405 ConvTranspose1dSingle(bias=False), # noqa: F405 + ConvTranspose1dSingle(dilation=2), # noqa: F405 ] - sample_input = (torch.randn([1, 1, 3]),) + sample_input = (torch.randn([1, 1, 33]),) for i, module in enumerate(modules): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) @@ -274,8 +297,11 @@ def test_qnn_backend_conv_transpose2d(self): modules = [ ConvTranspose2dSingle(), # noqa: F405 ConvTranspose2dSingle(bias=False), # noqa: F405 + ConvTranspose2dSingle(dilation=2), # noqa: F405 + ConvTranspose2dSingle(dilation=(2, 3)), # noqa: F405 + ConvTranspose2dSingle(dilation=(2, 1)), # noqa: F405 ] - sample_input = (torch.randn([1, 1, 3, 3]),) + sample_input = (torch.randn([1, 1, 33, 33]),) for i, module in enumerate(modules): with self.subTest(i=i): self.lower_module_and_test_output(module, sample_input) @@ -529,11 +555,40 @@ def test_qnn_backend_expm1(self): module = ExpM1() # noqa: F405 self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_flip(self): + sample_input = (torch.randn(3, 4, 5, 6),) + module = Flip() # noqa: F405 + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_floor(self): sample_input = (torch.randn(3, 4),) module = Floor() # noqa: F405 self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_floor_divide(self): + eps = 1e-03 + test_comb = [ + { + QCOM_MODULE: [FloorDiv()], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + (torch.randn(2, 5, 1, 3), eps + torch.randn(2, 5, 1, 3)), + (torch.randn([2, 5, 1, 3]), eps + torch.randn([4, 1])), + ], + }, + { + QCOM_MODULE: [FloorDivConstantFloat()], # noqa: F405 + QCOM_SAMPLE_INPUTS: [(torch.randn(2, 5, 1, 3),)], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + self.lower_module_and_test_output(module, sample_input) + index += 1 + def test_qnn_backend_fold(self): sample_input = (torch.randn(3, 512, 256),) module = Fold() # noqa: F405 @@ -555,16 +610,17 @@ def test_qnn_backend_gather(self): Gather(), # noqa: F405 # TODO: resolve accuracy problem # GatherArgmin(), # noqa: F405 - GatherWhere(), # noqa: F405 + # TODO: There is a accuracy regression after 2.37 + # GatherWhere(), # noqa: F405 ] - shape = (2, 2, 3, 4) + # shape = (2, 2, 3, 4) sample_inputs = [ ( torch.arange(128, dtype=torch.float32).view(64, 2), torch.ones(64, 2, dtype=torch.int64), ), # (torch.arange(128, dtype=torch.float32).view(64, 2),), - (torch.randn(shape), torch.randn(shape)), + # (torch.randn(shape), torch.randn(shape)), ] for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)): with self.subTest(i=i): @@ -728,6 +784,14 @@ def test_qnn_backend_index_put(self): skip_mutable_buffer=test[QCOM_MODULE].skip_mutable_buffer, ) + def test_qnn_backend_index_select(self): + module = IndexSelect(dim=1) # noqa: F405 + sample_input = ( + torch.randn(2, 3, 4, 5), + torch.tensor([0, 2]), + ) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_instance_norm_2d(self): modules = [InstanceNorm2d(32), InstanceNorm2d(32, affine=False)] # noqa: F405 sample_input = (torch.randn([4, 32, 16, 16]),) @@ -816,6 +880,11 @@ def test_qnn_backend_less_equal(self): test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS] ) + def test_qnn_backend_sign(self): + module = Sign() # noqa: F405 + sample_input = (torch.randn(3, 4),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_less_than(self): test_comb = [ { @@ -847,7 +916,7 @@ def test_qnn_backend_linalg_vector_norm(self): def test_qnn_backend_linear(self): module = Linear() # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_log(self): @@ -970,6 +1039,11 @@ def test_qnn_backend_relu(self): sample_input = (torch.randn([2, 5, 1, 3]),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_relu6(self): + module = Relu6() # noqa: F405 + sample_input = (torch.randn([2, 5, 1, 3]),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_repeat(self): module = Repeat() # noqa: F405 sample_input = (torch.randn([2, 2, 2, 2]),) @@ -1054,6 +1128,42 @@ def test_qnn_backend_slice_copy(self): for module, sample_input in zip(modules, sample_inputs): self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_slice_scatter(self): + test_comb = [ + { + QCOM_MODULE: [ + SliceScatter(dim=0, start=3, end=5, step=1) # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(8, 8), + torch.ones(2, 8), + ) + ], + }, + { + QCOM_MODULE: [ + SliceScatter(dim=1, start=2, end=6, step=2) # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [ + ( + ( + torch.zeros(8, 8), + torch.ones(8, 2), + ) + ) + ], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + self.lower_module_and_test_output(module, sample_input) + index += 1 + def test_qnn_backend_stack(self): module = Stack() # noqa: F405 sample_input = ( @@ -1115,18 +1225,47 @@ def test_qnn_backend_where(self): Where(), # noqa: F405 WhereConstant(torch.randn(3, 2), torch.randn(3, 2)), # noqa: F405 WhereConstantOther(), # noqa: F405 - WhereConstantAll(), # noqa: F405 + # TODO: There is a accuracy regression after 2.37 + # WhereConstantAll(), # noqa: F405 WhereConstantInf(), # noqa: F405 ] sample_inputs = [ (torch.randn(3, 2), torch.randn(3, 2), torch.randn(3, 2)), (torch.randn(3, 2),), (torch.randn(3, 2),), - (torch.randn(3, 2),), + # (torch.randn(3, 2),), (torch.randn(30, 20),), ] for i, module in enumerate(modules): - self.lower_module_and_test_output(module, sample_inputs[i]) + with self.subTest(i=i): + self.lower_module_and_test_output(module, sample_inputs[i]) + + def test_qnn_backend_element_wise_xor(self): + test_comb = [ + { + QCOM_MODULE: XorBitWise( # noqa: F405 + torch.tensor(1.7), torch.tensor(0.2) + ), + QCOM_SAMPLE_INPUTS: ( + torch.tensor([1, 0, 1, 0], dtype=torch.bool), + torch.tensor([1, 1, 0, 0], dtype=torch.bool), + ), + }, + { + QCOM_MODULE: XorOperator( # noqa: F405 + torch.tensor(1.5), torch.tensor(-1.2) + ), + QCOM_SAMPLE_INPUTS: ( + torch.full((3, 3), 1).triu(), + torch.full((3, 3), 1).tril(diagonal=0), + ), + }, + ] + for i, test in enumerate(test_comb): + with self.subTest(i=i): + self.lower_module_and_test_output( + test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS] + ) def test_qnn_backend_masked_fill(self): module = MaskedFill() # noqa: F405 @@ -1163,6 +1302,11 @@ def test_qnn_backend_argmin_view_squeeze_conv2d(self): sample_input = (torch.randn(32), torch.randn(32, 3, 32, 32)) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_causal_mask(self): + module = CausalMask() # noqa: F405 + sample_input = (torch.rand((1, 1, 1, 128)) < 0.5,) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_chunk_add(self): module = ChunkAdd() # noqa: F405 torch.manual_seed(8) @@ -1204,11 +1348,21 @@ def test_qnn_backend_conv2d_down_up_sample(self): sample_input = (torch.randn(1, 16, 224, 224),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_flip(self): + module = Conv2dFlip() # noqa: F405 + sample_input = (torch.randn(1, 16, 224, 224),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_max_pool2d(self): module = Conv2dMaxPool2d() # noqa: F405 sample_input = (torch.rand(1, 2, 14, 14),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_slice_copy(self): + module = Conv2dSliceCopy() # noqa: F405 + sample_input = (torch.randn([2, 1, 3, 3]),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_sum_reduce_dim(self): module = Conv2dSumReduceDim() # noqa: F405 sample_input = (torch.randn([1, 1, 3, 3]),) @@ -1372,7 +1526,7 @@ def test_qnn_backend_16a4w_layer_norm(self): def test_qnn_backend_16a4w_linear(self): module = Linear() # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) module = self.get_qdq_module( module, sample_input, @@ -1382,7 +1536,7 @@ def test_qnn_backend_16a4w_linear(self): def test_qnn_backend_16a4w_per_channel_linear(self): module = Linear(use_bias=False) # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) module = self.get_qdq_module( module, sample_input, @@ -1393,7 +1547,7 @@ def test_qnn_backend_16a4w_per_channel_linear(self): def test_qnn_backend_16a4w_per_channel_linear_with_bias(self): module = Linear() # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) module = self.get_qdq_module( module, sample_input, @@ -1427,15 +1581,31 @@ def test_qnn_backend_alias(self): self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_amax(self): - modules = [AMax(dim=1, keepdim=False), AMax(dim=1, keepdim=True)] # noqa: F405 + modules = [ + AMax(dim=1, keepdim=False), # noqa: F405 + AMax(dim=1, keepdim=True), # noqa: F405 + AMax(), # noqa: F405 + ] sample_input = (torch.randn(4, 4),) for i, module in enumerate(modules): with self.subTest(i=i): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_amax_conv(self): + sample_input = (torch.randn(2, 3, 64, 64),) # [batch, channels, height, width] + module = AMaxFollowingConv2D( # noqa: F405 + in_channels=3, out_channels=16, kernel_size=3, dim=-1, keepdim=False + ) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_amin(self): - modules = [AMin(dim=1, keepdim=False), AMin(dim=1, keepdim=True)] # noqa: F405 + modules = [ + AMin(dim=1, keepdim=False), # noqa: F405 + AMin(dim=1, keepdim=True), # noqa: F405 + AMin(), # noqa: F405 + ] sample_input = (torch.randn(4, 4),) for i, module in enumerate(modules): with self.subTest(i=i): @@ -1472,6 +1642,12 @@ def test_qnn_backend_argmin(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_asin(self): + module = Asin() # noqa: F405 + sample_input = (torch.rand([3, 4]) * 2 - 1,) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_atan(self): sample_input = (torch.randn(3, 4),) module = Atan() # noqa: F405 @@ -1610,6 +1786,7 @@ def test_qnn_backend_conv_transpose1d(self): modules = [ ConvTranspose1dSingle(), # noqa: F405 ConvTranspose1dSingle(bias=False), # noqa: F405 + ConvTranspose1dSingle(dilation=2), # noqa: F405 ] sample_input = (torch.randn([1, 1, 3]),) for i, module in enumerate(modules): @@ -1621,7 +1798,9 @@ def test_qnn_backend_conv_transpose2d(self): modules = [ ConvTranspose2dSingle(), # noqa: F405 ConvTranspose2dSingle(bias=False), # noqa: F405 - ] # noqa: F405 + ConvTranspose2dSingle(dilation=(2, 3)), # noqa: F405 + ConvTranspose2dSingle(dilation=(2, 1)), # noqa: F405 + ] sample_input = (torch.randn([1, 1, 3, 3]),) for i, module in enumerate(modules): with self.subTest(i=i): @@ -1876,12 +2055,43 @@ def test_qnn_backend_expm1(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_flip(self): + sample_input = (torch.randn(3, 4, 5, 6),) + module = Flip() # noqa: F405 + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_floor(self): sample_input = (torch.randn(3, 4),) module = Floor() # noqa: F405 module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_floor_divide(self): + eps = 1e-03 + test_comb = [ + { + QCOM_MODULE: [FloorDiv()], # noqa: F405 + QCOM_SAMPLE_INPUTS: [ + (torch.randn(2, 5, 1, 3), eps + torch.randn(2, 5, 1, 3)), + (torch.randn([2, 5, 1, 3]), eps + torch.randn([4, 1])), + ], + }, + { + QCOM_MODULE: [FloorDivConstantFloat()], # noqa: F405 + QCOM_SAMPLE_INPUTS: [(torch.randn(2, 5, 1, 3),)], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + gm = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(gm, sample_input) + index += 1 + def test_qnn_backend_fold(self): sample_input = (torch.randn(3, 512, 256),) module = Fold() # noqa: F405 @@ -1906,16 +2116,17 @@ def test_qnn_backend_gather(self): Gather(), # noqa: F405 # TODO: resolve accuracy problem # GatherArgmin(), # noqa: F405 - GatherWhere(), # noqa: F405 + # TODO: There is a accuracy regression after 2.37 + # GatherWhere(), # noqa: F405 ] - shape = (2, 2, 3, 4) + # shape = (2, 2, 3, 4) sample_inputs = [ ( torch.arange(128, dtype=torch.float32).view(64, 2), torch.ones(64, 2, dtype=torch.int64), ), # (torch.arange(128, dtype=torch.float32).view(64, 2),), - (torch.randn(shape), torch.randn(shape)), + # (torch.randn(shape), torch.randn(shape)), ] for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)): with self.subTest(i=i): @@ -2094,6 +2305,15 @@ def test_qnn_backend_index_put(self): skip_mutable_buffer=test[QCOM_MODULE].skip_mutable_buffer, ) + def test_qnn_backend_index_select(self): + module = IndexSelect(dim=1) # noqa: F405 + sample_input = ( + torch.randn(2, 3, 4, 5), + torch.tensor([0, 2]), + ) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_instance_norm_2d(self): modules = [InstanceNorm2d(32), InstanceNorm2d(32, affine=False)] # noqa: F405 sample_input = (torch.randn([4, 32, 16, 16]),) @@ -2223,16 +2443,37 @@ def test_qnn_backend_linalg_vector_norm(self): def test_qnn_backend_linear(self): module = Linear() # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + @unittest.skipIf(is_qnn_sdk_version_less_than("2.30"), "UT pass after QNN 2.30") + def test_qnn_backend_linear_block(self): + modules = [ + Linear(use_bias=False), # noqa: F405 + Linear(use_bias=True), # noqa: F405 + ] + + sample_input = (torch.randn([3, 512]),) + for i, module in enumerate(modules): + with self.subTest(i=i): + # update block size for linear weight (OI) + # channel dimension(O) is defaultly sliced in QNN + # divide dimension(I) into 16 groups + module = self.get_qdq_module( + module, + sample_input, + quant_dtype=QuantDtype.use_16a4w_block, + block_size_map={"linear": (1, 32)}, + ) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_linear_qat(self): """ Prototype to test qat model """ module = Linear() # noqa: F405 - sample_input = (torch.randn([3, 4]),) + sample_input = (torch.randn([3, 512]),) prepared = self.get_prepared_qat_module(module, sample_input) module = self.get_converted_sgd_trained_module(module, prepared, sample_input) self.lower_module_and_test_output(module, sample_input) @@ -2374,6 +2615,12 @@ def test_qnn_backend_relu(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_relu6(self): + module = Relu6() # noqa: F405 + sample_input = (torch.randn([2, 5, 1, 3]),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_repeat(self): module = Repeat() # noqa: F405 sample_input = (torch.randn([2, 2, 2, 2]),) @@ -2452,6 +2699,12 @@ def test_qnn_backend_sigmoid(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_sign(self): + module = Sign() # noqa: F405 + sample_input = (torch.randn(3, 4),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_sin(self): module = Sin() # noqa: F405 sample_input = (torch.randn(2, 5, 1, 3),) @@ -2473,6 +2726,43 @@ def test_qnn_backend_slice_copy(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_slice_scatter(self): + test_comb = [ + { + QCOM_MODULE: [ + SliceScatter(dim=0, start=3, end=5, step=1) # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [ + ( + torch.zeros(8, 8), + torch.ones(2, 8), + ) + ], + }, + { + QCOM_MODULE: [ + SliceScatter(dim=1, start=2, end=6, step=2) # noqa: F405 + ], + QCOM_SAMPLE_INPUTS: [ + ( + ( + torch.zeros(8, 8), + torch.ones(8, 2), + ) + ) + ], + }, + ] + + index = 0 + for comb in test_comb: + for module in comb[QCOM_MODULE]: + for sample_input in comb[QCOM_SAMPLE_INPUTS]: + with self.subTest(i=index): + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + index += 1 + def test_qnn_backend_softmax(self): modules = [Softmax(dim=1), Softmax(dim=-1)] # noqa: F405 sample_input = (torch.randn([1, 4, 8, 8]),) @@ -2544,19 +2834,21 @@ def test_qnn_backend_where(self): Where(), # noqa: F405 WhereConstant(torch.randn(3, 2), torch.randn(3, 2)), # noqa: F405 WhereConstantOther(), # noqa: F405 - WhereConstantAll(), # noqa: F405 + # TODO: There is a accuracy regression after 2.37 + # WhereConstantAll(), # noqa: F405 WhereConstantInf(), # noqa: F405 ] sample_inputs = [ (torch.randn(3, 2), torch.randn(3, 2), torch.randn(3, 2)), (torch.randn(3, 2),), (torch.randn(3, 2),), - (torch.randn(3, 2),), + # (torch.randn(3, 2),), (torch.randn(30, 20),), ] for i, module in enumerate(modules): - module = self.get_qdq_module(module, sample_inputs[i]) - self.lower_module_and_test_output(module, sample_inputs[i]) + with self.subTest(i=i): + module = self.get_qdq_module(module, sample_inputs[i]) + self.lower_module_and_test_output(module, sample_inputs[i]) def test_qnn_backend_masked_fill(self): module = MaskedFill() # noqa: F405 @@ -2571,6 +2863,34 @@ def test_qnn_backend_masked_fill(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_element_wise_xor(self): + test_comb = [ + { + QCOM_MODULE: XorBitWise( # noqa: F405 + torch.tensor(1.7), torch.tensor(0.2) + ), + QCOM_SAMPLE_INPUTS: ( + torch.tensor([1, 0, 1, 0], dtype=torch.bool), + torch.tensor([1, 1, 0, 0], dtype=torch.bool), + ), + }, + { + QCOM_MODULE: XorOperator( # noqa: F405 + torch.tensor(1.5), torch.tensor(-1.2) + ), + QCOM_SAMPLE_INPUTS: ( + torch.full((3, 3), 1).triu(), + torch.full((3, 3), 1).tril(diagonal=0), + ), + }, + ] + for i, test in enumerate(test_comb): + with self.subTest(i=i): + module = self.get_qdq_module( + test[QCOM_MODULE], test[QCOM_SAMPLE_INPUTS] + ) + self.lower_module_and_test_output(module, test[QCOM_SAMPLE_INPUTS]) + class TestQNNQuantizedModel(TestQNN): # TODO: refactor to support different backends @@ -2595,6 +2915,12 @@ def test_qnn_backend_argmin_view_squeeze_conv2d(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_causal_mask(self): + module = CausalMask() # noqa: F405 + sample_input = (torch.rand((1, 1, 1, 128)) < 0.5,) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_chunk_add(self): module = ChunkAdd() # noqa: F405 torch.manual_seed(8) @@ -2643,12 +2969,24 @@ def test_qnn_backend_conv2d_down_up_sample(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_flip(self): + module = Conv2dFlip() # noqa: F405 + sample_input = (torch.randn(1, 16, 224, 224),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_max_pool2d(self): module = Conv2dMaxPool2d() # noqa: F405 sample_input = (torch.rand(1, 2, 14, 14),) module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_slice_copy(self): + module = Conv2dSliceCopy() # noqa: F405 + sample_input = (torch.randn([2, 1, 3, 3]),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_conv2d_sum_reduce_dim(self): module = Conv2dSumReduceDim() # noqa: F405 sample_input = (torch.randn([1, 1, 3, 3]),) @@ -2707,7 +3045,7 @@ def test_qnn_backend_masked_softmax(self): edge_prog_mgr.write_to_file(f) adb = self.get_adb_tool(pte_path) binaries_trace = generate_optrace( - tmp_dir, self.chipset_table[self.model], adb, pte_path, sample_input + tmp_dir, self.chipset_table[self.model], adb, pte_path, [sample_input] ) has_masked_softmax = False for _, (_, qhas) in binaries_trace.items(): @@ -3054,27 +3392,125 @@ def test_qnn_backend_profile_op(self): expected_profile_events=30, ) - def test_qnn_backend_shared_buffer(self): - TestQNN.shared_buffer = True - backend_options = generate_htp_compiler_spec( - use_fp16=True, - ) + def test_qnn_backend_runtime_option_htp_performance(self): + backend_options = generate_htp_compiler_spec(use_fp16=True) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=self.chipset_table[TestQNN.model], backend_options=backend_options, - shared_buffer=True, ) module = SimpleModel() # noqa: F405 sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + + def output_callback(log_msg, is_burst): + msg = log_msg.stdout + # Refer to HtpDevice.cpp for the following values + min_voltage = ( + "coreVoltageCornerMin 160" if is_burst else "coreVoltageCornerMin 80" + ) + self.assertTrue(min_voltage in msg, f"Expecting '{min_voltage} ' in log") + + burst_runtime_commands = ( + " --htp_performance_mode 2 --log_level 4" # kHtpBurst, kLogLevelVerbose + ) self.lower_module_and_test_output( module, sample_input, - expected_partitions=1, + extra_cmds=burst_runtime_commands, + output_callback=partial(output_callback, is_burst=True), + save_inference_speed=True, ) + burst_speed = 1000 / self.inference_speed # inferences per second - def test_qnn_backend_online_prepare(self): - if self.enable_x86_64: - self.skipTest("TODO: add online_prepare support on host platform") + power_saver_runtime_commands = " --htp_performance_mode 6 --log_level 4" # kHtpHighPowerSaver, kLogLevelVerbose + self.lower_module_and_test_output( + module, + sample_input, + extra_cmds=power_saver_runtime_commands, + output_callback=partial(output_callback, is_burst=False), + save_inference_speed=True, + ) + power_saver_speed = 1000 / self.inference_speed # inferences per second + + # Only need to ensure device burst is faster than high power saver + if not self.enable_x86_64: + self.assertGreater( + burst_speed, + power_saver_speed, + f"Burst mode should be faster than high power saver mode, Burst: {burst_speed} inference / second, High Power Saver: {power_saver_speed} inference /second.", + ) + + def test_qnn_backend_runtime_option_log(self): + backend_options = generate_htp_compiler_spec(use_fp16=True) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + runtime_commands = " --log_level 4" # kLogLevelVerbose + + def output_callback(log_msg): + msg = log_msg.stdout + # Check log prefix, different QNN version will have slightly different message format. + self.assertTrue( + any( + sub in msg + for sub in [ + "[Qnn ExecuTorch]: QnnDsp ", + "[Qnn ExecuTorch]: ", + ] + ), + "Expecting Verbose message in log", + ) + + self.lower_module_and_test_output( + module, + sample_input, + extra_cmds=runtime_commands, + output_callback=output_callback, + ) + + def test_qnn_backend_runtime_option_profile(self): + TestQNN.enable_profile = True + backend_options = generate_htp_compiler_spec(use_fp16=True) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + profile=False, # Turn on using runtime command + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + runtime_commands = " --profile_level 2" # kProfileDetailed + # With same model, expected_profile events for this UT should match test_qnn_backend_profile_op + self.lower_module_and_test_output( + module, + sample_input, + expected_partitions=1, + expected_profile_events=30, + extra_cmds=runtime_commands, + ) + + def test_qnn_backend_shared_buffer(self): + TestQNN.shared_buffer = True + backend_options = generate_htp_compiler_spec( + use_fp16=True, + ) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + shared_buffer=True, + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + self.lower_module_and_test_output( + module, + sample_input, + expected_partitions=1, + ) + + def test_qnn_backend_online_prepare(self): + if self.enable_x86_64: + self.skipTest("TODO: add online_prepare support on host platform") backend_options = generate_htp_compiler_spec(use_fp16=True) TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( @@ -3319,7 +3755,6 @@ def test_qnn_backend_generate_optrace(self): for compiler_spec in compiler_specs: with tempfile.TemporaryDirectory() as tmp_dir: - edge_prog_mgr = to_edge_transform_and_lower_to_qnn( module, sample_input, compiler_spec ).to_executorch() @@ -3329,7 +3764,11 @@ def test_qnn_backend_generate_optrace(self): adb = self.get_adb_tool(pte_path) binaries_trace = generate_optrace( - tmp_dir, self.chipset_table[self.model], adb, pte_path, sample_input + tmp_dir, + self.chipset_table[self.model], + adb, + pte_path, + [sample_input], ) for _, (optrace, qhas) in binaries_trace.items(): with open(optrace, "r") as optrace_file: @@ -3774,6 +4213,107 @@ def test_qnn_backend_profile_op(self): expected_profile_events=30, ) + def test_qnn_backend_runtime_option_htp_performance(self): + backend_options = generate_htp_compiler_spec(use_fp16=False) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + module = self.get_qdq_module(module, sample_input) + + def output_callback(log_msg, is_burst): + msg = log_msg.stdout + # Refer to HtpDevice.cpp for the following values + min_voltage = ( + "coreVoltageCornerMin 160" if is_burst else "coreVoltageCornerMin 80" + ) + self.assertTrue(min_voltage in msg, f"Expecting '{min_voltage} ' in log") + + burst_runtime_commands = ( + " --htp_performance_mode 2 --log_level 4" # kHtpBurst, kLogLevelVerbose + ) + self.lower_module_and_test_output( + module, + sample_input, + extra_cmds=burst_runtime_commands, + output_callback=partial(output_callback, is_burst=True), + save_inference_speed=True, + ) + burst_speed = 1000 / self.inference_speed # num inference per second + + power_saver_runtime_commands = " --htp_performance_mode 6 --log_level 4" # kHtpHighPowerSaver, kLogLevelVerbose + self.lower_module_and_test_output( + module, + sample_input, + extra_cmds=power_saver_runtime_commands, + output_callback=partial(output_callback, is_burst=False), + save_inference_speed=True, + ) + power_saver_speed = 1000 / self.inference_speed # num inference per second + + # Only need to ensure device burst is faster than high power saver + if not self.enable_x86_64: + self.assertGreater( + burst_speed, + power_saver_speed, + f"Burst mode should be faster than high power saver mode, Burst: {burst_speed} inference / second, High Power Saver: {power_saver_speed} inference /second.", + ) + + def test_qnn_backend_runtime_option_log(self): + backend_options = generate_htp_compiler_spec(use_fp16=False) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + module = self.get_qdq_module(module, sample_input) + runtime_commands = " --log_level 4" # kLogLevelVerbose + + def output_callback(log_msg): + msg = log_msg.stdout + # Check log prefix, different QNN version will have slightly different message format. + self.assertTrue( + any( + sub in msg + for sub in [ + "[Qnn ExecuTorch]: QnnDsp ", + "[Qnn ExecuTorch]: ", + ] + ), + "Expecting Verbose message in log", + ) + + self.lower_module_and_test_output( + module, + sample_input, + extra_cmds=runtime_commands, + output_callback=output_callback, + ) + + def test_qnn_backend_runtime_option_profile(self): + TestQNN.enable_profile = True + backend_options = generate_htp_compiler_spec(use_fp16=False) + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.model], + backend_options=backend_options, + profile=False, # Turn on using runtime command + ) + module = SimpleModel() # noqa: F405 + sample_input = (torch.ones(1, 32, 28, 28), torch.ones(1, 32, 28, 28)) + module = self.get_qdq_module(module, sample_input) + runtime_commands = " --profile_level 2" # kProfileDetailed + # With same model, expected_profile events for this UT should match test_qnn_backend_profile_op + self.lower_module_and_test_output( + module, + sample_input, + expected_partitions=1, + expected_profile_events=30, + extra_cmds=runtime_commands, + ) + def test_qnn_backend_shared_buffer(self): TestQNN.shared_buffer = True backend_options = generate_htp_compiler_spec( @@ -4071,7 +4611,11 @@ def test_qnn_backend_generate_optrace(self): adb = self.get_adb_tool(pte_path) binaries_trace = generate_optrace( - tmp_dir, self.chipset_table[self.model], adb, pte_path, sample_input + tmp_dir, + self.chipset_table[self.model], + adb, + pte_path, + [sample_input], ) for _, (optrace, qhas) in binaries_trace.items(): with open(optrace, "r") as optrace_file: @@ -4092,9 +4636,107 @@ def test_qnn_backend_generate_optrace(self): qhas_data = json.load(qhas_file) self.assertIn("data", qhas_data) + def test_qnn_backend_seq_mse(self): + from executorch.backends.qualcomm._passes.seq_mse import SeqMSE + + o_ch, i_ch, kernel, padding = 32, 512, (1, 1), 0 + module = Conv2dSingle( # noqa: F405 + in_channel=i_ch, + out_channel=o_ch, + kernel_size=kernel, + padding=padding, + ) + sample_input = (torch.randn(1, i_ch, 1, o_ch),) + # per-channel / per-block + quantizers = [ + make_quantizer(), + make_quantizer(quant_dtype=QuantDtype.use_16a4w_block), + ] + quantizers[-1].set_block_size_map({"conv2d": (1, 32, 1, 1)}) + + for i, quantizer in enumerate(quantizers): + with self.subTest(i=i): + ep = torch.export.export(module, sample_input).module() + prepared = prepare_pt2e(ep, quantizer) + with SeqMSE(prepared, 100): + prepared(*sample_input) + converted = convert_pt2e(prepared) + self.lower_module_and_test_output(converted, sample_input) + class TestExampleLLMScript(TestQNN): - def test_llama3_2_1b(self): + def test_static_gemma3_1b(self): + if not self.required_envs(): + self.skipTest("missing required envs") + + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + "--prompt", + f"{prompt}", + "--ptq", + "16a4w_block", + "--temperature", + "0", + "--decoder_model", + "gemma3-1b", + "--model_mode", + "kv", + "--max_seq_len", + "1024", + "--eval_perplexity", + "--tasks", + "wikitext", + "--limit", + "1", + "--enable_masked_softmax", + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + if not self.compile_only: + self.assertLessEqual(msg["wiki_ppl"], 23) + if not self.enable_x86_64: + pte_size = msg["pte_size"] + self.assertLessEqual(pte_size, 1_200_000_000) # 1.2GB + inference_speed_ref = {"SM8650": 70, "SM8750": 100} + if ( + not self.compile_only + and not self.enable_x86_64 + and self.model in inference_speed_ref + ): + self.assertGreaterEqual( + msg["inference_speed"], inference_speed_ref[self.model] + ) + + def test_llama3_2_instruct(self): if not self.required_envs(): self.skipTest("missing required envs") assert ( @@ -4123,20 +4765,19 @@ def test_llama3_2_1b(self): str(self.port), "--prompt", f"{prompt}", - "--ptq", - "16a4w", "--temperature", "0", "--decoder_model", - "llama3_2", + "llama3_2-1b_instruct", "--model_mode", - "hybrid", - "--prefill_ar_len", - "32", + "kv", "--max_seq_len", - "512", - "--num_sharding", - "4", + "1024", + "--eval_perplexity", + "--tasks", + "wikitext", + "--limit", + "1", ] if self.compile_only: cmds.extend(["--compile_only"]) @@ -4149,7 +4790,6 @@ def test_llama3_2_1b(self): if self.pre_gen_pte: cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) - golden_start_with = "<|start_header_id|>user<|end_header_id|>" p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) with Listener((self.ip, self.port)) as listener: conn = listener.accept() @@ -4158,19 +4798,17 @@ def test_llama3_2_1b(self): if "Error" in msg: self.fail(msg["Error"]) else: - if not self.compile_only: - model_out = msg["result"][0] - self.assertTrue( - model_out.startswith(golden_start_with), - f"Expected Output: {golden_start_with}. Actual Output: {model_out}", + inference_speed_ref = {"SM8650": 37, "SM8750": 49} + if ( + not self.compile_only + and not self.enable_x86_64 + and self.model in inference_speed_ref + ): + self.assertLessEqual(msg["pte_size"], 1_500_000_000) + self.assertLessEqual(msg["wiki_ppl"], 15) + self.assertGreaterEqual( + msg["inference_speed"], inference_speed_ref[self.model] ) - # x86 does not allow weight sharing, so we don't check pte size. - # Inference speed on x86 is slow, so we only check when running on Android - if not self.enable_x86_64: - pte_size = msg["pte_size"] - self.assertLessEqual(pte_size, 1300000000) - if not self.compile_only and not self.enable_x86_64: - self.assertGreaterEqual(msg["inference_speed"], 66) # Lanai def test_llama_stories_260k(self): if not self.required_envs(): @@ -4203,8 +4841,6 @@ def test_llama_stories_260k(self): str(self.port), "--prompt", f"{prompt}", - "--ptq", - "16a4w", "--temperature", "0", "--decoder_model", @@ -4246,7 +4882,7 @@ def test_llama_stories_260k(self): # x86 does not allow weight sharing, so we don't check pte size if not self.enable_x86_64: pte_size = msg["pte_size"] - self.assertLessEqual(pte_size, 2020000) + self.assertLessEqual(pte_size, 2_020_000) # 2MB if not self.compile_only and not self.enable_x86_64: self.assertGreaterEqual(msg["inference_speed"], 1600) # Lanai @@ -4281,8 +4917,6 @@ def test_llama_stories_110m(self): str(self.port), "--prompt", f"{prompt}", - "--ptq", - "16a4w", "--temperature", "0", "--decoder_model", @@ -4323,10 +4957,69 @@ def test_llama_stories_110m(self): # x86 does not allow weight sharing, so we don't check pte size if not self.enable_x86_64: pte_size = msg["pte_size"] - self.assertLessEqual(pte_size, 130000000) + self.assertLessEqual(pte_size, 130_000_000) # 130MB if not self.compile_only and not self.enable_x86_64: self.assertGreaterEqual(msg["inference_speed"], 220) # Lanai + def test_static_phi4(self): + if not self.required_envs(): + self.skipTest("missing required envs") + + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + "--prompt", + f"{prompt}", + "--decoder_model", + "phi_4_mini", + "--model_mode", + "kv", + "--max_seq_len", + "1024", + "--eval_perplexity", + "--tasks", + "wikitext", + "--limit", + "1", + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + inference_speed_ref = {"SM8650": 14, "SM8750": 19} + self.assertLessEqual(msg["wiki_ppl"], 12) + self.assertLessEqual(msg["pte_size"], 4_000_000_000) # 4GB + if self.model in inference_speed_ref: + self.assertGreaterEqual( + msg["inference_speed"], inference_speed_ref[self.model] + ) + def test_static_qwen2_5(self): if not self.required_envs(): self.skipTest("missing required envs") @@ -4347,16 +5040,190 @@ def test_static_qwen2_5(self): str(self.port), "--prompt", f"{prompt}", + "--decoder_model", + "qwen2_5-0_5b", + "--model_mode", + "kv", + "--max_seq_len", + "1024", + "--eval_perplexity", + "--tasks", + "wikitext", + "--limit", + "1", + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + inference_speed_ref = {"SM8650": 115, "SM8750": 155} + self.assertLessEqual(msg["wiki_ppl"], 15) + self.assertLessEqual(msg["pte_size"], 600_000_000) # 600MB + if self.model in inference_speed_ref: + self.assertGreaterEqual( + msg["inference_speed"], inference_speed_ref[self.model] + ) + + def test_static_qwen3(self): + if not self.required_envs(): + self.skipTest("missing required envs") + + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + "--prompt", + f"{prompt}", + "--decoder_model", + "qwen3-0_6b", + "--model_mode", + "kv", + "--max_seq_len", + "1024", + "--eval_perplexity", + "--tasks", + "wikitext", + "--limit", + "1", + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + inference_speed_ref = {"SM8650": 38, "SM8750": 56} + self.assertLessEqual(msg["wiki_ppl"], 18) + self.assertLessEqual(msg["pte_size"], 950_000_000) # 950MB + if self.model in inference_speed_ref: + self.assertGreaterEqual( + msg["inference_speed"], inference_speed_ref[self.model] + ) + + def test_qwen2_5(self): + if not self.required_envs([]): + self.skipTest("missing required envs") + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/qwen2_5/qwen2_5.py", + "--prompt", + prompt, + "--decoder_model", + "qwen2.5_0.5B", "--ptq", "16a8w", + "--enable_spinquant_r3", + "--max_seq_len", + "128", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + golden_start_with = "My favourite condiment is iced tea." + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + if not self.compile_only: + model_out = msg["result"][0] + self.assertTrue( + model_out.startswith(golden_start_with), + f"Expected Output: '{golden_start_with}' Actual Output: '{model_out}'", + ) + + def test_static_smollm2(self): + if not self.required_envs(): + self.skipTest("missing required envs") + + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + "--prompt", + f"{prompt}", "--decoder_model", - "qwen2_5", + "smollm2_135m", "--model_mode", "kv", + "--temperature", + "0", + "--prefill_ar_len", + "128", "--max_seq_len", "1024", "--eval_perplexity", - "--tasks", + "--task", "wikitext", "--limit", "1", @@ -4380,9 +5247,65 @@ def test_static_qwen2_5(self): if "Error" in msg: self.fail(msg["Error"]) else: - inference_speed_ref = {"SM8650": 110, "SM8750": 130} self.assertLessEqual(msg["wiki_ppl"], 25) - self.assertLessEqual(msg["pte_size"], 800000000) # 800mb + self.assertGreaterEqual(msg["inference_speed"], 200) + + def test_static_smollm3(self): + if not self.required_envs(): + self.skipTest("missing required envs") + + prompt = "My favourite condiment is " + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py", + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--model", + self.model, + "--ip", + self.ip, + "--port", + str(self.port), + "--prompt", + f"{prompt}", + "--decoder_model", + "smollm3-3b", + "--model_mode", + "kv", + "--temperature", + "0", + "--max_seq_len", + "1024", + "--eval_perplexity", + "--task", + "wikitext", + "--limit", + "1", + ] + if self.compile_only: + cmds.extend(["--compile_only"]) + elif self.device: + cmds.extend(["--device", self.device]) + if self.host: + cmds.extend(["--host", self.host]) + elif self.enable_x86_64: + cmds.extend(["--enable_x86_64"]) + if self.pre_gen_pte: + cmds.extend(["--pre_gen_pte", self.pre_gen_pte]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + inference_speed_ref = {"SM8650": 23, "SM8750": 28} + self.assertLessEqual(msg["wiki_ppl"], 10) + self.assertLessEqual(msg["pte_size"], 2_600_000_000) # 2.6GB if self.model in inference_speed_ref: self.assertGreaterEqual( msg["inference_speed"], inference_speed_ref[self.model] @@ -4410,6 +5333,8 @@ def test_albert(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4422,7 +5347,7 @@ def test_albert(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["accuracy"], 0.8) + self.assertGreaterEqual(msg["accuracy"], 0.95) def test_bert(self): if not self.required_envs([self.sentence_dataset]): @@ -4444,6 +5369,8 @@ def test_bert(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4456,7 +5383,7 @@ def test_bert(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["accuracy"], 0.6) + self.assertGreaterEqual(msg["accuracy"], 0.55) def test_conv_former(self): if not self.required_envs([self.image_dataset]): @@ -4479,6 +5406,8 @@ def test_conv_former(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4491,8 +5420,8 @@ def test_conv_former(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 70) + self.assertGreaterEqual(msg["top_5"], 92) def test_cvt(self): if not self.required_envs([self.image_dataset]): @@ -4515,6 +5444,8 @@ def test_cvt(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4550,6 +5481,8 @@ def test_deit(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4562,8 +5495,8 @@ def test_deit(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 75) - self.assertGreaterEqual(msg["top_5"], 90) + self.assertGreaterEqual(msg["top_1"], 76) + self.assertGreaterEqual(msg["top_5"], 92) def test_dino_v2(self): if not self.required_envs([self.image_dataset]): @@ -4585,6 +5518,8 @@ def test_dino_v2(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4597,8 +5532,8 @@ def test_dino_v2(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 70) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_1"], 62) + self.assertGreaterEqual(msg["top_5"], 87) def test_distilbert(self): if not self.required_envs([self.sentence_dataset]): @@ -4620,6 +5555,8 @@ def test_distilbert(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4632,7 +5569,7 @@ def test_distilbert(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["accuracy"], 0.45) + self.assertGreaterEqual(msg["accuracy"], 0.48) def test_dit(self): if not self.required_envs(): @@ -4653,6 +5590,8 @@ def test_dit(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4665,15 +5604,15 @@ def test_dit(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 80) - self.assertGreaterEqual(msg["top_5"], 95) + self.assertGreaterEqual(msg["top_1"], 78) + self.assertGreaterEqual(msg["top_5"], 92) def test_efficientnet(self): if not self.required_envs([self.image_dataset]): self.skipTest("missing required envs") cmds = [ "python", - f"{self.executorch_root}/examples/qualcomm/oss_scripts/efficientnet.py" + f"{self.executorch_root}/examples/qualcomm/oss_scripts/efficientnet.py", "--dataset", self.image_dataset, "--artifact", @@ -4688,6 +5627,8 @@ def test_efficientnet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4700,8 +5641,8 @@ def test_efficientnet(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 70) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_1"], 61) + self.assertGreaterEqual(msg["top_5"], 88) def test_efficientSAM(self): if not self.required_envs( @@ -4729,6 +5670,8 @@ def test_efficientSAM(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4741,10 +5684,11 @@ def test_efficientSAM(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["MIoU"], 0.55) + # test with efficient_sam_vitt.pt + self.assertGreaterEqual(msg["MIoU"], 0.95) def test_esrgan(self): - if not self.required_envs(): + if not self.required_envs([self.oss_repo]): self.skipTest("missing required envs") cmds = [ @@ -4765,6 +5709,8 @@ def test_esrgan(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4777,8 +5723,8 @@ def test_esrgan(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["PSNR"], 24) - self.assertGreaterEqual(msg["SSIM"], 0.8) + self.assertGreaterEqual(msg["PSNR"], 23) + self.assertGreaterEqual(msg["SSIM"], 0.85) def test_eurobert(self): if not self.required_envs([self.sentence_dataset]): @@ -4800,6 +5746,8 @@ def test_eurobert(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4812,7 +5760,7 @@ def test_eurobert(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["accuracy"], 0.5) + self.assertGreaterEqual(msg["accuracy"], 0.54) def test_fastvit(self): if not self.required_envs( @@ -4840,6 +5788,8 @@ def test_fastvit(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4853,7 +5803,7 @@ def test_fastvit(self): self.fail(msg["Error"]) else: self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_5"], 77) def test_fbnet(self): if not self.required_envs([self.image_dataset]): @@ -4876,6 +5826,8 @@ def test_fbnet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4888,8 +5840,8 @@ def test_fbnet(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 90) + self.assertGreaterEqual(msg["top_1"], 67) + self.assertGreaterEqual(msg["top_5"], 88) def test_focalnet(self): if not self.required_envs([self.image_dataset]): @@ -4912,6 +5864,8 @@ def test_focalnet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4926,7 +5880,7 @@ def test_focalnet(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 55) + self.assertGreaterEqual(msg["top_1"], 54) self.assertGreaterEqual(msg["top_5"], 80) def test_gMLP(self): @@ -4950,6 +5904,8 @@ def test_gMLP(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -4962,8 +5918,8 @@ def test_gMLP(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_1"], 70) + self.assertGreaterEqual(msg["top_5"], 88) @unittest.skip("Only outputs good accuracy in QNN 2.29") def test_mobilevit_v2(self): @@ -4987,6 +5943,8 @@ def test_mobilevit_v2(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5004,13 +5962,13 @@ def test_mobilevit_v2(self): self.assertGreaterEqual(msg["top_1"], 50) self.assertGreaterEqual(msg["top_5"], 85) - def test_mobilevit1(self): + def test_mobilevit_v1(self): if not self.required_envs([self.image_dataset]): self.skipTest("missing required envs") cmds = [ "python", - f"{self.executorch_root}/examples/qualcomm/oss_scripts/mobilevit1.py", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/mobilevit_v1.py", "--dataset", self.image_dataset, "--artifact", @@ -5025,6 +5983,8 @@ def test_mobilevit1(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5037,8 +5997,8 @@ def test_mobilevit1(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 70) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_1"], 76) + self.assertGreaterEqual(msg["top_5"], 93) def test_pvt(self): if not self.required_envs([self.image_dataset]): @@ -5061,6 +6021,8 @@ def test_pvt(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5074,7 +6036,7 @@ def test_pvt(self): self.fail(msg["Error"]) else: self.assertGreaterEqual(msg["top_1"], 65) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_5"], 83) def test_regnet(self): if not self.required_envs([self.image_dataset]): @@ -5098,6 +6060,8 @@ def test_regnet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5113,8 +6077,8 @@ def test_regnet(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 85) + self.assertGreaterEqual(msg["top_1"], 64) + self.assertGreaterEqual(msg["top_5"], 86) def test_retinanet(self): if not self.required_envs([self.image_dataset]): @@ -5137,6 +6101,8 @@ def test_retinanet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5171,6 +6137,8 @@ def test_roberta(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5183,7 +6151,7 @@ def test_roberta(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["accuracy"], 0.5) + self.assertGreaterEqual(msg["accuracy"], 0.54) def test_squeezenet(self): if not self.required_envs([self.image_dataset]): @@ -5206,6 +6174,8 @@ def test_squeezenet(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5218,8 +6188,8 @@ def test_squeezenet(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 45) - self.assertGreaterEqual(msg["top_5"], 70) + self.assertGreaterEqual(msg["top_1"], 27) + self.assertGreaterEqual(msg["top_5"], 59) def test_ssd300_vgg16(self): if not self.required_envs([self.pretrained_weight, self.oss_repo]): @@ -5244,6 +6214,8 @@ def test_ssd300_vgg16(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5256,7 +6228,7 @@ def test_ssd300_vgg16(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["mAP"], 0.70) + self.assertGreaterEqual(msg["mAP"], 0.76) def test_swin_transformer(self): if not self.required_envs([self.image_dataset]): @@ -5278,6 +6250,8 @@ def test_swin_transformer(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5290,8 +6264,8 @@ def test_swin_transformer(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 71) + self.assertGreaterEqual(msg["top_5"], 90) def test_t5(self): if not self.required_envs([self.qa_dataset]): @@ -5300,7 +6274,7 @@ def test_t5(self): "python", f"{self.executorch_root}/examples/qualcomm/oss_scripts/t5/t5.py", "--dataset", - self.sentence_dataset, + self.qa_dataset, "--artifact", self.artifact_dir, "--build_folder", @@ -5313,6 +6287,8 @@ def test_t5(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5325,7 +6301,7 @@ def test_t5(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["f1"], 0.7) + self.assertGreaterEqual(msg["f1"], 0.72) def test_whisper(self): if not self.required_envs(): @@ -5346,6 +6322,8 @@ def test_whisper(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5611,6 +6589,8 @@ def test_mobilenet_v2(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5625,8 +6605,8 @@ def test_mobilenet_v2(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 52) + self.assertGreaterEqual(msg["top_5"], 84) def test_mobilenet_v3(self): if not self.required_envs([self.image_dataset]): @@ -5649,6 +6629,8 @@ def test_mobilenet_v3(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5663,8 +6645,8 @@ def test_mobilenet_v3(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 55) + self.assertGreaterEqual(msg["top_5"], 81) def test_inception_v3(self): if not self.required_envs([self.image_dataset]): @@ -5687,6 +6669,8 @@ def test_inception_v3(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5701,8 +6685,8 @@ def test_inception_v3(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 59) + self.assertGreaterEqual(msg["top_5"], 78) def test_inception_v4(self): if not self.required_envs([self.image_dataset]): @@ -5725,6 +6709,8 @@ def test_inception_v4(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5739,8 +6725,8 @@ def test_inception_v4(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 60) - self.assertGreaterEqual(msg["top_5"], 80) + self.assertGreaterEqual(msg["top_1"], 64) + self.assertGreaterEqual(msg["top_5"], 85) def test_vit(self): if not self.required_envs([self.image_dataset]): @@ -5763,6 +6749,8 @@ def test_vit(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5777,8 +6765,8 @@ def test_vit(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["top_1"], 65) - self.assertGreaterEqual(msg["top_5"], 90) + self.assertGreaterEqual(msg["top_1"], 69) + self.assertGreaterEqual(msg["top_5"], 91) def test_edsr(self): if not self.required_envs(): @@ -5800,6 +6788,8 @@ def test_edsr(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5814,8 +6804,8 @@ def test_edsr(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["PSNR"], 25) - self.assertGreaterEqual(msg["SSIM"], 0.8) + self.assertGreaterEqual(msg["PSNR"], 29) + self.assertGreaterEqual(msg["SSIM"], 0.94) def test_deeplab_v3(self): if not self.required_envs(): @@ -5837,6 +6827,8 @@ def test_deeplab_v3(self): self.ip, "--port", str(self.port), + "--seed", + str(1126), ] if self.host: cmds.extend(["--host", self.host]) @@ -5851,9 +6843,9 @@ def test_deeplab_v3(self): if "Error" in msg: self.fail(msg["Error"]) else: - self.assertGreaterEqual(msg["PA"], 0.85) - self.assertGreaterEqual(msg["MPA"], 0.70) - self.assertGreaterEqual(msg["MIoU"], 0.55) + self.assertGreaterEqual(msg["PA"], 0.88) + self.assertGreaterEqual(msg["MPA"], 0.79) + self.assertGreaterEqual(msg["MIoU"], 0.67) @unittest.skip("dynamic shape inputs appear in recent torch.export.export") def test_mobilebert(self): @@ -6167,6 +7159,11 @@ def setup_environment(): help="Location for imagenet dataset", type=str, ) + parser.add_argument( + "--qa_dataset", + help="Location for QA dataset", + type=str, + ) parser.add_argument( "--sentence_dataset", help="Location for sentence dataset", @@ -6230,6 +7227,7 @@ def setup_environment(): TestQNN.executorch_root = args.executorch_root TestQNN.artifact_dir = args.artifact_dir TestQNN.image_dataset = args.image_dataset + TestQNN.qa_dataset = args.qa_dataset TestQNN.sentence_dataset = args.sentence_dataset TestQNN.pretrained_weight = args.pretrained_weight TestQNN.model_name = args.model_name diff --git a/backends/qualcomm/tests/tester.py b/backends/qualcomm/tests/tester.py new file mode 100644 index 00000000000..812e8971115 --- /dev/null +++ b/backends/qualcomm/tests/tester.py @@ -0,0 +1,113 @@ +# 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 typing import Any, List, Optional, Sequence, Tuple + +import executorch +import executorch.backends.test.harness.stages as BaseStages + +import torch +from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager +from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner +from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_soc_to_chipset_map, +) +from executorch.backends.test.harness import Tester as TesterBase +from executorch.backends.test.harness.stages import StageType +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.backend.partitioner import Partitioner +from torch.export import ExportedProgram + + +class Quantize(BaseStages.Quantize): + def __init__( + self, + quantizer: QnnQuantizer, + quantization_config: Optional[Any] = None, + calibrate: bool = True, + calibration_samples: Optional[Sequence[Any]] = None, + is_qat: Optional[bool] = False, + ): + super().__init__( + quantizer=quantizer, + calibrate=calibrate, + calibration_samples=calibration_samples, + is_qat=is_qat, + set_global=False, + ) + + +class Partition(BaseStages.Partition): + def __init__(self, partitioner: Optional[Partitioner] = None): + super().__init__( + partitioner=partitioner or QnnPartitioner, + ) + + +class ToEdgeTransformAndLower(BaseStages.ToEdgeTransformAndLower): + def __init__( + self, + partitioners: Optional[List[Partitioner]] = None, + edge_compile_config: Optional[EdgeCompileConfig] = None, + soc_model: str = "SM8650", + use_fp16: bool = True, + ): + backend_options = generate_htp_compiler_spec(use_fp16=use_fp16) + self.chipset = get_soc_to_chipset_map()[soc_model] + self.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset, + backend_options=backend_options, + ) + + super().__init__( + partitioners=partitioners or [QnnPartitioner(self.compiler_specs)], + edge_compile_config=edge_compile_config + or EdgeCompileConfig(_check_ir_validity=False), + default_partitioner_cls=QnnPartitioner, + ) + + def run( + self, artifact: ExportedProgram, inputs=None, generate_etrecord: bool = False + ) -> None: + ep = QnnPassManager().transform_for_export_pipeline(artifact) + transform_passes = QnnPassManager().get_to_edge_transform_passes(ep) + + self.edge_dialect_program = to_edge_transform_and_lower( + ep, + transform_passes=transform_passes, + partitioner=self.partitioners, + compile_config=self.edge_compile_conf, + generate_etrecord=generate_etrecord, + ) + + +class QualcommTester(TesterBase): + def __init__( + self, + module: torch.nn.Module, + example_inputs: Tuple[torch.Tensor], + dynamic_shapes: Optional[Tuple[Any]] = None, + use_fp16: bool = True, + ): + def create_to_edge_transform_and_lower(*args, **kwargs): + kwargs["use_fp16"] = use_fp16 + return ToEdgeTransformAndLower(*args, **kwargs) + + # Specialize for Qualcomm + stage_classes = executorch.backends.test.harness.Tester.default_stage_classes() | { + StageType.PARTITION: Partition, + StageType.TO_EDGE_TRANSFORM_AND_LOWER: create_to_edge_transform_and_lower, + } + + super().__init__( + module=module, + stage_classes=stage_classes, + example_inputs=example_inputs, + dynamic_shapes=dynamic_shapes, + ) diff --git a/backends/qualcomm/tests/utils.py b/backends/qualcomm/tests/utils.py index fd2d10e2b93..93eee4dfc31 100644 --- a/backends/qualcomm/tests/utils.py +++ b/backends/qualcomm/tests/utils.py @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import collections -import copy import os import subprocess import tempfile @@ -30,7 +29,7 @@ get_soc_to_chipset_map, to_edge_transform_and_lower_to_qnn, ) -from executorch.devtools import generate_etrecord, Inspector +from executorch.devtools import Inspector from executorch.devtools.inspector._inspector_utils import TimeScale from executorch.examples.qualcomm.utils import ( generate_inputs, @@ -144,30 +143,6 @@ def validate_context_binary(ctx_bin: bytes): assert os.path.isfile(f"{tmp_dir}/ctx.json"), print(result.stderr) -def validate_qcir(qcir: bytes): - with tempfile.TemporaryDirectory() as tmp_dir: - with open(f"{tmp_dir}/qcir.bin", "wb") as binary_file: - binary_file.write(qcir) - - cmds = [ - "flatc", - "-o", - tmp_dir, - "--raw-binary", - "-t", - f"{os.path.dirname(__file__)}/../aot/ir/qcir.fbs", - "--", - f"{tmp_dir}/qcir.bin", - ] - result = subprocess.run( - " ".join(cmds), - shell=True, - executable="/bin/bash", - capture_output=True, - ) - assert os.path.isfile(f"{tmp_dir}/qcir.json"), print(result.stderr) - - class TestQNN(unittest.TestCase): rtol: float = 0 atol: float = 0 @@ -198,6 +173,8 @@ class TestQNN(unittest.TestCase): pre_gen_pte: str = "" llama_artifacts: str = "" dump_intermediate_outputs: bool = False + inference_speed: float = 0.0 + inference_speed_output_path = "outputs/inference_speed.txt" def _assert_outputs_equal(self, model_output, ref_output): self.assertTrue(len(ref_output) == len(model_output)) @@ -216,13 +193,6 @@ def _save_model_and_expected_output( inputs: Tuple[torch.Tensor], dir_name: str, ) -> None: - # Save the input data list to be executed - input_list = "" - for idx, _ in enumerate(inputs): - input_name = f"input_0_{idx}.raw" - input_list += input_name + " " - input_list = input_list.strip() + "\n" - ref_output = module(*inputs) # Save the expected output data to be verified @@ -239,7 +209,7 @@ def _save_model_and_expected_output( with open(pte_fname, "wb") as file: file.write(buffer) - return input_list, ref_outputs, pte_fname + return ref_outputs, pte_fname def required_envs(self, conditions=None) -> bool: conditions = [] if conditions is None else conditions @@ -264,10 +234,12 @@ def verify_output( # noqa: C901 output_encodings: Tuple = (), check_io_shape: bool = False, op_package_paths: List[str] = None, + extra_cmds: str = "", + output_callback: Optional[Callable[[str], None]] = None, + save_inference_speed: bool = False, ): with tempfile.TemporaryDirectory() as tmp_dir: ( - input_list, ref_outputs, pte_fname, ) = self._save_model_and_expected_output( @@ -287,7 +259,9 @@ def post_process(): torch_to_numpy_dtype_dict, ) - for i, f in enumerate(sorted(os.listdir(output_dir))): + for i, f in enumerate( + sorted(f for f in os.listdir(output_dir) if f.endswith(".raw")) + ): enc = output_encodings[i] if len(output_encodings) != 0 else None dtype = ( ref_outputs[i].numpy().dtype @@ -337,9 +311,7 @@ def validate_intermediate_tensor(): ) if self.enable_x86_64: - generate_inputs( - tmp_dir, "input_list.txt", [processed_inputs], input_list - ) + generate_inputs(tmp_dir, "input_list.txt", [processed_inputs]) make_output_dir(output_dir) target = "x86_64-linux-clang" @@ -368,6 +340,13 @@ def validate_intermediate_tensor(): ] if expected_intermediate_events != -1: cmd.append("--dump_intermediate_outputs") + cmd += extra_cmds.split() + + if save_inference_speed: + cmd += [ + "--performance_output_path", + self.inference_speed_output_path, + ] if check_io_shape: shape_info = { @@ -387,16 +366,19 @@ def validate_intermediate_tensor(): cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, env=env, cwd=tmp_dir, ) + if output_callback: + output_callback(proc) self.assertEqual( proc.returncode, 0, f"The process running qnn_executorch_runner return {proc.returncode}, " "STDOUT=\n" - f"{proc.stdout.decode('utf-8')}", + f"{proc.stdout}", ) # Verify the outputs @@ -409,6 +391,13 @@ def validate_intermediate_tensor(): if expected_intermediate_events != -1: validate_intermediate_tensor() + + if save_inference_speed: + with open( + f"{tmp_dir}/{self.inference_speed_output_path}", "r" + ) as f: + self.inference_speed = float(f.read()) + else: adb = SimpleADB( qnn_sdk=os.getenv("QNN_SDK_ROOT"), @@ -435,10 +424,14 @@ def validate_intermediate_tensor(): ) adb.push( inputs=[processed_inputs], - input_list=input_list, files=op_package_paths, ) - adb.execute(method_index=method_index) + adb.extra_cmds += extra_cmds + if save_inference_speed: + adb.extra_cmds += ( + f" --performance_output_path {self.inference_speed_output_path}" + ) + adb.execute(method_index=method_index, output_callback=output_callback) adb.pull(output_path=tmp_dir, callback=post_process) self._assert_outputs_equal(outputs, ref_outputs) @@ -451,6 +444,11 @@ def validate_intermediate_tensor(): debug_output_path, callback=validate_intermediate_tensor, ) + if save_inference_speed: + with open( + f"{tmp_dir}/{self.inference_speed_output_path}", "r" + ) as f: + self.inference_speed = float(f.read()) def lower_module_and_test_output( self, @@ -465,6 +463,9 @@ def lower_module_and_test_output( skip_node_op_set: set = None, skip_mutable_buffer: bool = False, dynamic_shapes: Dict = None, + extra_cmds: str = "", + output_callback: Optional[Callable[[str], None]] = None, + save_inference_speed: bool = False, ): delegated_program = to_edge_transform_and_lower_to_qnn( module, @@ -475,11 +476,9 @@ def lower_module_and_test_output( skip_node_id_set=skip_node_id_set, skip_node_op_set=skip_node_op_set, skip_mutable_buffer=skip_mutable_buffer, + generate_etrecord=self.enable_profile, ) - # this is needed for the ETRecord as lowering modifies the graph in-place - edge_copy = copy.deepcopy(delegated_program) - exec_prog = delegated_program.to_executorch( exir.ExecutorchBackendConfig( # For shared buffer, user must pass the memory address @@ -506,7 +505,7 @@ def lower_module_and_test_output( etrecord_path = "etrecord.bin" if self.enable_profile: - generate_etrecord(etrecord_path, edge_copy, exec_prog) + exec_prog.get_etrecord().save(etrecord_path) # Check numerics if ( assert_output_equal @@ -520,6 +519,9 @@ def lower_module_and_test_output( etrecord_path, expected_profile_events, expected_intermediate_events, + extra_cmds=extra_cmds, + output_callback=output_callback, + save_inference_speed=save_inference_speed, ) def get_qdq_module( @@ -574,7 +576,7 @@ def get_prepared_qat_module( quant_dtype: QuantDtype = QuantDtype.use_8a8w, submodule_qconfig_list: Optional[List[Tuple[Callable, ModuleQConfig]]] = None, ) -> torch.fx.GraphModule: - m = torch.export.export_for_training(module, inputs, strict=True).module() + m = torch.export.export(module, inputs, strict=True).module() quantizer = make_quantizer( quant_dtype=quant_dtype, @@ -658,7 +660,7 @@ def _insert_clone( users = list(node.users.keys()) inserted_node = graph_module.graph.create_node( "call_function", - exir_ops.edge.aten.clone.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, (node,), ) inserted_node.meta["val"] = node.meta["val"] diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index ff611385de5..be4e86de50f 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -333,7 +333,9 @@ def to_edge_transform_and_lower_to_qnn( passes_job: Optional[Union[OrderedDict, Dict[str, OrderedDict]]] = None, skip_node_id_set: Optional[set] = None, skip_node_op_set: Optional[set] = None, - skip_mutable_buffer: bool = False, + skip_mutable_buffer: Optional[bool] = False, + generate_etrecord: Optional[bool] = False, + convert_linear_to_conv2d: Optional[bool] = False, ) -> EdgeProgramManager: """ Transforms and lowers a given PyTorch module to the QNN backend. @@ -358,8 +360,10 @@ def to_edge_transform_and_lower_to_qnn( Set of node IDs to skip during partitioning. skip_node_op_set (Optional[set]): Set of node operations to skip during partitioning. - skip_mutable_buffer (Optional[set]): + skip_mutable_buffer (Optional[bool]): Whether to skip delegating the mutable buffer in QNN backend. + convert_linear_to_conv2d (Optional[bool]): + Whether to convert linear to conv2d in some cases to improve performance in HTP backend. Returns: EdgeProgramManager: @@ -431,7 +435,9 @@ def ensure_graph_specific_dict(value, graph_names): # If placed in the to_edge_transform_passes, it will be executed # after the lift_constant_tensor_pass, causing the operation builder # to fail to correctly retrieve the parameter by the get_parameter. - aten_programs[graph_name] = QnnPassManager().transform_for_export_pipeline(ep) + aten_programs[graph_name] = QnnPassManager().transform_for_export_pipeline( + ep, convert_linear_to_conv2d=convert_linear_to_conv2d + ) transform_passes[graph_name] = QnnPassManager().get_to_edge_transform_passes( ep, passes_job=passes_job[graph_name], dep_table=dep_table[graph_name] ) @@ -442,6 +448,7 @@ def ensure_graph_specific_dict(value, graph_names): partitioner=qnn_partitioners, constant_methods=constant_methods, compile_config=qnn_edge_config(), + generate_etrecord=generate_etrecord, ) @@ -1038,7 +1045,7 @@ def generate_qnn_executorch_compiler_spec( qnn_executorch_options.log_level = ( QnnExecuTorchLogLevel.kLogLevelDebug if debug - else QnnExecuTorchLogLevel.kLogLevelWarn + else QnnExecuTorchLogLevel.kLogLevelError ) qnn_executorch_options.dump_intermediate_outputs = dump_intermediate_outputs diff --git a/backends/samsung/CMakeLists.txt b/backends/samsung/CMakeLists.txt new file mode 100644 index 00000000000..fff3ece5239 --- /dev/null +++ b/backends/samsung/CMakeLists.txt @@ -0,0 +1,168 @@ +# 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) + +get_filename_component( + EXECUTORCH_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE +) + +if(NOT DEFINED EXYNOS_AI_LITECORE_ROOT) + message( + FATAL_ERROR + "Please define EXYNOS_AI_LIRECORE_PATH by adding cmake parameter -DEXYNOS_AI_LITECORE_ROOT=<...>" + ) +endif() +if(CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") + message(FATAL_ERROR "IOS is not supported on Exynos.") +endif() + +if(NOT FLATC_EXECUTABLE) + set(FLATC_EXECUTABLE flatc) +endif() + +add_compile_options(-Wall -Werror -fPIC) +if(CMAKE_BUILD_TYPE STREQUAL "Release") + # strip symbols + add_link_options("-s") + # hide dynamic symbols + set(CMAKE_C_VISIBILITY_PRESET hidden) + set(CMAKE_CXX_VISIBILITY_PRESET hidden) + add_definitions(-DNDEBUG) +endif() + +include_directories( + ${EXECUTORCH_SOURCE_DIR}/.. + ${EXECUTORCH_SOURCE_DIR}/runtime/core/portable_type/c10 + ${EXYNOS_AI_LITECORE_ROOT} +) +add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) + +if(${ANDROID}) + find_library(android_log log) +endif() + +# add logging library +add_library(enn_logging STATIC) + +if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64") + add_subdirectory( + ${EXECUTORCH_SOURCE_DIR}/third-party/pybind11 + ${CMAKE_CURRENT_BINARY_DIR}/pybind11 + ) + add_library(PyEnnWrapperAdaptor MODULE) + + find_library( + GG_API_LIB + NAMES graphgen_api + HINTS ${EXYNOS_AI_LITECORE_ROOT}/lib/x86_64-linux + ) + add_library(graphgen_api SHARED IMPORTED GLOBAL) + set_target_properties( + graphgen_api + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + "${EXYNOS_AI_LITECORE_ROOT}/include" IMPORTED_LOCATION + "${GG_API_LIB}" + ) + + set(_enn_compile_options_schema + ${CMAKE_CURRENT_SOURCE_DIR}/serialization/compile_options_def.fbs + ) + + set(_enn_schema_generate_dir + "${CMAKE_BINARY_DIR}/schema/include/executorch/backends/samsung" + ) + # Paths to headers generated from the .fbs files. + string(REGEX REPLACE "serialization/([^/]+)[.]fbs$" "\\1_generated.h" + generated_header "${fbs_file}" + ) + set(_enn_schema_output "${_enn_schema_generate_dir}/${generated_header}") + + # Generate the headers from the .fbs files. + add_custom_command( + OUTPUT ${_enn_schema_output} + COMMAND ${FLATC_EXECUTABLE} --cpp --cpp-std c++11 --scoped-enums -o + "${_enn_schema_generate_dir}" ${_enn_compile_options_schema} + DEPENDS ${_enn_compile_options_schema} + WORKING_DIRECTORY ${EXECUTORCH_SOURCE_DIR} + COMMENT "Generating enn compile options headers" + VERBATIM + ) + add_custom_target( + enn_compile_options_output ALL DEPENDS ${_enn_schema_output} + ) + + set_target_properties( + PyEnnWrapperAdaptor PROPERTIES CXX_VISIBILITY_PRESET hidden + ) + target_link_libraries( + PyEnnWrapperAdaptor PRIVATE pybind11::module pybind11::lto graphgen_api + enn_logging + ) + target_include_directories( + PyEnnWrapperAdaptor BEFORE + PRIVATE ${CMAKE_BINARY_DIR}/schema/include + ${EXECUTORCH_SOURCE_DIR}/third-party/flatbuffers/include + ) + add_dependencies(PyEnnWrapperAdaptor enn_compile_options_output) + pybind11_extension(PyEnnWrapperAdaptor) + + # PyGraphWrapperAdaptor + add_library(PyGraphWrapperAdaptor MODULE) + # + find_library( + GRAPH_WRAPPER_LIB + NAMES graph_wrapper + HINTS ${EXYNOS_AI_LITECORE_ROOT}/lib/x86_64-linux + ) + add_library(graph_wrapper SHARED IMPORTED GLOBAL) + set_target_properties( + graph_wrapper + PROPERTIES INTERFACE_INCLUDE_DIRECTORIES + "${EXYNOS_AI_LITECORE_ROOT}/include" IMPORTED_LOCATION + "${GRAPH_WRAPPER_LIB}" + ) + set_target_properties( + PyGraphWrapperAdaptor PROPERTIES CXX_VISIBILITY_PRESET hidden + ) + target_link_libraries( + PyGraphWrapperAdaptor PRIVATE pybind11::module pybind11::lto graph_wrapper + enn_logging + ) + pybind11_extension(PyGraphWrapperAdaptor) + + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/aot) +endif() + +if(${ANDROID}) + target_link_libraries(enn_logging PRIVATE ${android_log}) + add_library(enn_backend STATIC) + target_link_libraries(enn_backend PRIVATE enn_logging) + 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 + ) + set_target_properties( + enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden + ) + install( + TARGETS enn_backend enn_logging + EXPORT ExecuTorchTargets + DESTINATION lib + ) +endif() + +add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/runtime) diff --git a/backends/samsung/README.md b/backends/samsung/README.md new file mode 100644 index 00000000000..fcb7ef84b75 --- /dev/null +++ b/backends/samsung/README.md @@ -0,0 +1,72 @@ +# ExecuTorch Samsung Exynos Delegate + +The subtree contains Exynos delegate implementation for ExecuTorch. The target of delegation +is deploying torch model to run with exynos NPU/DSP. + +This backend is implemented on the top of [EXYNOS_LITECORE](https://soc-developer.semiconductor.samsung.com/global/development/ai-litecore) +Please prepare the SDK before you start, it is important to code compilation and runtime. + +## Delegate Options + +### Supported Chipset +- Exynos 2500 (E9955) + +### Supported Inference Type +- Quantized (i8/u8/i16/u16) +- FP16 + +## Directory Structure + +``` +backends/samsung +├── aot # Codes for generating binary buffer for ENN runtime. +├── builders # Codes for lowering each operators. +├── partition # ENN Partitioner. +├── passes # Various passes helping lower models to ENN backend. +├── python # Places to put pybind artifacts for accessing samsung libraries. +├── runtime # ENN runtime for executing lowered models. +├── scripts # Misc supporting scripts, not related to core functionality. +└── serialization # Codes for building Graph IR for Exynos and serializing. + +examples +└── samsung # Examples to run ENN backends. +``` + +## How to build +Please download Exynos AI LiteCore, and set the root path of SDK directory to `EXYNOS_AI_LITECORE_ROOT`.
+Please navigate to [Android NDK](https://developer.android.com/ndk) and download a version of NDK. +`ANDROID_NDK` refers the root path of NDK directory.
+ +### Set up environment variables +```bash +export LD_LIBRARY_PATH=${EXYNOS_AI_LITECORE_ROOT}/lib/x86_64-linux/ +``` + +### Build AOT Targets +Generates python artifacts that allow user call `Compile` interface to lower a model to Exynos backend in python script. +```bash +./backends/samsung/build.sh -b x86_64 +``` + +### Build ENN Executor Runner +```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. + +### Build Anroid 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 \ + -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK}/build/cmake/android.toolchain.cmake" \ + -DANDROID_ABI="${ANDROID_ABI}" \ + -DCMAKE_INSTALL_PREFIX=cmake-android-out \ + -Bcmake-android-out/extension/android + +cmake --build cmake-android-out/extension/android -j8 +``` + +## Examples +python -m executorch.examples.samsung.aot_compiler --chipset e9955 -m ic3 + +Please see this [README.md](../../examples/samsung/README.md). diff --git a/backends/samsung/_passes/conv1d_to_conv2d.py b/backends/samsung/_passes/conv1d_to_conv2d.py new file mode 100644 index 00000000000..57f1074b348 --- /dev/null +++ b/backends/samsung/_passes/conv1d_to_conv2d.py @@ -0,0 +1,88 @@ +# 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. + +import torch +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._export.utils import get_param + + +class Conv1dToConv2d(ExportPass): + + def __init__(self, edge_program: ExportedProgram): + super().__init__() + self.edge_program = edge_program + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + node_list = list(graph.nodes) + for node in node_list: + if node.op == "call_function": + if node.target == exir_ops.edge.aten.convolution.default: + stride = list(node.args[3]) + if len(stride) != 1: + continue + + # convert 3dim weight to 4dim + weight_node = node.args[1] + weight_3dim = get_param(self.edge_program, weight_node) + weight_4dim = torch.nn.Parameter( + data=weight_3dim.data.contiguous().unsqueeze(dim=-1), + requires_grad=False, + ) + parameter_name = ( + self.edge_program.graph_signature.inputs_to_parameters[ + weight_node.name + ] + ) + self.edge_program.state_dict[parameter_name] = weight_4dim + weight_node.meta["val"] = weight_node.meta["val"].data.unsqueeze( + dim=-1 + ) + + # Extend stride, padding, and dilation + node.args = ( + node.args[0], + node.args[1], + node.args[2], + node.args[3] + [1], # stride + node.args[4] + [0], # padding + node.args[5] + [1], # dilation + node.args[6], + node.args[7], + node.args[8], + ) + + # unsqueeze -> conv2d -> squeeze + with graph.inserting_before(node): + input_node = node.args[0] + unsqueeze_before = graph.create_node( + "call_function", exir_ops.edge.aten.unsqueeze_copy.default + ) + unsqueeze_before.args = ( + input_node, + -1, + ) + node.replace_input_with(input_node, unsqueeze_before) + + with graph.inserting_after(node): + squeeze_after = graph.create_node( + "call_function", exir_ops.edge.aten.squeeze_copy.dims + ) + squeeze_after.args = ( + node, + [-1], + ) + original_users = [ + user for user in node.users if user != squeeze_after + ] + for user in original_users: + user.replace_input_with(node, squeeze_after) + + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, True) diff --git a/backends/samsung/_passes/customized_constant_prop.py b/backends/samsung/_passes/customized_constant_prop.py new file mode 100644 index 00000000000..fa5bad3a056 --- /dev/null +++ b/backends/samsung/_passes/customized_constant_prop.py @@ -0,0 +1,40 @@ +# 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. + +import executorch.exir.passes.constant_prop_pass as constant_prop_module +from executorch.exir import ExportedProgram +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch.fx import GraphModule + + +class _constant_prop_context: + def __init__(self): + self.backup = constant_prop_module._DEFAULT_SKIP_TARGETS + + def __enter__(self): + constant_prop_module._DEFAULT_SKIP_TARGETS = ( + constant_prop_module._DEFAULT_SKIP_TARGETS_NO_QUANT + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + constant_prop_module._DEFAULT_SKIP_TARGETS = self.backup + + +class ConstantPropPass(ExportPass): + """ + Official constant_prop_pass will not fold Q-DQ + But we need to fold quantized constant tensor as well as non-quantized one + """ + + def __init__(self, edge_program: ExportedProgram): + super().__init__() + self.edge_program = edge_program + + def call(self, graph_module: GraphModule): + with _constant_prop_context(): + _ = constant_prop_pass(self.edge_program) + return PassResult(graph_module, True) diff --git a/backends/samsung/_passes/replace_scalar_ops.py b/backends/samsung/_passes/replace_scalar_ops.py new file mode 100644 index 00000000000..8ae54b0dc98 --- /dev/null +++ b/backends/samsung/_passes/replace_scalar_ops.py @@ -0,0 +1,46 @@ +# 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. + +from typing import Dict, Tuple + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from torch._export.pass_base import Argument +from torch._export.pass_infra.node_metadata import NodeMetadata +from torch._export.pass_infra.proxy_value import ProxyValue + + +class ReplaceOpsWithScalar(ExportPass): + # Replace binary ops with scalar into binary ops with tensor. + # Ops list below. + _ops_with_scalar = { + exir_ops.edge.aten.add.Scalar: exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.sub.Scalar: exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.div.Scalar: exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.pow.Tensor_Scalar: exir_ops.edge.aten.pow.Tensor_Tensor, + } + + def __init__(self): + super(ReplaceOpsWithScalar, self).__init__() + + def call_operator( + self, + op, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + meta: NodeMetadata, + ) -> ProxyValue: + if op not in self._ops_with_scalar: + return super().call_operator(op, args, kwargs, meta) + + return super().call_operator( + op=self._ops_with_scalar.get(op, op), + args=(args[0], torch.tensor(args[1])), + kwargs=kwargs, + meta=meta, + ) diff --git a/backends/samsung/aot/CMakeLists.txt b/backends/samsung/aot/CMakeLists.txt new file mode 100644 index 00000000000..1a4e6286f97 --- /dev/null +++ b/backends/samsung/aot/CMakeLists.txt @@ -0,0 +1,15 @@ +# 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. + +target_sources( + PyEnnWrapperAdaptor PUBLIC PyEnnWrapperAdaptor.cpp PyEnnWrapperAdaptor.h +) +target_sources( + PyGraphWrapperAdaptor + PUBLIC PyGraphWrapperAdaptor.cpp PyGraphWrapperAdaptor.h + wrappers/op_param_wrapper.h wrappers/op_wrapper.h + wrappers/tensor_wrapper.h +) diff --git a/backends/samsung/aot/PyEnnWrapperAdaptor.cpp b/backends/samsung/aot/PyEnnWrapperAdaptor.cpp new file mode 100644 index 00000000000..a44e47ac40b --- /dev/null +++ b/backends/samsung/aot/PyEnnWrapperAdaptor.cpp @@ -0,0 +1,29 @@ +/* + * 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 + +namespace torch { +namespace executor { +namespace enn { +PYBIND11_MODULE(PyEnnWrapperAdaptor, m) { + pybind11::class_>(m, "EnnWrapper") + .def(pybind11::init()) + .def("Init", &PyEnnWrapper::Init) + .def("IsNodeSupportedByBackend", &PyEnnWrapper::IsNodeSupportedByBackend) + .def( + "Compile", + &PyEnnWrapper::Compile, + "Ahead of time compilation for serialized graph.") + .def("Destroy", &PyEnnWrapper::Destroy, "Release resources."); +} +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/PyEnnWrapperAdaptor.h b/backends/samsung/aot/PyEnnWrapperAdaptor.h new file mode 100644 index 00000000000..953ec050174 --- /dev/null +++ b/backends/samsung/aot/PyEnnWrapperAdaptor.h @@ -0,0 +1,92 @@ +/* + * 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 +#include + +#include +#include + +#include +#include +#include + +namespace py = pybind11; + +namespace torch { +namespace executor { +namespace enn { + +class PyEnnWrapper { + public: + PyEnnWrapper() {} + + void Init(const py::bytes& compile_opts) { + graphgen_instance_ = graphgen_create(); + option_buf_ = enn_option::GetEnnExecuTorchOptions( + compile_opts.cast().data()); + } + + bool IsNodeSupportedByBackend() { + return false; + } + + py::array_t Compile(const py::array_t& model_buffer) { + if (graphgen_instance_ == nullptr) { + ENN_LOG_ERROR("Please call `Init()` first before compile."); + return py::array_t(); + } + auto soc_name = option_buf_->chipset(); + if (graphgen_initialize_context(graphgen_instance_, soc_name) != + GraphGenResult::SUCCESS) { + ENN_LOG_ERROR( + "Unsupported Soc (%d), please check your chipset version.", soc_name); + return py::array_t(); + } + + auto m_buf_info = model_buffer.request(); + auto* model_buf_ptr = reinterpret_cast(m_buf_info.ptr); + NNCBuffer* nnc_buffer = nullptr; + if (graphgen_generate( + graphgen_instance_, model_buf_ptr, m_buf_info.size, &nnc_buffer) != + GraphGenResult::SUCCESS) { + ENN_LOG_ERROR("Compile model failed."); + return py::array_t(); + } + + auto result = py::array_t({nnc_buffer->size}, {sizeof(char)}); + auto result_buf = result.request(); + memcpy(result_buf.ptr, nnc_buffer->addr, nnc_buffer->size); + + graphgen_release_buffer(graphgen_instance_, nnc_buffer); + + return result; + } + + void Destroy() { + graphgen_release(graphgen_instance_); + graphgen_instance_ = nullptr; + } + + ~PyEnnWrapper() { + Destroy(); + } + + private: + // pointer to enn software entry + void* graphgen_instance_ = nullptr; + // enn compilation option buf + const enn_option::EnnExecuTorchOptions* option_buf_ = nullptr; +}; +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/PyGraphWrapperAdaptor.cpp b/backends/samsung/aot/PyGraphWrapperAdaptor.cpp new file mode 100644 index 00000000000..bfde8cfe9bf --- /dev/null +++ b/backends/samsung/aot/PyGraphWrapperAdaptor.cpp @@ -0,0 +1,92 @@ +/* + * 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 "PyGraphWrapperAdaptor.h" + +namespace torch { +namespace executor { +namespace enn { + +PYBIND11_MODULE(PyGraphWrapperAdaptor, m) { + pybind11::class_>( + m, "OpParamWrapper") + .def(pybind11::init()) + .def("SetStringValue", &OpParamWrapper::SetStringValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetScalarValue", &OpParamWrapper::SetScalarValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue) + .def("SetVectorValue", &OpParamWrapper::SetVectorValue); + + pybind11::class_>( + m, "PyEnnTensorWrapper") + .def(pybind11::init< + std::string, + const std::vector&, + std::string, + std::string>()) + .def( + "AddQuantizeParam", + &EnnTensorWrapper::AddQuantizeParam, + "Add quantize parameter.") + .def( + "AddData", + &EnnTensorWrapper::AddData, + "Add data for constant tensor."); + + pybind11::class_>( + m, "PyEnnOpWrapper") + .def(pybind11::init< + std::string, + std::string, + const std::vector&, + const std::vector&>()) + .def( + "AddOpParam", + &EnnOpWrapper::AddOpParam, + "Add parameter for current op."); + + pybind11::class_>( + m, "PyEnnGraphWrapper") + .def(pybind11::init()) + .def("Init", &PyEnnGraphWrapper::Init, "Initialize Graph Wrapper.") + .def( + "DefineTensor", + &PyEnnGraphWrapper::DefineTensor, + "Define a tensor in graph.") + .def( + "DefineOpNode", + &PyEnnGraphWrapper::DefineOpNode, + "Define a op node in graph.") + .def( + "SetGraphInputTensors", + &PyEnnGraphWrapper::SetGraphInputTensors, + "Set inputs for Graph") + .def( + "SetGraphOutputTensors", + &PyEnnGraphWrapper::SetGraphOutputTensors, + "Set outputs for Graph") + .def( + "FinishBuild", + &PyEnnGraphWrapper::FinishBuild, + "Finish to build the graph.") + .def("Serialize", &PyEnnGraphWrapper::Serialize, "Serialize the graph."); +} + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/PyGraphWrapperAdaptor.h b/backends/samsung/aot/PyGraphWrapperAdaptor.h new file mode 100644 index 00000000000..baa082f4d69 --- /dev/null +++ b/backends/samsung/aot/PyGraphWrapperAdaptor.h @@ -0,0 +1,154 @@ +/* + * 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 +#include +#include +#include +#include + +#include "wrappers/op_param_wrapper.h" +#include "wrappers/op_wrapper.h" +#include "wrappers/tensor_wrapper.h" + +namespace py = pybind11; + +namespace torch { +namespace executor { +namespace enn { + +class PyEnnGraphWrapper { + public: + PyEnnGraphWrapper() {} + + void Init() { + graph_wrapper_ = create_graph(""); + } + + TENSOR_ID_T DefineTensor(std::shared_ptr tensor) const { + TENSOR_ID_T tensor_id; + auto result = define_tensor( + graph_wrapper_, + &tensor_id, + tensor->GetName().c_str(), + tensor->GetShape().data(), + tensor->GetShape().size(), + tensor->GetDataType().c_str(), + tensor->GetLayout().c_str()); + if (result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in define tensor"); + } + + if (tensor->HasConstantData()) { + auto set_data_result = set_data_for_constant_tensor( + graph_wrapper_, + tensor_id, + tensor->GetDataRawPtr(), + tensor->GetDataBytes()); + if (set_data_result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in define tensor"); + } + } + + auto* quantize_param = tensor->GetQuantizeParam(); + if (quantize_param != nullptr) { + auto set_qparam_result = set_quantize_param_for_tensor( + graph_wrapper_, + tensor_id, + quantize_param->GetQuantizeType().c_str(), + quantize_param->GetScales(), + quantize_param->GetZeroPoints()); + if (set_qparam_result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in define tensor"); + } + } + + return tensor_id; + } + + NODE_ID_T DefineOpNode(std::shared_ptr op) const { + NODE_ID_T op_id; + + auto result = define_op_node( + graph_wrapper_, + &op_id, + op->GetName().c_str(), + op->GetType().c_str(), + op->GetInputs().data(), + op->GetInputs().size(), + op->GetOutputs().data(), + op->GetOutputs().size()); + if (result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in define op"); + } + + for (const auto& param : op->GetParams()) { + add_op_parameter( + graph_wrapper_, op_id, param->getKeyName().c_str(), param->Dump()); + } + + return op_id; + } + + void SetGraphInputTensors(const std::vector& tensors) const { + auto result = + set_graph_input_tensors(graph_wrapper_, tensors.data(), tensors.size()); + if (result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in set graph inputs"); + } + } + + void SetGraphOutputTensors(const std::vector& tensors) const { + auto result = set_graph_output_tensors( + graph_wrapper_, tensors.data(), tensors.size()); + if (result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail in set graph outputs"); + } + } + + void FinishBuild() const { + auto result = finish_build_graph(graph_wrapper_); + + if (result != GraphWrapperReturn::SUCCESS) { + throw std::runtime_error("fail to build graph"); + } + } + + py::array_t Serialize() { + uint64_t nbytes = 0; + uint8_t* addr = nullptr; + auto result = serialize(graph_wrapper_, &addr, &nbytes); + + if (result != GraphWrapperReturn::SUCCESS || addr == nullptr || + nbytes == 0) { + throw std::runtime_error("fail to serialize"); + } + + auto serial_buf = py::array_t(nbytes); + auto serial_buf_block = serial_buf.request(); + char* serial_buf_ptr = (char*)serial_buf_block.ptr; + std::memcpy(serial_buf_ptr, addr, nbytes); + + return serial_buf; + } + + ~PyEnnGraphWrapper() { + release_graph(graph_wrapper_); + } + + private: + GraphHandler graph_wrapper_; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/wrappers/op_param_wrapper.h b/backends/samsung/aot/wrappers/op_param_wrapper.h new file mode 100644 index 00000000000..143bb80e478 --- /dev/null +++ b/backends/samsung/aot/wrappers/op_param_wrapper.h @@ -0,0 +1,123 @@ +/* + * 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 +#include +#include +#include + +#include + +namespace torch { +namespace executor { +namespace enn { + +template +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::UNKNOWN; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::UINT64; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::INT64; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::UINT32; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::INT32; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::FLOAT32; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::FLOAT64; +}; + +template <> +struct ScalarTypeCast { + constexpr static ScalarType value = ScalarType::BOOL; +}; + +class OpParamWrapper { + public: + OpParamWrapper(std::string key) : key_name_(std::move(key)) {} + + ~OpParamWrapper() = default; + + std::string getKeyName() const { + return key_name_; + } + + template + void SetScalarValue(T value) { + auto bytes = sizeof(T); + storage_ = std::unique_ptr(new uint8_t[bytes]); + memcpy(storage_.get(), &value, bytes); + size_ = 1; + is_scalar_ = true; + scalar_type_ = ScalarTypeCast::value; + } + + template + void SetVectorValue(const std::vector& value) { + auto bytes = sizeof(T) * value.size(); + storage_ = std::unique_ptr(new uint8_t[bytes]); + memcpy(storage_.get(), value.data(), bytes); + size_ = value.size(); + is_scalar_ = false; + scalar_type_ = ScalarTypeCast::value; + } + + void SetStringValue(const std::string& value) { + auto bytes = sizeof(std::string::value_type) * value.size(); + storage_ = std::unique_ptr(new uint8_t[bytes]); + memcpy(storage_.get(), value.data(), bytes); + size_ = value.size(); + is_scalar_ = false; + scalar_type_ = ScalarType::CHAR; + } + + ParamWrapper Dump() const { + ParamWrapper param; + param.data = storage_.get(); + param.size = size_; + param.is_scalar = is_scalar_; + param.type = scalar_type_; + + return param; + } + + private: + std::string key_name_; + std::unique_ptr storage_ = nullptr; + uint32_t size_ = 0; + bool is_scalar_ = false; + ScalarType scalar_type_ = ScalarType::UNKNOWN; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/wrappers/op_wrapper.h b/backends/samsung/aot/wrappers/op_wrapper.h new file mode 100644 index 00000000000..d5c1956affe --- /dev/null +++ b/backends/samsung/aot/wrappers/op_wrapper.h @@ -0,0 +1,67 @@ +/* + * 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 +#include + +#include +namespace torch { +namespace executor { +namespace enn { + +class EnnOpWrapper { + public: + EnnOpWrapper( + std::string op_name, + std::string op_type, + const std::vector& input_tensors_id, + const std::vector& output_tensors_id) + : name_(std::move(op_name)), + op_type_(std::move(op_type)), + input_tensors_(input_tensors_id), + output_tensors_(output_tensors_id) {} + + void AddOpParam(std::shared_ptr param) { + params_.emplace_back(std::move(param)); + } + + const std::string& GetName() const { + return name_; + } + + const std::string GetType() const { + return op_type_; + } + + const std::vector& GetInputs() const { + return input_tensors_; + } + + const std::vector& GetOutputs() const { + return output_tensors_; + } + + const std::vector>& GetParams() const { + return params_; + } + + private: + std::string name_; + std::string op_type_; + std::vector input_tensors_; + std::vector output_tensors_; + std::vector> params_; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/wrappers/quantize_param_wrapper.h b/backends/samsung/aot/wrappers/quantize_param_wrapper.h new file mode 100644 index 00000000000..7b9ae473b28 --- /dev/null +++ b/backends/samsung/aot/wrappers/quantize_param_wrapper.h @@ -0,0 +1,63 @@ +/* + * 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 +#include + +#include +namespace torch { +namespace executor { +namespace enn { + +class EnnQuantizeParamWrapper { + public: + EnnQuantizeParamWrapper( + std::string quantize_dtype, + const std::vector& scales, + const std::vector& zero_points) + : quantize_type_(std::move(quantize_dtype)), + scales_(scales), + zero_points_(zero_points) {} + + const std::string& GetQuantizeType() const { + return quantize_type_; + } + + ParamWrapper GetScales() const { + ParamWrapper param; + param.data = const_cast(scales_.data()); + param.size = scales_.size(); + param.is_scalar = false; + param.type = ScalarType::FLOAT32; + + return param; + } + + ParamWrapper GetZeroPoints() const { + ParamWrapper param; + param.data = const_cast(zero_points_.data()); + param.size = zero_points_.size(); + param.is_scalar = false; + param.type = ScalarType::INT32; + + return param; + } + + private: + std::string quantize_type_; + std::vector scales_; + std::vector zero_points_; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/aot/wrappers/tensor_wrapper.h b/backends/samsung/aot/wrappers/tensor_wrapper.h new file mode 100644 index 00000000000..0b50c839e04 --- /dev/null +++ b/backends/samsung/aot/wrappers/tensor_wrapper.h @@ -0,0 +1,92 @@ +/* + * 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 "quantize_param_wrapper.h" + +namespace py = pybind11; + +namespace torch { +namespace executor { +namespace enn { + +class EnnTensorWrapper { + public: + EnnTensorWrapper( + std::string tensor_name, + const std::vector& dims, + std::string data_type, + std::string layout) + : name_(std::move(tensor_name)), + shape_(dims), + data_type_(data_type), + layout_(std::move(layout)) {} + + void AddQuantizeParam( + std::string quantize_dtype, + const std::vector& scales, + const std::vector& zero_points) { + quantize_param_ = std::make_unique( + quantize_dtype, scales, zero_points); + } + + void AddData(py::array& data) { + data_bytes_ = data.nbytes(); + + if (data.data() == nullptr || data_bytes_ == 0) { + return; + } + data_ = std::unique_ptr(new uint8_t[data_bytes_]); + memcpy(data_.get(), data.data(), data_bytes_); + } + + const std::string& GetName() const { + return name_; + } + const std::vector& GetShape() const { + return shape_; + } + const std::string& GetDataType() const { + return data_type_; + } + const std::string& GetLayout() const { + return layout_; + } + const EnnQuantizeParamWrapper* GetQuantizeParam() const { + return quantize_param_.get(); + } + + bool HasConstantData() const { + return data_ != nullptr && data_bytes_ != 0; + } + + const uint8_t* GetDataRawPtr() const { + return data_.get(); + } + + uint32_t GetDataBytes() const { + return data_bytes_; + } + + private: + std::string name_; + std::vector shape_; + std::string data_type_; + std::string layout_; + std::unique_ptr quantize_param_ = nullptr; + std::unique_ptr data_ = nullptr; + uint32_t data_bytes_ = 0; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh new file mode 100755 index 00000000000..dfa6407ff50 --- /dev/null +++ b/backends/samsung/build.sh @@ -0,0 +1,133 @@ +#!/bin/bash +## Directory Info +BASE_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" +PROJECT_DIR=$(realpath ${BASE_DIR}/../../) + +X86_64_BUILD_DIR=${PROJECT_DIR}/build_samsung +ANDROID_BUILD_DIR=${PROJECT_DIR}/build_samsung_android +BUILD_ARCH=all +CLEAN_BUILD_DIR=false + + +EXYNOS_AI_LITECORE_ROOT=${EXYNOS_AI_LITECORE_ROOT} +ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT} + +function usage() { + echo "Usage build.sh + + + + --sdk The path of downloaded ENN SDK, which is required for building. + Or export EXYNOS_AI_LITECORE_ROOT=/path/to/xxx + --ndk The path of Android NDK, or export ANDROID_NDK_ROOT=/path/to/ndk. + + --build, -b [x86_64, android, all] Default is all, x86_64 target to offline compilation, + android target to online execution. + --clean, -c Clean the build cache. + --help, -h Print the usage information. + " +} + +function build_x86_64() { + if [[ -z ${EXYNOS_AI_LITECORE_ROOT} ]]; then + echo "Please export EXYNOS_AI_LITECORE_ROOT or set by command" + exit 1 + fi + + echo "EXYNOS_AI_LITECORE_ROOT: ${EXYNOS_AI_LITECORE_ROOT}" + echo "ANDROID_NDK_ROOT: ${ANDROID_NDK_ROOT}" + + cmake \ + -DCMAKE_INSTALL_PREFIX=${X86_64_BUILD_DIR} \ + -DEXYNOS_AI_LITECORE_ROOT=${EXYNOS_AI_LITECORE_ROOT} \ + -DEXECUTORCH_BUILD_ENN=ON \ + -DEXECUTORCH_BUILD_DEVTOOLS=ON \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -S ${PROJECT_DIR} \ + -B ${X86_64_BUILD_DIR} + + cmake --build ${X86_64_BUILD_DIR} -j $(nproc) --target install + + rm -f ${PROJECT_DIR}/backends/samsung/python/Py*so + cp -fv ${X86_64_BUILD_DIR}/backends/samsung/Py*so ${PROJECT_DIR}/backends/samsung/python/ + cp -fv ${PROJECT_DIR}/schema/*.fbs ${PROJECT_DIR}/exir/_serialize/ +} + +function build_android() { + if [[ -z ${EXYNOS_AI_LITECORE_ROOT} ]]; then + echo "Please export EXYNOS_AI_LITECORE_ROOT or set by command" + exit 1 + fi + if [[ -z ${ANDROID_NDK_ROOT} ]]; then + echo "Please export ANDROID_NDK_ROOT or set by command" + exit 1 + fi + + ANDROID_ABI=arm64-v8a + cmake \ + -DCMAKE_INSTALL_PREFIX=${ANDROID_BUILD_DIR} \ + -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake" \ + -DANDROID_ABI="${ANDROID_ABI}" \ + -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_ENN=ON \ + -DEXYNOS_AI_LITECORE_ROOT=${EXYNOS_AI_LITECORE_ROOT} \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_ENABLE_LOGGING=1 \ + -DEXECUTORCH_BUILD_DEVTOOLS=ON \ + -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ + -S ${PROJECT_DIR} \ + -B ${ANDROID_BUILD_DIR} + + cmake --build ${ANDROID_BUILD_DIR} -j $(nproc) --target install +} + +# Main +for (( i=1; i<=$#; i++)) +do + case "${!i}" in + "--sdk") + let i++ + EXYNOS_AI_LITECORE_ROOT="${!i}" + ;; + "--ndk") + let i++ + ANDROID_NDK_ROOT="${!i}" + ;; + "--clean"|"-c") + CLEAN_BUILD_DIR=true + ;; + "--build"|"-b") + let i++ + BUILD_ARCH="${!i}" + ;; + "--help"|"-h") + usage + exit 0 + ;; + *) + echo "Unknown option: ${!i}" + usage + exit 1 + ;; + esac +done + +cd ${PROJECT_DIR} +if [ "${CLEAN_BUILD_DIR}" = true ]; then + rm -rf ${X86_64_BUILD_DIR} + rm -rf ${ANDROID_BUILD_DIR} + exit 0 +fi + +if [[ "${BUILD_ARCH}" = "all" || "${BUILD_ARCH}" = "x86_64" ]]; then + build_x86_64 +fi +if [[ "${BUILD_ARCH}" = "all" || "${BUILD_ARCH}" = "android" ]]; then + build_android +fi diff --git a/backends/samsung/builders/__init__.py b/backends/samsung/builders/__init__.py new file mode 100644 index 00000000000..02a457fd06e --- /dev/null +++ b/backends/samsung/builders/__init__.py @@ -0,0 +1,91 @@ +# 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. + +from . import ( + node_visitor, + op_add, + op_avg_pool2d, + op_batch_norm, + op_bmm, + op_cat, + op_clamp, + op_constant_pad_nd, + op_conv2d, + op_div, + op_embedding, + op_expand_copy, + op_gelu, + op_getitem, + op_hardswish, + op_hardtanh, + op_layer_norm, + op_leaky_relu, + op_linear, + op_log_softmax, + op_max_pool2d, + op_maximum, + op_mean_dim, + op_minimum, + op_mul, + op_permute, + op_pixel_shuffle, + op_relu, + op_reshape, + op_rsqrt, + op_select, + op_slice_copy, + op_softmax, + op_sqrt, + op_squeeze, + op_sub, + op_to_copy, + op_unsqueeze, + op_upsample_bilinear2d, + op_upsample_nearest2d, +) + +__all__ = [ + node_visitor, + op_add, + op_avg_pool2d, + op_batch_norm, + op_bmm, + op_cat, + op_clamp, + op_conv2d, + op_constant_pad_nd, + op_div, + op_embedding, + op_expand_copy, + op_gelu, + op_getitem, + op_hardswish, + op_hardtanh, + op_layer_norm, + op_leaky_relu, + op_linear, + op_log_softmax, + op_max_pool2d, + op_maximum, + op_mean_dim, + op_minimum, + op_mul, + op_permute, + op_pixel_shuffle, + op_relu, + op_reshape, + op_rsqrt, + op_select, + op_slice_copy, + op_softmax, + op_sqrt, + op_squeeze, + op_sub, + op_to_copy, + op_unsqueeze, + op_upsample_bilinear2d, + op_upsample_nearest2d, +] diff --git a/backends/samsung/builders/node_visitor.py b/backends/samsung/builders/node_visitor.py new file mode 100644 index 00000000000..a35c0b4715d --- /dev/null +++ b/backends/samsung/builders/node_visitor.py @@ -0,0 +1,106 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# 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. + +from typing import Dict, Optional + +import torch +from executorch.backends.samsung.builders.utils import ( + get_map_dtype, + get_tensor, + get_tensor_type, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms.utils import is_param_node +from torch.export import ExportedProgram + + +class NodeVisitor: + """ + Node visitor pattern for visiting nodes in an edge IR graph + """ + + def __init__(self, exported_program: ExportedProgram) -> None: + self._exported_program = exported_program or None + + @property + def exported_program(self) -> ExportedProgram: + return self._exported_program + + def define_node(self, node: torch.fx.Node, enn_graph: EnnGraph): + raise NotImplementedError("NodeVisitor must be extended!") + + def define_tensor( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.fx.Node, int], + swap_nc_for_weights: bool = False, + output_idx: Optional[int] = None, + ) -> int: + if node in vals_to_ids and (output_idx is None or output_idx == 0): + return vals_to_ids[node] + + # Get tensor basic information + tensor = get_tensor(self.exported_program, node) + + if output_idx is not None: + tensor = tensor[output_idx] + + tensor_type = get_tensor_type(self.exported_program, node) + data_type = get_map_dtype(tensor.dtype) + + const_data = None + 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() + + dims = [1] if len(tensor.size()) == 0 else list(tensor.size()) + + enn_tensor_id = enn_graph.define_tensor( + node.name, + dims, + data_type, + tensor_type.name, + const_data, + ) + assert enn_tensor_id is not None + vals_to_ids[node] = enn_tensor_id + + return enn_tensor_id + + +_node_visitor_dict = {} + + +def register_node_visitor(visitor): + assert ( + isinstance(visitor, type) + and issubclass(visitor, NodeVisitor) + and hasattr(visitor, "target") + ), f"Illformed NodeVisitor subclass, can't register!, got: {visitor}" + if isinstance(visitor.target, str): + _node_visitor_dict[visitor.target] = visitor + elif isinstance(visitor.target, (list, tuple)): + for target in visitor.target: + _node_visitor_dict[target] = visitor + else: + raise TypeError( + f"target of vistor should be str|Tuple[str]|List[str], not{type(visitor.target)}" + ) + + +def get_node_visitors(*args) -> Dict[str, NodeVisitor]: + node_visitors = {} + """ + Create a new class instance at runtime, and put them in a dict + """ + for target, visitor in _node_visitor_dict.items(): + assert callable(visitor), "Expecting a callable class, " + f"but got {visitor} of type {type(visitor)}" + node_visitors[target] = visitor(*args) + return node_visitors diff --git a/backends/samsung/builders/op_add.py b/backends/samsung/builders/op_add.py new file mode 100644 index 00000000000..1b0dddb0d02 --- /dev/null +++ b/backends/samsung/builders/op_add.py @@ -0,0 +1,36 @@ +# 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. +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 AddVisitor(NodeVisitor): + target = "aten.add.Tensor" + + 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], + ) -> None: + 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) + + 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]) diff --git a/backends/samsung/builders/op_avg_pool2d.py b/backends/samsung/builders/op_avg_pool2d.py new file mode 100644 index 00000000000..ad7ccbac3ae --- /dev/null +++ b/backends/samsung/builders/op_avg_pool2d.py @@ -0,0 +1,70 @@ +# 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. +from typing import cast, Dict, List + +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 AvgPool2dVisitor(NodeVisitor): + target = "aten.avg_pool2d.default" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + kernel_size = cast(List[int], node.args[1]) + if len(kernel_size) == 1: + kernel_size = kernel_size * 2 + + stride = cast(List[int], node.args[2]) if len(node.args) > 2 else kernel_size + if len(stride) == 1: + stride = stride * 2 + + padding = cast(List[int], node.args[3]) if len(node.args) > 3 else [0, 0] + if len(padding) == 1: + padding = padding * 2 + explicit_padding = [padding[0], padding[1], padding[0], padding[1]] + + params = {} + params["kernel_h"] = kernel_size[0] + params["kernel_w"] = kernel_size[1] + params["stride_h"] = stride[0] + params["stride_w"] = stride[1] + params["padding"] = "EXPLICIT" + params["explicit_padding"] = explicit_padding + + 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: + params["count_include_pad"] = True + + if len(node.args) > 6: + divisor_override = cast(int, node.args[6]) + assert ( + divisor_override == kernel_size[0] * kernel_size[1] + ), "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) diff --git a/backends/samsung/builders/op_batch_norm.py b/backends/samsung/builders/op_batch_norm.py new file mode 100644 index 00000000000..e5373a8223a --- /dev/null +++ b/backends/samsung/builders/op_batch_norm.py @@ -0,0 +1,56 @@ +# 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. +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 BatchNormVisitor(NodeVisitor): + target = "aten._native_batch_norm_legit_no_training.default" + + 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], + ) -> None: + all_input_tensors = [] + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + all_input_tensors.append(input_id) + + weight_node, bias_node, mean_node, var_node = ( + node.args[1], + node.args[2], + node.args[3], + node.args[4], + ) + weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) + all_input_tensors.append(weight_id) + bias_id = self.define_tensor(bias_node, enn_graph, vals_to_ids) + all_input_tensors.append(bias_id) + mean_id = self.define_tensor(mean_node, enn_graph, vals_to_ids) + all_input_tensors.append(mean_id) + var_id = self.define_tensor(var_node, enn_graph, vals_to_ids) + all_input_tensors.append(var_id) + + eps = node.args[-1] + params = {"epsilon": eps} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids, output_idx=0) + + enn_graph.define_op( + node.name, "BatchNormalization", all_input_tensors, [output_id], params + ) diff --git a/backends/samsung/builders/op_bmm.py b/backends/samsung/builders/op_bmm.py new file mode 100644 index 00000000000..6ba8864ebb3 --- /dev/null +++ b/backends/samsung/builders/op_bmm.py @@ -0,0 +1,40 @@ +# 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. + +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 BMMVisitor(NodeVisitor): + target = "aten.bmm.default" + + 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], + ) -> None: + 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) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op( + node.name, "BATCH_MATMUL", [input_id_1, input_id_2], [output_id] + ) diff --git a/backends/samsung/builders/op_cat.py b/backends/samsung/builders/op_cat.py new file mode 100644 index 00000000000..e9c0a32b389 --- /dev/null +++ b/backends/samsung/builders/op_cat.py @@ -0,0 +1,42 @@ +# 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. +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class CatVisitor(NodeVisitor): + target = "aten.cat.default" + + def __init__(self, *args): + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> None: + tensors = cast(List[torch.fx.Node], node.args[0]) + input_tensor_ids = [] + + for in_tensor in tensors: + input_id = self.define_tensor(in_tensor, enn_graph, vals_to_ids) + input_tensor_ids.append(input_id) + + in_shape = get_shape(node) + axis = cast(int, node.args[1]) % len(in_shape) if len(node.args) >= 2 else 0 + params = {"axis": axis} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + enn_graph.define_op(node.name, "CONCAT", input_tensor_ids, [output_id], params) diff --git a/backends/samsung/builders/op_clamp.py b/backends/samsung/builders/op_clamp.py new file mode 100644 index 00000000000..c5670b80fa3 --- /dev/null +++ b/backends/samsung/builders/op_clamp.py @@ -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. +from typing import cast, 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 ClampVisitor(NodeVisitor): + target = "aten.clamp.default" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # The default value of lower bound and upper bound + output_min = torch.finfo(torch.float32).min + output_max = torch.finfo(torch.float32).max + if node.args[1] is not None: + output_min = cast(float, node.args[1]) + if len(node.args) > 2 and node.args[2] is not None: + output_max = cast(float, node.args[2]) + + params = {"minimum": output_min, "maximum": output_max} + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_constant_pad_nd.py b/backends/samsung/builders/op_constant_pad_nd.py new file mode 100644 index 00000000000..cc7cdc5751b --- /dev/null +++ b/backends/samsung/builders/op_constant_pad_nd.py @@ -0,0 +1,56 @@ +# 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. + +from typing import cast, Dict, List + +import numpy as np + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class ConstantPadNDVisitor(NodeVisitor): + target = "aten.constant_pad_nd.default" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # torch padding order starts from the last axis, change the order to fit samsung lite-core + paddings = np.reshape(cast(List[int], node.args[1]), (-1, 2))[::-1].astype( + np.uint32 + ) + in_shape = get_shape(input) + paddings = paddings.reshape(-1).tolist() + paddings = [0] * (2 * len(in_shape) - len(paddings)) + paddings + paddings = paddings[::2] + paddings[1::2] + + padding_value = node.args[2] + assert padding_value == 0.0, "Only Support pad constant 0 now." + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + params = { + "explicit_padding": paddings, + "padding": "EXPLICIT", + "padding_type": "CONSTANT", + } + + enn_graph.define_op(node.name, "PAD", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_conv2d.py b/backends/samsung/builders/op_conv2d.py new file mode 100644 index 00000000000..881a533801f --- /dev/null +++ b/backends/samsung/builders/op_conv2d.py @@ -0,0 +1,85 @@ +# 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. +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class Conv2dVisitor(NodeVisitor): + target = "aten.convolution.default" + + 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], + ) -> None: + all_input_tensors = [] + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + all_input_tensors.append(input_id) + + is_transpose_conv = cast(bool, node.args[6]) + weight_node = node.args[1] + weight_id = self.define_tensor( + weight_node, + enn_graph, + vals_to_ids, + is_transpose_conv, + ) + all_input_tensors.append(weight_id) + + if node.args[2] is not None: + bias_node = node.args[2] + bias_id = self.define_tensor(bias_node, enn_graph, vals_to_ids) + all_input_tensors.append(bias_id) + + stride = cast(List[int], node.args[3]) + padding = cast(List[int], node.args[4]) + dilation = cast(List[int], node.args[5]) + groups = cast(int, node.args[8]) + explicit_padding = [padding[0], padding[1], padding[0], padding[1]] + + input_shape = get_shape(input) + kernel_shape = get_shape(weight_node) + params = {} + params["kernel_h"] = kernel_shape[2] + params["kernel_w"] = kernel_shape[3] + params["stride_h"] = stride[0] + params["stride_w"] = stride[1] + params["dilation_h"] = dilation[0] + params["dilation_w"] = dilation[1] + params["groups"] = groups + params["padding"] = "EXPLICIT" + params["padding_type"] = "CONSTANT" # CONSTANT will be zero-padding + 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] + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + is_depthwise_conv = kernel_shape[1] == 1 and kernel_shape[0] / groups == 1 + if is_depthwise_conv: + conv_type = "DWDECONV2D" if is_transpose_conv else "DWCONV2D" + enn_graph.define_op( + node.name, conv_type, all_input_tensors, [output_id], params + ) + else: + conv_type = "DECONV2D" if is_transpose_conv else "CONV2D" + enn_graph.define_op( + node.name, conv_type, all_input_tensors, [output_id], params + ) diff --git a/backends/samsung/builders/op_div.py b/backends/samsung/builders/op_div.py new file mode 100644 index 00000000000..89d773ddb0e --- /dev/null +++ b/backends/samsung/builders/op_div.py @@ -0,0 +1,39 @@ +# 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. + +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 DivVisitor(NodeVisitor): + target = "aten.div.Tensor" + + 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], + ) -> None: + # 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) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "ELTDIV", [input_id_1, input_id_2], [output_id]) diff --git a/backends/samsung/builders/op_embedding.py b/backends/samsung/builders/op_embedding.py new file mode 100644 index 00000000000..f37c46a56d6 --- /dev/null +++ b/backends/samsung/builders/op_embedding.py @@ -0,0 +1,41 @@ +# 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. + +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 EmbeddingVisitor(NodeVisitor): + target = "aten.embedding.default" + + 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], + ) -> None: + weight_node = node.args[0] + weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) + + input = node.args[1] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + params = {"axis": 0, "input_type": "indices"} + enn_graph.define_op( + node.name, "GATHER", [input_id, weight_id], [output_id], params + ) diff --git a/backends/samsung/builders/op_expand_copy.py b/backends/samsung/builders/op_expand_copy.py new file mode 100644 index 00000000000..f4c707b8e62 --- /dev/null +++ b/backends/samsung/builders/op_expand_copy.py @@ -0,0 +1,77 @@ +# 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. + +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class ExpandVisitor(NodeVisitor): + target = "aten.expand_copy.default" + + 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], + ): + # inputs + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + in_shape = get_shape(input) + sizes = cast(List[int], node.args[1]) + expand_dims = self.check_expand_dims(sizes, in_shape) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + if len(expand_dims) == 0: + params = {"new_shape": [*node.meta["val"].shape]} + enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id], params) + elif len(expand_dims) == 1: + expand_dim = expand_dims[0] + params = {"axis": expand_dim} + enn_graph.define_op( + node.name, + "CONCAT", + [input_id] * sizes[expand_dim], + [output_id], + params, + ) + else: + raise NotImplementedError("Don't support expanding at more than one axes.") + + def check_expand_dims(self, sizes, in_shape): + expand_dims = [] + new_size_index = len(sizes) + in_shape_index = len(in_shape) + + while in_shape_index > 0 and new_size_index > 0: + in_shape_index -= 1 + new_size_index -= 1 + if ( + sizes[new_size_index] == -1 + or sizes[new_size_index] == in_shape[in_shape_index] + ): + continue + expand_dims.append(in_shape_index) + + while new_size_index > 0: + new_size_index -= 1 + assert sizes[new_size_index] == 1, "Current expand is unsupported!" + + return expand_dims diff --git a/backends/samsung/builders/op_gelu.py b/backends/samsung/builders/op_gelu.py new file mode 100644 index 00000000000..059a3b77850 --- /dev/null +++ b/backends/samsung/builders/op_gelu.py @@ -0,0 +1,34 @@ +# 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. + +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 GeluVisitor(NodeVisitor): + target = "aten.gelu.default" + + 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], + ) -> None: + 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, "GELU", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_getitem.py b/backends/samsung/builders/op_getitem.py new file mode 100644 index 00000000000..901ec73cf7d --- /dev/null +++ b/backends/samsung/builders/op_getitem.py @@ -0,0 +1,32 @@ +# 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. +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 OpGetItemVisitor(NodeVisitor): + target = "getitem" + """ + do nothing if node is getitem + """ + + 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], + ) -> None: + return diff --git a/backends/samsung/builders/op_hardswish.py b/backends/samsung/builders/op_hardswish.py new file mode 100644 index 00000000000..72a99d17b83 --- /dev/null +++ b/backends/samsung/builders/op_hardswish.py @@ -0,0 +1,35 @@ +# 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. + +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 HardSwishVisitor(NodeVisitor): + target = "aten.hardswish.default" + + 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], + ) -> None: + 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, "HARDSWISH", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_hardtanh.py b/backends/samsung/builders/op_hardtanh.py new file mode 100644 index 00000000000..4f667bf5299 --- /dev/null +++ b/backends/samsung/builders/op_hardtanh.py @@ -0,0 +1,38 @@ +# 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. +from typing import cast, 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 HardTanhVisitor(NodeVisitor): + target = "aten.hardtanh.default" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + output_min = cast(float, node.args[1]) if len(node.args) > 1 else -1 + output_max = cast(float, node.args[2]) if len(node.args) > 2 else 1 + params = {"minimum": output_min, "maximum": output_max} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_layer_norm.py b/backends/samsung/builders/op_layer_norm.py new file mode 100644 index 00000000000..e6f853178d8 --- /dev/null +++ b/backends/samsung/builders/op_layer_norm.py @@ -0,0 +1,54 @@ +# 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. + +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 + +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class LayerNormVisitor(NodeVisitor): + target = ["aten.layer_norm.default"] + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> None: + all_input_tensors = [] + input_node = node.args[0] + input_id = self.define_tensor(input_node, enn_graph, vals_to_ids) + all_input_tensors.append(input_id) + + normalized_shapes = node.args[1] + assert ( + len(normalized_shapes) == 1 + and normalized_shapes[0] == get_shape(input_node)[-1] + ), "Enn Backend only support norm at last axis." + + weight_node = node.args[2] + weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) + all_input_tensors.append(weight_id) + bias_node = node.args[3] + bias_id = self.define_tensor(bias_node, enn_graph, vals_to_ids) + all_input_tensors.append(bias_id) + + epsilon = node.args[4] if len(node.args) > 4 else 1e-5 + params = {"epsilon": epsilon} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op( + node.name, "LAYERNORM", all_input_tensors, [output_id], params + ) diff --git a/backends/samsung/builders/op_leaky_relu.py b/backends/samsung/builders/op_leaky_relu.py new file mode 100644 index 00000000000..c7ed37d12e5 --- /dev/null +++ b/backends/samsung/builders/op_leaky_relu.py @@ -0,0 +1,58 @@ +# 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. + +from typing import cast, Dict + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.builders.utils import get_map_dtype +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph + + +@register_node_visitor +class LeakyReluVisitor(NodeVisitor): + target = ["aten.leaky_relu.default", "aten.prelu.default"] + + 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], + ) -> None: + all_input_tensors = [] + input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) + all_input_tensors.append(input_id) + + if node.target.__name__ == "aten.prelu.default": + negative_slope = node.args[1] + negative_slope_id = self.define_tensor( + negative_slope, enn_graph, vals_to_ids + ) + else: + negative_slope = cast(float, node.args[1]) if len(node.args) > 1 else 0.01 + negative_slope_tensor = torch.tensor(negative_slope).to(torch.float32) + negative_slope_node_name = node.name + "_slope" + dims = list(negative_slope_tensor.size()) + data_type = get_map_dtype(negative_slope_tensor.dtype) + negative_slope_id = enn_graph.define_tensor( + negative_slope_node_name, + dims, + data_type, + "CONSTANT", + negative_slope_tensor.detach().numpy(), + ) + + all_input_tensors.append(negative_slope_id) + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "PRELU", all_input_tensors, [output_id]) diff --git a/backends/samsung/builders/op_linear.py b/backends/samsung/builders/op_linear.py new file mode 100644 index 00000000000..2f7aa1e6415 --- /dev/null +++ b/backends/samsung/builders/op_linear.py @@ -0,0 +1,49 @@ +# 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. +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 +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class LinearVisitor(NodeVisitor): + target = "aten.linear.default" + + 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], + ) -> None: + all_input_tensors = [] + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + all_input_tensors.append(input_id) + + weight_node = node.args[1] + weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) + all_input_tensors.append(weight_id) + + if len(node.args) > 2 and node.args[2] is not None: + bias_node = node.args[2] + bias_id = self.define_tensor(bias_node, enn_graph, vals_to_ids) + all_input_tensors.append(bias_id) + + weight_shape = get_shape(weight_node) + params = {"in_channels": weight_shape[1], "out_channels": weight_shape[0]} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "FC", all_input_tensors, [output_id], params) diff --git a/backends/samsung/builders/op_log_softmax.py b/backends/samsung/builders/op_log_softmax.py new file mode 100644 index 00000000000..f2d87601cbb --- /dev/null +++ b/backends/samsung/builders/op_log_softmax.py @@ -0,0 +1,39 @@ +# 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. + +from typing import cast, 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 LogSoftmaxVisitor(NodeVisitor): + target = "aten._log_softmax.default" + + 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], + ): + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + axis = cast(int, node.args[1]) + meta_data = {"axis": axis} + + enn_graph.define_op(node.name, "LOGSOFTMAX", [input_id], [output_id], meta_data) diff --git a/backends/samsung/builders/op_max_pool2d.py b/backends/samsung/builders/op_max_pool2d.py new file mode 100644 index 00000000000..d386dd30b1a --- /dev/null +++ b/backends/samsung/builders/op_max_pool2d.py @@ -0,0 +1,95 @@ +# 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. +from typing import cast, Dict, List + +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 MaxPool2dVisitor(NodeVisitor): + target = ["aten.max_pool2d.default", "aten.max_pool2d_with_indices.default"] + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + is_indices = False + if node.target.__name__ == "aten.max_pool2d_with_indices.default": + users = list(node.users.keys()) + + for user in users: + if user.target.__name__ == "getitem": + getitem_index = user.args[1] + is_indices = True + if getitem_index != 0: + raise AssertionError( + "Expected second argument of getitem" + f" node for {node.target.__name__ } to be 0, " + f"got {getitem_index}. ENN delegate currently " + "only supports getting just the max " + "values from the op, but doesn't support" + " getting the corresponding indices." + ) + + kernel_size = cast(List[int], node.args[1]) + if len(kernel_size) == 1: + kernel_size = kernel_size * 2 + + stride = cast(List[int], node.args[2]) if len(node.args) > 2 else kernel_size + if len(stride) == 1: + stride = stride * 2 + + padding = cast(List[int], node.args[3]) if len(node.args) > 3 else [0, 0] + if len(padding) == 1: + padding = padding * 2 + explicit_padding = [padding[0], padding[1], padding[0], padding[1]] + + dilation = cast(List[int], node.args[4]) if len(node.args) > 4 else [1, 1] + if len(dilation) == 1: + dilation = dilation * 2 + + params = {} + params["kernel_h"] = kernel_size[0] + params["kernel_w"] = kernel_size[1] + params["stride_h"] = stride[0] + params["stride_w"] = stride[1] + params["padding"] = "EXPLICIT" + params["explicit_padding"] = explicit_padding + params["dilation_h"] = dilation[0] + params["dilation_w"] = dilation[1] + + 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, + enn_graph, + vals_to_ids, + ) + else: + output_id = self.define_tensor( + node, + enn_graph, + vals_to_ids, + output_idx=0, + ) + + enn_graph.define_op(node.name, "MAXPOOL2D", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_maximum.py b/backends/samsung/builders/op_maximum.py new file mode 100644 index 00000000000..d3358d736f3 --- /dev/null +++ b/backends/samsung/builders/op_maximum.py @@ -0,0 +1,37 @@ +# 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. + +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 MaximumVisitor(NodeVisitor): + target = "aten.maximum.default" + + 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], + ) -> None: + # 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) + + # output + 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]) diff --git a/backends/samsung/builders/op_mean_dim.py b/backends/samsung/builders/op_mean_dim.py new file mode 100644 index 00000000000..2f07f870ec4 --- /dev/null +++ b/backends/samsung/builders/op_mean_dim.py @@ -0,0 +1,44 @@ +# 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. +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class ReduceMeanVisitor(NodeVisitor): + target = "aten.mean.dim" + + 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], + ) -> None: + 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) + + dims = cast(List[int], node.args[1]) + reduce_axes = [] + in_shape = get_shape(input) + for dim in dims: + reduce_axes.append(dim % len(in_shape)) + reduce_axes.sort() + + keep_dim = node.args[2] if len(node.args) >= 3 else False + params = {"keep_dims": keep_dim, "axis": reduce_axes} + enn_graph.define_op(node.name, "REDUCEMEAN", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_minimum.py b/backends/samsung/builders/op_minimum.py new file mode 100644 index 00000000000..a32b462d45f --- /dev/null +++ b/backends/samsung/builders/op_minimum.py @@ -0,0 +1,39 @@ +# 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. + +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 MinimumVisitor(NodeVisitor): + target = "aten.minimum.default" + + 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], + ) -> None: + # 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) + + # output + 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]) diff --git a/backends/samsung/builders/op_mul.py b/backends/samsung/builders/op_mul.py new file mode 100644 index 00000000000..dce531ff0b0 --- /dev/null +++ b/backends/samsung/builders/op_mul.py @@ -0,0 +1,33 @@ +# Copyright (c) 2024 Samsung Electronics Co. LTD +# All rights reserved +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 MulVisitor(NodeVisitor): + target = "aten.mul.Tensor" + + 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], + ) -> None: + 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) + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "ELTMUL", [input_id_1, input_id_2], [output_id]) diff --git a/backends/samsung/builders/op_permute.py b/backends/samsung/builders/op_permute.py new file mode 100644 index 00000000000..646eac4c06a --- /dev/null +++ b/backends/samsung/builders/op_permute.py @@ -0,0 +1,37 @@ +# 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. +from typing import cast, Dict, List + +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 PermuteVisitor(NodeVisitor): + target = "aten.permute_copy.default" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + # permutation + permute_order = cast(List[int], node.args[1]) + 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) diff --git a/backends/samsung/builders/op_pixel_shuffle.py b/backends/samsung/builders/op_pixel_shuffle.py new file mode 100644 index 00000000000..28259299c81 --- /dev/null +++ b/backends/samsung/builders/op_pixel_shuffle.py @@ -0,0 +1,38 @@ +# 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. +from typing import cast, 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 PixelShuffleVisitor(NodeVisitor): + target = "aten.pixel_shuffle.default" + + 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], + ) -> None: + input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) + + scale_factor = cast(int, node.args[1]) + params = {"block_size": scale_factor, "mode": "CRD"} + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op( + node.name, "DEPTH_TO_SPACE", [input_id], [output_id], params + ) diff --git a/backends/samsung/builders/op_relu.py b/backends/samsung/builders/op_relu.py new file mode 100644 index 00000000000..ba90116be1d --- /dev/null +++ b/backends/samsung/builders/op_relu.py @@ -0,0 +1,34 @@ +# 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. +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 ReluVisitor(NodeVisitor): + target = "aten.relu.default" + + 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], + ) -> None: + 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, "RELU", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_reshape.py b/backends/samsung/builders/op_reshape.py new file mode 100644 index 00000000000..1f4e85ac059 --- /dev/null +++ b/backends/samsung/builders/op_reshape.py @@ -0,0 +1,34 @@ +# Copyright (c) 2024 Samsung Electronics Co. LTD +# All rights reserved +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 ReshapeVisitor(NodeVisitor): + target = "aten.view_copy.default" + + 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], + ) -> None: + 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] + enn_graph.define_op( + node.name, "RESHAPE", [input_id], [output_id], {"new_shape": new_shape} + ) diff --git a/backends/samsung/builders/op_rsqrt.py b/backends/samsung/builders/op_rsqrt.py new file mode 100644 index 00000000000..b3600d41ee2 --- /dev/null +++ b/backends/samsung/builders/op_rsqrt.py @@ -0,0 +1,35 @@ +# 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. + +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 LinearVisitor(NodeVisitor): + target = "aten.rsqrt.default" + + 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], + ) -> None: + 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]) diff --git a/backends/samsung/builders/op_select.py b/backends/samsung/builders/op_select.py new file mode 100644 index 00000000000..26f455b2548 --- /dev/null +++ b/backends/samsung/builders/op_select.py @@ -0,0 +1,52 @@ +# 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. +from typing import cast, 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 +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class SelectVisitor(NodeVisitor): + target = ["aten.select_copy.int", "aten.select.int"] + + 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], + ): + 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) + + in_shape = get_shape(input) + dim = cast(int, node.args[1]) % len(in_shape) + index = cast(int, node.args[2]) % in_shape[dim] + + begin = [0] * len(in_shape) + begin[dim] = index + end = in_shape + end[dim] = index + 1 + strides = [1] * len(in_shape) + + params = { + "begin": begin, + "end": end, + "strides": strides, + "shrink_axis_mask": pow(2, dim), + } + + enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_slice_copy.py b/backends/samsung/builders/op_slice_copy.py new file mode 100644 index 00000000000..0d7a23118a0 --- /dev/null +++ b/backends/samsung/builders/op_slice_copy.py @@ -0,0 +1,59 @@ +# 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. + +from typing import cast, 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 +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class SliceCopyVisitor(NodeVisitor): + target = "aten.slice_copy.Tensor" + + 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], + ): + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + in_shape = get_shape(input) + dim = cast(int, node.args[1]) + if dim < 0: + dim = dim + len(in_shape) + start_val = cast(int, node.args[2]) + if start_val < 0: + start_val = start_val + in_shape[dim] + end_val = min(cast(int, node.args[3]), in_shape[dim]) + if end_val < 0: + end_val = end_val + in_shape[dim] + + step = cast(int, node.args[4]) if len(node.args) > 4 else 1 + + begin = [0] * len(in_shape) + begin[dim] = start_val + end = in_shape + end[dim] = end_val + strides = [1] * len(in_shape) + strides[dim] = step + + params = {"begin": begin, "end": end, "strides": strides} + + enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_softmax.py b/backends/samsung/builders/op_softmax.py new file mode 100644 index 00000000000..1e2e4a378dc --- /dev/null +++ b/backends/samsung/builders/op_softmax.py @@ -0,0 +1,39 @@ +# 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. + +from typing import cast, 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 SoftmaxVisitor(NodeVisitor): + target = ["aten._softmax.default", "aten._safe_softmax.default"] + + 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], + ): + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + axis = cast(int, node.args[1]) + params = {"axis": axis} + + enn_graph.define_op(node.name, "SOFTMAX", [input_id], [output_id], params) diff --git a/backends/samsung/builders/op_sqrt.py b/backends/samsung/builders/op_sqrt.py new file mode 100644 index 00000000000..3560542a0bc --- /dev/null +++ b/backends/samsung/builders/op_sqrt.py @@ -0,0 +1,35 @@ +# 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. + +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 SqrtVisitor(NodeVisitor): + target = "aten.sqrt.default" + + 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], + ) -> None: + 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]) diff --git a/backends/samsung/builders/op_squeeze.py b/backends/samsung/builders/op_squeeze.py new file mode 100644 index 00000000000..d165a22fcb3 --- /dev/null +++ b/backends/samsung/builders/op_squeeze.py @@ -0,0 +1,36 @@ +# 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. + +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 SqueezeVisitor(NodeVisitor): + target = "aten.squeeze_copy.dims" + + 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], + ) -> None: + input = node.args[0] + input_id = self.define_tensor(input, enn_graph, vals_to_ids) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_sub.py b/backends/samsung/builders/op_sub.py new file mode 100644 index 00000000000..af2931f298e --- /dev/null +++ b/backends/samsung/builders/op_sub.py @@ -0,0 +1,39 @@ +# 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. + +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 SubVisitor(NodeVisitor): + target = "aten.sub.Tensor" + + 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], + ) -> None: + # 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) + + # output + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "SUB", [input_id_1, input_id_2], [output_id]) diff --git a/backends/samsung/builders/op_to_copy.py b/backends/samsung/builders/op_to_copy.py new file mode 100644 index 00000000000..545672ef6a3 --- /dev/null +++ b/backends/samsung/builders/op_to_copy.py @@ -0,0 +1,39 @@ +# 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. + +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 ToCopyVisitor(NodeVisitor): + target = ["aten._to_copy.default", "dim_order_ops._to_dim_order_copy.default"] + + 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], + ) -> None: + 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" + + 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, "CAST", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_unsqueeze.py b/backends/samsung/builders/op_unsqueeze.py new file mode 100644 index 00000000000..942c3307de7 --- /dev/null +++ b/backends/samsung/builders/op_unsqueeze.py @@ -0,0 +1,34 @@ +# 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. +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 UnsqueezeVisitor(NodeVisitor): + target = "aten.unsqueeze_copy.default" + + 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], + ) -> None: + 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, "RESHAPE", [input_id], [output_id]) diff --git a/backends/samsung/builders/op_upsample_bilinear2d.py b/backends/samsung/builders/op_upsample_bilinear2d.py new file mode 100644 index 00000000000..a934b2789ba --- /dev/null +++ b/backends/samsung/builders/op_upsample_bilinear2d.py @@ -0,0 +1,52 @@ +# 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. + +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class UpsampleBilinear2dVisitor(NodeVisitor): + target = "aten.upsample_bilinear2d.vec" + + 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], + ) -> None: + 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]) + scale_factor = [ + output_size[0] * 1.0 / in_shape[-2], + output_size[1] * 1.0 / in_shape[-1], + ] + + align_corners = cast(bool, node.args[2]) + if len(node.args) > 3 and node.args[3]: + scale_factor = cast(List[float], node.args[3]) + + params = { + "align_corners": align_corners, + "upsampling_factor": scale_factor, + "half_pixel_centers": True, + } + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + enn_graph.define_op( + node.name, "RESIZE_BILINEAR", [input_id], [output_id], params + ) diff --git a/backends/samsung/builders/op_upsample_nearest2d.py b/backends/samsung/builders/op_upsample_nearest2d.py new file mode 100644 index 00000000000..9859cd8f07e --- /dev/null +++ b/backends/samsung/builders/op_upsample_nearest2d.py @@ -0,0 +1,52 @@ +# 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. + +from typing import cast, Dict, List + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph +from executorch.backends.transforms import get_shape + + +@register_node_visitor +class UpsampleNearest2dVisitor(NodeVisitor): + target = "aten.upsample_nearest2d.vec" + + 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], + ) -> None: + 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]) + scale_factor = [ + output_size[0] * 1.0 / in_shape[-2], + output_size[1] * 1.0 / in_shape[-1], + ] + + if len(node.args) > 2 and node.args[2]: + scale_factor = cast(List[float], node.args[2]) + + params = { + "align_corners": False, + "upsampling_factor": scale_factor, + "half_pixel_centers": True, + } + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + enn_graph.define_op( + node.name, "RESIZE_NEAREST_NEIGHBOR", [input_id], [output_id], params + ) diff --git a/backends/samsung/builders/utils.py b/backends/samsung/builders/utils.py new file mode 100644 index 00000000000..58c84ff6d31 --- /dev/null +++ b/backends/samsung/builders/utils.py @@ -0,0 +1,65 @@ +# 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. + +from enum import Enum + +import torch +from executorch.backends.samsung.utils.utils import is_graph_input, is_graph_output +from executorch.backends.transforms.utils import get_param_tensor, is_param_node + +from torch.export import ExportedProgram + +DATA_TYPE_STR_MAPPING = { + torch.int8: "INT8", + torch.uint8: "UINT8", + torch.int16: "INT16", + torch.uint16: "UINT16", + torch.int32: "INT32", + torch.int64: "INT64", + torch.float16: "FLOAT16", + torch.float32: "FLOAT32", +} + +TORCH_TYPE_QTYPE_MAPPING = { + torch.int8: torch.qint8, + torch.uint8: torch.quint8, + torch.int32: torch.qint32, +} + + +class TensorType(Enum): + INPUT = 0 + OUTPUT = 1 + CONSTANT = 2 + FEATUREMAP = 3 + + +def get_tensor_type(exported_program: ExportedProgram, tensor: torch.fx.Node) -> str: + if is_graph_input(exported_program, tensor): + return TensorType.INPUT + elif is_graph_output(tensor): + return TensorType.OUTPUT + elif is_param_node(exported_program, tensor): + return TensorType.CONSTANT + else: + return TensorType.FEATUREMAP + + +def get_map_dtype(dtype): + if dtype not in DATA_TYPE_STR_MAPPING: + raise RuntimeError("Data type cannot be decided: ", dtype) + return DATA_TYPE_STR_MAPPING[dtype] + + +def get_tensor(exported_program: ExportedProgram, node: torch.fx.Node): + if not is_param_node(exported_program, node): + return node.meta["val"] + tensor = get_param_tensor(exported_program, node) + return tensor.contiguous() + + +def affine_type_to_str(ttype: TensorType): + return str(ttype).removeprefix("TensorType.") diff --git a/backends/samsung/enn_preprocess.py b/backends/samsung/enn_preprocess.py new file mode 100644 index 00000000000..dde01bc09c7 --- /dev/null +++ b/backends/samsung/enn_preprocess.py @@ -0,0 +1,100 @@ +# 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. + +import logging +from typing import Dict, final, List + +import executorch.backends.samsung.python.PyEnnWrapperAdaptor as PyEnnWrapper +import torch +from executorch.backends.samsung._passes.conv1d_to_conv2d import Conv1dToConv2d +from executorch.backends.samsung._passes.customized_constant_prop import ( + ConstantPropPass, +) +from executorch.backends.samsung._passes.replace_scalar_ops import ReplaceOpsWithScalar +from executorch.backends.samsung.builders.node_visitor import get_node_visitors +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.backends.transforms.addmm_mm_to_linear import AddmmToLinearTransform +from executorch.backends.transforms.fuse_batch_norm_with_conv import ( + FuseBatchNormWithConvPass, +) + +from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass + +from executorch.exir.backend.backend_details import ( + BackendDetails, + CompileSpec, + PreprocessResult, +) + +from executorch.exir.passes import PassManager + +from torch.export.exported_program import ExportedProgram + + +@final +class EnnBackend(BackendDetails): + @staticmethod + def preprocess( + edge_program: ExportedProgram, + compile_specs: List[CompileSpec], + ) -> PreprocessResult: + enn_wrapper = PyEnnWrapper.EnnWrapper() + option_spec = get_compile_spec( + compile_specs, ENN_COMPILE_OPTION_TITLE, required=True + ) + enn_wrapper.Init(option_spec.value) + + enn_preprocess_passes = PassManager( + passes=[ + ConstantPropPass(edge_program), + Conv1dToConv2d(edge_program), + FuseBatchNormWithConvPass(edge_program), + AddmmToLinearTransform(), + ReplaceOpsWithScalar(), + RemoveGetItemPass(), + ] + ) + pass_result = enn_preprocess_passes(edge_program.graph_module) + assert pass_result is not None + + enn_graph = EnnGraph() + # node visitors + node_visitors = get_node_visitors(edge_program) + + vals_to_ids: Dict[torch.fx.Node, int] = {} + for node in pass_result.graph_module.graph.nodes: + if node.op == "call_function": + logging.info(f"Visiting: {node}, {node.target.__name__}") + if node.target.__name__ in node_visitors: + node_visitors[node.target.__name__].define_node( + node, enn_graph, vals_to_ids + ) + else: + raise RuntimeError( + f"{node.target.__name__}" " is not supported in ENN Delegate" + ) + elif node.op in [ + "get_attr", + "placeholder", + "output", + ]: + continue + else: + raise RuntimeError(f"{node.op}" " is not supported in ENN Delegate") + + # Compile Graph + enn_graph.finish() + ser_buf = enn_graph.serialize() + enn_context_binary = enn_wrapper.Compile(ser_buf) + assert enn_context_binary is not None and len(enn_context_binary) > 0 + enn_wrapper.Destroy() + return PreprocessResult( + processed_bytes=bytes(enn_context_binary), debug_handle_map={} + ) diff --git a/backends/samsung/partition/enn_partitioner.py b/backends/samsung/partition/enn_partitioner.py new file mode 100644 index 00000000000..952cb000429 --- /dev/null +++ b/backends/samsung/partition/enn_partitioner.py @@ -0,0 +1,133 @@ +# 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. +import logging +from typing import Any, Callable, Dict, List, Optional, Tuple + +import executorch.backends.samsung.builders.node_visitor as node_visitor + +import executorch.backends.samsung.python.PyEnnWrapperAdaptor as PyEnnWrapper + +import torch +from executorch.backends.samsung.enn_preprocess import EnnBackend +from executorch.backends.samsung.serialization.compile_options import ( + ENN_COMPILE_OPTION_TITLE, +) +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 ( + generate_partitions_from_list_of_nodes, +) +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.backend.utils import tag_constant_data + +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.passes.infra.partitioner import Partition +from torch.fx.passes.operator_support import OperatorSupportBase + +SUPPORTED_OPS = [ + # support because preprocess in backend + exir_ops.edge.aten.addmm.default, + exir_ops.edge.aten.add.Scalar, + exir_ops.edge.aten.sub.Scalar, + exir_ops.edge.aten.mul.Scalar, + exir_ops.edge.aten.div.Scalar, +] + + +class EnnOperatorSupport(OperatorSupportBase): + + def __init__( + self, + edge_program: torch.export.ExportedProgram, + compile_specs: List[CompileSpec], + ): + self.edge_program = edge_program + self.enn_wrapper = PyEnnWrapper.EnnWrapper() + self.node_visitors = node_visitor.get_node_visitors(edge_program) + option_spec = get_compile_spec( + compile_specs, ENN_COMPILE_OPTION_TITLE, required=True + ) + self.enn_wrapper.Init(option_spec.value) + + def is_node_supported(self, _, node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + + if node.op in [ + "get_attr", + "placeholder", + "output", + ]: + return False + + if node.target in SUPPORTED_OPS or node.target.__name__ in self.node_visitors: + return True + + supported = self.enn_wrapper.IsNodeSupportedByBackend() + return supported + + def __del__(self): + self.enn_wrapper.Destroy() + + +class EnnPartitioner(Partitioner): + def __init__(self, compile_specs: List[CompileSpec]): + # TODO(anyone): Add meaningful initialize + self.delegation_spec = DelegationSpec(EnnBackend.__name__, compile_specs) + self.partition_tags: Dict[str, DelegationSpec] = {} + self.compile_specs = compile_specs + + 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( + edge_program.graph_module, + op_support=self.op_support_checker, + ) + + def tag_nodes(self, partitions: List[Partition]) -> None: + partition_tags: Dict[str, DelegationSpec] = {} + for partition in partitions: + # Add delegation tags + for node in partition.nodes: + delegation_tag = f"enn_{partition.id}" + node.meta["delegation_tag"] = delegation_tag + partition_tags[delegation_tag] = self.delegation_spec + return partition_tags + + # override + def partition(self, edge_program: torch.export.ExportedProgram) -> PartitionResult: + partitions = self.generate_partitions(edge_program) + logging.info(f"Find {len(partitions)} " "subgraphs to partition and lowering.") + if len(partitions) != 0: + self.partition_tags = self.tag_nodes(partitions) + tag_constant_data(edge_program) + del self.op_support_checker + return PartitionResult( + tagged_exported_program=edge_program, partition_tags=self.partition_tags + ) + + # override + def ops_to_not_decompose( + self, ep: torch.export.ExportedProgram + ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]: + ops_not_to_decompose = [ + torch.ops.aten.hardswish.default, + 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, + ] + return (ops_not_to_decompose, None) diff --git a/backends/samsung/python/.gitignore b/backends/samsung/python/.gitignore new file mode 100644 index 00000000000..8e711cc19ff --- /dev/null +++ b/backends/samsung/python/.gitignore @@ -0,0 +1,5 @@ +# exclude the pybind libraries +* + +# keep folder for ease of use +!.gitignore diff --git a/backends/samsung/runtime/CMakeLists.txt b/backends/samsung/runtime/CMakeLists.txt new file mode 100644 index 00000000000..deb93f31bc8 --- /dev/null +++ b/backends/samsung/runtime/CMakeLists.txt @@ -0,0 +1,22 @@ +# 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 +target_sources( + enn_logging + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/logging.h + PRIVATE ${CMAKE_CURRENT_LIST_DIR}/logging.cpp +) + +if(${ANDROID}) + # backend + target_sources( + enn_backend + PRIVATE ${CMAKE_CURRENT_LIST_DIR}/enn_backend.cpp + ${CMAKE_CURRENT_LIST_DIR}/enn_executor.cpp + ${CMAKE_CURRENT_LIST_DIR}/enn_api_implementation.cpp + ) +endif() diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp new file mode 100644 index 00000000000..c3914dd2297 --- /dev/null +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -0,0 +1,100 @@ +/* + * 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 + +#define ENN_LOAD_API_FUNC(handle, name, enn_api_ptr) \ + enn_api_ptr->name = \ + reinterpret_cast(loadApiFunction(handle, #name, false)); \ + if (enn_api_ptr->name == nullptr) { \ + ENN_LOG_ERROR("Unable to access symbols in enn api library: %s", #name); \ + dlclose(handle); \ + return Error::Internal; \ + } + +namespace torch { +namespace executor { +namespace enn { +void* loadApiFunction(void* handle, const char* name, bool optional) { + if (handle == nullptr) { + return nullptr; + } + void* fn = dlsym(handle, name); + if (fn == nullptr && !optional) { + ENN_LOG_WARN("Failed to load function %s", name); + } + 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_api.initialize_ = true; + } + } + return &enn_api; +} + +EnnApi::~EnnApi() { + if (getInitialize()) { + unloadApiLib(); + } +} + +bool EnnApi::getInitialize() const { + return initialize_; +} + +Error EnnApi::loadApiLib() { + const char enn_api_lib_name[] = "libenn_public_api_cpp.so"; + libenn_public_api_ = dlopen(enn_api_lib_name, RTLD_NOW | RTLD_LOCAL); + ET_CHECK_OR_RETURN_ERROR( + libenn_public_api_ != nullptr, Internal, "Lib load failed.") + + ENN_LOAD_API_FUNC(libenn_public_api_, EnnInitialize, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetPreferencePerfMode, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnGetPreferencePerfMode, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModel, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromMemory, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetFastIpc, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnUnsetFastIpc, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnExecuteModelFastIpc, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnExecuteModel, this); + ENN_LOAD_API_FUNC( + libenn_public_api_, EnnExecuteModelWithSessionIdAsync, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnExecuteModelWithSessionIdWait, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnCloseModel, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnDeinitialize, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnAllocateAllBuffers, this); + ENN_LOAD_API_FUNC( + libenn_public_api_, EnnAllocateAllBuffersWithSessionId, this); + 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); + + return Error::Ok; +} + +Error EnnApi::unloadApiLib() { + if (dlclose(libenn_public_api_) != 0) { + ENN_LOG_ERROR("Failed to close enn public api library. %s", dlerror()); + return Error::Internal; + }; + return Error::Ok; +} + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/runtime/enn_api_implementation.h b/backends/samsung/runtime/enn_api_implementation.h new file mode 100644 index 00000000000..efc57f0d276 --- /dev/null +++ b/backends/samsung/runtime/enn_api_implementation.h @@ -0,0 +1,125 @@ +/* + * 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 +#include +#include +#include + +namespace torch { +namespace executor { +namespace enn { + +class EnnApi { + public: + EnnApi(const EnnApi&) = delete; + EnnApi& operator=(const EnnApi&) = delete; + EnnApi(EnnApi&&) = delete; + EnnApi& operator=(EnnApi&&) = delete; + ~EnnApi(); + + static EnnApi* getEnnApiInstance(); + + EnnReturn (*EnnInitialize)(void); + EnnReturn (*EnnSetPreferencePerfMode)(const uint32_t val); + EnnReturn (*EnnGetPreferencePerfMode)(uint32_t* val_ptr); + EnnReturn (*EnnOpenModel)(const char* model_file, EnnModelId* model_id); + EnnReturn (*EnnOpenModelFromMemory)( + const char* va, + const uint32_t size, + EnnModelId* model_id); + EnnReturn (*EnnSetFastIpc)(void); + EnnReturn (*EnnUnsetFastIpc)(void); + EnnReturn (*EnnExecuteModelFastIpc)( + const EnnModelId model_id, + int client_sleep_usec); + EnnReturn (*EnnExecuteModel)(const EnnModelId model_id); + EnnReturn (*EnnExecuteModelWithSessionIdAsync)( + const EnnModelId model_id, + const int session_id); + EnnReturn (*EnnCloseModel)(const EnnModelId model_id); + EnnReturn (*EnnDeinitialize)(void); + EnnReturn (*EnnAllocateAllBuffers)( + const EnnModelId model_id, + EnnBufferPtr** out_buffers, + NumberOfBuffersInfo* out_buffers_info); + EnnReturn (*EnnAllocateAllBuffersWithSessionId)( + const EnnModelId model_id, + EnnBufferPtr** out_buffers, + NumberOfBuffersInfo* out_buffers_info, + const int session_id, + const bool do_commit); + EnnReturn (*EnnExecuteModelWithSessionIdWait)( + const EnnModelId model_id, + const int session_id); + EnnReturn (*EnnBufferCommit)(const EnnModelId model_id); + EnnReturn (*EnnGetBuffersInfo)( + const EnnModelId model_id, + NumberOfBuffersInfo* buffers_info); + EnnReturn ( + *EnnReleaseBuffers)(EnnBufferPtr* buffers, const int32_t numOfBuffers); + + private: + static std::mutex instance_mutex_; + std::atomic_bool initialize_ = false; + // Pointer to the dlopen libs + void* libenn_public_api_ = nullptr; + + EnnApi() = default; + bool getInitialize() const; + Error loadApiLib(); + Error unloadApiLib(); +}; + +typedef EnnReturn (*EnnInitialize_fn)(void); +typedef EnnReturn (*EnnSetPreferencePerfMode_fn)(const uint32_t val); +typedef EnnReturn (*EnnGetPreferencePerfMode_fn)(uint32_t* val_ptr); +typedef EnnReturn ( + *EnnOpenModel_fn)(const char* model_file, EnnModelId* model_id); +typedef EnnReturn (*EnnOpenModelFromMemory_fn)( + const char* va, + const uint32_t size, + EnnModelId* model_id); +typedef EnnReturn (*EnnSetFastIpc_fn)(void); +typedef EnnReturn (*EnnUnsetFastIpc_fn)(void); +typedef EnnReturn (*EnnExecuteModelFastIpc_fn)( + const EnnModelId model_id, + int client_sleep_usec); +typedef EnnReturn (*EnnExecuteModel_fn)(const EnnModelId model_id); +typedef EnnReturn (*EnnExecuteModelWithSessionIdAsync_fn)( + const EnnModelId model_id, + const int session_id); +typedef EnnReturn (*EnnCloseModel_fn)(const EnnModelId model_id); +typedef EnnReturn (*EnnDeinitialize_fn)(void); +typedef EnnReturn (*EnnAllocateAllBuffers_fn)( + const EnnModelId model_id, + EnnBufferPtr** out_buffers, + NumberOfBuffersInfo* out_buffers_info); +typedef EnnReturn (*EnnAllocateAllBuffersWithSessionId_fn)( + const EnnModelId model_id, + EnnBufferPtr** out_buffers, + NumberOfBuffersInfo* out_buffers_info, + const int session_id, + const bool do_commit); +typedef EnnReturn (*EnnExecuteModelWithSessionIdWait_fn)( + const EnnModelId model_id, + const int session_id); +typedef EnnReturn (*EnnBufferCommit_fn)(const EnnModelId model_id); +typedef EnnReturn (*EnnGetBuffersInfo_fn)( + const EnnModelId model_id, + NumberOfBuffersInfo* buffers_info); +typedef EnnReturn ( + *EnnReleaseBuffers_fn)(EnnBufferPtr* buffers, const int32_t numOfBuffers); + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp new file mode 100644 index 00000000000..ec5ab89206a --- /dev/null +++ b/backends/samsung/runtime/enn_backend.cpp @@ -0,0 +1,97 @@ +/* + * 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 + +#pragma clang diagnostic ignored "-Wglobal-constructors" + +namespace torch { +namespace executor { + +class EnnBackend final : public PyTorchBackendInterface { + public: + ~EnnBackend() = default; + + bool is_available() const override { + return true; + } + + Result init( + BackendInitContext& context, + FreeableBuffer* processed, + ArrayRef compile_specs) const override { + MemoryAllocator* runtime_allocator = context.get_runtime_allocator(); + auto executor = runtime_allocator->allocateInstance(); + + const char* binary_buf_addr = + reinterpret_cast(processed->data()); + size_t buf_size = processed->size(); + Error err = executor->initialize(binary_buf_addr, buf_size); + if (err != Error::Ok) { + ENN_LOG_ERROR("Exynos backend initialize failed."); + executor->~EnnExecutor(); + return err; + } + return executor; + } + + Error execute( + BackendExecutionContext& context, + DelegateHandle* handle, + Span args) const override { + auto executor = static_cast(handle); + + std::vector inputs; + std::vector outputs; + for (int32_t index = 0; + index < executor->getInputSize() + executor->getOutputSize(); + index++) { + ET_CHECK_OR_RETURN_ERROR( + args[index]->isTensor(), + InvalidArgument, + "Expected argument to delegate at index %u to be a Tensor, but got %" PRIu32, + index, + static_cast(args[index]->tag)); + Tensor* tensor = &args[index]->toTensor(); + enn::DataBuffer data_buffer = { + .buf_ptr_ = tensor->mutable_data_ptr(), + .size_ = tensor->nbytes()}; + if (index < executor->getInputSize()) { + inputs.push_back(data_buffer); + } else { + outputs.push_back(data_buffer); + } + } + Error err = executor->eval(inputs, outputs); + return err; + } + + void destroy(DelegateHandle* handle) const override { + if (handle != nullptr) { + auto executor = static_cast(handle); + executor->~EnnExecutor(); + } + } +}; // namespace executor + +namespace { +auto cls = EnnBackend(); +Backend backend{"EnnBackend", &cls}; +static auto success_with_compiler = register_backend(backend); +} // namespace + +} // namespace executor +} // namespace torch diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp new file mode 100644 index 00000000000..bdada72f86f --- /dev/null +++ b/backends/samsung/runtime/enn_executor.cpp @@ -0,0 +1,107 @@ +/* + * 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 torch { +namespace executor { +namespace enn { + +Error EnnExecutor::initialize(const char* binary_buf_addr, size_t buf_size) { + 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."); + + 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_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); + ET_CHECK_OR_RETURN_ERROR( + ret == ENN_RET_SUCCESS, + Internal, + "Failed to allocate buffers for model_id = 0x%" PRIX64, + model_id_); + num_of_inputs_ = buffers_info.n_in_buf; + num_of_outputs_ = buffers_info.n_out_buf; + + return Error::Ok; +} + +Error EnnExecutor::eval( + const std::vector& inputs, + const std::vector& outputs) { + const EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); + ET_CHECK_OR_RETURN_ERROR( + inputs.size() == getInputSize(), + InvalidArgument, + "Invalid number of inputs, expect %" PRIu32 " while get %ld", + getInputSize(), + inputs.size()); + ET_CHECK_OR_RETURN_ERROR( + outputs.size() == getOutputSize(), + InvalidArgument, + "Invalid number of outputs, expected %" PRIu32 " while get %ld", + getOutputSize(), + outputs.size()); + + int relative_input_index = 0; + for (const auto& input : inputs) { + EnnBufferPtr* input_buffer_ptr = alloc_buffer_; + EnnBuffer& enn_buffer = *input_buffer_ptr[relative_input_index]; + memcpy(enn_buffer.va, input.buf_ptr_, input.size_); + relative_input_index++; + } + + ENN_LOG_DEBUG("Start to execute model."); + auto ret = enn_api_inst->EnnExecuteModel(model_id_); + if (ret != ENN_RET_SUCCESS) { + ENN_LOG_ERROR("EnnExecuteModel Failed"); + return Error::Internal; + } + + EnnBufferPtr* output_buffer_ptr = alloc_buffer_ + getInputSize(); + int relative_output_index = 0; + for (const auto& output : outputs) { + EnnBuffer& enn_buffer = *output_buffer_ptr[relative_output_index]; + memcpy(output.buf_ptr_, enn_buffer.va, output.size_); + relative_output_index++; + } + return Error::Ok; +} + +EnnExecutor::~EnnExecutor() { + const EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); + NumberOfBuffersInfo buffers_info; + if (enn_api_inst->EnnGetBuffersInfo(model_id_, &buffers_info) == + ENN_RET_SUCCESS) { + const int32_t num_of_buffers = + buffers_info.n_in_buf + buffers_info.n_out_buf; + enn_api_inst->EnnReleaseBuffers(alloc_buffer_, num_of_buffers); + } + enn_api_inst->EnnCloseModel(model_id_); +} + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h new file mode 100644 index 00000000000..902b420a036 --- /dev/null +++ b/backends/samsung/runtime/enn_executor.h @@ -0,0 +1,51 @@ +/* + * 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 torch { +namespace executor { +namespace enn { + +struct DataBuffer { + void* buf_ptr_ = nullptr; + size_t size_ = 0; +}; + +class EnnExecutor { + public: + Error initialize(const char* binary_buf_addr, size_t buf_size); + + Error eval( + const std::vector& inputs, + const std::vector& outputs); + + int32_t getInputSize() const { + return num_of_inputs_; + } + int32_t getOutputSize() const { + return num_of_outputs_; + } + + ~EnnExecutor(); + + private: + EnnModelId model_id_; + EnnBufferPtr* alloc_buffer_ = nullptr; + int32_t num_of_inputs_ = 0; + int32_t num_of_outputs_ = 0; +}; + +} // namespace enn +} // namespace executor +} // namespace torch diff --git a/backends/samsung/runtime/enn_type.h b/backends/samsung/runtime/enn_type.h new file mode 100644 index 00000000000..4b91fe2d08a --- /dev/null +++ b/backends/samsung/runtime/enn_type.h @@ -0,0 +1,55 @@ +/* + * 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 + +typedef uint64_t EnnModelId; // higher 32-bits are initialized with zero + +typedef enum _EnnReturn { + ENN_RET_SUCCESS = 0, + ENN_RET_FAILED, + ENN_RET_IO, + ENN_RET_INVAL, + ENN_RET_FILTERED, + ENN_RET_CANCELED, + ENN_RET_MEM_ERR, + ENN_RET_SIZE, + ENN_RET_FAILED_TIMEOUT_ENN = 10, + ENN_RET_FAILED_TIMEOUT_DD, + ENN_RET_FAILED_TIMEOUT_FW, + ENN_RET_FAILED_TIMEOUT_HW_NOTRECOVERED, + ENN_RET_FAILED_TIMEOUT_HW_RECOVERED, + ENN_RET_FAILED_SERVICE_NULL, + ENN_RET_FAILED_RESOURCE_BUSY, + ENN_RET_NOT_SUPPORTED = 0xFF, +} EnnReturn; + +/* NOTE: should be sync with types.hal */ +typedef enum _enn_buf_dir_e { + ENN_DIR_IN, + ENN_DIR_OUT, + ENN_DIR_EXT, + ENN_DIR_NONE, + ENN_DIR_SIZE +} enn_buf_dir_e; + +// data structure for user buffer +typedef struct _ennBuffer { + void* va; + uint32_t size; // requested size + uint32_t offset; +} EnnBuffer; + +typedef EnnBuffer* EnnBufferPtr; + +typedef struct _NumberOfBuffersInfo { + uint32_t n_in_buf; + uint32_t n_out_buf; +} NumberOfBuffersInfo; diff --git a/backends/samsung/runtime/logging.cpp b/backends/samsung/runtime/logging.cpp new file mode 100644 index 00000000000..04ab22ddc0d --- /dev/null +++ b/backends/samsung/runtime/logging.cpp @@ -0,0 +1,79 @@ +/* + * 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 +#ifdef __ANDROID__ +#include +#endif +#include + +void EnnLogImpl(ENN_LOG_LEVEL log_level, const char* format, ...) { + va_list args; + va_start(args, format); +#ifdef __ANDROID__ + int android_severity = ANDROID_LOG_DEBUG; + switch (log_level) { + case ENN_LOG_LEVEL::DEBUG: + android_severity = ANDROID_LOG_DEBUG; + break; + case ENN_LOG_LEVEL::INFO: + android_severity = ANDROID_LOG_INFO; + break; + case ENN_LOG_LEVEL::WARNING: + android_severity = ANDROID_LOG_WARN; + break; + case ENN_LOG_LEVEL::ERROR: + android_severity = ANDROID_LOG_ERROR; + break; + default: + android_severity = ANDROID_LOG_UNKNOWN; + break; + } + __android_log_print(android_severity, "[Exynos ExecuTorch]", format, args); +#endif + const char* serverity_name; + switch (log_level) { + case ENN_LOG_LEVEL::DEBUG: + serverity_name = "DEBUG"; + break; + case ENN_LOG_LEVEL::INFO: + serverity_name = "INFO"; + break; + case ENN_LOG_LEVEL::WARNING: + serverity_name = "WARN"; + break; + case ENN_LOG_LEVEL::ERROR: + serverity_name = "ERROR"; + break; + default: + serverity_name = "UNKNOWN"; + break; + } + fprintf(stderr, "[%s][Exynos ExecuTorch]", serverity_name); + vfprintf(stderr, format, args); + fputc('\n', stderr); + va_end(args); +} + +#if defined(NDEBUG) +ENN_LOG_LEVEL EnnLogManager::output_log_level_ = ENN_LOG_LEVEL::INFO; +#else +ENN_LOG_LEVEL EnnLogManager::output_log_level_ = ENN_LOG_LEVEL::DEBUG; +#endif +std::mutex EnnLogManager::log_mutex_; + +void EnnLogManager::setLogLevel(ENN_LOG_LEVEL log_level) { + std::lock_guard lgd(log_mutex_); + output_log_level_ = log_level; +} + +bool EnnLogManager::isLogOn(ENN_LOG_LEVEL log_level) { + std::lock_guard lgd(log_mutex_); + return static_cast(log_level) >= static_cast(output_log_level_); +} diff --git a/backends/samsung/runtime/logging.h b/backends/samsung/runtime/logging.h new file mode 100644 index 00000000000..6271017e126 --- /dev/null +++ b/backends/samsung/runtime/logging.h @@ -0,0 +1,56 @@ +/* + * 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 + +#define FILENAME(fp) (strrchr(fp, '/') ? strrchr(fp, '/') + 1 : fp) + +enum class ENN_LOG_LEVEL { + DEBUG = 0, + INFO = 1, + WARNING = 2, + ERROR = 3, + MAX_LEVEL = ERROR, +}; + +class EnnLogManager { + public: + static void setLogLevel(ENN_LOG_LEVEL log_level); + static bool isLogOn(ENN_LOG_LEVEL log_level); + + private: + static ENN_LOG_LEVEL output_log_level_; + static std::mutex log_mutex_; +}; + +void EnnLogImpl(ENN_LOG_LEVEL log_level, const char* format, ...); + +#ifdef NDEBUG +#define __LOGX(log_level, fmt, ...) \ + if (EnnLogManager::isLogOn(log_level)) { \ + EnnLogImpl(log_level, fmt, ##__VA_ARGS__); \ + } +#else +#define __LOGX(log_level, fmt, ...) \ + if (EnnLogManager::isLogOn(log_level)) { \ + EnnLogImpl( \ + log_level, \ + "[%s:%d]: " fmt, \ + FILENAME(__FILE__), \ + __LINE__, \ + ##__VA_ARGS__); \ + } +#endif + +#define ENN_LOG_DEBUG(fmt, ...) __LOGX(ENN_LOG_LEVEL::DEBUG, fmt, ##__VA_ARGS__) +#define ENN_LOG_INFO(fmt, ...) __LOGX(ENN_LOG_LEVEL::INFO, fmt, ##__VA_ARGS__) +#define ENN_LOG_WARN(fmt, ...) \ + __LOGX(ENN_LOG_LEVEL::WARNING, fmt, ##__VA_ARGS__) +#define ENN_LOG_ERROR(fmt, ...) __LOGX(ENN_LOG_LEVEL::ERROR, fmt, ##__VA_ARGS__) diff --git a/backends/samsung/serialization/compile_options.py b/backends/samsung/serialization/compile_options.py new file mode 100644 index 00000000000..1ad2350cfeb --- /dev/null +++ b/backends/samsung/serialization/compile_options.py @@ -0,0 +1,73 @@ +# 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. + +import json +import os +import tempfile + +from dataclasses import dataclass +from enum import IntEnum, unique + +import pkg_resources +from executorch.exir._serialize._dataclass import _DataclassEncoder +from executorch.exir._serialize._flatbuffer import _flatc_compile +from executorch.exir.backend.backend_details import CompileSpec + + +@unique +class SamsungChipset(IntEnum): + UNDEFINED_CHIP_V = 0 + E9955 = 9955 + + +@dataclass +class EnnExecuTorchOptions: + chipset: SamsungChipset = SamsungChipset.UNDEFINED_CHIP_V + + +ENN_COMPILE_OPTION_TITLE = "enn_compile_options" +COMPILE_OPTION_SCHEMA_NAME = "compile_options_def" + + +def gen_samsung_backend_compile_spec_core(options: EnnExecuTorchOptions) -> CompileSpec: + with tempfile.TemporaryDirectory() as d: + # schema + schema_path = os.path.join(d, "{}.fbs".format(COMPILE_OPTION_SCHEMA_NAME)) + with open(schema_path, "wb") as schema_file: + schema_file.write( + pkg_resources.resource_string( + __name__, "{}.fbs".format(COMPILE_OPTION_SCHEMA_NAME) + ) + ) + # dump json + json_path = os.path.join(d, "{}.json".format(COMPILE_OPTION_SCHEMA_NAME)) + enn_options_json = json.dumps(options, cls=_DataclassEncoder, indent=4) + with open(json_path, "wb") as json_file: + json_file.write(enn_options_json.encode("ascii")) + + _flatc_compile(d, schema_path, json_path) + output_path = os.path.join(d, "{}.eeto".format(COMPILE_OPTION_SCHEMA_NAME)) + with open(output_path, "rb") as output_file: + return CompileSpec(ENN_COMPILE_OPTION_TITLE, output_file.read()) + + +def gen_samsung_backend_compile_spec( + chipset: str, +): + """ + A function to generate an ExecuTorch binary for Samsung Backend. + + Attributes: + chipset (str): chipset name in SamsungChipset. For example, E9955 or e9955 both work. + + Returns: + CompileSpec: key is COMPILE_OPTION_SCHEMA_NAME, value is serialization binary of fb schema + """ + option = EnnExecuTorchOptions( + getattr(SamsungChipset, chipset.upper()), + ) + + return gen_samsung_backend_compile_spec_core(option) diff --git a/backends/samsung/serialization/compile_options_def.fbs b/backends/samsung/serialization/compile_options_def.fbs new file mode 100644 index 00000000000..d38c2772715 --- /dev/null +++ b/backends/samsung/serialization/compile_options_def.fbs @@ -0,0 +1,23 @@ +//============================================================================ +// +// Copyright (c) 2025 Samsung Electronics. 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. +// +//============================================================================ + +namespace enn_option; + +// Identifier of a valid executor schema. +file_identifier "EETO"; +// Extension of written files. +file_extension "eeto"; + + +table EnnExecuTorchOptions { + // The version of chipset. Specify the soc to compile and execute model. + chipset: int; +} + +root_type EnnExecuTorchOptions; diff --git a/backends/samsung/serialization/enn_graph_schema.py b/backends/samsung/serialization/enn_graph_schema.py new file mode 100644 index 00000000000..7e74182f9d7 --- /dev/null +++ b/backends/samsung/serialization/enn_graph_schema.py @@ -0,0 +1,85 @@ +# 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. + +import logging +from typing import Dict, List, Optional, Union + +import executorch.backends.samsung.python.PyGraphWrapperAdaptor as PyGraphWrapper + +import numpy as np + +import torch + + +class EnnGraph: + def __init__(self): + # default + self.version = "0.6.0" + self.graph = PyGraphWrapper.PyEnnGraphWrapper() + self.graph.Init() + + self.inputs = [] + self.outputs = [] + + def define_op( + self, + name, + type, + input_ids: List[int], + output_ids: List[int], + params: Optional[Dict] = None, + ): + op = PyGraphWrapper.PyEnnOpWrapper(name, type, input_ids, output_ids) + + if params is not None: + assert isinstance(params, dict), "Please pass op params as dict type." + for key in params: + py_param_wrapper = PyGraphWrapper.OpParamWrapper(key) + if isinstance(params[key], (list, tuple)): + py_param_wrapper.SetVectorValue(params[key]) + elif isinstance(params[key], str): + py_param_wrapper.SetStringValue(params[key]) + elif isinstance(params[key], (int, float, bool)): + py_param_wrapper.SetScalarValue(params[key]) + else: + logging.error("Unsupported param type.") + op.AddOpParam(py_param_wrapper) + + self.graph.DefineOpNode(op) + + def define_tensor( + self, + name: str, + shape: List, + data_type: str, + tensor_type: str, + data: Optional[Union[np.ndarray, torch.Tensor]] = None, + ) -> int: + layout = "NCHW" if len(shape) == 4 else "UNDEFINED" + + tensor = PyGraphWrapper.PyEnnTensorWrapper(name, shape, data_type, layout) + + if data is not None: + if isinstance(data, torch.Tensor): + data = data.detach().numpy() + tensor.AddData(data) + + id = self.graph.DefineTensor(tensor) + + if tensor_type == "INPUT": + self.inputs.append(id) + elif tensor_type == "OUTPUT": + self.outputs.append(id) + + return id + + def finish(self): + self.graph.SetGraphInputTensors(self.inputs) + self.graph.SetGraphOutputTensors(self.outputs) + self.graph.FinishBuild() + + def serialize(self): + return self.graph.Serialize() diff --git a/backends/samsung/test/ops/test_add.py b/backends/samsung/test/ops/test_add.py new file mode 100644 index 00000000000..8b1b4b4a770 --- /dev/null +++ b/backends/samsung/test/ops/test_add.py @@ -0,0 +1,62 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Add(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + +class AddConstant(torch.nn.Module): + def __init__(self, constant) -> None: + super().__init__() + self.constant = constant + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.constant + + +class TestAdd(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.add.Tensor": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_add_Tensor"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_simple_add(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Add(), inputs) + + def test_fp32_const_add(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(AddConstant(torch.randn(1, 3, 8, 8)), inputs) + + def test_fp32_add_broadcast(self): + inputs = (torch.randn(1, 1, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Add(), inputs) diff --git a/backends/samsung/test/ops/test_avg_pool2d.py b/backends/samsung/test/ops/test_avg_pool2d.py new file mode 100644 index 00000000000..2614516dd0d --- /dev/null +++ b/backends/samsung/test/ops/test_avg_pool2d.py @@ -0,0 +1,66 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class AvgPool2d(torch.nn.Module): + def __init__( + self, + kernel_size=2, + stride=1, + padding=0, + ) -> None: + super().__init__() + self.avg_pool = torch.nn.AvgPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + count_include_pad=False, + ceil_mode=False, + ).to(torch.float) + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(1, 16, 24, 24) + return (input_1,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.avg_pool(x) + + +class TestAvgPool2d(unittest.TestCase): + def _test(self, module: torch.nn.Module): + tester = SamsungTester( + module, + module.get_example_inputs(), + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.avg_pool2d.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_avg_pool2d_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_avg_pool2d(self): + self._test(AvgPool2d()) + + def test_fp32_avg_pool2d_with_stride(self): + self._test(AvgPool2d(stride=2)) + + def test_fp32_avg_pool2d_with_kernel_size(self): + self._test(AvgPool2d(kernel_size=4)) diff --git a/backends/samsung/test/ops/test_batch_norm.py b/backends/samsung/test/ops/test_batch_norm.py new file mode 100644 index 00000000000..7cb9db0e47c --- /dev/null +++ b/backends/samsung/test/ops/test_batch_norm.py @@ -0,0 +1,51 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class BatchNorm(torch.nn.Module): + def __init__(self, num_features: int) -> None: + super().__init__() + self.num_features = num_features + self.bn = torch.nn.BatchNorm2d(num_features=self.num_features) + self.bn.weight.data.uniform_(-0.1, 0.1) + self.bn.bias.data.uniform_(-0.1, 0.1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.bn(x) + + +class TestBatchNorm(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default" + ] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_batch_norm(self): + num_features = 16 + inputs = (torch.randn(4, num_features, 32, 32),) + self._test(BatchNorm(num_features), inputs) diff --git a/backends/samsung/test/ops/test_bmm.py b/backends/samsung/test/ops/test_bmm.py new file mode 100644 index 00000000000..7712a50a0c9 --- /dev/null +++ b/backends/samsung/test/ops/test_bmm.py @@ -0,0 +1,48 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class BatchMatmul(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(2, 16, 56) + input_2 = torch.randn(2, 56, 32) + return (input_1, input_2) + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.bmm(x, y) + + +class TestBatchMatmul(unittest.TestCase): + def _test(self, module: torch.nn.Module): + inputs = module.get_example_inputs() + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .check_count({"torch.ops.aten.bmm.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_bmm_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_bmm(self): + self._test(BatchMatmul()) diff --git a/backends/samsung/test/ops/test_cat.py b/backends/samsung/test/ops/test_cat.py new file mode 100644 index 00000000000..f744f9ca882 --- /dev/null +++ b/backends/samsung/test/ops/test_cat.py @@ -0,0 +1,48 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Concat(torch.nn.Module): + def __init__(self, axis) -> None: + super().__init__() + self.axis = axis + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.cat((x, y), dim=self.axis) + + +class TestConcat(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .check_count({"torch.ops.aten.cat.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_cat_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_concat_on_axis1(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Concat(axis=1), inputs) + + def test_fp32_concat_on_axis3(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Concat(axis=3), inputs) diff --git a/backends/samsung/test/ops/test_clamp.py b/backends/samsung/test/ops/test_clamp.py new file mode 100644 index 00000000000..00e3eb72690 --- /dev/null +++ b/backends/samsung/test/ops/test_clamp.py @@ -0,0 +1,49 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Clamp(torch.nn.Module): + def __init__( + self, + minimum=0.0, + maximum=0.0, + ) -> None: + super().__init__() + self.minimum = minimum + self.maximum = maximum + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.clamp(x, self.minimum, self.maximum) + + +class TestClamp(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .check_count({"torch.ops.aten.clamp.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_clamp_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_clamp(self): + inputs = (torch.randn(1, 16, 8, 8),) + self._test(Clamp(minimum=0, maximum=2.0), inputs) diff --git a/backends/samsung/test/ops/test_constant_pad_nd.py b/backends/samsung/test/ops/test_constant_pad_nd.py new file mode 100644 index 00000000000..dae24abb7d7 --- /dev/null +++ b/backends/samsung/test/ops/test_constant_pad_nd.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class ConstantPadND(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad = (0, 0, 1, 1) + return torch.nn.functional.pad(x, pad, mode="constant", value=0) + + +class TestConstantPadND(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not( + ["executorch_exir_dialects_edge__ops_aten_constant_pad_nd_default"] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_constant_pad_nd(self): + inputs = (torch.randn(1, 6, 8, 16),) + self._test(ConstantPadND(), inputs) diff --git a/backends/samsung/test/ops/test_conv2d.py b/backends/samsung/test/ops/test_conv2d.py new file mode 100644 index 00000000000..e832a5713dc --- /dev/null +++ b/backends/samsung/test/ops/test_conv2d.py @@ -0,0 +1,98 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Conv2d(torch.nn.Module): + def __init__( + self, + in_channels=3, + out_channels=16, + stride=(2, 2), + padding=(1, 1), + dilation=(1, 1), + groups=1, + bias=True, + ) -> None: + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=stride, + padding=padding, + dilation=dilation, + groups=groups, + bias=bias, + ).to(torch.float) + + self.in_channels = in_channels + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(1, self.in_channels, 24, 24) + return (input_1,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class TransposeConv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.ConvTranspose2d( + in_channels=32, + out_channels=8, + kernel_size=2, + stride=(2, 2), + padding=(0, 0), + dilation=(1, 1), + bias=True, + ) + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(1, 32, 24, 24) + return (input_1,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class TestConv2d(unittest.TestCase): + def _test(self, module: torch.nn.Module): + tester = SamsungTester( + module, + module.get_example_inputs(), + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_convolution_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_conv2d_without_bias(self): + self._test(Conv2d(bias=False)) + + def test_fp32_conv2d_with_bias(self): + self._test(Conv2d(bias=True)) + + def test_fp32_depthwise_conv2d(self): + self._test(Conv2d(in_channels=8, out_channels=8, groups=8)) + + def test_fp32_transpose_conv2d(self): + self._test(TransposeConv2d()) diff --git a/backends/samsung/test/ops/test_div.py b/backends/samsung/test/ops/test_div.py new file mode 100644 index 00000000000..31384afd896 --- /dev/null +++ b/backends/samsung/test/ops/test_div.py @@ -0,0 +1,49 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Div(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x / y + + +class TestDiv(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.div.Tensor": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_div_Tensor"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_simple_div(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8).abs() + 1e-3) + self._test(Div(), inputs) + + def test_fp32_div_broadcast(self): + inputs = (torch.randn(1, 1, 8, 8), torch.randn(1, 3, 8, 8).abs() + 1e-3) + self._test(Div(), inputs) diff --git a/backends/samsung/test/ops/test_embedding.py b/backends/samsung/test/ops/test_embedding.py new file mode 100644 index 00000000000..fb6aaaf7766 --- /dev/null +++ b/backends/samsung/test/ops/test_embedding.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Embedding(torch.nn.Module): + def __init__(self, num_embeddings, embedding_dim=256) -> None: + super().__init__() + self.embedding = torch.nn.Embedding(num_embeddings, embedding_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.embedding(x) + + +class TestEmbedding(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, inputs, [gen_samsung_backend_compile_spec("E9955")] + ) + ( + tester.export() + .check_count({"torch.ops.aten.embedding.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_embedding_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_embedding(self): + num_embeddings = 2048 + inputs = (torch.randint(0, num_embeddings, (1, 64), dtype=torch.int32),) + self._test(Embedding(num_embeddings=num_embeddings), inputs) diff --git a/backends/samsung/test/ops/test_expand_copy.py b/backends/samsung/test/ops/test_expand_copy.py new file mode 100644 index 00000000000..47df38be8e2 --- /dev/null +++ b/backends/samsung/test/ops/test_expand_copy.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class ExpandCopy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand(1, 16, -1, -1) + + +class TestExpand(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.expand.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_expand_copy_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_expand_copy(self): + inputs = (torch.randn(1, 1, 8, 8),) + self._test(ExpandCopy(), inputs) diff --git a/backends/samsung/test/ops/test_gelu.py b/backends/samsung/test/ops/test_gelu.py new file mode 100644 index 00000000000..4e6f2d971ab --- /dev/null +++ b/backends/samsung/test/ops/test_gelu.py @@ -0,0 +1,64 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class GELU(torch.nn.Module): + def __init__(self, with_conv=False) -> None: + super().__init__() + self.module = ( + torch.nn.Sequential( + torch.nn.Conv2d( + in_channels=3, + out_channels=16, + kernel_size=3, + stride=(2, 2), + padding=(1, 1), + dilation=(1, 1), + ).to(torch.float), + torch.nn.GELU(), + ) + if with_conv + else torch.nn.GELU() + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestGELU(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.gelu.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_gelu_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_single_gelu(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(GELU(with_conv=False), inputs) + + def test_fp32_conv_gelu(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(GELU(with_conv=True), inputs) diff --git a/backends/samsung/test/ops/test_leaky_relu.py b/backends/samsung/test/ops/test_leaky_relu.py new file mode 100644 index 00000000000..0af6ea0da90 --- /dev/null +++ b/backends/samsung/test/ops/test_leaky_relu.py @@ -0,0 +1,46 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class LeakyReLU(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.module = torch.nn.LeakyReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestLeakyReLU(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.leaky_relu.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_leaky_relu_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_leaky_relu(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(LeakyReLU(), inputs) diff --git a/backends/samsung/test/ops/test_linear.py b/backends/samsung/test/ops/test_linear.py new file mode 100644 index 00000000000..f327464fc0c --- /dev/null +++ b/backends/samsung/test/ops/test_linear.py @@ -0,0 +1,48 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Linear(torch.nn.Module): + def __init__(self, in_features) -> None: + super().__init__() + self.module = torch.nn.Linear(in_features, 8) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestLinear(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.linear.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_linear_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .run_method_and_compare_outputs() + ) + + def test_fp32_linear(self): + in_num_features = 24 + inputs = (torch.randn(128, in_num_features),) + self._test(Linear(in_num_features), inputs) diff --git a/backends/samsung/test/ops/test_log_softmax.py b/backends/samsung/test/ops/test_log_softmax.py new file mode 100644 index 00000000000..2e2b3ff0604 --- /dev/null +++ b/backends/samsung/test/ops/test_log_softmax.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class LogSoftmax(torch.nn.Module): + def __init__(self, dim) -> None: + super().__init__() + self.module = torch.nn.LogSoftmax(dim=dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestLogSoftmax(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten__log_softmax_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_log_softmax(self): + inputs = (torch.randn(1, 16, 56, 56),) + self._test(LogSoftmax(dim=1), inputs) diff --git a/backends/samsung/test/ops/test_max_pool2d.py b/backends/samsung/test/ops/test_max_pool2d.py new file mode 100644 index 00000000000..9ead91b2bff --- /dev/null +++ b/backends/samsung/test/ops/test_max_pool2d.py @@ -0,0 +1,71 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class MaxPool2d(torch.nn.Module): + def __init__( + self, + kernel_size=2, + stride=1, + padding=0, + dilation=1, + ) -> None: + super().__init__() + self.max_pool = torch.nn.MaxPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + return_indices=False, + ceil_mode=False, + ).to(torch.float) + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(1, 16, 24, 24) + return (input_1,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.max_pool(x) + + +class TestMaxPool2d(unittest.TestCase): + def _test(self, module: torch.nn.Module): + tester = SamsungTester( + module, + module.get_example_inputs(), + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.max_pool2d.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_max_pool2d_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_max_pool2d(self): + self._test(MaxPool2d()) + + def test_fp32_max_pool2d_with_padding(self): + self._test(MaxPool2d(padding=1)) + + def test_fp32_max_pool2d_with_kernel_size(self): + self._test(MaxPool2d(kernel_size=4)) + + def test_fp32_max_pool2d_with_dilation(self): + self._test(MaxPool2d(dilation=2)) diff --git a/backends/samsung/test/ops/test_mean_dim.py b/backends/samsung/test/ops/test_mean_dim.py new file mode 100644 index 00000000000..113e26c45b2 --- /dev/null +++ b/backends/samsung/test/ops/test_mean_dim.py @@ -0,0 +1,50 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class MeanDim(torch.nn.Module): + def __init__(self, keep_dims=True) -> None: + super().__init__() + self.keep_dims = keep_dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.mean(x, dim=[2, 3], keepdim=self.keep_dims) + + +class TestMeanDim(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.mean.dim": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_mean_dim"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_mean_with_keep_dims(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(MeanDim(), inputs) + + def test_fp32_mean_without_keep_dims(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(MeanDim(keep_dims=False), inputs) diff --git a/backends/samsung/test/ops/test_minimum.py b/backends/samsung/test/ops/test_minimum.py new file mode 100644 index 00000000000..de275cc4d46 --- /dev/null +++ b/backends/samsung/test/ops/test_minimum.py @@ -0,0 +1,44 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Minimum(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return torch.min(x, y) + + +class TestMinimum(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_minimum_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_minimum(self): + inputs = (torch.randn(1, 8, 16, 16), torch.randn(1, 8, 16, 16)) + self._test(Minimum(), inputs) diff --git a/backends/samsung/test/ops/test_mul.py b/backends/samsung/test/ops/test_mul.py new file mode 100644 index 00000000000..57d13c68b87 --- /dev/null +++ b/backends/samsung/test/ops/test_mul.py @@ -0,0 +1,62 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Mul(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x * y + + +class MulConstant(torch.nn.Module): + def __init__(self, constant) -> None: + super().__init__() + self.constant = constant + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * self.constant + + +class TestMul(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.mul.Tensor": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_mul_Tensor"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_simple_mul(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Mul(), inputs) + + def test_fp32_const_mul(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(MulConstant(torch.randn(1, 3, 8, 8)), inputs) + + def test_fp32_mul_broadcast(self): + inputs = (torch.randn(1, 1, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Mul(), inputs) diff --git a/backends/samsung/test/ops/test_permute.py b/backends/samsung/test/ops/test_permute.py new file mode 100644 index 00000000000..3889c803e85 --- /dev/null +++ b/backends/samsung/test/ops/test_permute.py @@ -0,0 +1,54 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Permute(torch.nn.Module): + def __init__(self, order=None) -> None: + super().__init__() + self.order = order + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.permute(x, self.order) + + +class TestPermute(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.permute.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_permute_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_permute_0231(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(Permute(order=[0, 2, 3, 1]), inputs) + + def test_fp32_permute_0312(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(Permute(order=[0, 3, 1, 2]), inputs) + + def test_fp32_permute_0321(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(Permute(order=[0, 3, 2, 1]), inputs) diff --git a/backends/samsung/test/ops/test_pixel_shuffle.py b/backends/samsung/test/ops/test_pixel_shuffle.py new file mode 100644 index 00000000000..bc7a53ff592 --- /dev/null +++ b/backends/samsung/test/ops/test_pixel_shuffle.py @@ -0,0 +1,48 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class PixelShuffle(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.module = torch.nn.PixelShuffle(upscale_factor=3) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestPixelShuffle(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.pixel_shuffle.default": 1}) + .to_edge_transform_and_lower() + .check_not( + ["executorch_exir_dialects_edge__ops_aten_pixel_shuffle_default"] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_pixel_shuffle(self): + inputs = (torch.randn(1, 9, 4, 4),) + self._test(PixelShuffle(), inputs) diff --git a/backends/samsung/test/ops/test_relu.py b/backends/samsung/test/ops/test_relu.py new file mode 100644 index 00000000000..386827109f0 --- /dev/null +++ b/backends/samsung/test/ops/test_relu.py @@ -0,0 +1,64 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class ReLU(torch.nn.Module): + def __init__(self, with_conv=False) -> None: + super().__init__() + self.module = ( + torch.nn.Sequential( + torch.nn.Conv2d( + in_channels=3, + out_channels=16, + kernel_size=3, + stride=(2, 2), + padding=(1, 1), + dilation=(1, 1), + ).to(torch.float), + torch.nn.ReLU(), + ) + if with_conv + else torch.nn.ReLU() + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.module(x) + + +class TestReLU(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.relu.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_relu_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_single_relu(self): + inputs = (torch.randn(1, 3, 56, 56),) + self._test(ReLU(with_conv=False), inputs) + + def test_fp32_conv_relu(self): + inputs = (torch.randn(1, 3, 56, 56),) + self._test(ReLU(with_conv=True), inputs) diff --git a/backends/samsung/test/ops/test_reshape.py b/backends/samsung/test/ops/test_reshape.py new file mode 100644 index 00000000000..8c89d946361 --- /dev/null +++ b/backends/samsung/test/ops/test_reshape.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Reshape(torch.nn.Module): + def __init__(self, new_shape) -> None: + super().__init__() + self.new_shape = new_shape + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.reshape(*self.new_shape) + + +class TestReshape(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_view_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_reshape(self): + inputs = (torch.randn(1, 16, 2, 8),) + self._test(Reshape(new_shape=[1, 32, 8]), inputs) diff --git a/backends/samsung/test/ops/test_rsqrt.py b/backends/samsung/test/ops/test_rsqrt.py new file mode 100644 index 00000000000..4bf302c867f --- /dev/null +++ b/backends/samsung/test/ops/test_rsqrt.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Rsqrt(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.rsqrt(x) + + +class TestRsqrt(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.rsqrt.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_rsqrt_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_rsqrt(self): + inputs = (torch.randn(16, 8).abs().add(1e-6),) + self._test(Rsqrt(), inputs) diff --git a/backends/samsung/test/ops/test_select.py b/backends/samsung/test/ops/test_select.py new file mode 100644 index 00000000000..dcb0667d036 --- /dev/null +++ b/backends/samsung/test/ops/test_select.py @@ -0,0 +1,51 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class SelectCopy(torch.nn.Module): + def __init__(self, axis, index) -> None: + super().__init__() + self.axis = axis + self.index = index + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.select(x, self.axis, self.index) + + +class TestSelectCopy(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.select.int": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_select_copy_int"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_select_on_axis1(self): + inputs = (torch.randn([1, 4, 16, 16]),) + self._test(SelectCopy(axis=1, index=0), inputs) + + def test_fp32_concat_on_axis3(self): + inputs = (torch.randn([1, 4, 16, 16]),) + self._test(SelectCopy(axis=3, index=6), inputs) diff --git a/backends/samsung/test/ops/test_slice_copy.py b/backends/samsung/test/ops/test_slice_copy.py new file mode 100644 index 00000000000..f31410b8a41 --- /dev/null +++ b/backends/samsung/test/ops/test_slice_copy.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class SliceCopy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x[:, :8, :, :8] + + +class TestSliceCopy(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.slice.Tensor": 2}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_slice_copy(self): + inputs = (torch.randn(1, 16, 16, 16),) + self._test(SliceCopy(), inputs) diff --git a/backends/samsung/test/ops/test_softmax.py b/backends/samsung/test/ops/test_softmax.py new file mode 100644 index 00000000000..a4c2f36acfc --- /dev/null +++ b/backends/samsung/test/ops/test_softmax.py @@ -0,0 +1,46 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Softmax(torch.nn.Module): + def __init__(self, dim=0) -> None: + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.softmax(x, dim=self.dim) + + +class TestSoftmax(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten__softmax_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .run_method_and_compare_outputs() + ) + + def test_fp32_softmax(self): + inputs = (torch.randn(1, 16, 8, 8),) + self._test(Softmax(dim=1), inputs) diff --git a/backends/samsung/test/ops/test_sqrt.py b/backends/samsung/test/ops/test_sqrt.py new file mode 100644 index 00000000000..e1a084c3611 --- /dev/null +++ b/backends/samsung/test/ops/test_sqrt.py @@ -0,0 +1,45 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Sqrt(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.sqrt(x) + + +class TestSqrt(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.sqrt.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_sqrt_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_sqrt(self): + inputs = (torch.randn(16, 8).abs(),) + self._test(Sqrt(), inputs) diff --git a/backends/samsung/test/ops/test_squeeze.py b/backends/samsung/test/ops/test_squeeze.py new file mode 100644 index 00000000000..ab93758f203 --- /dev/null +++ b/backends/samsung/test/ops/test_squeeze.py @@ -0,0 +1,46 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Squeeze(torch.nn.Module): + def __init__(self, dims) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.squeeze(x, self.dims) + + +class TestSqueeze(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.squeeze.dims": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_squeeze_dims"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_squeeze(self): + inputs = (torch.randn(1, 2, 1, 3, 1),) + self._test(Squeeze(dims=[2, 4]), inputs) diff --git a/backends/samsung/test/ops/test_sub.py b/backends/samsung/test/ops/test_sub.py new file mode 100644 index 00000000000..5541a52c80c --- /dev/null +++ b/backends/samsung/test/ops/test_sub.py @@ -0,0 +1,62 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class Sub(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x - y + + +class SubConstant(torch.nn.Module): + def __init__(self, constant) -> None: + super().__init__() + self.constant = constant + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x - self.constant + + +class TestSub(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.sub.Tensor": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_sub_Tensor"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_simple_sub(self): + inputs = (torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Sub(), inputs) + + def test_fp32_const_sub(self): + inputs = (torch.randn(1, 3, 8, 8),) + self._test(SubConstant(torch.randn(1, 3, 8, 8)), inputs) + + def test_fp32_sub_broadcast(self): + inputs = (torch.randn(1, 1, 8, 8), torch.randn(1, 3, 8, 8)) + self._test(Sub(), inputs) diff --git a/backends/samsung/test/ops/test_to_copy.py b/backends/samsung/test/ops/test_to_copy.py new file mode 100644 index 00000000000..d6917c9403f --- /dev/null +++ b/backends/samsung/test/ops/test_to_copy.py @@ -0,0 +1,48 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class ToCopy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float16) + + +class TestToCopy(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .to_edge_transform_and_lower() + .check_not( + [ + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default" + ] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_to_copy(self): + inputs = (torch.randn(1, 3, 56, 56),) + self._test(ToCopy(), inputs) diff --git a/backends/samsung/test/ops/test_unsqueeze.py b/backends/samsung/test/ops/test_unsqueeze.py new file mode 100644 index 00000000000..543fa0bc282 --- /dev/null +++ b/backends/samsung/test/ops/test_unsqueeze.py @@ -0,0 +1,50 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class UnSqueeze(torch.nn.Module): + def __init__(self, axis) -> None: + super().__init__() + self.axis = axis + + def get_example_inputs(self) -> tuple[torch.Tensor]: + input_1 = torch.randn(2, 3, 1, 4) # input should be positive + return (input_1,) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.unsqueeze(x, dim=self.axis) + + +class TestSqueeze(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.unsqueeze.default": 1}) + .to_edge_transform_and_lower() + .check_not(["executorch_exir_dialects_edge__ops_aten_unsqueeze_default"]) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_unsqueeze(self): + inputs = (torch.randn(2, 3, 1, 4),) + self._test(UnSqueeze(axis=1), inputs) diff --git a/backends/samsung/test/ops/test_upsample_bilinear2d.py b/backends/samsung/test/ops/test_upsample_bilinear2d.py new file mode 100644 index 00000000000..37dcb28df83 --- /dev/null +++ b/backends/samsung/test/ops/test_upsample_bilinear2d.py @@ -0,0 +1,53 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class UpsampleBilinear2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output_shape = [32, 32] + return torch.nn.functional.interpolate( + x, + size=output_shape, + mode="bilinear", + align_corners=False, + ) + + +class TestUpsampleBilinear2d(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.upsample_bilinear2d.vec": 1}) + .to_edge_transform_and_lower() + .check_not( + ["executorch_exir_dialects_edge__ops_aten_upsample_bilinear2d_vec"] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_upsample_bilinear2d(self): + inputs = (torch.randn(1, 16, 16, 16),) + self._test(UpsampleBilinear2d(), inputs) diff --git a/backends/samsung/test/ops/test_upsample_nearest2d.py b/backends/samsung/test/ops/test_upsample_nearest2d.py new file mode 100644 index 00000000000..e027ab23337 --- /dev/null +++ b/backends/samsung/test/ops/test_upsample_nearest2d.py @@ -0,0 +1,52 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + + +import unittest + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.test.tester import SamsungTester + + +class UpsampleNearest2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output_shape = [32, 32] + return torch.nn.functional.interpolate( + x, + size=output_shape, + mode="nearest", + ) + + +class TestUpsampleNearest2d(unittest.TestCase): + def _test(self, module: torch.nn.Module, inputs): + tester = SamsungTester( + module, + inputs, + [gen_samsung_backend_compile_spec("E9955")], + ) + ( + tester.export() + .check_count({"torch.ops.aten.upsample_nearest2d.vec": 1}) + .to_edge_transform_and_lower() + .check_not( + ["executorch_exir_dialects_edge__ops_aten_upsample_nearest2d_vec"] + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + ) + + def test_fp32_upsample_nearest2d(self): + inputs = (torch.randn(1, 4, 16, 16),) + self._test(UpsampleNearest2d(), inputs) diff --git a/backends/samsung/test/tester/__init__.py b/backends/samsung/test/tester/__init__.py new file mode 100644 index 00000000000..8cd333b2ac7 --- /dev/null +++ b/backends/samsung/test/tester/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. + +from executorch.backends.samsung.test.tester.samsung_tester import SamsungTester + + +__all__ = ["SamsungTester"] diff --git a/backends/samsung/test/tester/samsung_tester.py b/backends/samsung/test/tester/samsung_tester.py new file mode 100644 index 00000000000..1a595d5d77a --- /dev/null +++ b/backends/samsung/test/tester/samsung_tester.py @@ -0,0 +1,82 @@ +# Copyright (c) Samsung Electronics Co. LTD +# All rights reserved +# +# Licensed under the BSD License (the "License"); you may not use this file +# except in compliance with the License. See the license file in the root +# directory of this source tree for more details. +from typing import List, Optional, Tuple + +import executorch.backends.test.harness.stages as BaseStages +import torch +from executorch.backends.samsung.partition.enn_partitioner import EnnPartitioner +from executorch.backends.samsung.utils.export_utils import get_edge_compile_config + +from executorch.backends.test.harness import Tester as TesterBase +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.backend.backend_details import CompileSpec + +from torch.export import ExportedProgram + + +class Export(BaseStages.Export): + pass + + +class Quantize(BaseStages.Quantize): + pass + + +class ToEdgeTransformAndLower(BaseStages.ToEdgeTransformAndLower): + def __init__( + self, + compile_specs: Optional[List[CompileSpec]] = None, + edge_compile_config: Optional[EdgeCompileConfig] = None, + ): + compile_specs = compile_specs or [] + self.partitioners = [EnnPartitioner(compile_specs=compile_specs)] + self.edge_compile_config = edge_compile_config or get_edge_compile_config() + self.edge_dialect_program = None + + def run( + self, artifact: ExportedProgram, inputs=None, generate_etrecord: bool = False + ) -> None: + self.edge_dialect_program = to_edge_transform_and_lower( + artifact, + partitioner=self.partitioners, + compile_config=self.edge_compile_config, + ) + + +class ToExecutorch(BaseStages.ToExecutorch): + pass + + +class SamsungTester(TesterBase): + def __init__( + self, + module: torch.nn.Module, + example_inputs: Tuple[torch.Tensor], + compile_specs: Optional[List[CompileSpec]] = None, + ): + module.eval() + + super().__init__( + module=module, + example_inputs=example_inputs, + dynamic_shapes=None, + ) + + self.original_module = module + self.exported_module = module + self.example_inputs = example_inputs + self.compile_specs = compile_specs + + def to_edge_transform_and_lower( + self, + edge_compile_config: Optional[EdgeCompileConfig] = None, + ): + 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) diff --git a/backends/samsung/utils/export_utils.py b/backends/samsung/utils/export_utils.py new file mode 100644 index 00000000000..aaf407ef0b3 --- /dev/null +++ b/backends/samsung/utils/export_utils.py @@ -0,0 +1,52 @@ +# 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. + +from typing import Optional, Tuple + +import executorch.exir as exir +import torch +from executorch.backends.samsung.partition.enn_partitioner import EnnPartitioner +from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform +from executorch.exir import EdgeCompileConfig +from executorch.exir.backend.backend_details import CompileSpec + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.program._program import to_edge_transform_and_lower + + +def get_edge_compile_config(): + return EdgeCompileConfig( + _skip_dim_order=True, + _core_aten_ops_exception_list=[ + exir_ops.edge.aten.max_pool2d.default, + exir_ops.edge.aten.linear.default, + exir_ops.edge.aten.hardswish.default, + exir_ops.edge.aten.prelu.default, + exir_ops.edge.aten.pixel_shuffle.default, + exir_ops.edge.aten._safe_softmax.default, + exir_ops.edge.aten.layer_norm.default, + exir_ops.edge.aten.matmul.default, + ], + ) + + +def to_edge_transform_and_lower_to_enn( + module: torch.nn.Module, + inputs: Tuple[torch.Tensor], + compile_specs: Optional[CompileSpec] = None, +) -> exir.ExecutorchProgramManager: + assert ( + compile_specs is not None + ), "Please provide compile specifications for enn backend" + prog = torch.export.export(module, inputs) + + ahead_pass_list = [RemoveCloneOpsTransform()] + return to_edge_transform_and_lower( + prog, + ahead_pass_list, + {"forward": [EnnPartitioner(compile_specs)]}, + compile_config=get_edge_compile_config(), + ) diff --git a/backends/samsung/utils/utils.py b/backends/samsung/utils/utils.py new file mode 100644 index 00000000000..5da9808f38f --- /dev/null +++ b/backends/samsung/utils/utils.py @@ -0,0 +1,37 @@ +# 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. + +from typing import List + +import torch + +from executorch.backends.transforms.utils import is_param_node +from executorch.exir.backend.backend_details import CompileSpec + +from torch.export.exported_program import ExportedProgram + + +def get_compile_spec( + compile_specs: List[CompileSpec], spec_name: str, required=False +) -> CompileSpec: + for spec in compile_specs: + if spec_name == spec.key: + return spec + assert not required, f"Require {spec_name} but it doesn't exist." + + +def is_graph_input(exported_program: ExportedProgram, node: torch.fx.Node) -> bool: + return node.op == "placeholder" and not is_param_node(exported_program, node) + + +def is_graph_output(node: torch.fx.Node) -> bool: + # skip getitem node + for user in node.users.keys(): + if user.op == "output" or ( + user.target.__name__ == "getitem" and is_graph_output(user) + ): + return True + return False diff --git a/backends/test/facto/test_facto.py b/backends/test/facto/test_facto.py index dc2979a733c..405381f9643 100644 --- a/backends/test/facto/test_facto.py +++ b/backends/test/facto/test_facto.py @@ -8,7 +8,7 @@ # # This file contains logic to run generated operator tests using the FACTO -# library (https://github.com/pytorch-labs/FACTO). To run the tests, first +# library (https://github.com/meta-pytorch/FACTO). To run the tests, first # clone and install FACTO by running pip install . from the FACTO source # directory. Then, from the executorch root directory, run the following: # diff --git a/backends/test/harness/TARGETS b/backends/test/harness/TARGETS index 41d9a5b7682..d4edf9fb248 100644 --- a/backends/test/harness/TARGETS +++ b/backends/test/harness/TARGETS @@ -4,10 +4,7 @@ oncall("executorch") runtime.python_library( name = "tester", - srcs = [ - "__init__.py", - "tester.py", - ] + native.glob(["stages/*.py"]), + srcs = native.glob(["*.py", "stages/*.py"]), visibility = [ "//executorch/...", "@EXECUTORCH_CLIENTS", diff --git a/backends/test/harness/error_statistics.py b/backends/test/harness/error_statistics.py new file mode 100644 index 00000000000..db0ab7e3dd0 --- /dev/null +++ b/backends/test/harness/error_statistics.py @@ -0,0 +1,99 @@ +from dataclasses import dataclass + +import torch +from torch.ao.ns.fx.utils import compute_sqnr + + +@dataclass +class TensorStatistics: + """Contains summary statistics for a tensor.""" + + shape: torch.Size + """ The shape of the tensor. """ + + numel: int + """ The number of elements in the tensor. """ + + median: float + """ The median of the tensor. """ + + mean: float + """ The mean of the tensor. """ + + max: torch.types.Number + """ The maximum element of the tensor. """ + + min: torch.types.Number + """ The minimum element of the tensor. """ + + @classmethod + def from_tensor(cls, tensor: torch.Tensor) -> "TensorStatistics": + """Creates a TensorStatistics object from a tensor.""" + flattened = torch.flatten(tensor) + return cls( + shape=tensor.shape, + numel=tensor.numel(), + median=torch.quantile(flattened, q=0.5).item(), + mean=flattened.mean().item(), + max=flattened.max().item(), + min=flattened.min().item(), + ) + + +@dataclass +class ErrorStatistics: + """Contains statistics derived from the difference of two tensors.""" + + reference_stats: TensorStatistics + """ Statistics for the reference tensor. """ + + actual_stats: TensorStatistics + """ Statistics for the actual tensor. """ + + error_l2_norm: float | None + """ The L2 norm of the error between the actual and reference tensor. """ + + error_mae: float | None + """ The mean absolute error between the actual and reference tensor. """ + + error_max: float | None + """ The maximum absolute elementwise error between the actual and reference tensor. """ + + error_msd: float | None + """ The mean signed deviation between the actual and reference tensor. """ + + sqnr: float | None + """ The signal-to-quantization-noise ratio between the actual and reference tensor. """ + + @classmethod + def from_tensors( + cls, actual: torch.Tensor, reference: torch.Tensor + ) -> "ErrorStatistics": + """Creates an ErrorStatistics object from two tensors.""" + actual = actual.to(torch.float64) + reference = reference.to(torch.float64) + + if actual.shape != reference.shape: + return cls( + reference_stats=TensorStatistics.from_tensor(reference), + actual_stats=TensorStatistics.from_tensor(actual), + error_l2_norm=None, + error_mae=None, + error_max=None, + error_msd=None, + sqnr=None, + ) + + error = actual - reference + flat_error = torch.flatten(error) + + return cls( + reference_stats=TensorStatistics.from_tensor(reference), + actual_stats=TensorStatistics.from_tensor(actual), + error_l2_norm=torch.linalg.norm(flat_error).item(), + error_mae=torch.mean(torch.abs(flat_error)).item(), + error_max=torch.max(torch.abs(flat_error)).item(), + error_msd=torch.mean(flat_error).item(), + # Torch sqnr implementation requires float32 due to decorator logic + sqnr=compute_sqnr(actual.to(torch.float), reference.to(torch.float)).item(), + ) diff --git a/backends/test/harness/stages/quantize.py b/backends/test/harness/stages/quantize.py index b98c4faa3dd..9edb600e19f 100644 --- a/backends/test/harness/stages/quantize.py +++ b/backends/test/harness/stages/quantize.py @@ -7,7 +7,7 @@ DuplicateDynamicQuantChainPass, ) -from torch.export import export_for_training +from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import ( convert_pt2e, @@ -47,7 +47,7 @@ def run( assert inputs is not None if self.is_qat: artifact.train() - captured_graph = export_for_training(artifact, inputs, strict=True).module() + captured_graph = export(artifact, inputs, strict=True).module() assert isinstance(captured_graph, torch.fx.GraphModule) diff --git a/backends/test/harness/stages/serialize.py b/backends/test/harness/stages/serialize.py index 9d0bded0483..a5be1631d98 100644 --- a/backends/test/harness/stages/serialize.py +++ b/backends/test/harness/stages/serialize.py @@ -13,6 +13,7 @@ try: from executorch.extension.pybindings.portable_lib import ( # @manual _load_for_executorch_from_buffer, + Verification, ) except ImportError as e: logger.warning(f"{e=}") @@ -39,7 +40,9 @@ def graph_module(self) -> None: def run_artifact(self, inputs): inputs_flattened, _ = tree_flatten(inputs) - executorch_module = _load_for_executorch_from_buffer(self.buffer) + executorch_module = _load_for_executorch_from_buffer( + self.buffer, program_verification=Verification.Minimal + ) executorch_output = copy.deepcopy( executorch_module.run_method("forward", tuple(inputs_flattened)) ) diff --git a/backends/test/harness/stages/to_edge_transform_and_lower.py b/backends/test/harness/stages/to_edge_transform_and_lower.py index 6c5aa4b541b..19a6b6033c5 100644 --- a/backends/test/harness/stages/to_edge_transform_and_lower.py +++ b/backends/test/harness/stages/to_edge_transform_and_lower.py @@ -7,28 +7,38 @@ to_edge_transform_and_lower, ) from executorch.exir.backend.partitioner import Partitioner + from torch.export import ExportedProgram class ToEdgeTransformAndLower(Stage): def __init__( self, - default_partitioner_cls: Type, + default_partitioner_cls: Type | None = None, partitioners: Optional[List[Partitioner]] = None, edge_compile_config: Optional[EdgeCompileConfig] = None, ): - self.partitioners = partitioners or [default_partitioner_cls()] - self.edge_compile_conf = edge_compile_config or EdgeCompileConfig() + self.partitioners = ( + partitioners or [default_partitioner_cls()] + if default_partitioner_cls is not None + else [] + ) + self.edge_compile_conf = edge_compile_config or EdgeCompileConfig( + _check_ir_validity=False + ) self.edge_dialect_program = None def stage_type(self) -> StageType: return StageType.TO_EDGE_TRANSFORM_AND_LOWER - def run(self, artifact: ExportedProgram, inputs=None) -> None: + def run( + self, artifact: ExportedProgram, inputs=None, generate_etrecord: bool = False + ) -> None: self.edge_dialect_program = to_edge_transform_and_lower( artifact, compile_config=self.edge_compile_conf, partitioner=self.partitioners, + generate_etrecord=generate_etrecord, ) @property diff --git a/backends/test/harness/tester.py b/backends/test/harness/tester.py index 7019b734290..351bab4a605 100644 --- a/backends/test/harness/tester.py +++ b/backends/test/harness/tester.py @@ -4,6 +4,7 @@ import torch +from executorch.backends.test.harness.error_statistics import ErrorStatistics from executorch.backends.test.harness.stages import ( Export, Partition, @@ -33,12 +34,12 @@ def __init__( self, module: torch.nn.Module, example_inputs: Tuple[torch.Tensor], - stage_classes: Dict[StageType, Callable], + stage_classes: Dict[StageType, Callable] | None = None, dynamic_shapes: Optional[Tuple[Any]] = None, ): module.eval() - self.stage_classes = stage_classes + self.stage_classes = stage_classes or Tester.default_stage_classes() self.original_module = module self.example_inputs = example_inputs self.dynamic_shapes = dynamic_shapes @@ -182,10 +183,10 @@ def _post(self, stage): assert stage_type in self.stages self.stages[stage_type] = stage - def _run_stage(self, stage_instance, inputs=None): + def _run_stage(self, stage_instance, inputs=None, *args, **kwargs): assert isinstance(stage_instance, Stage) prev_stage_artifact = self._pre(stage_instance) - stage_instance.run(prev_stage_artifact, inputs=inputs) + stage_instance.run(prev_stage_artifact, inputs=inputs, *args, **kwargs) # noqa self._post(stage_instance) return self @@ -212,11 +213,14 @@ def to_edge(self, to_edge_stage: Optional[ToEdge] = None): return res def to_edge_transform_and_lower( - self, to_edge_and_transform_stage: Optional[ToEdgeTransformAndLower] = None + self, + to_edge_and_transform_stage: Optional[ToEdgeTransformAndLower] = None, + generate_etrecord: bool = False, ): return self._run_stage( to_edge_and_transform_stage - or self._get_default_stage(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + or self._get_default_stage(StageType.TO_EDGE_TRANSFORM_AND_LOWER), + generate_etrecord=generate_etrecord, ) def run_passes(self, run_passes_stage: Optional[RunPasses] = None): @@ -302,20 +306,15 @@ def run_method_and_compare_outputs( atol=1e-03, rtol=1e-03, qtol=0, + statistics_callback: Callable[[ErrorStatistics], None] | None = None, ): number_of_runs = 1 if inputs is not None else num_runs reference_stage = self.stages[StageType.EXPORT] stage = stage or self.cur - print(f"Comparing Stage {stage} with Stage {reference_stage}") - for run_iteration in range(number_of_runs): + for _ in range(number_of_runs): inputs_to_run = inputs if inputs else next(self.generate_random_inputs()) - input_shapes = [ - generated_input.shape if hasattr(generated_input, "shape") else None - for generated_input in inputs_to_run - ] - print(f"Run {run_iteration} with input shapes: {input_shapes}") # Reference output (and quantization scale) ( @@ -328,13 +327,25 @@ def run_method_and_compare_outputs( # Output from running artifact at stage stage_output = self.stages[stage].run_artifact(inputs_to_run) self._compare_outputs( - reference_output, stage_output, quantization_scale, atol, rtol, qtol + reference_output, + stage_output, + quantization_scale, + atol, + rtol, + qtol, + statistics_callback, ) return self @staticmethod - def _assert_outputs_equal(model_output, ref_output, atol=1e-03, rtol=1e-03): + def _assert_outputs_equal( + model_output, + ref_output, + atol=1e-03, + rtol=1e-03, + statistics_callback: Callable[[ErrorStatistics], None] | None = None, + ): """ Helper testing function that asserts that the model output and the reference output are equal with some tolerance. Due to numerical differences between eager mode and @@ -349,6 +360,11 @@ def _assert_outputs_equal(model_output, ref_output, atol=1e-03, rtol=1e-03): for i in range(len(model_output)): model = model_output[i] ref = ref_output[i] + + error_stats = ErrorStatistics.from_tensors(model, ref) + if statistics_callback is not None: + statistics_callback(error_stats) + assert ( ref.shape == model.shape ), f"Output {i} shape {model.shape} does not match reference output shape {ref.shape}" @@ -386,6 +402,7 @@ def _compare_outputs( atol=1e-03, rtol=1e-03, qtol=0, + statistics_callback: Callable[[ErrorStatistics], None] | None = None, ): """ Compares the original of the original nn module with the output of the generated artifact. @@ -408,6 +425,7 @@ def _compare_outputs( reference_output, atol=atol, rtol=rtol, + statistics_callback=statistics_callback, ) @staticmethod diff --git a/backends/test/harness/tests/test_error_statistics.py b/backends/test/harness/tests/test_error_statistics.py new file mode 100644 index 00000000000..fdff9c75b00 --- /dev/null +++ b/backends/test/harness/tests/test_error_statistics.py @@ -0,0 +1,65 @@ +import unittest + +import torch +from executorch.backends.test.harness.error_statistics import ErrorStatistics + + +class ErrorStatisticsTests(unittest.TestCase): + def test_error_stats_simple(self): + tensor1 = torch.tensor([1, 2, 3, 4]) + tensor2 = torch.tensor([2, 2, 2, 5]) + + error_stats = ErrorStatistics.from_tensors(tensor1, tensor2) + + # Check actual tensor statistics + self.assertEqual(error_stats.actual_stats.shape, torch.Size([4])) + self.assertEqual(error_stats.actual_stats.numel, 4) + self.assertEqual(error_stats.actual_stats.median, 2.5) + self.assertEqual(error_stats.actual_stats.mean, 2.5) + self.assertEqual(error_stats.actual_stats.max, 4) + self.assertEqual(error_stats.actual_stats.min, 1) + + # Check reference tensor statistics + self.assertEqual(error_stats.reference_stats.shape, torch.Size([4])) + self.assertEqual(error_stats.reference_stats.numel, 4) + self.assertEqual(error_stats.reference_stats.median, 2.0) + self.assertEqual(error_stats.reference_stats.mean, 2.75) + self.assertEqual(error_stats.reference_stats.max, 5) + self.assertEqual(error_stats.reference_stats.min, 2) + + # Check error statistics + self.assertAlmostEqual(error_stats.error_l2_norm, 1.732, places=3) + self.assertEqual(error_stats.error_mae, 0.75) + self.assertEqual(error_stats.error_max, 1.0) + self.assertEqual(error_stats.error_msd, -0.25) + self.assertAlmostEqual(error_stats.sqnr, 10.0, places=3) + + def test_error_stats_different_shapes(self): + # Create tensors with different shapes + tensor1 = torch.tensor([1, 2, 3, 4]) + tensor2 = torch.tensor([[2, 3], [4, 5]]) + + error_stats = ErrorStatistics.from_tensors(tensor1, tensor2) + + # Check actual tensor statistics + self.assertEqual(error_stats.actual_stats.shape, torch.Size([4])) + self.assertEqual(error_stats.actual_stats.numel, 4) + self.assertEqual(error_stats.actual_stats.median, 2.5) + self.assertEqual(error_stats.actual_stats.mean, 2.5) + self.assertEqual(error_stats.actual_stats.max, 4) + self.assertEqual(error_stats.actual_stats.min, 1) + + # Check reference tensor statistics + self.assertEqual(error_stats.reference_stats.shape, torch.Size([2, 2])) + self.assertEqual(error_stats.reference_stats.numel, 4) + self.assertEqual(error_stats.reference_stats.median, 3.5) + self.assertEqual(error_stats.reference_stats.mean, 3.5) + self.assertEqual(error_stats.reference_stats.max, 5) + self.assertEqual(error_stats.reference_stats.min, 2) + + # Check that all error values are None when shapes differ + self.assertIsNone(error_stats.error_l2_norm) + self.assertIsNone(error_stats.error_mae) + self.assertIsNone(error_stats.error_max) + self.assertIsNone(error_stats.error_msd) + self.assertIsNone(error_stats.sqnr) diff --git a/backends/test/suite/context.py b/backends/test/suite/context.py index 5f12284ae21..fd754737060 100644 --- a/backends/test/suite/context.py +++ b/backends/test/suite/context.py @@ -1,10 +1,16 @@ # Test run context management. This is used to determine the test context for reporting # purposes. class TestContext: - def __init__(self, test_name: str, flow_name: str, params: dict | None): + subtest_index: int + + def __init__( + self, test_name: str, test_base_name: str, flow_name: str, params: dict | None + ): self.test_name = test_name + self.test_base_name = test_base_name self.flow_name = flow_name self.params = params + self.subtest_index = 0 def __enter__(self): global _active_test_context diff --git a/backends/test/suite/discovery.py b/backends/test/suite/discovery.py index 92de356f550..34e588850ac 100644 --- a/backends/test/suite/discovery.py +++ b/backends/test/suite/discovery.py @@ -27,6 +27,9 @@ class TestFilter: backends: set[str] | None """ The set of backends to include. If None, all backends are included. """ + flows: set[str] | None + """ The set of test flows to include. If None, all backends are included. """ + name_regex: Pattern[str] | None """ A regular expression to filter test names. If None, all tests are included. """ @@ -86,6 +89,9 @@ def _is_test_enabled(test_case: unittest.TestCase, test_filter: TestFilter) -> b if test_filter.backends is not None and flow.backend not in test_filter.backends: return False + if test_filter.flows is not None and flow.name not in test_filter.flows: + return False + if test_filter.name_regex is not None and not test_filter.name_regex.search( test_case.id() ): diff --git a/backends/test/suite/flow.py b/backends/test/suite/flow.py index 2e2c2bf9391..b7a126eaf35 100644 --- a/backends/test/suite/flow.py +++ b/backends/test/suite/flow.py @@ -1,6 +1,6 @@ import logging -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Callable from executorch.backends.test.harness import Tester @@ -26,25 +26,38 @@ class TestFlow: tester_factory: Callable[..., Tester] """ A factory function that returns a Tester instance for this lowering flow. """ - quantize: bool = field(default=False) + quantize: bool = False """ Whether to tester should run the quantize stage on the model. """ quantize_stage_factory: Callable[..., Quantize] | None = None """ A factory function which instantiates a Quantize stage. Can be None to use the tester's default. """ + is_delegated: bool = True + """ Indicates whether the flow is expected to generate CALL_DELEGATE nodes. """ + def all_flows() -> dict[str, TestFlow]: flows = [] + from executorch.backends.test.suite.flows.portable import PORTABLE_TEST_FLOW + + flows += [ + PORTABLE_TEST_FLOW, + ] + try: from executorch.backends.test.suite.flows.xnnpack import ( + XNNPACK_DYNAMIC_INT8_PER_CHANNEL_TEST_FLOW, XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW, + XNNPACK_STATIC_INT8_PER_TENSOR_TEST_FLOW, XNNPACK_TEST_FLOW, ) flows += [ XNNPACK_TEST_FLOW, + XNNPACK_DYNAMIC_INT8_PER_CHANNEL_TEST_FLOW, XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW, + XNNPACK_STATIC_INT8_PER_TENSOR_TEST_FLOW, ] except Exception as e: logger.info(f"Skipping XNNPACK flow registration: {e}") @@ -63,12 +76,37 @@ def all_flows() -> dict[str, TestFlow]: logger.info(f"Skipping Core ML flow registration: {e}") try: - from executorch.backends.test.suite.flows.vulkan import VULKAN_TEST_FLOW + from executorch.backends.test.suite.flows.vulkan import ( + VULKAN_STATIC_INT8_PER_CHANNEL_TEST_FLOW, + VULKAN_TEST_FLOW, + ) flows += [ VULKAN_TEST_FLOW, + VULKAN_STATIC_INT8_PER_CHANNEL_TEST_FLOW, ] except Exception as e: logger.info(f"Skipping Vulkan flow registration: {e}") + try: + from executorch.backends.test.suite.flows.qualcomm import ( + QNN_16A16W_TEST_FLOW, + QNN_16A4W_BLOCK_TEST_FLOW, + QNN_16A4W_TEST_FLOW, + QNN_16A8W_TEST_FLOW, + QNN_8A8W_TEST_FLOW, + QNN_TEST_FLOW, + ) + + flows += [ + QNN_TEST_FLOW, + QNN_16A16W_TEST_FLOW, + QNN_16A8W_TEST_FLOW, + QNN_16A4W_TEST_FLOW, + QNN_16A4W_BLOCK_TEST_FLOW, + QNN_8A8W_TEST_FLOW, + ] + except Exception as e: + logger.info(f"Skipping QNN flow registration: {e}") + return {f.name: f for f in flows if f is not None} diff --git a/backends/test/suite/flows/portable.py b/backends/test/suite/flows/portable.py new file mode 100644 index 00000000000..ab176fb0e2d --- /dev/null +++ b/backends/test/suite/flows/portable.py @@ -0,0 +1,19 @@ +import logging + +from executorch.backends.test.harness import Tester +from executorch.backends.test.suite.flow import TestFlow + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +def _create_portable_flow() -> TestFlow: + return TestFlow( + "portable", + backend="portable", + tester_factory=Tester, + is_delegated=False, + ) + + +PORTABLE_TEST_FLOW = _create_portable_flow() diff --git a/backends/test/suite/flows/qualcomm.py b/backends/test/suite/flows/qualcomm.py new file mode 100644 index 00000000000..9998caa51b6 --- /dev/null +++ b/backends/test/suite/flows/qualcomm.py @@ -0,0 +1,61 @@ +from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype +from executorch.backends.qualcomm.tests.tester import QualcommTester, Quantize +from executorch.backends.test.suite.flow import TestFlow +from torchao.quantization.pt2e import MovingAverageMinMaxObserver + + +def _create_qnn_flow( + name: str, + quantize: bool = False, + quant_dtype: QuantDtype | None = None, + per_channel_conv=True, + per_channel_linear=False, + is_qat=False, + use_fp16=True, +) -> TestFlow: + if quantize and quant_dtype is None: + raise RuntimeError("Quant dtype must be provided when quantize is true.") + + def create_tester(*args, **kwargs) -> QualcommTester: + kwargs["use_fp16"] = (use_fp16,) + return QualcommTester(*args, **kwargs) + + def create_quantize_stage() -> Quantize: + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + quant_dtype, + is_qat=is_qat, + is_conv_per_channel=per_channel_conv, + is_linear_per_channel=per_channel_linear, + act_observer=MovingAverageMinMaxObserver, + ) + return Quantize(quantizer=quantizer) + + return TestFlow( + name, + backend="qualcomm", + tester_factory=create_tester, + quantize=quantize, + quantize_stage_factory=create_quantize_stage if quantize else None, + ) + + +QNN_TEST_FLOW = _create_qnn_flow("qnn") +QNN_16A16W_TEST_FLOW = _create_qnn_flow( + "qnn_16a16w", quantize=True, quant_dtype=QuantDtype.use_8a8w, use_fp16=False +) +QNN_16A8W_TEST_FLOW = _create_qnn_flow( + "qnn_16a8w", quantize=True, quant_dtype=QuantDtype.use_16a8w, use_fp16=False +) +QNN_16A4W_TEST_FLOW = _create_qnn_flow( + "qnn_16a4w", quantize=True, quant_dtype=QuantDtype.use_16a4w, use_fp16=False +) +QNN_16A4W_BLOCK_TEST_FLOW = _create_qnn_flow( + "qnn_16a4w_block", + quantize=True, + quant_dtype=QuantDtype.use_8a8w, + use_fp16=False, +) +QNN_8A8W_TEST_FLOW = _create_qnn_flow( + "qnn_8a8w", quantize=True, quant_dtype=QuantDtype.use_8a8w, use_fp16=False +) diff --git a/backends/test/suite/flows/vulkan.py b/backends/test/suite/flows/vulkan.py index 4d661efe3c7..2a8c4e506fa 100644 --- a/backends/test/suite/flows/vulkan.py +++ b/backends/test/suite/flows/vulkan.py @@ -1,17 +1,43 @@ +from typing import Callable + +from executorch.backends.test.harness.stages import Quantize from executorch.backends.test.suite.flow import TestFlow -from executorch.backends.vulkan.test.tester import VulkanTester +from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( + get_symmetric_quantization_config as get_symmetric_quantization_config_vulkan, +) +from executorch.backends.vulkan.test.tester import ( + Quantize as VulkanQuantize, + VulkanTester, +) -def _create_vulkan_flow( - name: str, - quantize: bool = False, +def _create_vulkan_flow_base( + name: str, quantize_stage_factory: Callable[..., Quantize] | None = None ) -> TestFlow: return TestFlow( name, backend="vulkan", tester_factory=VulkanTester, - quantize=quantize, + quantize=quantize_stage_factory is not None, + quantize_stage_factory=quantize_stage_factory, + ) + + +def _create_vulkan_flow() -> TestFlow: + return _create_vulkan_flow_base("vulkan") + + +def _create_vulkan_static_int8_per_channel_flow() -> TestFlow: + def create_quantize_stage() -> Quantize: + qparams = get_symmetric_quantization_config_vulkan() + return VulkanQuantize( + quantization_config=qparams, + ) + + return _create_vulkan_flow_base( + "vulkan_static_int8_per_channel", create_quantize_stage ) -VULKAN_TEST_FLOW = _create_vulkan_flow("vulkan") +VULKAN_TEST_FLOW = _create_vulkan_flow() +VULKAN_STATIC_INT8_PER_CHANNEL_TEST_FLOW = _create_vulkan_static_int8_per_channel_flow() diff --git a/backends/test/suite/flows/xnnpack.py b/backends/test/suite/flows/xnnpack.py index 9de071377ff..a181e2de711 100644 --- a/backends/test/suite/flows/xnnpack.py +++ b/backends/test/suite/flows/xnnpack.py @@ -31,6 +31,20 @@ def _create_xnnpack_flow() -> TestFlow: return _create_xnnpack_flow_base("xnnpack") +def _create_xnnpack_dynamic_int8_per_channel_flow() -> TestFlow: + def create_quantize_stage() -> Quantize: + qparams = get_symmetric_quantization_config( + is_per_channel=True, is_dynamic=True + ) + return XnnpackQuantize( + quantization_config=qparams, + ) + + return _create_xnnpack_flow_base( + "xnnpack_dynamic_int8_per_channel", create_quantize_stage + ) + + def _create_xnnpack_static_int8_per_channel_flow() -> TestFlow: def create_quantize_stage() -> Quantize: qparams = get_symmetric_quantization_config(is_per_channel=True) @@ -43,7 +57,23 @@ def create_quantize_stage() -> Quantize: ) +def _create_xnnpack_static_int8_per_tensor_flow() -> TestFlow: + def create_quantize_stage() -> Quantize: + qparams = get_symmetric_quantization_config(is_per_channel=False) + return XnnpackQuantize( + quantization_config=qparams, + ) + + return _create_xnnpack_flow_base( + "xnnpack_static_int8_per_tensor", create_quantize_stage + ) + + XNNPACK_TEST_FLOW = _create_xnnpack_flow() +XNNPACK_DYNAMIC_INT8_PER_CHANNEL_TEST_FLOW = ( + _create_xnnpack_dynamic_int8_per_channel_flow() +) XNNPACK_STATIC_INT8_PER_CHANNEL_TEST_FLOW = ( _create_xnnpack_static_int8_per_channel_flow() ) +XNNPACK_STATIC_INT8_PER_TENSOR_TEST_FLOW = _create_xnnpack_static_int8_per_tensor_flow() diff --git a/backends/test/suite/generate_markdown_summary.py b/backends/test/suite/generate_markdown_summary.py new file mode 100644 index 00000000000..37bf758fed0 --- /dev/null +++ b/backends/test/suite/generate_markdown_summary.py @@ -0,0 +1,124 @@ +import argparse +import csv +import sys + +# +# A standalone script to generate a Markdown representation of a test report. +# This is primarily intended to be used with GitHub actions to generate a nice +# representation of the test results when looking at the action run. +# +# Usage: python executorch/backends/test/suite/generate_markdown_summary.py +# Markdown is written to stdout. +# + + +def generate_markdown(csv_path: str, exit_code: int = 0): # noqa (C901) + # Print warning if exit code is non-zero + if exit_code != 0: + print("> [!WARNING]") + print( + f"> Exit code {exit_code} was non-zero. Test process may have crashed. Check the job logs for more information.\n" + ) + + with open(csv_path, newline="", encoding="utf-8") as f: + reader = csv.reader(f) + rows = list(reader) + + header = rows[0] + data_rows = rows[1:] + + # Find the Result and Result Detail column indices + result_column_index = None + result_detail_column_index = None + for i, col in enumerate(header): + if col.lower() == "result": + result_column_index = i + elif col.lower() == "result detail": + result_detail_column_index = i + + # Count results and prepare data + pass_count = 0 + fail_count = 0 + skip_count = 0 + failed_tests = [] + processed_rows = [] + result_detail_counts = {} + + for row in data_rows: + # Make a copy of the row to avoid modifying the original + processed_row = row.copy() + + # Count results and collect failed tests + if result_column_index is not None and result_column_index < len(row): + result_value = row[result_column_index].strip().lower() + if result_value == "pass": + pass_count += 1 + processed_row[result_column_index] = ( + 'Pass' + ) + elif result_value == "fail": + fail_count += 1 + processed_row[result_column_index] = ( + 'Fail' + ) + failed_tests.append(processed_row.copy()) + elif result_value == "skip": + skip_count += 1 + processed_row[result_column_index] = ( + 'Skip' + ) + + # Count result details (excluding empty ones) + if result_detail_column_index is not None and result_detail_column_index < len( + row + ): + result_detail_value = row[result_detail_column_index].strip() + if result_detail_value: # Only count non-empty result details + if result_detail_value in result_detail_counts: + result_detail_counts[result_detail_value] += 1 + else: + result_detail_counts[result_detail_value] = 1 + + processed_rows.append(processed_row) + + # Generate Summary section + total_rows = len(data_rows) + print("# Summary\n") + print(f"- **Pass**: {pass_count}/{total_rows}") + print(f"- **Fail**: {fail_count}/{total_rows}") + print(f"- **Skip**: {skip_count}/{total_rows}") + + print("## Failure Breakdown:") + total_rows_with_result_detail = sum(result_detail_counts.values()) + for detail, count in sorted(result_detail_counts.items()): + print(f"- **{detail}**: {count}/{total_rows_with_result_detail}") + + # Generate Failed Tests section + print("# Failed Tests\n") + if failed_tests: + print("| " + " | ".join(header) + " |") + print("|" + "|".join(["---"] * len(header)) + "|") + for row in failed_tests: + print("| " + " | ".join(row) + " |") + else: + print("No failed tests.\n") + + +def main(): + parser = argparse.ArgumentParser( + description="Generate a Markdown representation of a test report." + ) + parser.add_argument("csv_path", help="Path to the test report CSV file.") + parser.add_argument( + "--exit-code", type=int, default=0, help="Exit code from the test process." + ) + args = parser.parse_args() + try: + generate_markdown(args.csv_path, args.exit_code) + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/backends/test/suite/models/__init__.py b/backends/test/suite/models/__init__.py index e155e3382c5..65b546b0eb5 100644 --- a/backends/test/suite/models/__init__.py +++ b/backends/test/suite/models/__init__.py @@ -22,7 +22,6 @@ DTYPES: list[torch.dtype] = [ torch.float16, torch.float32, - torch.float64, ] @@ -42,19 +41,19 @@ def _create_test( dtype: torch.dtype, use_dynamic_shapes: bool, ): + dtype_name = str(dtype)[6:] # strip "torch." + test_name = f"{test_func.__name__}_{flow.name}_{dtype_name}" + if use_dynamic_shapes: + test_name += "_dynamic_shape" + def wrapped_test(self): params = { "dtype": dtype, "use_dynamic_shapes": use_dynamic_shapes, } - with TestContext(test_name, flow.name, params): + with TestContext(test_name, test_func.__name__, flow.name, params): test_func(self, flow, dtype, use_dynamic_shapes) - dtype_name = str(dtype)[6:] # strip "torch." - test_name = f"{test_func.__name__}_{flow.name}_{dtype_name}" - if use_dynamic_shapes: - test_name += "_dynamic_shape" - wrapped_test._name = test_func.__name__ # type: ignore wrapped_test._flow = flow # type: ignore @@ -118,6 +117,8 @@ def run_model_test( inputs, flow, context.test_name, + context.test_base_name, + 0, # subtest_index - currently unused for model tests context.params, dynamic_shapes=dynamic_shapes, ) diff --git a/backends/test/suite/operators/__init__.py b/backends/test/suite/operators/__init__.py index ec335562b39..6ceb9086f71 100644 --- a/backends/test/suite/operators/__init__.py +++ b/backends/test/suite/operators/__init__.py @@ -6,6 +6,7 @@ # pyre-unsafe +import copy import os import unittest @@ -90,12 +91,13 @@ def _expand_test(cls, test_name: str): def _make_wrapped_test( test_func: Callable, test_name: str, + test_base_name: str, flow: TestFlow, params: dict | None = None, ): def wrapped_test(self): - with TestContext(test_name, flow.name, params): - test_kwargs = params or {} + with TestContext(test_name, test_base_name, flow.name, params): + test_kwargs = copy.copy(params) or {} test_kwargs["flow"] = flow test_func(self, **test_kwargs) @@ -114,19 +116,22 @@ def _create_test_for_backend( test_type = getattr(test_func, "test_type", TestType.STANDARD) if test_type == TestType.STANDARD: - wrapped_test = _make_wrapped_test(test_func, test_func.__name__, flow) test_name = f"{test_func.__name__}_{flow.name}" + wrapped_test = _make_wrapped_test( + test_func, test_name, test_func.__name__, flow + ) setattr(cls, test_name, wrapped_test) elif test_type == TestType.DTYPE: for dtype in DTYPES: + dtype_name = str(dtype)[6:] # strip "torch." + test_name = f"{test_func.__name__}_{dtype_name}_{flow.name}" wrapped_test = _make_wrapped_test( test_func, + test_name, test_func.__name__, flow, {"dtype": dtype}, ) - dtype_name = str(dtype)[6:] # strip "torch." - test_name = f"{test_func.__name__}_{dtype_name}_{flow.name}" setattr(cls, test_name, wrapped_test) else: raise NotImplementedError(f"Unknown test type {test_type}.") @@ -146,12 +151,17 @@ def _test_op( inputs, flow, context.test_name, + context.test_base_name, + context.subtest_index, context.params, generate_random_test_inputs=generate_random_test_inputs, ) log_test_summary(run_summary) + # This is reset when a new test is started - it creates the context per-test. + context.subtest_index = context.subtest_index + 1 + if not run_summary.result.is_success(): if run_summary.result.is_backend_failure(): raise RuntimeError("Test failure.") from run_summary.error diff --git a/backends/test/suite/operators/test_adaptive_avgpool1d.py b/backends/test/suite/operators/test_adaptive_avgpool1d.py new file mode 100644 index 00000000000..f8858ecbc02 --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_avgpool1d.py @@ -0,0 +1,102 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=5, + ): + super().__init__() + self.adaptive_avgpool = torch.nn.AdaptiveAvgPool1d( + output_size=output_size, + ) + + def forward(self, x): + return self.adaptive_avgpool(x) + + +@operator_test +class AdaptiveAvgPool1d(OperatorTest): + @dtype_test + def test_adaptive_avgpool1d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, length) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 100) * 10).to(dtype),), + flow, + ) + + def test_adaptive_avgpool1d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(output_size=10), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(output_size=50), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_adaptive_avgpool1d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 100),), + flow, + ) + + def test_adaptive_avgpool1d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 50),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 200),), + flow, + ) diff --git a/backends/test/suite/operators/test_adaptive_avgpool2d.py b/backends/test/suite/operators/test_adaptive_avgpool2d.py new file mode 100644 index 00000000000..d0a456ccd9c --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_avgpool2d.py @@ -0,0 +1,112 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=(5, 5), + ): + super().__init__() + self.adaptive_avgpool = torch.nn.AdaptiveAvgPool2d( + output_size=output_size, + ) + + def forward(self, x): + return self.adaptive_avgpool(x) + + +@operator_test +class AdaptiveAvgPool2d(OperatorTest): + @dtype_test + def test_adaptive_avgpool2d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 20, 20) * 10).to(dtype),), + flow, + ) + + def test_adaptive_avgpool2d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(1, 1)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(10, 10)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(5, 10)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_adaptive_avgpool2d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 20, 20),), + flow, + ) + + def test_adaptive_avgpool2d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 10, 10),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 30, 30),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 15, 25),), + flow, + ) diff --git a/backends/test/suite/operators/test_adaptive_avgpool3d.py b/backends/test/suite/operators/test_adaptive_avgpool3d.py new file mode 100644 index 00000000000..658ded337f4 --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_avgpool3d.py @@ -0,0 +1,112 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=(4, 4, 4), + ): + super().__init__() + self.adaptive_avgpool = torch.nn.AdaptiveAvgPool3d( + output_size=output_size, + ) + + def forward(self, x): + return self.adaptive_avgpool(x) + + +@operator_test +class AdaptiveAvgPool3d(OperatorTest): + @dtype_test + def test_adaptive_avgpool3d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, depth, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 4, 8, 8, 8) * 10).to(dtype),), + flow, + ) + + def test_adaptive_avgpool3d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(1, 1, 1)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(6, 6, 6)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(2, 4, 6)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_adaptive_avgpool3d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 4, 8, 8, 8),), + flow, + ) + + def test_adaptive_avgpool3d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 2, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 6, 6, 6),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 10, 10, 10),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 7, 9, 11),), + flow, + ) diff --git a/backends/test/suite/operators/test_adaptive_maxpool1d.py b/backends/test/suite/operators/test_adaptive_maxpool1d.py new file mode 100644 index 00000000000..782bd1a5ea7 --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_maxpool1d.py @@ -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. + +# pyre-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=5, + return_indices=False, + ): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool1d( + output_size=output_size, + return_indices=return_indices, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + +@operator_test +class AdaptiveMaxPool1d(OperatorTest): + @dtype_test + def test_adaptive_maxpool1d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, length) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 100) * 10).to(dtype),), + flow, + ) + + def test_adaptive_maxpool1d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(output_size=10), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(output_size=50), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_adaptive_maxpool1d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool1d( + output_size=5, + return_indices=True, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + input_tensor = torch.randn(1, 8, 100) + + self._test_op( + ModelWithIndices(), + (input_tensor,), + flow, + ) + + def test_adaptive_maxpool1d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 100),), + flow, + ) + + def test_adaptive_maxpool1d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 50),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 200),), + flow, + ) diff --git a/backends/test/suite/operators/test_adaptive_maxpool2d.py b/backends/test/suite/operators/test_adaptive_maxpool2d.py new file mode 100644 index 00000000000..3ba98ed6c86 --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_maxpool2d.py @@ -0,0 +1,135 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=(5, 5), + return_indices=False, + ): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool2d( + output_size=output_size, + return_indices=return_indices, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + +@operator_test +class AdaptiveMaxPool2d(OperatorTest): + @dtype_test + def test_adaptive_maxpool2d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 20, 20) * 10).to(dtype),), + flow, + ) + + def test_adaptive_maxpool2d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(1, 1)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(10, 10)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(output_size=(5, 10)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_adaptive_maxpool2d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool2d( + output_size=(5, 5), + return_indices=True, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + input_tensor = torch.randn(1, 8, 20, 20) + + self._test_op( + ModelWithIndices(), + (input_tensor,), + flow, + ) + + def test_adaptive_maxpool2d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 20, 20),), + flow, + ) + + def test_adaptive_maxpool2d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 10, 10),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 30, 30),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 15, 25),), + flow, + ) diff --git a/backends/test/suite/operators/test_adaptive_maxpool3d.py b/backends/test/suite/operators/test_adaptive_maxpool3d.py new file mode 100644 index 00000000000..b2c507c12e1 --- /dev/null +++ b/backends/test/suite/operators/test_adaptive_maxpool3d.py @@ -0,0 +1,135 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + output_size=(4, 4, 4), + return_indices=False, + ): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool3d( + output_size=output_size, + return_indices=return_indices, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + +@operator_test +class AdaptiveMaxPool3d(OperatorTest): + @dtype_test + def test_adaptive_maxpool3d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, depth, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 4, 8, 8, 8) * 10).to(dtype),), + flow, + ) + + def test_adaptive_maxpool3d_output_size(self, flow: TestFlow) -> None: + # Test with different output sizes + self._test_op( + Model(output_size=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(1, 1, 1)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(6, 6, 6)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(output_size=(2, 4, 6)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_adaptive_maxpool3d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.adaptive_maxpool = torch.nn.AdaptiveMaxPool3d( + output_size=(4, 4, 4), + return_indices=True, + ) + + def forward(self, x): + return self.adaptive_maxpool(x) + + input_tensor = torch.randn(1, 4, 8, 8, 8) + + self._test_op( + ModelWithIndices(), + (input_tensor,), + flow, + ) + + def test_adaptive_maxpool3d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 4, 8, 8, 8),), + flow, + ) + + def test_adaptive_maxpool3d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 2, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 6, 6, 6),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 10, 10, 10),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 4, 7, 9, 11),), + flow, + ) diff --git a/backends/test/suite/operators/test_amax.py b/backends/test/suite/operators/test_amax.py index aff33476e69..0c9a8c06f0d 100644 --- a/backends/test/suite/operators/test_amax.py +++ b/backends/test/suite/operators/test_amax.py @@ -207,19 +207,19 @@ def test_amax_edge_cases(self, flow: TestFlow) -> None: AmaxModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AmaxModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AmaxModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) @@ -227,19 +227,19 @@ def test_amax_edge_cases(self, flow: TestFlow) -> None: AmaxModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AmaxModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AmaxModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) def test_amax_scalar(self, flow: TestFlow) -> None: diff --git a/backends/test/suite/operators/test_amin.py b/backends/test/suite/operators/test_amin.py index ab59d77d0be..f4b88b1dade 100644 --- a/backends/test/suite/operators/test_amin.py +++ b/backends/test/suite/operators/test_amin.py @@ -209,19 +209,19 @@ def test_amin_edge_cases(self, flow: TestFlow) -> None: AminModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AminModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AminModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) @@ -229,19 +229,19 @@ def test_amin_edge_cases(self, flow: TestFlow) -> None: AminModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AminModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( AminModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) def test_amin_scalar(self, flow: TestFlow) -> None: diff --git a/backends/test/suite/operators/test_argmax.py b/backends/test/suite/operators/test_argmax.py index adf1e43a340..dc8b57fc214 100644 --- a/backends/test/suite/operators/test_argmax.py +++ b/backends/test/suite/operators/test_argmax.py @@ -149,19 +149,19 @@ def test_argmax_edge_cases(self, flow: TestFlow) -> None: ArgmaxModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgmaxModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgmaxModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) @@ -169,19 +169,19 @@ def test_argmax_edge_cases(self, flow: TestFlow) -> None: ArgmaxModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgmaxModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgmaxModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([5.0]) diff --git a/backends/test/suite/operators/test_argmin.py b/backends/test/suite/operators/test_argmin.py index 0613c74a3ee..d7a24e24f5a 100644 --- a/backends/test/suite/operators/test_argmin.py +++ b/backends/test/suite/operators/test_argmin.py @@ -149,19 +149,19 @@ def test_argmin_edge_cases(self, flow: TestFlow) -> None: ArgminModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgminModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgminModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([[1.0, float("nan"), 3.0], [4.0, 5.0, float("nan")]]) @@ -169,19 +169,19 @@ def test_argmin_edge_cases(self, flow: TestFlow) -> None: ArgminModel(), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgminModel(dim=0), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) self._test_op( ArgminModel(dim=1), (x,), flow, - use_random_test_inputs=False, + generate_random_test_inputs=False, ) x = torch.tensor([5.0]) diff --git a/backends/test/suite/operators/test_avgpool1d.py b/backends/test/suite/operators/test_avgpool1d.py new file mode 100644 index 00000000000..0b2d001de01 --- /dev/null +++ b/backends/test/suite/operators/test_avgpool1d.py @@ -0,0 +1,155 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + ceil_mode=False, + count_include_pad=True, + ): + super().__init__() + self.avgpool = torch.nn.AvgPool1d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + + def forward(self, x): + return self.avgpool(x) + + +@operator_test +class AvgPool1d(OperatorTest): + @dtype_test + def test_avgpool1d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, length) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 100) * 10).to(dtype),), + flow, + ) + + def test_avgpool1d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=5), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_avgpool1d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(stride=3), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_avgpool1d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(padding=2), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_avgpool1d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_avgpool1d_count_include_pad(self, flow: TestFlow) -> None: + # Test with count_include_pad=False + self._test_op( + Model(padding=1, count_include_pad=False), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_avgpool1d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 100),), + flow, + ) + + def test_avgpool1d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 100),), + flow, + ) + + def test_avgpool1d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=2, stride=2, padding=1, count_include_pad=False), + (torch.randn(1, 8, 100),), + flow, + ) diff --git a/backends/test/suite/operators/test_avgpool2d.py b/backends/test/suite/operators/test_avgpool2d.py new file mode 100644 index 00000000000..97bcb00372a --- /dev/null +++ b/backends/test/suite/operators/test_avgpool2d.py @@ -0,0 +1,168 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + ceil_mode=False, + count_include_pad=True, + ): + super().__init__() + + # Create the avgpool layer with the given parameters + # torch.nn.AvgPool2d accepts both int and tuple types for kernel_size, stride, and padding + self.avgpool = torch.nn.AvgPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + + def forward(self, x): + return self.avgpool(x) + + +@operator_test +class AvgPool2d(OperatorTest): + @dtype_test + def test_avgpool2d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 20, 20) * 10).to(dtype),), + flow, + ) + + def test_avgpool2d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=5), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=(3, 2)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(stride=(2, 1)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(padding=(1, 2)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_count_include_pad(self, flow: TestFlow) -> None: + # Test with count_include_pad=False + self._test_op( + Model(padding=1, count_include_pad=False), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 20, 20),), + flow, + ) + + def test_avgpool2d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 20, 20),), + flow, + ) + + def test_avgpool2d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 8, 21, 21),), + flow, + ) + self._test_op( + Model( + kernel_size=(2, 3), + stride=(2, 1), + padding=(1, 0), + count_include_pad=False, + ), + (torch.randn(1, 8, 20, 20),), + flow, + ) diff --git a/backends/test/suite/operators/test_avgpool3d.py b/backends/test/suite/operators/test_avgpool3d.py new file mode 100644 index 00000000000..9e9b05907bc --- /dev/null +++ b/backends/test/suite/operators/test_avgpool3d.py @@ -0,0 +1,163 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + ceil_mode=False, + count_include_pad=True, + ): + super().__init__() + + # Create the avgpool layer with the given parameters + # torch.nn.AvgPool3d accepts both int and tuple types for kernel_size, stride, and padding + self.avgpool = torch.nn.AvgPool3d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + ceil_mode=ceil_mode, + count_include_pad=count_include_pad, + ) + + def forward(self, x): + return self.avgpool(x) + + +@operator_test +class AvgPool3d(OperatorTest): + @dtype_test + def test_avgpool3d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, depth, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 4, 8, 8, 8) * 10).to(dtype),), + flow, + ) + + def test_avgpool3d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(kernel_size=(1, 2, 2)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(stride=(1, 2, 2)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(padding=(0, 1, 1)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_count_include_pad(self, flow: TestFlow) -> None: + # Test with count_include_pad=False + self._test_op( + Model(padding=1, count_include_pad=False), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 4, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 2, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 8, 8, 8),), + flow, + ) + + def test_avgpool3d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 4, 10, 10, 10),), + flow, + ) + self._test_op( + Model( + kernel_size=(2, 2, 2), + stride=(1, 2, 2), + padding=(0, 1, 1), + count_include_pad=False, + ), + (torch.randn(1, 4, 8, 10, 10),), + flow, + ) diff --git a/backends/test/suite/operators/test_floor.py b/backends/test/suite/operators/test_floor.py index e5da5da63df..fcc834afa16 100644 --- a/backends/test/suite/operators/test_floor.py +++ b/backends/test/suite/operators/test_floor.py @@ -18,8 +18,8 @@ class FloorModel(torch.nn.Module): - def __init__(self): - super().__init__() + def forward(self, x): + return torch.floor(x) @operator_test diff --git a/backends/test/suite/operators/test_lstm.py b/backends/test/suite/operators/test_lstm.py new file mode 100644 index 00000000000..91dd73c9052 --- /dev/null +++ b/backends/test/suite/operators/test_lstm.py @@ -0,0 +1,208 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + input_size=64, + hidden_size=32, + num_layers=1, + bias=True, + batch_first=True, + dropout=0.0, + bidirectional=False, + ): + super().__init__() + self.lstm = torch.nn.LSTM( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + bias=bias, + batch_first=batch_first, + dropout=dropout, + bidirectional=bidirectional, + ) + + def forward(self, x): + return self.lstm(x)[0] # Return only the output, not the hidden states + + +@operator_test +class LSTM(OperatorTest): + @dtype_test + def test_lstm_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + Model(num_layers=2).to(dtype), + ((torch.rand(1, 10, 64) * 10).to(dtype),), # (batch=1, seq_len, input_size) + flow, + ) + + @dtype_test + def test_lstm_no_bias_dtype(self, flow: TestFlow, dtype) -> None: + self._test_op( + Model(num_layers=2, bias=False).to(dtype), + ((torch.rand(1, 10, 64) * 10).to(dtype),), + flow, + ) + + def test_lstm_feature_sizes(self, flow: TestFlow) -> None: + self._test_op( + Model(input_size=32, hidden_size=16), + (torch.randn(1, 8, 32),), # (batch=1, seq_len, input_size) + flow, + ) + self._test_op( + Model(input_size=128, hidden_size=64), + (torch.randn(1, 12, 128),), + flow, + ) + self._test_op( + Model(input_size=256, hidden_size=128), + (torch.randn(1, 6, 256),), + flow, + ) + self._test_op( + Model(input_size=16, hidden_size=32), + (torch.randn(1, 5, 16),), + flow, + ) + + def test_lstm_batch_sizes(self, flow: TestFlow) -> None: + self._test_op( + Model(), + (torch.randn(8, 10, 64),), + flow, + ) + self._test_op( + Model(), + (torch.randn(32, 10, 64),), + flow, + ) + self._test_op( + Model(), + (torch.randn(100, 10, 64),), + flow, + ) + + def test_lstm_seq_lengths(self, flow: TestFlow) -> None: + self._test_op( + Model(), + (torch.randn(1, 5, 64),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 20, 64),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 50, 64),), + flow, + ) + + def test_lstm_batch_first_false(self, flow: TestFlow) -> None: + self._test_op( + Model(batch_first=False), + (torch.randn(10, 1, 64),), # (seq_len, batch=1, input_size) + flow, + ) + + def test_lstm_num_layers(self, flow: TestFlow) -> None: + self._test_op( + Model(num_layers=2), + (torch.randn(1, 10, 64),), + flow, + ) + self._test_op( + Model(num_layers=3), + (torch.randn(1, 10, 64),), + flow, + ) + + def test_lstm_bidirectional(self, flow: TestFlow) -> None: + self._test_op( + Model(bidirectional=True), + (torch.randn(1, 10, 64),), + flow, + ) + + def test_lstm_with_dropout(self, flow: TestFlow) -> None: + # Note: Dropout is only effective with num_layers > 1 + self._test_op( + Model(num_layers=2, dropout=0.2), + (torch.randn(1, 10, 64),), + flow, + ) + + def test_lstm_with_initial_states(self, flow: TestFlow) -> None: + # Create a model that accepts initial states + class ModelWithStates(torch.nn.Module): + def __init__(self): + super().__init__() + self.lstm = torch.nn.LSTM( + input_size=64, + hidden_size=32, + num_layers=2, + batch_first=True, + ) + + def forward(self, x, h0, c0): + return self.lstm(x, (h0, c0))[0] # Return only the output + + batch_size = 1 + num_layers = 2 + hidden_size = 32 + + self._test_op( + ModelWithStates(), + ( + torch.randn(batch_size, 10, 64), # input + torch.randn(num_layers, batch_size, hidden_size), # h0 + torch.randn(num_layers, batch_size, hidden_size), # c0 + ), + flow, + ) + + def test_lstm_return_hidden_states(self, flow: TestFlow) -> None: + # Create a model that returns both output and hidden states + class ModelWithHiddenStates(torch.nn.Module): + def __init__(self): + super().__init__() + self.lstm = torch.nn.LSTM( + input_size=64, + hidden_size=32, + num_layers=2, + batch_first=True, + ) + + def forward(self, x): + # Return the complete output tuple: (output, (h_n, c_n)) + output, (h_n, c_n) = self.lstm(x) + return output, h_n, c_n + + batch_size = 1 + seq_len = 10 + input_size = 64 + + self._test_op( + ModelWithHiddenStates(), + (torch.randn(batch_size, seq_len, input_size),), + flow, + ) diff --git a/backends/test/suite/operators/test_maxpool1d.py b/backends/test/suite/operators/test_maxpool1d.py new file mode 100644 index 00000000000..e6de4dee2b7 --- /dev/null +++ b/backends/test/suite/operators/test_maxpool1d.py @@ -0,0 +1,185 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + dilation=1, + return_indices=False, + ceil_mode=False, + ): + super().__init__() + self.maxpool = torch.nn.MaxPool1d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + return_indices=return_indices, + ceil_mode=ceil_mode, + ) + + def forward(self, x): + return self.maxpool(x) + + +@operator_test +class MaxPool1d(OperatorTest): + @dtype_test + def test_maxpool1d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, length) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 100) * 10).to(dtype),), + flow, + ) + + def test_maxpool1d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=5), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_maxpool1d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(stride=3), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_maxpool1d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(padding=2), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_maxpool1d_dilation(self, flow: TestFlow) -> None: + # Test with different dilation values + self._test_op( + Model(dilation=2), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(dilation=3), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_maxpool1d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 8, 100),), + flow, + ) + + def test_maxpool1d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.maxpool = torch.nn.MaxPool1d( + kernel_size=3, + stride=2, + padding=1, + return_indices=True, + ) + + def forward(self, x): + return self.maxpool(x) + + input_tensor = torch.randn(1, 8, 100) + + self._test_op( + Model(kernel_size=3, stride=2, padding=1), + (input_tensor,), + flow, + ) + + def test_maxpool1d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 100),), + flow, + ) + + def test_maxpool1d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 100),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 100),), + flow, + ) + + def test_maxpool1d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 8, 100),), + flow, + ) + self._test_op( + Model(kernel_size=2, stride=2, padding=1, dilation=2), + (torch.randn(1, 8, 100),), + flow, + ) diff --git a/backends/test/suite/operators/test_maxpool2d.py b/backends/test/suite/operators/test_maxpool2d.py new file mode 100644 index 00000000000..f8112d3b7da --- /dev/null +++ b/backends/test/suite/operators/test_maxpool2d.py @@ -0,0 +1,191 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + dilation=1, + return_indices=False, + ceil_mode=False, + ): + super().__init__() + self.maxpool = torch.nn.MaxPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + return_indices=return_indices, + ceil_mode=ceil_mode, + ) + + def forward(self, x): + return self.maxpool(x) + + +@operator_test +class MaxPool2d(OperatorTest): + @dtype_test + def test_maxpool2d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 8, 20, 20) * 10).to(dtype),), + flow, + ) + + def test_maxpool2d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=5), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=(3, 2)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(stride=(2, 1)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(padding=(1, 2)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_dilation(self, flow: TestFlow) -> None: + # Test with different dilation values + self._test_op( + Model(dilation=2), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(dilation=(2, 1)), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.maxpool = torch.nn.MaxPool2d( + kernel_size=3, + stride=2, + padding=1, + return_indices=True, + ) + + def forward(self, x): + return self.maxpool(x) + + # Create a test input tensor + input_tensor = torch.randn(1, 8, 20, 20) + + self._test_op( + Model(kernel_size=3, stride=2, padding=1), + (input_tensor,), + flow, + ) + + def test_maxpool2d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 8, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 8, 20, 20),), + flow, + ) + + def test_maxpool2d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 4, 20, 20),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 16, 20, 20),), + flow, + ) + + def test_maxpool2d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 8, 20, 20),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 8, 21, 21),), + flow, + ) + self._test_op( + Model(kernel_size=(2, 3), stride=(2, 1), padding=(1, 0), dilation=2), + (torch.randn(1, 8, 20, 20),), + flow, + ) diff --git a/backends/test/suite/operators/test_maxpool3d.py b/backends/test/suite/operators/test_maxpool3d.py new file mode 100644 index 00000000000..3b231169371 --- /dev/null +++ b/backends/test/suite/operators/test_maxpool3d.py @@ -0,0 +1,189 @@ +# 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-unsafe + + +import torch +from executorch.backends.test.suite.flow import TestFlow + +from executorch.backends.test.suite.operators import ( + dtype_test, + operator_test, + OperatorTest, +) + + +class Model(torch.nn.Module): + def __init__( + self, + kernel_size=3, + stride=None, + padding=0, + dilation=1, + return_indices=False, + ceil_mode=False, + ): + super().__init__() + self.maxpool = torch.nn.MaxPool3d( + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + return_indices=return_indices, + ceil_mode=ceil_mode, + ) + + def forward(self, x): + return self.maxpool(x) + + +@operator_test +class MaxPool3d(OperatorTest): + @dtype_test + def test_maxpool3d_dtype(self, flow: TestFlow, dtype) -> None: + # Input shape: (batch_size, channels, depth, height, width) + self._test_op( + Model().to(dtype), + ((torch.rand(1, 4, 8, 8, 8) * 10).to(dtype),), + flow, + ) + + def test_maxpool3d_kernel_size(self, flow: TestFlow) -> None: + # Test with different kernel sizes + self._test_op( + Model(kernel_size=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(kernel_size=(1, 2, 2)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_stride(self, flow: TestFlow) -> None: + # Test with different stride values + self._test_op( + Model(stride=2), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(stride=(1, 2, 2)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_padding(self, flow: TestFlow) -> None: + # Test with different padding values + self._test_op( + Model(padding=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(padding=(0, 1, 1)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_dilation(self, flow: TestFlow) -> None: + # Test with different dilation values + self._test_op( + Model(dilation=2), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(dilation=(1, 2, 2)), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_ceil_mode(self, flow: TestFlow) -> None: + # Test with ceil_mode=True + self._test_op( + Model(ceil_mode=True), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_return_indices(self, flow: TestFlow) -> None: + # Test with return_indices=True + class ModelWithIndices(torch.nn.Module): + def __init__(self): + super().__init__() + self.maxpool = torch.nn.MaxPool3d( + kernel_size=3, + stride=2, + padding=1, + return_indices=True, + ) + + def forward(self, x): + # Return both output and indices + return self.maxpool(x) + + # Create a test input tensor + input_tensor = torch.randn(1, 4, 8, 8, 8) + + self._test_op( + Model(kernel_size=3, stride=2, padding=1), + (input_tensor,), + flow, + ) + + def test_maxpool3d_batch_sizes(self, flow: TestFlow) -> None: + # Test with batch inputs + self._test_op( + Model(), + (torch.randn(2, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(8, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(16, 4, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_input_sizes(self, flow: TestFlow) -> None: + # Test with different input sizes + self._test_op( + Model(), + (torch.randn(1, 2, 8, 8, 8),), + flow, + ) + self._test_op( + Model(), + (torch.randn(1, 8, 8, 8, 8),), + flow, + ) + + def test_maxpool3d_combinations(self, flow: TestFlow) -> None: + # Test with combinations of parameters + self._test_op( + Model(kernel_size=2, stride=2, padding=1), + (torch.randn(1, 4, 8, 8, 8),), + flow, + ) + self._test_op( + Model(kernel_size=3, stride=2, padding=1, ceil_mode=True), + (torch.randn(1, 4, 10, 10, 10),), + flow, + ) + self._test_op( + Model( + kernel_size=(2, 2, 2), stride=(1, 2, 2), padding=(0, 1, 1), dilation=2 + ), + (torch.randn(1, 4, 8, 10, 10),), + flow, + ) diff --git a/backends/test/suite/reporting.py b/backends/test/suite/reporting.py index ad32a8c74c9..ce8a48dcc12 100644 --- a/backends/test/suite/reporting.py +++ b/backends/test/suite/reporting.py @@ -1,6 +1,57 @@ +import csv + from collections import Counter -from dataclasses import dataclass +from dataclasses import dataclass, field +from datetime import timedelta from enum import IntEnum +from functools import reduce +from typing import Any, TextIO + +from executorch.backends.test.harness.error_statistics import ErrorStatistics +from torch.export import ExportedProgram + + +# The maximum number of model output tensors to log statistics for. Most model tests will +# only have one output, but some may return more than one tensor. This upper bound is needed +# upfront since the file is written progressively. Any outputs beyond these will not have stats logged. +MAX_LOGGED_MODEL_OUTPUTS = 2 + + +# Field names for the CSV report. +CSV_FIELD_NAMES = [ + "Test ID", + "Test Case", + "Subtest", + "Flow", + "Params", + "Result", + "Result Detail", + "Delegated", + "Quantize Time (s)", + "Lower Time (s)", + "Delegated Nodes", + "Undelegated Nodes", + "Delegated Ops", + "Undelegated Ops", + "PTE Size (Kb)", +] + +for i in range(MAX_LOGGED_MODEL_OUTPUTS): + CSV_FIELD_NAMES.extend( + [ + f"Output {i} Error Max", + f"Output {i} Error MAE", + f"Output {i} SNR", + ] + ) + + +# Operators that are excluded from the counts returned by count_ops. These are used to +# exclude operatations that are not logically relevant or delegatable to backends. +OP_COUNT_IGNORED_OPS = { + "executorch_call_delegate", + "getitem", +} class TestResult(IntEnum): @@ -12,50 +63,75 @@ class TestResult(IntEnum): SUCCESS_UNDELEGATED = 1 """ The test succeeded without the backend delegating anything. """ - EAGER_FAIL = 2 - """ The test failed due to the model failing to run in eager mode. """ + SKIPPED = 2 + """ The test was skipped due to a non-backend failure. """ QUANTIZE_FAIL = 3 """ The test failed due to the quantization stage failing. """ - EXPORT_FAIL = 4 - """ The test failed due to the model failing to export. """ - - LOWER_FAIL = 5 + LOWER_FAIL = 4 """ The test failed due to a failure in partitioning or lowering. """ - PTE_LOAD_FAIL = 6 + PTE_LOAD_FAIL = 5 """ The test failed due to the resulting PTE failing to load. """ - PTE_RUN_FAIL = 7 + PTE_RUN_FAIL = 6 """ The test failed due to the resulting PTE failing to run. """ - OUTPUT_MISMATCH_FAIL = 8 + OUTPUT_MISMATCH_FAIL = 7 """ The test failed due to a mismatch between runtime and reference outputs. """ - UNKNOWN_FAIL = 9 + UNKNOWN_FAIL = 8 """ The test failed in an unknown or unexpected manner. """ def is_success(self): return self in {TestResult.SUCCESS, TestResult.SUCCESS_UNDELEGATED} def is_non_backend_failure(self): - return self in {TestResult.EAGER_FAIL, TestResult.EAGER_FAIL} + return self in {TestResult.SKIPPED} def is_backend_failure(self): return not self.is_success() and not self.is_non_backend_failure() + def to_short_str(self): + if self in {TestResult.SUCCESS, TestResult.SUCCESS_UNDELEGATED}: + return "Pass" + elif self == TestResult.SKIPPED: + return "Skip" + else: + return "Fail" + + def to_detail_str(self): + if self == TestResult.SUCCESS: + return "" + elif self == TestResult.SUCCESS_UNDELEGATED: + return "" + elif self == TestResult.SKIPPED: + return "" + elif self == TestResult.QUANTIZE_FAIL: + return "Quantization Failed" + elif self == TestResult.LOWER_FAIL: + return "Lowering Failed" + elif self == TestResult.PTE_LOAD_FAIL: + return "PTE Load Failed" + elif self == TestResult.PTE_RUN_FAIL: + return "PTE Run Failed" + elif self == TestResult.OUTPUT_MISMATCH_FAIL: + return "Output Mismatch" + elif self == TestResult.UNKNOWN_FAIL: + return "Unknown Failure" + else: + raise ValueError(f"Invalid TestResult value: {self}.") + def display_name(self): if self == TestResult.SUCCESS: return "Success (Delegated)" elif self == TestResult.SUCCESS_UNDELEGATED: return "Success (Undelegated)" - elif self == TestResult.EAGER_FAIL: - return "Fail (Eager)" + elif self == TestResult.SKIPPED: + return "Skipped" elif self == TestResult.QUANTIZE_FAIL: return "Fail (Quantize)" - elif self == TestResult.EXPORT_FAIL: - return "Fail (Export)" elif self == TestResult.LOWER_FAIL: return "Fail (Lowering)" elif self == TestResult.PTE_LOAD_FAIL: @@ -76,12 +152,21 @@ class TestCaseSummary: Contains summary results for the execution of a single test case. """ - name: str - """ The qualified name of the test, not including the flow suffix. """ + backend: str + """ The name of the target backend. """ + + base_name: str + """ The base name of the test, not including flow or parameter suffixes. """ flow: str """ The backend-specific flow name. Corresponds to flows registered in backends/test/suite/__init__.py. """ + name: str + """ The full name of test, including flow and parameter suffixes. """ + + subtest_index: int + """ The subtest number. If a test case runs multiple tests, this field can be used to disambiguate. """ + params: dict | None """ Test-specific parameters, such as dtype. """ @@ -91,12 +176,46 @@ class TestCaseSummary: error: Exception | None """ The Python exception object, if any. """ + tensor_error_statistics: list[ErrorStatistics] + """ + Statistics about the error between the backend and reference outputs. Each element of this list corresponds to + a single output tensor. + """ + + quantize_time: timedelta | None = None + """ The total runtime of the quantization stage, or none, if the test did not run the quantize stage. """ + + lower_time: timedelta | None = None + """ The total runtime of the to_edge_transform_and_lower stage, or none, if the test did not run the quantize stage. """ + + delegated_op_counts: Counter | None = None + """ The number of delegated occurances of each operator in the graph. """ + + undelegated_op_counts: Counter | None = None + """ The number of undelegated occurances of each operator in the graph. """ + + pte_size_bytes: int | None = None + """ The size of the PTE file in bytes. """ + def is_delegated(self): + return ( + any(v > 0 for v in self.delegated_op_counts.values()) + if self.delegated_op_counts + else False + ) + + +@dataclass class TestSessionState: - test_case_summaries: list[TestCaseSummary] + seed: int + + # True if the CSV header has been written to report__path. + has_written_report_header: bool = False + + # The file path to write the detail report to, if enabled. + report_path: str | None = None - def __init__(self): - self.test_case_summaries = [] + test_case_summaries: list[TestCaseSummary] = field(default_factory=list) @dataclass @@ -140,11 +259,51 @@ def from_session(cls, session: TestSessionState) -> "RunSummary": _active_session: TestSessionState | None = None -def begin_test_session(): +def _get_target_name(target: Any) -> str: + """Retrieve a string representation of a node target.""" + if isinstance(target, str): + return target + elif hasattr(target, "name"): + return target.name() # Op overloads have this + elif hasattr(target, "__name__"): + return target.__name__ # Some builtins have this + else: + return str(target) + + +def _count_ops(program: ExportedProgram) -> Counter: + op_names = ( + _get_target_name(n.target) + for n in program.graph.nodes + if n.op == "call_function" + ) + + return Counter(op for op in op_names if op not in OP_COUNT_IGNORED_OPS) + + +def count_ops(program: dict[str, ExportedProgram] | ExportedProgram) -> Counter: + if isinstance(program, ExportedProgram): + return _count_ops(program) + else: + # Sum op counts for all methods in the program. + return reduce( + lambda a, b: a + b, + (_count_ops(p) for p in program.values()), + Counter(), + ) + + +def begin_test_session(report_path: str | None, seed: int): global _active_session assert _active_session is None, "A test session is already active." - _active_session = TestSessionState() + _active_session = TestSessionState(report_path=report_path, seed=seed) + + +def get_active_test_session() -> TestSessionState | None: + global _active_session + + return _active_session def log_test_summary(summary: TestCaseSummary): @@ -153,6 +312,15 @@ def log_test_summary(summary: TestCaseSummary): if _active_session is not None: _active_session.test_case_summaries.append(summary) + if _active_session.report_path is not None: + file_mode = "a" if _active_session.has_written_report_header else "w" + with open(_active_session.report_path, file_mode) as f: + if not _active_session.has_written_report_header: + write_csv_header(f) + _active_session.has_written_report_header = True + + write_csv_row(summary, f) + def complete_test_session() -> RunSummary: global _active_session @@ -162,3 +330,77 @@ def complete_test_session() -> RunSummary: _active_session = None return summary + + +def _sum_op_counts(counter: Counter | None) -> int | None: + """ + A utility function to count the total number of nodes in an op count dict. + """ + return sum(counter.values()) if counter is not None else None + + +def _serialize_params(params: dict[str, Any] | None) -> str: + if params is not None: + return str(dict(sorted(params.items()))) + else: + return "" + + +def _serialize_op_counts(counter: Counter | None) -> str: + """ + A utility function to serialize op counts to a string, for the purpose of including + in the test report. + """ + if counter is not None: + return str(dict(sorted(counter.items()))) + else: + return "" + + +def write_csv_header(output: TextIO): + writer = csv.DictWriter(output, CSV_FIELD_NAMES) + writer.writeheader() + + +def write_csv_row(record: TestCaseSummary, output: TextIO): + writer = csv.DictWriter(output, CSV_FIELD_NAMES) + + row = { + "Test ID": record.name, + "Test Case": record.base_name, + "Subtest": record.subtest_index, + "Flow": record.flow, + "Params": _serialize_params(record.params), + "Result": record.result.to_short_str(), + "Result Detail": record.result.to_detail_str(), + "Delegated": "True" if record.is_delegated() else "False", + "Quantize Time (s)": ( + f"{record.quantize_time.total_seconds():.3f}" + if record.quantize_time + else None + ), + "Lower Time (s)": ( + f"{record.lower_time.total_seconds():.3f}" if record.lower_time else None + ), + } + + for output_idx, error_stats in enumerate(record.tensor_error_statistics): + if output_idx >= MAX_LOGGED_MODEL_OUTPUTS: + print( + f"Model output stats are truncated as model has more than {MAX_LOGGED_MODEL_OUTPUTS} outputs. Consider increasing MAX_LOGGED_MODEL_OUTPUTS." + ) + break + + row[f"Output {output_idx} Error Max"] = f"{error_stats.error_max:.3f}" + row[f"Output {output_idx} Error MAE"] = f"{error_stats.error_mae:.3f}" + row[f"Output {output_idx} SNR"] = f"{error_stats.sqnr:.3f}" + + row["Delegated Nodes"] = _sum_op_counts(record.delegated_op_counts) + row["Undelegated Nodes"] = _sum_op_counts(record.undelegated_op_counts) + row["Delegated Ops"] = _serialize_op_counts(record.delegated_op_counts) + row["Undelegated Ops"] = _serialize_op_counts(record.undelegated_op_counts) + row["PTE Size (Kb)"] = ( + f"{record.pte_size_bytes / 1000.0:.3f}" if record.pte_size_bytes else "" + ) + + writer.writerow(row) diff --git a/backends/test/suite/runner.py b/backends/test/suite/runner.py index dd6e3586628..1f84db9c730 100644 --- a/backends/test/suite/runner.py +++ b/backends/test/suite/runner.py @@ -1,22 +1,39 @@ import argparse +import hashlib import importlib +import random import re +import time import unittest +from datetime import timedelta from typing import Any import torch +# Set of unsupported ops that should cause tests to be skipped +UNSUPPORTED_PORTABLE_OPS = { + "aten::_embedding_bag", + "aten::_adaptive_avg_pool2d", + "aten::median", + "aten::median.dim", + "aten::round.decimals", +} + +from executorch.backends.test.harness.error_statistics import ErrorStatistics from executorch.backends.test.harness.stages import StageType from executorch.backends.test.suite.discovery import discover_tests, TestFilter from executorch.backends.test.suite.flow import TestFlow from executorch.backends.test.suite.reporting import ( begin_test_session, complete_test_session, + count_ops, + get_active_test_session, RunSummary, TestCaseSummary, TestResult, ) +from executorch.exir import EdgeProgramManager # A list of all runnable test suites and the corresponding python package. @@ -26,11 +43,32 @@ } +def _get_test_seed(test_base_name: str) -> int: + # Set the seed based on the test base name to give consistent inputs between backends. Add the + # run seed to allow for reproducible results, but still allow for run-to-run variation. + # Having a stable hash between runs and across machines is a plus (builtin python hash is not). + # Using MD5 here because it's fast and we don't actually care about cryptographic properties. + test_session = get_active_test_session() + run_seed = ( + test_session.seed + if test_session is not None + else random.randint(0, 100_000_000) + ) + + hasher = hashlib.md5() + data = test_base_name.encode("utf-8") + hasher.update(data) + # Torch doesn't like very long seeds. + return (int.from_bytes(hasher.digest(), "little") % 100_000_000) + run_seed + + def run_test( # noqa: C901 model: torch.nn.Module, inputs: Any, flow: TestFlow, test_name: str, + test_base_name: str, + subtest_index: int, params: dict | None, dynamic_shapes: Any | None = None, generate_random_test_inputs: bool = True, @@ -40,23 +78,33 @@ def run_test( # noqa: C901 and reporting. """ + error_statistics: list[ErrorStatistics] = [] + extra_stats = {} + + torch.manual_seed(_get_test_seed(test_base_name)) + # Helper method to construct the summary. def build_result( result: TestResult, error: Exception | None = None ) -> TestCaseSummary: return TestCaseSummary( - name=test_name, + backend=flow.backend, + base_name=test_base_name, + subtest_index=subtest_index, flow=flow.name, + name=test_name, params=params, result=result, error=error, + tensor_error_statistics=error_statistics, + **extra_stats, ) # Ensure the model can run in eager mode. try: model(*inputs) except Exception as e: - return build_result(TestResult.EAGER_FAIL, e) + return build_result(TestResult.SKIPPED, e) try: tester = flow.tester_factory(model, inputs) @@ -64,11 +112,16 @@ def build_result( return build_result(TestResult.UNKNOWN_FAIL, e) if flow.quantize: + start_time = time.perf_counter() try: tester.quantize( flow.quantize_stage_factory() if flow.quantize_stage_factory else None ) + elapsed = time.perf_counter() - start_time + extra_stats["quantize_time"] = timedelta(seconds=elapsed) except Exception as e: + elapsed = time.perf_counter() - start_time + extra_stats["quantize_time"] = timedelta(seconds=elapsed) return build_result(TestResult.QUANTIZE_FAIL, e) try: @@ -77,23 +130,49 @@ def build_result( tester._get_default_stage(StageType.EXPORT, dynamic_shapes=dynamic_shapes), ) except Exception as e: - return build_result(TestResult.EXPORT_FAIL, e) + return build_result(TestResult.SKIPPED, e) + lower_start_time = time.perf_counter() try: - tester.to_edge_transform_and_lower() + tester.to_edge_transform_and_lower(generate_etrecord=True) + elapsed = time.perf_counter() - lower_start_time + extra_stats["lower_time"] = timedelta(seconds=elapsed) except Exception as e: + elapsed = time.perf_counter() - lower_start_time + extra_stats["lower_time"] = timedelta(seconds=elapsed) return build_result(TestResult.LOWER_FAIL, e) + # Compute delegation statistics. Use the ETRecord to access the edge dialect graph between + # to_edge and delegation. Note that ETRecord only stores the edge dialect graph for a single + # method currently and assumes it is called "forward". + edge_manager: EdgeProgramManager = tester.get_artifact() + edge_op_counts = count_ops({"forward": edge_manager._etrecord.edge_dialect_program}) + undelegated_op_counts = count_ops(edge_manager._edge_programs) + delegated_op_counts = edge_op_counts - undelegated_op_counts + + extra_stats["delegated_op_counts"] = delegated_op_counts + extra_stats["undelegated_op_counts"] = undelegated_op_counts + is_delegated = any( n.target == torch._higher_order_ops.executorch_call_delegate for n in tester.stages[tester.cur].graph_module.graph.nodes if n.op == "call_function" ) - # Only run the runtime portion if something was delegated. - if is_delegated: + # Check if any undelegated ops are in the unsupported ops set. + has_unsupported_ops = any( + op in UNSUPPORTED_PORTABLE_OPS for op in undelegated_op_counts.keys() + ) + + # Skip the test if there are unsupported portable ops remaining. + if has_unsupported_ops: + return build_result(TestResult.SKIPPED) + + # Only run the runtime portion if something was delegated (or the flow doesn't delegate) + if is_delegated or not flow.is_delegated: try: tester.to_executorch().serialize() + extra_stats["pte_size_bytes"] = len(tester.get_artifact()) except Exception as e: # We could introduce a result value for this, but I'm not sure it's necessary. # We can do this if we ever see to_executorch() or serialize() fail due a backend issue. @@ -104,13 +183,17 @@ def build_result( # AssertionErrors to catch output mismatches, but this might catch more than that. try: tester.run_method_and_compare_outputs( - inputs=None if generate_random_test_inputs else inputs + inputs=None if generate_random_test_inputs else inputs, + statistics_callback=lambda stats: error_statistics.append(stats), + atol=1e-1, + rtol=4e-2, ) except AssertionError as e: return build_result(TestResult.OUTPUT_MISMATCH_FAIL, e) except Exception as e: return build_result(TestResult.PTE_RUN_FAIL, e) else: + # Skip the test if nothing is delegated return build_result(TestResult.SUCCESS_UNDELEGATED) return build_result(TestResult.SUCCESS) @@ -168,15 +251,30 @@ def parse_args(): parser.add_argument( "-b", "--backend", nargs="*", help="The backend or backends to test." ) + parser.add_argument("-l", "--flow", nargs="*", help="The flow or flows to test.") parser.add_argument( "-f", "--filter", nargs="?", help="A regular expression filter for test names." ) + parser.add_argument( + "-r", + "--report", + nargs="?", + help="A file to write the test report to, in CSV format.", + default="backend_test_report.csv", + ) + parser.add_argument( + "--seed", + nargs="?", + help="The numeric seed value to use for random generation.", + type=int, + ) return parser.parse_args() def build_test_filter(args: argparse.Namespace) -> TestFilter: return TestFilter( backends=set(args.backend) if args.backend is not None else None, + flows=set(args.flow) if args.flow is not None else None, name_regex=re.compile(args.filter) if args.filter is not None else None, ) @@ -184,7 +282,10 @@ def build_test_filter(args: argparse.Namespace) -> TestFilter: def runner_main(): args = parse_args() - begin_test_session() + seed = args.seed or random.randint(0, 100_000_000) + print(f"Running with seed {seed}.") + + begin_test_session(args.report, seed=seed) if len(args.suite) > 1: raise NotImplementedError("TODO Support multiple suites.") diff --git a/backends/test/suite/tests/README.md b/backends/test/suite/tests/README.md new file mode 100644 index 00000000000..09117e1cd31 --- /dev/null +++ b/backends/test/suite/tests/README.md @@ -0,0 +1,3 @@ +# Tests + +This directory contains meta-tests for the backend test suite. As the test suite contains a non-neglible amount of logic, these tests are useful to ensure that the test suite itself is working correctly. diff --git a/backends/test/suite/tests/__init__.py b/backends/test/suite/tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/test/suite/tests/test_reporting.py b/backends/test/suite/tests/test_reporting.py new file mode 100644 index 00000000000..58ff76cba17 --- /dev/null +++ b/backends/test/suite/tests/test_reporting.py @@ -0,0 +1,141 @@ +import unittest + +from csv import DictReader +from io import StringIO + +import torch + +from executorch.exir import to_edge + +from ..reporting import ( + count_ops, + RunSummary, + TestCaseSummary, + TestResult, + TestSessionState, + write_csv_header, + write_csv_row, +) + +# Test data for simulated test results. +TEST_CASE_SUMMARIES = [ + TestCaseSummary( + backend="backend1", + base_name="test1", + flow="flow1", + name="test1_backend1_flow1", + subtest_index=0, + params=None, + result=TestResult.SUCCESS, + error=None, + tensor_error_statistics=[], + ), + TestCaseSummary( + backend="backend2", + base_name="test1", + flow="flow1", + name="test1_backend2_flow1", + subtest_index=0, + params=None, + result=TestResult.LOWER_FAIL, + error=None, + tensor_error_statistics=[], + ), + TestCaseSummary( + backend="backend1", + base_name="test2", + flow="flow1", + name="test2_backend1_flow1", + subtest_index=0, + params={"dtype": torch.float32}, + result=TestResult.SUCCESS_UNDELEGATED, + error=None, + tensor_error_statistics=[], + ), + TestCaseSummary( + backend="backend2", + base_name="test2", + flow="flow1", + name="test2_backend2_flow1", + subtest_index=0, + params={"use_dynamic_shapes": True}, + result=TestResult.SKIPPED, + error=None, + tensor_error_statistics=[], + ), +] + + +class Reporting(unittest.TestCase): + def test_csv_report_simple(self): + # Verify the format of a simple CSV run report. + session_state = TestSessionState(seed=0) + session_state.test_case_summaries.extend(TEST_CASE_SUMMARIES) + run_summary = RunSummary.from_session(session_state) + + strio = StringIO() + write_csv_header(strio) + for case_summary in run_summary.test_case_summaries: + write_csv_row(case_summary, strio) + + # Attempt to deserialize and validate the CSV report. + report = DictReader(StringIO(strio.getvalue())) + records = list(report) + self.assertEqual(len(records), 4) + + # Validate first record: test1, backend1, SUCCESS + self.assertEqual(records[0]["Test ID"], "test1_backend1_flow1") + self.assertEqual(records[0]["Test Case"], "test1") + self.assertEqual(records[0]["Flow"], "flow1") + self.assertEqual(records[0]["Result"], "Pass") + self.assertEqual(records[0]["Params"], "") + + # Validate second record: test1, backend2, LOWER_FAIL + self.assertEqual(records[1]["Test ID"], "test1_backend2_flow1") + self.assertEqual(records[1]["Test Case"], "test1") + self.assertEqual(records[1]["Flow"], "flow1") + self.assertEqual(records[1]["Result"], "Fail") + self.assertEqual(records[1]["Params"], "") + + # Validate third record: test2, backend1, SUCCESS_UNDELEGATED with dtype param + self.assertEqual(records[2]["Test ID"], "test2_backend1_flow1") + self.assertEqual(records[2]["Test Case"], "test2") + self.assertEqual(records[2]["Flow"], "flow1") + self.assertEqual(records[2]["Result"], "Pass") + self.assertEqual(records[2]["Params"], str({"dtype": torch.float32})) + + # Validate fourth record: test2, backend2, EXPORT_FAIL with use_dynamic_shapes param + self.assertEqual(records[3]["Test ID"], "test2_backend2_flow1") + self.assertEqual(records[3]["Test Case"], "test2") + self.assertEqual(records[3]["Flow"], "flow1") + self.assertEqual(records[3]["Result"], "Skip") + self.assertEqual(records[3]["Params"], str({"use_dynamic_shapes": True})) + + def test_count_ops(self): + """ + Verify that the count_ops function correctly counts operator occurances in the edge graph. + """ + + class Model1(torch.nn.Module): + def forward(self, x, y): + return x + y + + class Model2(torch.nn.Module): + def forward(self, x, y): + return x + y * y + + args = (torch.randn(2), torch.randn(2)) + ep1 = torch.export.export(Model1(), args) + ep2 = torch.export.export(Model2(), args) + + ep = to_edge({"forward1": ep1, "forward2": ep2}) + + op_counts = count_ops(ep._edge_programs) + + self.assertEqual( + op_counts, + { + "aten::add.Tensor": 2, + "aten::mul.Tensor": 1, + }, + ) diff --git a/backends/transforms/remove_clone_ops.py b/backends/transforms/remove_clone_ops.py index 2751dee2816..01fe2ee26a4 100644 --- a/backends/transforms/remove_clone_ops.py +++ b/backends/transforms/remove_clone_ops.py @@ -6,26 +6,73 @@ # pyre-strict +from typing import Set + import torch from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass +from executorch.exir.passes.remove_noop_pass import _DEQUANT_OPS, eliminate_dq_q -def remove_clone_ops(graph: torch.fx.Graph) -> torch.fx.Graph: +class RemoveCloneOpsTransform(ExportPass): """ - Remove clone op nodes and replace uses with parent node. + Trim the 'identity' operators to reduce the unnecessary copy overhead. """ - clone_op = exir_ops.edge.aten.clone.default - for node in graph.nodes: - if node.op == "call_function" and node.target == clone_op: - with graph.inserting_after(node): - node.replace_all_uses_with(node.args[0]) - graph.eliminate_dead_code() - return graph + clone_ops: Set[torch._ops.OpOverload] = { + exir_ops.edge.aten.clone.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + } + def __init__(self) -> None: + super().__init__() + + def _remove(self, graph_module: torch.fx.GraphModule) -> None: + dequant_nodes = [] + + for n in graph_module.graph.nodes: + if n.target not in self.clone_ops: + continue + + if self._is_non_identity_clone(n): + continue + + to_be_removed = n + for user_n in list(n.users.keys()): + user_n.replace_input_with(n, n.args[0]) + if n.args[0].target in _DEQUANT_OPS: + dequant_nodes += [n.args[0]] + graph_module.graph.erase_node(to_be_removed) + + eliminate_dq_q(graph_module, dequant_nodes) -class RemoveCloneOpsTransform(ExportPass): def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - graph_module.graph = remove_clone_ops(graph_module.graph) + self._remove(graph_module) + graph_module.recompile() + dead_code_elimination_pass(graph_module) return PassResult(graph_module, True) + + def _is_non_identity_clone(self, node: torch.fx.Node) -> bool: + """Return True if clone has modified memory layout or dim order.""" + + # aten.clone: check for memory_format changes + if node.target == exir_ops.edge.aten.clone.default: + memory_format = node.kwargs.get("memory_format") + if memory_format in (None, torch.preserve_format): + return False + input_meta = node.args[0].meta + return "val" in input_meta and not input_meta["val"].is_contiguous( + memory_format=memory_format + ) + + # _clone_dim_order: check for dim_order changes + if node.target == exir_ops.edge.dim_order_ops._clone_dim_order.default: + input_meta = node.args[0].meta + return ( + "val" in node.meta + and "val" in input_meta + and node.meta["val"].dim_order() != input_meta["val"].dim_order() + ) + + return False diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index ad6d93420e3..ca09d34c2fe 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -109,6 +109,7 @@ def define_common_targets(): srcs = ["remove_clone_ops.py"], visibility = [ "//executorch/backends/...", + "@EXECUTORCH_CLIENTS", ], deps = [ "//caffe2:torch", @@ -242,3 +243,16 @@ def define_common_targets(): ":rank_0_to_rank_1", ], ) + + runtime.python_test( + name = "test_remove_clone_ops", + srcs = [ + "test/test_remove_clone_ops.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + ":remove_clone_ops", + "//executorch/exir/tests:test_memory_format_ops_pass_utils", + ], + ) diff --git a/backends/transforms/test/test_create_mutable_buffer.py b/backends/transforms/test/test_create_mutable_buffer.py new file mode 100644 index 00000000000..4eb4d9538f2 --- /dev/null +++ b/backends/transforms/test/test_create_mutable_buffer.py @@ -0,0 +1,503 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Copyright 2025 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 copy +import unittest + +import executorch +import torch +from executorch.backends.transforms.utils import create_mutable_buffer +from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.extension.pybindings.portable_lib import ( # @manual + _load_for_executorch_from_buffer, + Verification, +) +from torch.export import export +from torch.utils._pytree import tree_flatten + + +class TestMutableBufferCreation(unittest.TestCase): + """ + Test suite for the create_mutable_buffer utility function. + """ + + def test_create_mutable_buffer(self): + """ + Tests the utility function create_mutable_buffer which creates a mutable buffer + that can be modified during execution. + """ + + class EmptyNetwork(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + test_data: torch.Tensor = (torch.zeros(1),) + + module = EmptyNetwork() + exported_program = export(module, args=module.test_data, strict=True) + exported_program = to_edge(exported_program).exported_program() + graph = exported_program.graph_module.graph + + assert len(graph.nodes) == 2 + assert exported_program.module()(torch.zeros(1)) == 0 + assert len(exported_program.graph_signature.input_specs) == 1 + assert len(exported_program.graph_signature.output_specs) == 1 + assert len(exported_program.state_dict) == 0 + + buffer_name = "b_test_mutable_buffer" + target_name = buffer_name[2:] # Remove the "b_" prefix + initial_data = torch.ones(1) * 2 # Initialize with value 2 + + # Create a mutable buffer using create_mutable_buffer + buffer_node = create_mutable_buffer( + exp_program=exported_program, + name=buffer_name, + data=initial_data, + ) + assert "val" in buffer_node.meta + + # Verify the buffer was created correctly + input_node = list(graph.nodes)[ + 1 + ] # Original input node (buffer_node is now first) + + # Create an add operation that uses the mutable buffer + with graph.inserting_after(input_node): + graph.create_node( + "call_function", + exir_ops.edge.aten.add.Tensor, + args=(input_node, buffer_node), + kwargs={}, + ) + + # We should now have four nodes: buffer, input, add, output + assert len(graph.nodes) == 4 + + assert target_name in exported_program.state_dict + assert torch.equal(exported_program.state_dict[target_name], initial_data) + + # Check that buffer is properly referenced in graph signature + assert buffer_name in exported_program.graph_signature.inputs_to_buffers + assert ( + exported_program.graph_signature.inputs_to_buffers[buffer_name] + == target_name + ) + assert ( + exported_program.graph_signature.buffers_to_mutate[buffer_name] + == target_name + ) + + +class TestRegisterMutableBufferPass(unittest.TestCase): + """ + This test is to test the `create_mutable_buffer` utility. + """ + + class HelperPass(ExportPass): + def __init__( + self, + exported_program: torch.export.ExportedProgram, + buf_data: torch.Tensor, + ): + super().__init__() + self.registered_buffers = set() + self.exported_program = exported_program + self.buf_data = buf_data + + def call(self, graph_module: torch.fx.GraphModule): + """ + This pass will register a mutable buffer for the add op(s) in the graph. + It will insert a new index_put_ op to the graph to update the buffer using the output of the add op. + And adjust the output of the graph to return the buffer or the index_put_ op. + """ + modified = False + + assert ( + len(graph_module.graph.output_node().args[0]) == 1 + ), "For this pass, expecting only one output i.e. add" + + for node in graph_module.graph.nodes: + if ( + node.op == "call_function" + and node.target == exir_ops.edge.aten.add.Tensor + ): + # To match what we do in export. + suffix = ( + "" + if len(self.registered_buffers) == 0 + else f"{len(self.registered_buffers) + 1}" + ) + buffer_name = f"b_my_buffer{suffix}" + self.registered_buffers.add(buffer_name) + + # Utility under test! + buf_node = create_mutable_buffer( + self.exported_program, + data=self.buf_data, + name=buffer_name, + ) + + # Assuming `indices` is always available + indices = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.name == "indices" + ] + with graph_module.graph.inserting_after(node): + index_put_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.aten.index_put_.default, + ( + buf_node, + indices, + node, + ), + {}, + ) + + # Replace the old add output with index_put_node + output_node = graph_module.graph.output_node() + outputs = list(output_node.args[0]) + if node in outputs: + outputs[-1] = ( + index_put_node # Replace the old add output with index_put_node + ) + output_node.args = (outputs,) + + # update the output node name in the graph signature + graph_signature = self.exported_program.graph_signature + graph_signature.replace_all_uses( + node.name, index_put_node.name + ) + self.exported_program._graph_signature = graph_signature + + modified = True + + if modified: + graph_module = super().call(graph_module).graph_module + graph_module.graph.lint() + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + def _test_edge_pass(self, model, example_inputs, num_lifted_args=1): + exported = export(model, example_inputs) + + edge_program = to_edge( + exported, + # for torch._export.verifier.SpecViolationError: + # Operator torch._ops.aten.index_put_.default is not in Core ATen opset + compile_config=executorch.exir.EdgeCompileConfig(_check_ir_validity=False), + ) + # Tensor data for the new buffer(s) + buffer_tensor = torch.zeros(2) + + transformed_example_inputs = ( + *(buffer_tensor for i in range(num_lifted_args)), + *example_inputs, + ) + transformed_ep = edge_program.transform( + passes=[ + TestRegisterMutableBufferPass.HelperPass( + edge_program.exported_program(), buffer_tensor + ) + ] + ) + transformed_edge_gm = transformed_ep.exported_program().graph_module + # Make sure it works + transformed_edge_gm(*transformed_example_inputs) + + # Explicitly passing inplace_pass to make sure it works with our manually inserted buffer and index_put_ node(s). + executorch_program_manager = edge_program.to_executorch( + ExecutorchBackendConfig( + emit_mutable_buffer_names=True, run_reinplace_pass=True + ) + ) + return executorch_program_manager + + def _test_eager(self, model, example_inputs, num_lifted_args=1): + exported = export(model, example_inputs, strict=True) + + edge_program = executorch.exir.to_edge(exported) + edge_gm = edge_program.exported_program().graph_module + + # Adding buffer as an extra pos[0] arg + buffer_tensor = torch.zeros(2) + edge_example_inputs = ( + *(buffer_tensor for i in range(num_lifted_args)), + *example_inputs, + ) + # Make sure it works + edge_gm(*edge_example_inputs) + + executorch_program_manager = edge_program.to_executorch( + ExecutorchBackendConfig( + emit_mutable_buffer_names=True, run_reinplace_pass=True + ) + ) + return executorch_program_manager + + def _compare_outputs(self, et_ep1, et_ep2, example_inputs): + + def run(et_pm, inputs): + buffer = et_pm.buffer + inputs_flattened, _ = tree_flatten(inputs) + executorch_module = _load_for_executorch_from_buffer( + buffer, program_verification=Verification.Minimal + ) + executorch_output = copy.deepcopy( + executorch_module.run_method("forward", tuple(inputs_flattened)) + ) + return executorch_output + + # compare the outputs of the two programs + output1 = run(et_ep1, example_inputs) + output2 = run(et_ep2, example_inputs) + assert len(output1) == len(output2) + for o1, o2 in zip(output1, output2): + self.assertTrue(torch.allclose(o1, o2)) + + def _compare_ep_state_dict(self, et_ep1, et_ep2): + # compare the state dict of the two programs + state_dict1 = et_ep1.exported_program().state_dict + state_dict2 = et_ep2.exported_program().state_dict + self.assertEqual(len(state_dict1), len(state_dict2)) + # a bit fragile comparing the names, but the util tries to match the names i.e. `b_my_buffer` and `b_my_buffer2` + for k, _ in state_dict1.items(): + self.assertTrue( + k in state_dict2, f"{state_dict1.keys()} != {state_dict2.keys()} @ {k}" + ) + + def _compare_signatures(self, et_ep1, et_ep2): + # compare the graph signatures + def _input_spec_compare(input_spec1, input_spec2): + self.assertEqual( + input_spec1.kind, + input_spec2.kind, + f"{input_spec1.kind} != {input_spec2.kind}", + ) + self.assertEqual(input_spec1.arg, input_spec2.arg) + self.assertEqual( + input_spec1.target, + input_spec2.target, + f"{input_spec1.target} != {input_spec2.target}", + ) + self.assertEqual( + input_spec1.persistent, + input_spec2.persistent, + f"{input_spec1.persistent} != {input_spec2.persistent}", + ) + + def _output_spec_compare(output_spec1, output_spec2): + self.assertEqual( + output_spec1.kind, + output_spec2.kind, + f"{output_spec1.kind} != {output_spec2.kind}", + ) + # TODO: Look into why the output names are different, + # and not updated by the buffer_write_back_pass when the buffer + # is inserted via a pass and used by an inplace op. + # self.assertEqual(output_spec1.arg, output_spec2.arg) + self.assertEqual( + output_spec1.target, + output_spec2.target, + f"{output_spec1.target} != {output_spec2.target}", + ) + + graph_signature1 = et_ep1.exported_program().graph_signature + graph_signature2 = et_ep2.exported_program().graph_signature + + # compare input spec order, kind and targets + self.assertEqual( + len(graph_signature1.input_specs), len(graph_signature2.input_specs) + ) + for i1, i2 in zip(graph_signature1.input_specs, graph_signature2.input_specs): + _input_spec_compare(i1, i2) + + # compare output spec order, kind and targets + self.assertEqual( + len(graph_signature1.output_specs), len(graph_signature2.output_specs) + ) + for o1, o2 in zip(graph_signature1.output_specs, graph_signature2.output_specs): + _output_spec_compare(o1, o2) + + def _compare_plan_ops(self, et1, et2): + operators1 = et1.executorch_program.execution_plan[0].operators + operators2 = et2.executorch_program.execution_plan[0].operators + for op1, op2 in zip(operators1, operators2): + self.assertEqual(op1.name, op2.name, f"{op1.name} != {op2.name}") + self.assertEqual( + op1.overload, op2.overload, f"{op1.overload} != {op2.overload}" + ) + + def compare(self, et1, et2, example_inputs): + self._compare_signatures(et1, et2) + self._compare_ep_state_dict(et1, et2) + self._compare_plan_ops(et1, et2) + self._compare_outputs(et1, et2, example_inputs) + + def test_basic(self): + class CustomModuleGraph(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, indices): + output = x + x + return output + + class CustomModuleSrc(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("my_buffer", torch.zeros(2)) + + def forward(self, x, indices): + output = x + x + self.my_buffer.index_put_((indices,), output) + return output + + example_inputs = (torch.ones(2), torch.tensor([0, 1])) + + with torch.no_grad(): + graph_model = CustomModuleGraph().eval() + et_1 = self._test_edge_pass(graph_model, example_inputs) + + src_model = CustomModuleSrc().eval() + et_2 = self._test_eager(src_model, example_inputs) + + self.compare(et_1, et_2, example_inputs) + + def test_basic_with_param(self): + input_tensor = torch.ones(2) + + class CustomModuleGraph(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_parameter("my_param", torch.nn.Parameter(input_tensor)) + + def forward(self, x, indices): + output = x + self.my_param + return output + + class CustomModuleSrc(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("my_buffer", torch.zeros(2)) + self.register_parameter("my_param", torch.nn.Parameter(input_tensor)) + + def forward(self, x, indices): + output = x + self.my_param + self.my_buffer.index_put_((indices,), output) + return output + + example_inputs = (input_tensor, torch.tensor([0, 1])) + with torch.no_grad(): + graph_model = CustomModuleGraph().eval() + et1 = self._test_edge_pass(graph_model, example_inputs, num_lifted_args=2) + + src_model = CustomModuleSrc().eval() + et2 = self._test_eager(src_model, example_inputs, num_lifted_args=2) + + self.compare(et1, et2, example_inputs) + + def test_basic_with_constant(self): + input_tensor = torch.ones(2) * 2 + + class CustomModuleGraph(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.constant = input_tensor + + def forward(self, x, indices): + output = x + self.constant + return output + + class CustomModuleSrc(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("my_buffer", torch.zeros(2)) + self.constant = input_tensor + + def forward(self, x, indices): + output = x + self.constant + self.my_buffer.index_put_((indices,), output) + return output + + example_inputs = (torch.ones(2), torch.tensor([0, 1])) + with torch.no_grad(): + graph_model = CustomModuleGraph().eval() + et1 = self._test_edge_pass(graph_model, example_inputs, num_lifted_args=2) + + src_model = CustomModuleSrc().eval() + et2 = self._test_eager(src_model, example_inputs, num_lifted_args=2) + + self.compare(et1, et2, example_inputs) + + def test_single(self): + class CustomModuleGraph(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, indices): + output = x + x + return output + + class CustomModuleSrc(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("my_buffer", torch.zeros(2)) + + def forward(self, x, indices): + output = x + x + self.my_buffer.index_put_((indices,), output) + return self.my_buffer + + example_inputs = (torch.ones(2), torch.tensor([0, 1])) + + graph_model = CustomModuleGraph().eval() + et1 = self._test_edge_pass(graph_model, example_inputs) + + src_model = CustomModuleSrc().eval() + et2 = self._test_eager(src_model, example_inputs) + + self.compare(et1, et2, example_inputs) + + def test_double(self): + class CustomModuleGraph(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x, indices): + output = x + x + output = output + x + return output + + class CustomModuleSrc(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("my_buffer", torch.zeros(2)) + self.register_buffer("my_buffer2", torch.zeros(2)) + + def forward(self, x, indices): + output = x + x + self.my_buffer.index_put_((indices,), output) + + output = output + x + self.my_buffer2.index_put_((indices,), output) + + return output + + example_inputs = (torch.ones(2), torch.tensor([0, 1])) + + graph_model = CustomModuleGraph().eval() + et1 = self._test_edge_pass(graph_model, example_inputs, num_lifted_args=2) + + src_model = CustomModuleSrc().eval() + et2 = self._test_eager(src_model, example_inputs, num_lifted_args=2) + + self.compare(et1, et2, example_inputs) diff --git a/backends/transforms/test/test_duplicate_dynamic_quant_chain.py b/backends/transforms/test/test_duplicate_dynamic_quant_chain.py index 79bc56f8780..4d4b5d8cd5a 100644 --- a/backends/transforms/test/test_duplicate_dynamic_quant_chain.py +++ b/backends/transforms/test/test_duplicate_dynamic_quant_chain.py @@ -58,7 +58,7 @@ def _test_duplicate_chain( # program capture m = copy.deepcopy(m_eager) - m = torch.export.export_for_training(m, example_inputs, strict=True).module() + m = torch.export.export(m, example_inputs, strict=True).module() m = prepare_pt2e(m, quantizer) # Calibrate diff --git a/backends/transforms/test/test_remove_clone_ops.py b/backends/transforms/test/test_remove_clone_ops.py new file mode 100644 index 00000000000..d34c522baaa --- /dev/null +++ b/backends/transforms/test/test_remove_clone_ops.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. + +import unittest + +import torch +from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dim_order_utils import is_channel_last_dim_order +from executorch.exir.tests.test_memory_format_ops_pass_utils import ( + SimpleCloneChannelsLastModule, +) +from torch.export import export +from torch.fx import GraphModule +from torch.testing import FileCheck +from torch.testing._internal.common_utils import TestCase + + +class TestRemoveCloneOpsTransform(TestCase): + # Clone ops can appear as either aten.clone or _clone_dim_order depending on the _skip_dim_order flag. + # _skip_dim_order=True tests aten.clone + # _skip_dim_order=False tests _clone_dim_order + CLONE_OP_CASES = [ + (True, "executorch_exir_dialects_edge__ops_aten_clone_default"), + ( + False, + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", + ), + ] + + def test_dq_clone_q_linear(self): + """ + Test RemoveCloneOpsTransform on a graph with d/q -> clone -> q -> linear pattern + + Before: Should contain all nodes + After: Should only have the linear operation + """ + + # Create a graph module directly with the pattern: quant -> clone -> dequant -> fp linear + class TestModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + # This will be replaced with our custom graph + return self.linear(x) + + # Create a module instance + module = TestModule() + + # Create a new graph with our desired pattern + graph = torch.fx.Graph() + + # Add placeholders + input_node = graph.placeholder("x") + + # Create nodes for our pattern: quant -> clone -> dequant -> fp linear + # Constants for quantization parameters + scale = graph.create_node( + "call_function", torch.tensor, args=([0.1],), kwargs={} + ) + zero_point = graph.create_node( + "call_function", torch.tensor, args=([0],), kwargs={} + ) + + # Dequantize node + dequant_node = graph.create_node( + "call_function", + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + args=(input_node, scale, zero_point, torch.int8), + kwargs={}, + ) + + # Clone node. + # Use Edge op as this is an executorch pass + clone_node = graph.create_node( + "call_function", + exir_ops.edge.aten.clone.default, + args=(dequant_node,), + kwargs={}, + ) + + # Quantize node + quant_node = graph.create_node( + "call_function", + torch.ops.quantized_decomposed.quantize_per_tensor.default, + args=(clone_node, scale, zero_point, torch.int8), + kwargs={}, + ) + + # Linear node (using the module's linear layer) + # Technically, should use quantized weight and bias + # but we are just inspecting graph patterns in this test + weight = graph.create_node("get_attr", "linear.weight") + bias = graph.create_node("get_attr", "linear.bias") + linear_node = graph.create_node( + "call_function", + torch.nn.functional.linear, + args=(quant_node, weight, bias), + kwargs={}, + ) + + # Output + graph.output(linear_node) + + # Create a GraphModule with our custom graph + gm = GraphModule(module, graph) + + # Verify we have the expected nodes before transformation using FileCheck + FileCheck().check( + "torch.ops.quantized_decomposed.dequantize_per_tensor.default", + ).check( + "executorch_exir_dialects_edge__ops_aten_clone_default", + ).check( + "torch.ops.quantized_decomposed.quantize_per_tensor.default", + ).check( + "torch._C._nn.linear", + ).run( + gm.code + ) + + # Apply the transform + transformed_gm = RemoveCloneOpsTransform()(gm).graph_module + + # Verify the dq -> clone -> q pattern is removed and linear op is still present using FileCheck + FileCheck().check_not( + "executorch_exir_dialects_edge__ops_aten_clone_default" + ).check_not("quantized_decomposed.dequantize_per_tensor.default").check_not( + "quantized_decomposed.quantize_per_tensor.default" + ).check_count( + "torch._C._nn.linear", + 1, + exactly=True, + ).run( + transformed_gm.code + ) + + def test_clone_non_identity_survives(self): + """Verify clone ops that modify memory_format are preserved by RemoveCloneOpsTransform.""" + + for skip_dim_order, clone_op_str in self.CLONE_OP_CASES: + model = SimpleCloneChannelsLastModule() + x = torch.randn(3, 4, 5, 6).to(memory_format=torch.contiguous_format) + + exported = export(model.eval(), (x,), strict=True) + before_epm = to_edge( + exported, + compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order), + ) + + updated_epm = before_epm.transform([RemoveCloneOpsTransform()]) + + FileCheck().check_count(clone_op_str, 1, exactly=True).run( + updated_epm.exported_program().graph_module.code + ) + + expected = before_epm.exported_program().module()(x) + actual = updated_epm.exported_program().module()(x) + assert torch.allclose(actual, expected) + assert is_channel_last_dim_order(actual) + + def test_clone_identity_removed(self): + """Verify identity clone ops are removed by RemoveCloneOpsTransform.""" + + for skip_dim_order, clone_op_str in self.CLONE_OP_CASES: + model = SimpleCloneChannelsLastModule() + x = torch.randn(3, 4, 5, 6).to(memory_format=torch.channels_last) + + exported = export(model.eval(), (x,), strict=True) + before_epm = to_edge( + exported, + compile_config=EdgeCompileConfig(_skip_dim_order=skip_dim_order), + ) + + FileCheck().check_count(clone_op_str, 1, exactly=True).run( + before_epm.exported_program().graph_module.code + ) + + updated_epm = before_epm.transform([RemoveCloneOpsTransform()]) + + FileCheck().check_not(clone_op_str).run( + updated_epm.exported_program().graph_module.code + ) + + expected = before_epm.exported_program().module()(x) + actual = updated_epm.exported_program().module()(x) + assert torch.allclose(actual, expected) + assert is_channel_last_dim_order(actual) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 4e451928ee4..ef19a937b0b 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -22,10 +22,43 @@ ExportGraphSignature, InputKind, InputSpec, + OutputKind, + OutputSpec, TensorArgument, ) +def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Tensor: + """ + Helper function to create a fake tensor using the fake_mode from existing nodes in the graph. + + Args: + graph: The graph to get fake_mode from + data: The tensor data to create fake tensor for + + Returns: + A fake tensor with the appropriate fake_mode + + Raises: + RuntimeError: If the graph has no nodes to extract fake_mode from + """ + nodes = list(graph.nodes) + if not nodes: + raise RuntimeError( + "Cannot create fake tensor: graph has no nodes to extract fake_mode from" + ) + + example_node = nodes[0] + if isinstance( + example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list) + ): + example_fake_tensor = example_node.meta["val"][0] + else: + example_fake_tensor = example_node.meta["val"] + + return FakeTensorConverter().from_real_tensor(example_fake_tensor.fake_mode, t=data) + + def is_get_attr_node(node: torch.fx.Node) -> bool: """ Returns true if the given node is a get attr node for a tensor of the model @@ -98,17 +131,7 @@ def create_constant_placeholder( case _: raise RuntimeError("Can only create constant input nodes.") - # Create fake tensor using the same fake_mode as the other fake tensors in the graph - example_node = list(graph.nodes)[0] - if isinstance( - example_node.meta["val"], (tuple, torch.fx.immutable_collections.immutable_list) - ): - example_fake_tensor = example_node.meta["val"][0] - else: - example_fake_tensor = example_node.meta["val"] - fake_tensor = FakeTensorConverter().from_real_tensor( - example_fake_tensor.fake_mode, t=data - ) + fake_tensor = _get_fake_tensor_mode(graph, data) # Create node node = graph.create_node(op="placeholder", name=name, target=name) @@ -187,3 +210,168 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod # Remove node from graph node.graph.erase_node(node) + + +def _validate_graph_signature(exp_program: ExportedProgram): + """ + Validates that the graph signature is up to date with the graph. + """ + placeholders = [n for n in exp_program.graph.nodes if n.op == "placeholder"] + if len(placeholders) != len(exp_program.graph_signature.input_specs): + raise RuntimeError( + f"Graph has {len(placeholders)} placeholder nodes but signature has " + f"{len(exp_program.graph_signature.input_specs)} input specs" + ) + for node, input_spec in zip(placeholders, exp_program.graph_signature.input_specs): + if node.name != input_spec.arg.name: + raise RuntimeError( + f"Input node {node.name} does not match input spec {input_spec.arg.name}" + ) + outputs = exp_program.graph.output_node().args[0] + if len(outputs) != len(exp_program.graph_signature.output_specs): + raise RuntimeError( + f"Graph has {len(outputs)} output nodes but signature has " + f"{len(exp_program.graph_signature.output_specs)} output specs" + ) + for node, output_spec in zip(outputs, exp_program.graph_signature.output_specs): + if node.name != output_spec.arg.name: + raise RuntimeError( + f"Output node {node.name} does not match output spec {output_spec.arg.name}" + ) + + +def _spec_to_node( + exp_program: ExportedProgram, spec: InputSpec | OutputSpec +) -> torch.fx.Node: + """ + Converts an InputSpec or OutputSpec to its corresponding node in the graph. + """ + # Extract the argument name from the spec + if hasattr(spec, "arg") and hasattr(spec.arg, "name"): + arg_name = spec.arg.name + else: + raise RuntimeError(f"Invalid spec format: {spec}") + + # Find the corresponding node in the graph + for node in exp_program.graph.nodes: + if node.name == arg_name: + return node + + raise RuntimeError(f"Could not find node with name '{arg_name}' in the graph") + + +def create_mutable_buffer( + exp_program: ExportedProgram, + name: str, + data: torch.Tensor, +) -> torch.fx.Node: + """ + Creates and returns a mutable buffer placeholder node. This is similar to + create_constant_placeholder but specifically for creating mutable buffers that + can be modified during execution. + + The difference between this and create_constant_placeholder is that this doesn't + expect user to set the correct position for the placeholder node to be inserted, + it finds the correct position automatically. + + It also updates the graph outputs to include the mutable buffer. + + Args: + exp_program: The exported program to modify + name: The name for the new buffer node (should start with "b_" prefix by convention) + data: The initial tensor data for the buffer + + Returns: + The created placeholder node to be used in the graph + """ + # Input validation + if not name or not name.strip(): + raise ValueError("Buffer name cannot be empty") + + if not isinstance(data, torch.Tensor): + raise ValueError("Data must be a torch.Tensor") + + # Extract target name (remove "b_" prefix if present, following export convention) + if name.startswith("b_"): + target = name[2:] + else: + target = name + + # Check if target already exists + if target in exp_program.state_dict: + raise RuntimeError(f"Buffer target '{target}' already exists in state_dict") + + _validate_graph_signature(exp_program) + + persistent_buffer = True + exp_program.state_dict[target] = data + + graph = exp_program.graph_module.graph + + # Create fake tensor using helper function + fake_tensor = _get_fake_tensor_mode(graph, data) + + # Signature ordering is as follows: + # Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs] + # ^^^^^^^ + # insert here (at the end of buffers) + # Outputs = [*mutated_inputs, *flattened_user_outputs] + # ^^^^^^^^^^^^^^^ + # insert here (at the end of mutated inputs) + + # Inputs + # Find const or user input node if any, and insert before it + node_index = 0 + node = None + + input_specs = exp_program.graph_signature.input_specs + if len(input_specs) == 0 or all( + spec.kind not in [InputKind.CONSTANT_TENSOR, InputKind.USER_INPUT] + for spec in input_specs + ): + # No const or user input nodes + node_index = len(input_specs) + node = graph.create_node(op="placeholder", name=name, target=name) + else: + # Find the first constant or user input node + for i, spec in enumerate(input_specs): + if spec.kind in [InputKind.CONSTANT_TENSOR, InputKind.USER_INPUT]: + node_index = i + with graph.inserting_before(_spec_to_node(exp_program, spec)): + node = graph.create_node(op="placeholder", name=name, target=name) + break + + assert node is not None, "node should be created at this point" + node.meta["val"] = fake_tensor + buffer_input_spec = InputSpec( + InputKind.BUFFER, TensorArgument(name), target, persistent_buffer + ) + input_specs.insert(node_index, buffer_input_spec) + + # Outputs + # Create output spec for the mutable buffer, and insert it at the beginning of output specs + user_output_indices = [ + i + for i, spec in enumerate(exp_program.graph_signature.output_specs) + if spec.kind == OutputKind.USER_OUTPUT + ] + + output_index = user_output_indices[0] if user_output_indices else 0 + + output_specs = exp_program.graph_signature.output_specs + mutation_output_spec = OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name), target + ) + output_specs.insert(output_index, mutation_output_spec) + + # Update the outputs to include the mutable buffer + output_node = graph.output_node() + args = list(output_node.args[0]) + args.insert(output_index, node) + output_node.args = (args,) + + # Update graph signature in the exported program + new_graph_signature = ExportGraphSignature(input_specs, output_specs) + exp_program._graph_signature = new_graph_signature + + return node diff --git a/backends/vulkan/CMakeLists.txt b/backends/vulkan/CMakeLists.txt index 72d5fb8d830..29ff90e7293 100644 --- a/backends/vulkan/CMakeLists.txt +++ b/backends/vulkan/CMakeLists.txt @@ -101,7 +101,7 @@ set_target_properties(vulkan_schema PROPERTIES LINKER_LANGUAGE CXX) target_include_directories( vulkan_schema INTERFACE - ${SCHEMA_INCLUDE_DIR} + $ $ ) diff --git a/backends/vulkan/_passes/TARGETS b/backends/vulkan/_passes/TARGETS index cfe20892994..aed41114ada 100644 --- a/backends/vulkan/_passes/TARGETS +++ b/backends/vulkan/_passes/TARGETS @@ -34,20 +34,6 @@ runtime.python_library( ], ) -runtime.python_library( - name = "int4_weight_only_quantizer", - srcs = [ - "int4_weight_only_quantizer.py", - ], - visibility = [ - "//executorch/backends/...", - ], - deps = [ - "//executorch/backends/vulkan:custom_ops_lib", - "//pytorch/ao:torchao", - ] -) - runtime.python_library( name = "squeeze_unsqueeze_inputs", srcs = [ @@ -118,6 +104,35 @@ runtime.python_library( ], ) +runtime.python_library( + name = "fold_qdq", + srcs = ["fold_qdq.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:utils_lib", + "//executorch/exir:pass_base", + ], +) + +runtime.python_library( + name = "fuse_patterns", + srcs = ["fuse_patterns.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/patterns:vulkan_patterns", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + typing = True, +) + runtime.python_library( name = "vulkan_passes", srcs = [ @@ -128,9 +143,10 @@ runtime.python_library( "//executorch/examples/...", ], deps = [ + ":fold_qdq", + ":fuse_patterns", ":fuse_quantized_ops", ":insert_prepack_nodes", - ":int4_weight_only_quantizer", ":remove_asserts", ":remove_local_scalar_dense", ":remove_redundant_ops", diff --git a/backends/vulkan/_passes/__init__.py b/backends/vulkan/_passes/__init__.py index 7ff93a6ee38..f4ef6b2ac0e 100644 --- a/backends/vulkan/_passes/__init__.py +++ b/backends/vulkan/_passes/__init__.py @@ -6,13 +6,12 @@ # pyre-strict +from executorch.backends.vulkan._passes.fold_qdq import FoldQDQPass +from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.backends.vulkan._passes.fuse_quantized_ops import ( FuseQuantizedOpsTransform, ) from executorch.backends.vulkan._passes.insert_prepack_nodes import insert_prepack_nodes -from executorch.backends.vulkan._passes.int4_weight_only_quantizer import ( - VkInt4WeightOnlyQuantizer, -) from executorch.backends.vulkan._passes.remove_asserts import ( remove_asserts, RemoveAssertsTransform, @@ -29,9 +28,10 @@ from executorch.backends.vulkan._passes.tag_memory_meta_pass import TagMemoryMetaPass __all__ = [ + "FoldQDQPass", + "FusePatternsPass", "FuseQuantizedOpsTransform", "insert_prepack_nodes", - "VkInt4WeightOnlyQuantizer", "remove_asserts", "RemoveAssertsTransform", "RemoveLocalScalarDenseOpsTransform", diff --git a/backends/vulkan/_passes/fold_qdq.py b/backends/vulkan/_passes/fold_qdq.py new file mode 100644 index 00000000000..3beccc2205c --- /dev/null +++ b/backends/vulkan/_passes/fold_qdq.py @@ -0,0 +1,41 @@ +# 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.vulkan.utils as utils +import torch + +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass + + +class FoldQDQPass(ExportPass): + """ + Erase Q/DQ chain introduced by PT2E quantization workflow. It is assumed that all + valid quant op patterns have already been fused before this pass. + """ + + def __init__(self, edge_program: torch.export.ExportedProgram): + super(FoldQDQPass, self).__init__() + self.edge_program = edge_program + + def call(self, graph_module: torch.fx.GraphModule): + for node in graph_module.graph.nodes: + if utils.is_quant_node(node): + original_node = node.args[0] + assert isinstance(original_node, torch.fx.Node) + # For each direct user that is a dequant node, connect the original + # node to the users of the dequant node. + for user in node.users: + if utils.is_dequant_node(user): + dq_node = user + dq_node.replace_all_uses_with(original_node) + + graph_module.recompile() + dead_code_elimination_pass(graph_module) + # Re-trace to validate everything is ok + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, True) diff --git a/backends/vulkan/_passes/fuse_patterns.py b/backends/vulkan/_passes/fuse_patterns.py new file mode 100644 index 00000000000..6ced1f32a7c --- /dev/null +++ b/backends/vulkan/_passes/fuse_patterns.py @@ -0,0 +1,30 @@ +# 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.vulkan.patterns as vk_patterns + +import torch + +from executorch.exir import ExportedProgram +from executorch.exir.pass_base import ExportPass, PassResult + + +class FusePatternsPass(ExportPass): + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.program = exported_program + + def call(self, graph_module: torch.fx.GraphModule): + total_replaced = vk_patterns.replace_all_fusable_subgraphs( + self.program, graph_module + ) + + if total_replaced > 0: + graph_module.recompile() + # Re-trace the graph + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, total_replaced > 0) diff --git a/backends/vulkan/_passes/fuse_quantized_ops.py b/backends/vulkan/_passes/fuse_quantized_ops.py index aa4829d9c90..ca9f7541159 100644 --- a/backends/vulkan/_passes/fuse_quantized_ops.py +++ b/backends/vulkan/_passes/fuse_quantized_ops.py @@ -210,278 +210,6 @@ def fuse_into_linear_qcnw_node( graph_module.graph.erase_node(dq_weight_node) -######################### -## linear_qta8a_qga4w ## -######################### - - -def _is_dequantize_affine_node(node: torch.fx.Node) -> bool: - """Check if a node is a dequantize_affine operation.""" - return ( - node.op == "call_function" - and node.target is not None - and hasattr(node.target, "__name__") - and "dequantize_affine" in getattr(node.target, "__name__", "") - ) - - -def _is_view_copy_node(node: torch.fx.Node) -> bool: - """Check if a node is a view_copy operation.""" - return ( - node.op == "call_function" - and node.target is not None - and hasattr(node.target, "__name__") - and "view_copy" in getattr(node.target, "__name__", "") - ) - - -def _validate_qta8a_qga4w_nodes( - input_node: torch.fx.node.Argument, weight_node: torch.fx.node.Argument -) -> Optional[torch.fx.Node]: - """ - Validate input and weight nodes for QTA8A_QGA4W pattern. - Returns the actual input node (after handling view operations) or None if invalid. - """ - # Type checking - ensure we have torch.fx.Node objects - if not isinstance(weight_node, torch.fx.Node) or not isinstance( - input_node, torch.fx.Node - ): - return None - - # Input may be preprocessed with a view node - actual_input_node = input_node - if _is_view_copy_node(input_node): - actual_input_node = input_node.args[0] - if not isinstance(actual_input_node, torch.fx.Node): - return None - - # Check if input is dequantized with dequantize_affine (from dynamic quantization) - if not _is_dequantize_affine_node(actual_input_node): - return None - - # Check if weight is dequantized with dequantize_affine - if not _is_dequantize_affine_node(weight_node): - return None - - return actual_input_node - - -def _extract_weight_params( - program: ExportedProgram, weight_node: torch.fx.Node -) -> Optional[Tuple[torch.fx.Node, torch.fx.Node, torch.fx.Node]]: - """Extract and validate weight parameters from dequantize_affine node.""" - # Get the original quantized weight and quantization parameters - if len(weight_node.args) < 4: - return None - - orig_weight = weight_node.args[0] - weight_scales = weight_node.args[2] - weight_zeros = weight_node.args[3] - - # Type checking - if not isinstance(orig_weight, torch.fx.Node) or not is_param_node( - program, orig_weight - ): - return None - if not isinstance(weight_scales, torch.fx.Node) or not is_param_node( - program, weight_scales - ): - return None - if not isinstance(weight_zeros, torch.fx.Node) or not is_param_node( - program, weight_zeros - ): - return None - - return orig_weight, weight_scales, weight_zeros - - -def _validate_4bit_quantization(weight_tensor: torch.Tensor) -> bool: - """Check if weight tensor is quantized to 4 bits (values in [-8, 7] range).""" - quant_min = weight_tensor.min().item() - quant_max = weight_tensor.max().item() - return quant_min >= -8 and quant_max <= 7 - - -def _calculate_group_size( - orig_weight_tensor: torch.Tensor, weight_scales_tensor: torch.Tensor -) -> Optional[int]: - """Calculate and validate group size from weight and scales tensors.""" - out_features, in_features = orig_weight_tensor.shape - - if len(weight_scales_tensor.shape) != 2: - return None - - scales_out_features, num_groups = weight_scales_tensor.shape - - if scales_out_features != out_features: - return None - - group_size = in_features // num_groups - if in_features % group_size != 0: - return None - - return group_size - - -def matches_linear_qta8a_qga4w_pattern( - program: ExportedProgram, node: torch.fx.Node -) -> Optional[Tuple[int, int]]: - """ - Checks if the nodes surrounding a linear node matches the pattern for dynamic - activation + grouped weight quantized linear (QTA8A_QGA4W). - - This pattern involves: - 1. Dynamic quantization of input activations (8-bit) - 2. Grouped quantization of weights (4-bit with group size) - - The expected pattern from Int8DynActInt4WeightQuantizer is: - scale, zero_point = choose_qparams_affine(input) - quantized_input = quantize_affine(input, scale, zero_point) - dequantized_input = dequantize_affine(quantized_input, ...) - dequantized_weight = dequantize_affine(weight, weight_scales, weight_zeros) - output = linear(dequantized_input, dequantized_weight) - - If the pattern matches, return (group_size, weight_bits), otherwise None. - """ - if not utils.is_linear_node(node): - return None - - input_node = node.args[0] - weight_node = node.args[1] - - # Validate nodes and get actual input node - actual_input_node = _validate_qta8a_qga4w_nodes(input_node, weight_node) - if actual_input_node is None: - return None - - # Extract weight parameters - if not isinstance(weight_node, torch.fx.Node): - return None - weight_params = _extract_weight_params(program, weight_node) - if weight_params is None: - return None - - orig_weight, weight_scales, weight_zeros = weight_params - - # Get tensors to analyze the quantization scheme - orig_weight_tensor = get_param_tensor(program, orig_weight) - weight_scales_tensor = get_param_tensor(program, weight_scales) - weight_zeros_tensor = get_param_tensor(program, weight_zeros) - - if not isinstance(orig_weight_tensor, torch.Tensor): - return None - if not isinstance(weight_scales_tensor, torch.Tensor): - return None - if not isinstance(weight_zeros_tensor, torch.Tensor): - return None - - # Check if weight is quantized to 4 bits - if not _validate_4bit_quantization(orig_weight_tensor): - return None - - # Calculate group size - group_size = _calculate_group_size(orig_weight_tensor, weight_scales_tensor) - if group_size is None: - return None - - # Verify this is 4-bit grouped quantization - weight_bits = 4 - - return group_size, weight_bits - - -def fuse_into_linear_qta8a_qga4w_node( - program: ExportedProgram, - graph_module: torch.fx.GraphModule, - linear_node: torch.fx.Node, - group_size: int, - weight_bits: int, -) -> None: - """ - Fuse the dynamic activation + grouped weight quantized linear pattern into - a single linear_qta8a_qga4w operator. - - The pattern: - dequantized_input = dequantize_affine(quantized_input, block_size, scale, zero_point, ...) - dequantized_weight = dequantize_affine(weight, block_size, weight_scales, weight_zeros, ...) - output = linear(dequantized_input, dequantized_weight) - - Becomes: - output = linear_qta8a_qga4w(quantized_input, input_scale, input_zero_point, - weight, group_size, weight_scales, weight_zeros) - """ - dq_input_node = linear_node.args[0] - dq_weight_node = linear_node.args[1] - - assert isinstance(dq_input_node, torch.fx.Node) - - input_view_node = None - # Input may be preprocessed with a view node - if ( - dq_input_node.op == "call_function" - and dq_input_node.target is not None - and hasattr(dq_input_node.target, "__name__") - and "view_copy" in getattr(dq_input_node.target, "__name__", "") - ): - input_view_node = dq_input_node - dq_input_node = dq_input_node.args[0] - assert isinstance(dq_input_node, torch.fx.Node) - - assert isinstance(dq_input_node, torch.fx.Node) - assert isinstance(dq_weight_node, torch.fx.Node) - - # Get the quantized input and quantization parameters from the input dequantize_affine node - # Args: (input, block_size, scale, zero_point, input_dtype, quant_min, quant_max, output_dtype) - quantized_input = dq_input_node.args[0] - input_scale = dq_input_node.args[2] # scale is the 3rd argument - input_zero_point = dq_input_node.args[3] if len(dq_input_node.args) > 3 else None - - # Get the weight and its quantization parameters from dequantize_affine - # Args: (weight, block_size, weight_scales, weight_zeros, input_dtype, quant_min, quant_max, output_dtype) - orig_weight = dq_weight_node.args[0] - weight_scales = dq_weight_node.args[2] - weight_zeros = dq_weight_node.args[3] - - # Pack the 4-bit weight tensor for efficient storage - assert isinstance(orig_weight, torch.fx.Node) - orig_weight_tensor = get_param_tensor(program, orig_weight) - assert isinstance(orig_weight_tensor, torch.Tensor) - packed_weight_tensor = pack_4bit_weight_tensor(orig_weight_tensor) - utils.update_program_state_dict( - program, - orig_weight.name, - packed_weight_tensor, - ) - # Update the metadata to reflect the new packed shape - orig_weight.meta["val"] = orig_weight.meta["val"][:, ::2].to(torch.uint8) - - # Create the linear_qta8a_qga4w node - with graph_module.graph.inserting_before(linear_node): - linear_qta8a_qga4w_node = graph_module.graph.create_node( - "call_function", - exir_ops.edge.et_vk.linear_qta8a_qga4w.default, - ( - quantized_input, # quantized input (int8) - input_scale, # mat1_scale - input_zero_point, # mat1_zero_point - orig_weight, # mat2_data (packed 4-bit weights) - group_size, # group_size (int) - weight_scales, # weight_scales - weight_zeros, # weight_zeros - ), - ) - - # Replace the linear node with the new fused node - linear_node.replace_all_uses_with(linear_qta8a_qga4w_node) - - # Erase nodes in the correct order (users first, then dependencies) - graph_module.graph.erase_node(linear_node) - if input_view_node is not None: - graph_module.graph.erase_node(input_view_node) - graph_module.graph.erase_node(dq_weight_node) - graph_module.graph.erase_node(dq_input_node) - - class FuseQuantizedOpsTransform(ExportPass): def __init__(self, exported_program: ExportedProgram) -> None: super().__init__() @@ -498,15 +226,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: ) continue - # Check for linear_qta8a_qga4w pattern (dynamic activation + grouped weight quantization) - qta8a_qga4w_details = matches_linear_qta8a_qga4w_pattern(self.program, node) - if qta8a_qga4w_details is not None: - group_size, weight_bits = qta8a_qga4w_details - fuse_into_linear_qta8a_qga4w_node( - self.program, graph_module, node, group_size, weight_bits - ) - continue - graph_module.recompile() dead_code_elimination_pass(graph_module) diff --git a/backends/vulkan/_passes/int4_weight_only_quantizer.py b/backends/vulkan/_passes/int4_weight_only_quantizer.py deleted file mode 100644 index 34ff5937822..00000000000 --- a/backends/vulkan/_passes/int4_weight_only_quantizer.py +++ /dev/null @@ -1,283 +0,0 @@ -# pyre-unsafe -import logging -from typing import Any, Callable, Dict, Optional, Type - -import executorch.backends.vulkan.custom_ops_lib # noqa - -import torch -import torch.nn.functional as F - -from torchao.quantization.unified import Quantizer -from torchao.quantization.utils import groupwise_affine_quantize_tensor - - -# TODO: import from from torchao.quantization.GPTQ.GPTQ import _check_linear_int4_k -# Once diff train catches up -def _check_linear_int4_k(k, group_size=1, inner_k_tiles=None): - """ - Check if the dimensions are compatible with int4 quantization. - - Args: - k: The dimension size to check - group_size: The group size for quantization - inner_k_tiles: The inner k tiles size - - Returns: - bool: Whether the dimensions are compatible - """ - k_divisible_by_group_size = k % group_size == 0 - if inner_k_tiles is not None: - k_divisible_by_16_times_inner_k_tiles = k % (inner_k_tiles * 16) == 0 - return k_divisible_by_group_size and k_divisible_by_16_times_inner_k_tiles - return k_divisible_by_group_size - - -# This module is copied from torchao.quantization.GPTQ.WeightOnlyInt4Linear with -# changes at the annotated lines. -class VkWeightOnlyInt4Linear(torch.nn.Module): - __constants__ = ["in_features", "out_features"] - in_features: int - out_features: int - weight: torch.Tensor - - def __init__( - self, - in_features: int, - out_features: int, - # TODO: remove dtype field, not used - bias=False, - device=None, - dtype=None, - groupsize: int = 128, - inner_k_tiles: int = 8, - precision: torch.dtype = torch.bfloat16, - scales_precision: torch.dtype = torch.bfloat16, - ) -> None: - super().__init__() - self.padding = not _check_linear_int4_k(in_features, groupsize, inner_k_tiles) - if self.padding: - from torchao.utils import find_multiple - - self.origin_in_features = in_features - # pyre-ignore[6]: Incompatible parameter type - in_features = find_multiple(in_features, 1024) - - self.use_bias = bias - self.in_features = in_features - self.out_features = out_features - self.device = device - self.groupsize = groupsize - self.inner_k_tiles = inner_k_tiles - self.precision = precision - self.scales_precision = scales_precision - - if dtype is not None: - raise ValueError("Please specify 'precision' instead of 'dtype'") - - assert out_features % 8 == 0, "require out_features % 8 == 0" - assert ( - in_features % (inner_k_tiles * 16) == 0 - ), "require in_features % (innerKTiles * 16) == 0" - # In the original implementation, the weight buffer is registered with the packed - # sizes, i.e. the result of calling the _convert_weight_to_int4pack operator. - # However, the Vulkan implementation does not expect the weights to be packed - # therefore the weight tensor is registered with the unpacked sizes instead. - # Note that in_features is divided by 2 because each `uint8` tensor element - # contains 2 4-bit packed values. - self.register_buffer( - "weight", - torch.empty( - (out_features, in_features // 2), - dtype=torch.uint8, - device=device, - ), - ) - self.dtype = dtype - self.register_buffer( - "scales_and_zeros", - torch.empty( - (in_features // groupsize, out_features, 2), - dtype=self.scales_precision, - device=device, - ), - ) - if bias: - self.register_buffer( - "bias", - torch.empty((out_features,), dtype=torch.float32, device=device), - ) - - def forward(self, input: torch.Tensor) -> torch.Tensor: - if self.padding: - input = F.pad(input, pad=(0, self.in_features - self.origin_in_features)) - # The forward method is replaced. In the original implementation, the forward - # method is torchao.quantization.GPTQ.linear_forward_int4; here a Vulkan custom - # operator is called instead. - r = torch.ops.et_vk.linear_weight_int4( - input, - self.weight, - self.groupsize, - self.scales_and_zeros, - self.inner_k_tiles, - ) - if self.use_bias: - return r + self.bias - return r - - -# This function is coped from torchao.quantization.GPTQ._replace_linear_int4 -# with small changes at the annotated locations. -def _vk_replace_linear_int4( - module: torch.nn.Module, - groupsize: int, - inner_k_tiles: Optional[int], - padding_allowed: bool, - skip_layer_func: Optional[Callable] = None, - precision: torch.dtype = torch.bfloat16, - scales_precision: torch.dtype = torch.bfloat16, - # Use custom vulkan linear layer as default - linear_class: Type[torch.nn.Module] = VkWeightOnlyInt4Linear, - copy_weights: bool = False, -): - for name, child in module.named_children(): - if isinstance(child, torch.nn.Linear) and ( - skip_layer_func is None or not skip_layer_func(child.weight) - ): - # Add an additional condition that the out/in features must not exceed the - # `feature_limit` argument. - if ( - _check_linear_int4_k(child.in_features, groupsize, inner_k_tiles) - or padding_allowed - ): - new_linear = linear_class( - child.in_features, - child.out_features, - bias=child.bias is not None, - device=child.weight.device, - groupsize=groupsize, - inner_k_tiles=inner_k_tiles, - precision=precision, - scales_precision=scales_precision, - ) - if copy_weights and child.weight.device != torch.device("meta"): - # pyre-fixme[16]: `Module` has no attribute `weight`. - new_linear.weight = child.weight - if child.bias is not None: - # pyre-fixme[16]: `Module` has no attribute `bias`. - new_linear.bias = child.bias - setattr(module, name, new_linear) - else: - _vk_replace_linear_int4( - child, - groupsize, - inner_k_tiles, - padding_allowed, - skip_layer_func, - precision, - scales_precision, - linear_class, - copy_weights, - ) - - -# This module is copied from torchao.quantization.GPTQ.Int4WeightOnlyQuantizer -# with some changes at the annotated lines. -class VkInt4WeightOnlyQuantizer(Quantizer): - def __init__( - self, - groupsize: int = 256, - padding_allowed: bool = True, - inner_k_tiles: Optional[int] = 8, - device: torch.device = torch.device("cpu"), # noqa - precision: torch.dtype = torch.float32, - ) -> None: - super().__init__() - assert inner_k_tiles in [2, 4, 8] - assert groupsize in [32, 64, 128, 256] - - self.inner_k_tiles = inner_k_tiles - self.groupsize: int = groupsize - self.padding_allowed: bool = padding_allowed - self.device: torch.device = device - self.precision: torch.dtype = precision - - @torch.no_grad() - def _create_quantized_state_dict( - self, model: torch.nn.Module - ) -> Dict[str, torch.Tensor]: - cur_state_dict = model.state_dict() - for fqn, mod in model.named_modules(): - # Add additional check to make sure features do not exceed feature limit - if isinstance(mod, torch.nn.Linear): - out_features = mod.out_features - in_features = mod.in_features - logging.info(f"linear: {fqn}, in={in_features}, out={out_features}") - - assert ( - in_features % self.groupsize == 0 - ), f"require in_features:{in_features} % self.groupsize:{self.groupsize} == 0" - - weight = mod.weight.data - if not _check_linear_int4_k( - in_features, self.groupsize, self.inner_k_tiles - ): - if self.padding_allowed: - import torch.nn.functional as F - - from torchao.utils import find_multiple - - logging.warn( - f"warning: {fqn} is padded to satisfy in_features % 1024 == 0" - ) - # pyre-ignore[6]: Incompatible parameter type - padded_in_features = find_multiple(in_features, 1024) - weight = F.pad( - weight, pad=(0, padded_in_features - in_features) - ) - else: - logging.warn( - f"warning: {fqn} is skipped, int4 requires that in_features is 32, 64, or is divisible by 1024, " - + "and that groupsize and inner_k_tiles*16 evenly divide into it" - ) - continue - (w_int4x8, scales_and_zeros) = groupwise_affine_quantize_tensor( - weight, - 4, # n_bit - self.groupsize, - self.precision, # dtype for scales_and_zeros - ) - # If the packing of 2 4-bit values into a single 8-bit value was not - # performed in the previous function call, then do it manually now. - if w_int4x8.shape == weight.shape: - w_int4x8 = (w_int4x8[::, ::2] << 4 | w_int4x8[::, 1::2]).to( - torch.uint8 - ) - # In the original implementation, w_int4x8 is packed via calling the - # _convert_weight_to_int4pack operator before storing the weight. However - # the Vulkan implementation does not expect the weights to be packed, so - # the w_int4x8 tensor is stored as the weight instead. - cur_state_dict[f"{fqn}.weight"] = w_int4x8.to(self.device) - cur_state_dict[f"{fqn}.scales_and_zeros"] = scales_and_zeros.to( - self.device - ) - return cur_state_dict - - def _convert_for_runtime(self, model: torch.nn.Module) -> torch.nn.Module: - _vk_replace_linear_int4( - model, - self.groupsize, - self.inner_k_tiles, - self.padding_allowed, - skip_layer_func=None, - precision=self.precision, - scales_precision=self.precision, - ) - return model - - def quantize( - self, model: torch.nn.Module, *args: Any, **kwargs: Any - ) -> torch.nn.Module: - state_dict = self._create_quantized_state_dict(model) - model = self._convert_for_runtime(model) - model.load_state_dict(state_dict, strict=False) - return model diff --git a/backends/vulkan/_passes/remove_redundant_ops.py b/backends/vulkan/_passes/remove_redundant_ops.py index 530505f7003..8e602dd17b4 100644 --- a/backends/vulkan/_passes/remove_redundant_ops.py +++ b/backends/vulkan/_passes/remove_redundant_ops.py @@ -30,6 +30,7 @@ class RemoveRedundantOpsTransform(ExportPass): 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, + exir_ops.edge.dim_order_ops._clone_dim_order.default, } def __init__(self) -> None: diff --git a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py index c415249383e..25b28ce3117 100644 --- a/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py +++ b/backends/vulkan/_passes/squeeze_unsqueeze_inputs.py @@ -22,7 +22,6 @@ class SqueezeUnsqueezeInputs(ExportPass): _squeezable_ops: Set[OpType] = { - exir_ops.edge.et_vk.linear_weight_int4.default, exir_ops.edge.aten.relu.default, exir_ops.edge.aten.gelu.default, } diff --git a/backends/vulkan/cmake/ShaderLibrary.cmake b/backends/vulkan/cmake/ShaderLibrary.cmake index c06a9d7097c..1b6838c4dfd 100644 --- a/backends/vulkan/cmake/ShaderLibrary.cmake +++ b/backends/vulkan/cmake/ShaderLibrary.cmake @@ -49,6 +49,15 @@ function(gen_vulkan_shader_lib_cpp shaders_path) set(VULKAN_SHADERGEN_ENV "") set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/vulkan_compute_shaders) + set(GEN_SPV_ARGS "--optimize") + if(DEFINED ENV{ETVK_USING_SWIFTSHADER}) + if("$ENV{ETVK_USING_SWIFTSHADER}" STREQUAL "1" + OR "$ENV{ETVK_USING_SWIFTSHADER}" STREQUAL "True" + ) + list(APPEND GEN_SPV_ARGS "--replace-u16vecn") + endif() + endif() + add_custom_command( COMMENT "Generating Vulkan Compute Shaders" OUTPUT ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp @@ -58,7 +67,7 @@ function(gen_vulkan_shader_lib_cpp shaders_path) ${shaders_path} --output-path ${VULKAN_SHADERGEN_OUT_PATH} --glslc-path=${GLSLC_PATH} --tmp-dir-path=${VULKAN_SHADERGEN_OUT_PATH}/shader_cache/ --env - ${VULKAN_GEN_ARG_ENV} --optimize + ${VULKAN_GEN_ARG_ENV} ${GEN_SPV_ARGS} DEPENDS ${shaders_path}/* ${EXECUTORCH_ROOT}/backends/vulkan/runtime/gen_vulkan_spv.py ) diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index c9b884e5b86..56e803b9127 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -4,6 +4,9 @@ # 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 Optional + +import executorch.backends.vulkan.patterns as vk_patterns import torch.library namespace = "et_vk" @@ -152,41 +155,9 @@ def grid_priors_out_impl( ) lib.impl(name, grid_priors_out_impl, "CompositeExplicitAutograd") -######################## -## linear_weight_int4 ## -######################## - - -def linear_weight_int4_impl( - x: torch.Tensor, - weights_4x8: torch.Tensor, - groupsize: int, - scales_and_zeros: torch.Tensor, - inner_k_tiles: int, -): - original_x_size = x.size() - out_features = weights_4x8.size(0) - x = x.reshape(-1, original_x_size[-1]) - weight_int4pack = torch.ops.aten._convert_weight_to_int4pack( - weights_4x8, inner_k_tiles - ) - out = torch.ops.aten._weight_int4pack_mm( - x, weight_int4pack, groupsize, scales_and_zeros - ) - out_shape = original_x_size[:-1] + (out_features,) - return out.reshape(out_shape) - - -name = "linear_weight_int4" -lib.define( - f"{name}(Tensor self, Tensor mat2, int qGroupSize, Tensor qScaleAndZeros, int inner_k_tiles) -> Tensor" -) -lib.impl(name, linear_weight_int4_impl, "CompositeExplicitAutograd") -linear_weight_int4_op = getattr(getattr(torch.ops, namespace), name) - -################# +################## ## linear_qcs4w ## -################# +################## def linear_qcs4w( @@ -231,136 +202,248 @@ def linear_qcs4w( lib.impl(name, linear_qcs4w, "CompositeExplicitAutograd") linear_qc4w_op = getattr(getattr(torch.ops, namespace), name) -######################## -## linear_qta8a_qga4w ## -######################## +################## +## linear_q4gsw ## +################## + + +def unpack_4bit_weight_tensor( + packed_weight_tensor: torch.Tensor, x: torch.Tensor +) -> torch.Tensor: + """ + Reverses the packing performed in quantized_linear.pack_4bit_weight_tensor + """ + # Each packed byte contains two 4-bit values: high nibble and low nibble + K, N_half = packed_weight_tensor.shape + N = N_half * 2 + + # Unpack high and low nibbles + high_nibble = (packed_weight_tensor >> 4) & 0x0F + low_nibble = packed_weight_tensor & 0x0F + + # Stack to shape (K, N) + unpacked = torch.empty( + (K, N), dtype=torch.uint8, device=packed_weight_tensor.device + ) + unpacked[:, ::2] = low_nibble + unpacked[:, 1::2] = high_nibble + + # Undo the +8 offset and convert to signed 4-bit range [-8, 7] + unpacked = unpacked.to(torch.int8) - 8 + + in_channels = x.shape[-1] + # Undo any padding that may have been added to input channels + if in_channels != unpacked.shape[-1]: + return unpacked[:, :in_channels] + return unpacked -def linear_qta8a_qga4w( - x_quantized: torch.Tensor, + +def linear_q4gsw( + x: torch.Tensor, + weights: torch.Tensor, + weight_scales: torch.Tensor, + group_size: int, + bias: Optional[torch.Tensor] = None, +): + # Unpack the packed weights + weights = unpack_4bit_weight_tensor(weights, x) + + # Un-transpose the weight scales + weight_scales = weight_scales.transpose(0, 1) + weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32) + + weights = torch.ops.torchao.dequantize_affine( + weights, [1, group_size], weight_scales, weight_zeros, torch.int8, -8, 7 + ) + + out = torch.nn.functional.linear(x, weights) + return out + + +def linear_dq8ca_q4gsw( + x: torch.Tensor, input_scale: torch.Tensor, input_zero_point: torch.Tensor, - weights_4bit: torch.Tensor, - group_size: int, + weights: torch.Tensor, + weight_sums: torch.Tensor, weight_scales: torch.Tensor, - weight_zeros: torch.Tensor, + group_size: int, + bias: Optional[torch.Tensor] = None, ): - """ - Dynamic activation + grouped weight quantized linear (QTA8A_QGA4W). - - Args: - x_quantized: Already quantized input tensor (int8, per-token quantized) - input_scale: Scale for per-token quantization of input (shape: [batch_size]) - input_zero_point: Zero point for per-token quantization of input (shape: [batch_size]) - weights_4bit: Packed 4-bit quantized weights - group_size: Group size for weight quantization (int) - weight_scales: Per-group scales for weights - weight_zeros: Per-group zero points for weights - """ - original_x_shape = x_quantized.shape - feature_dim = original_x_shape[-1] + return linear_q4gsw(x, weights, weight_scales, group_size) - # Reshape for processing - x_quantized_2d = x_quantized.reshape(-1, feature_dim) - # Unpack 4-bit weights - unpacked_weights_shape = weights_4bit.shape - out_features = unpacked_weights_shape[0] - in_features = unpacked_weights_shape[1] +name = "linear_q4gsw" +lib.define( + f""" + {name}( + Tensor self, + Tensor weights, + Tensor weight_scales, + int group_size, + Tensor? bias = None) -> Tensor + """ +) +lib.impl(name, linear_q4gsw, "CompositeExplicitAutograd") +linear_qc4w_op = getattr(getattr(torch.ops, namespace), name) - weights_unpacked = torch.empty( - (out_features, in_features * 2), dtype=torch.int8, device=weights_4bit.device - ) +name = "linear_dq8ca_q4gsw" +lib.define( + f""" + {name}( + Tensor input, + Tensor input_scales, + Tensor input_zp, + Tensor weights, + Tensor weight_sums, + Tensor weight_scales, + int group_size, + Tensor? bias = None) -> Tensor + """ +) +lib.impl(name, linear_dq8ca_q4gsw, "CompositeExplicitAutograd") +linear_dq8ca_q4gsw_op = getattr(getattr(torch.ops, namespace), name) + +################# +## qaqw_linear ## +################# - weights_unpacked[:, ::2] = weights_4bit >> 4 - weights_unpacked[:, 1::2] = weights_4bit & 0x0F - # Convert to signed 4-bit range [-8, 7] - weights_unpacked = torch.where( - weights_unpacked > 7, weights_unpacked - 16, weights_unpacked +def linear_q8ta_q8csw( + x: torch.Tensor, + input_scale: float, + input_zero_point: int, + weights: torch.Tensor, + weight_sums: torch.Tensor, + weight_scales: torch.Tensor, + bias: Optional[torch.Tensor] = None, +): + weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32) + weights = torch.ops.quantized_decomposed.dequantize_per_channel( + weights, + weight_scales, + weight_zeros, + 0, + -127, + 127, + torch.int8, ) - # Dequantize weights using grouped quantization - actual_in_features = in_features * 2 - num_groups = actual_in_features // group_size + # Perform linear operation + out = torch.nn.functional.linear(x, weights) + if bias is not None: + out = out + bias + + return out + - # Reshape weights for grouped dequantization - weights_grouped = weights_unpacked.view(out_features, num_groups, group_size) +name = "linear_q8ta_q8csw" +lib.define( + f""" + {name}( + Tensor x, + float input_scale, + int input_zero_point, + Tensor weights, + Tensor weight_sums, + Tensor weight_scales, + Tensor? bias = None) -> Tensor + """ +) +lib.impl(name, linear_q8ta_q8csw, "CompositeExplicitAutograd") +qa_q8csw_linear = getattr(getattr(torch.ops, namespace), name) - # Expand scales and zeros to match grouped weights - scales_expanded = weight_scales.unsqueeze(-1).expand(-1, -1, group_size) - zeros_expanded = weight_zeros.unsqueeze(-1).expand(-1, -1, group_size) +####################### +## conv2d_q8ta_q8csw ## +####################### - # Dequantize: (quantized - zero_point) * scale - dq_weights_grouped = (weights_grouped.float() - zeros_expanded) * scales_expanded - dq_weights = dq_weights_grouped.view(out_features, actual_in_features) - # Dequantize input (per-token) - # For per-token quantization, each token (row) has its own scale and zero_point - x_dequantized = torch.ops.quantized_decomposed.dequantize_per_token( - x_quantized_2d, - input_scale, - input_zero_point, - -128, +def conv2d_q8ta_q8csw( + x: torch.Tensor, + input_scale: float, + input_zero_point: int, + weights: torch.Tensor, + weight_sums: torch.Tensor, + weight_scales: torch.Tensor, + bias: Optional[torch.Tensor], + kernel_size: list, + stride: list, + padding: list, + dilation: list, + groups: int, +): + IC = x.shape[1] + K_h, K_w = kernel_size[0], kernel_size[1] + + canonical_weight_K_dim = K_h * K_w * IC + # Remove any padding added to output channels dim to align to a multiple of 4 + if weights.shape[-1] != canonical_weight_K_dim: + weights = weights[:, :canonical_weight_K_dim] + weight_scales = weight_scales[:canonical_weight_K_dim] + if bias is not None: + bias = bias[:canonical_weight_K_dim] + + weight_zeros = torch.zeros_like(weight_scales, dtype=torch.int32) + + # Calculate dimensions + OC = weights.shape[0] + in_features = weights.shape[1] + IC = in_features // (K_h * K_w) + + # Reshape to original 4D format (OC, IC, H, W) + weights = weights.view(OC, IC, K_h, K_w) + + # Dequantize weights + weights = torch.ops.quantized_decomposed.dequantize_per_channel( + weights, + weight_scales, + weight_zeros, + 0, # axis=0 for output channel quantization + -127, 127, torch.int8, - torch.float32, ) - # Perform linear operation - out = torch.nn.functional.linear(x_dequantized, dq_weights) - out_shape = original_x_shape[:-1] + (out_features,) - return out.reshape(out_shape) + # Perform convolution + out = torch.nn.functional.conv2d( + x, weights, bias, stride, padding, dilation, groups + ) + + return out -name = "linear_qta8a_qga4w" +name = "conv2d_q8ta_q8csw" lib.define( - f"{name}(Tensor self, Tensor input_scale, Tensor input_zero_point, Tensor weight, int group_size, Tensor weight_scales, Tensor weight_zeros) -> Tensor" + f""" + {name}( + Tensor x, + float input_scale, + int input_zero_point, + Tensor weights, + Tensor weight_sums, + Tensor weight_scales, + Tensor? bias, + SymInt[] kernel_size, + SymInt[] stride, + SymInt[] padding, + SymInt[] dilation, + SymInt groups) -> Tensor + """ ) -lib.impl(name, linear_qta8a_qga4w, "CompositeExplicitAutograd") -linear_qta8a_qga4w_op = getattr(getattr(torch.ops, namespace), name) +lib.impl(name, conv2d_q8ta_q8csw, "CompositeExplicitAutograd") +conv2d_q8ta_q8csw_op = getattr(getattr(torch.ops, namespace), name) ###################### ## apply_rotary_emb ## ###################### -# Note that this implementation is copied from executorch.examples.models.llama.rope -# but it is copied here to avoid introducing a dependency on the llama code. def apply_rotary_emb_impl( xq: torch.Tensor, xk: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor ): - def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): - ndim = x.ndim - freqs_cis_ndim = freqs_cis.ndim - if freqs_cis_ndim == 3: - # freqs_cis: (seq_len, n_heads, head_dim // 2) - assert freqs_cis.shape == (x.shape[-3], x.shape[-2], x.shape[-1]) - shape = [ - d if (i == ndim - 3 or i == ndim - 2 or i == ndim - 1) else 1 - for i, d in enumerate(x.shape) - ] - else: - # freqs_cis: (seq_len, head_dim // 2) - assert freqs_cis.shape == (x.shape[1], x.shape[-1]) - shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] - return freqs_cis.view(shape) - - xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1) - xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1) - - freqs_cos = reshape_for_broadcast(freqs_cos, xq_r) - freqs_sin = reshape_for_broadcast(freqs_sin, xq_r) - - xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin - xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos - xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin - xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos - - xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3) - xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3) - - return xq_out.type_as(xq), xk_out.type_as(xk) + pattern = vk_patterns.RotaryEmbeddingPattern() + return pattern.forward(xq, xk, freqs_cos, freqs_sin) name = "apply_rotary_emb" diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 2e0be1d68d7..4c686e0cfc5 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -16,6 +16,8 @@ import torch +from executorch.backends.vulkan.serialization.vulkan_graph_schema import VkMemoryLayout + from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload @@ -123,6 +125,7 @@ def update_features_impl(op: OpKey): operator.gt, operator.ge, operator.le, + operator.eq, # Guard and assert ops torch.ops.aten._assert_scalar.default, torch.ops.aten.sym_constrain_range_for_size.default, @@ -169,7 +172,6 @@ def register_affine_quantization_op(): @update_features( [ - exir_ops.edge.torchao.choose_qparams_affine.default, exir_ops.edge.quantized_decomposed.choose_qparams.tensor, exir_ops.edge.quantized_decomposed.choose_qparams_per_token_asymmetric.default, ] @@ -181,6 +183,20 @@ def register_torchao_quantization_op(): ) +@update_features( + exir_ops.edge.torchao.choose_qparams_affine.default, +) +def register_torchao_choose_qparams_affine(): + return OpFeatures( + inputs_storage=utils.CONTIGUOUS_ANY, + outputs_storage=[ + utils.WIDTH_PACKED_TEXTURE, # scales + utils.WIDTH_PACKED_TEXTURE, # zero_points + ], + supports_resize=True, + ) + + @update_features( [ exir_ops.edge.aten.add.Tensor, @@ -285,6 +301,32 @@ def check_dim_order_copy_node(node: torch.fx.Node) -> bool: ) +@update_features(exir_ops.edge.dim_order_ops._clone_dim_order.default) +def register_clone_dim_order_op(): + # Similar to to_dim_order_copy, _clone_dim_order can be removed as long as the + # operator is not changing the dtype, i.e. the operator call is modifying the dim + # order only. Therefore, check that the input and output dtypes are the same, if so + # the operator is safe to remove. + def check_clone_dim_order_node(node: torch.fx.Node) -> bool: + in_arg = node.args[0] + if not isinstance(in_arg, torch.fx.Node): + return False + + in_tensor = in_arg.meta.get("val", None) + out_tensor = node.meta.get("val", None) + + if in_tensor.dtype != out_tensor.dtype: + return False + + return True + + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + supports_resize=True, + are_node_inputs_supported_fn=check_clone_dim_order_node, + ) + + @update_features( [ exir_ops.edge.aten.bmm.default, @@ -317,34 +359,30 @@ def register_int8_mm_op(): @update_features( [ - exir_ops.edge.et_vk.linear_weight_int4.default, + exir_ops.edge.et_vk.linear_q8ta_q8csw.default, + exir_ops.edge.et_vk.linear_q4gsw.default, ] ) -def register_int4_mm_op(): +def register_quantized_linear_ops(): return OpFeatures( inputs_storage=utils.CONTIGUOUS_ANY, - supports_resize=True, supports_prepacking=True, ) -@update_features( - [ - exir_ops.edge.et_vk.linear_qta8a_qga4w.default, - ] -) -def register_dqlinear_op(): +@update_features(exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default) +def register_linear_dqa_qw_ops(): return OpFeatures( inputs_storage=[ utils.CONTIGUOUS_ANY, # input - utils.CONTIGUOUS_BUFFER, # mat1 scales - utils.CONTIGUOUS_BUFFER, # mat1 zeros + utils.WIDTH_PACKED_TEXTURE, # input_scale + utils.WIDTH_PACKED_TEXTURE, # input_zero_point utils.NO_STORAGE, # weight (prepacked) - utils.NO_STORAGE, # group size (non tensor) - utils.CONTIGUOUS_BUFFER, # mat2 scales - utils.CONTIGUOUS_BUFFER, # mat2 zeros + utils.NO_STORAGE, # weight_sums (prepacked) + utils.NO_STORAGE, # weight_scales (prepacked) + utils.NO_STORAGE, # group_size (scalar) + utils.NO_STORAGE, # bias (prepacked) ], - supports_resize=True, supports_prepacking=True, ) @@ -373,13 +411,38 @@ def register_softmax_op(): def register_reduce_op(): def check_reduce_node(node: torch.fx.Node) -> bool: dim_list = node.args[1] - if isinstance(dim_list, list) and len(dim_list) != 1: + if isinstance(dim_list, list) and len(dim_list) > 2: + return False + + if isinstance(dim_list, list) and len(dim_list) == 2: + # Try to get the memory layout for this node + try: + memory_layout = utils.get_node_memory_layout(node) + + # If we have memory layout information, check if any dimension in dim_list corresponds to a packed dimension + if ( + memory_layout is not None + and memory_layout != VkMemoryLayout.DEFAULT_LAYOUT + ): + # For now only default layout is supported for 2D reduction. + # Because we can't determine if the input is NCHW or NHWC here, + # assume the reduction dimension is packed so we cannot support it. + return False + except (AssertionError, KeyError, AttributeError): + # If we can't get memory layout information, we'll assume the dims aren't packed + pass + + def try_find_keepdim_arg(node: torch.fx.Node) -> bool: + for arg in node.args: + if isinstance(arg, bool): + return arg + + # Assume false by default return False - if len(node.args) > 2: - keepdim = node.args[2] - if isinstance(keepdim, bool) and not keepdim: - return False + keepdim = try_find_keepdim_arg(node) + if isinstance(keepdim, bool) and not keepdim: + return False return True @@ -393,6 +456,7 @@ def check_reduce_node(node: torch.fx.Node) -> bool: @update_features( [ exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.aten.max_pool2d.default, exir_ops.edge.aten.max_pool2d_with_indices.default, ] ) @@ -411,12 +475,51 @@ def register_2d_pool_op(): ) def register_convolution_op(): return OpFeatures( - inputs_storage=utils.CHANNELS_PACKED_TEXTURE, + inputs_storage=[ + utils.CHANNELS_PACKED_TEXTURE, # input + utils.NO_STORAGE, # weight (prepacked) + utils.NO_STORAGE, # bias (prepacked) + utils.NO_STORAGE, # stride (non tensor) + utils.NO_STORAGE, # padding (non tensor) + utils.NO_STORAGE, # dilation (non tensor) + utils.NO_STORAGE, # transposed (non tensor) + utils.NO_STORAGE, # output_padding (non tensor) + utils.NO_STORAGE, # groups (non tensor) + utils.NO_STORAGE, # output_min (non tensor) + utils.NO_STORAGE, # output_max (non tensor) + ], supports_resize=True, supports_prepacking=True, ) +@update_features( + [ + exir_ops.edge.et_vk.conv2d_q8ta_q8csw.default, + ] +) +def register_quantized_conv_op(): + return OpFeatures( + inputs_storage=[ + utils.CHANNELS_PACKED_TEXTURE, # input + utils.NO_STORAGE, # input_scale (non tensor) + utils.NO_STORAGE, # input_zero_point (non tensor) + utils.NO_STORAGE, # weight (prepacked) + utils.NO_STORAGE, # weight_sums (prepacked) + utils.NO_STORAGE, # weight_scales (prepacked) + utils.NO_STORAGE, # bias (prepacked) + utils.NO_STORAGE, # kernel_size (non tensor) + utils.NO_STORAGE, # stride (non tensor) + utils.NO_STORAGE, # padding (non tensor) + utils.NO_STORAGE, # dilation (non tensor) + utils.NO_STORAGE, # groups (non tensor) + utils.NO_STORAGE, # original OC count (non tensor) + ], + supports_resize=False, + supports_prepacking=True, + ) + + @update_features("llama::sdpa_with_kv_cache") def register_sdpa_with_kv_cache_op(): return OpFeatures( @@ -434,7 +537,7 @@ def register_sdpa_with_kv_cache_op(): ) def register_sdpa_ops(): return OpFeatures( - inputs_storage=utils.WIDTH_PACKED_TEXTURE, + inputs_storage=utils.CONTIGUOUS_ANY, supports_resize=True, ) @@ -449,10 +552,7 @@ def register_rotary_emb_op(): @update_features( [ - exir_ops.edge.aten.clone.default, exir_ops.edge.aten.permute.default, - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.view_copy.default, ] ) def register_view_ops(): @@ -462,22 +562,35 @@ def register_view_ops(): ) +@update_features( + [ + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.clone.default, + exir_ops.edge.aten.permute_copy.default, + ] +) +def register_view_ops_with_buffer_meta(): + return OpFeatures( + inputs_storage=utils.ANY_STORAGE, + supports_resize=True, + ) + + +@update_features(exir_ops.edge.aten.expand_copy.default) +def register_expand(): + return OpFeatures(inputs_storage=utils.ANY_BUFFER, supports_resize=False) + + # Fully featured transfer operators (i.e. operators that copy data from the input # tensor(s) to the output tensor(s)), which have memory layout agnostic implementations # for both texture and buffer storage types. @update_features(exir_ops.edge.aten.cat.default) def register_cat_op(): - def check_cat_node(node: torch.fx.Node) -> bool: - inputs = node.args[0] - if isinstance(inputs, (list, tuple)) and len(inputs) <= 3: - return True - - return False - return OpFeatures( inputs_storage=utils.ANY_STORAGE, supports_resize=True, - are_node_inputs_supported_fn=check_cat_node, ) @@ -530,9 +643,6 @@ def register_ported_op(): # Ops ported from PyTorch Vulkan backend. These ops are in a separate registry because they support all packed dimensions @update_features( [ - # Shape Manipulation - exir_ops.edge.aten.squeeze_copy.dims, - exir_ops.edge.aten.unsqueeze_copy.default, # Tensor combination exir_ops.edge.aten.repeat.default, exir_ops.edge.aten.split_with_sizes_copy.default, diff --git a/backends/vulkan/partitioner/TARGETS b/backends/vulkan/partitioner/TARGETS index 1d1d29f6fb0..986d872f730 100644 --- a/backends/vulkan/partitioner/TARGETS +++ b/backends/vulkan/partitioner/TARGETS @@ -15,10 +15,12 @@ runtime.python_library( "//executorch/backends/vulkan:op_registry", "//executorch/backends/vulkan:utils_lib", "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/backends/vulkan/patterns:vulkan_patterns", "//executorch/exir:delegate", "//executorch/exir:lib", "//executorch/exir/backend:partitioner", "//executorch/exir/backend:utils", "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib", ], + typing = True, ) diff --git a/backends/vulkan/partitioner/vulkan_partitioner.py b/backends/vulkan/partitioner/vulkan_partitioner.py index 776d1d6e168..e5b2d0f7864 100644 --- a/backends/vulkan/partitioner/vulkan_partitioner.py +++ b/backends/vulkan/partitioner/vulkan_partitioner.py @@ -7,8 +7,9 @@ # pyre-strict import logging -from typing import Any, Callable, Dict, final, List, Mapping, Optional, Tuple +from typing import Any, Callable, Dict, final, List, Mapping, Optional, Set, Tuple +import executorch.backends.vulkan.patterns as vk_patterns import executorch.backends.vulkan.utils as utils import torch @@ -17,9 +18,12 @@ get_op_features, has_impl, OpFeatures, + OpKey, vulkan_supported_ops, ) +from executorch.backends.vulkan.patterns import PatternMatch + from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( VkMemoryLayout, VkStorageType, @@ -36,8 +40,8 @@ from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import ExportedProgram -from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner +from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupportBase # pyre-ignore @@ -55,11 +59,30 @@ def __init__( texture_limits: utils.ImageExtents, buffer_limit: int, require_dynamic_shape: bool = False, + operator_blocklist: Optional[Set[OpKey]] = None, + operator_allowlist: Optional[Set[OpKey]] = None, + fusable_subgraphs: Optional[List[PatternMatch]] = None, + nn_module_blocklist: Optional[Set[str]] = None, + nn_module_allowlist: Optional[Set[str]] = None, ) -> None: super().__init__() self.texture_limits: utils.ImageExtents = texture_limits self.buffer_limit = buffer_limit self.require_dynamic_shapes = require_dynamic_shape + self.operator_blocklist: Set[OpKey] = ( + operator_blocklist if operator_blocklist is not None else set() + ) + self.operator_allowlist = operator_allowlist + self.fusable_subgraphs: List[PatternMatch] = ( + fusable_subgraphs if fusable_subgraphs is not None else [] + ) + # Create a set of all nodes that are part of fusable subgraphs for quick lookup + self.fusable_nodes: Set[torch.fx.Node] = set() + for match in self.fusable_subgraphs: + self.fusable_nodes.update(match.all_nodes) + + self.nn_module_blocklist = nn_module_blocklist + self.nn_module_allowlist = nn_module_allowlist def op_node_is_compatible( # noqa: C901: Function is too complex self, node: torch.fx.Node, features: Optional[OpFeatures] = None @@ -77,6 +100,17 @@ def op_node_is_compatible( # noqa: C901: Function is too complex assert isinstance(first_arg, torch._ops.OpOverload) target = first_arg.name() + # Operator allow list is only used for torch ops + if ( + utils.is_torch_op_node(node) + and (self.operator_allowlist is not None) + and (target not in self.operator_allowlist) + ): + return False, "op is not in allowlist" + + if target in self.operator_blocklist: + return False, "op is in blocklist" + # Extract the features for the node's operator, if no override was provided if features is None: if not has_impl(target): @@ -93,7 +127,7 @@ def op_node_is_compatible( # noqa: C901: Function is too complex if op_repsets.any_is_empty(): return ( False, - "No valid representations for a tensor in the operation", + f"no valid representations for op {utils.node_io_str(node)}", ) return True, "Op is compatible" @@ -176,7 +210,7 @@ def is_in_local_scalar_dense_chain(self, node: torch.fx.Node) -> Tuple[bool, boo def log_skip(self, node: torch.fx.Node, reason: str) -> None: if node.op == "call_function": logger.info( - f"[Vulkan Partitioner] Due to [{reason}], skipping {node.format_node()}" + f"[Vulkan Partitioner] Due to [{reason}], skipping {utils.node_io_str(node)}" ) def is_node_supported( @@ -185,7 +219,27 @@ def is_node_supported( r = self._is_node_supported(node) return r - def _is_node_supported(self, node: torch.fx.Node) -> bool: + def _is_node_supported(self, node: torch.fx.Node) -> bool: # noqa: C901 + if node.op == "call_function": + # Apply nn module allowlist and blocklist + if self.nn_module_allowlist is not None: + if not utils.node_comes_from_any_nn_module_in_set( + node, self.nn_module_allowlist + ): + self.log_skip(node, "source nn.Module is not in allowlist") + return False + + if self.nn_module_blocklist is not None: + if utils.node_comes_from_any_nn_module_in_set( + node, self.nn_module_blocklist + ): + self.log_skip(node, "source nn.Module is in blocklist") + return False + + # Check if this node is part of a fusable subgraph + if node in self.fusable_nodes: + return True + target = node.target if node.target == torch.ops.higher_order.auto_functionalized: first_arg = node.args[0] @@ -277,6 +331,10 @@ class VulkanPartitioner(Partitioner): def __init__( self, compile_options: Optional[Dict[str, Any]] = None, + operator_blocklist: Optional[List[OpKey]] = None, + operator_allowlist: Optional[List[OpKey]] = None, + nn_module_blocklist: Optional[List[str]] = None, + nn_module_allowlist: Optional[List[str]] = None, ) -> None: self.options: Dict[str, Any] = {} if compile_options is not None: @@ -285,6 +343,32 @@ def __init__( compile_spec = parse_compile_options(self.options) self.delegation_spec = DelegationSpec(VulkanBackend.__name__, compile_spec) + self.operator_blocklist: Set[OpKey] = set() + if operator_blocklist is not None: + for entry in operator_blocklist or []: + self.operator_blocklist.add(entry) + + self.operator_allowlist: Optional[Set[OpKey]] = None + if operator_allowlist is not None: + self.operator_allowlist = set() + for entry in operator_allowlist: + assert self.operator_allowlist is not None + self.operator_allowlist.add(entry) + + self.nn_module_blocklist: Optional[Set[str]] = None + if nn_module_blocklist is not None: + self.nn_module_blocklist = set() + for entry in nn_module_blocklist or []: + assert self.nn_module_blocklist is not None + self.nn_module_blocklist.add(entry) + + self.nn_module_allowlist: Optional[Set[str]] = None + if nn_module_allowlist is not None: + self.nn_module_allowlist = set() + for entry in nn_module_allowlist: + assert self.nn_module_allowlist is not None + self.nn_module_allowlist.add(entry) + def ops_to_not_decompose( self, ep: ExportedProgram ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]: @@ -298,6 +382,11 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: # subgraphs containing the nodes with the tags partition_tags = {} + # Get all fusable subgraphs from fuse_patterns + fusable_subgraphs = vk_patterns.get_all_fusable_subgraphs( + exported_program.graph_module + ) + texture_limits: utils.ImageExtents = self.options.get( "texture_limits", utils.DEFAULT_TEXTURE_LIMITS ) @@ -308,6 +397,11 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: texture_limits, buffer_limit, require_dynamic_shape=self.options.get("require_dynamic_shapes", False), + operator_blocklist=self.operator_blocklist, + operator_allowlist=self.operator_allowlist, + fusable_subgraphs=fusable_subgraphs, + nn_module_blocklist=self.nn_module_blocklist, + nn_module_allowlist=self.nn_module_allowlist, ), allows_single_node_partition=True, ) diff --git a/backends/vulkan/patterns/TARGETS b/backends/vulkan/patterns/TARGETS new file mode 100644 index 00000000000..791edf58984 --- /dev/null +++ b/backends/vulkan/patterns/TARGETS @@ -0,0 +1,26 @@ +load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_library( + name = "vulkan_patterns", + srcs = [ + "__init__.py", + "pattern_registry.py", + "rope.py", + "quantized_linear.py", + "quantized_convolution.py", + ], + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/backends/transforms:utils", + "//executorch/backends/vulkan:utils_lib", + ], + typing = True, +) diff --git a/backends/vulkan/patterns/__init__.py b/backends/vulkan/patterns/__init__.py new file mode 100644 index 00000000000..8ffad98b3c3 --- /dev/null +++ b/backends/vulkan/patterns/__init__.py @@ -0,0 +1,133 @@ +# 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 typing import List + +import executorch.backends.vulkan.patterns.quantized_convolution # noqa + +import executorch.backends.vulkan.patterns.quantized_linear # noqa + +import executorch.backends.vulkan.patterns.rope # noqa + +import torch + +from executorch.backends.vulkan.patterns.pattern_registry import ( + create_pattern_match_from_internal_match, + CreateReplacementFn, + DetectorFn, + fusable_patterns, + GetGraphFn, + PatternMatch, + register_pattern_detector, + register_pattern_graph, + register_pattern_replacement, +) + +from executorch.backends.vulkan.patterns.rope import RotaryEmbeddingPattern + +from executorch.exir import ExportedProgram + +from torch.fx.passes.utils.matcher_utils import SubgraphMatcher + + +__all__ = [ + "PatternMatch", + "GetGraphFn", + "DetectorFn", + "CreateReplacementFn", + "RotaryEmbeddingPattern", + "fusable_patterns", + "register_pattern_graph", + "register_pattern_detector", + "register_pattern_replacement", +] + + +def all_fusable_graph_patterns() -> List[torch.fx.GraphModule]: + all_patterns = [] + for entry in fusable_patterns.values(): + if entry.get_graphs_fn is not None: + all_patterns.extend(entry.get_graphs_fn()) + + return all_patterns + + +def get_all_fusable_subgraphs( + graph_module: torch.fx.GraphModule, +) -> List[PatternMatch]: + fusable_subgraphs = [] + + fuse_patterns = all_fusable_graph_patterns() + for pattern in fuse_patterns: + sm = SubgraphMatcher(pattern.graph, ignore_literals=True) + matches = list(sm.match(graph_module.graph)) + for match in matches: + fusable_subgraphs.append(create_pattern_match_from_internal_match(match)) + + for node in graph_module.graph.nodes: + for entry in fusable_patterns.values(): + if entry.detector_fn is not None: + maybe_match = entry.detector_fn(node) + if maybe_match is not None: + fusable_subgraphs.append(maybe_match) + + return fusable_subgraphs + + +def create_replacement_for_pattern( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + patterns: List[torch.fx.GraphModule], + create_replacement_func: CreateReplacementFn, +) -> int: + total_replaced = 0 + + for pattern in patterns: + sm = SubgraphMatcher(pattern.graph, ignore_literals=True) + matches = list(sm.match(graph_module.graph)) + + for partition_to_replace in matches: + pattern = create_pattern_match_from_internal_match(partition_to_replace) + create_replacement_func(ep, graph_module, pattern) + total_replaced += 1 + # Remove dead code so they won't be matched again + graph_module.graph.eliminate_dead_code() + + return total_replaced + + +def replace_all_fusable_subgraphs( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, +) -> int: + total_replaced = 0 + + # Handle patterns identified with SubgraphMatcher + for entry in fusable_patterns.values(): + if entry.get_graphs_fn is not None and entry.create_replacement_fn is not None: + total_replaced += create_replacement_for_pattern( + ep, + graph_module, + entry.get_graphs_fn(), + # pyre-ignore[6] + entry.create_replacement_fn, + ) + + # Handle patterns identified with custom detector function + for node in graph_module.graph.nodes: + for entry in fusable_patterns.values(): + if ( + entry.detector_fn is not None + and entry.create_replacement_fn is not None + ): + maybe_match = entry.detector_fn(node) + if maybe_match is not None: + assert entry.create_replacement_fn is not None + entry.create_replacement_fn(ep, graph_module, maybe_match) + total_replaced += 1 + + graph_module.graph.eliminate_dead_code() + return total_replaced diff --git a/backends/vulkan/patterns/pattern_registry.py b/backends/vulkan/patterns/pattern_registry.py new file mode 100644 index 00000000000..9a906cd8770 --- /dev/null +++ b/backends/vulkan/patterns/pattern_registry.py @@ -0,0 +1,116 @@ +# 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 typing import Callable, Dict, List, Optional + +import torch + +from executorch.exir import ExportedProgram + +from torch.fx.passes.utils.matcher_utils import InternalMatch + +GetGraphFn = Callable[[], List[torch.fx.GraphModule]] + + +class PatternMatch: + __slots__ = ("input_nodes", "output_nodes", "all_nodes", "anchor_node") + """ + The design of this class is based on InternalMatch from + torch.fx.passes.utils.matcher_utils. It represents nodes in a graph that + match a particular pattern. + + The reason to not use InternalMatch directly is to enable more (i.e. custom) + methods to detect and represent matches other than through SubgraphMatcher. + """ + + def __init__( + self, + input_nodes: List[torch.fx.Node], + output_nodes: List[torch.fx.Node], + all_nodes: List[torch.fx.Node], + anchor_node: Optional[torch.fx.Node] = None, + ): + self.input_nodes = input_nodes + self.output_nodes = output_nodes + self.all_nodes = all_nodes + self.anchor_node = anchor_node + + +def create_pattern_match_from_internal_match( + internal_match: InternalMatch, +) -> PatternMatch: + return PatternMatch( + internal_match.placeholder_nodes, + internal_match.returning_nodes, + list(internal_match.nodes_map.values()), + ) + + +CreateReplacementFn = Callable[ + [ExportedProgram, torch.fx.GraphModule, PatternMatch], None +] + + +DetectorFn = Callable[[torch.fx.Node], Optional[PatternMatch]] + + +class PatternEntry: + def __init__( + self, + get_graphs_fn: Optional[GetGraphFn] = None, + detector_fn: Optional[DetectorFn] = None, + create_replacement_fn: Optional[CreateReplacementFn] = None, + ): + self.get_graphs_fn = get_graphs_fn + self.detector_fn = detector_fn + self.create_replacement_fn = create_replacement_fn + + def is_valid(self): + return ( + self.get_graphs_fn is not None or self.detector_fn is not None + ) and self.create_replacement_fn is not None + + +fusable_patterns: Dict[str, PatternEntry] = {} + + +def register_pattern_graph(pattern_name: str): + def decorator(fn: GetGraphFn): + if pattern_name not in fusable_patterns: + fusable_patterns[pattern_name] = PatternEntry() + + # Cannot define both get_graphs_fn and detector_fn + assert fusable_patterns[pattern_name].detector_fn is None + fusable_patterns[pattern_name].get_graphs_fn = fn + + return fn + + return decorator + + +def register_pattern_detector(pattern_name: str): + def decorator(fn: DetectorFn): + if pattern_name not in fusable_patterns: + fusable_patterns[pattern_name] = PatternEntry() + + # Cannot define both get_graphs_fn and detector_fn + assert fusable_patterns[pattern_name].get_graphs_fn is None + fusable_patterns[pattern_name].detector_fn = fn + + return fn + + return decorator + + +def register_pattern_replacement(pattern_name: str): + def decorator(fn: CreateReplacementFn): + if pattern_name not in fusable_patterns: + fusable_patterns[pattern_name] = PatternEntry() + + fusable_patterns[pattern_name].create_replacement_fn = fn + return fn + + return decorator diff --git a/backends/vulkan/patterns/quantized_convolution.py b/backends/vulkan/patterns/quantized_convolution.py new file mode 100644 index 00000000000..65b51b5e103 --- /dev/null +++ b/backends/vulkan/patterns/quantized_convolution.py @@ -0,0 +1,224 @@ +# 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 typing import Optional + +import executorch.backends.vulkan.utils as utils + +import torch + +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + get_param_tensor, +) + +from executorch.backends.vulkan.patterns.pattern_registry import ( + PatternMatch, + register_pattern_detector, + register_pattern_replacement, +) + +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops + +from torch.export.graph_signature import InputKind + + +class QuantizedConvolutionMatch(PatternMatch): + def __init__(self, conv_node: torch.fx.Node) -> None: + self.anchor_node = conv_node + self.match_found = False + self.all_nodes = [self.anchor_node] + + # Extract convolution parameters + self.stride = conv_node.args[3] if len(conv_node.args) > 3 else [1, 1] + self.padding = conv_node.args[4] if len(conv_node.args) > 4 else [0, 0] + self.dilation = conv_node.args[5] if len(conv_node.args) > 5 else [1, 1] + self.groups = conv_node.args[8] if len(conv_node.args) > 8 else 1 + + const_node, arg_chain = utils.trace_args_until_placeholder( + self.anchor_node.args[1] + ) + + # weight is not a constant tensor - no match + if const_node is None: + return + + dequantize_weight_node = None + # Search for a dequantize node in the arg chain of weight + for node in arg_chain: + if isinstance(node, torch.fx.Node) and utils.is_dequant_node(node): + dequantize_weight_node = node + # weight is not quantized - no match + if dequantize_weight_node is None: + return + + self.weight_node = const_node + self.dequantize_weight_node = dequantize_weight_node + self.all_nodes.extend(arg_chain) + + # Identify weight quantization parameter nodes + self.weight_scales_node, arg_chain = utils.trace_args_until_placeholder( + self.dequantize_weight_node.args[1] + ) + assert self.weight_scales_node is not None + self.all_nodes.extend(arg_chain) + + self.weight_zeros_node, arg_chain = utils.trace_args_until_placeholder( + self.dequantize_weight_node.args[2] + ) + assert self.weight_zeros_node is not None + self.all_nodes.extend(arg_chain) + + # Identify output node + self.output_node = self.anchor_node + + out_channels = self.output_node.meta["val"].shape[-1] + # The implementation requires that for grouped convolutions, a group does not + # cross any texel boundary. The output channels per group must be a multiple of + # 4. If this is not true, then don't match the pattern. + if self.groups > 1 and (out_channels / self.groups) % 4 == 0: + return + + # Identify bias node, if applicable + self.bias_node = None + if len(self.anchor_node.args) > 2 and self.anchor_node.args[2] is not None: + self.bias_node, arg_chain = utils.trace_args_until_placeholder( + self.anchor_node.args[2] + ) + if self.bias_node is not None: + self.all_nodes.extend(arg_chain) + + # Identify input node + self.fp_input_node, self.quantize_input_node, dq_node = ( + utils.maybe_skip_q_dq_arg_chain(self.anchor_node.args[0]) + ) + assert self.fp_input_node is not None + self.all_nodes.append(self.fp_input_node) + assert self.quantize_input_node is not None + assert dq_node is not None + + self.input_scales_node = self.quantize_input_node.args[1] + self.input_zeros_node = self.quantize_input_node.args[2] + + self.all_nodes.extend( + [ + self.quantize_input_node, + dq_node, + ] + ) + + self.match_found = True + + +convolution_anchor_nodes = { + exir_ops.edge.aten.conv2d.default, + exir_ops.edge.aten.convolution.default, +} + + +@register_pattern_detector("quantized_convolution") +def find_quantized_convolution_patterns( + node: torch.fx.Node, +) -> Optional[QuantizedConvolutionMatch]: + if node.target not in convolution_anchor_nodes: + return None + + matched_pattern = QuantizedConvolutionMatch(node) + if matched_pattern.match_found: + return matched_pattern + + return None + + +## +## Pattern Replacement +## + + +@register_pattern_replacement("quantized_convolution") +def make_conv2d_q8ta_q8csw_custom_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedConvolutionMatch, +): + weight_tensor = get_param_tensor(ep, match.weight_node) + assert weight_tensor is not None + + assert match.weight_scales_node is not None + weight_scales_tensor = get_param_tensor(ep, match.weight_scales_node) + assert weight_scales_tensor is not None + + assert match.weight_zeros_node is not None + weight_zeros_tensor = get_param_tensor(ep, match.weight_zeros_node) + assert weight_zeros_tensor is not None + + bias_tensor = None + if match.bias_node is not None: + bias_tensor = get_param_tensor(ep, match.bias_node) + assert bias_tensor is not None + + OC, IC, H, W = weight_tensor.shape + + # Reshape weight tensor from (OC, IC, H, W) to (OC, H * W * IC) (i.e. matrix format) + # This prepares the weights for Im2Col-based convolution + weight_tensor = ( + weight_tensor.permute(0, 2, 3, 1).contiguous().view(OC, H * W * IC).contiguous() + ) + + # Need to make sure that OC dim is a multiple of 4 so that data load/stores are well + # aligned with texel boundaries. Add padding to align to the next multiple of 4 if + # needed. + utils.align_width_and_update_state_dict( + ep, match.weight_node, weight_tensor, force_update=True + ) + utils.align_width_and_update_state_dict( + ep, match.weight_scales_node, weight_scales_tensor + ) + if bias_tensor is not None: + utils.align_width_and_update_state_dict(ep, match.bias_node, bias_tensor) + + first_graph_node = list(graph_module.graph.nodes)[0] + with graph_module.graph.inserting_before(first_graph_node): + qweight_tensor_name = utils.get_tensor_name(ep, match.weight_node) + # Pre-compute the weight sums which are needed to apply activation zero point + # when using integer accumulation. For the reshaped 2D weight matrix (IC * H * W, OC), + # sum over dimension 0 to get sums per output channel + sum_per_output_channel = weight_tensor.sum(dim=1).to(torch.int32).contiguous() + sums_name = qweight_tensor_name + "_sums" + # Sanitize the name + sums_name = sums_name.replace(".", "_") + + weight_sums_node = create_constant_placeholder( + exp_program=ep, + graph=graph_module.graph, + kind=InputKind.CONSTANT_TENSOR, + name=sums_name, + data=sum_per_output_channel, + ) + + with graph_module.graph.inserting_before(match.output_node): + qconv_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.conv2d_q8ta_q8csw.default, + args=( + match.fp_input_node, + match.input_scales_node, + match.input_zeros_node, + match.weight_node, + weight_sums_node, + match.weight_scales_node, + match.bias_node, # Add bias after weight_scales + [H, W], # Pass kernel size information before stride + match.stride, + match.padding, + match.dilation, + match.groups, + ), + ) + + qconv_node.meta["val"] = match.output_node.meta["val"] + match.output_node.replace_all_uses_with(qconv_node) diff --git a/backends/vulkan/patterns/quantized_linear.py b/backends/vulkan/patterns/quantized_linear.py new file mode 100644 index 00000000000..882d0d41e6d --- /dev/null +++ b/backends/vulkan/patterns/quantized_linear.py @@ -0,0 +1,498 @@ +# 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 operator + +from typing import Optional + +import executorch.backends.vulkan.utils as utils + +import torch +import torch.nn.functional as F + +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + get_param_tensor, +) + +from executorch.backends.vulkan.patterns.pattern_registry import ( + PatternMatch, + register_pattern_detector, + register_pattern_replacement, +) + +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops + +from torch.export.graph_signature import InputKind + + +class QuantizedLinearMatch(PatternMatch): + def __init__(self, mm_node: torch.fx.Node) -> None: + self.anchor_node = mm_node + self.match_found = False + self.all_nodes = [self.anchor_node] + + const_node, arg_chain = utils.trace_args_until_placeholder( + self.anchor_node.args[1] + ) + + # mat2 is not a constant tensor - no match + if const_node is None: + return + + dequantize_weight_node = None + # Search for a dequantize node in the arg chain of weight + for node in arg_chain: + if isinstance(node, torch.fx.Node) and utils.is_dequant_node(node): + dequantize_weight_node = node + # weight is not quantized - no match + if dequantize_weight_node is None: + return + + self.weight_node = const_node + self.dequantize_weight_node = dequantize_weight_node + self.all_nodes.extend(arg_chain) + + # By default, assume dequant node is from quantized_decomposed namespace + scales_arg_idx = 1 + zeros_arg_idx = 2 + # torchao dequantize has a different function schema than quantized_decomposed + if ( + self.dequantize_weight_node.target + == exir_ops.edge.torchao.dequantize_affine.default + ): + scales_arg_idx = 2 + zeros_arg_idx = 3 + + # Identify weight quantization parameter nodes + self.weight_scales_node, arg_chain = utils.trace_args_until_placeholder( + self.dequantize_weight_node.args[scales_arg_idx] + ) + assert self.weight_scales_node is not None + self.all_nodes.extend(arg_chain) + + self.weight_zeros_node, arg_chain = utils.trace_args_until_placeholder( + self.dequantize_weight_node.args[zeros_arg_idx] + ) + assert self.weight_zeros_node is not None + self.all_nodes.extend(arg_chain) + + # Identify output node + self.output_node = self.anchor_node + + # The implementation has a limitation that output channels must be a + # multiple of 4. This is to ensure that data loads are aligned well with + # texel boundaries. If this is not true, then don't match the pattern. + out_channels = self.output_node.meta["val"].shape[-1] + if out_channels % 4 != 0: + return + + # Identify input node + self.fp_input_node, self.quantize_input_node, dq_node = ( + utils.maybe_skip_q_dq_arg_chain(self.anchor_node.args[0]) + ) + assert self.fp_input_node is not None + self.all_nodes.append(self.fp_input_node) + + # The implementation has a limitation that input channels must be a + # multiple of 4. This is to ensure that data loads are aligned well with + # texel boundaries. If this is not true, then don't match the pattern. + in_channels = self.fp_input_node.meta["val"].shape[-1] + if in_channels % 4 != 0: + return + + # Identify bias node, if applicable + self.bias_node = None + if self.anchor_node.target == exir_ops.edge.aten.addmm.default: + self.bias_node, arg_chain = utils.trace_args_until_placeholder( + self.anchor_node.args[2] + ) + assert self.bias_node is not None + self.all_nodes.extend(arg_chain) + + # If input is not quantized, then we are done + if self.quantize_input_node is None: + self.match_found = True + return + + scales_arg_idx = 1 + zeros_arg_idx = 2 + + # torchao op has a slightly different function schema + if ( + self.quantize_input_node.target + == exir_ops.edge.torchao.quantize_affine.default + ): + scales_arg_idx = 2 + zeros_arg_idx = 3 + + self.input_scales_node = self.quantize_input_node.args[scales_arg_idx] + self.input_zeros_node = self.quantize_input_node.args[zeros_arg_idx] + + assert dq_node is not None + self.all_nodes.extend( + [ + self.quantize_input_node, + dq_node, + ] + ) + + self.match_found = True + + def is_weight_only_quantized(self) -> bool: + return self.quantize_input_node is None + + def is_weight_pergroup_quantized(self) -> bool: + weight_shape = self.weight_node.meta["val"].shape + scales_shape = self.weight_scales_node.meta["val"].shape + if len(scales_shape) != 2: + return False + + # Check that: + # height dim of scales is same as height dim of weight (N / output channels dim) + # width dim of weight (K / in channels dim) is divisible by width dim of scales + # (number of quantization groups) + return scales_shape[-2] == weight_shape[-2] and ( + weight_shape[-1] % scales_shape[-1] == 0 + ) + + def is_weight_perchannel_quantized(self) -> bool: + weight_shape = self.weight_node.meta["val"].shape + scales_shape = self.weight_scales_node.meta["val"].shape + if len(scales_shape) != 1: + return False + + # scales should have same size as weight's output channels dim + return scales_shape[0] == weight_shape[-2] + + def is_input_static_per_tensor_quantized(self) -> bool: + if self.quantize_input_node is None: + return False + + # For static quantization per tensor quantization, the scales and zeros + # are scalars. + return isinstance(self.input_scales_node, float) + + def is_input_dynamic_perchannel_quantized(self) -> bool: + if self.quantize_input_node is None: + return False + + if not isinstance(self.input_scales_node, torch.fx.Node): + return False + + # For dynamic quantization, input scale node should be a getitem operator + # retrieving the output of a choose_qparams op + if self.input_scales_node.target != operator.getitem: + return False + + # The getitem node should be retrieving from a choose_qparams op + if not utils.is_choose_qparams_node(self.input_scales_node.args[0]): + return False + + scales_shape = self.input_scales_node.meta["val"].shape + input_shape = self.fp_input_node.meta["val"].shape + + return input_shape[-2] == scales_shape[-1] + + +linear_anchor_nodes = { + exir_ops.edge.aten.linear.default, + exir_ops.edge.aten.mm.default, + exir_ops.edge.aten.addmm.default, +} + + +@register_pattern_detector("quantized_linear") +def find_quantized_linear_patterns( + node: torch.fx.Node, +) -> Optional[QuantizedLinearMatch]: + if node.target not in linear_anchor_nodes: + return None + + matched_pattern = QuantizedLinearMatch(node) + if matched_pattern.match_found: + return matched_pattern + + return None + + +## +## Constant tensor manipulation +## + + +def pack_4bit_weight_tensor(weight_tensor: torch.Tensor) -> torch.Tensor: + """ + Given a 8-bit weight tensor containing values quantized to 4 bits, create a packed + weight tensor by transposing the weight tensor, then packing 2 4-bit values in one + 8-bit value. + + An input weight tensor of shape (N, K) will produce a packed weight tensor of shape + (K, N / 2). + """ + + # Assert we got a properly quantized tensor. + min_val, max_val = weight_tensor.min().item(), weight_tensor.max().item() + assert ( + max_val <= 7 and min_val >= -8 + ), f"pack_4bit_weight_tensor: [min_val,max_val] out of [-8, 7] range, got [{min_val}, {max_val}]" + + # Assuming we have a 2d tensor + if weight_tensor.ndim != 2: + weight_tensor = weight_tensor.squeeze() + assert ( + weight_tensor.ndim == 2 + ), f"pack_4bit_weight_tensor: expecting input tensor to be 2d, got {weight_tensor.ndim}" + + # Need to pad innermost dim to be a multiple of 8, since the minimum load granularity + # is int32 (4 bytes), which contains 8 4-bit values. + if weight_tensor.shape[-1] % 8 != 0: + num_pad = 8 - (weight_tensor.shape[-1] % 8) + weight_tensor = F.pad(input=weight_tensor, pad=(0, num_pad)) + + # Shape after padding + _, in_channels = weight_tensor.shape + assert in_channels % 8 == 0, "convert_to_qc4w: expecting ic to be divisible by 8" + + # Adjust weight_tensor tensor for zp + weight_tensor = weight_tensor.to(dtype=torch.uint8) + 8 + # Pack each 4-bit value into a single 8-bit value + return weight_tensor[::, 1::2] << 4 | weight_tensor[::, ::2] + + +def compute_per_group_sums(weight_tensor: torch.Tensor, group_size: int): + """ + Compute the sum of weights per quantization group. + + Args: + weight_tensor (torch.Tensor): Tensor of shape [out_channels, in_channels], dtype int8. + group_size (int): Number of input channels per quantization group. + + Returns: + torch.Tensor: Tensor of shape [num_groups, out_channels], where num_groups = in_channels // group_size. + """ + out_channels, in_channels = weight_tensor.shape + num_groups = in_channels // group_size + # Reshape to [out_channels, num_groups, group_size] + reshaped = weight_tensor.view(out_channels, num_groups, group_size) + # Sum over group_size dimension to get [out_channels, num_groups] + sums = reshaped.sum(dim=2) + # Transpose to [num_groups, out_channels] + sums = sums.transpose(0, 1).contiguous() + # Pad out_channels dim (dim=1) to be a multiple of 8 if needed + out_channels = sums.shape[1] + if out_channels % 8 != 0: + num_pad = 8 - (out_channels % 8) + sums = F.pad(sums, (0, num_pad)) + + return sums.to(torch.int32).contiguous() + + +## +## Pattern Replacement +## + + +def make_linear_q4gsw_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedLinearMatch, + weight_tensor: torch.Tensor, + weight_scales_tensor: torch.Tensor, +): + num_groups = weight_scales_tensor.shape[-1] + in_channels = weight_tensor.shape[-1] + group_size = in_channels // num_groups + + weight_tensor = pack_4bit_weight_tensor(weight_tensor) + # Use this function for convenience to update the state dict with the packed + # weight tensor. Alignment will already have been done in the above function. + weight_tensor = utils.align_width_and_update_state_dict( + ep, match.weight_node, weight_tensor, align_to=1, force_update=True + ) + + # Also transpose the weight scales tensor to shape [num_groups, N] + weight_scales_tensor = weight_scales_tensor.transpose(0, 1).contiguous() + # Align to multiple of 8 to ensure that data loads from the weight scales + # tensor do not go out of bounds. Each thread computes 8 output channels. + utils.align_width_and_update_state_dict( + ep, + match.weight_scales_node, + weight_scales_tensor, + align_to=8, + force_update=True, + ) + + with graph_module.graph.inserting_before(match.output_node): + linear_q4gsw_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.linear_q4gsw.default, + args=( + match.fp_input_node, + match.weight_node, + match.weight_scales_node, + group_size, + ), + ) + + linear_q4gsw_node.meta["val"] = match.output_node.meta["val"] + match.output_node.replace_all_uses_with(linear_q4gsw_node) + + +def make_linear_dq8ca_q4gsw_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedLinearMatch, + weight_tensor: torch.Tensor, + weight_scales_tensor: torch.Tensor, +): + num_groups = weight_scales_tensor.shape[-1] + in_channels = weight_tensor.shape[-1] + group_size = in_channels // num_groups + + # Compute per quant group sums before packing the weight tensor + sum_per_quant_group = compute_per_group_sums(weight_tensor, group_size) + + weight_tensor = pack_4bit_weight_tensor(weight_tensor) + # Use this function for convenience to update the state dict with the packed + # weight tensor. Alignment will already have been done in the above function. + weight_tensor = utils.align_width_and_update_state_dict( + ep, match.weight_node, weight_tensor, align_to=1, force_update=True + ) + + # Also transpose the weight scales tensor to shape [num_groups, N] + weight_scales_tensor = weight_scales_tensor.transpose(0, 1).contiguous() + utils.align_width_and_update_state_dict( + ep, + match.weight_scales_node, + weight_scales_tensor, + align_to=1, + force_update=True, + ) + + first_graph_node = list(graph_module.graph.nodes)[0] + with graph_module.graph.inserting_before(first_graph_node): + weight_tensor_name = utils.get_tensor_name(ep, match.weight_node) + # Pre-compute the weight sums which are needed to apply activation zero point + # when using integer accumulation. + sums_name = weight_tensor_name + "_sums" + # Sanitize the name + sums_name = sums_name.replace(".", "_") + + weight_sums_node = create_constant_placeholder( + exp_program=ep, + graph=graph_module.graph, + kind=InputKind.CONSTANT_TENSOR, + name=sums_name, + data=sum_per_quant_group, + ) + + with graph_module.graph.inserting_before(match.output_node): + qlinear_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default, + args=( + match.fp_input_node, + match.input_scales_node, + match.input_zeros_node, + match.weight_node, + weight_sums_node, + match.weight_scales_node, + group_size, + ), + ) + + qlinear_node.meta["val"] = match.output_node.meta["val"] + match.output_node.replace_all_uses_with(qlinear_node) + + +def make_linear_q8ta_q8csw_custom_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedLinearMatch, + weight_tensor: torch.Tensor, +): + first_graph_node = list(graph_module.graph.nodes)[0] + with graph_module.graph.inserting_before(first_graph_node): + weight_tensor_name = utils.get_tensor_name(ep, match.weight_node) + # Pre-compute the weight sums which are needed to apply activation zero point + # when using integer accumulation. + sum_per_output_channel = weight_tensor.sum(dim=1).to(torch.int32).contiguous() + sums_name = weight_tensor_name + "_sums" + # Sanitize the name + sums_name = sums_name.replace(".", "_") + + weight_sums_node = create_constant_placeholder( + exp_program=ep, + graph=graph_module.graph, + kind=InputKind.CONSTANT_TENSOR, + name=sums_name, + data=sum_per_output_channel, + ) + + with graph_module.graph.inserting_before(match.output_node): + qlinear_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.linear_q8ta_q8csw.default, + args=( + match.fp_input_node, + match.input_scales_node, + match.input_zeros_node, + match.weight_node, + weight_sums_node, + match.weight_scales_node, + ), + ) + + qlinear_node.meta["val"] = match.output_node.meta["val"] + match.output_node.replace_all_uses_with(qlinear_node) + + +@register_pattern_replacement("quantized_linear") +def replace_quantized_linear_patterns( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: QuantizedLinearMatch, +): + # Extract relevant tensors + weight_tensor = get_param_tensor(ep, match.weight_node) + assert weight_tensor is not None + + assert match.weight_scales_node is not None + weight_scales_tensor = get_param_tensor(ep, match.weight_scales_node) + assert weight_scales_tensor is not None + + assert match.weight_zeros_node is not None + weight_zeros_tensor = get_param_tensor(ep, match.weight_zeros_node) + assert weight_zeros_tensor is not None + + # Biases not supported at the moment + if match.bias_node is not None: + return + + # Route to appropriate custom op + if ( + match.is_weight_only_quantized() + and match.is_weight_pergroup_quantized() + and utils.is_in_4bit_range(weight_tensor) + ): + make_linear_q4gsw_op( + ep, graph_module, match, weight_tensor, weight_scales_tensor + ) + elif ( + match.is_input_dynamic_perchannel_quantized() + and match.is_weight_pergroup_quantized() + and utils.is_in_4bit_range(weight_tensor) + ): + make_linear_dq8ca_q4gsw_op( + ep, graph_module, match, weight_tensor, weight_scales_tensor + ) + elif ( + match.is_input_static_per_tensor_quantized() + and match.is_weight_perchannel_quantized() + ): + make_linear_q8ta_q8csw_custom_op(ep, graph_module, match, weight_tensor) diff --git a/backends/vulkan/patterns/rope.py b/backends/vulkan/patterns/rope.py new file mode 100644 index 00000000000..b174224ab78 --- /dev/null +++ b/backends/vulkan/patterns/rope.py @@ -0,0 +1,173 @@ +# 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 operator + +from functools import lru_cache +from typing import List, Optional + +import torch + +from executorch.backends.vulkan.patterns.pattern_registry import ( + PatternMatch, + register_pattern_graph, + register_pattern_replacement, +) + +from executorch.exir import EdgeCompileConfig, ExportedProgram, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + +from torch.export import export + + +class RotaryEmbeddingPattern(torch.nn.Module): + """ + Implementation of rotary embedding pattern that matches the one + in examples/model/llama/rope.py + """ + + def __init__(self): + super().__init__() + + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + # This implementation matches the apply_rotary_emb function in rope.py + # Split into real and imaginary parts + xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1) + xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1) + + # Reshape frequencies for broadcasting + freqs_cos = self._reshape_for_broadcast(freqs_cos, xq_r) + freqs_sin = self._reshape_for_broadcast(freqs_sin, xq_r) + + # Apply rotary embedding + xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin + xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos + xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin + xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos + + # Recombine real and imaginary parts + xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3) + xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3) + + return xq_out.type_as(xq), xk_out.type_as(xk) + + def _reshape_for_broadcast(self, freqs_cis: torch.Tensor, x: torch.Tensor): + ndim = x.ndim + freqs_cis_ndim = freqs_cis.ndim + if freqs_cis_ndim == 3: + # freqs_cis: (seq_len, n_heads, head_dim // 2) + assert freqs_cis.shape == (x.shape[-3], x.shape[-2], x.shape[-1]) + shape = [ + d if (i == ndim - 3 or i == ndim - 2 or i == ndim - 1) else 1 + for i, d in enumerate(x.shape) + ] + else: + # freqs_cis: (seq_len, head_dim // 2) + assert freqs_cis.shape == (x.shape[1], x.shape[-1]) + shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] + return freqs_cis.view(shape) + + +@lru_cache(maxsize=2) +@register_pattern_graph("export_llama_rope") +def get_rope_graphs() -> List[torch.fx.GraphModule]: + batch_size = 1 + seq_len = 1 + n_heads = 4 + n_kv_heads = 2 + head_dim = 32 + + graphs = [] + dtype = torch.float32 + + xq = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=dtype) + xk = torch.randn(batch_size, seq_len, n_kv_heads, head_dim, dtype=dtype) + freqs_cos = torch.randn(seq_len, head_dim // 2, dtype=dtype) + freqs_sin = torch.randn(seq_len, head_dim // 2, dtype=dtype) + + edge = to_edge( + export( + RotaryEmbeddingPattern(), + (xq, xk, freqs_cos, freqs_sin), + strict=True, + ), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + gm = edge.exported_program().graph_module + graphs.append(gm) + + return graphs + + +def identify_rotary_emb_io_nodes( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: PatternMatch, +) -> Optional[List[torch.fx.Node]]: + # Get the input inputs (xq, xk, freqs_cos, freqs_sin) + input_nodes = match.input_nodes + if len(input_nodes) != 4: + return None + + xq, xk, freqs_cos, freqs_sin = input_nodes + + output_nodes = match.output_nodes + if len(output_nodes) != 2: + return None + + xq_out, xk_out = output_nodes + + return [xq, xk, freqs_cos, freqs_sin, xq_out, xk_out] + + +@register_pattern_replacement("export_llama_rope") +def create_rotary_emb_custom_op( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: PatternMatch, +): + io_nodes = identify_rotary_emb_io_nodes(ep, graph_module, match) + if io_nodes is None: + return + + assert len(io_nodes) == 6 + xq, xk, freqs_cos, freqs_sin, xq_out, xk_out = io_nodes + + # Create the custom op node + with graph_module.graph.inserting_before(xq_out): + rotary_emb_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.et_vk.apply_rotary_emb.default, + args=(xq, xk, freqs_cos, freqs_sin), + ) + + # The custom op returns a tuple (xq_out, xk_out) + # We need to extract the individual outputs + with graph_module.graph.inserting_after(rotary_emb_node): + getitem_0 = graph_module.graph.create_node( + "call_function", + operator.getitem, + args=(rotary_emb_node, 0), + ) + getitem_1 = graph_module.graph.create_node( + "call_function", + operator.getitem, + args=(rotary_emb_node, 1), + ) + + if hasattr(xq_out, "meta") and "val" in xq_out.meta: + getitem_0.meta["val"] = xq_out.meta["val"] + if hasattr(xk_out, "meta") and "val" in xk_out.meta: + getitem_1.meta["val"] = xk_out.meta["val"] + + xq_out.replace_all_uses_with(getitem_0) + xk_out.replace_all_uses_with(getitem_1) diff --git a/backends/vulkan/quantizer/TARGETS b/backends/vulkan/quantizer/TARGETS index 2c3ae37923a..be70bafb670 100644 --- a/backends/vulkan/quantizer/TARGETS +++ b/backends/vulkan/quantizer/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "vulkan_quantizer", srcs = ["vulkan_quantizer.py"], deps = [ @@ -11,7 +11,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "vulkan_quantizer_utils", srcs = ["vulkan_quantizer_utils.py"], deps = [ diff --git a/backends/vulkan/runtime/VulkanBackend.cpp b/backends/vulkan/runtime/VulkanBackend.cpp index 4ff0f9e93d6..67b646ae1a8 100644 --- a/backends/vulkan/runtime/VulkanBackend.cpp +++ b/backends/vulkan/runtime/VulkanBackend.cpp @@ -22,6 +22,7 @@ #include #endif // ET_EVENT_TRACER_ENABLED #include +#include #include #include @@ -47,7 +48,9 @@ using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::kTensorDimensionLimit; +using executorch::runtime::NamedDataMap; using executorch::runtime::Result; +using executorch::runtime::Span; using namespace vkcompute; @@ -65,14 +68,6 @@ using BytesVector = const flatbuffers::Vector>*; using UIntVector = const flatbuffers::Vector*; -const uint8_t* get_constant_data_ptr( - VkGraphPtr flatbuffer_graph, - const int32_t buffer_idx, - const uint8_t* constant_data) { - VkBytesPtr constant_bytes = flatbuffer_graph->constants()->Get(buffer_idx); - return constant_data + constant_bytes->offset(); -} - vkapi::ScalarType get_scalar_type(const vkgraph::VkDataType& vk_datatype) { switch (vk_datatype) { case vkgraph::VkDataType::BOOL: @@ -91,6 +86,32 @@ vkapi::ScalarType get_scalar_type(const vkgraph::VkDataType& vk_datatype) { return vkapi::kFloat; case vkgraph::VkDataType::FLOAT64: return vkapi::kDouble; + default: + VK_THROW("Invalid VkDataType type encountered!"); + } +} + +vkapi::ScalarType equivalent_scalar_type( + const executorch::runtime::etensor::ScalarType& et_datatype) { + switch (et_datatype) { + case executorch::runtime::etensor::ScalarType::Byte: + return vkapi::kByte; + case executorch::runtime::etensor::ScalarType::Char: + return vkapi::kChar; + case executorch::runtime::etensor::ScalarType::Int: + return vkapi::kInt; + case executorch::runtime::etensor::ScalarType::Long: + return vkapi::kLong; + case executorch::runtime::etensor::ScalarType::Half: + return vkapi::kHalf; + case executorch::runtime::etensor::ScalarType::Float: + return vkapi::kFloat; + case executorch::runtime::etensor::ScalarType::Double: + return vkapi::kDouble; + case executorch::runtime::etensor::ScalarType::Bool: + return vkapi::kBool; + default: + VK_THROW("Invalid etensor::ScalarType encountered!"); } } @@ -165,6 +186,8 @@ class GraphBuilder { ComputeGraph* compute_graph_; VkGraphPtr flatbuffer_; const uint8_t* constant_data_; + const NamedDataMap* named_data_map_; + std::vector loaded_buffers_from_map_; std::vector ref_mapping_; @@ -172,10 +195,13 @@ class GraphBuilder { explicit GraphBuilder( ComputeGraph* compute_graph, VkGraphPtr flatbuffer, - const uint8_t* constant_data) + const uint8_t* constant_data, + const NamedDataMap* named_data_map) : compute_graph_(compute_graph), flatbuffer_(flatbuffer), constant_data_(constant_data), + named_data_map_(named_data_map), + loaded_buffers_from_map_(), ref_mapping_() {} void resize(uint32_t size) { @@ -211,10 +237,27 @@ class GraphBuilder { ValueRef ref; if (tensor_fb->constant_id() >= 0) { - const uint8_t* tensor_data = get_constant_data_ptr( - flatbuffer_, tensor_fb->constant_id(), constant_data_); + VkBytesPtr constant_bytes = + flatbuffer_->constants()->Get(tensor_fb->constant_id()); + + if (constant_bytes->named_key() != nullptr && + constant_bytes->offset() == UINT64_MAX && + named_data_map_ != nullptr) { + const std::string& data_name = constant_bytes->named_key()->str(); + Result buffer = + named_data_map_->get_data(data_name.c_str()); - ref = compute_graph_->add_tensorref(dims_vector, dtype, tensor_data); + VK_CHECK_COND( + buffer.ok(), + "Failed to get constant data for key %s from named_data_map. Error code: %u", + data_name.c_str(), + static_cast(buffer.error())); + ref = compute_graph_->add_tensorref( + dims_vector, dtype, std::move(buffer.get())); + } else { + const uint8_t* tensor_data = constant_data_ + constant_bytes->offset(); + ref = compute_graph_->add_tensorref(dims_vector, dtype, tensor_data); + } } else { ref = compute_graph_->add_tensor( dims_vector, @@ -326,6 +369,15 @@ class GraphBuilder { } } + vkapi::ScalarType get_staging_scalar_type_of(const uint32_t fb_id) { + VkTensorPtr tensor_fb = + flatbuffer_->values()->Get(fb_id)->value_as_VkTensor(); + if (tensor_fb->staging_datatype() == vkgraph::VkDataType::UNSET) { + return get_scalar_type(tensor_fb->datatype()); + } + return get_scalar_type(tensor_fb->staging_datatype()); + } + void build_graph() { // Resize the mapping to the number of values in the flatbuffer resize(flatbuffer_->values()->size()); @@ -342,7 +394,8 @@ class GraphBuilder { for (const uint32_t fb_id : *flatbuffer_->input_ids()) { const ValueRef ref = get_fb_id_valueref(fb_id); if (compute_graph_->val_is_tensor(ref)) { - compute_graph_->set_input_tensor(ref); + compute_graph_->set_input_tensor( + ref, get_staging_scalar_type_of(fb_id)); } else { compute_graph_->set_val_as_input(ref); } @@ -367,7 +420,12 @@ class GraphBuilder { // values as well if the source graph returns parameter nodes. for (const uint32_t fb_id : *flatbuffer_->output_ids()) { const ValueRef ref = get_fb_id_valueref(fb_id); - compute_graph_->set_output_value(ref); + if (compute_graph_->val_is_tensor(ref)) { + compute_graph_->set_output_tensor( + ref, get_staging_scalar_type_of(fb_id)); + } else { + compute_graph_->set_output_value(ref); + } } if (compute_graph_->graphconfig().enable_querypool) { @@ -390,18 +448,20 @@ bool maybe_resize_input( const size_t input_i, executorch::aten::Tensor& et_tensor) { ValueRef in_tensor_ref = graph->inputs()[input_i].value; - vTensorPtr in_tensor = graph->get_tensor(in_tensor_ref); + + const std::vector in_tensor_vk_sizes = + graph->sizes_of(in_tensor_ref); ET_CHECK_MSG( - et_tensor.dim() == in_tensor->sizes().size(), + et_tensor.dim() == in_tensor_vk_sizes.size(), "Cannot resize input tensor: old ndim %zu does not match new ndim %zu", - static_cast(in_tensor->sizes().size()), + static_cast(in_tensor_vk_sizes.size()), static_cast(et_tensor.dim())); bool should_resize = false; std::vector new_sizes(et_tensor.dim()); for (size_t i = 0; i < et_tensor.dim(); i++) { - if (in_tensor->sizes()[i] != et_tensor.sizes()[i]) { + if (in_tensor_vk_sizes[i] != et_tensor.sizes()[i]) { should_resize = true; } new_sizes.at(i) = et_tensor.sizes()[i]; @@ -411,10 +471,11 @@ bool maybe_resize_input( graph->resize_input(input_i, new_sizes); } + const size_t in_tensor_vk_numel = graph->numel_of(in_tensor_ref); ET_CHECK_MSG( - in_tensor->numel() == et_tensor.numel(), + in_tensor_vk_numel == et_tensor.numel(), "Vulkan tensor numel %zu does not match ET tensor numel %zu", - static_cast(in_tensor->numel()), + static_cast(in_tensor_vk_numel), static_cast(et_tensor.numel())); return should_resize; @@ -445,12 +506,14 @@ void maybe_resize_output( const size_t output_i, executorch::aten::Tensor& et_tensor) { ValueRef out_tensor_ref = graph->outputs()[output_i].value; - vTensorPtr out_tensor = graph->get_tensor(out_tensor_ref); + + const std::vector out_tensor_vk_sizes = + graph->sizes_of(out_tensor_ref); executorch::aten::SizesType new_output_size[kTensorDimensionLimit]; - size_t ndim = out_tensor->sizes().size(); + size_t ndim = out_tensor_vk_sizes.size(); for (int i = 0; i < ndim; ++i) { - new_output_size[i] = out_tensor->sizes()[i]; + new_output_size[i] = out_tensor_vk_sizes[i]; } executorch::aten::ArrayRef output_size{ @@ -473,8 +536,10 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { return true; } - ET_NODISCARD Error - compileModel(const void* buffer_pointer, ComputeGraph* compute_graph) const { + ET_NODISCARD Error compileModel( + const void* buffer_pointer, + ComputeGraph* compute_graph, + const NamedDataMap* named_data_map) const { Result header = VulkanDelegateHeader::parse(buffer_pointer); @@ -500,7 +565,8 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { VkGraphPtr flatbuffer_graph = vkgraph::GetVkGraph(flatbuffer_data); - GraphBuilder builder(compute_graph, flatbuffer_graph, constant_data); + GraphBuilder builder( + compute_graph, flatbuffer_graph, constant_data, named_data_map); builder.build_graph(); @@ -526,7 +592,8 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { graph_config.external_adapter = vkapi::set_and_get_external_adapter(); new (compute_graph) ComputeGraph(graph_config); - Error err = compileModel(processed->data(), compute_graph); + const NamedDataMap* named_data_map = context.get_named_data_map(); + Error err = compileModel(processed->data(), compute_graph, named_data_map); // This backend does not need its processed data after compiling the // model. @@ -542,7 +609,7 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { Error execute( ET_UNUSED BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const override { + Span args) const override { EXECUTORCH_SCOPE_PROF("VulkanBackend::execute"); ComputeGraph* compute_graph = static_cast(handle); @@ -556,10 +623,11 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { bool was_resized = maybe_resize_input(compute_graph, i, args[i]->toTensor()); should_propagate_resize = should_propagate_resize || was_resized; - compute_graph->copy_into_staging( + compute_graph->maybe_cast_and_copy_into_staging( compute_graph->inputs()[i].staging, args[i]->toTensor().const_data_ptr(), - args[i]->toTensor().numel()); + args[i]->toTensor().numel(), + equivalent_scalar_type(args[i]->toTensor().scalar_type())); } else if (compute_graph->val_is_symint(iref)) { VK_CHECK_COND( args[i]->isTensor(), @@ -577,13 +645,7 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { } } - // propagate_resize() will re-encode the command buffer so that push - // constants are updated and DynamicDispatchNode can update the compute - // shader, global workgroup size, and local workgroup size to perform the - // model inference. - if (should_propagate_resize || - (compute_graph->graphconfig().expect_dynamic_shapes && - compute_graph->execute_count() == 0u)) { + if (should_propagate_resize) { compute_graph->propagate_resize(); } @@ -597,10 +659,11 @@ class VulkanBackend final : public ::executorch::runtime::BackendInterface { maybe_resize_output(compute_graph, i, args[o]->toTensor()); // args holds inputs directly followed by outputs, so the i'th output // for compute_graph corresponds to the o'th arg - compute_graph->copy_from_staging( + compute_graph->maybe_cast_and_copy_from_staging( compute_graph->outputs()[i].staging, args[o]->toTensor().mutable_data_ptr(), - args[o]->toTensor().numel()); + args[o]->toTensor().numel(), + equivalent_scalar_type(args[o]->toTensor().scalar_type())); } // TensorRef values represent constant tensors which will not have been // modified by the graph execution. Therefore, if a constant tensor is diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp index 68db37b866e..adb8409d28c 100644 --- a/backends/vulkan/runtime/api/Context.cpp +++ b/backends/vulkan/runtime/api/Context.cpp @@ -111,6 +111,24 @@ void Context::check_device_capabilities(const vkapi::ShaderInfo& shader) { shader.kernel_name, vkapi::VulkanExtension::INT8_STORAGE); } } + if (shader.requires_integer_dot_product) { + if (!adapter_p_->supports_int8_dot_product()) { + throw vkapi::ShaderNotSupportedError( + shader.kernel_name, vkapi::VulkanExtension::INTEGER_DOT_PRODUCT); + } + } + if (shader.requires_shader_int64) { + if (!adapter_p_->supports_int64_shader_types()) { + throw vkapi::ShaderNotSupportedError( + shader.kernel_name, vkapi::VulkanExtension::SHADER_INT64); + } + } + if (shader.requires_shader_float64) { + if (!adapter_p_->supports_float64_shader_types()) { + throw vkapi::ShaderNotSupportedError( + shader.kernel_name, vkapi::VulkanExtension::SHADER_FLOAT64); + } + } } vkapi::DescriptorSet Context::get_descriptor_set( diff --git a/backends/vulkan/runtime/api/containers/StagingBuffer.h b/backends/vulkan/runtime/api/containers/StagingBuffer.h index 1e9f569fc4a..6d0e5a4a457 100644 --- a/backends/vulkan/runtime/api/containers/StagingBuffer.h +++ b/backends/vulkan/runtime/api/containers/StagingBuffer.h @@ -48,7 +48,7 @@ class StagingBuffer final { context_p_->register_buffer_cleanup(vulkan_buffer_); } - inline vkapi::ScalarType dtype() { + inline vkapi::ScalarType dtype() const { return dtype_; } @@ -81,6 +81,15 @@ class StagingBuffer final { VK_WHOLE_SIZE); } + template + void cast_and_copy_from(const SRC_T* src, const size_t numel) { + VK_CHECK_COND(numel <= this->numel()); + DST_T* dst = reinterpret_cast(data()); + for (size_t i = 0; i < numel; ++i) { + dst[i] = static_cast(src[i]); + } + } + inline void copy_to(void* dst, const size_t nbytes) { VK_CHECK_COND(nbytes <= this->nbytes()); vmaInvalidateAllocation( @@ -91,6 +100,15 @@ class StagingBuffer final { memcpy(dst, data(), nbytes); } + template + void cast_and_copy_to(DST_T* dst, const size_t numel) { + VK_CHECK_COND(numel <= this->numel()); + const SRC_T* src = reinterpret_cast(data()); + for (size_t i = 0; i < numel; ++i) { + dst[i] = static_cast(src[i]); + } + } + inline void set_staging_zeros() { memset(data(), 0, nbytes()); } diff --git a/backends/vulkan/runtime/api/containers/Tensor.cpp b/backends/vulkan/runtime/api/containers/Tensor.cpp index 64f330de59c..433ae15db4e 100644 --- a/backends/vulkan/runtime/api/containers/Tensor.cpp +++ b/backends/vulkan/runtime/api/containers/Tensor.cpp @@ -14,6 +14,10 @@ namespace vkcompute { namespace api { +/* + * Used to infer the sizes of a tensor that would correspond to a given + * VulkanImage. + */ std::vector calculate_sizes( const vkapi::VulkanImage& image, const utils::GPUMemoryLayout memory_layout) { @@ -143,58 +147,19 @@ bool dim_order_is_valid(const std::vector& dim_order) { return sum == n * (n + 1) / 2; } -/* - * Applies the following transformations to a tensor's dim_order vector: - * 1. Reverse the order of elements so that the fastest moving dimensions are - * first. - * 2. Convert NCHW dimension indices to WHCN indices, so that 0 represents the - * width dimension, 1 represents the height dimension, and 2 represents the - * channels dimension. - * 3. Unsqueeze the dim_order vector to the next multiple of 4. - - * These transformations make it easier to use the dim order in a compute shader - */ -std::vector create_whcn_dim_order( - const std::vector& dim_order) { - size_t ndim = dim_order.size(); - std::vector whcn_order(ndim); - - // Convert from NCHW to WHCN index, and flip the dim order so that the fastest - // moving dimension is first. - // example: { 1, 2, 0} -> { 2, 0, 1} - // {height, width, channels} -> {channels, width, height} - for (size_t whcn_i = 0, nchw_i = (ndim - 1); whcn_i < ndim; - ++whcn_i, --nchw_i) { - whcn_order.at(whcn_i) = ndim - 1 - dim_order.at(nchw_i); - } - - // Unsqueeze to the next multiple of 4 - size_t ndim_up4 = utils::align_up_4(ndim); - whcn_order.resize(ndim_up4); - - // Append unsqueezed dimensions - for (size_t i = ndim; i < ndim_up4; ++i) { - whcn_order.at(i) = i; - } - - return whcn_order; -} - -std::vector unsqueeze_strides( - const std::vector& strides, - const int64_t numel) { - const size_t ndim = strides.size(); - const size_t ndim_up4 = utils::align_up_4(strides.size()); - std::vector unsqueezed_strides(ndim_up4); - for (int32_t i = 1; i <= ndim; ++i) { - int64_t dim_stride = strides.at(ndim - i); - unsqueezed_strides.at(ndim_up4 - i) = dim_stride; - } - - for (int32_t i = ndim + 1; i <= ndim_up4; ++i) { - unsqueezed_strides.at(ndim_up4 - i) = numel; - } - return unsqueezed_strides; +utils::ivec4 flip_and_unsqueeze_ivec4( + const std::vector& tensor_metadata, + const vTensor::Attribute metadata_type, + const size_t numel) { + VK_CHECK_COND(tensor_metadata.size() <= 4); + std::vector flipped_metadata = + flip_and_unsqueeze(tensor_metadata, metadata_type, numel); + return { + flipped_metadata.at(0), + flipped_metadata.at(1), + flipped_metadata.at(2), + flipped_metadata.at(3), + }; } std::vector calculate_padded_sizes( @@ -224,10 +189,14 @@ utils::uvec3 calculate_image_extents( const std::vector& padded_sizes, const std::vector& axis_map, const int32_t packed_dim) { - VK_CHECK_COND(padded_sizes.size() == 4); - VK_CHECK_COND(axis_map.size() == 4); - utils::uvec3 extents({1, 1, 1}); + + // For high dimensional tensors, buffer storage must be used. No need to + // compute image extents in this case. + if (padded_sizes.size() > 4) { + return extents; + } + // First three elements of axis_map indicate which (X,Y,Z) image axis the // width, height, and channels dim of the tensor maps to. for (int whcn_dim = 0; whcn_dim < 3; ++whcn_dim) { @@ -309,7 +278,8 @@ int64_t calculate_gpu_buffer_numel( return numel; } -int32_t pack_into_int32(const std::vector& vec, const int32_t extra) { +template ::value>> +int32_t pack_into_int32(const std::vector& vec, const int32_t extra) { int32_t packed = static_cast( vec.at(0) + (vec.at(1) << 4) + (vec.at(2) << 8) + (vec.at(3) << 12) + (extra << 16)); @@ -322,22 +292,24 @@ int32_t create_hashed_layout( const int32_t packed_dim, const utils::StorageType storage_type) { if (storage_type == utils::kBuffer) { - return pack_into_int32(create_whcn_dim_order(dim_order), 0); + return pack_into_int32( + flip_and_unsqueeze(dim_order, kTensorDimOrder, 0), 0); } return pack_into_int32(axis_map, packed_dim); } size_t calculate_max_ubo_nbytes( - const size_t nbytes_per_ubo, + const size_t min_nbytes_per_ubo, const utils::StorageType storage_type) { - // For texture backed tensors, the metadata fields needed are: - // sizes, logical limits - size_t max_metadata_field_count = 2u; + size_t ivec4_ubo_nbytes = utils::align_up(size_t(16), min_nbytes_per_ubo); + size_t uvec3_ubo_nbytes = utils::align_up(size_t(12), min_nbytes_per_ubo); + size_t int32_ubo_nbytes = utils::align_up(size_t(4), min_nbytes_per_ubo); if (storage_type == utils::kBuffer) { // sizes, strides, dim order, numel - max_metadata_field_count = 4u; + return 3 * ivec4_ubo_nbytes + int32_ubo_nbytes; } - return max_metadata_field_count * nbytes_per_ubo; + // sizes, logical limits + return ivec4_ubo_nbytes + uvec3_ubo_nbytes; } // @@ -517,6 +489,7 @@ void vTensorStorage::transition( vkapi::MemoryAccessFlags prev_access = last_access_.access; const bool prev_written = (prev_access & vkapi::MemoryAccessType::WRITE) != 0; + const bool cur_written = (cur_access & vkapi::MemoryAccessType::WRITE) != 0; VkImageLayout cur_layout = VK_IMAGE_LAYOUT_UNDEFINED; VkImageLayout new_layout = VK_IMAGE_LAYOUT_UNDEFINED; @@ -528,7 +501,13 @@ void vTensorStorage::transition( layout_changed = cur_layout != new_layout; } - if (prev_written || layout_changed) { + // RAW: need to make sure current read sees previous writes + // WAW: need to make sure the current write occurs after previous write so + // the final value is correct. + // WAR: need to make sure previous read does not read the value from the + // current write. + // RAR: no need for synchronization + if (prev_written || cur_written || layout_changed) { VkPipelineStageFlags src_stage = vkapi::vk_stage(prev_stage); if (0u == src_stage) { src_stage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT; @@ -588,9 +567,11 @@ vTensor::vTensor( packed_dim_, storage_type)), // Related to tensor metadata UBOs - nbytes_per_ubo_{context->adapter_ptr()->min_ubo_alignment()}, - max_ubo_nbytes_{calculate_max_ubo_nbytes(nbytes_per_ubo_, storage_type)}, + min_nbytes_per_ubo_{context->adapter_ptr()->min_ubo_alignment()}, + max_ubo_nbytes_{ + calculate_max_ubo_nbytes(min_nbytes_per_ubo_, storage_type)}, uniforms_(), + buffer_meta_(), // Construct Tensor storage storage_(std::make_shared( context, @@ -600,23 +581,16 @@ vTensor::vTensor( sizes, dtype_, allocate_memory)) { - // Derived metadata - std::vector whcn_dim_order(4, 0); - std::vector unsqueezed_strides(4, 0); - // Only calculate derived metadata if needed for the desired storage type. - // Note that logical limits may be used by buffer storage as well in order to - // set global work group sizes for some compute shaders. - if (storage_type == utils::kBuffer) { - whcn_dim_order = create_whcn_dim_order(dim_order_); - unsqueezed_strides = unsqueeze_strides(strides_, numel_); + // uniform_data_ only valid for low dim tensors + if (sizes.size() <= 4) { + uniform_data_ = std::make_shared(UniformData{ + numel_, + sizes_, + dim_order_, + strides_, + calculate_logical_limits(storage_->image_extents_, axis_map_)}); } - uniform_data_ = std::make_shared(UniformData{ - sizes_, - whcn_dim_order, - unsqueezed_strides, - calculate_logical_limits(storage_->image_extents_, axis_map_), - numel_}); VK_CHECK_COND( dim_order_is_valid(dim_order_), "computed dim order is invalid"); } @@ -641,18 +615,19 @@ vTensor::vTensor( packed_dim_, utils::kTexture3D)), // Related to tensor metadata UBOs - nbytes_per_ubo_{context->adapter_ptr()->min_ubo_alignment()}, + min_nbytes_per_ubo_{context->adapter_ptr()->min_ubo_alignment()}, max_ubo_nbytes_{ - calculate_max_ubo_nbytes(nbytes_per_ubo_, utils::kTexture3D)}, + calculate_max_ubo_nbytes(min_nbytes_per_ubo_, utils::kTexture3D)}, uniforms_(), + buffer_meta_(), // Construct Tensor storage storage_(std::make_shared(context, image)) { uniform_data_ = std::make_shared(UniformData{ + numel_, sizes_, {0, 0, 0, 0}, {0, 0, 0, 0}, - calculate_logical_limits(storage_->image_extents_, axis_map_), - numel_}); + calculate_logical_limits(storage_->image_extents_, axis_map_)}); } vTensor::vTensor(vTensor& other) @@ -665,9 +640,10 @@ vTensor::vTensor(vTensor& other) strides_(other.strides_.begin(), other.strides_.end()), numel_(other.numel_), hashed_layout_(other.hashed_layout_), - nbytes_per_ubo_{other.nbytes_per_ubo_}, + min_nbytes_per_ubo_{other.min_nbytes_per_ubo_}, max_ubo_nbytes_{other.max_ubo_nbytes_}, uniforms_(), + buffer_meta_(), // Copy Tensor storage storage_(other.storage_) { uniform_data_ = std::make_shared(*other.get_uniform_data()); @@ -690,22 +666,36 @@ vTensor::vTensor( axis_map_, packed_dim_, other.storage_type())), - nbytes_per_ubo_{other.nbytes_per_ubo_}, + min_nbytes_per_ubo_{other.min_nbytes_per_ubo_}, max_ubo_nbytes_{other.max_ubo_nbytes_}, uniforms_(), + buffer_meta_(), // Copy Tensor storage storage_(other.storage_) { uniform_data_ = std::make_shared(UniformData{ + static_cast(utils::multiply_integers(sizes_)), sizes_, - create_whcn_dim_order(dim_order_), - unsqueeze_strides(strides_, numel_), - other.logical_limits(), - static_cast(utils::multiply_integers(sizes_))}); + dim_order_, + strides_, + other.logical_limits()}); VK_CHECK_COND( dim_order_is_valid(dim_order_), "new dim order provided is invalid"); } +vTensor::UniformData::UniformData( + const size_t numel_ll, + const std::vector& sizes, + const std::vector& dim_order, + const std::vector& strides, + const utils::uvec3& limits) + : numel(utils::safe_downcast(numel_ll)), + sizes_v(flip_and_unsqueeze_ivec4(sizes, kTensorSizes, numel_ll)), + dim_order_v( + flip_and_unsqueeze_ivec4(dim_order, kTensorDimOrder, numel_ll)), + strides_v(flip_and_unsqueeze_ivec4(strides, kTensorStrides, numel_ll)), + logical_limits(limits) {} + uint32_t vTensor::UniformData::write_attribute( void* dst, const uint32_t dst_offset, @@ -720,11 +710,11 @@ uint32_t vTensor::UniformData::write_attribute( return sizeof(member_name); \ } switch (attr) { + WRITE_ATTRIBUTE_CASE(NUMEL, numel); WRITE_ATTRIBUTE_CASE(SIZES, sizes_v); - WRITE_ATTRIBUTE_CASE(WHCN_DIM_ORDER, whcn_dim_order_v); + WRITE_ATTRIBUTE_CASE(WHCN_DIM_ORDER, dim_order_v); WRITE_ATTRIBUTE_CASE(STRIDES, strides_v); WRITE_ATTRIBUTE_CASE(LOGICAL_LIMITS, logical_limits); - WRITE_ATTRIBUTE_CASE(NUMEL, numel); default: VK_THROW("Invalid Attribute"); } @@ -732,6 +722,38 @@ uint32_t vTensor::UniformData::write_attribute( return 0; } +vTensor::BufferMetadata::BufferMetadata( + std::vector& src_sizes, + std::vector& src_dim_order, + std::vector& src_strides, + size_t src_numel) { + update(src_sizes, src_dim_order, src_strides, src_numel); +} + +void vTensor::BufferMetadata::update( + std::vector& src_sizes, + std::vector& src_dim_order, + std::vector& src_strides, + size_t src_numel) { + int32_t fixed_ndim = utils::safe_downcast(kTensorDimLimit); + + std::vector fu_sizes = flip_and_unsqueeze( + src_sizes, kTensorSizes, src_numel, fixed_ndim); + std::vector fu_dim_order = flip_and_unsqueeze( + src_dim_order, kTensorDimOrder, src_numel, fixed_ndim); + std::vector fu_strides = flip_and_unsqueeze( + src_strides, kTensorStrides, src_numel, fixed_ndim); + + for (int i = 0; i < fixed_ndim; ++i) { + sizes[i] = fu_sizes.at(i); + dim_order[i] = fu_dim_order.at(i); + strides[i] = fu_strides.at(i); + } + + ndim = utils::safe_downcast(src_sizes.size()); + numel = utils::safe_downcast(src_numel); +} + vkapi::VulkanImage& vTensor::image( vkapi::PipelineBarrier& pipeline_barrier, const vkapi::PipelineStageFlags stage) & { @@ -799,84 +821,39 @@ size_t vTensor::get_max_ubo_nbytes(const size_t nbytes_per_ubo) const { } const vkapi::BufferBindInfo vTensor::sizes_ubo() { - if (!uniforms_.buffer()) { - uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); - } - if (sizes_uniform_offset_ == kUniformOffsetUnset) { - VK_CHECK_COND( - (uniforms_size_ + nbytes_per_ubo_) <= max_ubo_nbytes_, - "Uniform data allocation has exceeded Tensor uniform buffer size"); - sizes_uniform_offset_ = uniforms_size_; - uniforms_size_ += nbytes_per_ubo_; - uniforms_.update(utils::make_whcn_ivec4(sizes_), sizes_uniform_offset_); - } - return vkapi::BufferBindInfo( - uniforms_.buffer(), sizes_uniform_offset_, nbytes_per_ubo_); + VK_CHECK_COND(sizes_.size() <= 4); + return metadata_ubo_impl(&sizes_uniform_offset_, uniform_data_->sizes_v); } const vkapi::BufferBindInfo vTensor::dim_order_ubo() { - if (!uniforms_.buffer()) { - uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); - } - if (dim_order_uniform_offset_ == kUniformOffsetUnset) { - VK_CHECK_COND( - (uniforms_size_ + nbytes_per_ubo_) <= max_ubo_nbytes_, - "Uniform data allocation has exceeded Tensor uniform buffer size"); - dim_order_uniform_offset_ = uniforms_size_; - uniforms_size_ += nbytes_per_ubo_; - uniforms_.update( - uniform_data_->whcn_dim_order_v, dim_order_uniform_offset_); - } - return vkapi::BufferBindInfo( - uniforms_.buffer(), dim_order_uniform_offset_, nbytes_per_ubo_); + VK_CHECK_COND(sizes_.size() <= 4); + return metadata_ubo_impl( + &dim_order_uniform_offset_, uniform_data_->dim_order_v); } const vkapi::BufferBindInfo vTensor::strides_ubo() { - if (!uniforms_.buffer()) { - uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); - } - if (strides_uniform_offset == kUniformOffsetUnset) { - VK_CHECK_COND( - (uniforms_size_ + nbytes_per_ubo_) <= max_ubo_nbytes_, - "Uniform data allocation has exceeded Tensor uniform buffer size"); - strides_uniform_offset = uniforms_size_; - uniforms_size_ += nbytes_per_ubo_; - uniforms_.update(uniform_data_->strides_v, strides_uniform_offset); - } - return vkapi::BufferBindInfo( - uniforms_.buffer(), strides_uniform_offset, nbytes_per_ubo_); + VK_CHECK_COND(sizes_.size() <= 4); + return metadata_ubo_impl(&strides_uniform_offset, uniform_data_->strides_v); } const vkapi::BufferBindInfo vTensor::logical_limits_ubo() { - if (!uniforms_.buffer()) { - uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); - } - if (logical_limits_uniform_offset_ == kUniformOffsetUnset) { - VK_CHECK_COND( - (uniforms_size_ + nbytes_per_ubo_) <= max_ubo_nbytes_, - "Uniform data allocation has exceeded Tensor uniform buffer size"); - logical_limits_uniform_offset_ = uniforms_size_; - uniforms_size_ += nbytes_per_ubo_; - uniforms_.update(logical_limits(), logical_limits_uniform_offset_); - } - return vkapi::BufferBindInfo( - uniforms_.buffer(), logical_limits_uniform_offset_, nbytes_per_ubo_); + VK_CHECK_COND(sizes_.size() <= 4); + return metadata_ubo_impl( + &logical_limits_uniform_offset_, uniform_data_->logical_limits); } const vkapi::BufferBindInfo vTensor::numel_ubo() { - if (!uniforms_.buffer()) { - uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); - } - if (numel_uniform_offset_ == kUniformOffsetUnset) { - VK_CHECK_COND( - (uniforms_size_ + nbytes_per_ubo_) <= max_ubo_nbytes_, - "Uniform data allocation has exceeded Tensor uniform buffer size"); - numel_uniform_offset_ = uniforms_size_; - uniforms_size_ += nbytes_per_ubo_; - uniforms_.update(numel(), numel_uniform_offset_); + VK_CHECK_COND(sizes_.size() <= 4); + return metadata_ubo_impl(&numel_uniform_offset_, uniform_data_->numel); +} + +const vkapi::BufferBindInfo vTensor::buffer_meta_ubo() { + size_t ubo_nbytes = sizeof(BufferMetadata); + if (!buffer_meta_.buffer()) { + BufferMetadata data(sizes_, dim_order_, strides_, numel_); + buffer_meta_ = ParamsBuffer(storage_->context_, data); } - return vkapi::BufferBindInfo( - uniforms_.buffer(), numel_uniform_offset_, nbytes_per_ubo_); + return vkapi::BufferBindInfo(buffer_meta_.buffer(), 0, ubo_nbytes); } VkMemoryRequirements vTensor::get_memory_requirements() const { @@ -890,6 +867,16 @@ VkMemoryRequirements vTensor::get_memory_requirements() const { return {}; } +bool vTensor::memory_is_bound() const { + switch (storage_type()) { + case utils::kBuffer: + return storage_->buffer_.has_memory(); + case utils::kTexture2D: + case utils::kTexture3D: + return storage_->image_.has_memory(); + } +} + void vTensor::bind_allocation(const vkapi::Allocation& allocation) { switch (storage_type()) { case utils::kBuffer: @@ -902,37 +889,55 @@ void vTensor::bind_allocation(const vkapi::Allocation& allocation) { } } +void vTensor::acquire_allocation(vkapi::Allocation&& allocation) { + switch (storage_type()) { + case utils::kBuffer: + storage_->buffer_.acquire_allocation(std::move(allocation)); + break; + case utils::kTexture2D: + case utils::kTexture3D: + storage_->image_.acquire_allocation(std::move(allocation)); + break; + } +} + void vTensor::update_metadata() { numel_ = utils::multiply_integers(sizes_); strides_ = calculate_strides(sizes_, dim_order_); // Update uniform data if it has been modified - uniform_data_->numel = numel_; - uniform_data_->sizes_v = utils::make_whcn_ivec4(sizes_); - uniform_data_->whcn_dim_order_v = - utils::make_ivec4(create_whcn_dim_order(dim_order_)); - uniform_data_->strides_v = - utils::make_whcn_ivec4(unsqueeze_strides(strides_, numel_)); - uniform_data_->numel = utils::safe_downcast(numel_); - uniform_data_->logical_limits.limits = - calculate_logical_limits(sizes_, axis_map_, packed_dim_); - - if (sizes_uniform_offset_ != kUniformOffsetUnset) { - uniforms_.update(uniform_data_->sizes_v, sizes_uniform_offset_); - } - if (dim_order_uniform_offset_ != kUniformOffsetUnset) { - uniforms_.update( - uniform_data_->whcn_dim_order_v, dim_order_uniform_offset_); - } - if (strides_uniform_offset != kUniformOffsetUnset) { - uniforms_.update(uniform_data_->strides_v, strides_uniform_offset); - } - if (numel_uniform_offset_ != kUniformOffsetUnset) { - uniforms_.update(numel_, numel_uniform_offset_); + if (sizes_.size() <= 4) { + uniform_data_->numel = utils::safe_downcast(numel_); + uniform_data_->sizes_v = + flip_and_unsqueeze_ivec4(sizes_, kTensorSizes, numel_); + uniform_data_->dim_order_v = + flip_and_unsqueeze_ivec4(dim_order_, kTensorDimOrder, numel_); + uniform_data_->strides_v = + flip_and_unsqueeze_ivec4(strides_, kTensorStrides, numel_); + uniform_data_->logical_limits.limits = + calculate_logical_limits(sizes_, axis_map_, packed_dim_); + + if (sizes_uniform_offset_ != kUniformOffsetUnset) { + uniforms_.update(uniform_data_->sizes_v, sizes_uniform_offset_); + } + if (dim_order_uniform_offset_ != kUniformOffsetUnset) { + uniforms_.update(uniform_data_->dim_order_v, dim_order_uniform_offset_); + } + if (strides_uniform_offset != kUniformOffsetUnset) { + uniforms_.update(uniform_data_->strides_v, strides_uniform_offset); + } + if (numel_uniform_offset_ != kUniformOffsetUnset) { + uniforms_.update(numel_, numel_uniform_offset_); + } + if (logical_limits_uniform_offset_ != kUniformOffsetUnset) { + uniforms_.update( + uniform_data_->logical_limits.limits, logical_limits_uniform_offset_); + } } - if (logical_limits_uniform_offset_ != kUniformOffsetUnset) { - uniforms_.update( - uniform_data_->logical_limits.limits, logical_limits_uniform_offset_); + + if (buffer_meta_.buffer()) { + BufferMetadata data(sizes_, dim_order_, strides_, numel_); + buffer_meta_.update(data); } } diff --git a/backends/vulkan/runtime/api/containers/Tensor.h b/backends/vulkan/runtime/api/containers/Tensor.h index 0e1a1526d88..66c1fd1e4da 100644 --- a/backends/vulkan/runtime/api/containers/Tensor.h +++ b/backends/vulkan/runtime/api/containers/Tensor.h @@ -19,6 +19,8 @@ namespace vkcompute { namespace api { +static constexpr size_t kTensorDimLimit = 8; + /* * Given a GPUMemoryLayout value, produce a dim order vector that matches the * given memory layout. The produced dim order vector will be in the NCHW @@ -36,10 +38,6 @@ std::vector calculate_strides( const std::vector& sizes, const std::vector& dim_order); -std::vector unsqueeze_strides( - const std::vector& strides, - const int64_t numel); - /* * When stored on the GPU, tensor data is stored using texels (i.e. a vector of * 4 scalar values) in order to take advantage of the GPU's native vectorization @@ -236,28 +234,23 @@ class vTensor final { }; class UniformData { + // Contains the number of elements in the tensor according to the canonical + // sizes. + int32_t numel; utils::ivec4 sizes_v; - utils::ivec4 whcn_dim_order_v; + utils::ivec4 dim_order_v; utils::ivec4 strides_v; // See the comments documenting logical_limits() for more context. TextureLimits logical_limits; - // Contains the number of elements in the tensor according to the canonical - // sizes. - int32_t numel; friend class vTensor; UniformData( + const size_t numel_ll, const std::vector& sizes, - const std::vector& whcn_dim_order, + const std::vector& dim_order, const std::vector& strides, - const utils::uvec3& logical_limits, - const size_t numel_ll) - : sizes_v(utils::make_whcn_ivec4(sizes)), - whcn_dim_order_v(utils::make_ivec4(whcn_dim_order)), - strides_v(utils::make_whcn_ivec4(strides)), - logical_limits(logical_limits), - numel(utils::safe_downcast(numel_ll)) {} + const utils::uvec3& limits); public: /* @@ -271,6 +264,26 @@ class vTensor final { const Attribute attr); }; + struct BufferMetadata { + uint32_t sizes[kTensorDimLimit]; + uint32_t dim_order[kTensorDimLimit]; + uint32_t strides[kTensorDimLimit]; + uint32_t ndim; + uint32_t numel; + + BufferMetadata( + std::vector& sizes, + std::vector& dim_order, + std::vector& strides, + size_t numel); + + void update( + std::vector& sizes, + std::vector& dim_order, + std::vector& strides, + size_t numel); + }; + private: /* * "Core" tensor metadata. They are the minimum amount of information required @@ -326,7 +339,7 @@ class vTensor final { int32_t hashed_layout_; // Pre-compute these quantities to avoid frequent re-computation - size_t nbytes_per_ubo_; + size_t min_nbytes_per_ubo_; size_t max_ubo_nbytes_; /* @@ -341,6 +354,11 @@ class vTensor final { */ ParamsBuffer uniforms_; + /* + * Used to store data for BufferMetadata to pass to shaders as buffer_meta_ubo + */ + ParamsBuffer buffer_meta_; + uint32_t uniforms_size_ = 0u; uint32_t sizes_uniform_offset_ = kUniformOffsetUnset; uint32_t dim_order_uniform_offset_ = kUniformOffsetUnset; @@ -523,6 +541,26 @@ class vTensor final { size_t get_max_ubo_nbytes(const size_t nbytes_per_ubo) const; + template + const vkapi::BufferBindInfo metadata_ubo_impl( + uint32_t* param_buffer_offset, + const T& data) { + if (!uniforms_.buffer()) { + uniforms_ = ParamsBuffer(storage_->context_, max_ubo_nbytes_, true); + } + size_t ubo_nbytes = utils::align_up(sizeof(data), min_nbytes_per_ubo_); + if (*param_buffer_offset == kUniformOffsetUnset) { + VK_CHECK_COND( + (uniforms_size_ + ubo_nbytes) <= max_ubo_nbytes_, + "Uniform data allocation has exceeded Tensor uniform buffer size"); + *param_buffer_offset = uniforms_size_; + uniforms_size_ += ubo_nbytes; + uniforms_.update(data, *param_buffer_offset); + } + return vkapi::BufferBindInfo( + uniforms_.buffer(), *param_buffer_offset, ubo_nbytes); + } + public: /* * The functions below return the buffer binding info for a UBO that contains @@ -546,6 +584,8 @@ class vTensor final { const vkapi::BufferBindInfo numel_ubo(); + const vkapi::BufferBindInfo buffer_meta_ubo(); + public: inline size_t staging_buffer_numel() const { return storage_->buffer_len(); @@ -560,6 +600,12 @@ class vTensor final { */ VmaAllocationCreateInfo get_allocation_create_info() const; + /* + * Checks if the tensor's underlying buffer or image resource is bound to a + * memory allocation. + */ + bool memory_is_bound() const; + /* * Return the VkMemoryRequirements of the underlying resource */ @@ -570,6 +616,11 @@ class vTensor final { */ void bind_allocation(const vkapi::Allocation& allocation); + /* + * Binds and acquires a rvalue memory allocation + */ + void acquire_allocation(vkapi::Allocation&& allocation); + private: /* * Assuming sizes, dim order, or axis mapping was modified, recompute all @@ -625,6 +676,7 @@ class vTensor final { } const std::shared_ptr& get_uniform_data() const { + VK_CHECK_COND(sizes_.size() <= 4); return uniform_data_; } }; @@ -638,5 +690,70 @@ static constexpr vTensor::Attribute kTensorLogicalLimits = vTensor::Attribute::LOGICAL_LIMITS; static constexpr vTensor::Attribute kTensorNumel = vTensor::Attribute::NUMEL; +/* + * Prepare tensor metadata vector for consumption on the GPU: + * 1. Convert NCHW dim order and indexes to WCHN dim order and indexes + * 2. Unsqueeze to the next multiple of 4 dims + * 3. Convert to requested output dtype + */ +template < + typename T, + typename std::enable_if::value, int>::type = 0> +std::vector flip_and_unsqueeze( + const std::vector& tensor_metadata, + const vTensor::Attribute metadata_type, + const size_t numel, + const int32_t fixed_ndim = -1) { + const size_t ndim = tensor_metadata.size(); + size_t ndim_up4 = + std::max(utils::align_up_4(tensor_metadata.size()), size_t(4)); + + if (fixed_ndim > 0) { + VK_CHECK_COND(fixed_ndim >= ndim); + ndim_up4 = static_cast(fixed_ndim); + } + + std::vector flipped_metadata(ndim_up4); + + for (int flipped_i = 0; flipped_i < ndim; ++flipped_i) { + T val_at_dim = + utils::safe_downcast(tensor_metadata.at(ndim - 1 - flipped_i)); + if (metadata_type == kTensorDimOrder) { + val_at_dim = utils::safe_downcast(ndim - 1 - val_at_dim); + } + flipped_metadata.at(flipped_i) = val_at_dim; + } + + switch (metadata_type) { + case kTensorStrides: + for (int unsqueezed_i = ndim; unsqueezed_i < ndim_up4; ++unsqueezed_i) { + flipped_metadata.at(unsqueezed_i) = utils::safe_downcast(numel); + } + break; + case kTensorDimOrder: + for (int unsqueezed_i = ndim; unsqueezed_i < ndim_up4; ++unsqueezed_i) { + flipped_metadata.at(unsqueezed_i) = + utils::safe_downcast(unsqueezed_i); + } + break; + // Default: unsqueeze with ones + default: + for (int unsqueezed_i = ndim; unsqueezed_i < ndim_up4; ++unsqueezed_i) { + flipped_metadata.at(unsqueezed_i) = utils::safe_downcast(1); + } + break; + } + + return flipped_metadata; +} + +/* + * Same as flip and unsqueeze, but returns the metadata as an `ivec4`. + */ +utils::ivec4 flip_and_unsqueezed_ivec4( + const std::vector& tensor_metadata, + const vTensor::Attribute metadata_type, + const size_t numel); + } // namespace api } // namespace vkcompute diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py index d42d7ab33be..6db2e01d7e2 100644 --- a/backends/vulkan/runtime/gen_vulkan_spv.py +++ b/backends/vulkan/runtime/gen_vulkan_spv.py @@ -670,7 +670,7 @@ def addSrcAndYamlFiles(self, src_dir_paths: List[str]) -> None: if len(file) > 1: self.template_yaml_files.append(file) - def generateVariantCombinations( + def generateVariantCombinations( # noqa: C901 self, iterated_params: Dict[str, Any], exclude_params: Optional[Set[str]] = None, @@ -679,7 +679,25 @@ def generateVariantCombinations( exclude_params = set() all_iterated_params = [] for param_name, value_list in iterated_params.items(): - if param_name not in exclude_params: + if re.match(r"^combination\d*$", param_name): + param_values = [] + param_names = value_list["parameter_names"] + combos = value_list["combos"] + for combo in combos: + parameter_values = combo["parameter_values"] + if "suffix" in combo: + suffix = combo["suffix"] + else: + suffix = "" + for param_value in parameter_values: + if len(str(param_value)) > 0: + suffix += "_" + str(param_value) + suffix = suffix[1:] + param_values.append((param_names, suffix, parameter_values)) + + all_iterated_params.append(param_values) + + elif param_name not in exclude_params: param_values = [] for value in value_list: if "RANGE" in value: @@ -713,7 +731,7 @@ def generateVariantCombinations( return list(product(*all_iterated_params)) - def parseTemplateYaml(self, yaml_file: str) -> None: + def parseTemplateYaml(self, yaml_file: str) -> None: # noqa: C901 with open(yaml_file) as f: contents = yaml.load(f, Loader=UniqueKeyLoader) for template_name, params_dict in contents.items(): @@ -762,10 +780,21 @@ def parseTemplateYaml(self, yaml_file: str) -> None: default_params_copy[key] = variant[key] variant_name = variant["NAME"] - for param_value in combination: - default_params_copy[param_value[0]] = param_value[2] - if len(str(param_value[1])) > 0: - variant_name = f"{variant_name}_{param_value[1]}" + + for setting in combination: + param_names = setting[0] + suffix = setting[1] + param_values = setting[2] + if isinstance(param_names, list): + for param_name, param_value in zip( + param_names, param_values + ): + default_params_copy[param_name] = param_value + else: + default_params_copy[param_names] = param_values + + if len(str(suffix)) > 0: + variant_name = f"{variant_name}_{suffix}" default_params_copy["NAME"] = variant_name default_params_copy["VARIANT_NAME"] = variant["NAME"] @@ -1083,7 +1112,6 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]: for spv_out_path, glsl_out_path in pool.map( compile_spirv, self.output_file_map.items() ): - print(spv_to_glsl_map) spv_to_glsl_map[spv_out_path] = glsl_out_path return spv_to_glsl_map @@ -1104,6 +1132,9 @@ class ShaderInfo: requires_shader_int16_ext: bool = False requires_16bit_storage_ext: bool = False requires_8bit_storage_ext: bool = False + requires_integer_dot_product_ext: bool = False + requires_shader_int64_ext: bool = False + requires_shader_float64_ext: bool = False def getName(filePath: str) -> str: @@ -1193,7 +1224,7 @@ def determineDescriptorType(lineStr: str) -> str: ) -def getShaderInfo(srcFilePath: str) -> ShaderInfo: +def getShaderInfo(srcFilePath: str) -> ShaderInfo: # noqa: C901 shader_info = ShaderInfo([], [], "") with open(srcFilePath) as srcFile: for line in srcFile: @@ -1214,6 +1245,12 @@ def getShaderInfo(srcFilePath: str) -> ShaderInfo: shader_info.requires_16bit_storage_ext = True if "GL_EXT_shader_8bit_storage" in line: shader_info.requires_8bit_storage_ext = True + if "GL_EXT_integer_dot_product" in line: + shader_info.requires_integer_dot_product_ext = True + if "GL_EXT_shader_explicit_arithmetic_types_int64" in line: + shader_info.requires_shader_int64_ext = True + if "GL_EXT_shader_explicit_arithmetic_types_float64" in line: + shader_info.requires_shader_float64_ext = True return shader_info @@ -1289,6 +1326,9 @@ def to_cpp_str(val: bool): to_cpp_str(shader_info.requires_shader_int16_ext), to_cpp_str(shader_info.requires_16bit_storage_ext), to_cpp_str(shader_info.requires_8bit_storage_ext), + to_cpp_str(shader_info.requires_integer_dot_product_ext), + to_cpp_str(shader_info.requires_shader_int64_ext), + to_cpp_str(shader_info.requires_shader_float64_ext), ] shader_info_str = textwrap.indent( diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index 14328027362..2ec63a89df8 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -151,6 +151,15 @@ ComputeGraph::ComputeGraph(GraphConfig config) config_.prepack_threshold_nbytes = 10 * MB; config_.prepack_initial_threshold_nbytes = 10 * MB; } + if (config_.execute_threshold_node_count == 0) { + config_.execute_threshold_node_count = 128; + config_.execute_initial_threshold_node_count = 64; + } + + // Check if the underlying GPU can access accelerated integer dot product + // instructions + can_use_int8_dot_product_ = + context_->adapter_ptr()->supports_int8_dot_product(); } ComputeGraph::~ComputeGraph() { @@ -202,6 +211,29 @@ 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; + } + + // Check if this ValueRef itself was updated + if (updated_values_.find(idx) != updated_values_.end()) { + return true; + } + + // 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; + } + } + } + + return false; +} + utils::GPUMemoryLayout ComputeGraph::suggested_memory_layout( const std::vector& sizes) { if (config_.enable_memory_layout_override) { @@ -232,6 +264,10 @@ void ComputeGraph::check_no_active_value_ptrs() { "invalidated."); } +bool ComputeGraph::is_valid_value_idx(const ValueRef idx) const noexcept { + return idx >= 0 && idx < static_cast(values_.size()); +} + std::vector ComputeGraph::sizes_of(const ValueRef idx) const { const Value& val = values_.at(idx); if (val.isTensor()) { @@ -274,6 +310,8 @@ vkapi::ScalarType ComputeGraph::dtype_of(const ValueRef idx) const { return val.toConstTensor().dtype(); } else if (val.isTensorRef()) { return val.toConstTensorRef().dtype; + } else if (val.isStaging()) { + return val.toConstStaging().dtype(); } else if (val.isBool()) { return vkapi::ScalarType::Bool; } else if (val.isDouble()) { @@ -296,6 +334,16 @@ bool ComputeGraph::is_contiguous_buffer_tensor(const ValueRef idx) const { return is_contiguous(idx); } +bool ComputeGraph::is_contiguous_texture_tensor(const ValueRef idx) const { + if (!val_is_tensor(idx)) { + return false; + } + if (is_buffer_storage(idx)) { + return false; + } + return has_standard_axis_map(idx) && packed_dim_of(idx) == 0; +} + bool ComputeGraph::is_standard_channels_packed_texture_tensor( const ValueRef idx) const { if (!val_is_tensor(idx)) { @@ -307,15 +355,50 @@ bool ComputeGraph::is_standard_channels_packed_texture_tensor( return has_standard_axis_map(idx) && packed_dim_of(idx) == 2; } -bool ComputeGraph::is_standard_width_packed_texture_tensor( +bool ComputeGraph::is_2d_matrix(const ValueRef idx) const { + std::vector sizes = sizes_of(idx); + const size_t ndim = sizes.size(); + if (sizes.size() < 2) { + return false; + } + if (sizes.size() == 2) { + return true; + } + + // Check that outermost dims have size of 1 + for (int d = 0; d < ndim - 2; d++) { + if (sizes[d] != 1) { + return false; + } + } + + return true; +} + +bool ComputeGraph::is_vectorizable_contiguous_2d_matrix( const ValueRef idx) const { - if (!val_is_tensor(idx)) { + if (!is_2d_matrix(idx)) { return false; } if (is_buffer_storage(idx)) { + return is_contiguous_buffer_tensor(idx) && + size_at(-1, idx) % 4 == 0; + } + return is_contiguous_texture_tensor(idx); +} + +bool ComputeGraph::is_vectorizable_width_packed_tensor( + const ValueRef idx) const { + // Not a tensor - return false + if (!val_is_tensor(idx)) { return false; } - return has_standard_axis_map(idx) && packed_dim_of(idx) == 0; + if (is_buffer_storage(idx)) { + return is_contiguous_buffer_tensor(idx) && + size_at(-1, idx) % 4 == 0; + } + + return is_standard_channels_packed_texture_tensor(idx); } ValueRef ComputeGraph::add_tensor( @@ -325,8 +408,6 @@ ValueRef ComputeGraph::add_tensor( const utils::GPUMemoryLayout memory_layout, const int64_t shared_object_idx, const utils::AxisMapLayout axis_map_layout) { - bool allocate_memory = shared_object_idx < 0; - ValueRef idx(static_cast(values_.size())); check_no_active_value_ptrs(); values_.emplace_back(api::vTensor( @@ -335,10 +416,10 @@ ValueRef ComputeGraph::add_tensor( dtype, storage_type, memory_layout, - allocate_memory, + false, axis_map_layout)); - if (!allocate_memory) { + if (shared_object_idx >= 0) { get_shared_object(shared_object_idx).add_user(this, idx); } return idx; @@ -449,6 +530,17 @@ ValueRef ComputeGraph::add_tensorref( return idx; } +ValueRef ComputeGraph::add_tensorref( + const std::vector& sizes, + const vkapi::ScalarType dtype, + executorch::runtime::FreeableBuffer&& buffer) { + ValueRef idx(static_cast(values_.size())); + check_no_active_value_ptrs(); + values_.emplace_back(TensorRef(sizes, dtype, std::move(buffer))); + total_constant_nbytes_ += values_.back().toConstTensorRef().nbytes(); + return idx; +} + ValueRef ComputeGraph::add_staging( const vkapi::ScalarType dtype, const size_t numel) { @@ -495,21 +587,45 @@ ValueRef ComputeGraph::get_or_add_value_for_int(const int64_t val) { return add_scalar(val); } +ValueRef ComputeGraph::set_input_tensor( + const ValueRef idx, + vkapi::ScalarType staging_dtype) { + // For texture storage, the buffer size needs to account for the zero + // padding applied by unused texel elements. + size_t buf_numel = get_tensor(idx)->staging_buffer_numel(); + ValueRef staging_idx = add_staging(staging_dtype, buf_numel); + add_staging_to_tensor_node(*this, staging_idx, idx); + inputs_.push_back({idx, staging_idx}); + return staging_idx; +} + ValueRef ComputeGraph::set_input_tensor( const ValueRef idx, const bool use_staging) { if (use_staging) { vkapi::ScalarType dtype = get_tensor(idx)->dtype(); - // For texture storage, the buffer size needs to account for the zero - // padding applied by unused texel elements. - size_t buf_numel = get_tensor(idx)->staging_buffer_numel(); - ValueRef staging_idx = add_staging(dtype, buf_numel); - add_staging_to_tensor_node(*this, staging_idx, idx); - inputs_.push_back({idx, staging_idx}); - return staging_idx; - } - inputs_.push_back({idx, kDummyValueRef}); - return idx; + return set_input_tensor(idx, dtype); + } else { + inputs_.push_back({idx, kDummyValueRef}); + return idx; + } +} + +ValueRef ComputeGraph::set_output_tensor( + const ValueRef idx, + vkapi::ScalarType staging_dtype) { + // For texture storage, the buffer size needs to account for the zero + // padding applied by unused texel elements. + size_t buf_numel = get_tensor(idx)->staging_buffer_numel(); + ValueRef staging_idx = add_staging(staging_dtype, buf_numel); + // We only run this when the tensor is non-empty. When the underlying + // tensor is empty (e.g. padded_numel == 0), we do not allocate a VkImage to + // tensor, we will not be able to bind the node for execution. + if (buf_numel > 0) { + add_tensor_to_staging_node(*this, idx, staging_idx); + } + outputs_.push_back({idx, staging_idx}); + return staging_idx; } ValueRef ComputeGraph::set_output_tensor( @@ -517,21 +633,11 @@ ValueRef ComputeGraph::set_output_tensor( const bool use_staging) { if (use_staging) { vkapi::ScalarType dtype = get_tensor(idx)->dtype(); - // For texture storage, the buffer size needs to account for the zero - // padding applied by unused texel elements. - size_t buf_numel = get_tensor(idx)->staging_buffer_numel(); - ValueRef staging_idx = add_staging(dtype, buf_numel); - // We only run this when the tensor is non-empty. When the underlying - // tensor is empty (e.g. padded_numel == 0), we do not allocate a VkImage to - // tensor, we will not be able to bind the node for execution. - if (buf_numel > 0) { - add_tensor_to_staging_node(*this, idx, staging_idx); - } - outputs_.push_back({idx, staging_idx}); - return staging_idx; + return set_output_tensor(idx, dtype); + } else { + outputs_.push_back({idx, kDummyValueRef}); + return idx; } - outputs_.push_back({idx, kDummyValueRef}); - return idx; } ValueRef ComputeGraph::set_output_value(const ValueRef idx) { @@ -565,7 +671,12 @@ vkapi::BufferBindInfo ComputeGraph::get_or_create_int_param_buffer( } void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) { - get_symint(idx)->set(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); + } } int32_t ComputeGraph::read_symint(const ValueRef idx) { @@ -579,6 +690,17 @@ SharedObject& ComputeGraph::get_shared_object(const int64_t idx) { return shared_objects_.at(idx); } +void ComputeGraph::create_dedicated_allocation_for(const ValueRef idx) { + vTensorPtr tensor = get_tensor(idx); + if (!tensor->memory_is_bound()) { + VmaAllocationCreateInfo alloc_create_info = + context()->adapter_ptr()->vma().gpuonly_resource_create_info(); + tensor->acquire_allocation( + context()->adapter_ptr()->vma().create_allocation( + tensor->get_memory_requirements(), alloc_create_info)); + } +} + void ComputeGraph::update_descriptor_counts( const vkapi::ShaderInfo& shader_info, bool execute) { @@ -700,6 +822,38 @@ utils::uvec3 ComputeGraph::create_local_wg_size(const ValueRef idx) { return create_local_wg_size(create_global_wg_size(idx)); } +void ComputeGraph::bind_tensor_to_descriptor_set( + const ValueRef ref, + vkapi::PipelineBarrier& pipeline_barrier, + const vkapi::MemoryAccessFlags access_type, + vkapi::DescriptorSet& descriptor_set, + const uint32_t idx) { + vTensorPtr tensor = get_tensor(ref); + if (tensor->buffer()) { + vkapi::VulkanBuffer& buffer = tensor->buffer( + pipeline_barrier, vkapi::PipelineStage::COMPUTE, access_type); + descriptor_set.bind(idx, buffer); + } else { + vkapi::VulkanImage& image = tensor->image( + pipeline_barrier, vkapi::PipelineStage::COMPUTE, access_type); + descriptor_set.bind(idx, image); + } +} + +void ComputeGraph::bind_value_to_descriptor_set( + const ValueRef ref, + vkapi::PipelineBarrier& pipeline_barrier, + const vkapi::MemoryAccessFlags access_type, + vkapi::DescriptorSet& descriptor_set, + const uint32_t idx) { + if (val_is_tensor(ref)) { + bind_tensor_to_descriptor_set( + ref, pipeline_barrier, access_type, descriptor_set, idx); + } else if (val_is_staging(ref)) { + descriptor_set.bind(idx, get_staging(ref)->buffer()); + } +} + void ComputeGraph::copy_into_staging( const ValueRef idx, const void* data, @@ -709,6 +863,36 @@ void ComputeGraph::copy_into_staging( staging->copy_from(data, nbytes); } +void ComputeGraph::maybe_cast_and_copy_into_staging( + const ValueRef idx, + const void* data, + const size_t numel, + const vkapi::ScalarType src_data_dtype) { + StagingPtr staging = get_staging(idx); + vkapi::ScalarType staging_dtype = staging->dtype(); + if (src_data_dtype == staging_dtype) { + size_t nbytes = numel * vkapi::element_size(staging_dtype); + staging->copy_from(data, nbytes); + return; + } else { + // Hard-coded type conversion cases + if (src_data_dtype == vkapi::kLong && staging_dtype == vkapi::kInt) { + const int64_t* casted_data = reinterpret_cast(data); + staging->cast_and_copy_from(casted_data, numel); + } else if ( + src_data_dtype == vkapi::kDouble && staging_dtype == vkapi::kFloat) { + const double* casted_data = reinterpret_cast(data); + staging->cast_and_copy_from(casted_data, numel); + } else { + VK_THROW( + "Unsupported type conversion from ", + src_data_dtype, + " to staging dtype ", + staging_dtype); + } + } +} + void ComputeGraph::copy_from_staging( const ValueRef idx, void* data, @@ -718,6 +902,36 @@ void ComputeGraph::copy_from_staging( staging->copy_to(data, nbytes); } +void ComputeGraph::maybe_cast_and_copy_from_staging( + const ValueRef idx, + void* data, + const size_t numel, + const vkapi::ScalarType dst_data_dtype) { + StagingPtr staging = get_staging(idx); + vkapi::ScalarType staging_dtype = staging->dtype(); + if (dst_data_dtype == staging_dtype) { + size_t nbytes = numel * vkapi::element_size(staging_dtype); + staging->copy_to(data, nbytes); + return; + } else { + // Hard-coded type conversion cases + if (dst_data_dtype == vkapi::kLong && staging_dtype == vkapi::kInt) { + int64_t* casted_data = reinterpret_cast(data); + staging->cast_and_copy_to(casted_data, numel); + } else if ( + dst_data_dtype == vkapi::kDouble && staging_dtype == vkapi::kFloat) { + double* casted_data = reinterpret_cast(data); + staging->cast_and_copy_to(casted_data, numel); + } else { + VK_THROW( + "Unsupported type conversion from staging dtype ", + staging_dtype, + " to ", + dst_data_dtype); + } + } +} + void ComputeGraph::prepare() { #define MERGE_FIELD(field) \ static_cast(std::ceil( \ @@ -745,10 +959,34 @@ void ComputeGraph::prepare() { context_->initialize_querypool(); } - for (SharedObject& shared_object : shared_objects_) { - shared_object.allocate(this); - shared_object.bind_users(this); + // Calculate the threshold at which a new command buffer should be created + // during execute() + const size_t total_node_count = execute_nodes_.size(); + size_t init_threshold = config_.execute_initial_threshold_node_count; + size_t count_threshold = config_.execute_threshold_node_count; + + // If max command buffer count is set, we need to adjust the thresholds to + // accommodate execution within the limit, if total command buffers with + // current thresholds would exceed execute_max_cmds + if (config_.execute_max_cmds > 0) { + // Worse case scenario we have one command buffer for nodes before init + // threshold and config_.execute_max_cmds - 1 command buffers for the rest + // of dispatches + + // If command buffers created after offsetting init_threshold would exceed + // max command buffer count, we need to adjust init and count thresholds + const bool slicing_exceeds_max_cmds = (total_node_count - init_threshold) > + count_threshold * (config_.execute_max_cmds - 1); + if (total_node_count > init_threshold && slicing_exceeds_max_cmds) { + // Increase count threshold so remaining nodes after offsetting init fits + // in config_.execute_max_cmds - 1 + count_threshold = static_cast(ceil( + (total_node_count - init_threshold) / + double(config_.execute_max_cmds - 1))); + } } + + execute_threshold_node_count_ = count_threshold; } void ComputeGraph::prepare_pipelines() { @@ -844,6 +1082,18 @@ void ComputeGraph::prepack() { submit_current_cmd_and_wait(/*final_use=*/true); context_->flush(); staging_nbytes_in_cmd_ = 0; + + // Initialize allocations for intermediate tensors + for (SharedObject& shared_object : shared_objects_) { + shared_object.allocate(this); + 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); + } + } } void ComputeGraph::execute() { @@ -852,37 +1102,89 @@ void ComputeGraph::execute() { context_->set_cmd(/*reusable = */ true); context_->cmd_reset_querypool(); + const size_t total_node_count = execute_nodes_.size(); + uint32_t encoded_node_count = 0; for (std::unique_ptr& node : execute_nodes_) { node->encode(this); + encoded_node_count++; + + // Threshold is reached when the node count reached + // execute_initial_threshold_node_count or if its a multiple of + // execute_threshold_node_count. + const bool reached_threshold = + encoded_node_count >= config_.execute_initial_threshold_node_count && + ((encoded_node_count - config_.execute_initial_threshold_node_count) % + execute_threshold_node_count_ == + 0); + + // Create a new command buffer when threashold is reached + // But avoid it if this is the last node, since last cmd buf is submitted + // after the loop + if (reached_threshold && encoded_node_count != total_node_count) { + context_->submit_cmd_to_gpu(VK_NULL_HANDLE, false); + deferred_cmd_list_.emplace_back(std::move(context_->extract_cmd())); + context_->set_cmd(true); + } } + vkapi::VulkanFence fence = context_->fences().get_fence(); + context_->submit_cmd_to_gpu(fence.get_submit_handle(), false); + fence.wait(); + context_->fences().return_fence(fence); deferred_cmd_list_.emplace_back(std::move(context_->extract_cmd())); + } else { + submit_deferred_cmds_and_wait(); } - submit_deferred_cmds_and_wait(); execute_count_++; + + // Clear the set of updated values at the end of inference + updated_values_.clear(); + + // Reset the re-encoding flag at the end of inference + requires_reencode_ = false; +} + +void ComputeGraph::virtual_clone(const ValueRef dst, const ValueRef src) { + get_tensor(dst)->virtual_clone(*get_tensor(src)); +} + +void ComputeGraph::virtual_transpose( + const ValueRef tensor, + const int64_t dim0, + const int64_t dim1) { + get_tensor(tensor)->virtual_transpose(dim0, dim1); } void ComputeGraph::resize_input( const int64_t idx, const std::vector& new_sizes) { IOValueRef io_val = inputs_.at(idx); - get_tensor(io_val.value)->virtual_resize(new_sizes); + virtual_resize(io_val.value, new_sizes); + updated_values_.insert(io_val.staging); } void ComputeGraph::virtual_resize( const ValueRef idx, const std::vector& new_sizes) { - get_tensor(idx)->virtual_resize(new_sizes); + 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); + } } void ComputeGraph::propagate_resize() { for (std::unique_ptr& node : execute_nodes_) { node->trigger_resize(this); } - // Only re-encode on resize if dynamic shapes are expected - if (config_.expect_dynamic_shapes) { + // A command buffer re-encode will be needed if: + // 1. Any push constant data (used for tensor metadata) was updated + // 2. Compute shader dispatch parameters (i.e. compute shader, global and + // local work group sizes) were updated + if (requires_reencode_) { clear_deferred_cmds(); } } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 886e2c5ccea..baa15233a00 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -196,6 +196,12 @@ 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_; + + // Flag to indicate if re-encoding is required + bool requires_reencode_ = false; + protected: size_t values_in_use_ = 0; size_t execute_count_ = 0; @@ -207,6 +213,18 @@ class ComputeGraph final { // current Context's command buffer is submitted now. size_t staging_nbytes_in_cmd_ = 0; + // Represents the nodes to wait before submitting commands. + // If command buffers created with config.execute_threshold_node_count exceeds + // config.execute_max_cmds, then execute_threshold_node_count will be + // increased to fit command buffers within the limit. Otherwise, + // execute_threshold_node_count will be set to + // config.execute_threshold_node_count. + size_t execute_threshold_node_count_ = 0; + + // Whether the underlying GPU support accelerated integer dot product + // extensions + bool can_use_int8_dot_product_ = false; + public: // // Accessors @@ -236,6 +254,9 @@ class ComputeGraph final { return config_; } + // Check if the ComputeGraph has a value at the specified index + bool is_valid_value_idx(const ValueRef idx) const noexcept; + // // Value Extraction // @@ -248,7 +269,16 @@ class ComputeGraph final { return values_.at(idx).is##type_name(); \ } - GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(vTensorPtr, tensor, Tensor) + protected: + inline vTensorPtr get_tensor(const ValueRef idx) { + return vTensorPtr(this, idx); + } + + public: + inline bool val_is_tensor(const ValueRef idx) const { + return values_.at(idx).isTensor(); + } + GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(TensorRefPtr, tref, TensorRef) GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(StagingPtr, staging, Staging) GET_AND_CHECK_VAL_AS_PTR_TYPE_FNS(IntListPtr, int_list, IntList) @@ -278,6 +308,10 @@ class ComputeGraph final { return idx == kDummyValueRef ? true : values_.at(idx).isNone(); } + inline bool val_is_not_none(const ValueRef idx) { + return !val_is_none(idx); + } + inline TypeTag get_val_type(const ValueRef idx) { return values_.at(idx).type(); } @@ -319,6 +353,10 @@ class ComputeGraph final { return values_.at(idx).toConstTensor().numel(); } + inline size_t staging_buffer_numel_of(const ValueRef idx) const { + return values_.at(idx).toConstTensor().staging_buffer_numel(); + } + inline utils::StorageType storage_type_of(const ValueRef idx) const { return values_.at(idx).toConstTensor().storage_type(); } @@ -327,6 +365,10 @@ class ComputeGraph final { return values_.at(idx).toConstTensor().has_buffer_storage(); } + inline bool is_texture_storage(const ValueRef idx) const { + return !is_buffer_storage(idx); + } + /* * Checks that the following is true: * 1. The value at `idx` is a tensor @@ -340,18 +382,40 @@ class ComputeGraph final { * 1. The value at `idx` is a tensor * 2. The tensor at `idx` has texture storage * 3. The texture backed tensor at `idx` has a standard axis mapping - * 4. The texture backed tensor at `idx` is channels packed + * 4. The texture backed tensor at `idx` is width packed */ - bool is_standard_channels_packed_texture_tensor(const ValueRef idx) const; + bool is_contiguous_texture_tensor(const ValueRef idx) const; /* * Checks that the following is true: * 1. The value at `idx` is a tensor * 2. The tensor at `idx` has texture storage * 3. The texture backed tensor at `idx` has a standard axis mapping - * 4. The texture backed tensor at `idx` is width packed + * 4. The texture backed tensor at `idx` is channels packed */ - bool is_standard_width_packed_texture_tensor(const ValueRef idx) const; + bool is_standard_channels_packed_texture_tensor(const ValueRef idx) const; + + /* + * Checks that the value at `idx` is either a 2D tensor, or if the tensor has + * more than 2 dims, the outermost dims have size of 1, i.e. can be squeezed + * to be a 2D tensor. + */ + bool is_2d_matrix(const ValueRef idx) const; + + /* + * Same as the above, but also requires that the tensor is a contiguous + * buffer with a width divisible by 4 or a standard width packed texture. + */ + bool is_vectorizable_contiguous_2d_matrix(const ValueRef idx) const; + + /* + * Checks that the following is true: + * 1. The value at `idx` is a tensor + * 2. The tensor at `idx` is width packed + * 3. The tensor at `idx` has a standard axis mapping or is a contiguous + * buffer + */ + bool is_vectorizable_width_packed_tensor(const ValueRef idx) const; inline bool val_is_view_of(const ValueRef maybe_view, const ValueRef base) const { @@ -381,6 +445,10 @@ class ComputeGraph final { return values_.at(idx).toTensor().sizes_ubo(); } + inline vkapi::BufferBindInfo buffer_meta_ubo(const ValueRef idx) { + return values_.at(idx).toTensor().buffer_meta_ubo(); + } + inline vkapi::BufferBindInfo strides_ubo(const ValueRef idx) { return values_.at(idx).toTensor().strides_ubo(); } @@ -406,31 +474,41 @@ class ComputeGraph final { } inline PushConstantDataInfo sizes_pc_of(const ValueRef idx) const { - return PushConstantDataInfo( + PushConstantDataInfo pc_data = PushConstantDataInfo( values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorSizes); + pc_data.set_value(idx); + return pc_data; } inline PushConstantDataInfo dim_order_pc_of(const ValueRef idx) const { - return PushConstantDataInfo( + PushConstantDataInfo pc_data = PushConstantDataInfo( values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorDimOrder); + pc_data.set_value(idx); + return pc_data; } inline PushConstantDataInfo strides_pc_of(const ValueRef idx) const { - return PushConstantDataInfo( + PushConstantDataInfo pc_data = PushConstantDataInfo( values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorStrides); + pc_data.set_value(idx); + return pc_data; } inline PushConstantDataInfo logical_limits_pc_of(const ValueRef idx) const { - return PushConstantDataInfo( + PushConstantDataInfo pc_data = PushConstantDataInfo( values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorLogicalLimits); + pc_data.set_value(idx); + return pc_data; } inline PushConstantDataInfo numel_pc_of(const ValueRef idx) const { - return PushConstantDataInfo( + PushConstantDataInfo pc_data = PushConstantDataInfo( values_.at(idx).toConstTensor().get_uniform_data(), api::kTensorNumel); + pc_data.set_value(idx); + return pc_data; } // @@ -653,6 +731,16 @@ class ComputeGraph final { const vkapi::ScalarType dtype, const void* const data); + /* + * Add a `TensorRef` value to the graph with the specific properties. A + * `TensorRef` is a reference to a `api::vTensor` whose data is stored in a + * FreeableBuffer. The TensorRef will take ownership of the FreeableBuffer. + */ + ValueRef add_tensorref( + const std::vector& sizes, + const vkapi::ScalarType dtype, + executorch::runtime::FreeableBuffer&& buffer); + /* * Add a staging buffer to the graph. Staging buffers are data buffers that * use memory that is visible to both the CPU and GPU, and therefore is used @@ -683,7 +771,16 @@ class ComputeGraph final { */ ValueRef get_or_add_value_for_int(const int64_t val); + ValueRef set_input_tensor( + const ValueRef idx, + vkapi::ScalarType staging_dtype); + ValueRef set_input_tensor(const ValueRef idx, const bool use_staging = true); + + ValueRef set_output_tensor( + const ValueRef idx, + vkapi::ScalarType staging_dtype); + ValueRef set_output_tensor(const ValueRef idx, const bool use_staging = true); ValueRef set_output_value(const ValueRef idx); @@ -777,6 +874,13 @@ class ComputeGraph final { SharedObject& get_shared_object(const int64_t idx); + /* + * Creates a dedicated memory allocation for a vTensor value, and have the + * tensor acquire the allocation object. If the tensor is already bound to a + * memory allocation, this function will be a no-op. + */ + void create_dedicated_allocation_for(const ValueRef idx); + // // Graph Preparation // @@ -832,14 +936,41 @@ class ComputeGraph final { */ utils::uvec3 create_local_wg_size(const ValueRef idx); + void bind_tensor_to_descriptor_set( + const ValueRef ref, + vkapi::PipelineBarrier& pipeline_barrier, + const vkapi::MemoryAccessFlags accessType, + vkapi::DescriptorSet& descriptor_set, + const uint32_t idx); + + void bind_value_to_descriptor_set( + const ValueRef ref, + vkapi::PipelineBarrier& pipeline_barrier, + const vkapi::MemoryAccessFlags access_type, + vkapi::DescriptorSet& descriptor_set, + const uint32_t idx); + // // Input/Output // void copy_into_staging(const ValueRef idx, const void* data, const size_t numel); + + void maybe_cast_and_copy_into_staging( + const ValueRef idx, + const void* data, + const size_t numel, + const vkapi::ScalarType src_data_dtype); + void copy_from_staging(const ValueRef idx, void* data, const size_t numel); + void maybe_cast_and_copy_from_staging( + const ValueRef idx, + void* data, + const size_t numel, + const vkapi::ScalarType dst_data_dtype); + protected: // Command Buffer Management @@ -890,16 +1021,38 @@ class ComputeGraph final { void execute(); + // + // Tensor View + // + + void virtual_clone(const ValueRef dst, const ValueRef src); + + void virtual_transpose( + const ValueRef tensor, + const int64_t dim0, + const int64_t dim1); + // // Dynamic Shape support // void resize_input(const int64_t idx, const std::vector& new_sizes); + void virtual_resize( const ValueRef idx, const std::vector& new_sizes); + void propagate_resize(); + // Check if a specific ValueRef (or ValueList) was updated, with recursive + // handling + bool was_value_updated(const ValueRef idx) const noexcept; + + // Set the flag to indicate that re-encoding is required + inline void set_requires_reencode() noexcept { + requires_reencode_ = true; + } + // // Miscellaneous Utilities // @@ -912,6 +1065,10 @@ class ComputeGraph final { return execute_count_; } + inline bool can_use_int8_dot_product() const { + return can_use_int8_dot_product_; + } + /* * Check whether the GPU supports 8 bit buffers. */ @@ -939,6 +1096,8 @@ class ComputeGraph final { friend class SymIntPtr; friend struct TmpTensor; + friend struct SharedObject; + friend class BlitNode; }; template diff --git a/backends/vulkan/runtime/graph/GraphConfig.h b/backends/vulkan/runtime/graph/GraphConfig.h index 33c7ae73e62..aa5cd8f8c4e 100644 --- a/backends/vulkan/runtime/graph/GraphConfig.h +++ b/backends/vulkan/runtime/graph/GraphConfig.h @@ -50,6 +50,21 @@ struct GraphConfig final { // by taking more advantage of parallelism between the CPU and GPU. size_t prepack_initial_threshold_nbytes = 0; + // During execute, once this node count is reached, submit the current + // command buffer for execution. This allows the work to be distributed over + // multiple command buffer submissions, which can improve execution + // performance. + size_t execute_threshold_node_count = 0; + // Execute node count used for the first command buffer submission during + // execute. This can be set to be lower than execute_threshold_nbytes to + // submit a command buffer for execution earlier which can improve performance + // by taking more advantage of parallelism between the CPU and GPU. + size_t execute_initial_threshold_node_count = 0; + + // If this number is greater than 0 then, during execute create at most this + // many command buffers. + size_t execute_max_cmds = 0; + vkapi::Adapter* external_adapter; // Generate a default graph config with pre-configured settings diff --git a/backends/vulkan/runtime/graph/Logging.cpp b/backends/vulkan/runtime/graph/Logging.cpp index 7102345773c..081083e3a63 100644 --- a/backends/vulkan/runtime/graph/Logging.cpp +++ b/backends/vulkan/runtime/graph/Logging.cpp @@ -86,7 +86,7 @@ void ComputeGraph::print_readable() { ss << v_tensor.sizes(); std::cout << ss.str(); } else if (val.isTensorRef()) { - const TensorRef tensor_ref = val.toTensorRef(); + const TensorRef& tensor_ref = val.toTensorRef(); std::stringstream ss; ss << tensor_ref.sizes; std::cout << ss.str(); diff --git a/backends/vulkan/runtime/graph/containers/Constant.cpp b/backends/vulkan/runtime/graph/containers/Constant.cpp index cb43295a42a..4dc2cdda8f5 100644 --- a/backends/vulkan/runtime/graph/containers/Constant.cpp +++ b/backends/vulkan/runtime/graph/containers/Constant.cpp @@ -14,7 +14,22 @@ TensorRef::TensorRef( const std::vector& t_sizes, vkapi::ScalarType t_dtype, const void* const t_data) - : sizes{}, dtype{t_dtype}, data{t_data} { + : sizes{}, dtype{t_dtype}, data{t_data}, buffer{} { + size_t ndim = t_sizes.size(); + sizes.resize(ndim); + for (int i = 0; i < ndim; ++i) { + sizes[i] = t_sizes.at(i); + } +} + +TensorRef::TensorRef( + const std::vector& t_sizes, + vkapi::ScalarType t_dtype, + executorch::runtime::FreeableBuffer&& t_buffer) + : sizes{}, + dtype{t_dtype}, + data{t_buffer.data()}, + buffer{std::move(t_buffer)} { size_t ndim = t_sizes.size(); sizes.resize(ndim); for (int i = 0; i < ndim; ++i) { diff --git a/backends/vulkan/runtime/graph/containers/Constant.h b/backends/vulkan/runtime/graph/containers/Constant.h index aaa92360a9e..a18c284a219 100644 --- a/backends/vulkan/runtime/graph/containers/Constant.h +++ b/backends/vulkan/runtime/graph/containers/Constant.h @@ -9,6 +9,7 @@ #pragma once #include +#include namespace vkcompute { @@ -24,14 +25,30 @@ struct TensorRef final { vkapi::ScalarType dtype; const void* data; + // Optional FreeableBuffer for managing memory lifecycle + // This will be empty (default constructed) for the raw pointer constructor + executorch::runtime::FreeableBuffer buffer; + explicit TensorRef( const std::vector& t_sizes, vkapi::ScalarType t_dtype, const void* const t_data); + // Constructor that takes ownership of a FreeableBuffer + explicit TensorRef( + const std::vector& t_sizes, + vkapi::ScalarType t_dtype, + executorch::runtime::FreeableBuffer&& t_buffer); + inline size_t nbytes() const { return utils::multiply_integers(sizes) * vkapi::element_size(dtype); } + + // Manually free the buffer if needed (though it will be freed automatically + // on destruction) + void free_buffer() { + buffer.Free(); + } }; } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/containers/PushConstantData.h b/backends/vulkan/runtime/graph/containers/PushConstantData.h index 39cde4722a7..c86232983ea 100644 --- a/backends/vulkan/runtime/graph/containers/PushConstantData.h +++ b/backends/vulkan/runtime/graph/containers/PushConstantData.h @@ -10,6 +10,8 @@ #include +#include + namespace vkcompute { class ComputeGraph; @@ -33,6 +35,9 @@ class PushConstantDataInfo { }; Payload payload_; + // The value in a compute graph that this push constant data is associated + // with, if any. + ValueRef value_ = kDummyValueRef; public: explicit PushConstantDataInfo( @@ -60,6 +65,18 @@ class PushConstantDataInfo { void* dst, const uint32_t dst_offset, const uint32_t max_dst_size) const; + + inline bool is_tensor_metadata() const noexcept { + return tensorUniformData != nullptr; + } + + inline void set_value(ValueRef value) noexcept { + value_ = value; + } + + inline ValueRef value() const noexcept { + return value_; + } }; } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/BlitNode.cpp b/backends/vulkan/runtime/graph/ops/BlitNode.cpp index 03ee4caa51a..de1ad596069 100644 --- a/backends/vulkan/runtime/graph/ops/BlitNode.cpp +++ b/backends/vulkan/runtime/graph/ops/BlitNode.cpp @@ -26,11 +26,9 @@ BlitNode::BlitNode( } void BlitNode::encode(ComputeGraph* graph) { - auto src_tensor = graph->get_tensor(src_); - auto dst_tensor = graph->get_tensor(dst_); VK_CHECK_COND( - src_tensor->storage_type() != utils::kBuffer && - dst_tensor->storage_type() != utils::kBuffer, + graph->storage_type_of(src_) != utils::kBuffer && + graph->storage_type_of(dst_) != utils::kBuffer, "BlitNode: Only texture backed tensors are supported."); api::Context* const context = graph->context(); @@ -41,18 +39,18 @@ void BlitNode::encode(ComputeGraph* graph) { // Hack to get timing data for non shader op std::string kernel_name("Blit_"); kernel_name.reserve(32); - kernel_name += vkapi::to_string(src_tensor->dtype()); + kernel_name += vkapi::to_string(graph->dtype_of(src_)); kernel_name += "_to_"; - kernel_name += vkapi::to_string(dst_tensor->dtype()); + kernel_name += vkapi::to_string(graph->dtype_of(dst_)); context->report_shader_dispatch_start( kernel_name, utils::uvec3(), utils::WorkgroupSize(), node_id_); context->register_blit( pipeline_barrier, - src_tensor->image( + graph->get_tensor(src_)->image( pipeline_barrier, vkapi::PipelineStage::TRANSFER, vkapi::kRead), - dst_tensor->image( + graph->get_tensor(dst_)->image( pipeline_barrier, vkapi::PipelineStage::TRANSFER, vkapi::kWrite)); context->report_shader_dispatch_end(); diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp index b5644cf3dcd..d1add8227de 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.cpp +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.cpp @@ -44,6 +44,13 @@ void DispatchNode::encode(ComputeGraph* graph) { if (!shader_) { return; } + + // If any global wg size element is 0, then skip encoding this shader + if (global_workgroup_size_[0] == 0 || global_workgroup_size_[1] == 0 || + global_workgroup_size_[2] == 0) { + return; + } + api::Context* const context = graph->context(); vkapi::PipelineBarrier pipeline_barrier{}; @@ -89,4 +96,21 @@ void DispatchNode::write_push_constant_data() { } } +bool DispatchNode::trigger_resize(ComputeGraph* graph) { + const bool any_arg_updated = ExecuteNode::trigger_resize(graph); + + if (any_arg_updated) { + // If this shader uses push constants, and the tensor metadata associated + // with the push constants has changed, then the command buffer needs to be + // re-encoded since push constants cannot be updated. + for (const auto& push_constant : push_constants_) { + if (push_constant.is_tensor_metadata() && + graph->was_value_updated(push_constant.value())) { + graph->set_requires_reencode(); + } + } + } + return any_arg_updated; +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/DispatchNode.h b/backends/vulkan/runtime/graph/ops/DispatchNode.h index b6eb8624c26..89d24a77d6e 100644 --- a/backends/vulkan/runtime/graph/ops/DispatchNode.h +++ b/backends/vulkan/runtime/graph/ops/DispatchNode.h @@ -44,6 +44,8 @@ class DispatchNode : public ExecuteNode { void encode(ComputeGraph* graph) override; + bool trigger_resize(ComputeGraph* graph) override; + protected: vkapi::ShaderInfo shader_; utils::uvec3 global_workgroup_size_; diff --git a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp index b8c0fcbbf79..5a88bba88c9 100644 --- a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp +++ b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.cpp @@ -41,6 +41,12 @@ DynamicDispatchNode::DynamicDispatchNode( pick_global_wg_fn(&graph, shader_, args, resize_args); local_workgroup_size_ = utils::WorkgroupSize(pick_local_wg_fn( &graph, shader_, global_workgroup_size_, args, resize_args)); + + // Calculate dispatch grid similar to Context.cpp register_shader_dispatch + wg_dispatch_grid_ = { + utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), + utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), + utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; } DynamicDispatchNode::DynamicDispatchNode( @@ -57,13 +63,8 @@ DynamicDispatchNode::DynamicDispatchNode( : DispatchNode( graph, shader, - pick_global_wg_fn(&graph, shader, args, resize_args), - pick_local_wg_fn( - &graph, - shader, - pick_global_wg_fn(&graph, shader, args, resize_args), - args, - resize_args), + {1u, 1u, 1u}, + {8u, 8u, 1u}, args, params, push_constants, @@ -72,21 +73,79 @@ DynamicDispatchNode::DynamicDispatchNode( resize_fn), pick_shader_fn_{nullptr}, pick_global_wg_fn_(pick_global_wg_fn), - pick_local_wg_fn_(pick_local_wg_fn) {} + pick_local_wg_fn_(pick_local_wg_fn) { + global_workgroup_size_ = + pick_global_wg_fn(&graph, shader_, args, resize_args); + local_workgroup_size_ = utils::WorkgroupSize(pick_local_wg_fn( + &graph, shader_, global_workgroup_size_, args, resize_args)); + // Calculate the work group grid that will be dispatched + wg_dispatch_grid_ = { + utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), + utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), + utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; +} + +bool DynamicDispatchNode::trigger_resize(ComputeGraph* graph) { + // DispatchNode::trigger_resize() will return true if any of the values + // participating in this operation were updated. + const bool any_arg_updated = DispatchNode::trigger_resize(graph); + // Only re-compute the shader, global workgroup size, and local workgroup size + // if any of the values participating in this operation were updated. + // Otherwise, assume that these will not have changed. + if (!any_arg_updated) { + return false; + } + + // Indicates if the shader dispatch should be changed since the last time the + // command buffer was encoded. + bool dispatch_params_changed = false; -void DynamicDispatchNode::encode(ComputeGraph* graph) { if (pick_shader_fn_) { - shader_ = pick_shader_fn_(graph, args_, resize_args_); + vkapi::ShaderInfo new_shader = pick_shader_fn_(graph, args_, resize_args_); + // Compare shader kernel names as a proxy for shader equality + if (shader_.kernel_name != new_shader.kernel_name) { + shader_ = new_shader; + dispatch_params_changed = true; + } } if (pick_global_wg_fn_) { + // Note that if global workgroup size changes, then the dispatch params + // may not actually be different. The actual value to check is the + // work group grid size that will be dispatched, which is calculated + // below. global_workgroup_size_ = pick_global_wg_fn_(graph, shader_, args_, resize_args_); } if (pick_local_wg_fn_) { - local_workgroup_size_ = utils::WorkgroupSize(pick_local_wg_fn_( - graph, shader_, global_workgroup_size_, args_, resize_args_)); + utils::uvec3 new_local_wg_uvec3 = pick_local_wg_fn_( + graph, shader_, global_workgroup_size_, args_, resize_args_); + utils::WorkgroupSize new_local_wg = + utils::WorkgroupSize(new_local_wg_uvec3); + if (local_workgroup_size_ != new_local_wg) { + local_workgroup_size_ = new_local_wg; + dispatch_params_changed = true; + } } - DispatchNode::encode(graph); + + // Always recompute the new dispatch grid and check if it's different + utils::uvec3 new_wg_dispatch_grid = { + utils::div_up(global_workgroup_size_[0], local_workgroup_size_[0]), + utils::div_up(global_workgroup_size_[1], local_workgroup_size_[1]), + utils::div_up(global_workgroup_size_[2], local_workgroup_size_[2])}; + + // Check if the new dispatch grid is different from the old one + if (wg_dispatch_grid_ != new_wg_dispatch_grid) { + dispatch_params_changed = true; + } + wg_dispatch_grid_ = new_wg_dispatch_grid; + + // If any of the dispatch params have changed, then the command buffer must + // be re-encoded. + if (dispatch_params_changed) { + graph->set_requires_reencode(); + } + + return true; } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h index 005151272c3..d3b82968eb2 100644 --- a/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h +++ b/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h @@ -68,13 +68,15 @@ class DynamicDispatchNode final : public DispatchNode { ~DynamicDispatchNode() override = default; - void encode(ComputeGraph* graph) override; + bool trigger_resize(ComputeGraph* graph) override; protected: const PickShaderFn pick_shader_fn_; const PickGlobalFn pick_global_wg_fn_; const PickLocalFn pick_local_wg_fn_; + utils::uvec3 wg_dispatch_grid_{1u, 1u, 1u}; + public: operator bool() const { return shader_; diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp index 7335ce2703b..953f15e7b4d 100644 --- a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp +++ b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include namespace vkcompute { @@ -18,4 +19,33 @@ ExecuteNode::ExecuteNode( resize_args_(resize_args), args_(args), name_(name) {} + +bool ExecuteNode::trigger_resize(ComputeGraph* graph) { + const bool any_arg_updated = was_any_arg_updated(graph); + if (resize_fn_ && any_arg_updated) { + resize_fn_(graph, args_, resize_args_); + } + return any_arg_updated; +} + +bool ExecuteNode::was_any_arg_updated(const ComputeGraph* const graph) const { + // Check all ValueRefs in ArgGroups + for (const auto& arg_group : args_) { + for (const auto& value_ref : arg_group.refs) { + if (graph->was_value_updated(value_ref)) { + return true; + } + } + } + + // Check all ValueRefs in resize_args + for (const auto& value_ref : resize_args_) { + if (graph->was_value_updated(value_ref)) { + return true; + } + } + + return false; +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.h b/backends/vulkan/runtime/graph/ops/ExecuteNode.h index 6a815b246ef..323036cef90 100644 --- a/backends/vulkan/runtime/graph/ops/ExecuteNode.h +++ b/backends/vulkan/runtime/graph/ops/ExecuteNode.h @@ -43,7 +43,7 @@ class ExecuteNode { friend class ComputeGraph; public: - using ResizeFunction = const std::function&, const std::vector&)>; @@ -69,11 +69,9 @@ class ExecuteNode { (void)graph; } - virtual inline void trigger_resize(ComputeGraph* graph) { - if (resize_fn_ != nullptr) { - resize_fn_(graph, args_, resize_args_); - } - } + virtual bool trigger_resize(ComputeGraph* graph); + + bool was_any_arg_updated(const ComputeGraph* const graph) const; inline void set_node_id(uint32_t node_id) { node_id_ = node_id; diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp index 05729172420..62e1dc86f43 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.cpp +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.cpp @@ -18,9 +18,8 @@ namespace vkcompute { vkapi::ShaderInfo get_noop_shader(ComputeGraph& graph, const ValueRef packed) { std::string noop_shader_name("no_op"); - vTensorPtr t_packed = graph.get_tensor(packed); - add_dtype_suffix(noop_shader_name, *t_packed); - add_storage_type_suffix(noop_shader_name, *t_packed); + add_dtype_suffix(noop_shader_name, graph.dtype_of(packed)); + add_storage_type_suffix(noop_shader_name, graph.storage_type_of(packed)); return VK_KERNEL_FROM_STR(noop_shader_name); } @@ -48,13 +47,13 @@ PrepackNode::PrepackNode( } api::StagingBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) { - vTensorPtr packed = graph->get_tensor(packed_); - - // If no TensorRef is provided, create a staging buffer of zeros according to - // the vkapi::vTensor metadata. + // If no TensorRef is provided, create a staging buffer of zeros based on the + // Tensor metadata. if (graph->val_is_none(tref_)) { - size_t numel = utils::multiply_integers(packed->sizes()); - api::StagingBuffer staging(graph->context(), packed->dtype(), numel); + const std::vector packed_sizes = graph->sizes_of(packed_); + size_t numel = utils::multiply_integers(packed_sizes); + api::StagingBuffer staging( + graph->context(), graph->dtype_of(packed_), numel); staging.set_staging_zeros(); return staging; } @@ -65,6 +64,9 @@ api::StagingBuffer PrepackNode::create_staging_buffer(ComputeGraph* graph) { graph->update_staging_nbytes_in_cmd(staging.buffer().mem_size_as_size_t()); size_t nbytes = numel * vkapi::element_size(tref->dtype); staging.copy_from(tref->data, nbytes); + // Once the staging buffer is copied, if the TensorRef owns a FreeableBuffer, + // it can be freed. + tref->free_buffer(); return staging; } @@ -80,7 +82,6 @@ void PrepackNode::encode(ComputeGraph* graph) { context->check_device_capabilities(shader_); - vTensorPtr packed = graph->get_tensor(packed_); api::StagingBuffer staging = create_staging_buffer(graph); std::unique_lock cmd_lock = context->dispatch_lock(); @@ -96,13 +97,17 @@ void PrepackNode::encode(ComputeGraph* graph) { } { + // If the vTensor is not yet bound to a memory allocation, create a new one + // and aquire it. + graph->create_dedicated_allocation_for(packed_); + vkapi::PipelineBarrier pipeline_barrier{}; vkapi::DescriptorSet descriptor_set = context->get_descriptor_set( shader_, local_workgroup_size_, spec_vars_, push_constants_offset); uint32_t idx = 0; - bind_tensor_to_descriptor_set( - *packed, + graph->bind_tensor_to_descriptor_set( + packed_, pipeline_barrier, vkapi::MemoryAccessType::WRITE, descriptor_set, @@ -128,8 +133,8 @@ void PrepackNode::encode(ComputeGraph* graph) { vkapi::DescriptorSet descriptor_set = context->get_descriptor_set( noop_shader_, utils::WorkgroupSize(1, 1, 1)); - bind_tensor_to_descriptor_set( - *packed, + graph->bind_tensor_to_descriptor_set( + packed_, pipeline_barrier, vkapi::MemoryAccessType::READ, descriptor_set, diff --git a/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl b/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl index a4ed494fe6d..bd210e210ce 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/addmm_naive_texture3d.glsl @@ -136,6 +136,9 @@ vec4 matmul_naive_k_dim_packed_row_dim_packed(const ivec3 out_lpos) { const vec4 mat1_tex = texelFetch(mat1_tensor, mat1_pos, 0); for (int r = 0; r < 4; ++r) { + if (4 * i + r >= mat2_sizes.y) { + continue; + } // On-demand construction of mat2_pos appears to provide the lowest // latency. Surprisingly, this doesn't translate to mat1_pos. ivec3 mat2_pos = ivec3(0); diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl b/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl index f2a9e9cfdac..6f2a93667ea 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl @@ -34,6 +34,8 @@ $if IS_COMPARISON_OP: layout(std430) buffer; +#include "indexing.glslh" + $if IS_COMPARISON_OP: ${layout_declare_tensor(B, "w", "t_out", "uint8", STORAGE)} $else: @@ -43,13 +45,11 @@ ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_tensor(B, "r", "t_other", DTYPE, STORAGE)} $if STORAGE == "buffer": + ${layout_declare_ubo(B, "BufferMetadata", "outp")} + ${layout_declare_ubo(B, "BufferMetadata", "inp")} + ${layout_declare_ubo(B, "BufferMetadata", "other")} + layout(push_constant) uniform restrict Block { - ivec4 in_sizes; - ivec4 other_sizes; - ivec4 out_strides; - ivec4 in_strides; - ivec4 other_strides; - int out_numel; float alpha; }; $else: @@ -83,25 +83,30 @@ $else: #ifdef USING_BUFFER void main() { - const int out_bufi = ivec3(gl_GlobalInvocationID).x; - if (out_bufi >= out_numel) { + const uint out_bufi = gl_GlobalInvocationID.x; + if (out_bufi >= numel(outp)) { return; } // Simple case; no broadcasting - if (in_sizes == other_sizes) { + if (are_equal(inp, other)) { t_out[out_bufi] = T(op(t_in[out_bufi], t_other[out_bufi], T(alpha))); return; } - const ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, out_dim_order); - const ivec4 in_tidx = min(out_tidx, in_sizes - 1); - const ivec4 other_tidx = min(out_tidx, other_sizes - 1); + TensorIndex outp_tidx; + linear_idx_to_tensor_idx(outp, out_bufi, outp_tidx); + + TensorIndex inp_tidx = outp_tidx; + clamp_tensor_idx(inp, inp_tidx); + + TensorIndex other_tidx = outp_tidx; + clamp_tensor_idx(other, other_tidx); - const int in_bufi = tidx_to_bufi(in_tidx, in_strides); - const int other_bufi = tidx_to_bufi(other_tidx, other_strides); + uint inp_bufi = tensor_idx_to_linear_idx(inp, inp_tidx); + uint other_bufi = tensor_idx_to_linear_idx(other, other_tidx); - t_out[out_bufi] = T(op(t_in[in_bufi], t_other[other_bufi], T(alpha))); + t_out[out_bufi] = T(op(t_in[inp_bufi], t_other[other_bufi], T(alpha))); } #else // USING_TEXTURE diff --git a/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.glsl b/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.glsl index 423c4df2679..f61081d33b7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.glsl @@ -3,41 +3,36 @@ #define PRECISION ${PRECISION} #define T ${buffer_scalar_type(DTYPE)} - -#include "indexing_utils.h" +#define DST_T ${buffer_scalar_type(BUF_DTYPE)} ${define_required_extensions(DTYPE)} +${define_required_extensions(BUF_DTYPE)} layout(std430) buffer; -${layout_declare_tensor(0, "w", "nchw_buf", DTYPE, STORAGE)} -${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "nchw_buf", BUF_DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_inp", DTYPE, STORAGE)} -$if USE_PUSH_CONST: - layout(push_constant) uniform restrict Block { - ivec4 in_sizes; - ivec4 in_strides; - int numel; - }; -$else: - ${layout_declare_ubo(2, "ivec4", "in_sizes")} - ${layout_declare_ubo(3, "ivec4", "in_strides")} - ${layout_declare_ubo(4, "int", "numel")} +${layout_declare_ubo(B, "BufferMetadata", "inp")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; // This constant is unused in this shader but is kept so that the signature is // consistent with image_to_nchw. -layout(constant_id = 3) const int UNUSED_packed_dim = W_DIM; +${layout_declare_spec_const(C, "int", "unused", "0")} void main() { - int nchwi = int(gl_GlobalInvocationID.x); - if (nchwi >= numel) { + uint inp_bufi = gl_GlobalInvocationID.x; + if (inp_bufi>= numel(inp)) { return; } - ivec4 in_tidx = nchwi_to_tidx(nchwi, in_sizes); - const int in_bufi = tidx_to_bufi(in_tidx, in_strides); + TensorIndex inp_tidx; + linear_idx_to_tensor_idx(inp, inp_bufi, inp_tidx); + + uint nchwi = tensor_idx_to_contiguous_idx(inp, inp_tidx); - nchw_buf[nchwi] = t_in[in_bufi]; + nchw_buf[nchwi] = DST_T(t_inp[inp_bufi]); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.yaml b/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.yaml index 679e686dc2f..1ee7d2db8c1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/buffer_to_nchw.yaml @@ -7,17 +7,19 @@ buffer_to_nchw: parameter_names_with_default_values: DTYPE: float + BUF_DTYPE: float STORAGE: buffer USE_PUSH_CONST: True generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: double - - VALUE: int8 - - VALUE: uint8 - - VALUE: int32 + combination: + parameter_names: [DTYPE, BUF_DTYPE] + combos: + - parameter_values: [half, half] + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [double, double] + - parameter_values: [int8, int8] + - parameter_values: [uint8, uint8] + - parameter_values: [int32, int32] shader_variants: - NAME: buffer_to_nchw - - NAME: buffer_to_nchw_no_pc - USE_PUSH_CONST: False diff --git a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.glsl index 99a64c3589e..7e21bcf0eba 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.glsl @@ -11,18 +11,22 @@ #define PRECISION ${PRECISION} #define IN_T ${buffer_scalar_type(IN_DTYPE)} +#define SCALE_OUT_T ${buffer_scalar_type(SCALE_OUT_DTYPE)} +#define ZP_OUT_T ${buffer_scalar_type(ZP_OUT_DTYPE)} #define ${MODE} ${define_active_storage_type("buffer")} ${define_required_extensions(IN_DTYPE)} +${define_required_extensions(SCALE_OUT_DTYPE)} +${define_required_extensions(ZP_OUT_DTYPE)} #extension GL_EXT_control_flow_attributes : require layout(std430) buffer; -${layout_declare_tensor(B, "w", "t_scale", "float", "buffer")} -${layout_declare_tensor(B, "w", "t_zero_point", "int", "buffer")} +${layout_declare_tensor(B, "w", "t_scale", SCALE_OUT_DTYPE, "buffer")} +${layout_declare_tensor(B, "w", "t_zero_point", ZP_OUT_DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "buffer")} $if MODE == "per_tensor": @@ -254,8 +258,8 @@ void choose_qparams_per_tensor() { // Use default values: mapping_type=0 (ASYMMETRIC), eps from push constant calc_scale_zp(global_min, global_max, quant_min, quant_max, 0, eps, scale_val, zero_point_val); - t_scale[0] = scale_val; - t_zero_point[0] = zero_point_val; + t_scale[0] = SCALE_OUT_T(scale_val); + t_zero_point[0] = ZP_OUT_T(zero_point_val); } } @@ -306,8 +310,8 @@ void choose_qparams_per_token() { calc_scale_zp(lo, hi, quant_min, quant_max, 0, 1e-5, scale_val, zero_point_val); // Write results - t_scale[token_id] = scale_val; - t_zero_point[token_id] = zero_point_val; + t_scale[token_id] = SCALE_OUT_T(scale_val); + t_zero_point[token_id] = ZP_OUT_T(zero_point_val); } } @@ -380,12 +384,12 @@ void choose_qparams_block_wise() { hi = 0.0; } - float scale; - int zp; - calc_scale_zp(lo, hi, quant_min, quant_max, mapping_type, eps, scale, zp); + float scale_val; + int zero_point_val; + calc_scale_zp(lo, hi, quant_min, quant_max, mapping_type, eps, scale_val, zero_point_val); - t_zero_point[block_id] = zp; - t_scale[block_id] = scale; + t_scale[block_id] = SCALE_OUT_T(scale_val); + t_zero_point[block_id] = ZP_OUT_T(zero_point_val); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.yaml index ee900750e16..8459b043baa 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_buffer.yaml @@ -1,10 +1,18 @@ choose_qparams_buffer: parameter_names_with_default_values: IN_DTYPE: float + SCALE_OUT_DTYPE: float + ZP_OUT_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: - VALUE: float + SCALE_OUT_DTYPE: + - VALUE: float + ZP_OUT_DTYPE: + - VALUE: int32 + - VALUE: int8 + - VALUE: float shader_variants: - NAME: choose_qparams_tensor_buffer MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.glsl b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.glsl new file mode 100644 index 00000000000..7234b50a3f5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.glsl @@ -0,0 +1,202 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, STORAGE)} +#define T ${texel_load_component_type(DTYPE, STORAGE)} + +#define NUM_OUTPUTS_PER_WG 1 +#define NUM_WORKERS_PER_OUTPUT 64 + +// Maximum total threads in a work group +#define MAX_THREADS 256 + +${define_active_storage_type(STORAGE)} + +${define_required_extensions(DTYPE)} +${define_required_extensions("int8")} + +#extension GL_EXT_control_flow_attributes : require + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "w", "t_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_input", DTYPE, STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +layout(push_constant) uniform PushConstants { + int quant_min; + int quant_max; +}; + +// Shared memory for cooperative min/max finding +shared T shared_min[NUM_OUTPUTS_PER_WG][NUM_WORKERS_PER_OUTPUT]; +shared T shared_max[NUM_OUTPUTS_PER_WG][NUM_WORKERS_PER_OUTPUT]; + +const float SMALL_SCALE_THRESHOLD = 6.1e-5; + +void calculate_scale_and_zero_point( + float min_val, + float max_val, + int qmin, + int qmax, + out float scale, + out int zero_point) { + + // Extend the [min, max] interval to ensure it contains 0 + min_val = min(min_val, 0.0); + max_val = max(max_val, 0.0); + + // Calculate scale + scale = (max_val - min_val) / float(qmax - qmin); + + // Handle special cases for scale + if (scale == 0.0 || isinf(1.0 / scale)) { + scale = 0.1; + } + + // Cut off small scale + if (scale < SMALL_SCALE_THRESHOLD) { + float org_scale = scale; + scale = SMALL_SCALE_THRESHOLD; + // Adjust the min and max based on the new scale + if (min_val == 0.0) { + max_val = SMALL_SCALE_THRESHOLD * float(qmax - qmin); + } else if (max_val == 0.0) { + min_val = -SMALL_SCALE_THRESHOLD * float(qmax - qmin); + } else { + float amplifier = SMALL_SCALE_THRESHOLD / org_scale; + min_val *= amplifier; + max_val *= amplifier; + } + } + + // Zero-point computation + float zero_point_from_min = float(qmin) - min_val / scale; + float zero_point_from_max = float(qmax) - max_val / scale; + float zero_point_from_min_error = abs(float(qmin)) - abs(min_val / scale); + float zero_point_from_max_error = abs(float(qmax)) - abs(max_val / scale); + + float initial_zero_point = zero_point_from_min_error < zero_point_from_max_error + ? zero_point_from_min + : zero_point_from_max; + + // Nudge zero point to be an integer + int nudged_zero_point; + if (initial_zero_point < float(qmin)) { + nudged_zero_point = qmin; + } else if (initial_zero_point > float(qmax)) { + nudged_zero_point = qmax; + } else { + nudged_zero_point = int(round(initial_zero_point)); + } + + zero_point = nudged_zero_point; +} + +VEC4_T load_input_x4(const int x4, const int y, const int ntexels_x) { +#ifdef USING_BUFFER + return t_input[(y * ntexels_x) + x4]; +#else + return texelFetch(t_input, ivec3(x4, y, 0), 0); +#endif +} + +void find_min_max_for_row(const int output_y) { + const int worker_id = int(gl_LocalInvocationID.x); + const int output_id = int(gl_LocalInvocationID.y); + + if (output_y >= input_sizes.y) { + return; + } + + // Input is 2D tensor (height x width), width-packed + // Each channel corresponds to a row in the tensor + const int X4 = div_4(input_sizes.x); + + // Initialize thread-local min/max + T local_min = T(1e30); + T local_max = T(-1e30); + + // Each thread processes elements along their assigned output_id with stride + // NUM_WORKERS_PER_OUTPUT + for (int x4 = worker_id; x4 < X4; x4 += NUM_WORKERS_PER_OUTPUT) { + VEC4_T in_texel = load_input_x4(x4, output_y, X4); + for (int i = 0; i < 4; i++) { + local_min = min(local_min, in_texel[i]); + local_max = max(local_max, in_texel[i]); + } + } + + // Store thread-local results in shared memory + shared_min[output_id][worker_id] = local_min; + shared_max[output_id][worker_id] = local_max; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_OUTPUT / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_min[output_id][worker_id] = min( + shared_min[output_id][worker_id], + shared_min[output_id][worker_id + i]); + shared_max[output_id][worker_id] = max( + shared_max[output_id][worker_id], + shared_max[output_id][worker_id + i]); + } + memoryBarrierShared(); + barrier(); + } +} + +void main() { + const int worker_id = int(gl_LocalInvocationID.x); + const int output_id = int(gl_LocalInvocationID.y); + + const int output_y4 = int(gl_GlobalInvocationID.y); + const int output_y = mul_4(output_y4); + + + VEC4_T scales_out = VEC4_T(0.0); + ivec4 zps_out = ivec4(0); + + int limit = min(input_sizes.y - output_y, 4); + for (int i = 0; i < limit; i++) { + find_min_max_for_row(output_y + i); + + // Only the first thread in the work group will compute the result + if (worker_id == 0) { + float local_min = shared_min[output_id][0]; + float local_max = shared_max[output_id][0]; + + float scale; + int zero_point; + + calculate_scale_and_zero_point( + local_min, local_max, quant_min, quant_max, scale, zero_point); + + scales_out[i] = T(scale); + zps_out[i] = zero_point; + } + } + + if (worker_id == 0) { + imageStore(t_scales, ivec3(output_y4, 0, 0), scales_out); + imageStore(t_zps, ivec3(output_y4, 0, 0), zps_out); + } + +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.yaml b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.yaml similarity index 83% rename from backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.yaml rename to backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.yaml index e2a96234465..5dbf3d7adaa 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_per_row.yaml @@ -4,16 +4,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -kv_cache_update: +choose_qparams_per_row: parameter_names_with_default_values: DTYPE: float - STORAGE: buffer + STORAGE: texture3d generate_variant_forall: STORAGE: - - VALUE: buffer - VALUE: texture3d + - VALUE: buffer DTYPE: - - VALUE: half - VALUE: float + - VALUE: half shader_variants: - - NAME: kv_cache_update + - NAME: choose_qparams_per_row diff --git a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.glsl index 62ea7099f8c..a17a3ae41dd 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.glsl @@ -12,22 +12,26 @@ #define IN_T ${buffer_scalar_type(IN_DTYPE)} #define FVEC4_T ${texel_load_type(IN_DTYPE, "texture3d")} +#define SCALE_OUT_T ${buffer_scalar_type(SCALE_OUT_DTYPE)} +#define ZP_OUT_T ${buffer_scalar_type(ZP_OUT_DTYPE)} #define ${MODE} ${define_active_storage_type("texture3d")} ${define_required_extensions(IN_DTYPE)} +${define_required_extensions(SCALE_OUT_DTYPE)} +${define_required_extensions(ZP_OUT_DTYPE)} #extension GL_EXT_control_flow_attributes : require layout(std430) buffer; $if MODE != "block_wise": - ${layout_declare_tensor(B, "w", "t_scale", "float", "texture3d")} - ${layout_declare_tensor(B, "w", "t_zero_point", "int", "texture3d")} + ${layout_declare_tensor(B, "w", "t_scale", SCALE_OUT_DTYPE, "texture3d")} + ${layout_declare_tensor(B, "w", "t_zero_point", ZP_OUT_DTYPE, "texture3d")} $else: - ${layout_declare_tensor(B, "w", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "w", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "w", "t_scale", SCALE_OUT_DTYPE, "buffer")} + ${layout_declare_tensor(B, "w", "t_zero_point", ZP_OUT_DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "texture3d")} @@ -273,8 +277,8 @@ void choose_qparams_per_tensor() { int zero_point_val; calc_scale_zp(global_min, global_max, quant_min, quant_max, 0, eps, scale_val, zero_point_val); - write_texel(t_scale, ivec3(0, 0, 0), vec4(scale_val, 0.0, 0.0, 0.0)); - write_texel(t_zero_point, ivec3(0, 0, 0), ivec4(zero_point_val, 0, 0, 0)); + write_texel(t_scale, ivec3(0, 0, 0), vec4(SCALE_OUT_T(scale_val), 0.0, 0.0, 0.0)); + write_texel(t_zero_point, ivec3(0, 0, 0), ivec4(ZP_OUT_T(zero_point_val), 0, 0, 0)); } } @@ -419,8 +423,8 @@ void choose_qparams_per_token() { uint out_x = out_remainder % uint(t_scale_limits.x); ivec3 out_pos = ivec3(int(out_x), int(out_y), int(out_z)); - write_texel(t_scale, out_pos, vec4(scale_val, 0.0, 0.0, 0.0)); - write_texel(t_zero_point, out_pos, ivec4(zero_point_val, 0, 0, 0)); + write_texel(t_scale, out_pos, vec4(SCALE_OUT_T(scale_val), 0.0, 0.0, 0.0)); + write_texel(t_zero_point, out_pos, ivec4(ZP_OUT_T(zero_point_val), 0, 0, 0)); } // Synchronize before processing next token @@ -517,8 +521,8 @@ void choose_qparams_block_wise() { calc_scale_zp(vmin, vmax, quant_min, quant_max, mapping_type, eps, scale, zp); // Write the scalar values directly to buffer using linear index - t_scale[blkIdx] = scale; - t_zero_point[blkIdx] = zp; + t_scale[blkIdx] = SCALE_OUT_T(scale); + t_zero_point[blkIdx] = ZP_OUT_T(zp); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.yaml index a097ce0da48..12228822d4b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/choose_qparams_texture.yaml @@ -1,10 +1,18 @@ choose_qparams_texture: parameter_names_with_default_values: IN_DTYPE: float + SCALE_OUT_DTYPE: float + ZP_OUT_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: - VALUE: float + SCALE_OUT_DTYPE: + - VALUE: float + ZP_OUT_DTYPE: + - VALUE: int32 + - VALUE: int8 + - VALUE: float shader_variants: - NAME: choose_qparams_tensor_texture3d MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/col2im.glsl b/backends/vulkan/runtime/graph/ops/glsl/col2im.glsl new file mode 100644 index 00000000000..c105ef18719 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/col2im.glsl @@ -0,0 +1,94 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, OUTPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, OUTPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#define TILE_M4 1 +#define TILE_N4 1 +#define TILE_K4 1 + +#define TILE_M 4 +#define TILE_N 4 +#define TILE_K 4 + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "conv2d_common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} + +// Sizes of the convolution output image +${layout_declare_ubo(B, "ivec4", "output_sizes")} +// Sizes of the convolution input image +${layout_declare_ubo(B, "ivec4", "input_sizes")} +// Sizes of the im2col matrix of the convolution output +${layout_declare_ubo(B, "ivec4", "matrix_sizes")} + +${layout_declare_ubo(B, "Conv2DParams", "conv2d_params")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "conv2d_fp_im2col_block_store.glslh" + +#ifdef INPUT_BUFFER + +void load_matrix_tile( + out FPOutTile tile, + const int n4, + const int m_start, + const int N4) { + [[unroll]] for (int m = 0; m < TILE_M; m++) { + tile.data[m][0] = t_input[(m_start + m) * N4 + n4]; + } +} + +#else // INPUT_TEXTURE + +void load_matrix_tile( + out FPOutTile tile, + const int n4, + const int m_start, + const int N4) { + [[unroll]] for (int m = 0; m < TILE_M; m++) { + tile.data[m][0] = texelFetch( + t_input, ivec3(n4, m_start + m, 0), 0); + } +} + +#endif // INPUT_BUFFER + +void main() { + // Each thread loads and writes a 4 wide x 4 high block of the matrix + const int n4 = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + const int n = mul_4(n4); + const int m = mul_4(m4); + + if (n >= matrix_sizes.x || m >= matrix_sizes.y) { + return; + } + + FPOutTile tile; + + const int N4 = div_4(matrix_sizes.x); + load_matrix_tile(tile, n4, m, N4); + write_im2col_tile_as_image(tile, n4, m); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/col2im.yaml b/backends/vulkan/runtime/graph/ops/glsl/col2im.yaml new file mode 100644 index 00000000000..b6d0972271a --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/col2im.yaml @@ -0,0 +1,19 @@ +# 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. + +col2im: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: texture3d + INPUT_STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: col2im_texture3d_buffer + - NAME: col2im_texture3d_texture3d + INPUT_STORAGE: texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/common.glslh b/backends/vulkan/runtime/graph/ops/glsl/common.glslh new file mode 100644 index 00000000000..732b7006c2c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/common.glslh @@ -0,0 +1,51 @@ +/* + * 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. + */ + +#ifndef COMMON_GLSLH +#define COMMON_GLSLH + +#define mul_2(x) ((x) << 1) +#define mul_4(x) ((x) << 2) +#define mul_8(x) ((x) << 3) + +#define div_2(x) ((x) >> 1) +#define div_4(x) ((x) >> 2) +#define div_8(x) ((x) >> 3) + +#define div_up_2(x) (((x) + 1) >> 1) +#define div_up_4(x) (((x) + 3) >> 2) +#define div_up_8(x) (((x) + 7) >> 3) + +#define align_up_2(x) ((x + 1) & -2) +#define align_up_4(x) ((x + 3) & -4) +#define align_up_8(x) ((x + 7) & -8) + +#define mod_2(x) ((x) & 1) +#define mod_4(x) ((x) & 3) +#define mod_8(x) ((x) & 7) + +struct TensorIndex4D { + ivec4 data; +}; + +#ifdef DEBUG_MODE + +#extension GL_EXT_debug_printf : require + +void printTensorIndex4D(const TensorIndex4D index) { + debugPrintfEXT( + "tensor_idx: %d, %d, %d, %d\\n", + index.data.x, + index.data.y, + index.data.z, + index.data.w); +} + +#endif // DEBUG_MODE + +#endif // COMMON_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/concat_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/concat_buffer.glsl index 895cecb413a..e34ecaf8309 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/concat_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/concat_buffer.glsl @@ -20,10 +20,12 @@ layout(std430) buffer; #include "indexing_utils.h" -${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} +${layout_declare_tensor(B, "rw", "t_out", DTYPE, "buffer")} $for i in range(NUM_INPUTS): - ${layout_declare_tensor(B, "r", "t_in" + str(i + 1), DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_inp" + str(i), DTYPE, "buffer")} + +${layout_declare_tensor(B, "r", "t_concat_offset", "int", "buffer")} ${layout_declare_ubo(B, "int", "concat_dim")} @@ -31,8 +33,8 @@ ${layout_declare_ubo(B, "ivec4", "out_sizes")} ${layout_declare_ubo(B, "ivec4", "out_strides")} $for i in range(NUM_INPUTS): - ${layout_declare_ubo(B, "ivec4", "in" + str(i+1) + "_sizes")} - ${layout_declare_ubo(B, "ivec4", "in" + str(i+1) + "_strides")} + ${layout_declare_ubo(B, "ivec4", "inp" + str(i) + "_sizes")} + ${layout_declare_ubo(B, "ivec4", "inp" + str(i) + "_strides")} ${layout_declare_ubo(B, "int", "out_numel")} @@ -42,28 +44,53 @@ const lowp ivec4 out_dim_order = unhash_dim_order(out_layout); layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +#define NUM_INPUTS ${NUM_INPUTS} + +#include "concat_utils.glslh" + +/* + * This shader template concatenates up to NUM_INPUT input tensors to the + * output tensor along the concat_dim. Elements from the input tensor will + * be inserted along the output's concat_dim starting at concat_offset. + */ void main() { - const int out_bufi = ivec3(gl_GlobalInvocationID).x; - if (out_bufi >= out_numel) { + const int tid = ivec3(gl_GlobalInvocationID).x; + + // The 1-3 input tensors are interpreted as one concatenated tensor ("volume") + // along the concat_dim for the purposes of tensor indexing. Each thread is + // responsible for reading one item from this volume and writing it to the + // appropriate output location. + ivec4 inp_volume_sizes = out_sizes; + inp_volume_sizes[concat_dim] = total_concat_dim_numel(); + + // Account for 0 size input tensors + if (any(lessThanEqual(inp_volume_sizes, ivec4(0)))) { + return; + } + + ivec4 inp_volume_tidx = nchwi_to_tidx(tid, inp_volume_sizes); + + // bounds check + if (any(greaterThanEqual(inp_volume_tidx, inp_volume_sizes))) { return; } - // Convert buffer linear index to 4-D tensor index for output - const ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, out_dim_order); + int concat_offset = t_concat_offset[0]; + + ivec4 out_tidx = inp_volume_tidx; + out_tidx[concat_dim] += concat_offset; - // Determine which input tensor to read from - ivec4 in_tidx = out_tidx; + const uint out_bufi = tidx_to_bufi(out_tidx, out_strides); + // Go through the list of input tensors, and find which input this output + // element should be read from. $for i in range(NUM_INPUTS): - // Check if the index at the concat dim is within bounds of the input tensor - // If so, read from that input tensor and write to output - if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) { - int in_bufi = tidx_to_bufi(in_tidx, in${i+1}_strides); - t_out[out_bufi] = t_in${i+1}[in_bufi]; + if (inp_volume_tidx[concat_dim] < inp${i}_sizes[concat_dim]) { + int inp_bufi = tidx_to_bufi(inp_volume_tidx, inp${i}_strides); + t_out[out_bufi] = t_inp${i}[inp_bufi]; return; } - // otherwise, decrement the index at the concat dim else { - in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim]; + inp_volume_tidx[concat_dim] -= inp${i}_sizes[concat_dim]; } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/concat_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/concat_texture.glsl index dac6266bf67..afab0c524d6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/concat_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/concat_texture.glsl @@ -19,16 +19,18 @@ layout(std430) buffer; #include "indexing_utils.h" -${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} +${layout_declare_tensor(B, "rw", "t_out", DTYPE, "texture3d")} $for i in range(NUM_INPUTS): - ${layout_declare_tensor(B, "r", "t_in" + str(i + 1), DTYPE, "texture3d")} + ${layout_declare_tensor(B, "r", "t_inp" + str(i), DTYPE, "texture3d")} + +${layout_declare_tensor(B, "r", "t_concat_offset", "int", "buffer")} ${layout_declare_ubo(B, "int", "concat_dim")} $in_metadata = "" $for i in range(NUM_INPUTS): - $in_metadata += "ivec4 in" + str(i + 1) + "_sizes;\n" + $in_metadata += "ivec4 inp" + str(i) + "_sizes;\n" layout(push_constant) uniform restrict Block { ivec4 out_sizes; @@ -40,90 +42,135 @@ const lowp ivec4 out_axis_map = unhash_axis_map(out_layout); const lowp int out_packed_dim = unhash_packed_dim(out_layout); $for i in range(NUM_INPUTS): - ${layout_declare_spec_const(C, "int", "in" + str(i+1) + "_layout", "DEFAULT_LAYOUT")} - const lowp ivec4 in${i+1}_axis_map = unhash_axis_map(in${i+1}_layout); - const lowp int in${i+1}_packed_dim = unhash_packed_dim(in${i+1}_layout); + ${layout_declare_spec_const(C, "int", "inp" + str(i) + "_layout", "DEFAULT_LAYOUT")} + const lowp ivec4 inp${i}_axis_map = unhash_axis_map(inp${i}_layout); + const lowp int inp${i}_packed_dim = unhash_packed_dim(inp${i}_layout); layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -// Check if we can use the fast path (no texel merging required) -bool can_use_fast_path() { - // Fast path is possible when: - // 1. The concat dimension is not the packed dimension, or - // 2. The concat dimension is the packed dimension but both input tensors have dimensions - // that are multiples of 4 along the packed dimension - if (concat_dim != out_packed_dim) { - return true; - } - - // Check if all input tensors have dimensions that are multiples of 4 along the packed dimension - bool all_concat_dim_size_multiple_of_4 = true; - $for i in range(NUM_INPUTS): - all_concat_dim_size_multiple_of_4 = - all_concat_dim_size_multiple_of_4 && - (in${i+1}_sizes[concat_dim] % 4 == 0); +#define NUM_INPUTS ${NUM_INPUTS} - return all_concat_dim_size_multiple_of_4; -} +#include "concat_utils.glslh" +/* + * This shader template concatenates up to NUM_INPUT input tensors to the + * output tensor along the concat_dim. Elements from the input tensor will + * be inserted along the output's concat_dim starting at concat_offset. + * + * Each thread is responsible for writing out one output texel. The data + * required for the output texel may be read from multiple input texels of one + * input tensor. + */ void main() { - const ivec3 lpos = ivec3(gl_GlobalInvocationID); - ivec4 out_tidx = lpos_to_tidx(lpos, out_sizes, out_axis_map.w, out_packed_dim); - - if (any(greaterThanEqual(out_tidx, out_sizes))) { + const int tid = ivec3(gl_GlobalInvocationID).x; + + // Sum of the sizes of all input tensors along the concat_dim + const int concat_numel = total_concat_dim_numel(); + + // The 1-3 input tensors are interpreted as one concatenated tensor ("volume") + // along the concat_dim for the purposes of tensor indexing. Each thread is + // responsible for writing out 4 elements along the packed dim of the output + // tensor by reading the source data from the input tensor(s). + ivec4 inp_volume_sizes = out_sizes; + inp_volume_sizes[concat_dim] = total_concat_dim_numel(); + + // Reconstruct inp_volume_texel_sizes from Concat.cpp + ivec4 inp_volume_texel_sizes = inp_volume_sizes; + inp_volume_texel_sizes[out_packed_dim] = DIV_UP_4( + inp_volume_texel_sizes[out_packed_dim] + ) + 1; + + // tensor index of the first element that will be read from the input volume + ivec4 inp_volume_start_tidx = nchwi_to_tidx(tid, inp_volume_texel_sizes); + inp_volume_start_tidx[out_packed_dim] = MUL_4( + inp_volume_start_tidx[out_packed_dim] + ); + + int concat_offset = t_concat_offset[0]; + + // tensor index of the first element that will be written to the output tensor + ivec4 out_write_start_tidx = inp_volume_start_tidx; + out_write_start_tidx[concat_dim] += concat_offset; + + // To write to the the desired output element, we will need to load the texel + // to which the element belongs. Calculate the tensor index of the first + // element of that texel. + ivec4 out_read_start_tidx = out_write_start_tidx; + out_read_start_tidx[out_packed_dim] = ALIGN_DOWN_4( + out_write_start_tidx[out_packed_dim]); + + // bounds check + if (any(greaterThanEqual(out_read_start_tidx, out_sizes))) { return; } - if (can_use_fast_path()) { - // Fast path: No texel merging required - ivec4 in_tidx = out_tidx; + ivec3 out_pos = tidx_to_pos( + out_read_start_tidx, + out_sizes, + out_axis_map, + out_packed_dim + ); - $for i in range(NUM_INPUTS): - // For each input tensor, check if the tensor index is within bounds. If - // so, read the texel from the input tensor and write it to the output - if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) { - const ivec3 in_pos = tidx_to_pos(in_tidx, in${i+1}_sizes, in${i+1}_axis_map, in${i+1}_packed_dim); - const VEC4_T in_texel = load_texel(t_in${i+1}, in_pos); - write_texel_lpos(t_out, lpos, in_texel, out_axis_map); - return; - } - // Otherwise, adjust the index along the concat dimension and try the next - // input tensor. - else { - in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim]; - } - } - else { - // Slow path: Texel merging required - VEC4_T out_texel = VEC4_T(0); + VEC4_T out_texel = imageLoad(t_out, out_pos); - // Process each element in the output texel individually - for (int texel_i = 0; texel_i < 4; ++texel_i) { - ivec4 curr_out_tidx = out_tidx; - curr_out_tidx[out_packed_dim] += texel_i; + VEC4_T test_texel = VEC4_T(-1.0); - // Skip if we're out of bounds - if (curr_out_tidx[out_packed_dim] >= out_sizes[out_packed_dim]) { - continue; - } + for (int comp = 0; comp < 4; ++comp) { + ivec4 out_tidx = out_read_start_tidx; + out_tidx[out_packed_dim] += comp; - ivec4 in_tidx = curr_out_tidx; - $for i in range(NUM_INPUTS): - // For each input tensor, check if the tensor index is within bounds. If - // so, read the corresponding texel element from the input tensor and - // write it to the output texel. - if (in_tidx[concat_dim] < in${i+1}_sizes[concat_dim]) { - const ivec4 in_posi = tidx_to_posi(in_tidx, in${i+1}_sizes, in${i+1}_axis_map, in${i+1}_packed_dim); - out_texel[texel_i] = load_texel(t_in${i+1}, in_posi.xyz)[in_posi.w]; - continue; - } - // Otherwise, adjust the index along the concat dimension and try the - // next input tensor. - else { - in_tidx[concat_dim] -= in${i+1}_sizes[concat_dim]; - } + + // It's possible that the current texel element has been written to as part + // of the previous input batch; if so, then don't overwrite this texel + // element + if (out_tidx[concat_dim] < concat_offset) { + test_texel[comp] = -5.0; + continue; } - write_texel_lpos(t_out, lpos, out_texel, out_axis_map); + // Calculate the tidx of the input volume that corresponds to this output + // element + ivec4 inp_volume_tidx = out_tidx; + inp_volume_tidx[concat_dim] -= concat_offset; + + // go through the list of input tensors, and figure out which input this + // output element should be read from. + $for i in range(NUM_INPUTS): + if (inp_volume_tidx[concat_dim] < inp${i}_sizes[concat_dim]) { + // Special fast path case if, for the first output texel element, the + // corresponding input element is at the start of the texel it belongs + // to. In this case, the input texel can be written as-is to the output + // texel. Also require that The entire input texel is valid and does not + // contain any padding elements. + if (comp == 0 && + out_tidx[out_packed_dim] % 4 == 0 && + inp_volume_tidx[inp${i}_packed_dim] % 4 == 0 && + inp_volume_tidx[inp${i}_packed_dim] + 3 < inp${i}_sizes[inp${i}_packed_dim]) { + const ivec3 in_pos = tidx_to_pos( + inp_volume_tidx, + inp${i}_sizes, + inp${i}_axis_map, + inp${i}_packed_dim); + + out_texel = texelFetch(t_inp${i}, in_pos, 0); + break; + } + + // Otherwise, locate the specific input element required + const ivec4 in_posi = tidx_to_posi( + inp_volume_tidx, + inp${i}_sizes, + inp${i}_axis_map, + inp${i}_packed_dim); + + out_texel[comp] = texelFetch(t_inp${i}, in_posi.xyz, 0)[in_posi.w]; + test_texel[comp] = out_texel[comp]; + continue; + } + else { + inp_volume_tidx[concat_dim] -= inp${i}_sizes[concat_dim]; + } } + + imageStore(t_out, out_pos, out_texel); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/concat_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/concat_utils.glslh new file mode 100644 index 00000000000..000b86a7fce --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/concat_utils.glslh @@ -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. + */ + +#ifndef CONCAT_UTILS_H +#define CONCAT_UTILS_H + + +/********************************** + * Concatenation utililty functions + * + */ + +/* + * Returns the total number of elements along the concatenation dim that will + * be concatenated in this input batch. + */ +$for N in range(1, 4): + #if NUM_INPUTS == ${N} + int total_concat_dim_numel() { + int total = 0; + $for i in range(N): + total += inp${i}_sizes[concat_dim]; + + return total; + } + #endif + +#endif // CONCAT_UTILS_H diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl index c0ed9204227..0f5dbc41273 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d.glsl @@ -30,6 +30,8 @@ ${layout_declare_ubo(8, "float", "out_min", "float", "out_max")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +${layout_declare_spec_const(C, "int", "ngroups", "1")} + /* * Computes a 2D convolution. Each shader invocation calculates the output at * a single output location. @@ -74,7 +76,18 @@ void main() { // Perform the convolution by iterating over the overlay region. VEC4_T sum = texelFetch(t_bias, ivec2(pos.z, 0), 0); const int ic4 = in_group_size / 4; - for (int z4 = 0; z4 < ic4; ++z4, kstart.x += kernel_size.x * 4) { + + int z_start = 0; + int z_end = ic4; + if (ngroups > 1) { + const int group_size = (out_limits.z) / ngroups; + const int group_idx = pos.z / group_size; + + z_start = ic4 * group_idx; + z_end = z_start + ic4; + } + + for (int z4 = z_start; z4 < z_end; ++z4, kstart.x += kernel_size.x * 4) { for (int y = start.y, ky = kstart.y; y < end.y; y += dilation.y, ++ky) { for (int x = start.x, kx = kstart.x; x < end.x; x += dilation.x, kx += 4) { const VEC4_T in_texel = texelFetch(t_in, ivec3(x, y, z4), 0); diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_common.glslh b/backends/vulkan/runtime/graph/ops/glsl/conv2d_common.glslh new file mode 100644 index 00000000000..41825cba867 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_common.glslh @@ -0,0 +1,51 @@ +/* + * 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. + */ + +#ifndef CONV2D_COMMON_GLSLH +#define CONV2D_COMMON_GLSLH + +#include "common.glslh" + +struct Conv2DParams { + ivec2 kernel_size; + ivec2 stride; + ivec2 padding; + ivec2 dilation; + int groups; + int out_channels_per_group; + int in_channels_per_group; + int logical_K_per_group; + int K_per_group; + int K4_per_group; + int logical_K; + int K; + int K4; +}; + +#ifdef DEBUG_MODE + +void printConv2DParams(const Conv2DParams params) { + debugPrintfEXT("Conv2DParams: \\n"); + debugPrintfEXT( + " kernel_size: %d, %d\\n", params.kernel_size.x, params.kernel_size.y); + debugPrintfEXT(" stride: %d, %d\\n", params.stride.x, params.stride.y); + debugPrintfEXT(" padding: %d, %d\\n", params.padding.x, params.padding.y); + debugPrintfEXT(" dilation: %d, %d\\n", params.dilation.x, params.dilation.y); + debugPrintfEXT(" groups: %d\\n", params.groups); + debugPrintfEXT( + " out_channels_per_group: %d\\n", params.out_channels_per_group); + debugPrintfEXT( + " in_channels_per_group: %d\\n", params.in_channels_per_group); + debugPrintfEXT(" logical_K_per_group: %d\\n", params.logical_K_per_group); + debugPrintfEXT(" K_per_group: %d\\n", params.K_per_group); + debugPrintfEXT(" K4_per_group: %d\\n", params.K4_per_group); +} + +#endif // DEBUG_MODE + +#endif // CONV2D_COMMON_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl index 8a845b6a8a6..02fbef29b75 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_dw.glsl @@ -30,6 +30,8 @@ ${layout_declare_ubo(8, "float", "out_min", "float", "out_max")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +${layout_declare_spec_const(C, "int", "ngroups", "1")} + /* * Computes a depthwise convolution. Each shader invocation calculates the * output at a single output location. diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block.glslh new file mode 100644 index 00000000000..7add8c4cd16 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block.glslh @@ -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. + */ + +#ifndef CONV2D_FP_IM2COL_BLOCK +#define CONV2D_FP_IM2COL_BLOCK + +/* + * Defines utilities to convert between (col, row) indices of an im2col matrix + * and 4-dimension tensor indices of image tensors. + * + * Requires: + * - output_sizes to be defined in the shader layout, corresponding to the sizes + * of the output image of the convolution op. + * - image_sizes to be defined in the shader layout, corresponding to the sizes + * of the input image of the convolution op. + * - conv2d_params to be defined in the shader layout + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "common.glslh" +#include "conv2d_common.glslh" + +struct Im2ColMatrixIdx { + int row; + int col; + // Relevant for grouped convolution. This indicates the column index relative + // to the first column in the group. + int col_idx_in_group; + int group_idx; +}; + +void unwrap_m(out TensorIndex4D out_tidx_base, const int m) { + out_tidx_base.data[3] = m / (output_sizes.y * output_sizes.x); + out_tidx_base.data[1] = (m / output_sizes.x) % output_sizes.y; + out_tidx_base.data[0] = m % output_sizes.x; + + // Initialize channels to 0; assume it will be set later on + out_tidx_base.data[2] = 0; +} + +void im2col_tidx_to_output_tidx( + out TensorIndex4D output_tidx, + const Im2ColMatrixIdx im2col_tidx) { + unwrap_m(output_tidx, im2col_tidx.row); + // Set channels + output_tidx.data.z = im2col_tidx.col; +} + +/* + * Converts im2col matrix position to corresponding 4D tensor index, accounting + * for grouped convolutions. The conversion should ensure that all data within + * the same group occupy a contiguous block in memory. + */ +void im2col_idx_to_input_tidx( + out TensorIndex4D input_tidx, + const Im2ColMatrixIdx im2col_idx) { + TensorIndex4D output_tidx; + unwrap_m(output_tidx, im2col_idx.row); + + const int in_channels_per_group = conv2d_params.in_channels_per_group; + // Determine the corresponding position within the convolution window based + // on the col index (more specifically, the col index within the group) + const int channel_within_group = + im2col_idx.col_idx_in_group % in_channels_per_group; + const int kernel_x = (im2col_idx.col_idx_in_group / in_channels_per_group) % + conv2d_params.kernel_size.x; + const int kernel_y = im2col_idx.col_idx_in_group / + (in_channels_per_group * conv2d_params.kernel_size.x); + + // Calculate the actual input channel index + const int channel_idx = + im2col_idx.group_idx * conv2d_params.in_channels_per_group + + channel_within_group; + + // Calculate corresponding input coordinates based on output position + // associated with the row index. + const int input_y = int(output_tidx.data.y * conv2d_params.stride.y) - + int(conv2d_params.padding.y) + int(kernel_y * conv2d_params.dilation.y); + const int input_x = int(output_tidx.data.x * conv2d_params.stride.x) - + int(conv2d_params.padding.x) + int(kernel_x * conv2d_params.dilation.x); + + input_tidx.data = ivec4(input_x, input_y, channel_idx, output_tidx.data.w); +} + +// 4x4 block of the im2col matrix +struct FPIm2ColBlock { + VEC4_T data[4]; +}; + +#endif // CONV2D_FP_IM2COL_BLOCK diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_load.glslh new file mode 100644 index 00000000000..c02b070e17e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_load.glslh @@ -0,0 +1,169 @@ +/* + * 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. + */ + +#ifndef CONV2D_FP_IM2COL_BLOCK_LOAD +#define CONV2D_FP_IM2COL_BLOCK_LOAD + +/* + * Defines utilities to load data for a 4x4 im2col matrix block from an + * input image and store the data as a FPInputTile. + * + * Requires: + * - t_input to be defined in the shader layout, representing the texture of the + * source image + * - conv2d_params to be defined in the shader layout + */ + +#extension GL_EXT_control_flow_attributes : require + +#extension GL_EXT_debug_printf : require + +#include "common.glslh" +#include "conv2d_common.glslh" +#include "conv2d_fp_im2col_block.glslh" +#include "linear_fp_input_tile.glslh" + +VEC4_T load_input_texel(const TensorIndex4D tidx) { + // Assumes batch size is 1 and channels packing + return texelFetch( + t_input, ivec3(tidx.data.x, tidx.data.y, div_4(tidx.data.z)), 0); +} + +T load_input_texel_element(const TensorIndex4D tidx) { + const int channels_texel_idx = div_4(tidx.data.z); + const int texel_comp = mod_4(tidx.data.z); + // Assumes batch size is 1 and channels packing + return texelFetch( + t_input, + ivec3(tidx.data.x, tidx.data.y, channels_texel_idx), + 0)[texel_comp]; +} + +// k4 -> group of 4 input channels idx +// m -> flattened batch, output width, output height dim idx +/* + * Fast impl for when the input image's channels per group is a multiple of 4. + * In this case, it is guaranteed that a texel loaded from the input can be + * stored directly to the output without any additional filtering. + */ +void load_im2col_block_fast( + out FPIm2ColBlock block, + const int k4, + const int m4, + const int logical_K, + const int M) { + Im2ColMatrixIdx im2col_idx; + im2col_idx.col = mul_4(k4); // k + im2col_idx.row = mul_4(m4); // m + + // Due to the assumption that in_channels_per_group % 4 == 0, it is + // guaranteed that the next 4 columns (including this one) is part of the + // same group. + im2col_idx.group_idx = im2col_idx.col / conv2d_params.K_per_group; + im2col_idx.col_idx_in_group = im2col_idx.col % conv2d_params.K_per_group; + + [[unroll]] for (int m_off = 0; m_off < 4; ++m_off) { + if (im2col_idx.row >= M) { + block.data[m_off] = VEC4_T(0); + continue; + } + + TensorIndex4D input_tidx; + im2col_idx_to_input_tidx(input_tidx, im2col_idx); + + // Load the texel + block.data[m_off] = load_input_texel(input_tidx); + + im2col_idx.row++; + } +} + +/* + * If input image channels is not a multiple of 4, then it is likely that for + * some matrix texels, the source data is split between different texels of the + * source image. In this case it's better to retreive each element individually. + */ +void load_im2col_block_slow( + out FPIm2ColBlock block, + const int k4, + const int m4, + const int logical_K, + const int M) { + Im2ColMatrixIdx im2col_idx_base; + im2col_idx_base.col = mul_4(k4); + im2col_idx_base.row = mul_4(m4); + + im2col_idx_base.group_idx = im2col_idx_base.col / conv2d_params.K_per_group; + im2col_idx_base.col_idx_in_group = + im2col_idx_base.col % conv2d_params.K_per_group; + + [[unroll]] for (int m_off = 0; m_off < 4; ++m_off) { + [[unroll]] for (int k_off = 0; k_off < 4; ++k_off) { + Im2ColMatrixIdx im2col_idx = im2col_idx_base; + im2col_idx.row += m_off; + im2col_idx.col_idx_in_group += k_off; + + // bounds checking + if (im2col_idx.col_idx_in_group >= conv2d_params.logical_K_per_group || + im2col_idx.row >= M) { + block.data[m_off][k_off] = T(0); + continue; + } + + TensorIndex4D input_tidx; + im2col_idx_to_input_tidx(input_tidx, im2col_idx); + + block.data[m_off][k_off] = load_input_texel_element(input_tidx); + } + } +} + +void load_im2col_block( + out FPIm2ColBlock block, + const int k4, + const int m4, + const int logical_K, + const int M) { + if (mod_4(conv2d_params.in_channels_per_group) == 0) { + load_im2col_block_fast(block, k4, m4, logical_K, M); + } else { + load_im2col_block_slow(block, k4, m4, logical_K, M); + } +} + +void load_input_im2col_tile( + out FPInputTile tile, + const int k4_start, + const int m4_start, + const int logical_K, + const int M) { + FPIm2ColBlock block; +#if TILE_K4 == 1 + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + load_im2col_block(block, k4_start, m4_start + m4, logical_K, M); + for (int row = 0; row < 4; ++row) { + const int m = mul_4(m4) + row; + tile.data[m][0] = block.data[row]; + } + } + +#else + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + load_im2col_block(block, k4_start + k4, m4_start + m4, logical_K, M); + for (int row = 0; row < 4; ++row) { + const int m = mul_4(m4) + row; + tile.data[m][k4] = block.data[row]; + } + } + } + +#endif +} + +#endif // CONV2D_FP_IM2COL_BLOCK_LOAD diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_store.glslh new file mode 100644 index 00000000000..2171d75c628 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_fp_im2col_block_store.glslh @@ -0,0 +1,68 @@ +/* + * 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. + */ + +#ifndef CONV2D_FP_IM2COL_BLOCK_STORE +#define CONV2D_FP_IM2COL_BLOCK_STORE + +/* + * Defines utilities to store data for a 4x4 im2col output matrix block computed + * from matrix multiplication to an output image. + * + * Requires: + * - t_output to be defined in the shader layout, representing the texture of + * the output image + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "common.glslh" +#include "conv2d_common.glslh" +#include "conv2d_fp_im2col_block.glslh" +#include "linear_fp_output_tile.glslh" + +// TODO: implement buffer support +void write_output_texel(const VEC4_T out_texel, const TensorIndex4D tidx) { + // Assume batch size is 1 + imageStore( + t_output, ivec3(tidx.data.x, tidx.data.y, div_4(tidx.data.z)), out_texel); +} + +void write_im2col_tile_as_image( + const FPOutTile tile, + const int n4_start, + const int m_start) { + Im2ColMatrixIdx im2col_tidx; + im2col_tidx.col = mul_4(n4_start); + im2col_tidx.row = m_start; +#if TILE_K4 == 1 + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + TensorIndex4D output_tidx; + im2col_tidx_to_output_tidx(output_tidx, im2col_tidx); + + if (any(greaterThanEqual(output_tidx.data, output_sizes))) { + continue; + } + write_output_texel(tile.data[m][0], output_tidx); + im2col_tidx.row++; + } + +#else + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + TensorIndex4D output_tidx; + im2col_tidx_to_output_tidx(output_tidx, im2col_tidx); + + write_output_texel(tile.data[m][k4], output_tidx); + im2col_tidx.row++; + } + } + +#endif +} + +#endif // CONV2D_FP_IM2COL_BLOCK_STORE diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl index cf9714ca468..4c6031152ee 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw.glsl @@ -38,6 +38,8 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +${layout_declare_spec_const(C, "int", "ngroups", "1")} + #extension GL_EXT_control_flow_attributes : require /* diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl index a46f1e3b99c..9f84afeb1a1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_pw_s1p0.glsl @@ -40,6 +40,8 @@ layout(push_constant) uniform restrict Block { layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +${layout_declare_spec_const(C, "int", "ngroups", "1")} + #extension GL_EXT_control_flow_attributes : require /* diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.glsl new file mode 100644 index 00000000000..3615d423230 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.glsl @@ -0,0 +1,120 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, OUTPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, OUTPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "conv2d_common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, 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_scales", DTYPE, "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", DTYPE, "buffer", is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} +${layout_declare_ubo(B, "Conv2DParams", "conv2d_params")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "1")} + +#include "linear_fp_input_tile_load.glslh" +#include "linear_int8_weight_tile_load.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_fp_bias_load.glslh" +#include "linear_fp_output_tile_fp_int8_compute.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "conv2d_fp_im2col_block_store.glslh" + +void main() { + // Each thread writes out a 4 wide x 4 high tile of output values + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = int(out_tile_x * TILE_N); + const int m = int(out_tile_y * TILE_M); + + const int n4 = div_4(n); + const int m4 = div_4(m); + + // M = flattened output width, height, batches dims + const int M = output_sizes.x * output_sizes.y * output_sizes.w; + // N = output channels + const int N = output_sizes.z; + + if (n >= N || m >= M) { + return; + } + + const int group_idx = n / conv2d_params.out_channels_per_group; + const int input_k4_offset = conv2d_params.K4_per_group * group_idx; + + const int K4 = conv2d_params.K4; + const int N4 = div_up_4(N); + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile in_tile; + Int8WeightTile int8_weight_tile; + + const bool dont_check_bounds = (M - m) >= TILE_M; + + if (dont_check_bounds) { + for (int k4 = 0; k4 < conv2d_params.K4_per_group; k4++) { + load_input_tile_no_checks(in_tile, k4 + input_k4_offset, m, K4, M); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + fp_accumulate_with_int8_weight(out_tile, in_tile, int8_weight_tile); + } + } else { + for (int k4 = 0; k4 < conv2d_params.K4_per_group; k4++) { + load_input_tile_with_checks(in_tile, k4 + input_k4_offset, m, K4, M); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + fp_accumulate_with_int8_weight(out_tile, in_tile, int8_weight_tile); + } + } + + FPPerOutChannelParams weight_scales_tile; + load_weight_scales_tile(weight_scales_tile, n4); + + if (apply_bias > 0) { + FPPerOutChannelParams bias_tile; + load_bias_tile(bias_tile, n4); + + apply_weight_scales_and_biases(out_tile, weight_scales_tile, bias_tile); + } + else { + apply_weight_scales(out_tile, weight_scales_tile); + } + + write_im2col_tile_as_image(out_tile, n4, m); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.yaml new file mode 100644 index 00000000000..9b3b5aa2c0a --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8csw_linear_tiled.yaml @@ -0,0 +1,22 @@ +# 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. + +conv2d_q8csw_linear_tiled: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: texture3d + INPUT_STORAGE: buffer + WEIGHT_STORAGE: texture2d + TILE_M4: 1 + TILE_N4: 1 + TILE_K4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: conv2d_q8csw_linear_tiled_texture3d_buffer_texture2d + - NAME: conv2d_q8csw_linear_tiled_texture3d_buffer_buffer + WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.glsl new file mode 100644 index 00000000000..f74a1311095 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.glsl @@ -0,0 +1,137 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, OUTPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, OUTPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if PACKED_INT8_INPUT_STORAGE == "buffer": + #define PACKED_INT8_INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +#extension GL_EXT_integer_dot_product : require + +layout(std430) buffer; + +#include "conv2d_common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", PACKED_INT8_INPUT_STORAGE, 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)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} +${layout_declare_ubo(B, "Conv2DParams", "conv2d_params")} + +layout(push_constant) uniform restrict Block { + float input_scale; + int input_zp; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "1")} + +#include "linear_int8_input_tile_load.glslh" +#include "linear_int8_weight_tile_load.glslh" +#include "linear_int_weight_sums_load.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_fp_output_tile_int8_int8_compute.glslh" +#include "linear_fp_bias_load.glslh" +#include "conv2d_fp_im2col_block_store.glslh" + +void main() { + // Each thread writes out a 4 wide x 4 high tile of output values + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = int(out_tile_x * TILE_N); + const int m = int(out_tile_y * TILE_M); + + const int n4 = div_4(n); + const int m4 = div_4(m); + + // M = flattened output width, height, batches dims + const int M = output_sizes.x * output_sizes.y * output_sizes.w; + // N = output channels + const int N = output_sizes.z; + + if (n >= N || m >= M) { + return; + } + + const int group_idx = n / conv2d_params.out_channels_per_group; + const int input_k4_offset = conv2d_params.K4_per_group * group_idx; + + const int K4 = conv2d_params.K4; + const int N4 = div_up_4(N); + + Int32Accum out_accum; + initialize(out_accum); + + Int8InputTile int8_in_tile; + Int8WeightTile int8_weight_tile; + + for (int k4 = 0; k4 < conv2d_params.K4_per_group; k4++) { + load_int8_input_tile(int8_in_tile, k4 + input_k4_offset, m4, K4); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + + int_accumulate_with_int8_weight(out_accum, int8_in_tile, int8_weight_tile); + } + + FPPerOutChannelParams weight_scales_tile; + load_weight_scales_tile(weight_scales_tile, n4); + + IntPerOutChannelParams weight_sums_tile; + load_weight_sums_tile(weight_sums_tile, n4); + + FPOutTile out_tile; + initialize(out_tile); + if (apply_bias > 0) { + FPPerOutChannelParams bias_tile; + load_bias_tile(bias_tile, int(n4)); + + accumulate_out_tile_with_int_accum( + out_tile, + out_accum, + input_scale, + input_zp, + weight_sums_tile, + weight_scales_tile, + bias_tile); + } + else { + accumulate_out_tile_with_int_accum( + out_tile, + out_accum, + input_scale, + input_zp, + weight_sums_tile, + weight_scales_tile); + } + + write_im2col_tile_as_image(out_tile, n4, m); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.yaml new file mode 100644 index 00000000000..629001765c1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/conv2d_q8ta_q8csw_linear_tiled.yaml @@ -0,0 +1,22 @@ +# 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. + +conv2d_q8ta_q8csw_linear_tiled: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: texture3d + PACKED_INT8_INPUT_STORAGE: buffer + WEIGHT_STORAGE: texture2d + TILE_M4: 1 + TILE_N4: 1 + TILE_K4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: conv2d_q8ta_q8csw_linear_tiled_texture3d_buffer_texture2d + - NAME: conv2d_q8ta_q8csw_linear_tiled_texture3d_buffer_buffer + WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.glsl index 43e62eadeee..57dc2d53fff 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.glsl @@ -12,12 +12,16 @@ #define IN_T ${buffer_scalar_type(IN_DTYPE)} #define OUT_T ${buffer_scalar_type(OUT_DTYPE)} +#define SCALE_T ${buffer_scalar_type(SCALE_DTYPE)} +#define ZP_T ${buffer_scalar_type(ZP_DTYPE)} #define ${MODE} ${define_active_storage_type("buffer")} ${define_required_extensions(IN_DTYPE)} ${define_required_extensions(OUT_DTYPE)} +${define_required_extensions(SCALE_DTYPE)} +${define_required_extensions(ZP_DTYPE)} layout(std430) buffer; @@ -27,16 +31,16 @@ ${layout_declare_tensor(B, "w", "t_out", OUT_DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "buffer")} $if MODE == "per_tensor": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int quant_min; int quant_max; }; $if MODE == "per_token": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int num_tokens; @@ -44,8 +48,8 @@ $if MODE == "per_token": int quant_max; }; $if MODE == "per_channel": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int axis; @@ -54,8 +58,8 @@ $if MODE == "per_channel": int quant_max; }; $if MODE == "block_wise": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { ivec4 blockSize; // bW, bH, bC, bN @@ -150,7 +154,7 @@ void dequantize_per_tensor() { const int in_bufi = tidx_to_bufi(out_tidx, t_in_strides); IN_T qvalue = t_in[in_bufi]; - OUT_T value = dequantize_val(qvalue, t_scale[0], t_zero_point[0]); + OUT_T value = dequantize_val(qvalue, float(t_scale[0]), int(t_zero_point[0])); t_out[out_bufi] = value; } @@ -185,7 +189,7 @@ void dequantize_per_token() { token_idx = min(token_idx, num_tokens - 1); - OUT_T value = dequantize_val(qvalue, t_scale[token_idx], t_zero_point[token_idx]); + OUT_T value = dequantize_val(qvalue, float(t_scale[token_idx]), int(t_zero_point[token_idx])); t_out[out_bufi] = value; } @@ -224,7 +228,7 @@ void dequantize_per_channel() { channel_idx = min(channel_idx, num_channels - 1); - OUT_T value = dequantize_val(qvalue, t_scale[channel_idx], t_zero_point[channel_idx]); + OUT_T value = dequantize_val(qvalue, float(t_scale[channel_idx]), int(t_zero_point[channel_idx])); t_out[out_bufi] = value; } @@ -247,7 +251,7 @@ void dequantize_block_wise() { const int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w; - const OUT_T value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id]); + const OUT_T value = dequantize_val(qvalue, float(t_scale[block_id]), int(t_zero_point[block_id])); t_out[out_bufi] = value; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.yaml index 999c59d3b79..a4375038a75 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/dequantize_buffer.yaml @@ -2,6 +2,8 @@ dequantize_buffer: parameter_names_with_default_values: IN_DTYPE: int32 OUT_DTYPE: float + SCALE_DTYPE: float + ZP_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: @@ -12,6 +14,12 @@ dequantize_buffer: - VALUE: half - VALUE: float - VALUE: double + SCALE_DTYPE: + - VALUE: float + ZP_DTYPE: + - VALUE: int8 + - VALUE: int32 + - VALUE: float shader_variants: - NAME: dequantize_per_tensor_buffer MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.glsl index 20bf6c87e26..19276cd8f7f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.glsl @@ -15,12 +15,16 @@ #define OUT_T ${buffer_scalar_type(OUT_DTYPE)} #define FVEC4_T ${texel_load_type(OUT_DTYPE, "texture3d")} +#define SCALE_T ${buffer_scalar_type(SCALE_DTYPE)} +#define ZP_T ${buffer_scalar_type(ZP_DTYPE)} #define ${MODE} ${define_active_storage_type("texture3d")} ${define_required_extensions(IN_DTYPE)} ${define_required_extensions(OUT_DTYPE)} +${define_required_extensions(SCALE_DTYPE)} +${define_required_extensions(ZP_DTYPE)} #extension GL_EXT_control_flow_attributes : require @@ -30,16 +34,16 @@ ${layout_declare_tensor(B, "w", "t_out", OUT_DTYPE, "texture3d")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "texture3d")} $if MODE == "per_tensor": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int quant_min; int quant_max; }; $if MODE == "per_token": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int num_tokens; @@ -47,8 +51,8 @@ $if MODE == "per_token": int quant_max; }; $if MODE == "per_channel": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int axis; @@ -57,8 +61,8 @@ $if MODE == "per_channel": int quant_max; }; $if MODE == "block_wise": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { ivec4 blockSize; // bW, bH, bC, bN @@ -160,7 +164,7 @@ void dequantize_per_tensor() { [[unroll]] for (int i = 0; i < 4; ++i) { IN_T qvalue = IN_T(intex[i]); - OUT_T value = dequantize_val(qvalue, t_scale[0], t_zero_point[0]); + OUT_T value = dequantize_val(qvalue, float(t_scale[0]), int(t_zero_point[0])); $if OUT_DTYPE == "double": outtex[i] = float(value); @@ -196,8 +200,8 @@ void dequantize_per_token() { token_idx = min(token_idx, num_tokens - 1); // Scale and zero_point are prepacked as buffers, so direct access - float scale_val = t_scale[token_idx]; - int zero_point_val = t_zero_point[token_idx]; + float scale_val = float(t_scale[token_idx]); + int zero_point_val = int(t_zero_point[token_idx]); FVEC4_T outtex; [[unroll]] for (int i = 0; i < 4; ++i) { @@ -238,8 +242,8 @@ void dequantize_per_channel() { int channel_idx = pos.x * 4 + i; channel_idx = min(channel_idx, num_channels - 1); - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); OUT_T value = dequantize_val(qvalue, scale_val, zero_point_val); $if OUT_DTYPE == "double": outtex[i] = float(value); @@ -249,8 +253,8 @@ void dequantize_per_channel() { } else if (axis == 1) { int channel_idx = pos.y; channel_idx = min(channel_idx, num_channels - 1); - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T qvalue = IN_T(intex[i]); @@ -267,8 +271,8 @@ void dequantize_per_channel() { int folded_idx = pos.z; int channel_idx = folded_idx % num_channels; - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T qvalue = IN_T(intex[i]); @@ -287,8 +291,8 @@ void dequantize_per_channel() { // the C dimension N(C)HW int channel_idx = folded_idx / num_channels; - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T qvalue = IN_T(intex[i]); @@ -326,7 +330,7 @@ void dequantize_block_wise() { int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w; IN_T qvalue = IN_T(intex[i]); - OUT_T value = dequantize_val(qvalue, t_scale[block_id], t_zero_point[block_id]); + OUT_T value = dequantize_val(qvalue, float(t_scale[block_id]), int(t_zero_point[block_id])); $if OUT_DTYPE == "double": outtex[i] = float(value); $else: diff --git a/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.yaml index 9b624762192..7a58e9410d3 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/dequantize_texture.yaml @@ -2,6 +2,8 @@ dequantize_texture: parameter_names_with_default_values: IN_DTYPE: int32 OUT_DTYPE: float + SCALE_DTYPE: float + ZP_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: @@ -12,6 +14,12 @@ dequantize_texture: - VALUE: half - VALUE: float - VALUE: double + SCALE_DTYPE: + - VALUE: float + ZP_DTYPE: + - VALUE: int8 + - VALUE: int32 + - VALUE: float shader_variants: - NAME: dequantize_per_tensor_texture3d MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/expand_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/expand_buffer.glsl new file mode 100644 index 00000000000..ce433040b66 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/expand_buffer.glsl @@ -0,0 +1,51 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define VEC4_T ${texel_type(DTYPE)} +#define T ${buffer_scalar_type(DTYPE)} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "t_outp", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_inp", DTYPE, "buffer")} + +${layout_declare_ubo(B, "BufferMetadata", "outp")} +${layout_declare_ubo(B, "BufferMetadata", "inp")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const uint outp_bufi = gl_GlobalInvocationID.x; + if (outp_bufi >= numel(outp)) { + return; + } + + TensorIndex outp_tidx; + linear_idx_to_tensor_idx(outp, outp_bufi, outp_tidx); + + // Map output tensor index to input tensor index by taking modulo + // with input tensor sizes for each dimension + TensorIndex inp_tidx = outp_tidx; + for (int d = 0; d < ndim(inp); ++d) { + uint inp_size = size_at(inp, d); + uint outp_idx = idx_at(outp_tidx, d); + inp_tidx.data[div_4(d)][mod_4(d)] = outp_idx % inp_size; + } + + const uint inp_bufi = tensor_idx_to_linear_idx(inp, inp_tidx); + // Copy data from input to output + t_outp[outp_bufi] = t_inp[inp_bufi]; +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/flash_attention.yaml b/backends/vulkan/runtime/graph/ops/glsl/expand_buffer.yaml similarity index 58% rename from backends/vulkan/runtime/graph/ops/glsl/flash_attention.yaml rename to backends/vulkan/runtime/graph/ops/glsl/expand_buffer.yaml index 2314b701040..6d90e1fa8b1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/flash_attention.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/expand_buffer.yaml @@ -1,10 +1,10 @@ -flash_attention: +expand_buffer: parameter_names_with_default_values: DTYPE: float - STORAGE: buffer generate_variant_forall: DTYPE: + - VALUE: half - VALUE: float + - VALUE: int32 shader_variants: - - NAME: flash_attention_buffer - STORAGE: buffer + - NAME: expand_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/flash_attention.glsl b/backends/vulkan/runtime/graph/ops/glsl/flash_attention.glsl deleted file mode 100644 index 1b5f47f3f3c..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/flash_attention.glsl +++ /dev/null @@ -1,226 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} -#define T ${buffer_scalar_type(DTYPE)} -${define_required_extensions(DTYPE)} - -layout(std430) buffer; - -#include "indexing_utils.h" - -// Flash Attention inputs: Query, Key, Value tensors -${layout_declare_tensor(B, "rw", "t_O", DTYPE, "buffer")} -${layout_declare_tensor(B, "rw", "t_l", "float", "buffer")} -${layout_declare_tensor(B, "rw", "t_m", "float", "buffer")} -${layout_declare_tensor(B, "r", "t_Q", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_K", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_V", DTYPE, "buffer")} - -${layout_declare_ubo(B, "ivec4", "Q_sizes")} // [B, H, N, D] -${layout_declare_ubo(B, "ivec4", "K_sizes")} -${layout_declare_ubo(B, "ivec4", "V_sizes")} -${layout_declare_ubo(B, "ivec4", "O_sizes")} - -${layout_declare_ubo(B, "ivec3", "l_sizes")} // [B, H, N] -${layout_declare_ubo(B, "ivec3", "m_sizes")} // [B, H, N] - -${layout_declare_ubo(B, "float", "scale")} -${layout_declare_ubo(B, "int", "block_size_r")} // Br (num rows in Q block) -${layout_declare_ubo(B, "int", "block_size_c")} // Bc (num cols in K/V block) -${layout_declare_ubo(B, "int", "input_pos")} // Starting position for causal masking -${layout_declare_ubo(B, "int", "num_heads")} // Number of query heads -${layout_declare_ubo(B, "int", "num_kv_heads")} // Number of key/value heads -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -// Maximum block sizes to prevent array overflow -#define MAX_BR 64 -#define MAX_BC 128 - -void main() { - // Each thread processes one row block - const int thread_id = int(gl_GlobalInvocationID.x); - - // Tensor dimensions: Q_sizes = [D, H, N, B] from graph.sizes_ubo() - // The UBO layout is different from the PyTorch tensor layout - const int head_dim = Q_sizes.x; // D (head dim) - const int num_heads = Q_sizes.y; // H (num heads) - const int seq_len = Q_sizes.z; // N (sequence length) - const int batch_size = Q_sizes.w; // B (batch) - - // Block sizes - const int Br = block_size_r; - const int Bc = block_size_c; - - const int Tr = (seq_len + Br - 1) / Br; // Number of row blocks - const int total_row_blocks = batch_size * num_heads * Tr; - - if (thread_id >= total_row_blocks) { - return; - } - - // Decode thread_id to (batch, head, row_block) - const int batch = thread_id / (num_heads * Tr); - const int remaining = thread_id % (num_heads * Tr); - const int head = remaining / Tr; - const int row_block = remaining % Tr; - - // Calculate row range for this block - const int row_start = row_block * Br; - const int row_end = min(row_start + Br, seq_len); - const int actual_Br = row_end - row_start; - - // Base indices for this batch - const int q_base = batch * (seq_len * num_heads * head_dim); - const int k_base = batch * (seq_len * num_heads * head_dim); - const int v_base = batch * (seq_len * num_heads * head_dim); - const int o_base = batch * (seq_len * num_heads * head_dim); - const int lm_base = batch * (seq_len * num_heads); - - // STEP 2: Initialize O = 0, l = 0, m = -inf for this row block - for (int r = 0; r < actual_Br; r++) { - const int seq_pos = row_start + r; - const int lm_idx = lm_base + head * seq_len + seq_pos; - - t_l[lm_idx] = 0.0; - t_m[lm_idx] = -1.0 / 0.0; // -infinity - - for (int dim = 0; dim < head_dim; dim++) { - const int o_idx = o_base + seq_pos * (num_heads * head_dim) + head * head_dim + dim; - t_O[o_idx] = T(0.0); - } - } - - // STEP 5: Outer loop over column blocks (For K, V tensors) - const int Tc = (seq_len + Bc - 1) / Bc; // Number of column blocks - for (int j = 0; j < Tc; j++) { - const int col_start = j * Bc; - const int col_end = min(col_start + Bc, seq_len); - const int actual_Bc = col_end - col_start; - - // STEP 6-8 done implicitly below - - // Load current statistics for all rows in this block - float m_i[MAX_BR]; - float l_i[MAX_BR]; - for (int r = 0; r < actual_Br; r++) { - const int seq_pos = row_start + r; - const int lm_idx = lm_base + head * seq_len + seq_pos; - m_i[r] = t_m[lm_idx]; - l_i[r] = t_l[lm_idx]; - } - - // STEP 9: Compute Sij = Qi * Kj^T - T S_block[MAX_BR][MAX_BC]; // Use MAX_BR and MAX_BC constants - float m_tilde_ij[MAX_BR]; // Row maxes (float to match l/m) - float l_tilde_ij[MAX_BR]; // Row sums (float to match l/m) - - // Initialize row statistics - for (int r = 0; r < actual_Br; r++) { - m_tilde_ij[r] = -1.0 / 0.0; // -infinity - l_tilde_ij[r] = 0.0; - } - - // Compute attention scores Sij = Qi @ Kj^T - for (int r = 0; r < actual_Br; r++) { - const int global_row = row_start + r; - for (int c = 0; c < actual_Bc; c++) { - const int global_col = col_start + c; - - // For multi-query attention: map query head to KV head - const int kv_head = (head * num_kv_heads) / num_heads; - - // Dot product: Q[seq_pos, :] · K[col_pos, :] - T score = T(0.0); - for (int dim = 0; dim < head_dim; dim++) { - const int q_idx = q_base + global_row * (num_heads * head_dim) + head * head_dim + dim; - const int k_idx = k_base + global_col * (num_kv_heads * head_dim) + kv_head * head_dim + dim; - score += t_Q[q_idx] * t_K[k_idx]; - } - score *= scale; - - // Apply causal masking: mask if global_col > global_row + input_pos - if (global_col > global_row + input_pos) { - score = T(-1.0 / 0.0); // Set to negative infinity - } - - S_block[r][c] = score; - - // Track row maximum (after masking) - m_tilde_ij[r] = max(m_tilde_ij[r], float(score)); - } - } - - // STEP 10: Compute P'ij = exp(Sij − m'ij) and l'ij = rowsum(P'ij) - for (int r = 0; r < actual_Br; r++) { - // Handle the case where all scores are -inf (fully masked row) - if (isinf(m_tilde_ij[r]) && m_tilde_ij[r] < 0.0) { - // All scores are -inf, so all probabilities are 0 - for (int c = 0; c < actual_Bc; c++) { - S_block[r][c] = T(0.0); - } - l_tilde_ij[r] = 0.0; - } else { - // Normal case: compute softmax - for (int c = 0; c < actual_Bc; c++) { - S_block[r][c] = exp(S_block[r][c] - T(m_tilde_ij[r])); - l_tilde_ij[r] += float(S_block[r][c]); - } - } - } - - // STEP 11: Softmax update - float m_new_i[MAX_BR]; - float l_new_i[MAX_BR]; - for (int r = 0; r < actual_Br; r++) { - m_new_i[r] = max(m_i[r], m_tilde_ij[r]); - - l_new_i[r] = exp(m_i[r] - m_new_i[r]) * l_i[r] + exp(m_tilde_ij[r] - m_new_i[r]) * l_tilde_ij[r]; - } - - // STEP 12: Update Oi - for (int r = 0; r < actual_Br; r++) { - const int global_row = row_start + r; - float alpha = exp(m_i[r] - m_new_i[r]); - float beta = exp(m_tilde_ij[r] - m_new_i[r]); - - // For multi-query attention: map query head to KV head - const int kv_head = (head * num_kv_heads) / num_heads; - - for (int dim = 0; dim < head_dim; dim++) { - const int o_idx = o_base + global_row * (num_heads * head_dim) + head * head_dim + dim; - - // Compute P'ij @ Vj for this dimension - T pv_sum = T(0.0); - for (int c = 0; c < actual_Bc; c++) { - const int global_col = col_start + c; - const int v_idx = v_base + global_col * (num_kv_heads * head_dim) + kv_head * head_dim + dim; - pv_sum += S_block[r][c] * t_V[v_idx]; - } - - // Check for division by zero before updating output - if (l_new_i[r] <= 0.0) { - t_O[o_idx] = T(0.0); // Set to zero to avoid NaN - } else { - // Oi = (alpha * l_i * Oi + beta * P'ij @ Vj) / l_new_i - t_O[o_idx] = (T(alpha) * T(l_i[r]) * t_O[o_idx] + T(beta) * pv_sum) / T(l_new_i[r]); - } - } - } - - // STEP 13: Update li, mi - for (int r = 0; r < actual_Br; r++) { - const int seq_pos = row_start + r; - const int lm_idx = lm_base + head * seq_len + seq_pos; - t_l[lm_idx] = l_new_i[r]; - t_m[lm_idx] = m_new_i[r]; - } - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/im2col.glsl new file mode 100644 index 00000000000..f006ec993fe --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/im2col.glsl @@ -0,0 +1,108 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#define TILE_M4 1 +#define TILE_N4 1 +#define TILE_K4 1 + +#define TILE_M 4 +#define TILE_N 4 +#define TILE_K 4 + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "conv2d_common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} + +// Sizes of the im2col matrix of the convolution input +${layout_declare_ubo(B, "ivec4", "matrix_sizes")} +// Sizes of the input image +${layout_declare_ubo(B, "ivec4", "input_sizes")} +// Sizes of the output image +${layout_declare_ubo(B, "ivec4", "output_sizes")} + +${layout_declare_ubo(B, "Conv2DParams", "conv2d_params")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "conv2d_fp_im2col_block_load.glslh" + +#ifdef OUTPUT_BUFFER + +void write_tile( + const FPInputTile in_tile, + const int k4, + const int m_start, + const int K4) { + [[unroll]] for (int m = 0; m < TILE_M; m++) { + t_output[(m_start + m) * K4 + k4] = in_tile.data[m][0]; + } +} + +#else // OUTPUT_TEXTURE + +void write_tile( + const FPInputTile in_tile, + const int k4, + const int m_start, + const int K4) { + [[unroll]] for (int m = 0; m < TILE_M; m++) { + imageStore(t_output, ivec3(k4, m_start + m, 0), vec4(in_tile.data[m][0])); + } +} + +#endif // OUTPUT_BUFFER + +void main() { + // Each thread writes out a 4 wide x 4 high block of the output matrix. The + // thread position corresponds to the block index. + const int k4 = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + // Convert block idx to tensor idx + const int k = mul_4(k4); + const int m = mul_4(m4); + + const int in_channels_per_group = input_sizes.z / conv2d_params.groups; + + // Logical K dim size (unpadded) + const int logical_K = conv2d_params.logical_K; + // Physical K dim, which contains padding elements + const int K = matrix_sizes.x; + + // M dim, which represents the number of flattened output width, height, + // batches. Unlike K, there is no difference between the physical and logical + // sizes. + const int M = matrix_sizes.y; + + if (k >= K || m >= M) { + return; + } + + FPInputTile in_tile; + load_input_im2col_tile(in_tile, k4, m4, logical_K, M); + + // Number of texels in the x dim of the output matrix + const int K4 = div_4(K); + write_tile(in_tile, k4, m, K4); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/im2col.yaml new file mode 100644 index 00000000000..dd486b0e1a6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/im2col.yaml @@ -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. + +im2col: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: im2col_buffer_texture3d + - NAME: im2col_texture3d_texture3d + OUTPUT_STORAGE: texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.glsl b/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.glsl index d7bef9f0163..1498ed01aef 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.glsl @@ -16,10 +16,11 @@ ${define_active_storage_type(STORAGE)} ${define_required_extensions(DTYPE)} +${define_required_extensions(BUF_DTYPE)} layout(std430) buffer; -${layout_declare_buffer(B, "w", "buf_out", DTYPE)} +${layout_declare_buffer(B, "w", "buf_out", BUF_DTYPE)} ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} $if USE_PUSH_CONST: diff --git a/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml b/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml index 646d8f1be81..ebbc55dd9dc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/image_to_nchw.yaml @@ -7,17 +7,21 @@ image_to_nchw: parameter_names_with_default_values: DTYPE: float + BUF_DTYPE: float STORAGE: texture3d TO_STAGING: True USE_PUSH_CONST: True generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: double - - VALUE: int8 - - VALUE: uint8 - - VALUE: int32 + combination: + parameter_names: [DTYPE, BUF_DTYPE] + combos: + - parameter_values: [half, half] + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [double, double] + - parameter_values: [int8, int8] + - parameter_values: [uint8, uint8] + - parameter_values: [int32, int32] shader_variants: - NAME: image_to_nchw_texture3d - NAME: image_to_nchw_texture2d diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh b/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh new file mode 100644 index 00000000000..81783422ab4 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing.glslh @@ -0,0 +1,216 @@ +/* + * 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. + */ + +#ifndef INDEXING_GLSLH +#define INDEXING_GLSLH + +#define DIMLIMIT 8 +#define DIMLIMIT_DIV4 2 + +#define mul_4(x) ((x) << 2) +#define div_4(x) ((x) >> 2) + +#define mod_4(x) ((x) & 3) + +// +// BufferMetadata +// + +struct BufferMetadata { + uvec4 sizes[DIMLIMIT_DIV4]; + uvec4 dim_order[DIMLIMIT_DIV4]; + uvec4 strides[DIMLIMIT_DIV4]; + uvec2 ndim_numel; +}; + +uint ndim(const BufferMetadata meta) { + return meta.ndim_numel[0]; +} + +int int_ndim(const BufferMetadata meta) { + return int(meta.ndim_numel[0]); +} + +uint numel(const BufferMetadata meta) { + return meta.ndim_numel[1]; +} + +uint dim_order_at(const BufferMetadata meta, const int dim) { + return meta.dim_order[div_4(dim)][mod_4(dim)]; +} + +uint dim_order_at(const BufferMetadata meta, const uint dim) { + return meta.dim_order[div_4(dim)][mod_4(dim)]; +} + +uint stride_at(const BufferMetadata meta, const int dim) { + return meta.strides[div_4(dim)][mod_4(dim)]; +} + +uint stride_at(const BufferMetadata meta, const uint dim) { + return meta.strides[div_4(dim)][mod_4(dim)]; +} + +uint size_at(const BufferMetadata meta, const int dim) { + return meta.sizes[div_4(dim)][mod_4(dim)]; +} + +uint size_at(const BufferMetadata meta, const uint dim) { + return meta.sizes[div_4(dim)][mod_4(dim)]; +} + +bool are_equal(const BufferMetadata meta1, const BufferMetadata meta2) { + // sizes and strides must be the same to be considered equal + if (meta1.sizes[0] != meta2.sizes[0]) { + return false; + } + if (meta1.sizes[1] != meta2.sizes[1]) { + return false; + } + if (meta1.strides[0] != meta2.strides[0]) { + return false; + } + if (meta1.strides[1] != meta2.strides[1]) { + return false; + } + return true; +} + +// +// TensorIndex +// + +struct TensorIndex { + uvec4 data[DIMLIMIT_DIV4]; +}; + +void initialize(out TensorIndex tidx) { + tidx.data[0] = uvec4(0); + tidx.data[1] = uvec4(0); +} + +uint idx_at(const TensorIndex tidx, const int dim) { + return tidx.data[div_4(dim)][mod_4(dim)]; +} + +void permute(inout TensorIndex tidx, const ivec4 permute_order[DIMLIMIT_DIV4]) { + TensorIndex new_tidx = tidx; + for (int d = 0; d < DIMLIMIT; ++d) { + int src_dim = permute_order[div_4(d)][mod_4(d)]; + new_tidx.data[div_4(d)][mod_4(d)] = idx_at(tidx, src_dim); + } + tidx = new_tidx; +} + +// +// Index Conversions +// + +void contiguous_idx_to_tensor_idx( + const BufferMetadata meta, + uint contiguous_idx, + out TensorIndex tidx) { + initialize(tidx); + int dim = int_ndim(meta); + int i = 0; + + uint contiguous_strides[DIMLIMIT]; + contiguous_strides[0] = 1; + for (int d = 1; d < DIMLIMIT; ++d) { + contiguous_strides[d] = size_at(meta, d - 1) * contiguous_strides[d - 1]; + } + + for (int d = max(dim - 1, 0); d >= 0; d--) { + uint dim_stride = contiguous_strides[d]; + + tidx.data[div_4(d)][mod_4(d)] = contiguous_idx / dim_stride; + contiguous_idx = contiguous_idx % dim_stride; + } +} + +uint tensor_idx_to_contiguous_idx( + const BufferMetadata meta, + const TensorIndex tidx) { + uint contiguous_strides[DIMLIMIT]; + contiguous_strides[0] = 1; + for (int d = 1; d < DIMLIMIT; ++d) { + contiguous_strides[d] = size_at(meta, d - 1) * contiguous_strides[d - 1]; + } + + uint contig_idx = 0; + for (int d = 0; d < ndim(meta); ++d) { + contig_idx += contiguous_strides[d] * idx_at(tidx, d); + } + return contig_idx; +} + +void linear_idx_to_tensor_idx( + const BufferMetadata meta, + uint linear_idx, + out TensorIndex tidx) { + initialize(tidx); + int dim = int_ndim(meta); + int i = 0; + for (int d = max(dim - 1, 0); d >= 0; d--) { + uint dim_idx = dim_order_at(meta, d); + uint dim_stride = stride_at(meta, dim_idx); + + tidx.data[div_4(dim_idx)][mod_4(dim_idx)] = linear_idx / dim_stride; + linear_idx = linear_idx % dim_stride; + } +} + +uint tensor_idx_to_linear_idx( + const BufferMetadata meta, + const TensorIndex tidx) { + uint lin_idx = 0; + for (int d = 0; d < ndim(meta); ++d) { + lin_idx += stride_at(meta, d) * idx_at(tidx, d); + } + return lin_idx; +} + +void clamp_tensor_idx(const BufferMetadata meta, inout TensorIndex tidx) { + tidx.data[0] = min(tidx.data[0], meta.sizes[0] - 1); + tidx.data[1] = min(tidx.data[1], meta.sizes[1] - 1); +} + +// +// Debug utilities +// + +#ifdef DEBUG_MODE + +void printTensorIndex(const TensorIndex tidx) { + debugPrintfEXT( + "TensorIndex: tidx=[%u %u %u %u %u %u %u %u]\\n", + tidx.data[0][0], tidx.data[0][1], tidx.data[0][2], tidx.data[0][3], + tidx.data[1][0], tidx.data[1][1], tidx.data[1][2], tidx.data[1][3] + ); +} + +void printBufferMetadata(const BufferMetadata meta) { + debugPrintfEXT( + "BufferMetadata: ndim=%u numel=%u\\n sizes=[%u %u %u %u %u %u %u %u]\\n dim_order=[%u %u %u %u %u %u %u %u]\\n strides=[%u %u %u %u %u %u %u %u]\\n", + meta.ndim_numel[0], meta.ndim_numel[1], + meta.sizes[0][0], meta.sizes[0][1], meta.sizes[0][2], meta.sizes[0][3], + meta.sizes[1][1], meta.sizes[1][1], meta.sizes[1][2], meta.sizes[1][3], + meta.dim_order[0][0], meta.dim_order[0][1], + meta.dim_order[0][2], meta.dim_order[0][3], + meta.dim_order[1][0], meta.dim_order[1][1], + meta.dim_order[1][2], meta.dim_order[1][3], + meta.strides[0][0], meta.strides[0][1], + meta.strides[0][2], meta.strides[0][3], + meta.strides[1][1], meta.strides[1][1], + meta.strides[1][2], meta.strides[1][3] + ); +} + +#endif + +#endif // INDEXING_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h index 72650bb7040..fdb6f514a3e 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h +++ b/backends/vulkan/runtime/graph/ops/glsl/indexing_utils.h @@ -68,6 +68,8 @@ */ #define mod4(x) ((x) & 3) +#define ALIGN_DOWN_4(x) ((x) & ~3) + #define ALIGN_UP_4(x) (((x) + 3) & ~3) #define DIV_UP_8(x) (((x) + 7) >> 3) @@ -110,6 +112,10 @@ ivec4 tidx_to_4bufi( return base_i + ivec4(0, 1, 2, 3) * strides[packed_dim]; } +/* + * Given a buffer index to a contiguous tensor and the tensor's sizes, return + * the tensor index that corresponds to the buffer index. + */ ivec4 nchwi_to_tidx(const int nchwi, const ivec4 sizes) { const int nchwi_div_x = nchwi / sizes.x; const int nchwi_div_y = nchwi_div_x / sizes.y; diff --git a/backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.glsl b/backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.glsl deleted file mode 100644 index 8028362c3e5..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/kv_cache_update.glsl +++ /dev/null @@ -1,80 +0,0 @@ -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${buffer_scalar_type(DTYPE)} - -${define_active_storage_type(STORAGE)} -${define_required_extensions(DTYPE)} - -layout(std430) buffer; - -#include "indexing_utils.h" - -${layout_declare_tensor(B, "w", "cache", DTYPE, STORAGE)} -${layout_declare_tensor(B, "r", "projected", DTYPE, STORAGE)} -$if STORAGE == "buffer": - ${layout_declare_ubo(B, "int", "projected_numel")} - ${layout_declare_ubo(B, "ivec4", "cache_strides")} - ${layout_declare_ubo(B, "int", "input_pos")} -$else: - ${layout_declare_ubo(B, "ivec3", "projected_limits")} - ${layout_declare_ubo(B, "int", "input_pos")} - - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -/* - * t_cache will have sizes of (max_batch_size, max_seq_len, n_heads, head_dim). - * t_projected will have sizes of (batch_size, seq_len, n_heads, head_dim). - * - * The cache update inserts the values of t_projected into t_cache at the index - * specified by input_pos at the seq_len dimension. It is equivalent to calling - - * t_cache = t_cache.slice_scatter( - * t_projected, dim=1, start=input_pos, end=input_pos+seq_len) - * - * Note that this shader is implemented assuming that max_batch_size is 1. - */ - -#ifdef USING_BUFFER - -/*************************** - ** Buffer Implementation ** - ***************************/ - -void main() { - int projected_bufi = int(gl_GlobalInvocationID.x); - // Bump cache index forward by input_pos elements along the seq_len dimension. - // cache_strides contains the strides of the cache tensor. - int cache_bufi = input_pos * cache_strides.z + projected_bufi; - if (projected_bufi >= projected_numel) { - return; - } - cache[cache_bufi] = projected[projected_bufi]; -} - -#else - -/**************************** - ** Texture Implementation ** - ****************************/ - -// Note that this shader assumes the that tensors are width packed, i.e. -// packed_dim = 0 -void main() { - const ivec3 projected_pos = ivec3(gl_GlobalInvocationID); - - if (any(greaterThanEqual(projected_pos, projected_limits))) { - return; - } - - const ivec3 cache_pos = ivec3( - projected_pos.x, - projected_pos.y, - projected_pos.z + input_pos); - - write_texel(cache, cache_pos, load_texel(projected, projected_pos)); -} - -#endif // USING_BUFFER diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_common.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_common.glslh new file mode 100644 index 00000000000..da326b26e93 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_common.glslh @@ -0,0 +1,47 @@ +/* + * 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. + */ + +/* + * Defines common functions and structs to be used across matrix multiplication + * operators. + */ + +#ifndef LINEAR_COMMON_GLSLH +#define LINEAR_COMMON_GLSLH + +#include "common.glslh" + +int sign_extend_8bit(const int val) { + if ((val & 0x80) != 0) { + return val | (~0xFF); + } + return val; +} + +int extract_8bit_from_packed_int_le(const int packed, const int i) { + // account for little endian + int byte = sign_extend_8bit(packed >> (8 * i) & 0xFF); + return byte; +} + +// Extract a 4-bit value from a packed int (little endian) +// It is assumed that the 4-bit value is in the range [0, 15] +int extract_4bit_from_packed_int_le(const int packed, const int col) { + // Extract the 4-bit value from the 8-bit value + int val = packed >> (4 * col) & 0xF; + return val; +} + +// Convenience overload for packed uint +int extract_4bit_from_packed_uint_le(const uint packed, const int col) { + // Extract the 4-bit value from the 8-bit value + int val = int(packed >> (4 * col)) & 0xF; + return val; +} + +#endif // LINEAR_COMMON_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl new file mode 100644 index 00000000000..d918e4c9c7f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.glsl @@ -0,0 +1,144 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if PACKED_INT8_INPUT_STORAGE == "buffer": + #define PACKED_INT8_INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_N8 ${TILE_N8} + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N8 * 2} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N8 * 8} + +${define_required_extensions(DTYPE)} +${define_required_extensions("int8")} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", PACKED_INT8_INPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_sums", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} +${layout_declare_tensor(B, "r", "t_packed_int4_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)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +#include "linear_fp_input_tile_load.glslh" +#include "linear_int8_input_tile_load.glslh" +#include "linear_int8_input_scales_zps_load.glslh" +#include "linear_int4_weight_tile_load.glslh" +#include "linear_int_weight_sums_load.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_int8_input_sums_load.glslh" +#include "linear_fp_output_tile_int8_int8_compute.glslh" +#include "linear_fp_output_tile_int8_int4_compute.glslh" +#include "linear_fp_output_tile_fp_int4_compute.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "linear_fp_output_tile_store.glslh" +#include "linear_fp_bias_load.glslh" + +void main() { + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = out_tile_x * TILE_N; + const int m = out_tile_y * TILE_M; + + const int n8 = div_8(n); + const int n4 = div_4(n); + const int m4 = div_4(m); + + if (n >= output_sizes.x || m >= output_sizes.y) { + return; + } + + const int M = input_sizes.y; + const int K4 = div_up_4(input_sizes.x); + const int M4 = div_up_4(M); + const int N4 = div_up_4(output_sizes.x); + const int N8 = div_up_8(output_sizes.x); + + FPOutTile out_tile; + initialize(out_tile); + + Int32Accum out_accum; + initialize(out_accum); + + Int8InputTile int8_in_tile; + Int4WeightTile int4_weight_tile; + + Int8InputScales input_scales; + Int8InputZeroPoints input_zps; + load_int8_input_scales_and_zps(input_scales, input_zps, m4); + + FPPerOutChannelParams weight_scales_tile; + IntPerOutChannelParams weight_sums_tile; + + IntPerInChannelParams int8_input_sums_tile; + + const int num_groups = K4 / K4_per_group; + const int group_size = mul_4(K4_per_group); + + for (int group_i = 0; group_i < num_groups; ++group_i) { + // Reset int accumulator + initialize(out_accum); + for (int k4_inner = 0; k4_inner < K4_per_group; k4_inner++) { + const int k4 = group_i * K4_per_group + k4_inner; + + load_int8_input_tile(int8_in_tile, k4, m4, K4); + load_int4_weight_tile(int4_weight_tile, k4, n8, K4); + + int_accumulate_with_int4_weight( + out_accum, int8_in_tile, int4_weight_tile); + } + + load_weight_scales_tile_for_group(weight_scales_tile, n4, group_i, N4); + load_weight_sums_tile_for_group(weight_sums_tile, n4, group_i, N4); + load_int8_input_sums_tile_for_group(int8_input_sums_tile, m4, group_i, M4); + + accumulate_out_tile_with_int_accum_from_int4_weights( + out_tile, + out_accum, + int8_input_sums_tile, + input_scales, + input_zps, + weight_sums_tile, + weight_scales_tile, + group_size); + } + + write_output_tile_with_checks(out_tile, n4, m, N4, M); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.yaml new file mode 100644 index 00000000000..a252055ed40 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dq8ca_q4gsw_tiled.yaml @@ -0,0 +1,29 @@ +# 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. + +linear_dq8ca_q4gsw_tiled: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + PACKED_INT8_INPUT_STORAGE: buffer + TILE_M4: 1 + TILE_K4: 1 + TILE_N8: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: linear_dq8ca_q4gsw_tiled_texture3d_texture2d + - NAME: linear_dq8ca_q4gsw_tiled_texture3d_buffer + WEIGHT_STORAGE: buffer + - NAME: linear_dq8ca_q4gsw_tiled_buffer_texture2d + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + - NAME: linear_dq8ca_q4gsw_tiled_buffer_buffer + IO_STORAGE: buffer + WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_bias_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_bias_load.glslh new file mode 100644 index 00000000000..c1ab8999420 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_bias_load.glslh @@ -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. + */ + +#ifndef LINEAR_FP_BIAS_LOAD_GLSLH +#define LINEAR_FP_BIAS_LOAD_GLSLH + +#include "linear_fp_per_out_channel_params.glslh" + +VEC4_T load_bias_x4(const int n4) { + return t_bias[n4]; +} + +void load_bias_tile(out FPPerOutChannelParams bias, const int n4_start) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + bias.data[n4] = load_bias_x4(n4_start + n4); + } +} + +#endif // LINEAR_FP_BIAS_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile.glslh new file mode 100644 index 00000000000..72b5bdb812e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile.glslh @@ -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. + */ + +#ifndef LINEAR_FP_INPUT_TILE_GLSLH +#define LINEAR_FP_INPUT_TILE_GLSLH + +/* + * Macro Settings: + * - TILE_M + * - TILE_K4 + */ + +#extension GL_EXT_control_flow_attributes : require + +struct FPInputTile { + VEC4_T data[TILE_M][TILE_K4]; +}; + +#ifdef DEBUG_MODE + +void printFPInputTile(const FPInputTile in_tile) { + debugPrintfEXT("input_tile: \\n"); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + debugPrintfEXT( + " %f, %f, %f, %f, \\n", + in_tile.data[m][k4].x, + in_tile.data[m][k4].y, + in_tile.data[m][k4].z, + in_tile.data[m][k4].w); + } + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_FP_INPUT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile_load.glslh new file mode 100644 index 00000000000..358379b3efd --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_input_tile_load.glslh @@ -0,0 +1,62 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_input + * + * Macro Settings: + * - INPUT_BUFFER + */ + +#ifndef LINEAR_FP_INPUT_TILE_LOAD_GLSLH +#define LINEAR_FP_INPUT_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_input_tile.glslh" + +VEC4_T load_input_x4(const int k4, const int m, const int ntexels_k) { +#ifdef INPUT_BUFFER + return t_input[(m * ntexels_k) + k4]; +#else + return texelFetch(t_input, ivec3(k4, m, 0), 0); +#endif +} + +void load_input_tile_no_checks( + out FPInputTile in_tile, + const int k4_start, + const int m_start, + const int K4, + const int M) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + in_tile.data[m][k4] = load_input_x4(k4_start + k4, m_start + m, K4); + } + } +} + +void load_input_tile_with_checks( + out FPInputTile in_tile, + const int k4_start, + const int m_start, + const int K4, + const int M) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + if (m_start + m < M && k4_start + k4 < K4) { + in_tile.data[m][k4] = load_input_x4(k4_start + k4, m_start + m, K4); + } else { + in_tile.data[m][k4] = VEC4_T(0.0); + } + } + } +} + +#endif // LINEAR_FP_INPUT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile.glslh new file mode 100644 index 00000000000..ca466447084 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile.glslh @@ -0,0 +1,51 @@ +/* + * 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. + */ + +/* + * Macro Settings: + * - TILE_M + * - TILE_N4 + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +struct FPOutTile { + VEC4_T data[TILE_M][TILE_N4]; +}; + +void initialize(out FPOutTile out_tile) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + out_tile.data[m][n4] = VEC4_T(0); + } + } +} + +#ifdef DEBUG_MODE + +void printFPOutTile(const FPOutTile tile) { + debugPrintfEXT("output_tile: \\n"); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT( + " %f, %f, %f, %f,", + tile.data[m][n4].x, + tile.data[m][n4].y, + tile.data[m][n4].z, + tile.data[m][n4].w); + } + debugPrintfEXT("\\n"); + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_FP_OUTPUT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_compute.glslh new file mode 100644 index 00000000000..01b3c762e39 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_compute.glslh @@ -0,0 +1,86 @@ +/* + * 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. + */ + +/* + * Defines functions to compute a FPOutTile using fp input and weight tiles. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_FP_COMPUTE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_FP_COMPUTE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_common.glslh" +#include "linear_fp_input_tile.glslh" +#include "linear_fp_output_tile.glslh" +#include "linear_fp_per_out_channel_params.glslh" +#include "linear_fp_weight_tile.glslh" + +void fp_accumulate_with_fp_weight( + inout FPOutTile accum, + FPInputTile in_tile, + FPWeightTile w_tile) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][0]), + w_tile.data[mul_4(k4)][n4], + accum.data[m][n4]); + + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][1]), + w_tile.data[mul_4(k4) + 1][n4], + accum.data[m][n4]); + + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][2]), + w_tile.data[mul_4(k4) + 2][n4], + accum.data[m][n4]); + + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][3]), + w_tile.data[mul_4(k4) + 3][n4], + accum.data[m][n4]); + } + } + } +} + +void apply_weight_scales( + inout FPOutTile tile, + const FPPerOutChannelParams scales) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + tile.data[m][n4] = tile.data[m][n4] * scales.data[n4]; + } + } +} + +void apply_weight_scales_and_biases( + inout FPOutTile tile, + const FPPerOutChannelParams scales, + const FPPerOutChannelParams bias) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + tile.data[m][n4] = tile.data[m][n4] * scales.data[n4] + bias.data[n4]; + } + } +} + +void accumulate_out_tile_with_out_tile( + inout FPOutTile accum, + const FPOutTile other) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + accum.data[m][n4] = accum.data[m][n4] + other.data[m][n4]; + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_FP_COMPUTE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int4_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int4_compute.glslh new file mode 100644 index 00000000000..1f99b6672c0 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int4_compute.glslh @@ -0,0 +1,90 @@ +/* + * 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. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_FP_INT4_COMPUTE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_FP_INT4_COMPUTE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_common.glslh" +#include "linear_fp_input_tile.glslh" +#include "linear_fp_output_tile.glslh" +#include "linear_fp_per_out_channel_params.glslh" +#include "linear_int4_weight_tile.glslh" + +VEC4_T unpack_packed_4xint4(const int int8x4, const int n4_group) { + return VEC4_T( + extract_4bit_from_packed_int_le(int8x4, n4_group + 0), + extract_4bit_from_packed_int_le(int8x4, n4_group + 2), + extract_4bit_from_packed_int_le(int8x4, n4_group + 4), + extract_4bit_from_packed_int_le(int8x4, n4_group + 6)); +} + +T extract_4bit_from_weight_block( + const ivec4 block, + const int col, + const int row) { + return T(((block[row] >> (4 * col)) & 0xF) - 8); +} + +void fp_accumulate_with_int4_weight( + inout FPOutTile accum, + FPInputTile in_tile, + Int4WeightTile w_tile, + FPPerOutChannelParams scales_tile, + FPPerOutChannelParams zeros_tile) { + // Accum tile is indexed as accum[m][n4][n4i] + // -> gives fp accumulator for output tile element at (x = n, y = m) + // Input tile is indexed as in_tile.data[m][k4] + // -> gives vec4 containing the fp inputs at index + // (k, m), (k + 1, m), (k + 2, m), (k + 3, m) + // Weight tile is indexed as w_tile.data[k4][n8][n4i] + // -> gives packed integer containing the 8x 4-bit quantized values at index + // (n, k), (n, k + 1), (n, k + 2), (n, k + 3), + // (n + 4, k), (n + 4, k + 1), (n + 4, k + 2), (n + 4, k + 3) + VEC4_T weight_texels[2]; + [[unroll]] for (int n8 = 0; n8 < TILE_N8; ++n8) { + const int n4 = mul_2(n8); + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + [[unroll]] for (int k4i = 0; k4i < 4; ++k4i) { + const int base_col_1 = mul_2(k4i); + const int base_col_2 = base_col_1 + 1; + weight_texels[0] = VEC4_T( + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_1, 0), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_1, 1), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_1, 2), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_1, 3)); + weight_texels[1] = VEC4_T( + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_2, 0), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_2, 1), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_2, 2), + extract_4bit_from_weight_block(w_tile.data[k4][n8], base_col_2, 3)); + + weight_texels[0] = + fma(weight_texels[0], scales_tile.data[n4], zeros_tile.data[n4]); + weight_texels[1] = + fma(weight_texels[1], + scales_tile.data[n4 + 1], + zeros_tile.data[n4 + 1]); + + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][k4i]), + weight_texels[0], + accum.data[m][n4]); + accum.data[m][n4 + 1] = + fma(VEC4_T(in_tile.data[m][k4][k4i]), + weight_texels[1], + accum.data[m][n4 + 1]); + } + } + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_FP_INT4_COMPUTE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int8_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int8_compute.glslh new file mode 100644 index 00000000000..cb0ded2f9ec --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_fp_int8_compute.glslh @@ -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. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_FP_INT8_COMPUTE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_FP_INT8_COMPUTE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_common.glslh" +#include "linear_fp_input_tile.glslh" +#include "linear_fp_output_tile.glslh" +#include "linear_int8_weight_tile.glslh" + +VEC4_T unpack_packed_4xint8(int int8x4) { + return VEC4_T( + extract_8bit_from_packed_int_le(int8x4, 0), + extract_8bit_from_packed_int_le(int8x4, 1), + extract_8bit_from_packed_int_le(int8x4, 2), + extract_8bit_from_packed_int_le(int8x4, 3)); +} + +void fp_accumulate_with_int8_weight( + inout FPOutTile accum, + FPInputTile in_tile, + Int8WeightTile w_tile) { + // Accum tile is indexed as accum[m][n4][n4i] + // -> gives fp accumulator for output tile element at (x = n, y = m) + // Input tile is indexed as in_tile.data[m][k4] + // -> gives vec4 containing the fp inputs at index + // (k, m), (k + 1, m), (k + 2, m), (k + 3, m) + // Weight tile is indexed as w_tile.data[k4][n4][n4i] + // -> gives packed integer containing the 4x 8-bit quantized values at index + // (n, k), (n, k + 1), (n, k + 2), (n, k + 3) + VEC4_T weight_texel; + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + [[unroll]] for (int k4i = 0; k4i < 4; ++k4i) { + // Unpack one column of weights + weight_texel = VEC4_T( + extract_8bit_from_packed_int_le(w_tile.data[k4][n4][0], k4i), + extract_8bit_from_packed_int_le(w_tile.data[k4][n4][1], k4i), + extract_8bit_from_packed_int_le(w_tile.data[k4][n4][2], k4i), + extract_8bit_from_packed_int_le(w_tile.data[k4][n4][3], k4i)); + + for (int m = 0; m < TILE_M; ++m) { + accum.data[m][n4] = + fma(VEC4_T(in_tile.data[m][k4][k4i]), + weight_texel, + accum.data[m][n4]); + } + } + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_FP_INT8_COMPUTE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int4_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int4_compute.glslh new file mode 100644 index 00000000000..0a11ed6f482 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int4_compute.glslh @@ -0,0 +1,86 @@ +/* + * 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. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_INT8_INT4_COMPUTE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_INT8_INT4_COMPUTE_GLSLH + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_integer_dot_product : require + +#include "linear_int4_weight_tile.glslh" +#include "linear_int8_input_tile.glslh" +#include "linear_int_accumulator.glslh" + +void int_accumulate_with_int4_weight( + inout Int32Accum accum, + Int8InputTile in_tile, + Int4WeightTile w_tile) { + // Accum tile is indexed as accum[m][n4][n4i] + // -> gives integer accumulator for output tile element at (x = n, y = m) + // Input tile is indexed as in_tile.data[m4][k4][m4i] + // -> gives packed integer containing the 4x 8-bit quantized values at index + // (k, m), (k + 1, m), (k + 2, m), (k + 3, m) + // Weight tile is indexed as w_tile.data[k4][n8][n4i] + // -> gives packed integer containing the 8x 4-bit quantized values + [[unroll]] for (int n8 = 0; n8 < TILE_N8; ++n8) { + const int n4_base = mul_2(n8); + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + ivec4 unpacked_block_1 = w_tile.data[k4][n8] & 0x0F0F0F0F; + ivec4 unpacked_block_2 = (w_tile.data[k4][n8] >> 4) & 0x0F0F0F0F; + + [[unroll]] for (int n4i = 0; n4i < 4; ++n4i) { + // Accumulate unpacked_block_1[n4i] and unpacked_block_2[n4i] with + // each row of the input tile + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + const int m4 = div_4(m); + const int m4i = mod_4(m); + accum.data[m][n4_base][n4i] = dotPacked4x8AccSatEXT( + in_tile.data[m4][k4][m4i], + unpacked_block_1[n4i], + accum.data[m][n4_base][n4i]); + accum.data[m][n4_base + 1][n4i] = dotPacked4x8AccSatEXT( + in_tile.data[m4][k4][m4i], + unpacked_block_2[n4i], + accum.data[m][n4_base + 1][n4i]); + } + } + } + } +} + +void accumulate_out_tile_with_int_accum_from_int4_weights( + inout FPOutTile out_tile, + const Int32Accum accum, + const IntPerInChannelParams input_sums, + const Int8InputScales input_scales, + const Int8InputZeroPoints input_zps, + const IntPerOutChannelParams weight_sums, + const FPPerOutChannelParams weight_scales, + const int group_size) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + const int m4 = div_4(m); + const int m4i = mod_4(m); + + float input_scale_m = input_scales.data[m4][m4i]; + int input_zp_m = input_zps.data[m4][m4i]; + int input_sum_m = input_sums.data[m4][m4i]; + + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + ivec4 accum_adjusted = accum.data[m][n4] - + (weight_sums.data[n4] + 8 * group_size) * input_zp_m + + mul_8(group_size * input_zp_m - input_sum_m); + + out_tile.data[m][n4] = + fma(VEC4_T(accum_adjusted), + VEC4_T(input_scale_m * weight_scales.data[n4]), + out_tile.data[m][n4]); + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_INT8_INT4_COMPUTE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int8_compute.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int8_compute.glslh new file mode 100644 index 00000000000..ca25e406ac1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_int8_int8_compute.glslh @@ -0,0 +1,108 @@ +/* + * 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. + */ + +/* + * Defines functions to compute a FPOutTile using int8 input and weight tiles. + * + * Settings: + * - TILE_M: The number of rows in the output tile. + * - TILE_N4: The number of (groups of 4) columns in the output tile. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_INT8_INT8_COMPUTE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_INT8_INT8_COMPUTE_GLSLH + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_integer_dot_product : require + +#include "linear_common.glslh" +#include "linear_fp_output_tile.glslh" +#include "linear_fp_per_out_channel_params.glslh" +#include "linear_int8_input_scales_zps.glslh" +#include "linear_int8_input_tile.glslh" +#include "linear_int8_weight_tile.glslh" +#include "linear_int_accumulator.glslh" +#include "linear_int_per_in_channel_params.glslh" +#include "linear_int_per_out_channel_params.glslh" + +// Accumulate int8 input and weight tiles into integer accumulator tile +void int_accumulate_with_int8_weight( + inout Int32Accum accum, + Int8InputTile in_tile, + Int8WeightTile w_tile) { + // Accum tile is indexed as accum[m][n4][n4i] + // -> gives integer accumulator for output tile element at (x = n, y = m) + // Input tile is indexed as in_tile.data[m4][k4][m4i] + // -> gives packed integer containing the 4x 8-bit quantized values at index + // (k, m), (k + 1, m), (k + 2, m), (k + 3, m) + // Weight tile is indexed as w_tile.data[k4][n4][n4i] + // -> gives packed integer containing the 4x 8-bit quantized values at index + // (n, k), (n, k + 1), (n, k + 2), (n, k + 3) + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + const int m4 = div_4(m); + const int m4i = mod_4(m); + [[unroll]] for (int n = 0; n < TILE_N; ++n) { + const int n4 = div_4(n); + const int n4i = mod_4(n); + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + accum.data[m][n4][n4i] = dotPacked4x8AccSatEXT( + in_tile.data[m4][k4][m4i], + w_tile.data[k4][n4][n4i], + accum.data[m][n4][n4i]); + } + } + } +} + +void accumulate_out_tile_with_int_accum( + inout FPOutTile out_tile, + const Int32Accum accum, + const float input_q_scale, + const int input_q_zp, + const IntPerOutChannelParams weight_sums, + const FPPerOutChannelParams weight_scales) { + ivec4 input_zp_vec = ivec4(-input_q_zp); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + // Unfortunately fma doesn't work with ivec4. Prefer to preserve integer + // format for as long as possible to avoid precision loss. + ivec4 accum_adjusted = + input_zp_vec * weight_sums.data[n4] + accum.data[m][n4]; + out_tile.data[m][n4] = + fma(VEC4_T(accum_adjusted), + VEC4_T(input_q_scale * weight_scales.data[0]), + out_tile.data[m][n4]); + } + } +} + +// overload of the above but with bias +void accumulate_out_tile_with_int_accum( + inout FPOutTile out_tile, + const Int32Accum accum, + const float input_q_scale, + const int input_q_zp, + const IntPerOutChannelParams weight_sums, + const FPPerOutChannelParams weight_scales, + const FPPerOutChannelParams bias) { + ivec4 input_zp_vec = ivec4(-input_q_zp); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + // Apply scale and zero points to the int accumulator + ivec4 accum_adjusted = + input_zp_vec * weight_sums.data[n4] + accum.data[m][n4]; + out_tile.data[m][n4] = + fma(VEC4_T(accum_adjusted), + VEC4_T(input_q_scale * weight_scales.data[n4]), + out_tile.data[m][n4]); + out_tile.data[m][n4] += bias.data[n4]; + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_INT8_INT8_COMPUTE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_store.glslh new file mode 100644 index 00000000000..6fb399ff99b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_output_tile_store.glslh @@ -0,0 +1,81 @@ +/* + * 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. + */ + +/* + * Defines functions store a FpOutTile to output buffer/texture. + * + * Requires: + * - t_output to be declared in the shader layout + * + * Settings: + * - OUTPUT_BUFFER to indicate t_output is a vec4 buffer, otherwise texture + * storage is assumed. + */ + +#ifndef LINEAR_FP_OUTPUT_TILE_STORE_GLSLH +#define LINEAR_FP_OUTPUT_TILE_STORE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_output_tile.glslh" + +void write_output_x4( + const VEC4_T out_texel, + const int n4, + const int m, + const int N4) { +#ifdef OUTPUT_BUFFER + t_output[m * N4 + n4] = out_texel; +#else + imageStore(t_output, ivec3(n4, m, 0), out_texel); +#endif +} + +void write_output_tile( + const FPOutTile out_tile, + const int n4_start, + const int m_start, + const int N4) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + write_output_x4(out_tile.data[m][n4], n4_start + n4, m_start + m, N4); + } + } +} + +// To be used if M - m >= TILE_M && N4 - n4 >= TILE_N4 +void write_output_tile_no_checks( + const FPOutTile out_tile, + const int n4_start, + const int m_start, + const int N4, + const int M) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + write_output_x4(out_tile.data[m][n4], n4_start + n4, m_start + m, N4); + } + } +} + +// To be used if close to tensor boundaries +void write_output_tile_with_checks( + const FPOutTile out_tile, + const int n4_start, + const int m_start, + const int N4, + const int M) { + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + if (m_start + m < M && n4_start + n4 < N4) { + write_output_x4(out_tile.data[m][n4], n4_start + n4, m_start + m, N4); + } + } + } +} + +#endif // LINEAR_FP_OUTPUT_TILE_STORE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_per_out_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_per_out_channel_params.glslh new file mode 100644 index 00000000000..96fb3e9900b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_per_out_channel_params.glslh @@ -0,0 +1,38 @@ +/* + * 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. + */ + +#ifndef LINEAR_FP_PER_OUT_CHANNEL_PARAMS_GLSLH +#define LINEAR_FP_PER_OUT_CHANNEL_PARAMS_GLSLH + +#include "common.glslh" + +#extension GL_EXT_control_flow_attributes : require + +// Represents floating point parameter tensors where each element is associated +// with an output channel, such as weight scales, biases, etc. +struct FPPerOutChannelParams { + VEC4_T data[TILE_N4]; +}; + +#ifdef DEBUG_MODE + +void printFPPerOutChannelParams(const FPPerOutChannelParams params) { + debugPrintfEXT("per_out_channel_params: \\n"); + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT( + " %f, %f, %f, %f, \\n", + params.data[n4].x, + params.data[n4].y, + params.data[n4].z, + params.data[n4].w); + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_FP_PER_OUT_CHANNEL_PARAMS_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_scales_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_scales_load.glslh new file mode 100644 index 00000000000..6011a7c6201 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_scales_load.glslh @@ -0,0 +1,40 @@ +/* + * 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. + */ + +#ifndef LINEAR_FP_WEIGHT_SCALES_LOAD_GLSLH +#define LINEAR_FP_WEIGHT_SCALES_LOAD_GLSLH + +#include "linear_fp_per_out_channel_params.glslh" + +VEC4_T load_weight_scale_x4(const int n4) { + return t_weight_scales[n4]; +} + +VEC4_T load_scale_x4(const int n4, const int quant_group_idx, const int N4) { + return t_weight_scales[quant_group_idx * N4 + n4]; +} + +void load_weight_scales_tile( + out FPPerOutChannelParams scales, + const int n4_start) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + scales.data[n4] = load_weight_scale_x4(n4_start + n4); + } +} + +void load_weight_scales_tile_for_group( + out FPPerOutChannelParams scales, + const int n4_start, + const int quant_group_idx, + const int N4) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + scales.data[n4] = load_scale_x4(n4_start + n4, quant_group_idx, N4); + } +} + +#endif // LINEAR_FP_WEIGHT_SCALES_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_tile.glslh new file mode 100644 index 00000000000..5e010442540 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_fp_weight_tile.glslh @@ -0,0 +1,45 @@ +/* + * 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. + */ + +/* + * Macro Settings: + * - TILE_K + * - TILE_N4 + */ + +#ifndef LINEAR_FP_WEIGHT_TILE_GLSLH +#define LINEAR_FP_WEIGHT_TILE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "common.glslh" + +struct FPWeightTile { + VEC4_T data[TILE_K][TILE_N4]; +}; + +#ifdef DEBUG_MODE + +void printFPWeightTile(const FPWeightTile tile) { + debugPrintfEXT("weight_tile: \\n"); + [[unroll]] for (int k = 0; k < TILE_K; ++k) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT( + " %f, %f, %f, %f, ", + tile.data[k][n4].x, + tile.data[k][n4].y, + tile.data[k][n4].z, + tile.data[k][n4].w); + } + debugPrintfEXT("\\n"); + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_FP_WEIGHT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_block.glslh new file mode 100644 index 00000000000..38b805a7c49 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_block.glslh @@ -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. + */ + +#ifndef LINEAR_INT4_WEIGHT_BLOCK_GLSLH +#define LINEAR_INT4_WEIGHT_BLOCK_GLSLH + +/* + * Assumes the following variables are defined in shader layout + * - t_packed_int4_weight + * - t_int4_weight + * + * Macro Settings: + * - USING_BUFFER + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_common.glslh" + +// Represents source data for 2 8Kx4N block of the weight matrix read from the +// input buffer. Each int element contains 8 packed 4-bit values along the K +// dimension. Overall the data represents 8Kx8N block. +struct Int4Weight2xBlockSourceData { + uint data[8]; +}; + +// Represents data for a packed 4Kx8N block of the weight matrix to be written +// out to output texture/buffer. An individual block was originally a 4Kx8N +// block in the original weight tensor, and then the top and bottom halves are +// concatenated along the width dim. +struct Int4WeightBlockPacked { + ivec4 data; +}; + +void load_block_source_data_no_checks( + out Int4Weight2xBlockSourceData src_data, + const int k8, + const int n_start, + const int ntexels_K, + const int N) { + [[unroll]] for (int n = 0; n < 8; ++n) { + src_data.data[n] = t_int4_weight[(n_start + n) * ntexels_K + k8]; + } +} + +// To be used if K - k_start < 4 +void load_block_source_data_with_checks( + out Int4Weight2xBlockSourceData src_data, + const int k8, + const int n_start, + const int ntexels_K, + const int N) { + [[unroll]] for (int n = 0; n < 8; ++n) { + if (n_start + n < N) { + src_data.data[n] = t_int4_weight[(n_start + n) * ntexels_K + k8]; + } else { + // Equivalent to a row of zeros since int4 weights have an implicit zero + // point of -8. + src_data.data[n] = 0x88888888; + } + } +} + +int pack_8x4bit_signed_into_int( + const int val0, + const int val1, + const int val2, + const int val3, + const int val4, + const int val5, + const int val6, + const int val7) { + return int( + ((val7 & 0xF) << 28) | ((val6 & 0xF) << 24) | ((val5 & 0xF) << 20) | + ((val4 & 0xF) << 16) | ((val3 & 0xF) << 12) | ((val2 & 0xF) << 8) | + ((val1 & 0xF) << 4) | ((val0 & 0xF))); +} + +void create_packed_blocks( + out Int4WeightBlockPacked block1, + out Int4WeightBlockPacked block2, + const Int4Weight2xBlockSourceData src_data) { + [[unroll]] for (int row = 0; row < 4; ++row) { + const int row_idx_1 = row; + const int row_idx_2 = row + 4; + block1.data[row] = pack_8x4bit_signed_into_int( + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 0), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 0), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 1), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 1), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 2), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 2), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 3), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 3)); + + block2.data[row] = pack_8x4bit_signed_into_int( + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 4), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 4), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 5), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 5), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 6), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 6), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_1], 7), + extract_4bit_from_packed_uint_le(src_data.data[row_idx_2], 7)); + } +} + +void write_packed_block( + const Int4WeightBlockPacked block, + const int k4, + const int n8, + const int nblocks_K) { +#ifdef USING_BUFFER + t_packed_int4_weight[n8 * nblocks_K + k4] = block.data; +#else + imageStore(t_packed_int4_weight, ivec2(k4, n8), block.data); +#endif // USING_BUFFER +} + +#ifdef DEBUG_MODE + +void printInt4Weight2xBlockSourceData( + const Int4Weight2xBlockSourceData src_data) { + debugPrintfEXT("int4_weight_block_source_data: \\n"); + [[unroll]] for (int row = 0; row < 8; ++row) { + debugPrintfEXT("row %i (raw: %u): ", row, src_data.data[row]); + // Extract and print individual 4-bit values directly from packed int + [[unroll]] for (int col = 0; col < 8; ++col) { + int val_4bit = extract_4bit_from_packed_uint_le(src_data.data[row], col); + debugPrintfEXT("[%i] ", val_4bit); + } + debugPrintfEXT("\\n"); + } +} + +void printInt4WeightBlockPacked(const Int4WeightBlockPacked block) { + debugPrintfEXT("int4_weight_block_packed: \\n"); + // Print unpacked 4-bit values for each int in block.data + [[unroll]] for (int i = 0; i < 4; ++i) { + debugPrintfEXT("block.data[%i] 4-bit values: ", i); + [[unroll]] for (int col = 0; col < 8; ++col) { + int val_4bit = extract_4bit_from_packed_int_le(block.data[i], col); + debugPrintfEXT("[%i] ", val_4bit); + } + debugPrintfEXT("\\n"); + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT4_WEIGHT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile.glslh new file mode 100644 index 00000000000..c51259e1ae8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile.glslh @@ -0,0 +1,51 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT4_WEIGHT_TILE_GLSLH +#define LINEAR_INT4_WEIGHT_TILE_GLSLH + +#include "linear_common.glslh" +#include "linear_fp_weight_tile.glslh" +#include "linear_int8_weight_tile.glslh" + +/* + * Defines the Int4WeightTile struct, which is used to represent a tile of the + * quantized int4 weight matrix of a quantized matrix multiplication operation. + * + * Settings: + * - TILE_K4: number of (groups of 4) rows in the weight tile + * - TILE_N8: number of (groups of 8) columns in the weight tile + */ + +#extension GL_EXT_control_flow_attributes : require + +struct Int4WeightTile { + ivec4 data[TILE_K4][TILE_N8]; +}; + +#ifdef DEBUG_MODE + +void printInt4WeightTile(const Int4WeightTile block) { + debugPrintfEXT("int4_weight_tile: \\n"); + // Print unpacked 4-bit values for each int in block.data + [[unroll]] for (int i = 0; i < TILE_K; ++i) { + const int k4 = div_4(i); + const int k4i = mod_4(i); + debugPrintfEXT("block.data[%i] 4-bit values: ", i); + [[unroll]] for (int col = 0; col < TILE_N; ++col) { + int val_4bit = + extract_4bit_from_packed_int_le(block.data[k4][0][k4i], col); + debugPrintfEXT("[%i] ", val_4bit); + } + debugPrintfEXT("\\n"); + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT4_WEIGHT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile_load.glslh new file mode 100644 index 00000000000..965496d70d7 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int4_weight_tile_load.glslh @@ -0,0 +1,48 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT4_WEIGHT_TILE_LOAD_GLSLH +#define LINEAR_INT4_WEIGHT_TILE_LOAD_GLSLH + +/* + * Assumes the following variables are declared in shader layout: + * - t_packed_int4_weight + * + * Macro Settings: + * - WEIGHT_BUFFER + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_int4_weight_tile.glslh" + +ivec4 load_int4_weight_block( + const int block_x, + const int block_y, + const int nblocks_x) { +#ifdef WEIGHT_BUFFER + return t_packed_int4_weight[(block_y * nblocks_x) + block_x]; +#else + return texelFetch(t_packed_int4_weight, ivec2(block_x, block_y), 0); +#endif +} + +void load_int4_weight_tile( + out Int4WeightTile weight_tile, + const int block_x, + const int block_y, + const int nblocks_x) { + [[unroll]] for (int y = 0; y < TILE_K4; ++y) { + [[unroll]] for (int x = 0; x < TILE_N8; ++x) { + weight_tile.data[y][x] = + load_int4_weight_block(block_x + x, block_y + y, nblocks_x); + } + } +} + +#endif // LINEAR_INT4_WEIGHT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh new file mode 100644 index 00000000000..a6dbd7e78a2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_block.glslh @@ -0,0 +1,91 @@ +/* + * 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. + */ + +/* + * This file defines utilties to perform int8 quantization and block packing of + * matrix multiplation inputs. It also defines utilities to store packed block + * data to an output buffer or texture. + * + * Requires: + * - t_packed_int8_input to be defined in shader layout (output buffer/texture) + * + * Settings: + * - OUTPUT_BUFFER to indicate if output resource is a buffer. Otherwise texture + * is assumed. + */ + +#ifndef LINEAR_INT8_INPUT_BLOCK_GLSLH +#define LINEAR_INT8_INPUT_BLOCK_GLSLH + +#define TILE_M4 1 +#define TILE_K4 1 + +#define TILE_M 4 + +#include "linear_fp_input_tile.glslh" +#include "linear_int8_input_scales_zps.glslh" + +struct Int8InputBlock { + ivec4 data; +}; + +ivec4 quantize( + const VEC4_T val, + const float q_inv_scale, + const int q_zero_point) { + vec4 quantized = round(vec4(val) * q_inv_scale) + q_zero_point; + // hard-code 8 bit quantization range + return clamp(ivec4(quantized), -128, 127); +} + +int pack_into_int32(const ivec4 quant_vals) { + int packed = ((quant_vals[0] & 0xFF) << 0) | ((quant_vals[1] & 0xFF) << 8) | + ((quant_vals[2] & 0xFF) << 16) | ((quant_vals[3] & 0xFF) << 24); + + return packed; +} + +void quantize_and_pack( + out Int8InputBlock packed, + const FPInputTile in_block, + const float q_inv_scale, + const int q_zero_point) { + for (int row = 0; row < 4; ++row) { + ivec4 quantized_inputs = + quantize(in_block.data[row][0], q_inv_scale, q_zero_point); + packed.data[row] = pack_into_int32(quantized_inputs); + } +} + +void quantize_and_pack( + out Int8InputBlock packed, + const FPInputTile in_block, + const Int8InputScales scales, + const Int8InputZeroPoints zps) { + for (int m = 0; m < 4; ++m) { + const float q_inv_scale = 1.0 / float(scales.data[0][m]); + const int q_zero_point = zps.data[0][m]; + ivec4 quantized_inputs = + quantize(in_block.data[m][0], q_inv_scale, q_zero_point); + packed.data[m] = pack_into_int32(quantized_inputs); + } +} + +void write_block( + const Int8InputBlock block, + const int block_x, + const int block_y, + const int nblocks_x) { +#ifdef OUTPUT_BUFFER + t_packed_int8_input[block_y * nblocks_x + block_x] = block.data; +#else // OUTPUT_TEXTURE + imageStore(t_packed_int8_input, ivec3(block_x, block_y, 0), block.data); +#endif // OUTPUT_BUFFER +} + +#endif // LINEAR_INT8_INPUT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps.glslh new file mode 100644 index 00000000000..6c0bb0c2add --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps.glslh @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * Defines common functions and structs to be used across matrix multiplication + * operators. + */ + +#ifndef LINEAR_INT8_INPUT_SCALES_ZPS_GLSLH +#define LINEAR_INT8_INPUT_SCALES_ZPS_GLSLH + +#include "common.glslh" + +#extension GL_EXT_control_flow_attributes : require + +struct Int8InputScales { + VEC4_T data[TILE_M4]; +}; + +struct Int8InputZeroPoints { + ivec4 data[TILE_M4]; +}; + +#ifdef DEBUG_MODE + +void printInt8InputScales(const Int8InputScales scales) { + debugPrintfEXT("input_scales: \\n"); + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + debugPrintfEXT( + " %f, %f, %f, %f, ", + scales.data[m4].x, + scales.data[m4].y, + scales.data[m4].z, + scales.data[m4].w); + } + debugPrintfEXT("\\n"); +} + +void printInt8InputZeroPoints(const Int8InputZeroPoints zero_points) { + debugPrintfEXT("input_zero_points: \\n"); + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + debugPrintfEXT( + " %d, %d, %d, %d, ", + zero_points.data[m4].x, + zero_points.data[m4].y, + zero_points.data[m4].z, + zero_points.data[m4].w); + } + debugPrintfEXT("\\n"); +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT8_INPUT_SCALES_ZPS_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps_load.glslh new file mode 100644 index 00000000000..e1a570622c2 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_scales_zps_load.glslh @@ -0,0 +1,27 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT8_INPUT_SCALES_ZPS_LOAD_GLSLH +#define LINEAR_INT8_INPUT_SCALES_ZPS_LOAD_GLSLH + +#include "linear_int8_input_scales_zps.glslh" + +#extension GL_EXT_control_flow_attributes : require + +void load_int8_input_scales_and_zps( + out Int8InputScales scales, + out Int8InputZeroPoints zps, + const int m4_start) { + [[unroll]] for (int m4 = 0; m4 < TILE_M4; m4++) { + scales.data[m4] = + VEC4_T(texelFetch(t_int8_input_scales, ivec3(m4_start + m4, 0, 0), 0)); + zps.data[m4] = texelFetch(t_int8_input_zps, ivec3(m4_start + m4, 0, 0), 0); + } +} + +#endif // LINEAR_INT8_INPUT_SCALES_ZPS_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_load.glslh new file mode 100644 index 00000000000..4bb2fa12fcb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_load.glslh @@ -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. + */ + +#ifndef LINEAR_INT_INPUT_SUMS_LOAD_GLSLH +#define LINEAR_INT_INPUT_SUMS_LOAD_GLSLH + +#include "linear_int_per_in_channel_params.glslh" + +#extension GL_EXT_control_flow_attributes : require + +ivec4 load_int8_input_sum_x4( + const int m4, + const int quant_group_idx, + const int M4) { + return t_int8_input_sums[quant_group_idx * M4 + m4]; +} + +void load_int8_input_sums_tile_for_group( + out IntPerInChannelParams sums, + const int m4_start, + const int quant_group_idx, + const int M4) { + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + sums.data[m4] = load_int8_input_sum_x4(m4_start + m4, quant_group_idx, M4); + } +} + +#endif // LINEAR_INT_INPUT_SUMS_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_store.glslh new file mode 100644 index 00000000000..c9066be8e74 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_sums_store.glslh @@ -0,0 +1,26 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT_INPUT_SUMS_LOAD_GLSLH +#define LINEAR_INT_INPUT_SUMS_LOAD_GLSLH + +#include "linear_int_per_in_channel_params.glslh" + +#extension GL_EXT_control_flow_attributes : require + +void store_int8_input_sums_tile_for_group( + out IntPerInChannelParams sums, + const int m4_start, + const int quant_group_idx, + const int M4) { + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + t_int8_input_sums[quant_group_idx * M4 + m4] = sums.data[m4]; + } +} + +#endif // LINEAR_INT_INPUT_SUMS_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile.glslh new file mode 100644 index 00000000000..5d8f78bae7c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile.glslh @@ -0,0 +1,59 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/* + * Macro Settings: + * - TILE_M4 + * - TILE_K4 + */ + +#ifndef LINEAR_INT8_INPUT_TILE_GLSLH +#define LINEAR_INT8_INPUT_TILE_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +struct Int8InputTile { + ivec4 data[TILE_M4][TILE_K4]; +}; + +#ifdef DEBUG_MODE + +#include "linear_common.glslh" + +void printInt8InputTile(const Int8InputTile tile) { + debugPrintfEXT( + "Int8InputTile [TILE_M4=%d][TILE_K4=%d]:\\n", TILE_M4, TILE_K4); + + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + debugPrintfEXT(" tile[%d][%d] (ivec4): ", m4, k4); + + // Each ivec4 contains 4 packed integers, each integer contains 4 8-bit + // values + [[unroll]] for (int vec_idx = 0; vec_idx < 4; ++vec_idx) { + int packed_int = tile.data[m4][k4][vec_idx]; + debugPrintfEXT("packed_int[%d]=%d -> [", vec_idx, packed_int); + + // Extract 4 8-bit values from this packed integer + [[unroll]] for (int byte_idx = 0; byte_idx < 4; ++byte_idx) { + int val = extract_8bit_from_packed_int_le(packed_int, byte_idx); + if (byte_idx < 3) { + debugPrintfEXT("%d, ", val); + } else { + debugPrintfEXT("%d] ", val); + } + } + } + debugPrintfEXT("\\n"); + } + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT8_INPUT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile_load.glslh new file mode 100644 index 00000000000..fc7c3c9c181 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_input_tile_load.glslh @@ -0,0 +1,48 @@ +/* + * 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. + */ + +/* + * Assumes the following variables are defined in the shader layout: + * - t_packed_int8_input + * + * Macro Settings: + * - PACKED_INT8_INPUT_BUFFER + */ + +#ifndef LINEAR_INT8_INPUT_TILE_LOAD_GLSLH +#define LINEAR_INT8_INPUT_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_int8_input_tile.glslh" + +ivec4 load_int8_input_block( + const int block_x, + const int block_y, + const int nblocks_x) { +#ifdef PACKED_INT8_INPUT_BUFFER + return t_packed_int8_input[(block_y * nblocks_x) + block_x]; +#else + return texelFetch(t_packed_int8_input, ivec3(block_x, block_y, 0), 0); +#endif // PACKED_INT8_INPUT_BUFFER +} + +void load_int8_input_tile( + out Int8InputTile in_tile, + const int block_x, + const int block_y, + const int nblocks_x) { + [[unroll]] for (int y = 0; y < TILE_M4; ++y) { + [[unroll]] for (int x = 0; x < TILE_K4; ++x) { + in_tile.data[y][x] = + load_int8_input_block(block_x + x, block_y + y, nblocks_x); + } + } +} + +#endif // LINEAR_INT8_INPUT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_block.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_block.glslh new file mode 100644 index 00000000000..6fa51452057 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_block.glslh @@ -0,0 +1,83 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT8_WEIGHT_BLOCK_GLSLH +#define LINEAR_INT8_WEIGHT_BLOCK_GLSLH + +/* + * Assumes the following variables are defined in the shader layout + * - t_packed_int8_weight + * - t_int8_weight + * + * Macro Settings: + * - USING_BUFFER + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_common.glslh" + +// Represents data for a 4x4 block of the weight matrix read from the input +// buffer. +struct Int8WeightBlock { + ivec4 data; +}; + +void load_block_data_no_checks( + out Int8WeightBlock block, + const int k4, + const int n_start, + const int ntexels_K, + const int N) { + [[unroll]] for (int n = 0; n < 4; ++n) { + block.data[n] = t_int8_weight[(n_start + n) * ntexels_K + k4]; + } +} + +void load_block_data_with_checks( + out Int8WeightBlock block, + const int k4, + const int n_start, + const int ntexels_K, + const int N) { + [[unroll]] for (int n = 0; n < 4; ++n) { + if (n_start + n < N) { + block.data[n] = t_int8_weight[(n_start + n) * ntexels_K + k4]; + } else { + block.data[n] = 0; + } + } +} + +void write_weight_block( + const Int8WeightBlock block, + const int n4, + const int k4, + const int ntexels_N) { +#ifdef USING_BUFFER + t_packed_int8_weight[k4 * ntexels_N + n4] = block.data; +#else + imageStore(t_packed_int8_weight, ivec2(n4, k4), block.data); +#endif +} + +#ifdef DEBUG_MODE + +void printInt8WeightBlock(const Int8WeightBlockPacked block) { + debugPrintfEXT("int8_weight_block_packed: \\n"); + debugPrintfEXT( + "%i %i %i %i \\n", + block.data[0], + block.data[1], + block.data[2], + block.data[3]); +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT8_WEIGHT_BLOCK_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile.glslh new file mode 100644 index 00000000000..6b3d4d4419e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile.glslh @@ -0,0 +1,57 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT8_WEIGHT_TILE_GLSLH +#define LINEAR_INT8_WEIGHT_TILE_GLSLH + +/* + * Macro Settings: + * - TILE_K4 + * - TILE_N4 + */ + +#extension GL_EXT_control_flow_attributes : require + +struct Int8WeightTile { + ivec4 data[TILE_K4][TILE_N4]; +}; + +#ifdef DEBUG_MODE + +void printInt8WeightTile(const Int8WeightTile tile) { + debugPrintfEXT( + "Int8WeightTile [TILE_K4=%d][TILE_N4=%d]:\\n", TILE_K4, TILE_N4); + + [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT(" tile[%d][%d] (ivec4): ", m4, n4); + + // Each ivec4 contains 4 packed integers, each integer contains 4 8-bit + // values + [[unroll]] for (int vec_idx = 0; vec_idx < 4; ++vec_idx) { + int packed_int = tile.data[m4][n4][vec_idx]; + debugPrintfEXT("packed_int[%d]=%d -> [", vec_idx, packed_int); + + // Extract 4 8-bit values from this packed integer + [[unroll]] for (int byte_idx = 0; byte_idx < 4; ++byte_idx) { + int val = extract_8bit_from_packed_int_le(packed_int, byte_idx); + if (byte_idx < 3) { + debugPrintfEXT("%d, ", val); + } else { + debugPrintfEXT("%d] ", val); + } + } + } + debugPrintfEXT("\\n"); + } + } +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT8_WEIGHT_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile_load.glslh new file mode 100644 index 00000000000..5021fe24d6f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int8_weight_tile_load.glslh @@ -0,0 +1,48 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT8_WEIGHT_TILE_LOAD_GLSLH +#define LINEAR_INT8_WEIGHT_TILE_LOAD_GLSLH + +/* + * Assumes the following variables are defined in the shader layout: + * - t_packed_int8_weight + * + * Macro Settings: + * - WEIGHT_BUFFER + */ + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_int8_weight_tile.glslh" + +ivec4 load_int8_weight_block( + const int block_x, + const int block_y, + const int nblocks_x) { +#ifdef WEIGHT_BUFFER + return t_packed_int8_weight[(block_y * nblocks_x) + block_x]; +#else + return texelFetch(t_packed_int8_weight, ivec2(block_x, block_y), 0); +#endif // WEIGHT_BUFFER +} + +void load_int8_weight_tile( + out Int8WeightTile weight_tile, + const int block_x, + const int block_y, + const int nblocks_x) { + [[unroll]] for (int y = 0; y < TILE_K4; ++y) { + [[unroll]] for (int x = 0; x < TILE_N4; ++x) { + weight_tile.data[y][x] = + load_int8_weight_block(block_x + x, block_y + y, nblocks_x); + } + } +} + +#endif // LINEAR_INT8_WEIGHT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_accumulator.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int_accumulator.glslh new file mode 100644 index 00000000000..3bfc27e6f54 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int_accumulator.glslh @@ -0,0 +1,46 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT_ACCUM_GLSLH +#define LINEAR_INT_ACCUM_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +// Stores integer accumulators for an output tile. +struct Int32Accum { + ivec4 data[TILE_M][TILE_N4]; +}; + +void initialize(out Int32Accum out_accum) { + [[unroll]] for (int y = 0; y < TILE_M; ++y) { + [[unroll]] for (int x4 = 0; x4 < TILE_N4; ++x4) { + out_accum.data[y][x4] = ivec4(0); + } + } +} + +#ifdef DEBUG_MODE + +void printInt32Accum(const Int32Accum tile) { + debugPrintfEXT("int accum: \\n"); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT( + " %d, %d, %d, %d,", + tile.data[m][n4].x, + tile.data[m][n4].y, + tile.data[m][n4].z, + tile.data[m][n4].w); + } + debugPrintfEXT("\\n"); + } +} + +#endif + +#endif // LINEAR_INT_ACCUM_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_in_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_in_channel_params.glslh new file mode 100644 index 00000000000..06ac91ca92f --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_in_channel_params.glslh @@ -0,0 +1,22 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT_PER_IN_CHANNEL_PARAMS_GLSLH +#define LINEAR_INT_PER_IN_CHANNEL_PARAMS_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +struct IntPerInChannelParams { + ivec4 data[TILE_M4]; +}; + +void initialize(out IntPerInChannelParams params) { + [[unroll]] for (int i = 0; i < TILE_M4; i++) { params.data[i] = ivec4(0); } +} + +#endif // LINEAR_INT_PER_IN_CHANNEL_PARAMS_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_out_channel_params.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_out_channel_params.glslh new file mode 100644 index 00000000000..73a079d2479 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int_per_out_channel_params.glslh @@ -0,0 +1,39 @@ +/* + * 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. + */ + +#ifndef LINEAR_INT_PER_OUT_CHANNEL_PARAMS_GLSLH +#define LINEAR_INT_PER_OUT_CHANNEL_PARAMS_GLSLH + +#include "common.glslh" + +#extension GL_EXT_control_flow_attributes : require + +// Represents floating point parameter tensors where each element is associated +// with an output channel, such as weight sums +struct IntPerOutChannelParams { + ivec4 data[TILE_N4]; +}; + +#ifdef DEBUG_MODE + +void printIntPerOutChannelParams(const IntPerOutChannelParams params) { + debugPrintfEXT("per_out_channel_params: \\n"); + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + debugPrintfEXT( + " %d, %d, %d, %d, ", + params.data[n4].x, + params.data[n4].y, + params.data[n4].z, + params.data[n4].w); + } + debugPrintfEXT("\\n"); +} + +#endif // DEBUG_MODE + +#endif // LINEAR_INT_PER_OUT_CHANNEL_PARAMS_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_int_weight_sums_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/linear_int_weight_sums_load.glslh new file mode 100644 index 00000000000..2fbc13e1bdb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_int_weight_sums_load.glslh @@ -0,0 +1,43 @@ +/* + * 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. + */ + +#ifndef LINEAR_FP_WEIGHT_SUMS_LOAD_GLSLH +#define LINEAR_FP_WEIGHT_SUMS_LOAD_GLSLH + +#include "linear_int_per_out_channel_params.glslh" + +ivec4 load_weight_sum_x4(const int n4) { + return ivec4(t_weight_sums[n4]); +} + +ivec4 load_weight_sum_x4( + const int n4, + const int quant_group_idx, + const int N4) { + return t_weight_sums[quant_group_idx * N4 + n4]; +} + +void load_weight_sums_tile( + out IntPerOutChannelParams sums, + const int n4_start) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + sums.data[n4] = load_weight_sum_x4(n4_start + n4); + } +} + +void load_weight_sums_tile_for_group( + out IntPerOutChannelParams sums, + const int n4_start, + const int quant_group_idx, + const int N4) { + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + sums.data[n4] = load_weight_sum_x4(n4_start + n4, quant_group_idx, N4); + } +} + +#endif // LINEAR_FP_WEIGHT_SUMS_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl new file mode 100644 index 00000000000..72c77eec704 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.glsl @@ -0,0 +1,146 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_N8 ${TILE_N8} + +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N8 * 2} + +#define TILE_M ${TILE_M} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N8 * 8} + +#define WGS ${WGS} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +$if DYNAMIC_QUANT_VARIANT: + ${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", PACKED_INPUT_STORAGE, is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_int_input_sums", "int", "buffer", is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_input_scale", DTYPE, "texture3d")} + ${layout_declare_tensor(B, "r", "t_input_zp", "int", "texture3d")} + ${layout_declare_tensor(B, "r", "t_packed_int4_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)} +$else: + ${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} + ${layout_declare_tensor(B, "r", "t_packed_int4_weight", "int", WEIGHT_STORAGE, 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)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +#include "common.glslh" +#include "linear_fp_input_tile_load.glslh" +#include "linear_int4_weight_tile_load.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_fp_output_tile_fp_int4_compute.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "linear_fp_output_tile_store.glslh" +#include "linear_fp_bias_load.glslh" + +shared FPOutTile partial_sums[WGS]; + +void main() { + const int lid = int(gl_LocalInvocationID.z); + const int n8 = int(gl_GlobalInvocationID.x); + + // The output tensor will have a shape of [n, 1, 1, 1]. Each thread computes + // 8 output elements, so each thread will write to 8 elements starting at the + // tensor index (gid.x * 8, 0, 0, 0). + const int n = mul_8(n8); + const int n4 = mul_2(n8); + const int K4 = div_up_4(input_sizes.x); + const int N4 = div_up_4(output_sizes.x); + + const int group_size = mul_4(K4_per_group); + + if (n >= output_sizes.x) { + return; + } + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile in_tile; + Int4WeightTile int4_weight_tile; + + FPPerOutChannelParams weight_scales_tile; + FPPerOutChannelParams weight_zeros_tile; + weight_zeros_tile.data[0] = VEC4_T(0.0); + weight_zeros_tile.data[1] = VEC4_T(0.0); + + // initialize the group index to a value larger than the largest possible + int cur_group_idx = input_sizes.x; + + for (int k4 = lid; k4 < div_up_4(input_sizes.x); k4 += WGS) { + const int group_idx = k4 / K4_per_group; + + // Only update the scales/zeros if the current iteration is now working on a + // new quantization group. + if (group_idx != cur_group_idx) { + load_weight_scales_tile_for_group(weight_scales_tile, n4, group_idx, N4); + cur_group_idx = group_idx; + } + + load_input_tile_no_checks(in_tile, k4, 0, K4, 1); + load_int4_weight_tile(int4_weight_tile, k4, n8, K4); + + fp_accumulate_with_int4_weight( + out_tile, + in_tile, + int4_weight_tile, + weight_scales_tile, + weight_zeros_tile); + } + + partial_sums[lid] = out_tile; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result. + for (int i = WGS / 2; i > 0; i /= 2) { + if (lid < i) { + accumulate_out_tile_with_out_tile( + partial_sums[lid], partial_sums[lid + i]); + } + memoryBarrierShared(); + barrier(); + } + + // Only the first thread will write out result + if (lid == 0) { + out_tile = partial_sums[0]; + write_output_tile_with_checks(out_tile, n4, 0, N4, 1); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.yaml new file mode 100644 index 00000000000..2c5001fdd17 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_coop.yaml @@ -0,0 +1,43 @@ +# 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. + +linear_q4gsw_coop: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + PACKED_INPUT_STORAGE: buffer + TILE_M: 1 + TILE_K4: 1 + TILE_N8: 1 + WGS: 64 + DYNAMIC_QUANT_VARIANT: false + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: linear_q4gsw_coop_texture3d_texture2d + - NAME: linear_q4gsw_coop_texture3d_buffer + WEIGHT_STORAGE: buffer + - NAME: linear_q4gsw_coop_buffer_texture2d + IO_STORAGE: buffer + - NAME: linear_q4gsw_coop_buffer_buffer + IO_STORAGE: buffer + WEIGHT_STORAGE: buffer + - NAME: linear_dq8ca_q4gsw_coop_texture3d_texture2d + DYNAMIC_QUANT_VARIANT: true + - NAME: linear_dq8ca_q4gsw_coop_texture3d_buffer + WEIGHT_STORAGE: buffer + DYNAMIC_QUANT_VARIANT: true + - NAME: linear_dq8ca_q4gsw_coop_buffer_texture2d + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + DYNAMIC_QUANT_VARIANT: true + - NAME: linear_dq8ca_q4gsw_coop_buffer_buffer + IO_STORAGE: buffer + WEIGHT_STORAGE: buffer + DYNAMIC_QUANT_VARIANT: true diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl new file mode 100644 index 00000000000..0ad91643219 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.glsl @@ -0,0 +1,116 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_N8 ${TILE_N8} + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N8 * 2} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N8 * 8} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int4_weight", "int", WEIGHT_STORAGE, 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)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +#include "linear_fp_input_tile_load.glslh" +#include "linear_int4_weight_tile_load.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_fp_bias_load.glslh" +#include "linear_fp_output_tile_fp_int4_compute.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "linear_fp_output_tile_store.glslh" + +void main() { + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = out_tile_x * TILE_N; + const int m = out_tile_y * TILE_M; + + const int n8 = div_8(n); + const int n4 = div_4(n); + const int m4 = div_4(m); + + if (n >= output_sizes.x || m >= output_sizes.y) { + return; + } + + const int M = input_sizes.y; + const int K4 = div_up_4(input_sizes.x); + const int N4 = div_up_4(output_sizes.x); // number of texels in each row + const int N8 = div_up_8(output_sizes.x); // number of texels in each row + + bool should_print = (n8 == 0) && (m4 == 0); + should_print = false; + + // VEC4_T out_texels[4][2]; + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile in_tile; + Int4WeightTile int4_weight_tile; + + FPPerOutChannelParams weight_scales_tile; + FPPerOutChannelParams weight_zeros_tile; + weight_zeros_tile.data[0] = VEC4_T(0.0); + weight_zeros_tile.data[1] = VEC4_T(0.0); + + const int num_groups = K4 / K4_per_group; + + for (int group_i = 0; group_i < num_groups; ++group_i) { + // Load quantization scales and zeros for the current group + load_weight_scales_tile_for_group(weight_scales_tile, n4, group_i, N4); + + for (int k4_inner = 0; k4_inner < K4_per_group; k4_inner++) { + const int k4 = group_i * K4_per_group + k4_inner; + + load_input_tile_no_checks(in_tile, k4, m, K4, M); + load_int4_weight_tile(int4_weight_tile, k4, n8, K4); + + fp_accumulate_with_int4_weight( + out_tile, + in_tile, + int4_weight_tile, + weight_scales_tile, + weight_zeros_tile); + } + } + + write_output_tile_with_checks(out_tile, n4, m, N4, M); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.yaml new file mode 100644 index 00000000000..5a6bcb711bb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_tiled.yaml @@ -0,0 +1,28 @@ +# 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. + +linear_q4gsw_tiled: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + TILE_M4: 1 + TILE_K4: 1 + TILE_N8: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: linear_q4gsw_tiled_texture3d_texture2d + - NAME: linear_q4gsw_tiled_texture3d_buffer + WEIGHT_STORAGE: buffer + - NAME: linear_q4gsw_tiled_buffer_texture2d + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + - NAME: linear_q4gsw_tiled_buffer_buffer + IO_STORAGE: buffer + WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl new file mode 100644 index 00000000000..84be82f64e3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.glsl @@ -0,0 +1,113 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, 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_scales", DTYPE, "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_bias", DTYPE, "buffer", is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "uint", "apply_bias", "0")} + +#include "linear_fp_input_tile_load.glslh" +#include "linear_int8_weight_tile_load.glslh" +#include "linear_fp_weight_tile.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "linear_fp_output_tile_fp_int8_compute.glslh" +#include "linear_fp_output_tile_store.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_fp_bias_load.glslh" + +void main() { + // Each thread writes out a 4 wide x 4 high tile of output values + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = out_tile_x * TILE_N; + const int m = out_tile_y * TILE_M; + + const int n4 = div_4(n); + const int m4 = div_4(m); + + if (n >= output_sizes.x || m >= output_sizes.y) { + return; + } + + const int M = input_sizes.y; + const int K4 = div_up_4(input_sizes.x); + const int N4 = div_up_4(output_sizes.x); + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile in_tile; + Int8WeightTile int8_weight_tile; + + const bool dont_check_bounds = (M - m) >= TILE_M; + if (dont_check_bounds) { + for (int k4 = 0; k4 < K4; k4 += TILE_K4) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + fp_accumulate_with_int8_weight(out_tile, in_tile, int8_weight_tile); + } + } else { + for (int k4 = 0; k4 < K4; k4 += TILE_K4) { + load_input_tile_with_checks(in_tile, k4, m, K4, M); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + fp_accumulate_with_int8_weight(out_tile, in_tile, int8_weight_tile); + } + } + + FPPerOutChannelParams weight_scales_tile; + load_weight_scales_tile(weight_scales_tile, n4); + + if (apply_bias > 0) { + FPPerOutChannelParams bias_tile; + load_bias_tile(bias_tile, n4); + + apply_weight_scales_and_biases(out_tile, weight_scales_tile, bias_tile); + } + else { + apply_weight_scales(out_tile, weight_scales_tile); + } + + if (dont_check_bounds) { + write_output_tile_no_checks(out_tile, n4, m, N4, M); + } else { + write_output_tile_with_checks(out_tile, n4, m, N4, M); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.yaml new file mode 100644 index 00000000000..242c4471b3d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q8csw_tiled.yaml @@ -0,0 +1,28 @@ +# 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. + +linear_q8csw_tiled: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + WEIGHT_STORAGE: texture2d + TILE_M4: 1 + TILE_N4: 1 + TILE_K4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: linear_q8csw_tiled_texture3d_texture2d + - NAME: linear_q8csw_tiled_texture3d_buffer + WEIGHT_STORAGE: buffer + - NAME: linear_q8csw_tiled_buffer_texture2d + IO_STORAGE: buffer + WEIGHT_STORAGE: texture2d + - NAME: linear_q8csw_tiled_buffer_buffer + IO_STORAGE: buffer + WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl new file mode 100644 index 00000000000..9f7e00e3317 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.glsl @@ -0,0 +1,132 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, OUTPUT_STORAGE)} +#define T int + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if PACKED_INT8_INPUT_STORAGE == "buffer": + #define PACKED_INT8_INPUT_BUFFER +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_output", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", PACKED_INT8_INPUT_STORAGE, 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)} + +${layout_declare_spec_const(C, "int", "apply_bias", "0")} + +${layout_declare_ubo(B, "ivec4", "output_sizes")} +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(push_constant) uniform restrict Block { + float input_scale; + int input_zp; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "linear_int8_input_tile_load.glslh" +#include "linear_int8_weight_tile_load.glslh" +#include "linear_fp_output_tile_int8_int8_compute.glslh" +#include "linear_fp_output_tile_store.glslh" +#include "linear_fp_weight_scales_load.glslh" +#include "linear_int_weight_sums_load.glslh" +#include "linear_fp_bias_load.glslh" + +void main() { + // Each thread writes out a 4 wide x 4 high tile of output values + const int out_tile_x = int(gl_GlobalInvocationID.x); + const int out_tile_y = int(gl_GlobalInvocationID.y); + + const int n = out_tile_x * TILE_N; + const int m = out_tile_y * TILE_M; + + const int n4 = div_4(n); + const int m4 = div_4(m); + + if (n >= output_sizes.x || m >= output_sizes.y) { + return; + } + + const int M = output_sizes.y; + const int K4 = div_up_4(input_sizes.x); + const int N4 = div_up_4(output_sizes.x); + + Int32Accum out_accum; + initialize(out_accum); + + Int8InputTile int8_in_tile; + Int8WeightTile int8_weight_tile; + + // No checks are needed since packed input and weight are structured in units + // of 4x4 blocks. + for (int k4 = 0; k4 < K4; k4 += TILE_K4) { + load_int8_input_tile(int8_in_tile, k4, m4, K4); + load_int8_weight_tile(int8_weight_tile, n4, k4, N4); + + int_accumulate_with_int8_weight(out_accum, int8_in_tile, int8_weight_tile); + } + + FPPerOutChannelParams weight_scales_tile; + load_weight_scales_tile(weight_scales_tile, n4); + + IntPerOutChannelParams weight_sums_tile; + load_weight_sums_tile(weight_sums_tile, n4); + + FPOutTile out_tile; + initialize(out_tile); + + if (apply_bias > 0) { + FPPerOutChannelParams bias_tile; + load_bias_tile(bias_tile, n4); + + accumulate_out_tile_with_int_accum( + out_tile, + out_accum, + input_scale, + input_zp, + weight_sums_tile, + weight_scales_tile, + bias_tile); + } + else { + accumulate_out_tile_with_int_accum( + out_tile, + out_accum, + input_scale, + input_zp, + weight_sums_tile, + weight_scales_tile); + } + + if (M - m >= TILE_M) { + write_output_tile_no_checks(out_tile, n4, m, N4, M); + } else { + write_output_tile_with_checks(out_tile, n4, m, N4, M); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.yaml new file mode 100644 index 00000000000..aa1de3077fc --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q8ta_q8csw_tiled.yaml @@ -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. + +linear_q8ta_q8csw_tiled: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: texture3d + PACKED_INT8_INPUT_STORAGE: buffer + WEIGHT_STORAGE: texture2d + TILE_M4: 1 + TILE_N4: 1 + TILE_K4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: linear_q8ta_q8csw_tiled_texture3d_buffer_texture2d + - NAME: linear_q8ta_q8csw_tiled_buffer_buffer_texture2d + OUTPUT_STORAGE: buffer + PACKED_INT8_INPUT_STORAGE: buffer + WEIGHT_STORAGE: texture2d diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.glsl deleted file mode 100644 index 81d2a5f0aed..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.glsl +++ /dev/null @@ -1,156 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} -#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} - -#define WGS ${WGS} - -${define_required_extensions(DTYPE)} -${define_required_extensions("uint8")} - -layout(std430) buffer; - -#include "indexing_utils.h" - -${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qmat2", "uint", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qparams", DTYPE, "buffer", is_scalar_array=False)} - -layout(push_constant) uniform restrict Block { - ivec4 output_sizes; - ivec4 input_sizes; - ivec4 weight_sizes; -}; - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -layout(constant_id = 3) const int group_size = 64; - -shared VEC4_T partial_sums[WGS][2]; - -$if IO_STORAGE == "buffer": - #define BUFFER_IO -$if WEIGHT_STORAGE == "buffer": - #define BUFFER_WEIGHT - -#include "qlinear_utils.glslh" - -void main() { - const uint lid = gl_LocalInvocationID.x; - const uint n8 = gl_GlobalInvocationID.y; - // The output tensor will have a shape of [n, 1, 1, 1]. Each thread computes - // 8 output elements, so each thread will write to 8 elements starting at the - // tensor index (gid.x * 8, 0, 0, 0). - const uint n = MUL_8(n8); - const uint K4 = DIV_UP_4(input_sizes.x); - - if (n >= output_sizes.x) { - return; - } - - VEC4_T out_texels[2]; - out_texels[0] = VEC4_T(0); - out_texels[1] = VEC4_T(0); - - // initialize the group index to a value larger than the largest possible - uint cur_group_idx = input_sizes.x; - - // Each thread in the work group accumulates a partial result. - for (uint k4 = lid; k4 < DIV_UP_4(input_sizes.x); k4 += WGS) { - const uint k = MUL_4(k4); - const uint group_idx = k / group_size; - - VEC4_T scales[2]; - VEC4_T zeros[2]; - - // Only update the scales/zeros if the current iteration is now working on a - // new quantization group. - if (group_idx != cur_group_idx) { - // The qparams tensor contains the quantization scales and zeros, with - // shape [2, N, K / group_size, 1]. - // Loading a texel from the qparams tensor will return 2 scales and 2 - // zeros for 2 adjacent output channels. - uint qparams_bufi = group_idx * DIV_2(output_sizes.x) + DIV_2(n); - VEC4_T scales_zeros_texels[4]; - $for comp in range(4): - scales_zeros_texels[${comp}] = t_qparams[qparams_bufi++]; - - scales[0] = VEC4_T(scales_zeros_texels[0].xz, scales_zeros_texels[1].xz); - zeros[0] = VEC4_T(scales_zeros_texels[0].yw, scales_zeros_texels[1].yw); - - scales[1] = VEC4_T(scales_zeros_texels[2].xz, scales_zeros_texels[3].xz); - zeros[1] = VEC4_T(scales_zeros_texels[2].yw, scales_zeros_texels[3].yw); - - cur_group_idx = group_idx; - } - // The input tensor will have a shape of [K, 1, 1, 1]; in each iteration, - // load 4 elements starting from the tensor index (k, 0, 0, 0). - VEC4_T in_texel = load_input_texel_1d(k4); - // Extract each element of the in_texel into a separate vectorized variable; - // these are used to "broadcast" the input values in subsequent fma calls. - VEC4_T in_texel_val[4]; - $for comp in range(4): - in_texel_val[${comp}] = VEC4_T(in_texel[${comp}]); - - uvec4 packed_weight_block = load_transposed_weight_block(k4, n8, K4); - - VEC4_T weight_texels[2]; - $for comp in range(4): - { - weight_texels[0].x = extract_4bit_from_transposed_block(packed_weight_block, 0, ${comp}); - weight_texels[0].y = extract_4bit_from_transposed_block(packed_weight_block, 1, ${comp}); - weight_texels[0].z = extract_4bit_from_transposed_block(packed_weight_block, 2, ${comp}); - weight_texels[0].w = extract_4bit_from_transposed_block(packed_weight_block, 3, ${comp}); - - weight_texels[1].x = extract_4bit_from_transposed_block(packed_weight_block, 4, ${comp}); - weight_texels[1].y = extract_4bit_from_transposed_block(packed_weight_block, 5, ${comp}); - weight_texels[1].z = extract_4bit_from_transposed_block(packed_weight_block, 6, ${comp}); - weight_texels[1].w = extract_4bit_from_transposed_block(packed_weight_block, 7, ${comp}); - - weight_texels[0] = fma(weight_texels[0], scales[0], zeros[0]); - weight_texels[1] = fma(weight_texels[1], scales[1], zeros[1]); - - out_texels[0] = fma(in_texel_val[${comp}], weight_texels[0], out_texels[0]); - out_texels[1] = fma(in_texel_val[${comp}], weight_texels[1], out_texels[1]); - } - } - - partial_sums[lid][0] = out_texels[0]; - partial_sums[lid][1] = out_texels[1]; - - memoryBarrierShared(); - barrier(); - - // Tree reduction to compute the overall result. - for (int i = WGS / 2; i > 0; i /= 2) { - if (lid < i) { - partial_sums[lid][0] = partial_sums[lid][0] + partial_sums[lid + i][0]; - partial_sums[lid][1] = partial_sums[lid][1] + partial_sums[lid + i][1]; - } - memoryBarrierShared(); - barrier(); - } - - // Only the first thread will write out result - if (lid == 0) { - out_texels[0] = partial_sums[0][0]; - out_texels[1] = partial_sums[0][1]; - - uint n4 = DIV_4(n); - write_output_texel_1d(out_texels[0], n4); - if (n + 4 < output_sizes.x) { - write_output_texel_1d(out_texels[1], n4 + 1); - } - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.glsl deleted file mode 100644 index 97327ea5818..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.glsl +++ /dev/null @@ -1,130 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} -#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} - -${define_required_extensions(DTYPE)} - -layout(std430) buffer; - -#include "indexing_utils.h" - -${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_input", DTYPE, IO_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qmat2", "uint", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qparams", DTYPE, "buffer", is_scalar_array=False)} - -layout(push_constant) uniform restrict Block { - ivec4 output_sizes; - ivec4 input_sizes; - ivec4 weight_sizes; -}; - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -layout(constant_id = 3) const int group_size = 64; - -$if IO_STORAGE == "buffer": - #define BUFFER_IO -$if WEIGHT_STORAGE == "buffer": - #define BUFFER_WEIGHT - -#include "qlinear_utils.glslh" - -void main() { - // Each thread writes out a 8 wide x 4 high tile of output values - const uint n8 = gl_GlobalInvocationID.x; - const uint m4 = gl_GlobalInvocationID.y; - - const uint n = MUL_8(n8); // output col idx - const uint m = MUL_4(m4); // output row idx - const uint n4 = MUL_2(n8); // output col texel idx - - const uint group_num = input_sizes.x / group_size; - const uint group_ntexels = DIV_UP_4(group_size); - - if (n >= output_sizes.x || m >= output_sizes.y) { - return; - } - - const uint K4 = DIV_UP_4(input_sizes.x); - const uint N4 = DIV_UP_4(output_sizes.x); // number of texels in each row - - VEC4_T out_texels[4][2]; - // Initialize to 0 - $for row_i in range(4): - $for col_i in range(2): - out_texels[${row_i}][${col_i}] = VEC4_T(0.00); - - for (uint group_i = 0; group_i < group_num; ++group_i) { - // Load quantization scales and zeros for the current group - VEC4_T scales[2]; - VEC4_T zeros[2]; - { - uint qparams_bufi = group_i * DIV_2(output_sizes.x) + DIV_2(n); - - VEC4_T scales_zeros_texels[4]; - $for comp in range(4): - scales_zeros_texels[${comp}] = t_qparams[qparams_bufi++]; - - scales[0] = VEC4_T(scales_zeros_texels[0].xz, scales_zeros_texels[1].xz); - zeros[0] = VEC4_T(scales_zeros_texels[0].yw, scales_zeros_texels[1].yw); - - scales[1] = VEC4_T(scales_zeros_texels[2].xz, scales_zeros_texels[3].xz); - zeros[1] = VEC4_T(scales_zeros_texels[2].yw, scales_zeros_texels[3].yw); - } - - for (uint inner_k4 = 0; inner_k4 < group_ntexels; inner_k4++) { - const uint k4 = group_i * group_ntexels + inner_k4; - - // Load 4x4 block of the input tensor, with the top left corner of the - // block at (k, m) - VEC4_T in_texels[4]; - $for comp in range(4): - in_texels[${comp}] = load_input_texel_2d(k4, m + ${comp}, K4); - - uvec4 packed_weight_block = load_transposed_weight_block(k4, n8, K4); - - VEC4_T weight_texels[2]; - $for tile_k in range(4): - // Process weight row k + comp - { - // Weight columns n + 0, 1, 2, 3 - weight_texels[0].x = extract_4bit_from_transposed_block(packed_weight_block, 0, ${tile_k}); - weight_texels[0].y = extract_4bit_from_transposed_block(packed_weight_block, 1, ${tile_k}); - weight_texels[0].z = extract_4bit_from_transposed_block(packed_weight_block, 2, ${tile_k}); - weight_texels[0].w = extract_4bit_from_transposed_block(packed_weight_block, 3, ${tile_k}); - - // Weight colums n + 4, 5, 6, 7 - weight_texels[1].x = extract_4bit_from_transposed_block(packed_weight_block, 4, ${tile_k}); - weight_texels[1].y = extract_4bit_from_transposed_block(packed_weight_block, 5, ${tile_k}); - weight_texels[1].z = extract_4bit_from_transposed_block(packed_weight_block, 6, ${tile_k}); - weight_texels[1].w = extract_4bit_from_transposed_block(packed_weight_block, 7, ${tile_k}); - - weight_texels[0] = fma(weight_texels[0], scales[0], zeros[0]); - weight_texels[1] = fma(weight_texels[1], scales[1], zeros[1]); - - $for tile_m in range(4): - out_texels[${tile_m}][0] = fma(VEC4_T(in_texels[${tile_m}][${tile_k}]), weight_texels[0], out_texels[${tile_m}][0]); - out_texels[${tile_m}][1] = fma(VEC4_T(in_texels[${tile_m}][${tile_k}]), weight_texels[1], out_texels[${tile_m}][1]); - } - } - } - - for (uint row_i = 0; row_i < 4 && m + row_i < output_sizes.y; ++row_i) { - write_output_texel_2d(out_texels[row_i][0], n4, m + row_i, N4); - if (n + 4 < output_sizes.x) { - write_output_texel_2d(out_texels[row_i][1], n4 + 1, m + row_i, N4); - } - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.glsl deleted file mode 100644 index 174ea1cc9bb..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.glsl +++ /dev/null @@ -1,232 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${buffer_scalar_type(DTYPE)} -#define VEC4_T ${buffer_gvec_type(DTYPE, 4)} - -#define TILE_ROWS ${TILE_ROWS} - -#define NGROUPS 8 -#define NWORKERS 8 - -${define_required_extensions(DTYPE)} -$if IN_STORAGE == "buffer": - ${define_required_extensions("int8")} -$if WEIGHT_STORAGE == "buffer": - ${define_required_extensions("uint8")} - -#extension GL_EXT_control_flow_attributes : require - -layout(std430) buffer; - -${layout_declare_tensor(B, "w", "t_out", DTYPE, OUT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_mat1", "int8", IN_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qmat2", "uint8", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_scales", "float", PARAMS_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_zeros", "int", PARAMS_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_input_scale", "float", PARAMS_STORAGE, is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_input_zero_point", "int", PARAMS_STORAGE, is_scalar_array=True)} - -layout(push_constant) uniform restrict Block { - ivec4 out_sizes; - ivec4 mat1_sizes; - ivec4 qmat2_sizes; -}; - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -layout(constant_id = 3) const int group_size = 64; - -shared vec4 partial_results[NGROUPS][NWORKERS][TILE_ROWS][2]; - -/* - * This shader computes a linear operator between a quantized int8 input matrix - * x and a weights matrix that is quantized to 4 bits, producing a float output. - * - * This shader implements a co-operative algorithm to compute the output. The - * work group size is {NGROUP, 1, NWORKERS}, and each group of NWORKERS threads - * cooperative to compute TILE_ROWS * 2 output texels. Therefore, - * NGROUP * TILE_ROWS * 2 output texels are computed across one work group. - * - * The threads co-operate by each thread computing a partial reduction along the - * K dimension. To illustrate the computation, consider a scalar variant of the - * algorithm that computes the dot product of 2 vectors. Also assume that - * NWORKERS is 8. - * - * Thread 1 in each group will compute: - * (mat1[0] * mat2[0]) + (mat1[8] * mat2[8]) + (mat1[16] * mat2[16]) + ... - * - * Thread 2 in each group will compute: - * (mat1[1] * mat2[1]) + (mat2[9] * mat2[9]) + (mat1[17] * mat2[17]) + ... - * - * Thread 3 in each group will compute: - * (mat1[2] * mat2[2]) + (mat2[10] * mat2[10]) + (mat1[18] * mat2[18]) + ... - * - * The partial accumulations is structured such that memory accesses in each - * loop iteration can be coalesced. - * - * Then, at the end first thread in each group will accumulate the partial - * accumulations computed by each thread to obtain the final result. - * - * Note that this shader assumes that all tensors are width packed. - */ - -void main() { - const uint out_row = gl_GlobalInvocationID.y * TILE_ROWS; - const uint out_col = gl_GlobalInvocationID.x << 3; - const int out_col_texel_idx = int(gl_GlobalInvocationID.x) << 1; - - const uint gid = gl_LocalInvocationID.x; // group id - const uint wid = gl_LocalInvocationID.z; // worker id - - if (out_col >= out_sizes.x || out_row >= out_sizes.y) { - return; - } - - const int num_blocks = mat1_sizes.x / group_size; - - ivec4 mat1_quantized[TILE_ROWS]; - ivec4 qmat2_quantized[4][2]; - vec4 final_result[TILE_ROWS][2]; - - // Initialize accumulators - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - final_result[r][0] = vec4(0.0); - final_result[r][1] = vec4(0.0); - } - - vec4 scales[2]; - vec4 zeros[2]; - - $if WEIGHT_STORAGE == "buffer": - const int qmat2_stride = qmat2_sizes.x >> 2; - $if PARAMS_STORAGE == "buffer": - const int qparams_stride = out_sizes.x >> 2; - - for (int block_idx = 0; block_idx < num_blocks; ++block_idx) { - $if PARAMS_STORAGE == "buffer": - scales[0] = t_weight_scales[block_idx * qparams_stride + out_col_texel_idx]; - scales[1] = t_weight_scales[block_idx * qparams_stride + out_col_texel_idx + 1]; - - zeros[0] = vec4(t_weight_zeros[block_idx * qparams_stride + out_col_texel_idx]); - zeros[1] = vec4(t_weight_zeros[block_idx * qparams_stride + out_col_texel_idx + 1]); - $else: - scales[0] = texelFetch(t_weight_scales, ivec3(out_col_texel_idx, block_idx, 0), 0); - scales[1] = texelFetch(t_weight_scales, ivec3(out_col_texel_idx + 1, block_idx, 0), 0); - - zeros[0] = vec4(texelFetch(t_weight_zeros, ivec3(out_col_texel_idx, block_idx, 0), 0)); - zeros[1] = vec4(texelFetch(t_weight_zeros, ivec3(out_col_texel_idx + 1, block_idx, 0), 0)); - - ivec4 int32_sums[TILE_ROWS][2]; - int input_sums[TILE_ROWS]; - - // Initialize accumulators for this block - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - int32_sums[r][0] = ivec4(0); - int32_sums[r][1] = ivec4(0); - input_sums[r] = 0; - } - - for (int g_idx = 4 * int(wid); g_idx < group_size; g_idx += (4 * NWORKERS)) { - const int k = block_idx * group_size + g_idx; - - // Preload B (weights) - keep as quantized integers - [[unroll]] for (int r = 0; r < 4; ++r) { - $if WEIGHT_STORAGE == "buffer": - const u8vec4 packed_weight_tex = t_qmat2[(k + r) * qmat2_stride + gl_GlobalInvocationID.x]; - $else: - const uvec4 packed_weight_tex = texelFetch( - t_qmat2, - ivec2(gl_GlobalInvocationID.x, k + r), - 0); - - // Unpack 4-bit weights to integers and subtract zero point (8 for 4-bit) - qmat2_quantized[r][0] = ivec4((packed_weight_tex & 0xF0) >> 4) - 8; - qmat2_quantized[r][1] = ivec4(packed_weight_tex & 0x0F) - 8; - } - - // Preload A (quantized input) - keep as quantized integers - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - $if IN_STORAGE == "buffer": - mat1_quantized[r] = t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r]; - $else: - mat1_quantized[r] = texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r]; - } - - // Accumulate in integer arithmetic: (input_quantized - input_zero_point) * (weight_quantized - weight_zero_point) - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w; - - int32_sums[r][0] += mat1_quantized[r].x * qmat2_quantized[0][0] - + mat1_quantized[r].y * qmat2_quantized[1][0] - + mat1_quantized[r].z * qmat2_quantized[2][0] - + mat1_quantized[r].w * qmat2_quantized[3][0]; - - int32_sums[r][1] += mat1_quantized[r].x * qmat2_quantized[0][1] - + mat1_quantized[r].y * qmat2_quantized[1][1] - + mat1_quantized[r].z * qmat2_quantized[2][1] - + mat1_quantized[r].w * qmat2_quantized[3][1]; - } - } - - // Incorporates this block's results into the final accumulation - // Following proper quantization paradigm: result = input_scale * weight_scale * - // Sum((input_quantized - input_zero) * (weight_quantized - weight_zero)) - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - if (out_row + r >= out_sizes.y) { - continue; - } - - float input_scale = t_input_scale[int(out_row) + r]; - float input_sum_scalar = float(input_sums[r]); - - // Apply proper quantization paradigm: input_scale * weight_scale * (accumulator - weight_zero * input_sum) - final_result[r][0] += input_scale * scales[0] * (vec4(int32_sums[r][0]) - zeros[0] * input_sum_scalar); - final_result[r][1] += input_scale * scales[1] * (vec4(int32_sums[r][1]) - zeros[1] * input_sum_scalar); - } - } - - // Store worker results in shared memory - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - partial_results[gid][wid][r][0] = final_result[r][0]; - partial_results[gid][wid][r][1] = final_result[r][1]; - } - - memoryBarrierShared(); - barrier(); - - // Only the first worker in each group accumulates and writes output - if (wid != 0) { - return; - } - - vec4 cooperative_result[TILE_ROWS][2]; - - for (int r = 0; r < TILE_ROWS; ++r) { - cooperative_result[r][0] = vec4(0.0); - cooperative_result[r][1] = vec4(0.0); - [[unroll]] for (int worker = 0; worker < NWORKERS; ++worker) { - cooperative_result[r][0] += partial_results[gid][worker][r][0]; - cooperative_result[r][1] += partial_results[gid][worker][r][1]; - } - } - - // Apply final output quantization - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - $if OUT_STORAGE == "buffer": - t_out[((out_row + r) * out_sizes.x + out_col) >> 2] = cooperative_result[r][0]; - t_out[((out_row + r) * out_sizes.x + out_col + 4) >> 2] = cooperative_result[r][1]; - $else: - imageStore(t_out, ivec3(out_col_texel_idx, out_row + r, 0), cooperative_result[r][0]); - imageStore(t_out, ivec3(out_col_texel_idx + 1, out_row + r, 0), cooperative_result[r][1]); - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.yaml deleted file mode 100644 index 9f6db77094a..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_coop.yaml +++ /dev/null @@ -1,26 +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. - -linear_qta8a_qga4w_coop: - parameter_names_with_default_values: - DTYPE: float - OUT_STORAGE: texture3d - IN_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - PARAMS_STORAGE: buffer - TILE_ROWS: 1 - shader_variants: - - NAME: linear_qta8a_qga4w_coop_texture3d_texture3d_texture2d_float - - NAME: linear_qta8a_qga4w_coop_buffer_buffer_texture2d_float - OUT_STORAGE: buffer - IN_STORAGE: buffer - - NAME: linear_qta8a_qga4w_coop_buffer_buffer_buffer_float - OUT_STORAGE: buffer - IN_STORAGE: buffer - WEIGHT_STORAGE: buffer - - NAME: linear_qta8a_qga4w_coop_buffer_texture2d_buffer_float - OUT_STORAGE: buffer - WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.glsl deleted file mode 100644 index dbb7da998f4..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.glsl +++ /dev/null @@ -1,196 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${buffer_scalar_type(DTYPE)} -#define VEC4_T ${buffer_gvec_type(DTYPE, 4)} - -#define TILE_ROWS ${TILE_ROWS} - -${define_required_extensions(DTYPE)} -$if IN_STORAGE == "buffer": - ${define_required_extensions("int8")} -$if WEIGHT_STORAGE == "buffer": - ${define_required_extensions("uint8")} - -#extension GL_EXT_control_flow_attributes : require - -layout(std430) buffer; - -${layout_declare_tensor(B, "w", "t_out", DTYPE, OUT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_mat1", "int8", IN_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_qmat2", "uint8", WEIGHT_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_scales", "float", PARAMS_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_weight_zeros", "int", PARAMS_STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_input_scale", "float", "buffer", is_scalar_array=True)} -${layout_declare_tensor(B, "r", "t_input_zero_point", "int", "buffer", is_scalar_array=True)} - -layout(push_constant) uniform restrict Block { - ivec4 out_sizes; - ivec4 mat1_sizes; - ivec4 qmat2_sizes; -}; - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -layout(constant_id = 3) const int group_size = 64; - -/* - * This shader computes a linear operator between a quantized int8 input matrix - * x and a weights matrix that is quantized to 4 bits, producing a float output. - * - * The (W, H, C) shape of each tensor is: - * - x: (K, M) - quantized int8 input with per-token quantization - * - weights: (N / 2, K) - * - The weights tensor has a data type of `uint8`. Each element in the tensor - * contains 2 4-bit values packed into a uint8. - * - See the pack_int4_linear_weight_transposed_interleave shader to see more - * details on how the weight tensor is stored. - * - qparams: (2, N, number_of_groups) - * - This tensor contains the scales and zeros quantization parameters for the - * weights tensor. The weight tensor is quantized group-wise, which means - * that every `group_size` elements along the K dimension of the weights - * tensor has independent quantization parameters. Along the width dim, the - * first value contains the scale for the group and the second value - * contains the zero point for the group. - * - input_scale: (num_tokens,) - per-token scale values for input quantization - * - input_zero_point: (num_tokens,) - per-token zero points for input quantization - * - output: (N, M) - float output - * - * Each thread computes a tile of TILE_ROWS * 2 texels of the output tensor. - * - * Note that this shader assumes that all tensors are width packed. - */ - -void main() { - const uint out_row = gl_GlobalInvocationID.y * TILE_ROWS; - const uint out_col = gl_GlobalInvocationID.x << 3; - const int out_col_texel_idx = int(gl_GlobalInvocationID.x) << 1; - - if (out_col >= out_sizes.x || out_row >= out_sizes.y) { - return; - } - - const int num_blocks = mat1_sizes.x / group_size; - - ivec4 mat1_quantized[TILE_ROWS]; - ivec4 qmat2_quantized[4][2]; - vec4 final_result[TILE_ROWS][2]; - - // Initialize accumulatoxrs - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - final_result[r][0] = vec4(0.0); - final_result[r][1] = vec4(0.0); - } - - vec4 scales[2]; - vec4 zeros[2]; - - $if WEIGHT_STORAGE == "buffer": - const int qmat2_stride = qmat2_sizes.x >> 2; - $if PARAMS_STORAGE == "buffer": - const int qparams_stride = out_sizes.x >> 2; - - for (int block_idx = 0; block_idx < num_blocks; ++block_idx) { - $if PARAMS_STORAGE == "buffer": - scales[0] = t_weight_scales[block_idx * qparams_stride + out_col_texel_idx]; - scales[1] = t_weight_scales[block_idx * qparams_stride + out_col_texel_idx + 1]; - - zeros[0] = vec4(t_weight_zeros[block_idx * qparams_stride + out_col_texel_idx]); - zeros[1] = vec4(t_weight_zeros[block_idx * qparams_stride + out_col_texel_idx + 1]); - $else: - scales[0] = texelFetch(t_weight_scales, ivec3(out_col_texel_idx, block_idx, 0), 0); - scales[1] = texelFetch(t_weight_scales, ivec3(out_col_texel_idx + 1, block_idx, 0), 0); - - zeros[0] = vec4(texelFetch(t_weight_zeros, ivec3(out_col_texel_idx, block_idx, 0), 0)); - zeros[1] = vec4(texelFetch(t_weight_zeros, ivec3(out_col_texel_idx + 1, block_idx, 0), 0)); - - ivec4 int32_sums[TILE_ROWS][2]; - int input_sums[TILE_ROWS]; - - // Initialize accumulators - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - int32_sums[r][0] = ivec4(0); - int32_sums[r][1] = ivec4(0); - input_sums[r] = 0; - } - - for (int g_idx = 0; g_idx < group_size; g_idx += 4) { - const int k = block_idx * group_size + g_idx; - - // Preload B (weights) - keep as quantized integers - [[unroll]] for (int r = 0; r < 4; ++r) { - $if WEIGHT_STORAGE == "buffer": - const u8vec4 packed_weight_tex = t_qmat2[(k + r) * qmat2_stride + gl_GlobalInvocationID.x]; - $else: - const uvec4 packed_weight_tex = texelFetch( - t_qmat2, - ivec2(gl_GlobalInvocationID.x, k + r), - 0); - - // Unpack 4-bit weights to integers (subtract 8 as the 4-bit zero point) - qmat2_quantized[r][0] = ivec4((packed_weight_tex & 0xF0) >> 4) - 8; - qmat2_quantized[r][1] = ivec4(packed_weight_tex & 0x0F) - 8; - } - - // Preload A (quantized input) - keep as quantized integers - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - $if IN_STORAGE == "buffer": - mat1_quantized[r] = t_mat1[((out_row + r) * mat1_sizes.x + k) >> 2] - t_input_zero_point[int(out_row) + r]; - $else: - mat1_quantized[r] = texelFetch(t_mat1, ivec3(k >> 2, out_row + r, 0), 0) - t_input_zero_point[int(out_row) + r]; - } - - // Accumulate in integer arithmetic: (input_quantized - input_zero_point) * (weight_quantized - weight_zero_point) - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - input_sums[r] += mat1_quantized[r].x + mat1_quantized[r].y + mat1_quantized[r].z + mat1_quantized[r].w; - - int32_sums[r][0] += mat1_quantized[r].x * qmat2_quantized[0][0] - + mat1_quantized[r].y * qmat2_quantized[1][0] - + mat1_quantized[r].z * qmat2_quantized[2][0] - + mat1_quantized[r].w * qmat2_quantized[3][0]; - - int32_sums[r][1] += mat1_quantized[r].x * qmat2_quantized[0][1] - + mat1_quantized[r].y * qmat2_quantized[1][1] - + mat1_quantized[r].z * qmat2_quantized[2][1] - + mat1_quantized[r].w * qmat2_quantized[3][1]; - } - } - - // Incorporates this block's results into the final accumulation - // Following proper quantization paradigm: result = input_scale * weight_scale * - // Sum((input_quantized - input_zero) * (weight_quantized - weight_zero)) - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - if (out_row + r >= out_sizes.y) { - continue; - } - - float input_scale = t_input_scale[int(out_row) + r]; - float input_sum_scalar = float(input_sums[r]); - - // Apply proper quantization paradigm: input_scale * weight_scale * (accumulator - weight_zero * input_sum) - final_result[r][0] += input_scale * scales[0] * (vec4(int32_sums[r][0]) - zeros[0] * input_sum_scalar); - final_result[r][1] += input_scale * scales[1] * (vec4(int32_sums[r][1]) - zeros[1] * input_sum_scalar); - } - } - - // Apply ALL scaling at the very end - [[unroll]] for (int r = 0; r < TILE_ROWS; ++r) { - $if OUT_STORAGE == "buffer": - if (out_row + r < out_sizes.y) { - t_out[((out_row + r) * out_sizes.x + out_col) >> 2] = final_result[r][0]; - t_out[((out_row + r) * out_sizes.x + out_col + 4) >> 2] = final_result[r][1]; - } - $else: - imageStore(t_out, ivec3(out_col_texel_idx, out_row + r, 0), final_result[r][0]); - imageStore(t_out, ivec3(out_col_texel_idx + 1, out_row + r, 0), final_result[r][1]); - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.yaml deleted file mode 100644 index c96d693834b..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qta8a_qga4w_tiled.yaml +++ /dev/null @@ -1,26 +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. - -linear_qta8a_qga4w_tiled: - parameter_names_with_default_values: - DTYPE: float - OUT_STORAGE: texture3d - IN_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - PARAMS_STORAGE: buffer - TILE_ROWS: 3 - shader_variants: - - NAME: linear_qta8a_qga4w_tiled_texture3d_texture3d_texture2d_float - - NAME: linear_qta8a_qga4w_tiled_buffer_buffer_texture2d_float - OUT_STORAGE: buffer - IN_STORAGE: buffer - - NAME: linear_qta8a_qga4w_tiled_buffer_buffer_buffer_float - OUT_STORAGE: buffer - IN_STORAGE: buffer - WEIGHT_STORAGE: buffer - - NAME: linear_qta8a_qga4w_tiled_buffer_texture2d_buffer_float - OUT_STORAGE: buffer - WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/max_pool2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/max_pool2d.glsl index 9d78b7a6a6e..28afe5a822f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/max_pool2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/max_pool2d.glsl @@ -24,6 +24,8 @@ ${layout_declare_ubo(B, "ivec2", "kernel_size", "ivec2", "stride", "ivec2", "pad layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; +${layout_declare_spec_const(C, "int", "write_indices", "1")} + void main() { const ivec3 pos = ivec3(gl_GlobalInvocationID); @@ -55,5 +57,7 @@ void main() { } imageStore(t_out, pos, out_texel); - imageStore(t_idx, pos, idx_texel); + if (write_indices > 0) { + imageStore(t_idx, pos, idx_texel); + } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.glsl index 62cd0610ffb..a16f5405cbb 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.glsl @@ -4,46 +4,46 @@ #define T ${buffer_scalar_type(DTYPE)} -#include "indexing_utils.h" - ${define_required_extensions(DTYPE)} +${define_required_extensions(BUF_DTYPE)} layout(std430) buffer; -${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_tensor(B, "r", "nchw_in", DTYPE, STORAGE)} +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "t_outp", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "nchw_in", BUF_DTYPE, STORAGE)} -$if USE_PUSH_CONST: - layout(push_constant) uniform restrict Block { - ivec4 out_sizes; - ivec4 out_strides; - int numel; - }; -$else: - ${layout_declare_ubo(B, "ivec4", "out_sizes")} - ${layout_declare_ubo(B, "ivec4", "out_strides")} - ${layout_declare_ubo(B, "int", "numel")} +${layout_declare_ubo(B, "BufferMetadata", "outp")} 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", "DEFAULT_DIM_ORDER")} -const lowp ivec4 out_dim_order = unhash_dim_order(out_layout); +// This constant is unused in this shader but is kept so that the signature is +// consistent with nchw_to_image. +${layout_declare_spec_const(C, "int", "unused", "0")} ${layout_declare_spec_const(C, "int", "transpose_hw", "0")} void main() { - int out_bufi = int(gl_GlobalInvocationID.x); - if (out_bufi >= numel) { + const uint outp_bufi = int(gl_GlobalInvocationID.x); + if (outp_bufi >= numel(outp)) { return; } - ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, out_dim_order); + TensorIndex outp_tidx; + uint nchwi; + + linear_idx_to_tensor_idx(outp, outp_bufi, outp_tidx); - ivec4 sizes = out_sizes; if (transpose_hw == 1) { - sizes.xy = sizes.yx; - out_tidx.xy = out_tidx.yx; + BufferMetadata transposed_meta = outp; + transposed_meta.sizes[0].xy = transposed_meta.sizes[0].yx; + outp_tidx.data[0].xy = outp_tidx.data[0].yx; + nchwi = tensor_idx_to_contiguous_idx(transposed_meta, outp_tidx); + } + // Normal case + else { + nchwi = tensor_idx_to_contiguous_idx(outp, outp_tidx); } - const int in_nchwi = tidx_to_nchwi(out_tidx, sizes); - t_out[out_bufi] = nchw_in[in_nchwi]; + t_outp[outp_bufi] = T(nchw_in[nchwi]); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.yaml index 99e41a0ab6f..602fd1bc65a 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_buffer.yaml @@ -7,17 +7,19 @@ nchw_to_buffer: parameter_names_with_default_values: DTYPE: float + BUF_DTYPE: float STORAGE: buffer USE_PUSH_CONST: True generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: double - - VALUE: int8 - - VALUE: uint8 - - VALUE: int32 + combination: + parameter_names: [DTYPE, BUF_DTYPE] + combos: + - parameter_values: [half, half] + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [double, double] + - parameter_values: [int8, int8] + - parameter_values: [uint8, uint8] + - parameter_values: [int32, int32] shader_variants: - NAME: nchw_to_buffer - - NAME: nchw_to_buffer_no_pc - USE_PUSH_CONST: False diff --git a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl index f3f604e10cd..15676fb0500 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.glsl @@ -16,11 +16,12 @@ ${define_active_storage_type(STORAGE)} ${define_required_extensions(DTYPE)} +${define_required_extensions(BUF_DTYPE)} layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_buffer(B, "r", "buf_in", DTYPE)} +${layout_declare_buffer(B, "r", "buf_in", BUF_DTYPE)} $if USE_PUSH_CONST: layout(push_constant) uniform restrict Block { diff --git a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.yaml b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.yaml index 85119c8d508..f6809e4024a 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/nchw_to_image.yaml @@ -11,13 +11,16 @@ nchw_to_image: FROM_STAGING: True USE_PUSH_CONST: True generate_variant_forall: - DTYPE: - - VALUE: half - - VALUE: float - - VALUE: double - - VALUE: int8 - - VALUE: uint8 - - VALUE: int32 + combination: + parameter_names: [DTYPE, BUF_DTYPE] + combos: + - parameter_values: [half, half] + - parameter_values: [half, float] + - parameter_values: [float, float] + - parameter_values: [double, double] + - parameter_values: [int8, int8] + - parameter_values: [uint8, uint8] + - parameter_values: [int32, int32] shader_variants: - NAME: nchw_to_image_texture3d - NAME: nchw_to_image_texture2d diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.glsl deleted file mode 100644 index e42cf05dd7f..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.glsl +++ /dev/null @@ -1,154 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -layout(std430) buffer; - -#include "indexing_utils.h" - -${layout_declare_tensor(B, "w", "t_qmat2", "uint", STORAGE, is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_input", "uint", "buffer")} - -layout(push_constant) uniform restrict Block { - ivec4 qmat2_sizes; - ivec2 orig_sizes; -}; - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -$if STORAGE == "buffer": - #define BUFFER_WEIGHT - -#include "qlinear_weight_pack_utils.glslh" - -#define extract_4bit(input_block_data, col, row) \ - (extract_4bit_from_packed_uint_le(input_block_data[row], col)) - -/* - * This shader packs the weight tensor into blocks for efficient consumption. - * - * The input tensor has shape [K/2, N] where each element is a uint8 containing - * 2 packed 4-bit values. The logical tensor shape is [K, N] of 4-bit values. - * - * The transformation partitions the tensor into blocks of size 4x8 (4-bit values) - * and transposes each block to 8x4, then packs the result so that each uvec4 - * contains an entire transposed block. - * - * Original block (4x8 4-bit values, shown as 2x8 uint8 values): - * w00|w10, w20|w30, - * w01|w11, w21|w31, - * w02|w12, w22|w32, - * w03|w13, w23|w33, - * w04|w14, w24|w34, - * w05|w15, w25|w35, - * w06|w16, w26|w36, - * w07|w17, w27|w37, - * - * Transposed block (8x4 4-bit values, packed into uvec4): - * w00|w01, w02|w03, w04|w05, w06|w07 - * w10|w11, w12|w13, w14|w15, w16|w17 - * w20|w21, w22|w23, w24|w25, w26|w27 - * w30|w31, w32|w33, w34|w35, w36|w37 - */ -void main() { - // Each thread writes out 2 adjacent 8 wide x 4 high transposed block. Each - // block is packed as one uvec4. - ivec2 block_pos = ivec2( - MUL_2(gl_GlobalInvocationID.x), - gl_GlobalInvocationID.y); - - // There are K wide x N high 4-bit values in the original weight tensor - const int input_width = orig_sizes.x; // K - const int input_height = orig_sizes.y; // N - - const int input_width_uint = DIV_UP_8(input_width); - - // Original block spans 4 wide x 8 high 4-bit values. Since uint is used to - // read the input tensor, each block spans 0.5 wide x 8 high uint values. - const ivec2 block_start = ivec2( - DIV_2(block_pos.x), - MUL_8(block_pos.y)); - - // Check bounds - if (block_start.x >= input_width_uint || block_start.y >= input_height) { - return; - } - - // Read input block. Note that this block will contain the source data for - // both output blocks, as it contains 1 wide x 8 high uint values, which is - // equivalent to 8 wide x 8 high 4-bit values. - uint input_block_data[8]; - - // Read in 8 rows along the same column of uints, each uint contains 4 4-bit - // values. This will be the source data for the transposed block. - for (int i = 0; i < 8; ++i) { - uint input_bufi = (block_start.y + i) * input_width_uint + block_start.x; - input_block_data[i] = t_input[input_bufi]; - } - - for (int col_offset = 0; col_offset <= 4; col_offset+=4) { - uvec4 output_block; - - output_block.x = pack_8x4bit_into_uint( - extract_4bit(input_block_data, col_offset, 0), - extract_4bit(input_block_data, col_offset, 1), - extract_4bit(input_block_data, col_offset, 2), - extract_4bit(input_block_data, col_offset, 3), - extract_4bit(input_block_data, col_offset, 4), - extract_4bit(input_block_data, col_offset, 5), - extract_4bit(input_block_data, col_offset, 6), - extract_4bit(input_block_data, col_offset, 7)); - - output_block.y = pack_8x4bit_into_uint( - extract_4bit(input_block_data, col_offset + 1, 0), - extract_4bit(input_block_data, col_offset + 1, 1), - extract_4bit(input_block_data, col_offset + 1, 2), - extract_4bit(input_block_data, col_offset + 1, 3), - extract_4bit(input_block_data, col_offset + 1, 4), - extract_4bit(input_block_data, col_offset + 1, 5), - extract_4bit(input_block_data, col_offset + 1, 6), - extract_4bit(input_block_data, col_offset + 1, 7)); - - output_block.z = pack_8x4bit_into_uint( - extract_4bit(input_block_data, col_offset + 2, 0), - extract_4bit(input_block_data, col_offset + 2, 1), - extract_4bit(input_block_data, col_offset + 2, 2), - extract_4bit(input_block_data, col_offset + 2, 3), - extract_4bit(input_block_data, col_offset + 2, 4), - extract_4bit(input_block_data, col_offset + 2, 5), - extract_4bit(input_block_data, col_offset + 2, 6), - extract_4bit(input_block_data, col_offset + 2, 7)); - - output_block.w = pack_8x4bit_into_uint( - extract_4bit(input_block_data, col_offset + 3, 0), - extract_4bit(input_block_data, col_offset + 3, 1), - extract_4bit(input_block_data, col_offset + 3, 2), - extract_4bit(input_block_data, col_offset + 3, 3), - extract_4bit(input_block_data, col_offset + 3, 4), - extract_4bit(input_block_data, col_offset + 3, 5), - extract_4bit(input_block_data, col_offset + 3, 6), - extract_4bit(input_block_data, col_offset + 3, 7)); - - const uint qmat2_texel_stride_x = DIV_UP_4(qmat2_sizes.x); - write_transposed_weight_block( - output_block, - block_pos.x, - block_pos.y, - qmat2_texel_stride_x); - - if (MUL_8(block_start.x) + 4 >= input_width) { - return; - } - // Otherwise, implement the block position to write to the next block in the - // following iteration. - block_pos.x += 1; - } -} diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl new file mode 100644 index 00000000000..b9f5c994910 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.glsl @@ -0,0 +1,68 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_packed_int4_weight", "int", STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int4_weight", "uint", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 qmat2_sizes; + ivec2 orig_sizes; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "common.glslh" +#include "linear_int4_weight_block.glslh" + +void main() { + const int k8 = int(gl_GlobalInvocationID.x); + const int n8 = int(gl_GlobalInvocationID.y); + + const int K = orig_sizes.x; + const int N = orig_sizes.y; + + // Each shader invocation processes a 4x8 block of the input data. + const int K4 = div_up_4(K); + const int K8 = div_up_8(K); + const int N8 = div_up_8(N); + + // Check bounds + if (n8 >= N8 || k8 >= K8) { + return; + } + + Int4Weight2xBlockSourceData src_data; + const int n = mul_8(n8); + if (N - n >= 8) { + load_block_source_data_no_checks(src_data, k8, n, K8, N); + } else { + load_block_source_data_with_checks(src_data, k8, n, K8, N); + } + + // A 8Kx8K block of the weight matrix is loaded into memory. This will be + // split into two blocks each holding 4Kx8N worth of data. + // The first block contains data for k + (0, 1, 2, 3) i.e. the first 4 columns + // of the loaded weight block. + Int4WeightBlockPacked packed_block_1; + // The second block contains data for k + (4, 5, 6, 7) i.e. the second 4 cols + // of the loaded weight block + Int4WeightBlockPacked packed_block_2; + create_packed_blocks(packed_block_1, packed_block_2, src_data); + + const int k4 = mul_2(k8); + write_packed_block(packed_block_1, k4, n8, K4); + write_packed_block(packed_block_2, k4 + 1, n8, K4); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.yaml similarity index 65% rename from backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.yaml rename to backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.yaml index c72a2cc1df6..7a145ec95d7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_int4_linear_weight_transposed_block_4x8.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q4_linear_weight.yaml @@ -4,11 +4,11 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -pack_int4_linear_weight_transposed_block_4x8: +pack_q4_linear_weight: parameter_names_with_default_values: STORAGE: buffer shader_variants: - - NAME: pack_int4_linear_weight_transposed_block_4x8_buffer + - NAME: pack_q4_linear_weight_buffer STORAGE: buffer - - NAME: pack_int4_linear_weight_transposed_block_4x8_texture2d + - NAME: pack_q4_linear_weight_texture2d STORAGE: texture2d 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 new file mode 100644 index 00000000000..f2c74b67283 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl @@ -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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_packed_int8_weight", "int", STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_weight", "int", "buffer")} + +layout(push_constant) uniform restrict Block { + ivec4 qmat2_sizes; + ivec2 orig_sizes; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "common.glslh" +#include "linear_int8_weight_block.glslh" + +void main() { + // The size of the source weight tensor is [W=K, H=N]. Each shader invocation + // processes a 4x4 block. The thread position corresponds to the block index. + int n4 = int(gl_GlobalInvocationID.x); + int k4 = int(gl_GlobalInvocationID.y); + + const int K = orig_sizes.x; + const int N = orig_sizes.y; + + // Determine the total number of blocks and check bounds + const int N4 = div_up_4(N); + const int K4 = div_up_4(K); + if (n4 >= N4 || k4 >= K4) { + return; + } + + // Each block is represented as an ivec4. Each int corresponds to a row i.e. + // N dim of the weight tensor and contains data for 4 columns i.e. K dim. + Int8WeightBlock block; + const int n = mul_4(n4); + if (N - n >= 4) { + load_block_data_no_checks(block, k4, n, K4, N); + } else { + load_block_data_with_checks(block , k4, n, K4, N); + } + + // 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. + write_weight_block(block, n4, k4, N4); +} 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 new file mode 100644 index 00000000000..13e6d43b2c5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml @@ -0,0 +1,14 @@ +# 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. + +pack_q8_linear_weight: + parameter_names_with_default_values: + STORAGE: buffer + shader_variants: + - NAME: pack_q8_linear_weight_buffer + STORAGE: buffer + - NAME: pack_q8_linear_weight_texture2d + STORAGE: texture2d diff --git a/backends/vulkan/runtime/graph/ops/glsl/permute_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/permute_buffer.glsl index 55b9e3dc9ea..3447ab07552 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/permute_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/permute_buffer.glsl @@ -18,55 +18,31 @@ ${define_required_extensions(DTYPE)} layout(std430) buffer; -#include "indexing_utils.h" +#include "indexing.glslh" -${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")} +${layout_declare_tensor(B, "w", "t_outp", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_inp", DTYPE, "buffer")} -${layout_declare_ubo(B, "ivec4", "in_sizes")} -${layout_declare_ubo(B, "ivec4", "out_strides")} -${layout_declare_ubo(B, "int", "out_numel")} +${layout_declare_ubo(B, "BufferMetadata", "outp")} +${layout_declare_ubo(B, "BufferMetadata", "inp")} -layout(push_constant) uniform restrict Block { - ivec4 in_strides; - ivec4 permute_dims; // Permutation mapping: permute_dims[i] = j means output dim i comes from input dim j -}; - -${layout_declare_spec_const(C, "int", "out_layout", "DEFAULT_LAYOUT")} -${layout_declare_spec_const(C, "int", "in_layout", "DEFAULT_LAYOUT")} - -const lowp ivec4 out_dim_order = unhash_dim_order(out_layout); +${layout_declare_ubo(B, "ivec4[DIMLIMIT_DIV4]", "permute_order")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; -// Convert output tensor index to input tensor index based on permutation -ivec4 out_tidx_to_in_tidx(const ivec4 out_tidx) { - ivec4 in_tidx; - - // Apply the permutation mapping: in_tidx[permute_dims[i]] = out_tidx[i] - in_tidx[permute_dims.x] = out_tidx.x; - in_tidx[permute_dims.y] = out_tidx.y; - in_tidx[permute_dims.z] = out_tidx.z; - in_tidx[permute_dims.w] = out_tidx.w; - - return in_tidx; -} - void main() { - const int out_bufi = ivec3(gl_GlobalInvocationID).x; - if (out_bufi >= out_numel) { + const uint inp_bufi = gl_GlobalInvocationID.x; + if (inp_bufi >= numel(inp)) { return; } - // Convert buffer index to tensor index for output - const ivec4 out_tidx = bufi_to_tidx(out_bufi, out_strides, out_dim_order); - - // Convert output tensor index to input tensor index using permutation - const ivec4 in_tidx = out_tidx_to_in_tidx(out_tidx); + TensorIndex inp_tidx; + linear_idx_to_tensor_idx(inp, inp_bufi, inp_tidx); - // Convert input tensor index back to buffer index - const int in_bufi = tidx_to_bufi(in_tidx, in_strides); + TensorIndex outp_tidx = inp_tidx; + permute(outp_tidx, permute_order); + const uint outp_bufi = tensor_idx_to_linear_idx(outp, outp_tidx); // Copy data from input to output - t_out[out_bufi] = t_in[in_bufi]; + t_outp[outp_bufi] = t_inp[inp_bufi]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/qlinear_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/qlinear_utils.glslh deleted file mode 100644 index 80ec44c153a..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/qlinear_utils.glslh +++ /dev/null @@ -1,116 +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. - */ - -#ifndef QLINEAR_UTILS_H -#define QLINEAR_UTILS_H - -/*********************************** - * Packed Weight data read/write functions - * - * These functions assume that t_qmat2 is declared in the shader layout as a storage - * buffer or storage image. - */ - -#ifdef BUFFER_WEIGHT - -uvec4 load_transposed_weight_block(const uint k4, const uint n8, const uint K4) { - return t_qmat2[n8 * K4 + k4]; -} - -#else // TEXTURE_WEIGHT - -uvec4 load_transposed_weight_block(const uint k4, const uint n8, const uint K4) { - return texelFetch(t_qmat2, ivec2(k4, n8), 0); -} - -#endif // BUFFER_WEIGHT - -/*********************************** - * Packed weight data extraction functions - */ - -/* - * uvec4 block contains a packed 4 high x 8 wide matrix of 4-bit signed integers. This - * function extracts the 4-bit values at the given column and row index. - * - * Each uint in the uvec4 corresponds to one row; thus the desired row can be extracted - * via block[row]. From there, column 0 is packed in bits 28-31, column 1 is packed into - * bits 24-27, column 3 is packed into bits 20-23, and so on. To extract the desired - * value: - * - * 1. First, shift the row uint by 4 * (7 - col) bits - * 2. Apply a mask of 0b1111 = 15 - * - * Finally, convert the masked value to int and subtract it by int to obtain the desired - * signed integer. - */ -T extract_4bit_from_transposed_block(const uvec4 block, const uint col, const uint row) { - return T(int((block[row] >> (4 * (7 - col))) & 15) - 8); -} - -/*********************************** - * Input/Output read/write functions - * - * These functions assume that t_input and t_output are declared in the shader layout as - * storage buffers or storage images. - */ - -#ifdef BUFFER_IO - -VEC4_T load_input_texel_1d(const uint k4) { - return t_input[k4]; -} - -VEC4_T load_input_texel_2d( - const uint k4, - const uint m, - const uint K4) { - return t_input[(m * K4) + k4]; -} - -void write_output_texel_1d(const VEC4_T out_texel, const uint n4) { - t_output[n4] = out_texel; -} - -void write_output_texel_2d( - const VEC4_T out_texel, - const uint n4, - const uint m, - const uint N4) { - t_output[m * N4 + n4] = out_texel; -} - -#else // TEXTURE_IO - -VEC4_T load_input_texel_1d(const uint k4) { - return texelFetch(t_input, ivec3(k4, 0, 0), 0); -} - -VEC4_T load_input_texel_2d( - const uint k4, - const uint m, - const uint K4) { - return texelFetch(t_input, ivec3(k4, m, 0), 0); -} - - -void write_output_texel_1d(const VEC4_T out_texel, const uint n4) { - imageStore(t_output, ivec3(n4, 0, 0), out_texel); -} - -void write_output_texel_2d( - const VEC4_T out_texel, - const uint n4, - const uint m, - const uint N4) { - imageStore(t_output, ivec3(n4, m, 0), out_texel); -} - -#endif // BUFFER_IO - -#endif // QLINEAR_UTILS_H diff --git a/backends/vulkan/runtime/graph/ops/glsl/qlinear_weight_pack_utils.glslh b/backends/vulkan/runtime/graph/ops/glsl/qlinear_weight_pack_utils.glslh deleted file mode 100644 index 1f481f4f859..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/qlinear_weight_pack_utils.glslh +++ /dev/null @@ -1,58 +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. - */ - -#ifndef QLINEAR_WEIGHT_PACK_UTILS_H -#define QLINEAR_WEIGHT_PACK_UTILS_H - -/*********************************** - * Packed Weight data write functions - * - * These functions assume that t_qmat2 has been defined in the shader layout as either - * a storage buffer or a storage image. - */ - -#ifdef BUFFER_WEIGHT - -void write_transposed_weight_block(const uvec4 block, const uint k4, const uint n8, const uint K4) { - t_qmat2[n8 * K4 + k4] = block; -} - -#else // TEXTURE_WEIGHT - -void write_transposed_weight_block(const uvec4 block, const uint k4, const uint n8, const uint K4) { - imageStore(t_qmat2, ivec2(k4, n8), block); -} - -#endif // BUFFER_WEIGHT - -/*********************************** - * Utilities for packing weight data - */ - -uint extract_4bit_from_packed_uint_le(const uint packed, const uint i) { - // account for little endian - uint byte = packed >> (8 * (i / 2)) & 255; - return (byte >> (4 - 4 * (i % 2))) & 15; -} - -uint pack_8x4bit_into_uint( - const uint val0, - const uint val1, - const uint val2, - const uint val3, - const uint val4, - const uint val5, - const uint val6, - const uint val7) { - return uint( - (val0 << 28) | (val1 << 24) | (val2 << 20) | (val3 << 16) | (val4 << 12) | - (val5 << 8) | (val6 << 4) | val7 - ); -} - -#endif // QLINEAR_WEIGHT_PACK_UTILS_H diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.glsl new file mode 100644 index 00000000000..450d6376537 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.glsl @@ -0,0 +1,89 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +#define TILE_M4 1 +#define TILE_N4 1 +#define TILE_K4 1 + +#define TILE_M 4 +#define TILE_N 4 +#define TILE_K 4 + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "conv2d_common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} + +// Sizes of the im2col matrix of the convolution input +${layout_declare_ubo(B, "ivec4", "matrix_sizes")} +// Sizes of the input image +${layout_declare_ubo(B, "ivec4", "input_sizes")} +// Sizes of the output image +${layout_declare_ubo(B, "ivec4", "output_sizes")} + +${layout_declare_ubo(B, "Conv2DParams", "conv2d_params")} + +layout(push_constant) uniform restrict Block { + float inv_scale; + int zp; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "conv2d_fp_im2col_block_load.glslh" +#include "linear_int8_input_block.glslh" + +void main() { + // The quantized and packed im2col matrix can be conceptualized as a 2D matrix + // with K/4 columns and M/4 rows. Each element of the matrix is a ivec4 which + // contains packed data for a 4 wide x 4 high block of the original im2col + // matrix. Each shader invocation works on writing out one ivec4, i.e. one + // block of the quantized and packed matrix. + + // Thread id corresponds to the block index + const int k4 = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + // Convert block idx to tensor idx + const int k = mul_4(k4); + const int m = mul_4(m4); + + const int logical_K = conv2d_params.logical_K; + // Similarly, compute the logical size of the M dim. + const int logical_M = output_sizes.x * output_sizes.y * output_sizes.w; + + // Check if tensor indices are out of bounds + if (k >= logical_K || m >= logical_M) { + return; + } + + FPInputTile in_tile; + load_input_im2col_tile(in_tile, k4, m4, logical_K, logical_M); + + Int8InputBlock packed_block; + quantize_and_pack(packed_block, in_tile, inv_scale, zp); + + // Number of texels in the x dim of the output matrix + const int K4 = div_4(matrix_sizes.x); + write_block(packed_block, k4, m4, K4); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.yaml new file mode 100644 index 00000000000..93f8269d607 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_im2col.yaml @@ -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. + +quantize_and_pack_im2col: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: quantize_and_pack_im2col_buffer_texture3d + - NAME: quantize_and_pack_im2col_texture3d_texture3d + OUTPUT_STORAGE: texture3d diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.glsl new file mode 100644 index 00000000000..6ba9343f10d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.glsl @@ -0,0 +1,79 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} + +$if GRANULARITY == "per_channel": + ${layout_declare_tensor(B, "r", "t_scale", DTYPE, "buffer")} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(push_constant) uniform restrict Block { + float inv_scale; + int zp; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "linear_int8_input_block.glslh" +#include "linear_fp_input_tile_load.glslh" + +void main() { + // Each input block contains 4x4 int8 quantized values, which are packed into + // a ivec4. k4 and m4 represent the "block index" of the current block being + // processed. + int k4 = int(gl_GlobalInvocationID.x); + int m4 = int(gl_GlobalInvocationID.y); + + const int K = input_sizes.x; + const int M = input_sizes.y; + + // K4 and M4 represent the number of blocks in each dimension. + const int K4 = div_up_4(K); + const int M4 = div_up_4(M); + + if (k4 >= K4 || m4 >= M4) { + return; + } + + // row of the input tensor to start loading from. Note the input tensor is + // interpreted as a t + const int m = mul_4(m4); + + const bool dont_check_bounds = (M - m) >= 4; + + FPInputTile in_tile; + if (dont_check_bounds) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + } else { + load_input_tile_with_checks(in_tile, k4, m, K4, M); + } + + Int8InputBlock packed_block; + quantize_and_pack(packed_block, in_tile, inv_scale, zp); + + write_block(packed_block, k4, m4, K4); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.yaml new file mode 100644 index 00000000000..37721db1ba8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input.yaml @@ -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. + +quantize_and_pack_linear_input: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: texture3d + INPUT_STORAGE: texture3d + STORAGE: texture3d + GRANULARITY: per_tensor + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: quantize_and_pack_linear_input_per_tensor_texture3d_texture3d + - NAME: quantize_and_pack_linear_input_per_tensor_buffer_texture3d + OUTPUT_STORAGE: buffer + - NAME: quantize_and_pack_linear_input_per_tensor_buffer_buffer + OUTPUT_STORAGE: buffer + INPUT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.glsl new file mode 100644 index 00000000000..792d8be908b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.glsl @@ -0,0 +1,122 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${texel_load_component_type(DTYPE, INPUT_STORAGE)} + +$if OUTPUT_STORAGE == "buffer": + #define OUTPUT_BUFFER +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +${define_required_extensions(DTYPE)} +${define_required_extensions("int8")} + +#extension GL_EXT_integer_dot_product : require + +#define NUM_GROUPS_PER_WG ${NUM_GROUPS_PER_WG} +#define NUM_WORKERS_PER_GROUP ${NUM_WORKERS_PER_GROUP} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_packed_int8_input", "int", OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "w", "t_int8_input_sums", "int", "buffer", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_input", DTYPE, INPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_int8_input_scales", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_int8_input_zps", "int8", "texture3d")} + +${layout_declare_ubo(B, "ivec4", "input_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "int", "K4_per_group", "0")} + +shared ivec4 shared_sums[NUM_GROUPS_PER_WG][NUM_WORKERS_PER_GROUP]; + +#define TILE_M4 1 +#define TILE_K4 1 + +#define TILE_M 4 + +#include "linear_int8_input_block.glslh" +#include "linear_int8_input_scales_zps_load.glslh" +#include "linear_fp_input_tile_load.glslh" + +void main() { + const int group_idx = int(gl_GlobalInvocationID.x); + const int m4 = int(gl_GlobalInvocationID.y); + + const int worker_id = int(gl_LocalInvocationID.z); + const int group_offset = int(gl_LocalInvocationID.x); + + const int K = input_sizes.x; + const int M = input_sizes.y; + + // K4 and M4 represent the number of blocks in each dimension. + const int K4 = div_up_4(K); + const int M4 = div_up_4(M); + + const int num_groups = K4 / K4_per_group;; + + if (group_idx >= num_groups || m4 >= M4) { + return; + } + + const int start_k4 = group_idx * K4_per_group + worker_id; + const int end_k4 = (group_idx + 1) * K4_per_group; + + Int8InputScales input_scales; + Int8InputZeroPoints input_zps; + load_int8_input_scales_and_zps(input_scales, input_zps, m4); + + // row of the input tensor to start loading from + const int m = mul_4(m4); + + FPInputTile in_tile; + Int8InputBlock packed_block; + + ivec4 local_sum = ivec4(0, 0, 0, 0); + const int packed_ones = 0x01010101; + + for (int k4 = start_k4; k4 < end_k4; k4 += NUM_WORKERS_PER_GROUP) { + load_input_tile_no_checks(in_tile, k4, m, K4, M); + quantize_and_pack(packed_block, in_tile, input_scales, input_zps); + + // Sum the quantized values in the block + [[unroll]] for (int m = 0; m < TILE_M; m++) { + local_sum[m] += dotPacked4x8AccSatEXT( + packed_block.data[m], packed_ones, local_sum[m]); + } + write_block(packed_block, k4, m4, K4); + } + + shared_sums[group_offset][worker_id] = local_sum; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_GROUP / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_sums[group_offset][worker_id] = + shared_sums[group_offset][worker_id] + + shared_sums[group_offset][worker_id + i]; + } + memoryBarrierShared(); + barrier(); + } + + if (worker_id == 0) { + t_int8_input_sums[group_idx * M4 + m4] = shared_sums[group_offset][0]; + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.yaml new file mode 100644 index 00000000000..3fc66db2718 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_and_pack_linear_input_with_sums.yaml @@ -0,0 +1,30 @@ +# 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. + +quantize_and_pack_linear_input_with_sums: + parameter_names_with_default_values: + DTYPE: float + OUTPUT_STORAGE: buffer + INPUT_STORAGE: texture3d + NUM_GROUPS_PER_WG: 2 + NUM_WORKERS_PER_GROUP: 32 + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: quantize_and_pack_linear_input_with_sums_o2w32_buffer_texture3d + - NAME: quantize_and_pack_linear_input_with_sums_o2w32_buffer_buffer + OUTPUT_STORAGE: buffer + INPUT_STORAGE: buffer + - NAME: quantize_and_pack_linear_input_with_sums_o4w16_buffer_texture3d + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 + - NAME: quantize_and_pack_linear_input_with_sums_o4w16_buffer_buffer + NUM_GROUPS_PER_WG: 4 + NUM_WORKERS_PER_GROUP: 16 + OUTPUT_STORAGE: buffer + INPUT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.glsl index 9a342d8e057..7bf3a932c6c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.glsl @@ -12,12 +12,16 @@ #define IN_T ${buffer_scalar_type(IN_DTYPE)} #define OUT_T ${buffer_scalar_type(OUT_DTYPE)} +#define SCALE_T ${buffer_scalar_type(SCALE_DTYPE)} +#define ZP_T ${buffer_scalar_type(ZP_DTYPE)} #define ${MODE} ${define_active_storage_type("buffer")} ${define_required_extensions(IN_DTYPE)} ${define_required_extensions(OUT_DTYPE)} +${define_required_extensions(SCALE_DTYPE)} +${define_required_extensions(ZP_DTYPE)} layout(std430) buffer; @@ -27,16 +31,16 @@ ${layout_declare_tensor(B, "w", "t_out", OUT_DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "buffer")} $if MODE == "per_tensor": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int quant_min; int quant_max; }; $if MODE == "per_token": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int num_tokens; @@ -44,8 +48,8 @@ $if MODE == "per_token": int quant_max; }; $if MODE == "per_channel": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int axis; @@ -54,8 +58,8 @@ $if MODE == "per_channel": int quant_max; }; $if MODE == "block_wise": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { ivec4 blockSize; // bW, bH, bC, bN @@ -144,7 +148,7 @@ void quantize_per_tensor() { const int in_bufi = tidx_to_bufi(out_tidx, t_in_strides); IN_T value = t_in[in_bufi]; - OUT_T qvalue = quantize_val(value, t_scale[0], t_zero_point[0]); + OUT_T qvalue = quantize_val(value, float(t_scale[0]), int(t_zero_point[0])); t_out[out_bufi] = qvalue; } @@ -179,7 +183,7 @@ void quantize_per_token() { token_idx = min(token_idx, num_tokens - 1); - OUT_T qvalue = quantize_val(value, t_scale[token_idx], t_zero_point[token_idx]); + OUT_T qvalue = quantize_val(value, float(t_scale[token_idx]), int(t_zero_point[token_idx])); t_out[out_bufi] = qvalue; } @@ -218,7 +222,7 @@ void quantize_per_channel() { channel_idx = min(channel_idx, num_channels - 1); - OUT_T qvalue = quantize_val(value, t_scale[channel_idx], t_zero_point[channel_idx]); + OUT_T qvalue = quantize_val(value, float(t_scale[channel_idx]), int(t_zero_point[channel_idx])); t_out[out_bufi] = qvalue; } @@ -241,7 +245,7 @@ void quantize_block_wise() { const int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w; - const OUT_T qvalue = quantize_val(value, t_scale[block_id], t_zero_point[block_id]); + const OUT_T qvalue = quantize_val(value, float(t_scale[block_id]), int(t_zero_point[block_id])); t_out[out_bufi] = qvalue; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.yaml index 5b479c2f90f..fb5853ecd20 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_buffer.yaml @@ -2,6 +2,8 @@ quantize_buffer: parameter_names_with_default_values: IN_DTYPE: float OUT_DTYPE: int32 + SCALE_DTYPE: float + ZP_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: @@ -12,6 +14,12 @@ quantize_buffer: - VALUE: uint8 - VALUE: int8 - VALUE: int32 + SCALE_DTYPE: + - VALUE: float + ZP_DTYPE: + - VALUE: int8 + - VALUE: int32 + - VALUE: float shader_variants: - NAME: quantize_per_tensor_buffer MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.glsl index 69f219ef329..12e5769f50d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.glsl @@ -15,12 +15,16 @@ #define OUT_T ${buffer_scalar_type(OUT_DTYPE)} #define IVEC4_T ${texel_load_type(OUT_DTYPE, "texture3d")} +#define SCALE_T ${buffer_scalar_type(SCALE_DTYPE)} +#define ZP_T ${buffer_scalar_type(ZP_DTYPE)} #define ${MODE} ${define_active_storage_type("texture3d")} ${define_required_extensions(IN_DTYPE)} ${define_required_extensions(OUT_DTYPE)} +${define_required_extensions(SCALE_DTYPE)} +${define_required_extensions(ZP_DTYPE)} #extension GL_EXT_control_flow_attributes : require @@ -32,16 +36,16 @@ ${layout_declare_tensor(B, "w", "t_out", OUT_DTYPE, "texture3d")} ${layout_declare_tensor(B, "r", "t_in", IN_DTYPE, "texture3d")} $if MODE == "per_tensor": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int quant_min; int quant_max; }; $if MODE == "per_token": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int num_tokens; @@ -49,8 +53,8 @@ $if MODE == "per_token": int quant_max; }; $if MODE == "per_channel": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict Block { int axis; @@ -59,8 +63,8 @@ $if MODE == "per_channel": int quant_max; }; $if MODE == "block_wise": - ${layout_declare_tensor(B, "r", "t_scale", "float", "buffer")} - ${layout_declare_tensor(B, "r", "t_zero_point", "int", "buffer")} + ${layout_declare_tensor(B, "r", "t_scale", SCALE_DTYPE, "buffer")} + ${layout_declare_tensor(B, "r", "t_zero_point", ZP_DTYPE, "buffer")} layout(push_constant) uniform restrict BlockPC { ivec4 blockSize; // WHCN @@ -148,7 +152,7 @@ void quantize_per_tensor() { [[unroll]] for (int i = 0; i < 4; ++i) { IN_T value = IN_T(intex[i]); - OUT_T qvalue = quantize_val(value, t_scale[0], t_zero_point[0]); + OUT_T qvalue = quantize_val(value, float(t_scale[0]), int(t_zero_point[0])); outtex[i] = qvalue; } write_texel(t_out, pos, outtex); @@ -180,8 +184,8 @@ void quantize_per_token() { token_idx = min(token_idx, num_tokens - 1); // Scale and zero_point are prepacked as buffers, so direct access - float scale_val = t_scale[token_idx]; - int zero_point_val = t_zero_point[token_idx]; + float scale_val = float(t_scale[token_idx]); + int zero_point_val = int(t_zero_point[token_idx]); IVEC4_T outtex; [[unroll]] for (int i = 0; i < 4; ++i) { @@ -219,8 +223,8 @@ void quantize_per_channel() { int channel_idx = pos.x * 4 + i; channel_idx = min(channel_idx, num_channels - 1); - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); OUT_T qvalue = quantize_val(value, scale_val, zero_point_val); outtex[i] = qvalue; } @@ -228,8 +232,8 @@ void quantize_per_channel() { // Height dimension - all texel components use same channel index int channel_idx = pos.y; channel_idx = min(channel_idx, num_channels - 1); - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T value = IN_T(intex[i]); @@ -243,8 +247,8 @@ void quantize_per_channel() { int folded_idx = pos.z; int channel_idx = folded_idx % num_channels; - float scale_val = t_scale[channel_idx]; - int zero_point_val = t_zero_point[channel_idx]; + float scale_val = float(t_scale[channel_idx]); + int zero_point_val = int(t_zero_point[channel_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T value = IN_T(intex[i]); @@ -258,8 +262,8 @@ void quantize_per_channel() { int folded_idx = pos.z; int batch_idx = folded_idx / num_channels; - float scale_val = t_scale[batch_idx]; - int zero_point_val = t_zero_point[batch_idx]; + float scale_val = float(t_scale[batch_idx]); + int zero_point_val = int(t_zero_point[batch_idx]); [[unroll]] for (int i = 0; i < 4; ++i) { IN_T value = IN_T(intex[i]); @@ -294,7 +298,7 @@ void quantize_block_wise() { int block_id = bcoord.x * blockStride.x + bcoord.y * blockStride.y + bcoord.z * blockStride.z + bcoord.w * blockStride.w; IN_T value = IN_T(intex[i]); - OUT_T qvalue = quantize_val(value, t_scale[block_id], t_zero_point[block_id]); + OUT_T qvalue = quantize_val(value, float(t_scale[block_id]), int(t_zero_point[block_id])); outtex[i] = qvalue; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.yaml index 2e40ac90794..03d418ff2f7 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/quantize_texture.yaml @@ -2,6 +2,8 @@ quantize_texture: parameter_names_with_default_values: IN_DTYPE: float OUT_DTYPE: int32 + SCALE_DTYPE: float + ZP_DTYPE: int32 MODE: per_tensor generate_variant_forall: IN_DTYPE: @@ -12,6 +14,12 @@ quantize_texture: - VALUE: uint8 - VALUE: int8 - VALUE: int32 + SCALE_DTYPE: + - VALUE: float + ZP_DTYPE: + - VALUE: int8 + - VALUE: int32 + - VALUE: float shader_variants: - NAME: quantize_per_tensor_texture3d MODE: per_tensor diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl new file mode 100644 index 00000000000..98370a9bcde --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl @@ -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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, STORAGE)} + +${define_active_storage_type(STORAGE)} + +#extension GL_EXT_control_flow_attributes : require + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "tout", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "tin", DTYPE, STORAGE)} + +${layout_declare_ubo(B, "ivec3", "tin_limits")} +${layout_declare_ubo(B, "ivec4", "tin_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +layout(constant_id = 3) const int packed_dim = 0; +layout(constant_id = 4) const int reduce_dim1 = 0; +layout(constant_id = 5) const int reduce_dim2 = 1; +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 + +// 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 + + +shared vec4 shared_vecs[MAX_NTHREADS]; + +#include "indexing_utils.h" + +int tid_to_smi(const ivec2 tid) { + return tid.x + tid.y * NWORKERS; +} + +// Initializing the accumulator accepts the first value in the reduction row, +// since some reduction operations (i.e. amax, amin) prefer to initialize with +// a data point instead of a static value. +#define INIT_ACCUM(first_val) ${INIT_ACCUM} +#define UPDATE_ACCUM(accum, new_val) ${UPDATE_ACCUM} +// Useful for operators such as mean which want to perform a final calculation +// with the accumulator. +#define POSTPROCESS(accum) ${POSTPROCESS} + +void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { + // 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 < tin_sizes[reduce_dim1]; + i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { + + // Second dimension reduction + scan_pos[reduce_dim2] = 0; + for (int j = 0; j < tin_sizes[reduce_dim2]; j++, scan_pos[reduce_dim2]++) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } + } + + // Write partial output to shared memory and synchronize + shared_vecs[smi] = accum; + barrier(); + + // Main thread aggregates results + if (tid.x == 0) { + // Iterate over the partial outputs to obtain the overall output + int group_i = tid.y * NWORKERS; + accum = shared_vecs[group_i++]; + for (int i = 1; i < NWORKERS; i++, group_i++) { + accum = UPDATE_ACCUM(accum, shared_vecs[group_i]); + } + + // Determine if there are any padding elements in the final texel of the + // packed dimension + const int nspill = mod4(tin_sizes[packed_dim]); + // Detect if this thread is working on the final texels of the packed + // dimension, which may have padding elements + const bool is_last_texel = + scan_pos[packed_dim] == (tin_limits[packed_dim] - 1); + + // Explicitly set padding elements to 0 + if (is_last_texel && nspill > 0) { + [[unroll]] for (int i = nspill; i < 4; i++) { + accum[i] = 0; + } + } + scan_pos[reduce_dim1] = 0; + scan_pos[reduce_dim2] = 0; + write_texel(tout, scan_pos, POSTPROCESS(accum)); + } +} + +void main() { + ivec3 scan_pos = ivec3(gl_GlobalInvocationID); + scan_pos[reduce_dim1] = 0; + scan_pos[reduce_dim2] = 0; + + const ivec2 tid = ivec2( + gl_LocalInvocationID[reduce_dim1], + gl_LocalInvocationID[group_dim]); + + if (any(greaterThanEqual(scan_pos, tin_limits))) { + return; + } + + reduce_2d_non_packed_dim(tid, scan_pos); +} \ No newline at end of file diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml new file mode 100644 index 00000000000..fdc5eb9f105 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.yaml @@ -0,0 +1,29 @@ +# 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. + +reduce2d: + parameter_names_with_default_values: + DTYPE: float + STORAGE: texture3d + INIT_ACCUM: VEC4_T(0) + UPDATE_ACCUM: accum + new_val + POSTPROCESS: accum + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: sum2d + - NAME: mean2d + POSTPROCESS: (accum / (tin_sizes[reduce_dim1] * tin_sizes[reduce_dim2])) + - NAME: amax2d + INIT_ACCUM: first_val + UPDATE_ACCUM: max(accum, new_val) + POSTPROCESS: accum + - NAME: amin2d + INIT_ACCUM: first_val + UPDATE_ACCUM: min(accum, new_val) + POSTPROCESS: accum diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.glsl deleted file mode 100644 index 1e854bf7f85..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.glsl +++ /dev/null @@ -1,120 +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. - */ - -#version 450 core - -#define PRECISION ${PRECISION} - -#define T ${buffer_scalar_type(DTYPE)} - -${define_active_storage_type(STORAGE)} -${define_required_extensions(DTYPE)} - -#extension GL_EXT_control_flow_attributes : require - -layout(std430) buffer; - -${layout_declare_tensor(B, "rw", "attn_weight", DTYPE, STORAGE)} - -$if STORAGE == "buffer": - ${layout_declare_ubo(B, "ivec4", "attn_weight_sizes")} - ${layout_declare_ubo(B, "ivec4", "attn_weight_strides")} -$else: - ${layout_declare_ubo(B, "ivec3", "attn_weight_limits")} - -${layout_declare_ubo(B, "int", "input_pos")} -${layout_declare_ubo(B, "float", "scale")} - - -#include "indexing_utils.h" - -layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; - -// Negative infinity is represented by having sign bit be 1, all exponent bits -// be 1, all mantissa bits be 0. -#define NEGATIVE_INF_BITS 0xFF800000 -const float negative_infinity = NEGATIVE_INF_BITS; - -#ifdef USING_BUFFER - -/* - * This implementations applies a scale and mask to the attention weight tensor - * of an SDPA block. The sizes of the attention weight is - * (batch_size, n_heads, seq_len, input_pos + seq_len) - * Conceptually the weights represent the relationship between each token in the - * sequence with each token preceding it. - * - * The scale applied is 1.0 / sqrt(head_dim_length) - * - * The mask applied is a bit more complicated. Imagine you create a square - * matrix of size (input_pos + seq_len, input_pos + seq_len), and then set the - * lower triangular section of the matrix to -inf. Then, slice the matrix along - * the row dimension starting from input_pos to input_pos + seq_len. You end up - * with a partial mask with size (seq_len, input_pos + seq_len). This is the - * mask that is applied to the attention weight. - * - * In the shader, instead of generating the mask, the index of the elment is - * inspected to determine if it would have been masked. Given an element at - * tensor index (n, c, h, w), it would be masked if w < h + input_pos. - */ - -/*************************** - ** Buffer Implementation ** - ***************************/ - -void main() { - const ivec4 attn_weight_idx = ivec4( - gl_GlobalInvocationID.x, - gl_GlobalInvocationID.y, - gl_GlobalInvocationID.z, - 0); - - if (any(greaterThanEqual(attn_weight_idx, attn_weight_sizes))) { - return; - } - - const T scale_conv = T(scale); - - const int attn_weight_id = tidx_to_bufi(attn_weight_idx, attn_weight_strides); - if (attn_weight_idx.x <= attn_weight_idx.y + input_pos) { - attn_weight[attn_weight_id] = attn_weight[attn_weight_id] * scale_conv; - } else { - attn_weight[attn_weight_id] = T(negative_infinity); - } -} - -#else - -/**************************** - ** Texture Implementation ** - ****************************/ - -/* - * This implementation assumes that the attention weight is width packed, i.e. - * the packed dim of the attn_weight is 0. - */ -void main() { - const ivec3 attn_weight_pos = ivec3(gl_GlobalInvocationID); - - if (any(greaterThanEqual(attn_weight_pos, attn_weight_limits))) { - return; - } - - vec4 outtex = imageLoad(attn_weight, attn_weight_pos) * scale; - - // Mask out the upper triangular of attn_weight to -inf - [[unroll]] for (int i = 0; i < 4; ++i) { - if (attn_weight_pos.x * 4 + i > attn_weight_pos.y + input_pos) { - outtex[i] = negative_infinity; - } - } - - write_texel(attn_weight, attn_weight_pos, outtex); -} - -#endif // USING_BUFFER diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.yaml deleted file mode 100644 index ca8806fe000..00000000000 --- a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weight_scale_and_mask.yaml +++ /dev/null @@ -1,13 +0,0 @@ -sdpa_attn_weight_scale_and_mask: - parameter_names_with_default_values: - DTYPE: float - STORAGE: buffer - generate_variant_forall: - STORAGE: - - VALUE: buffer - - VALUE: texture3d - DTYPE: - - VALUE: half - - VALUE: float - shader_variants: - - NAME: sdpa_attn_weight_scale_and_mask diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl new file mode 100644 index 00000000000..67d9c100f68 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.glsl @@ -0,0 +1,166 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, STORAGE)} +#define T ${texel_load_component_type(DTYPE, STORAGE)} + +#define NUM_WORKERS_PER_WG 64 + +${define_active_storage_type(STORAGE)} + +${define_required_extensions(DTYPE)} + +#extension GL_EXT_control_flow_attributes : require + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_attn_weights_softmax", DTYPE, STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_attn_weights", DTYPE, STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_projected_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +// Shared memory for cooperative exp sum finding +shared T shared_exp_sum[NUM_WORKERS_PER_WG]; + +VEC4_T load_attn_weights_c4( + const int c4, + const int s, + const int q_h, + const int C4, + const int S, + const int Q_H) { +#ifdef USING_BUFFER + return t_attn_weights[(q_h * S * C4) + (s * C4) + c4]; +#else + return texelFetch(t_attn_weights, ivec3(c4, s, q_h), 0); +#endif +} + +void store_attn_weights_softmax_c4( + const VEC4_T out_texel, + const int c4, + const int s, + const int q_h, + const int C4, + const int S, + const int Q_H) { +#ifdef USING_BUFFER + t_attn_weights_softmax[(q_h * S * C4) + (s * C4) + c4] = out_texel; +#else + imageStore(t_attn_weights_softmax, ivec3(c4, s, q_h), out_texel); +#endif +} + +void main() { + const int worker_id = int(gl_LocalInvocationID.x); + + // Index along attention weight's sequence_len dim + const int s = int(gl_GlobalInvocationID.y); + // idx along attention weight's num_q_heads dim + const int q_h = int(gl_GlobalInvocationID.z); + + // number of Q heads + const int Q_H = q_projected_sizes.y; + // sequence length + const int S = q_projected_sizes.z; + // manually determine size of the context_len dim of the attention weight. + // The "actual" tensor sizes may have been aligned to a multiple of 4 to allow + // memory loads to be aligned to texel boundaries. + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + if (s >= S || q_h >= Q_H) { + return; + } + + // Initialize thread-local min/max + T local_exp_sum = T(0); + + const int context_len_aligned_down = context_len - mod_4(context_len); + const int C4_limit = div_4(context_len_aligned_down); + + // Each thread processes elements along a context_len row with a stride of the + // number of threads in the work group. + for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_WG) { + VEC4_T in_texel = load_attn_weights_c4( + c4, s, q_h, context_texel_len, S, Q_H); + + for (int comp = 0; comp < 4; comp++) { + local_exp_sum += exp(in_texel[comp]); + } + } + // First thread in the work group responsible for handling last texel if it + // contains any padded elements + if (worker_id == 0) { + for (int c4 = C4_limit; c4 < context_texel_len; ++c4) { + const int c_base = mul_4(c4); + VEC4_T in_texel = load_attn_weights_c4( + c4, s, q_h, context_texel_len, S, Q_H); + + [[unroll]] for (int comp = 0; comp < 4; comp++) { + if (c_base + comp < context_len) { + local_exp_sum += exp(in_texel[comp]); + } + } + } + } + + // Store thread-local results in shared memory + shared_exp_sum[worker_id] = local_exp_sum; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result + for (int i = NUM_WORKERS_PER_WG / 2; i > 0; i >>= 1) { + if (worker_id < i) { + shared_exp_sum[worker_id] = shared_exp_sum[worker_id] + + shared_exp_sum[worker_id + i]; + } + memoryBarrierShared(); + barrier(); + } + + local_exp_sum = shared_exp_sum[0]; + // Now go back through each element in the row and normalize + for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_WG) { + VEC4_T in_texel = load_attn_weights_c4( + c4, s, q_h, context_texel_len, S, Q_H); + + VEC4_T out_texel = exp(in_texel) / local_exp_sum; + store_attn_weights_softmax_c4( + out_texel, c4, s, q_h, context_texel_len, S, Q_H); + } + // First thread in the work group responsible for handling last texel if it + // contains any padded elements + if (worker_id == 0) { + for (int c4 = C4_limit; c4 < context_texel_len; ++c4) { + const int c_base = mul_4(c4); + VEC4_T in_texel = load_attn_weights_c4( + c4, s, q_h, context_texel_len, S, Q_H); + + // Ensure that padding elements are set to 0. + VEC4_T out_texel = VEC4_T(0); + [[unroll]] for (int comp = 0; comp < 4; comp++) { + if (c_base + comp < context_len) { + out_texel[comp] = exp(in_texel[comp]) / local_exp_sum; + } + } + store_attn_weights_softmax_c4( + out_texel, c4, s, q_h, context_texel_len, S, Q_H); + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.yaml new file mode 100644 index 00000000000..66ec030680e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_attn_weights_softmax.yaml @@ -0,0 +1,19 @@ +# 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. + +sdpa_attn_weights_softmax: + parameter_names_with_default_values: + DTYPE: float + STORAGE: texture3d + generate_variant_forall: + STORAGE: + - VALUE: texture3d + - VALUE: buffer + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: sdpa_attn_weights_softmax diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl new file mode 100644 index 00000000000..2900d63666b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl @@ -0,0 +1,211 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if K_CACHE_STORAGE == "buffer": + #define K_CACHE_BUFFER + +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M 1 +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +#define NUM_WORKERS_PER_OUT 64 + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_q_projected", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_k_cache", DTYPE, K_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_projected_sizes")} +${layout_declare_ubo(B, "ivec4", "k_cache_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} + +#include "sdpa_fp_q_projected_tile_load.glslh" +#include "sdpa_fp_k_cache_tile_load.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "sdpa_fp_attn_weight_tile_store.glslh" + +shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; + +/* + * See the tiled variant of this shader for the implemented behavior. This + * shader is implements an optimization for cases where sequence length is 1; in + * these cases, the matrix multiplication being performed is akin to gemv, which + * benefits from using a co-operative algorithm for reduction. For this shader + * the entire work group co-operates to compute one reduction output. + */ + +void main() { + const int worker_id = int(gl_LocalInvocationID.y); + + const int tile_idx_x = int(gl_GlobalInvocationID.x); + // idx along output num_q_heads dim + const int q_h = int(gl_GlobalInvocationID.z); + + // idx along the output context_len dim + const int c = tile_idx_x * TILE_N; + const int c4 = div_4(c); + + // idx along the output seq_len dim. Note that for this shader seq_len will be + // 1. + const int s = 0; + + // texel size of head_dim, over which the dot product is accumulated + const int D4 = div_up_4(q_projected_sizes.x); + // number of Q heads + const int Q_H = q_projected_sizes.y; + // sequence length + const int S = q_projected_sizes.z; + + // number of K/V heads + const int KV_H = k_cache_sizes.y; + // Max context length + const int C = k_cache_sizes.z; + const int C4 = div_up_4(C); + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + // current context length + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + // bounds check + if (c >= context_len || s >= S || q_h >= Q_H) { + return; + } + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile q_tile; + FPWeightTile w_tile; + + // If the tile is completely inside the mask region, then there is no need to + // compute the output tile. All the elements in the output tile can be set to + // negative infinity. + bool tile_in_mask_region = c > (input_pos + s + (TILE_M - 1)); + if (tile_in_mask_region) { + const VEC4_T negative_infinity_vec = VEC4_T(negative_infinity_val); + set_out_tile_to_vec(out_tile, negative_infinity_vec); + } + // Otherwise, need to actually compute output tile + else { + const bool dont_check_bounds = (S - s) >= TILE_M && + (context_len - c) >= TILE_N; + + if (dont_check_bounds) { + for (int d4 = worker_id; d4 < D4; d4 += NUM_WORKERS_PER_OUT) { + load_q_projected_tile_no_checks( + q_tile, + d4, + s, + q_h, + D4, + Q_H, + S); + + load_k_cache_tile_no_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile); + } + } else { + for (int d4 = worker_id; d4 < D4; d4 += NUM_WORKERS_PER_OUT) { + load_q_projected_tile_with_checks( + q_tile, + d4, + s, + q_h, + D4, + Q_H, + S); + + load_k_cache_tile_with_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile); + } + } + } + + partial_sums[worker_id] = out_tile; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result. + for (int i = NUM_WORKERS_PER_OUT / 2; i > 0; i /= 2) { + if (worker_id < i) { + accumulate_out_tile_with_out_tile( + partial_sums[worker_id], partial_sums[worker_id + i]); + } + memoryBarrierShared(); + barrier(); + } + + // Only the first thread will write out the result + if (worker_id == 0) { + out_tile = partial_sums[0]; + // Apply scale and mask if the tile was not entirely in the mask region + if (!tile_in_mask_region) { + VEC4_T inv_scale_vec = VEC4_T(inv_scale); + apply_scale_and_mask( + out_tile, + inv_scale_vec, + input_pos, + c, + s); + } + + store_attn_weight_tile_with_checks( + out_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml similarity index 52% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.yaml rename to backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml index 04e803a2e94..6a4cffcc913 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_coop.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml @@ -4,16 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -linear_qga4w_coop: +sdpa_compute_attn_weights_coop: parameter_names_with_default_values: DTYPE: float IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d - WGS: 64 + K_CACHE_STORAGE: texture3d + TILE_K4: 1 + TILE_N4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half shader_variants: - - NAME: linear_qga4w_coop_texture3d_texture3d_texture2d_float - - NAME: linear_qga4w_coop_buffer_buffer_texture2d_float + - NAME: sdpa_compute_attn_weights_coop_texture3d_texture3d + - NAME: sdpa_compute_attn_weights_coop_buffer_texture3d IO_STORAGE: buffer - - NAME: linear_qga4w_coop_buffer_buffer_buffer_float - IO_STORAGE: buffer - WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl new file mode 100644 index 00000000000..95c22d91b80 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl @@ -0,0 +1,201 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if K_CACHE_STORAGE == "buffer": + #define K_CACHE_BUFFER + +#define TILE_M4 ${TILE_M4} +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_q_projected", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_k_cache", DTYPE, K_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_projected_sizes")} +${layout_declare_ubo(B, "ivec4", "k_cache_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} + +#include "sdpa_fp_q_projected_tile_load.glslh" +#include "sdpa_fp_k_cache_tile_load.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "sdpa_fp_attn_weight_tile_store.glslh" + +/* + * Compute attention weights given the q_projected and k_cache tensors. + * q_projected has shape (batches, seq_len, num_q_heads, head_dim) + * k_cache has shape (batches, max_context_len, num_kv_heads, head_dim) + * output has shape (batches, num_q_heads, seq_len, context_len) + * + * This shader also applies scales and masking to the computed attention + * weights. + * + * The scale applied is 1.0 / sqrt(head_dim_length). + * + * The mask applied is a bit more complicated. Imagine you create a square + * matrix of size (input_pos + seq_len, input_pos + seq_len), and then set the + * lower triangular section of the matrix to -inf. Then, slice the matrix along + * the row dimension starting from input_pos to input_pos + seq_len. You end up + * with a partial mask with size (seq_len, input_pos + seq_len). This is the + * mask that is applied to the attention weight. + * + * In the shader, instead of generating the mask, the index of the elment is + * inspected to determine if it would have been masked. Given an element at + * tensor index (n, c, h, w), it would be masked if w < h + input_pos. + * + */ + +void main() { + const int tile_idx_x = int(gl_GlobalInvocationID.x); + const int tile_idx_y = int(gl_GlobalInvocationID.y); + // idx along output num_q_heads dim + const int q_h = int(gl_GlobalInvocationID.z); + + // idx along the output context_len dim + const int c = tile_idx_x * TILE_N; + const int c4 = div_4(c); + + // idx along the output seq_len dim + const int s = tile_idx_y * TILE_M; + const int s4 = div_4(s); + + // texel size of head_dim, over which the dot product is accumulated + const int D4 = div_up_4(q_projected_sizes.x); + // number of Q heads + const int Q_H = q_projected_sizes.y; + // sequence length + const int S = q_projected_sizes.z; + + // number of K/V heads + const int KV_H = k_cache_sizes.y; + // Max context length + const int C = k_cache_sizes.z; + const int C4 = div_up_4(C); + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + // bounds check + if (c >= context_len || s >= S || q_h >= Q_H) { + return; + } + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile q_tile; + FPWeightTile w_tile; + + // If the tile is completely inside the mask region, then there is no need to + // compute the output tile. All the elements in the output tile can be set to + // negative infinity. + bool tile_in_mask_region = c > (input_pos + s + (TILE_M - 1)); + if (tile_in_mask_region) { + const VEC4_T negative_infinity_vec = VEC4_T(negative_infinity_val); + set_out_tile_to_vec(out_tile, negative_infinity_vec); + } + // Otherwise, need to actually compute output tile + else { + const bool dont_check_bounds = (S - s) >= TILE_M && + (context_len - c) >= TILE_N; + + if (dont_check_bounds) { + for (int d4 = 0; d4 < D4; d4++) { + load_q_projected_tile_no_checks( + q_tile, + d4, + s, + q_h, + D4, + Q_H, + S); + + load_k_cache_tile_no_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile); + } + } else { + for (int d4 = 0; d4 < D4; d4++) { + load_q_projected_tile_with_checks( + q_tile, + d4, + s, + q_h, + D4, + Q_H, + S); + + load_k_cache_tile_with_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile); + } + } + + // Apply scale and mask + VEC4_T inv_scale_vec = VEC4_T(inv_scale); + apply_scale_and_mask( + out_tile, + inv_scale_vec, + input_pos, + c, + s); + } + + store_attn_weight_tile_with_checks( + out_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml new file mode 100644 index 00000000000..6aadbbc379e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml @@ -0,0 +1,22 @@ +# 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. + +sdpa_compute_attn_weights_tiled: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + K_CACHE_STORAGE: texture3d + TILE_M4: 1 + TILE_K4: 1 + TILE_N4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: sdpa_compute_attn_weights_tiled_texture3d_texture3d + - NAME: sdpa_compute_attn_weights_tiled_buffer_texture3d + IO_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl new file mode 100644 index 00000000000..5f408b7581d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.glsl @@ -0,0 +1,193 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if V_CACHE_STORAGE == "buffer": + #define V_CACHE_BUFFER + +#define TILE_K4 ${TILE_K4} +#define TILE_N4 ${TILE_N4} + +#define TILE_M 1 +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +#define NUM_WORKERS_PER_OUT 64 + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_v_cache", DTYPE, V_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_projected_sizes")} +${layout_declare_ubo(B, "ivec4", "v_cache_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +${layout_declare_spec_const(C, "float", "inv_scale", "1.0")} + +#include "sdpa_fp_attn_weight_tile_load.glslh" +#include "sdpa_fp_v_cache_tile_load.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "sdpa_fp_out_tile_store.glslh" + +shared FPOutTile partial_sums[NUM_WORKERS_PER_OUT]; + +/* + * See the tiled variant of this shader for the implemented behavior. This + * shader is implements an optimization for cases where sequence length is 1; in + * these cases, the matrix multiplication being performed is akin to gemv, which + * benefits from using a co-operative algorithm for reduction. For this shader + * the entire work group co-operates to compute one reduction output. + */ + +void main() { + const int worker_id = int(gl_LocalInvocationID.y); + + const int tile_idx_x = int(gl_GlobalInvocationID.x); + // idx along output num_q_heads dim + const int q_h = int(gl_GlobalInvocationID.z); + + // idx along the output head_dim dim + const int d = tile_idx_x * TILE_N; + const int d4 = div_4(d); + + // idx along the output seq_len dim. Note that for this shader seq_len will be + // 1. + const int s = 0; + + // texel size of head_dim + const int D4 = div_up_4(q_projected_sizes.x); + // number of Q heads + const int Q_H = q_projected_sizes.y; + // sequence length + const int S = q_projected_sizes.z; + + // number of K/V heads + const int KV_H = v_cache_sizes.y; + // Max context length + const int C = v_cache_sizes.z; + const int C4 = div_up_4(C); + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + // current context length + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + // bounds check + if (d4 >= D4 || s >= S || q_h >= Q_H) { + return; + } + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile attn_weight_tile; + FPWeightTile w_tile; + + const int context_len_aligned_down = context_len - mod_4(context_len); + const int C4_limit = div_up_4(context_len_aligned_down); + + for (int c4 = worker_id; c4 < C4_limit; c4 += NUM_WORKERS_PER_OUT) { + const int c = mul_4(c4); + + load_attn_weight_tile_no_checks( + attn_weight_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); + + load_v_cache_tile_no_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); + } + // first worker in the work group will handle final texel, which may contain + // padding elements. + if (worker_id == 0) { + for (int c4 = C4_limit; c4 < context_texel_len; c4++) { + const int c = mul_4(c4); + load_attn_weight_tile_with_checks( + attn_weight_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); + + load_v_cache_tile_with_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); + } + } + + partial_sums[worker_id] = out_tile; + + memoryBarrierShared(); + barrier(); + + // Tree reduction to compute the overall result. + for (int i = NUM_WORKERS_PER_OUT / 2; i > 0; i /= 2) { + if (worker_id < i) { + accumulate_out_tile_with_out_tile( + partial_sums[worker_id], partial_sums[worker_id + i]); + } + memoryBarrierShared(); + barrier(); + } + + // Only the first thread will write out the result + if (worker_id == 0) { + out_tile = partial_sums[0]; + store_sdpa_out_tile_with_checks( + out_tile, + d4, + s, + q_h, + D4, + S, + Q_H); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml similarity index 53% rename from backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.yaml rename to backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml index 94d10dcf978..ccebf8f7c1c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/linear_qga4w_tiled.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_coop.yaml @@ -4,15 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -linear_qga4w_tiled: +sdpa_compute_out_coop: parameter_names_with_default_values: DTYPE: float IO_STORAGE: texture3d - WEIGHT_STORAGE: texture2d + V_CACHE_STORAGE: texture3d + TILE_K4: 1 + TILE_N4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half shader_variants: - - NAME: linear_qga4w_tiled_texture3d_texture3d_texture2d_float - - NAME: linear_qga4w_tiled_buffer_buffer_texture2d_float + - NAME: sdpa_compute_out_coop_texture3d_texture3d + - NAME: sdpa_compute_out_coop_buffer_texture3d IO_STORAGE: buffer - - NAME: linear_qga4w_tiled_buffer_buffer_buffer_float - IO_STORAGE: buffer - WEIGHT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl new file mode 100644 index 00000000000..0063ebf9d38 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.glsl @@ -0,0 +1,163 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} +#define VEC4_T ${texel_load_type(DTYPE, IO_STORAGE)} +#define T ${texel_load_component_type(DTYPE, IO_STORAGE)} + +$if IO_STORAGE == "buffer": + #define OUTPUT_BUFFER + #define INPUT_BUFFER +$if V_CACHE_STORAGE == "buffer": + #define V_CACHE_BUFFER + +#define TILE_M4 ${TILE_M4} +// Equvalent to K4 in matrix multiplication +#define TILE_K4 ${TILE_K4} +// Equvalent to N4 in matrix multiplication +#define TILE_N4 ${TILE_N4} + +#define TILE_M ${TILE_M4 * 4} +#define TILE_K ${TILE_K4 * 4} +#define TILE_N ${TILE_N4 * 4} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_output", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_v_cache", DTYPE, V_CACHE_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "q_projected_sizes")} +${layout_declare_ubo(B, "ivec4", "v_cache_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#include "sdpa_fp_attn_weight_tile_load.glslh" +#include "sdpa_fp_v_cache_tile_load.glslh" +#include "linear_fp_output_tile_fp_compute.glslh" +#include "sdpa_fp_out_tile_store.glslh" + +/* + * Compute SDPA output given the attention weights and v_cache tensors. + * attention weights has shape (batches, num_q_heads, seq_len, context_len) + * v_cache has shape (batches, max_context_len, num_kv_heads, head_dim) + * output has shape (batches, seq_len, num_q_heads, head_dim) + */ + +void main() { + const int tile_idx_x = int(gl_GlobalInvocationID.x); + const int tile_idx_y = int(gl_GlobalInvocationID.y); + // idx along output num_q_heads dim + const int q_h = int(gl_GlobalInvocationID.z); + + // idx along the output head_dim dim + const int d = tile_idx_x * TILE_N; + const int d4 = div_4(d); + + // idx along the output seq_len dim + const int s = tile_idx_y * TILE_M; + + // texel size of head_dim + const int D4 = div_up_4(q_projected_sizes.x); + // number of Q heads + const int Q_H = q_projected_sizes.y; + // sequence length + const int S = q_projected_sizes.z; + + // number of K/V heads + const int KV_H = v_cache_sizes.y; + // Max context length + const int C = v_cache_sizes.z; + const int C4 = div_up_4(C); + + int kv_h = q_h; + if (KV_H < Q_H) { + kv_h = q_h / (Q_H / KV_H); + } + + // current context length + const int context_len = input_pos + S; + const int context_texel_len = div_up_4(context_len); + + // bounds check + if (d4 >= D4 || s >= S || q_h >= Q_H) { + return; + } + + FPOutTile out_tile; + initialize(out_tile); + + FPInputTile attn_weight_tile; + FPWeightTile w_tile; + + const int context_len_aligned_down = context_len - mod_4(context_len); + const int C4_limit = div_4(context_len_aligned_down); + + for (int c4 = 0; c4 < C4_limit; c4++) { + const int c = mul_4(c4); + load_attn_weight_tile_no_checks( + attn_weight_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); + + load_v_cache_tile_no_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); + } + for (int c4 = C4_limit; c4 < context_texel_len; c4++) { + const int c = mul_4(c4); + load_attn_weight_tile_with_checks( + attn_weight_tile, + c4, + s, + q_h, + context_texel_len, + S, + Q_H); + + load_v_cache_tile_with_checks( + w_tile, + d4, + c, + kv_h, + D4, + context_len, + C, + KV_H); + + fp_accumulate_with_fp_weight(out_tile, attn_weight_tile, w_tile); + } + + store_sdpa_out_tile_with_checks( + out_tile, + d4, + s, + q_h, + D4, + S, + Q_H); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml new file mode 100644 index 00000000000..7fbce29e908 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_out_tiled.yaml @@ -0,0 +1,22 @@ +# 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. + +sdpa_compute_out_tiled: + parameter_names_with_default_values: + DTYPE: float + IO_STORAGE: texture3d + V_CACHE_STORAGE: texture3d + TILE_M4: 1 + TILE_K4: 1 + TILE_N4: 1 + generate_variant_forall: + DTYPE: + - VALUE: float + - VALUE: half + shader_variants: + - NAME: sdpa_compute_out_tiled_texture3d_texture3d + - NAME: sdpa_compute_out_tiled_buffer_texture3d + IO_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh new file mode 100644 index 00000000000..12b2292fa45 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_load.glslh @@ -0,0 +1,74 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_attn_weights + * + * Macro Settings: + * - INPUT_BUFFER + */ + +#ifndef SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH +#define SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_input_tile.glslh" + +VEC4_T load_attn_weight_c4( + const int c4, + const int s, + const int q_h, + const int C4, + const int S, + const int Q_H) { +#ifdef INPUT_BUFFER + return t_attn_weights[(q_h * S * C4) + (s * C4) + c4]; +#else + return texelFetch(t_attn_weights, ivec3(c4, s, q_h), 0); +#endif +} + +void load_attn_weight_tile_no_checks( + out FPInputTile tile, + const int c4_start, + const int s_start, + const int q_h, + const int C4, + const int S, + const int Q_H) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { + tile.data[s][c4] = + load_attn_weight_c4(c4_start + c4, s_start + s, q_h, C4, S, Q_H); + } + } +} + +void load_attn_weight_tile_with_checks( + out FPInputTile tile, + const int c4_start, + const int s_start, + const int q_h, + const int C4, + const int S, + const int Q_H) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { + if (c4_start + c4 < C4 && s_start + s < S) { + tile.data[s][c4] = + load_attn_weight_c4(c4_start + c4, s_start + s, q_h, C4, S, Q_H); + } else { + tile.data[s][c4] = VEC4_T(0.0); + } + } + } +} + +#endif // SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh new file mode 100644 index 00000000000..c64d9af8cfb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh @@ -0,0 +1,105 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_attn_weights + * + * Macro Settings: + * - OUTPUT_BUFFER + */ + +#ifndef SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH +#define SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_output_tile.glslh" + +T negative_infinity_val = T(-1.0 / 0.0); + +void store_attn_weight_c4( + const VEC4_T out_texel, + const int c4, + const int s, + const int q_h, + const int C4, + const int S, + const int Q_H) { +#ifdef OUTPUT_BUFFER + t_attn_weights[(q_h * S * C4) + (s * C4) + c4] = out_texel; +#else + imageStore(t_attn_weights, ivec3(c4, s, q_h), out_texel); +#endif +} + +void store_attn_weight_tile_no_checks( + const FPOutTile tile, + const int c4_start, + const int s_start, + const int q_h, + const int C4, + const int S, + const int Q_H) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { + store_attn_weight_c4( + tile.data[s][c4], c4_start + c4, s_start + s, q_h, C4, S, Q_H); + } + } +} + +void store_attn_weight_tile_with_checks( + const FPOutTile tile, + const int c4_start, + const int s_start, + const int q_h, + const int C4, + const int S, + const int Q_H) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { + if (c4_start + c4 < C4 && s_start + s < S) { + store_attn_weight_c4( + tile.data[s][c4], c4_start + c4, s_start + s, q_h, C4, S, Q_H); + } + } + } +} + +void set_out_tile_to_vec(out FPOutTile tile, const VEC4_T vec) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { tile.data[s][c4] = vec; } + } +} + +void apply_scale_and_mask( + inout FPOutTile tile, + const VEC4_T inv_scale_vec, + const int input_pos, + const int c_idx_start, + const int s_idx_start) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { + tile.data[s][c4] = tile.data[s][c4] * inv_scale_vec; + + const int c_base = mul_4(c4); + [[unroll]] for (int c4i = 0; c4i < 4; ++c4i) { + const int c = c_base + c4i; + // Indices of the tile element in the overall output tensor + const int c_idx = c_idx_start + c; + const int s_idx = s_idx_start + s; + if (c_idx > s_idx + input_pos) { + tile.data[s][c4][c4i] = negative_infinity_val; + } + } + } + } +} + +#endif // SDPA_FP_ATTN_WEIGHT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh new file mode 100644 index 00000000000..03132db1348 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_k_cache_tile_load.glslh @@ -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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_k_cache + * + * Macro Settings: + * - K_CACHE_BUFFER + */ + +#ifndef SDPA_FP_K_CACHE_TILE_LOAD_GLSLH +#define SDPA_FP_K_CACHE_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_weight_tile.glslh" + +VEC4_T load_k_cache_d4( + const int d4, + const int c, + const int kv_h, + const int D4, + const int C, + const int KV_H) { +#ifdef K_CACHE_BUFFER + return VEC4_T(t_k_cache[(c * KV_H * D4) + (kv_h * D4) + d4]); +#else + return VEC4_T(texelFetch(t_k_cache, ivec3(d4, kv_h, c), 0)); +#endif +} + +void load_k_cache_tile_no_checks( + out FPWeightTile tile, + const int d4_start, + const int c_start, + const int kv_h, + const int D4, + const int context_len, + const int C, + const int KV_H) { + bool should_print = d4_start == 0 && c_start == 0 && kv_h == 0; + [[unroll]] for (int c = 0; c < TILE_N; ++c) { + const int c4 = div_4(c); + const int c4i = mod_4(c); + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + VEC4_T d4_row = + load_k_cache_d4(d4_start + d4, c_start + c, kv_h, D4, C, KV_H); + + // Transpose in-place + const int d_base = mul_4(d4); + tile.data[d_base][c4][c4i] = d4_row[0]; + tile.data[d_base + 1][c4][c4i] = d4_row[1]; + tile.data[d_base + 2][c4][c4i] = d4_row[2]; + tile.data[d_base + 3][c4][c4i] = d4_row[3]; + } + } +} + +void load_k_cache_tile_with_checks( + out FPWeightTile tile, + const int d4_start, + const int c_start, + const int kv_h, + const int D4, + const int context_len, + const int C, + const int KV_H) { + [[unroll]] for (int c = 0; c < TILE_N; ++c) { + const int c4 = div_4(c); + const int c4i = mod_4(c); + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + VEC4_T d4_row = VEC4_T(0.0); + if (d4_start + d4 < D4 && c_start + c < context_len) { + d4_row = load_k_cache_d4(d4_start + d4, c_start + c, kv_h, D4, C, KV_H); + } + + // Transpose in-place + const int d_base = mul_4(d4); + tile.data[d_base][c4][c4i] = d4_row[0]; + tile.data[d_base + 1][c4][c4i] = d4_row[1]; + tile.data[d_base + 2][c4][c4i] = d4_row[2]; + tile.data[d_base + 3][c4][c4i] = d4_row[3]; + } + } +} + +#endif // SDPA_FP_K_CACHE_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh new file mode 100644 index 00000000000..17e0988a6a4 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_out_tile_store.glslh @@ -0,0 +1,57 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_attn_weights + * + * Macro Settings: + * - OUTPUT_BUFFER + */ + +#ifndef SDPA_FP_OUT_TILE_LOAD_GLSLH +#define SDPA_FP_OUT_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_output_tile.glslh" + +void store_out_d4( + const VEC4_T out_texel, + const int d4, + const int q_h, + const int s, + const int D4, + const int Q_H, + const int S) { +#ifdef OUTPUT_BUFFER + t_output[(s * Q_H * D4) + (q_h * D4) + d4] = out_texel; +#else + imageStore(t_output, ivec3(d4, q_h, s), out_texel); +#endif +} + +void store_sdpa_out_tile_with_checks( + const FPOutTile tile, + const int d4_start, + const int s_start, + const int q_h, + const int D4, + const int S, + const int Q_H) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int d4 = 0; d4 < TILE_N4; ++d4) { + if (d4_start + d4 < D4 && s_start + s < S) { + store_out_d4( + tile.data[s][d4], d4_start + d4, q_h, s_start + s, D4, Q_H, S); + } + } + } +} + +#endif // SDPA_FP_OUT_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh new file mode 100644 index 00000000000..a304e5019e9 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_q_projected_tile_load.glslh @@ -0,0 +1,74 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_input + * + * Macro Settings: + * - INPUT_BUFFER + */ + +#ifndef SDPA_FP_Q_PROJECTED_TILE_LOAD_GLSLH +#define SDPA_FP_Q_PROJECTED_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_input_tile.glslh" + +VEC4_T load_q_projected_d4( + const int d4, + const int q_h, + const int s, + const int D4, + const int Q_H, + const int S) { +#ifdef INPUT_BUFFER + return t_q_projected[(s * Q_H * D4) + (q_h * D4) + d4]; +#else + return texelFetch(t_q_projected, ivec3(d4, q_h, s), 0); +#endif +} + +void load_q_projected_tile_no_checks( + out FPInputTile tile, + const int d4_start, + const int s_start, + const int q_h, + const int D4, + const int Q_H, + const int S) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + tile.data[s][d4] = + load_q_projected_d4(d4_start + d4, q_h, s_start + s, D4, Q_H, S); + } + } +} + +void load_q_projected_tile_with_checks( + out FPInputTile tile, + const int d4_start, + const int s_start, + const int q_h, + const int D4, + const int Q_H, + const int S) { + [[unroll]] for (int s = 0; s < TILE_M; ++s) { + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + if (d4_start + d4 < D4 && s_start + s < S) { + tile.data[s][d4] = + load_q_projected_d4(d4_start + d4, q_h, s_start + s, D4, Q_H, S); + } else { + tile.data[s][d4] = VEC4_T(0.0); + } + } + } +} + +#endif // SDPA_FP_Q_PROJECTED_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh new file mode 100644 index 00000000000..bf94b251c43 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_v_cache_tile_load.glslh @@ -0,0 +1,76 @@ +/* + * 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. + */ + +/* + * Assume the following variables are defined in the shader layout: + * - t_v_cache + * + * Macro Settings: + * - V_CACHE_BUFFER + */ + +#ifndef SDPA_FP_V_CACHE_TILE_LOAD_GLSLH +#define SDPA_FP_V_CACHE_TILE_LOAD_GLSLH + +#extension GL_EXT_control_flow_attributes : require + +#include "linear_fp_weight_tile.glslh" + +VEC4_T load_v_cache_d4( + const int d4, + const int c, + const int kv_h, + const int D4, + const int C, + const int KV_H) { +#ifdef V_CACHE_BUFFER + return VEC4_T(t_v_cache[(c * KV_H * D4) + (kv_h * D4) + d4]); +#else + return VEC4_T(texelFetch(t_v_cache, ivec3(d4, kv_h, c), 0)); +#endif +} + +void load_v_cache_tile_no_checks( + out FPWeightTile tile, + const int d4_start, + const int c_start, + const int kv_h, + const int D4, + const int context_len, + const int C, + const int KV_H) { + [[unroll]] for (int c = 0; c < TILE_N; ++c) { + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + tile.data[c][d4] = + load_v_cache_d4(d4_start + d4, c_start + c, kv_h, D4, C, KV_H); + } + } +} + +void load_v_cache_tile_with_checks( + out FPWeightTile tile, + const int d4_start, + const int c_start, + const int kv_h, + const int D4, + const int context_len, + const int C, + const int KV_H) { + [[unroll]] for (int c = 0; c < TILE_N; ++c) { + [[unroll]] for (int d4 = 0; d4 < TILE_K4; ++d4) { + if (d4_start + d4 < D4 && c_start + c < context_len) { + tile.data[c][d4] = + load_v_cache_d4(d4_start + d4, c_start + c, kv_h, D4, C, KV_H); + } else { + tile.data[c][d4] = VEC4_T(0.0); + } + } + } +} + +#endif // SDPA_FP_V_CACHE_TILE_LOAD_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl new file mode 100644 index 00000000000..932696fff02 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.glsl @@ -0,0 +1,90 @@ +#version 450 core + +#define PRECISION ${PRECISION} + +#define IN_VEC4_T ${texel_load_type(DTYPE, INPUT_STORAGE)} +#define T ${buffer_scalar_type(DTYPE)} + +$if INPUT_STORAGE == "buffer": + #define INPUT_BUFFER + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "common.glslh" + +${layout_declare_tensor(B, "w", "t_cache", DTYPE, OUTPUT_STORAGE, is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_projected", DTYPE, INPUT_STORAGE, is_scalar_array=False)} + +${layout_declare_ubo(B, "ivec4", "cache_sizes")} +${layout_declare_ubo(B, "ivec4", "projected_sizes")} +${layout_declare_ubo(B, "int", "input_pos")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +/* + * t_cache will have sizes of (batches, n_heads, max_context_len, head_dim). + * t_projected will have sizes of (batches, seq_len, n_heads, head_dim). + * + * Note that the cache tensor swaps the order of the n_heads and seq len + * dimensions. This is to faciliate more optimal memory access patterns when + * using the caches to compute matrix multiplications. + * + * The cache update inserts the values of t_projected into t_cache at the index + * specified by input_pos at the seq_len dimension. It is equivalent to calling + + * t_cache = t_cache.slice_scatter( + * t_projected, dim=1, start=input_pos, end=input_pos+seq_len) + * + * Note that this shader is implemented assuming that max_batch_size is 1. + */ + +IN_VEC4_T read_projected_d4( + const int d4, + const int h, + const int s, + const int D4, + const int H, + const int S) { +#ifdef INPUT_BUFFER + return t_projected[(s * H * D4) + (h * D4) + d4]; +#else + return texelFetch(t_projected, ivec3(d4, h, s), 0); +#endif +} + +void write_cache_d4( + const IN_VEC4_T texel, + const int d4, + const int c, + const int h, + const int D4, + const int C, + const int H) { +#ifdef OUTPUT_BUFFER + t_cache[(c * H * D4) + (h * D4) + d4] = texel; +#else + imageStore(t_cache, ivec3(d4, h, c), texel); +#endif +} + +void main() { + const int d4 = int(gl_GlobalInvocationID.x); // idx along the head_dim dim + const int s = int(gl_GlobalInvocationID.y); // idx along the seq_len dim + const int h = int(gl_GlobalInvocationID.z); // idx along the n_heads dim + + const int D4 = div_up_4(projected_sizes.x); + const int S = projected_sizes.z; + const int H = projected_sizes.y; + + if (d4 >= D4 || s >= S || h >= H) { + return; + } + + const int c = s + input_pos; // idx along max_context_len dim + const int C = cache_sizes.y; + + IN_VEC4_T in_texel = read_projected_d4(d4, h, s, D4, H, S); + write_cache_d4(in_texel, d4, c, h, D4, C, H); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml new file mode 100644 index 00000000000..85f4ce090f8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_kv_cache_update.yaml @@ -0,0 +1,19 @@ +# 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. + +sdpa_kv_cache_update: + parameter_names_with_default_values: + DTYPE: float + INPUT_STORAGE: texture3d + OUTPUT_STORAGE: texture3d + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: sdpa_kv_cache_update_texture3d + - NAME: sdpa_kv_cache_update_buffer + INPUT_STORAGE: buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/sdpa_q_projected_input_tile.glslh b/backends/vulkan/runtime/graph/ops/glsl/sdpa_q_projected_input_tile.glslh new file mode 100644 index 00000000000..da5dcd63b31 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/sdpa_q_projected_input_tile.glslh @@ -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. + */ + +#ifndef SDPA_FP_Q_PROJECTED_TILE_GLSLH +#define SDPA_FP_Q_PROJECTED_TILE_GLSLH + +/* + * Macro Settings: + * - TILE_S + * - TILE_D4 + */ + +#extension GL_EXT_control_flow_attributes : require + +struct FPQProjectedTile { + VEC4_T data[TILE_S][TILE_D4]; +}; + +#ifdef DEBUG_MODE + +void printFPQProjectedTile(const FPQProjectedTile in_tile) { + debugPrintfEXT("input_tile: \\n"); + [[unroll]] for (int m = 0; m < TILE_M; ++m) { + [[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) { + debugPrintfEXT( + " %f, %f, %f, %f, \\n", + in_tile.data[m][k4].x, + in_tile.data[m][k4].y, + in_tile.data[m][k4].z, + in_tile.data[m][k4].w); + } + } +} + +#endif // DEBUG_MODE + +#endif // SDPA_FP_Q_PROJECTED_TILE_GLSLH diff --git a/backends/vulkan/runtime/graph/ops/glsl/set_zero.glsl b/backends/vulkan/runtime/graph/ops/glsl/set_zero.glsl new file mode 100644 index 00000000000..d01780b9e30 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/set_zero.glsl @@ -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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} + +${layout_declare_ubo(B, "int", "out_numel")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const int out_bufi = ivec3(gl_GlobalInvocationID).x; + if (out_bufi >= out_numel) { + return; + } + + t_out[out_bufi] = T(0); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/set_zero.yaml b/backends/vulkan/runtime/graph/ops/glsl/set_zero.yaml new file mode 100644 index 00000000000..cee87c468b1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/set_zero.yaml @@ -0,0 +1,8 @@ +set_zero: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: int32 + shader_variants: + - NAME: set_zero diff --git a/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.glsl b/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.glsl new file mode 100644 index 00000000000..ba02da1c301 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.glsl @@ -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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "concat_offset", DTYPE, "buffer")} + +${layout_declare_ubo(B, "int", "concat_dim")} + +$for i in range(NUM_INPUTS): + ${layout_declare_ubo(B, "ivec4", "in" + str(i+1) + "_sizes")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + // Only one thread needs to update the offset + if (gl_GlobalInvocationID.x != 0) { + return; + } + + // Sum up the sizes along the concat dimension for all input tensors + int total_size = 0; + $for i in range(NUM_INPUTS): + total_size += in${i+1}_sizes[concat_dim]; + + // Add to the current offset + concat_offset[0] += T(total_size); +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.yaml b/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.yaml new file mode 100644 index 00000000000..35e8740e0a3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/update_concat_offset.yaml @@ -0,0 +1,13 @@ +update_concat_offset: + parameter_names_with_default_values: + DTYPE: float + NUM_INPUTS: 2 + generate_variant_forall: + DTYPE: + - VALUE: int32 + shader_variants: + - NAME: update_concat_offset_1 + NUM_INPUTS: 1 + - NAME: update_concat_offset_2 + - NAME: update_concat_offset_3 + NUM_INPUTS: 3 diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/view_buffer.glsl new file mode 100644 index 00000000000..2c02803a9b1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/view_buffer.glsl @@ -0,0 +1,44 @@ +#version 450 core + +#define PRECISION ${PRECISION} + +#define T ${buffer_scalar_type(DTYPE)} + +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +#include "indexing.glslh" + +${layout_declare_tensor(B, "w", "t_outp", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_inp", DTYPE, STORAGE)} + +${layout_declare_ubo(B, "BufferMetadata", "outp")} +${layout_declare_ubo(B, "BufferMetadata", "inp")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +/* + * The insight behind the view operation is that the contiguous index of each + * tensor element in the input and output tensors are the same. + */ +void main() { + const uint outp_bufi = gl_GlobalInvocationID.x; + if (outp_bufi >= numel(outp)) { + return; + } + + TensorIndex outp_tidx; + linear_idx_to_tensor_idx(outp, outp_bufi, outp_tidx); + + // To map the output to the input, find the input element that has the same + // contiguous index as the output element. + const uint contig_idx = tensor_idx_to_contiguous_idx(outp, outp_tidx); + + TensorIndex inp_tidx; + contiguous_idx_to_tensor_idx(inp, contig_idx, inp_tidx); + + const uint inp_bufi = tensor_idx_to_linear_idx(inp, inp_tidx); + + t_outp[outp_bufi] = t_inp[inp_bufi]; +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/view_buffer.yaml new file mode 100644 index 00000000000..ec92bf483c8 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/view_buffer.yaml @@ -0,0 +1,20 @@ +# 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. + +view_buffer: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + - VALUE: double + - VALUE: int8 + - VALUE: uint8 + - VALUE: int32 + shader_variants: + - NAME: view_buffer diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index 490def4860a..3171fbeb488 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -20,22 +21,22 @@ void resize_arange_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); + const ValueRef out = args.at(0).refs.at(0); int start_val = 0; int step_val = 1; - if (!graph->val_is_none(extra_args[0])) { - start_val = graph->extract_scalar(extra_args[0]); + if (!graph->val_is_none(extra_args.at(0))) { + start_val = graph->extract_scalar(extra_args.at(0)); } - int end_val = graph->extract_scalar(extra_args[1]); - if (!graph->val_is_none(extra_args[2])) { - step_val = graph->extract_scalar(extra_args[2]); + const int 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)); } - std::vector out_sizes = { + const std::vector out_sizes = { utils::div_up(end_val - start_val, step_val)}; - out->virtual_resize(out_sizes); + graph->virtual_resize(out, out_sizes); } void check_arange_input( @@ -82,21 +83,19 @@ void add_arange_node( } } - vTensorPtr t_out = graph.get_tensor(out); - std::string kernel_name("arange"); kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers - {t_out->sizes_ubo(), + {graph.sizes_ubo(out), graph.create_params_buffer(start_val), graph.create_params_buffer(step_val)}, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp index 81cbd62d90c..757afd06849 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BatchNorm.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -46,53 +47,51 @@ void add_native_batch_norm_node( ValueRef var_ref, ValueRef eps_ref, ValueRef out_tuple_ref) { - std::vector in_sizes = graph.get_tensor(in_ref)->sizes(); - std::vector out_sizes = graph.get_tensor(in_ref)->sizes(); + const std::vector in_sizes = graph.sizes_of(in_ref); + const std::vector out_sizes = graph.sizes_of(in_ref); VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor"); VK_CHECK_COND(out_sizes.size() == 4, "BatchNorm only support 4d tensor"); // Only the first element of the return value is propagated. The remaining 2 // elements are zero-size dummy tensor. - ValueRef out_ref = graph.get_value_list(out_tuple_ref)->at(0); + const ValueRef out_ref = graph.get_value_list(out_tuple_ref)->at(0); - utils::StorageType stype = graph.storage_type_of(out_ref); + const utils::StorageType stype = graph.storage_type_of(out_ref); - int64_t num_channels = dim_at(in_sizes); + const int64_t num_channels = dim_at(in_sizes); - ValueRef arg_weight = + const ValueRef arg_weight = check_and_prepack_arg(graph, weight_ref, stype, num_channels, "weight"); - ValueRef arg_bias = + const ValueRef arg_bias = check_and_prepack_arg(graph, bias_ref, stype, num_channels, "bias"); - ValueRef arg_mean = + const ValueRef arg_mean = check_and_prepack_arg(graph, mean_ref, stype, num_channels, "mean"); - ValueRef arg_var = + const ValueRef arg_var = check_and_prepack_arg(graph, var_ref, stype, num_channels, "var"); - float epsilon = graph.extract_scalar(eps_ref); - - vTensorPtr t_in = graph.get_tensor(in_ref); + const float epsilon = graph.extract_scalar(eps_ref); VK_CHECK_COND(!graph.val_is_tref(out_ref), "Output should not be tref"); - vTensorPtr t_out = graph.get_tensor(out_ref); + const std::vector out_tensor_sizes = graph.sizes_of(out_ref); VK_CHECK_COND( - dim_at(t_out->sizes()) == num_channels, + dim_at(out_tensor_sizes) == num_channels, "out channel must match in channel"); std::string kernel_name = "batchnorm"; - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out_ref)); - int32_t num_texel_per_batch = - utils::div_up_4((dim_at(t_in->sizes()))); + const int32_t num_texel_per_batch = + utils::div_up_4((dim_at(in_sizes))); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out_ref), - graph.create_local_wg_size(out_ref), + default_pick_global_wg_size, + default_pick_local_wg_size, {{out_ref, vkapi::kWrite}, {{in_ref, arg_weight, arg_bias, arg_mean, arg_var}, vkapi::kRead}}, - {t_out->logical_limits_ubo(), + {graph.logical_limits_ubo(out_ref), graph.create_params_buffer(epsilon), graph.create_params_buffer(num_texel_per_batch)}, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp index 18a1aacf323..025b483eab7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp @@ -19,13 +19,20 @@ namespace vkcompute { void check_binary_op_args( - const api::vTensor& self, - const api::vTensor& other, - const api::vTensor& out) { - VK_CHECK_COND(check_same_packed_dim(self, other, out)); + ComputeGraph& graph, + const ValueRef self, + const ValueRef other, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(self) == graph.packed_dim_of(other)); + VK_CHECK_COND(graph.packed_dim_of(self) == graph.packed_dim_of(out)); + + const std::vector self_sizes = graph.sizes_of(self); + const std::vector other_sizes = graph.sizes_of(other); + const std::vector out_sizes = graph.sizes_of(out); + std::vector broadcasted_sizes = - calculate_broadcasted_output_size(self, other); - VK_CHECK_COND(out.sizes() == broadcasted_sizes); + calculate_broadcasted_output_size(self_sizes, other_sizes); + VK_CHECK_COND(out_sizes == broadcasted_sizes); } void resize_binary_op_node( @@ -33,16 +40,18 @@ void resize_binary_op_node( const std::vector& args, const std::vector& resize_args) { (void)resize_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); + const ValueRef out = args.at(0).refs.at(0); // TODO(T183442143): Verify tensors are broadcastable. - vTensorPtr self = graph->get_tensor(args[1].refs[0]); - vTensorPtr other = graph->get_tensor(args[1].refs[1]); + const ValueRef self = args.at(1).refs.at(0); + const ValueRef other = args.at(1).refs.at(1); - std::vector new_out_sizes = - calculate_broadcasted_output_size(*self, *other); + const std::vector self_sizes = graph->sizes_of(self); + const std::vector other_sizes = graph->sizes_of(other); + const std::vector new_out_sizes = + calculate_broadcasted_output_size(self_sizes, other_sizes); - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } void add_binary_op_texture_node( @@ -55,11 +64,7 @@ void add_binary_op_texture_node( ValueRef arg1 = prepack_standard_like(graph, in1, out, true); ValueRef arg2 = prepack_standard_like(graph, in2, out, true); - vTensorPtr t_in1 = graph.get_tensor(arg1); - vTensorPtr t_in2 = graph.get_tensor(arg2); - vTensorPtr t_out = graph.get_tensor(out); - - check_binary_op_args(*t_in1, *t_in2, *t_out); + check_binary_op_args(graph, arg1, arg2, out); float alpha_val = 1.0f; // String is checked since floor_div passes in an unused string argument in @@ -71,12 +76,12 @@ void add_binary_op_texture_node( const struct BinaryOpsParams { const utils::ivec2 broadcast_params; const float alpha_val; - } binary_ops_params{create_broadcast_params(*t_in1, *t_in2), alpha_val}; + } binary_ops_params{create_broadcast_params(graph, arg1, arg2), alpha_val}; std::string kernel_name("binary_"); kernel_name.reserve(kShaderNameReserve); kernel_name += op_name; - add_storage_type_suffix(kernel_name, *t_out); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(in1)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( @@ -94,7 +99,9 @@ void add_binary_op_texture_node( graph.sizes_pc_of(arg2), PushConstantDataInfo(&binary_ops_params, sizeof(binary_ops_params))}}, // Specialization Constants - {t_out->hashed_layout(), t_in1->hashed_layout(), t_in2->hashed_layout()}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(arg1), + graph.hashed_layout_of(arg2)}, // Resize Args {}, // Resizing Logic @@ -132,15 +139,11 @@ void add_binary_op_buffer_node( // Inputs and Outputs {{out, vkapi::kWrite}, {{in1, in2}, vkapi::kRead}}, // Shader params buffers - {}, + {graph.buffer_meta_ubo(out), + graph.buffer_meta_ubo(in1), + graph.buffer_meta_ubo(in2)}, // Push Constants {{ - graph.sizes_pc_of(in1), - graph.sizes_pc_of(in2), - graph.strides_pc_of(out), - graph.strides_pc_of(in1), - graph.strides_pc_of(in2), - graph.numel_pc_of(out), PushConstantDataInfo(&alpha_val, sizeof(float)), }}, // Specialization Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp index 76d352334e3..a4a96ffdb88 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp @@ -14,6 +14,26 @@ namespace vkcompute { +void resize_choose_qparams_per_row( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + (void)extra_args; + + ValueRef input_scales = args.at(0).refs.at(0); + ValueRef input_zeros = args.at(0).refs.at(1); + ValueRef input = args.at(1).refs.at(0); + + std::vector new_sizes = graph->sizes_of(input_scales); + const size_t ndim = new_sizes.size(); + + const int64_t input_height = graph->size_at(-2, input); + new_sizes.at(ndim - 1) = input_height; + + graph->virtual_resize(input_scales, new_sizes); + graph->virtual_resize(input_zeros, new_sizes); +} + utils::uvec3 choose_qparams_pick_global_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -158,6 +178,50 @@ utils::uvec3 choose_qparams_block_wise_pick_local_wg_size( } } +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); + + 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)); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +utils::uvec3 pick_choose_qparams_per_row_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + + const ValueRef input = args.at(1).refs.at(0); + const uint32_t height = graph->size_at(-2, input); + return {1u, utils::div_up_4(height), 1u}; +} + +utils::uvec3 pick_choose_qparams_per_row_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)global_workgroup_size; + (void)args; + (void)resize_args; + + uint32_t outputs_per_wg = 1u; + uint32_t workers_per_output = 64u; + + return {workers_per_output, outputs_per_wg, 1u}; +} + void add_choose_qparams_tensor_node( ComputeGraph& graph, const ValueRef& input, @@ -169,9 +233,35 @@ void add_choose_qparams_tensor_node( std::string kernel_name("choose_qparams_tensor"); add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale_out)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point_out)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(zero_point_out)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } float eps_val = static_cast(graph.get_double(eps)); vkapi::ParamsBindList param_ubos; @@ -227,6 +317,8 @@ void add_choose_qparams_per_token_asymmetric_node( std::string kernel_name("choose_qparams_per_token_asymmetric"); add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale_out)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point_out)); // Calculate number of tokens (product of all dimensions except the last one) int64_t num_tokens = 1; @@ -284,6 +376,57 @@ void add_choose_qparams_per_token_asymmetric_node( nullptr)); } +void add_choose_qparams_per_row_node( + ComputeGraph& graph, + const ValueRef& input, + const ValueRef& quant_min, + const ValueRef& quant_max, + const ValueRef& input_scales, + const ValueRef& input_zps) { + int32_t quant_min_val = -128; + int32_t quant_max_val = 127; + + // Int8 range by default + if (graph.val_is_none(quant_min)) { + quant_min_val = -128; + } else { + quant_min_val = graph.extract_scalar(quant_min); + } + + // Int8 range by default + if (graph.val_is_none(quant_min)) { + quant_max_val = 127; + } else { + quant_max_val = graph.extract_scalar(quant_max); + } + + vkapi::ParamsBindList param_ubos = { + graph.sizes_ubo(input), + }; + std::vector push_constants = { + PushConstantDataInfo(&quant_min_val, sizeof(int32_t)), + PushConstantDataInfo(&quant_max_val, sizeof(int32_t)), + }; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_choose_qparams_per_row_shader, + pick_choose_qparams_per_row_global_wg_size, + pick_choose_qparams_per_row_local_wg_size, + // Inputs and Outputs + {{{input_scales, input_zps}, vkapi::kWrite}, {input, vkapi::kRead}}, + // Shader param buffers + param_ubos, + // Push Constants + push_constants, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_choose_qparams_per_row)); +} + void add_choose_qparams_block_wise_node( ComputeGraph& graph, ValueRef input, @@ -317,9 +460,26 @@ void add_choose_qparams_block_wise_node( num_blocks_vec[0] * num_blocks_vec[1], num_blocks_vec[0] * num_blocks_vec[1] * num_blocks_vec[2]}; - int qmin = static_cast(graph.get_int(quant_min)); - int qmax = static_cast(graph.get_int(quant_max)); - float eps_val = static_cast(graph.get_double(eps)); + // Handle optional quant_min and quant_max parameters + int qmin, qmax; + if (graph.val_is_none(quant_min) || graph.val_is_none(quant_max)) { + // Use default values based on target_dtype (similar to + // _get_and_check_qmin_qmax) For now, assume int8 range as default - this + // should match the Python implementation + qmin = -128; + qmax = 127; + } else { + qmin = static_cast(graph.get_int(quant_min)); + qmax = static_cast(graph.get_int(quant_max)); + } + + float eps_val; + if (graph.val_is_none(eps)) { + // Use default eps value (similar to Python implementation) + eps_val = 1.192092896e-07f; // torch.finfo(torch.float32).eps + } else { + eps_val = static_cast(graph.get_double(eps)); + } // Create push constants vector std::vector push_constants = { @@ -334,6 +494,8 @@ void add_choose_qparams_block_wise_node( std::string kernel_name("choose_qparams_block_wise"); add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale_out)); + add_dtype_suffix(kernel_name, graph.dtype_of(zp_out)); vkapi::ParamsBindList param_ubos; @@ -408,9 +570,18 @@ void choose_qparams_tensor_impl( // Verify input is a floating point type VK_CHECK_COND(graph.dtype_of(input) == vkapi::kFloat); - // Verify output types - VK_CHECK_COND(graph.dtype_of(scale_out) == vkapi::kFloat); - VK_CHECK_COND(graph.dtype_of(zero_point_out) == vkapi::kInt); + // Get scale and zero point output dtypes + vkapi::ScalarType scale_out_dtype = graph.dtype_of(scale_out); + vkapi::ScalarType zero_point_out_dtype = graph.dtype_of(zero_point_out); + + // Verify supported output types for scale (fp32 only for now) + VK_CHECK_COND(scale_out_dtype == vkapi::kFloat); + + // Verify supported output types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_out_dtype == vkapi::kInt || + zero_point_out_dtype == vkapi::kChar || + zero_point_out_dtype == vkapi::kFloat); // Check that texture storage is width packed if (!graph.is_buffer_storage(input)) { @@ -449,9 +620,18 @@ void choose_qparams_per_token_asymmetric_impl( // Verify input is a floating point type VK_CHECK_COND(graph.dtype_of(input) == vkapi::kFloat); - // Verify output types - VK_CHECK_COND(graph.dtype_of(scale_out) == vkapi::kFloat); - VK_CHECK_COND(graph.dtype_of(zero_point_out) == vkapi::kInt); + // Get scale and zero point output dtypes + vkapi::ScalarType scale_out_dtype = graph.dtype_of(scale_out); + vkapi::ScalarType zero_point_out_dtype = graph.dtype_of(zero_point_out); + + // Verify supported output types for scale (fp32 only for now) + VK_CHECK_COND(scale_out_dtype == vkapi::kFloat); + + // Verify supported output types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_out_dtype == vkapi::kInt || + zero_point_out_dtype == vkapi::kChar || + zero_point_out_dtype == vkapi::kFloat); // Check that texture storage is width packed if (!graph.is_buffer_storage(input)) { @@ -462,6 +642,32 @@ void choose_qparams_per_token_asymmetric_impl( graph, input, scale_out, zero_point_out); } +bool can_use_choose_qparams_per_row( + ComputeGraph& graph, + const ValueRef input, + const ValueRef block_size, + const ValueRef input_zero_point) { + if (!graph.is_vectorizable_contiguous_2d_matrix(input)) { + return false; + } + + std::vector input_sizes = graph.sizes_of(input); + const IntListPtr block_size_vals = graph.get_int_list(block_size); + const size_t ndim = block_size_vals->size(); + + // Check for per y - dim quantization + if (utils::val_at(-1, input_sizes) != utils::val_at(-1, *block_size_vals)) { + return false; + } + + for (int d = 0; d < ndim - 1; ++d) { + if (block_size_vals->at(d) != 1) { + return false; + } + } + return true; +} + void choose_qparams_affine_impl( ComputeGraph& graph, const std::vector& args) { @@ -491,6 +697,13 @@ void choose_qparams_affine_impl( zero_point_out = out_tuple->at(1); } + // Use fast path if certain conditions are met + if (can_use_choose_qparams_per_row( + graph, input, block_size, zero_point_out)) { + return add_choose_qparams_per_row_node( + graph, input, quant_min, quant_max, scale_out, zero_point_out); + } + // Check tensor types VK_CHECK_COND(graph.val_is_tensor(input)); VK_CHECK_COND(graph.val_is_tensor(scale_out)); @@ -499,9 +712,18 @@ void choose_qparams_affine_impl( // Verify input is a floating point type VK_CHECK_COND(graph.dtype_of(input) == vkapi::kFloat); - // Verify output types - VK_CHECK_COND(graph.dtype_of(scale_out) == vkapi::kFloat); - VK_CHECK_COND(graph.dtype_of(zero_point_out) == vkapi::kInt); + // Get scale and zero point dtypes from arguments + vkapi::ScalarType scale_out_dtype = graph.dtype_of(scale_out); + vkapi::ScalarType zero_point_out_dtype = graph.dtype_of(zero_point_out); + + // Verify supported output types for scale (fp32 only for now) + VK_CHECK_COND(scale_out_dtype == vkapi::kFloat); + + // Verify supported output types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_out_dtype == vkapi::kInt || + zero_point_out_dtype == vkapi::kChar || + zero_point_out_dtype == vkapi::kFloat); // Check that texture storage is width packed if (!graph.is_buffer_storage(input)) { @@ -515,12 +737,14 @@ void choose_qparams_affine_impl( std::string mapping_type_str = graph.get_string(mapping_type); int mapping_type_val = 0; // Default to ASYMMETRIC - if (mapping_type_str == "ASYMMETRIC") { - mapping_type_val = 0; + if (mapping_type_str == "ASYMMETRIC" || mapping_type_str.empty()) { + mapping_type_val = 0; // ASYMMETRIC } else if (mapping_type_str == "SYMMETRIC") { mapping_type_val = 1; } else if (mapping_type_str == "SYMMETRIC_NO_CLIPPING_ERR") { mapping_type_val = 2; + } else { + VK_THROW("Unsupported mapping_type: ", mapping_type_str); } add_choose_qparams_block_wise_node( @@ -535,6 +759,30 @@ void choose_qparams_affine_impl( zero_point_out); } +void choose_qparams_per_row( + ComputeGraph& graph, + const std::vector& args) { + int arg_idx = 0; + const ValueRef input = args[arg_idx++]; + const ValueRef quant_min = args[arg_idx++]; + const ValueRef quant_max = args[arg_idx++]; + const ValueRef input_scales = args[arg_idx++]; + const ValueRef input_zps = args[arg_idx++]; + + // ValueRef scale_out = kDummyValueRef; + // ValueRef zero_point_out = kDummyValueRef; + // + // { + // const ValueListPtr out_tuple = graph.get_value_list(out_tuple_ref); + // scale_out = out_tuple->at(0); + // zero_point_out = out_tuple->at(1); + // } + // + + add_choose_qparams_per_row_node( + graph, input, quant_min, quant_max, input_scales, input_zps); +} + REGISTER_OPERATORS { VK_REGISTER_OP( quantized_decomposed.choose_qparams.tensor, choose_qparams_tensor_impl); @@ -542,6 +790,9 @@ REGISTER_OPERATORS { quantized_decomposed.choose_qparams_per_token_asymmetric.default, choose_qparams_per_token_asymmetric_impl); + // Register the per-channel quantization operator + VK_REGISTER_OP(etvk.choose_qparams_per_row.default, choose_qparams_per_row); + // TorchAO affine choose_qparams operators VK_REGISTER_OP( torchao.choose_qparams_affine.default, choose_qparams_affine_impl); diff --git a/backends/vulkan/runtime/graph/ops/impl/Clone.cpp b/backends/vulkan/runtime/graph/ops/impl/Clone.cpp index fcbac2df0fc..059aade5b04 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Clone.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Clone.cpp @@ -24,12 +24,12 @@ void resize_clone_node( const std::vector& args, const std::vector& resize_args) { (void)resize_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); // TODO: support for when dimensionality doesn't match, i.e. clone is used to // implement squeeze. - if (out->dim() == in->dim()) { - out->virtual_resize(in->sizes()); + if (graph->dim_of(out) == graph->dim_of(in)) { + graph->virtual_resize(out, graph->sizes_of(in)); } } @@ -37,10 +37,8 @@ void add_clone_node( ComputeGraph& graph, const ValueRef in, const ValueRef out) { - vTensorPtr t_out = graph.get_tensor(out); - std::string kernel_name = "clone"; - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -50,7 +48,7 @@ void add_clone_node( // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Parameter Buffers - {t_out->logical_limits_ubo()}, + {graph.logical_limits_ubo(out)}, // Push Constants {}, // Specialization Constants @@ -78,6 +76,7 @@ void add_image_to_buffer_node( const ValueRef buffer) { std::string kernel_name = "clone_image_to_buffer"; add_dtype_suffix(kernel_name, graph.dtype_of(image)); + add_dtype_suffix(kernel_name, graph.dtype_of(buffer)); vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); graph.execute_nodes().emplace_back(new DynamicDispatchNode( @@ -105,6 +104,7 @@ void add_buffer_to_image_node( const ValueRef image) { std::string kernel_name = "clone_buffer_to_image"; add_dtype_suffix(kernel_name, graph.dtype_of(image)); + add_dtype_suffix(kernel_name, graph.dtype_of(buffer)); vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); graph.execute_nodes().emplace_back(new DynamicDispatchNode( @@ -145,7 +145,11 @@ void clone(ComputeGraph& graph, const std::vector& args) { if (src_storage == utils::kBuffer && dst_storage == utils::kTexture3D) { return add_buffer_to_image_node(graph, src, dst); } - VK_THROW("Buffer to buffer memory layout transition not supported yet!"); + + std::vector extra_args = {}; + // Buffer to buffer copy + return add_view_copy_buffer_node( + graph, src, dst, extra_args, resize_clone_node); } // Clone node is not the most efficient implementation for the aten.clone diff --git a/backends/vulkan/runtime/graph/ops/impl/Common.cpp b/backends/vulkan/runtime/graph/ops/impl/Common.cpp index 4c3c16417b5..6c701224f7f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Common.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Common.cpp @@ -33,4 +33,27 @@ utils::uvec3 default_pick_local_wg_size( return graph->create_local_wg_size(global_workgroup_size); } +utils::uvec3 pick_hw_square_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)args; + (void)resize_args; + // Some inactive invocations are okay; set 6 as the threshold to use the + // a square wg size. + if (global_workgroup_size[0u] >= 6 && global_workgroup_size[1u] >= 6) { + return {8u, 8u, 1u}; + } + // If width dim is sufficiently small, then bias towards height dim to reduce + // the number of inactive invocations. + if (global_workgroup_size[0u] < 6u) { + return {4u, 16u, 1u}; + } + return {16u, 4u, 1u}; +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Common.h b/backends/vulkan/runtime/graph/ops/impl/Common.h index 662fb07095a..1831ab2a845 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Common.h +++ b/backends/vulkan/runtime/graph/ops/impl/Common.h @@ -36,4 +36,22 @@ utils::uvec3 default_pick_local_wg_size( const std::vector& args, const std::vector& resize_args); +/** + * Constructs a local work group size with the shape {W, H, 1}. The function + * will try to set W == H == sqrt(num_invocations), where num_invocations is + * typically 64. This configuration is good for ops like matrix multiplication + * as it reduces the total volume of unique data that the entire work group + * will need to read from input tensors in order to produce the output data. + * To compute an output tile of {W, H, 1}, the work group will need to read + * H unique rows = H * K unique elements from the input tensor and W unique cols + * = W * K elements from the weight tensor, resulting in (W + H) * K unique + * elements in total. + */ +utils::uvec3 pick_hw_square_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args); + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Concat.cpp b/backends/vulkan/runtime/graph/ops/impl/Concat.cpp index 315dabdb1d5..0a4acb6cef3 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Concat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Concat.cpp @@ -19,15 +19,16 @@ namespace vkcompute { std::vector get_concat_sizes( ComputeGraph& graph, - const std::vector& in_value_refs, - const int64_t dim) { + ValueRef all_input_refs, + const int64_t concat_dim) { + ValueListPtr in_value_refs = graph.get_value_list(all_input_refs); // Get the sizes of the first input tensor as a starting point - std::vector new_out_sizes = graph.sizes_of(in_value_refs.at(0)); + std::vector new_out_sizes = graph.sizes_of(in_value_refs->at(0)); // Sum up the sizes along the concatenation dimension - for (size_t i = 1; i < in_value_refs.size(); ++i) { - const std::vector in_sizes = graph.sizes_of(in_value_refs.at(i)); - new_out_sizes.at(dim) += in_sizes.at(dim); + for (size_t i = 1; i < in_value_refs->size(); ++i) { + const std::vector in_sizes = graph.sizes_of(in_value_refs->at(i)); + new_out_sizes.at(concat_dim) += in_sizes.at(concat_dim); } return new_out_sizes; @@ -37,24 +38,122 @@ void resize_concat_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - // Extract relevant ValueRefs - const ValueRef out_ref = args.at(0).refs.at(0); - const std::vector& in_value_refs = args.at(1).refs; + const ValueRef out = args.at(0).refs.at(0); + const ValueRef all_inputs = extra_args.at(0); - int64_t dim = graph->extract_scalar(extra_args.at(0)); + int64_t concat_dim = graph->extract_scalar(extra_args.at(1)); - // Normalize dim if negative - const int64_t ndim = graph->dim_of(out_ref); - if (dim < 0) { - dim += ndim; + // Normalize concat_dim if negative + const int64_t ndim = graph->dim_of(out); + if (concat_dim < 0) { + concat_dim += ndim; } // Calculate the new sizes std::vector new_out_sizes = - get_concat_sizes(*graph, in_value_refs, dim); + get_concat_sizes(*graph, all_inputs, concat_dim); // Resize the output tensor - graph->virtual_resize(out_ref, new_out_sizes); + graph->virtual_resize(out, new_out_sizes); +} + +utils::uvec3 concat_pick_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& extra_args) { + (void)shader; + (void)extra_args; + + const ValueRef out = args.at(0).refs.at(0); + const std::vector inputs_in_batch = args.at(1).refs; + + int64_t concat_dim = graph->extract_scalar(extra_args.at(1)); + + // Normalize concat_dim if negative + const int64_t ndim = graph->dim_of(out); + if (concat_dim < 0) { + concat_dim += ndim; + } + + // The concat shader concatenates N input tensors at a time to the output + // tensor. Since the shader may need to be invoked multiple times to finish + // concatenation when the number of input tensors is >N, the global workgroup + // is based on the volume of input data being concatenated in this batch, + // as opposed to the overall size of the output tensor. Conceptually, the + // global work group size represents which elements of the output tensor will + // be written to during this dispatch. + + uint32_t total_input_numel = 0; + int64_t concat_dim_numel = 0; + for (const ValueRef input : inputs_in_batch) { + total_input_numel += graph->numel_of(input); + concat_dim_numel += graph->size_at(concat_dim, input); + } + + if (graph->is_buffer_storage(out)) { + return {total_input_numel, 1, 1}; + } + + // The texture implementation is similar, except each invocation writes out 4 + // output elements along the packed dim (i.e. one texel). In this case, the + // global work group size represents the number of output texels that will be + // written to in this batch, rather than the number of output elements. Note + // that to update an element of the output, the entire texel that contains it + // will need to be loaded, updated, then written back. + + std::vector inp_volume_sizes = graph->sizes_of(out); + inp_volume_sizes.at(concat_dim) = concat_dim_numel; + + // Calculate what the image extents would be of a tensor with the input + // volume's sizes. This produces the number of texels that would need to be + // written to. + const int32_t packed_dim = graph->packed_dim_of(out); + std::vector inp_volume_texel_sizes = + api::calculate_padded_sizes(inp_volume_sizes, packed_dim); + // If the concat_dim is the same as the packed dim, and the concat_offset for + // this input batch is not a multiple of 4, then the data from an input texel + // may be split up between two output texels. For example: + // I0 , I1 , I2 , I2 + // O0 , O1 , O2 , X | X , X , X , X + // Therefore, 1 texel is added to the packed dim to account for this. + inp_volume_texel_sizes.at(3 - packed_dim) = + utils::div_up_4(inp_volume_texel_sizes.at(3 - packed_dim)) + 1; + + const uint32_t inp_volume_texel_numel = + utils::multiply_integers(inp_volume_texel_sizes); + + return {inp_volume_texel_numel, 1, 1}; + + // The texture implementation is similar, expect each thread is responsible + // for writing out an entire output texel. Therefore, the overall global work + // group size will be the concatenation of the texture extents of the input + // tensors in this batch. + + // One complication is when the previous concatenation batch does not write + // up to a texel boundary. An example is if the previous concatenation batch + // only wrote 7 elements along the concatenation dim. The first input element + // would then have to be inserted at the last element of the final texel + // written by the previous batch. To account for this, initialize the + // workgroup size at the concatenation dim to 1 (need to read N total texels + // along the concat dim for input tensors + up to 1 texel from the output + // tensor). + + // The axis along which to concatenate the input texture extents + int64_t extent_concat_axis = nchw_dim_to_whcn_dim(concat_dim, ndim); + // For batch concatenation, the concat axis is the batch-concatenation axis + if (concat_dim == 4) { + extent_concat_axis = graph->concat_dim_of(out); + } + + utils::uvec3 global_workgroup_size = graph->create_global_wg_size(out); + global_workgroup_size[concat_dim] = 0; + for (const ValueRef input : inputs_in_batch) { + utils::uvec3 texture_extents = graph->logical_limits_of(input); + global_workgroup_size[extent_concat_axis] += texture_extents[concat_dim]; + } + + return global_workgroup_size; } void add_concat_node( @@ -67,10 +166,6 @@ void add_concat_node( { const ValueListPtr tensors = graph.get_value_list(tensors_ref); - VK_CHECK_COND( - tensors->size() <= 3, - "Currently only concatenation of <= 3 tensors is supported"); - for (const ValueRef in : *tensors) { in_value_refs.push_back(in); } @@ -87,68 +182,161 @@ void add_concat_node( const int64_t dim_whcn = nchw_dim_to_whcn_dim(normalized_dim, ndim); const ValueRef dim_whcn_ref = graph.get_or_add_value_for_int(dim_whcn); - vkapi::ParamsBindList param_buffers = { - graph.get_or_create_int_param_buffer(dim_whcn_ref, 0)}; + // Create a temporary tensor to hold the concat offset + TmpTensor concat_offset( + &graph, {1}, vkapi::kInt, utils::kBuffer, utils::kWidthPacked); - std::vector push_constants; - vkapi::SpecVarList spec_vars; - - if (graph.is_buffer_storage(out)) { - param_buffers.append(graph.sizes_ubo(out)); - param_buffers.append(graph.strides_ubo(out)); + // Add node to set concat_offset to 0 + { + std::string kernel_name = "set_zero"; + add_dtype_suffix(kernel_name, graph.dtype_of(concat_offset)); + + vkapi::ParamsBindList param_buffers = {graph.numel_ubo(concat_offset)}; + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + {1, 1, 1}, + {1, 1, 1}, + // Inputs and Outputs + {{concat_offset, vkapi::kWrite}}, + // Parameter buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + nullptr)); + } - for (const ValueRef in_ref : in_value_refs) { - param_buffers.append(graph.sizes_ubo(in_ref)); - param_buffers.append(graph.strides_ubo(in_ref)); + // Process inputs in batches of up to 3 tensors + const size_t batch_size = 3; + for (size_t batch_start = 0; batch_start < in_value_refs.size(); + batch_start += batch_size) { + const size_t batch_end = + std::min(batch_start + batch_size, in_value_refs.size()); + const size_t current_batch_size = batch_end - batch_start; + + std::vector batch_inputs; + for (size_t i = batch_start; i < batch_end; ++i) { + batch_inputs.push_back(in_value_refs.at(i)); } - param_buffers.append(graph.numel_ubo(out)); - - spec_vars = {graph.hashed_layout_of(out)}; - } else { - push_constants = {graph.sizes_pc_of(out)}; - - spec_vars = {graph.hashed_layout_of(out)}; - - for (const ValueRef in_ref : in_value_refs) { - push_constants.push_back(graph.sizes_pc_of(in_ref)); - spec_vars.append(graph.hashed_layout_of(in_ref)); + // Add concat node for this batch + { + vkapi::ParamsBindList param_buffers = { + graph.get_or_create_int_param_buffer(dim_whcn_ref, 0)}; + + std::vector push_constants; + vkapi::SpecVarList spec_vars; + + if (graph.is_buffer_storage(out)) { + param_buffers.append(graph.sizes_ubo(out)); + param_buffers.append(graph.strides_ubo(out)); + + for (const ValueRef in_ref : batch_inputs) { + param_buffers.append(graph.sizes_ubo(in_ref)); + param_buffers.append(graph.strides_ubo(in_ref)); + } + + param_buffers.append(graph.numel_ubo(out)); + + spec_vars = {graph.hashed_layout_of(out)}; + } else { + push_constants = {graph.sizes_pc_of(out)}; + + spec_vars = {graph.hashed_layout_of(out)}; + + for (const ValueRef in_ref : batch_inputs) { + push_constants.push_back(graph.sizes_pc_of(in_ref)); + spec_vars.append(graph.hashed_layout_of(in_ref)); + } + } + + std::string kernel_name = "concat"; + if (current_batch_size == 1) { + kernel_name += "_1"; + } else if (current_batch_size == 2) { + kernel_name += "_2"; + } else if (current_batch_size == 3) { + kernel_name += "_3"; + } + if (graph.is_buffer_storage(out)) { + kernel_name += "_buffer"; + } else { + kernel_name += "_texture3d"; + } + + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + DispatchNode::ResizeFunction resize_fn = nullptr; + if (batch_start == 0) { + resize_fn = resize_concat_node; + } + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + concat_pick_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{out, vkapi::kReadWrite}, + {batch_inputs, vkapi::kRead}, + {concat_offset, vkapi::kRead}}, + // Parameter buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + spec_vars, + // Resize Args + {tensors_ref, dim_ref}, + // Resizing Logic + resize_fn)); } - } - std::string kernel_name = "concat"; - if (in_value_refs.size() == 1) { - kernel_name += "_1"; - } else if (in_value_refs.size() == 2) { - kernel_name += "_2"; - } else if (in_value_refs.size() == 3) { - kernel_name += "_3"; - } - if (graph.is_buffer_storage(out)) { - kernel_name += "_buffer"; - } else { - kernel_name += "_texture3d"; + // Add node to update concat_offset (except for the last batch) + if (batch_end < in_value_refs.size()) { + vkapi::ParamsBindList param_buffers = { + graph.get_or_create_int_param_buffer(dim_whcn_ref, 0)}; + + for (const ValueRef in_ref : batch_inputs) { + param_buffers.append(graph.sizes_ubo(in_ref)); + } + + std::string kernel_name = "update_concat_offset"; + if (current_batch_size == 1) { + kernel_name += "_1"; + } else if (current_batch_size == 2) { + kernel_name += "_2"; + } else if (current_batch_size == 3) { + kernel_name += "_3"; + } + add_dtype_suffix(kernel_name, graph.dtype_of(concat_offset)); + + vkapi::SpecVarList spec_vars = {}; + + graph.execute_nodes().emplace_back(new DispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + {1u, 1u, 1u}, + {1u, 1u, 1u}, + // Inputs and Outputs + {{concat_offset, vkapi::kWrite}}, + // Parameter buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + spec_vars, + // Resize Args + {}, + // Resizing Logic + nullptr)); + } } - - add_dtype_suffix(kernel_name, graph.dtype_of(out)); - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - default_pick_global_wg_size, - default_pick_local_wg_size, - // Inputs and Outputs - {{out, vkapi::kWrite}, {in_value_refs, vkapi::kRead}}, - // Parameter buffers - param_buffers, - // Push Constants - push_constants, - // Specialization Constants - spec_vars, - // Resize Args - {dim_ref}, - // Resizing Logic - resize_concat_node)); } void cat_tensor(ComputeGraph& graph, const std::vector& args) { diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index d85bd9d841e..b83164f27d2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -19,23 +20,31 @@ namespace vkcompute { +enum class Conv2dMethod : uint8_t { + Depthwise, + Pointwise, + SlidingWindow, + Transposed, +}; + void resize_conv2d_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); - size_t ndim = self->sizes().size(); + size_t ndim = graph->dim_of(self); std::vector new_out_sizes(ndim); - const bool transposed = graph->get_bool(extra_args[4]); + const bool transposed = graph->get_bool(extra_args.at(4)); + std::vector self_sizes = graph->sizes_of(self); // Batch, Channel if (ndim == 4) { - new_out_sizes.at(ndim - 4) = self->sizes().at(ndim - 4); + new_out_sizes.at(ndim - 4) = self_sizes.at(ndim - 4); } - TensorRefPtr weight_ref = graph->get_tref(extra_args[0]); + TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); const auto& weight_sizes = weight_ref->sizes; new_out_sizes.at(ndim - 3) = transposed ? weight_sizes.at(ndim - 3) : weight_sizes.at(ndim - 4); @@ -43,44 +52,44 @@ void resize_conv2d_node( // Height, Width const auto& new_out_sizes_hw = calc_out_sizes_hw( *graph, - self->sizes(), - extra_args[0], + self_sizes, + extra_args.at(0), /*kernel_size_only = */ false, - {extra_args[1], extra_args[2], extra_args[3], extra_args[5]}, + {extra_args.at(1), extra_args.at(2), extra_args.at(3), extra_args.at(5)}, transposed); new_out_sizes.at(ndim - 2) = new_out_sizes_hw.at(0); new_out_sizes.at(ndim - 1) = new_out_sizes_hw.at(1); - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } void resize_conv1d_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); - TensorRefPtr weight_ref = graph->get_tref(extra_args[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); + TensorRefPtr weight_ref = graph->get_tref(extra_args.at(0)); - int64_t stride_size = graph->get_int_list(extra_args[1])->at(0); - int64_t padding_size = graph->get_int_list(extra_args[2])->at(0); - int64_t dilation_size = graph->get_int_list(extra_args[3])->at(0); + const int64_t stride_size = graph->get_int_list(extra_args.at(1))->at(0); + const int64_t padding_size = graph->get_int_list(extra_args.at(2))->at(0); + const int64_t dilation_size = graph->get_int_list(extra_args.at(3))->at(0); const std::vector& weight_sizes = weight_ref->sizes; - const std::vector& in_sizes = self->sizes(); - size_t ndim = in_sizes.size(); + const std::vector in_sizes = graph->sizes_of(self); + const size_t ndim = in_sizes.size(); std::vector new_out_sizes(ndim); - int64_t kernel_size = weight_sizes.at(2); - int64_t in_length = in_sizes.at(2); + const int64_t kernel_size = weight_sizes.at(2); + const int64_t in_length = in_sizes.at(2); new_out_sizes.at(0) = in_sizes.at(0); new_out_sizes.at(1) = weight_sizes.at(0); new_out_sizes.at(2) = calc_out_size( in_length, kernel_size, stride_size, padding_size, dilation_size, false); - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } ValueRef prepack_biases( @@ -95,9 +104,9 @@ ValueRef prepack_biases( ValueRef v = graph.add_tensor( {out_channels}, graph.dtype_of(weight), storage_type, memory_layout); - vTensorPtr t = graph.get_tensor(v); - vkapi::ShaderInfo shader = get_nchw_to_tensor_shader(*t); + vkapi::ShaderInfo shader = + get_nchw_to_tensor_shader(graph, v, graph.dtype_of(weight)); graph.prepack_nodes().emplace_back(new PrepackNode( graph, @@ -108,22 +117,15 @@ ValueRef prepack_biases( v, {}, // Specialization constants - {t->hashed_layout()}, + {graph.hashed_layout_of(v)}, {graph.sizes_pc_of(v)})); return v; } -enum class Conv2dMethod : uint8_t { - Depthwise, - Pointwise, - SlidingWindow, - Transposed, -}; - vkapi::ShaderInfo get_conv2d_shader( ComputeGraph& graph, - const api::vTensor& t_out, + const ValueRef out, const bool prepack_weights, const Conv2dMethod method, const ValueRef weight, @@ -167,7 +169,7 @@ vkapi::ShaderInfo get_conv2d_shader( } else if (clamp_out) { kernel_name += "_clamp"; } - add_dtype_suffix(kernel_name, t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); return VK_KERNEL_FROM_STR(kernel_name); } @@ -206,10 +208,9 @@ ValueRef prepack_weights( graph.dtype_of(vref), utils::kTexture2D, utils::kChannelsPacked); - vTensorPtr t = graph.get_tensor(v); vkapi::ShaderInfo shader = - get_conv2d_shader(graph, *t, /*prepack_weights = */ true, method, vref); + get_conv2d_shader(graph, v, /*prepack_weights = */ true, method, vref); const auto original_sizes_pc = utils::make_ivec4(original_sizes, /*reverse = */ true); @@ -222,16 +223,19 @@ ValueRef prepack_weights( v, {}, // Specialization constants - {SV(t->packed_dim())}, + {graph.packed_dim_of(v)}, {graph.sizes_pc_of(v), PushConstantDataInfo(&original_sizes_pc, sizeof(original_sizes_pc))})); return v; } -void check_conv_args(const api::vTensor& in, const api::vTensor& out) { - VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim)); +void check_conv_args( + ComputeGraph& graph, + const ValueRef in, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } struct Conv2dParams final { @@ -277,9 +281,6 @@ Conv2dMethod get_conv2d_method( if (!transposed && weight_sizes.at(0) == groups && weight_sizes.at(1) == 1) { return Conv2dMethod::Depthwise; } - if (groups > 1) { - VK_THROW("aten.convolution.default: groups > 1 is not supported yet!"); - } if (transposed) { return Conv2dMethod::Transposed; } @@ -325,6 +326,108 @@ utils::uvec3 create_conv2d_global_wg_size( } } +// Custom global workgroup size function for conv2d +utils::uvec3 conv2d_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + 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_dw") != std::string::npos) { + method = Conv2dMethod::Depthwise; + } else 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; + } + + // Determine stride_equals_dilation from shader name + bool stride_equals_dilation = + shader.kernel_name.find("_sned") == std::string::npos; + + utils::uvec3 wg_size = create_conv2d_global_wg_size( + *graph, method, out, weight_data, stride_equals_dilation); + + if (method == Conv2dMethod::Depthwise || method == Conv2dMethod::Pointwise) { + wg_size = {wg_size[0] * wg_size[1], wg_size[2], 1}; + } + + return wg_size; +} + +// Custom local workgroup size function for conv2d +utils::uvec3 conv2d_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + 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_dw") != std::string::npos) { + method = Conv2dMethod::Depthwise; + } else 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; + } + + if (method == Conv2dMethod::Pointwise) { + uint32_t local_wg_size_y = 1; + if (global_workgroup_size[1] % 8 == 0) { + local_wg_size_y = 8; + } else if (global_workgroup_size[1] % 4 == 0) { + local_wg_size_y = 4; + } else if (global_workgroup_size[1] % 2 == 0) { + local_wg_size_y = 2; + } + return {64 / local_wg_size_y, local_wg_size_y, 1}; + } else if (method == Conv2dMethod::Depthwise) { + return {64, 1, 1}; + } else { + return graph->create_local_wg_size(global_workgroup_size); + } +} + +// Custom global workgroup size function for conv1d +utils::uvec3 conv1d_global_wg_size( + 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 {// out length + graph->size_at(-1, out), + // out channels + static_cast(graph->size_at(-2, out)), + // out batches + utils::div_up_4(graph->size_at(-3, out))}; +} + void add_conv2d_node( ComputeGraph& graph, const ValueRef in, @@ -365,12 +468,12 @@ void add_conv2d_node( /* storage_type = */ utils::kTexture2D, /* memory_layout = */ utils::kWidthPacked); - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - if (t_in->sizes().at(0) > 1) { + const std::vector in_sizes = graph.sizes_of(in); + if (in_sizes.at(0) > 1) { VK_THROW("conv2d: input batch size > 1 is not supported yet!"); } - check_conv_args(*t_in, *t_out); + + check_conv_args(graph, in, out); Kernel2dParams kernel_params = create_kernel2d_params( graph, @@ -396,7 +499,7 @@ void add_conv2d_node( vkapi::ShaderInfo shader = get_conv2d_shader( graph, - *t_out, + out, /*prepack_weights = */ false, method, weight_data, @@ -476,19 +579,19 @@ void add_conv2d_node( }; } else { param_buffers = { - t_out->logical_limits_ubo(), - t_in->sizes_ubo(), + graph.logical_limits_ubo(out), + graph.sizes_ubo(in), graph.create_params_buffer(kernel_params), graph.create_params_buffer(extra_params), graph.create_params_buffer(out_params), }; } - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, shader, - wg_size, - local_wg_size, + conv2d_global_wg_size, + conv2d_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers @@ -496,7 +599,7 @@ void add_conv2d_node( // Push Constants push_constants, // Specialization Constants - {}, + {utils::safe_downcast(groups_val)}, // Resize Args {weight_data, stride, padding, dilation, transposed, output_padding}, // Resizing Logic @@ -540,17 +643,13 @@ void add_conv1d_node( out_max_val = graph.extract_scalar(out_max); } - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_weight = graph.get_tensor(arg_weight); - vTensorPtr t_bias = graph.get_tensor(arg_bias); - vTensorPtr t_out = graph.get_tensor(out); const int64_t groups_val = graph.get_int(groups); - std::vector in_sizes = t_in->sizes(); - std::vector weight_sizes = t_weight->sizes(); - std::vector out_sizes = t_out->sizes(); + const std::vector in_sizes = graph.sizes_of(in); + const std::vector weight_sizes = graph.sizes_of(arg_weight); + const std::vector out_sizes = graph.sizes_of(out); - check_conv_args(*t_in, *t_out); + check_conv_args(graph, in, out); const int32_t in_channels = in_sizes.at(1); const int32_t out_channels = weight_sizes.at(0); @@ -562,15 +661,6 @@ void add_conv1d_node( const int32_t out_group_size = static_cast(out_channels / groups_val); - const utils::uvec3 global_size = { - // out length - graph.size_at(-1, out), - // out channels - static_cast(out_channels), - // out batches - utils::div_up_4(graph.size_at(-3, out))}; - const utils::uvec3 local_size = graph.create_local_wg_size(global_size); - Kernel1dParams kernel_params = { kernel_size, stride_size, @@ -587,29 +677,29 @@ void add_conv1d_node( } kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - local_size, + conv1d_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers { - t_out->logical_limits_ubo(), - t_in->sizes_ubo(), + graph.logical_limits_ubo(out), + graph.sizes_ubo(in), graph.create_params_buffer(kernel_params), graph.create_params_buffer(out_params), }, // Push Constants {}, // Specialization Constants - {t_out->hashed_layout(), - t_in->hashed_layout(), - t_weight->hashed_layout(), - t_bias->hashed_layout()}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(in), + graph.hashed_layout_of(arg_weight), + graph.hashed_layout_of(arg_bias)}, // Resize Args {weight, stride, padding, dilation}, // Resizing Logic @@ -617,7 +707,7 @@ void add_conv1d_node( } void conv(ComputeGraph& graph, const std::vector& args) { - int64_t in_ndim = graph.get_tensor(args[0])->sizes().size(); + int64_t in_ndim = graph.dim_of(args[0]); if (in_ndim == 4) { if (args.size() == 10) { // ordinary conv2d diff --git a/backends/vulkan/runtime/graph/ops/impl/Copy.cpp b/backends/vulkan/runtime/graph/ops/impl/Copy.cpp index c4f37bd9386..bd648dbae2d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Copy.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Copy.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -28,21 +29,18 @@ void add_copy_offset_node( const ValueRef out, bool calc_out_pos_using_src_chnl, bool calc_in_pos_using_dst_chnl) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - std::string kernel_name = "copy_offset"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); - add_storage_type_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); auto shader = VK_KERNEL_FROM_STR(kernel_name); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs { {out, vkapi::kWrite}, @@ -75,27 +73,27 @@ void add_copy_packed_dim_offset_node( const ivec4& src_offset, const ivec4& dst_offset, const ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - // Check the packed dimension is same for both tensors, also check if the // packed dimension is Width or Height. Since the function does not support // channel packing. VK_CHECK_COND( - check_same_packed_dim(*t_in, *t_out) && - (check_packed_dim_is(*t_in, WHCN::kWidthDim) || - check_packed_dim_is(*t_in, WHCN::kHeightDim))); + graph.packed_dim_of(in) == graph.packed_dim_of(out) && + (graph.packed_dim_of(in) == WHCN::kWidthDim || + graph.packed_dim_of(in) == WHCN::kHeightDim)); std::string kernel_name = "copy_packed_dim_offset"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + const std::vector in_sizes = graph.sizes_of(in); + const std::vector out_sizes = graph.sizes_of(out); // A copy of range with the last element set to batch size of the input tensor ivec4 final_range = { - range[0], range[1], range[2], dim_at(t_in->sizes(), kBatch4D)}; - ivec3 global_wg_size = t_out->logical_limits(); + range[0], range[1], range[2], dim_at(in_sizes, kBatch4D)}; + ivec3 global_wg_size = graph.logical_limits_of(out); - const auto packed_dim = t_in->packed_dim(); + const auto packed_dim = graph.packed_dim_of(in); // The starting offset in a texel where this tensor will start copying from const auto src_lane_offset = src_offset[packed_dim] & 0x3; // The starting offset in a texel where this tensor will start copying to @@ -106,16 +104,14 @@ void add_copy_packed_dim_offset_node( // remaining lanes from current source Hence (4 - src_lane_offset) is added // to tensor size in packed dimension const auto src_packed_size = utils::div_up_4( - (4 - src_lane_offset) + - dim_at(t_out->sizes(), normalize_to_dim_index(*t_out, packed_dim))); + (4 - src_lane_offset) + utils::val_at(-packed_dim, out_sizes)); // The total packed texels this tensor will be copied to // The first texel of tensor data in packed dimension will be copied to // remaining lanes from previous write Hence (4 - dst_lane_offset) is added // to tensor size in packed dimension const auto dst_packed_size = utils::div_up_4( - (4 - dst_lane_offset) + - dim_at(t_in->sizes(), normalize_to_dim_index(*t_in, packed_dim))); + (4 - dst_lane_offset) + utils::val_at(-packed_dim, in_sizes)); // If the starting src offset is not 0, and the total packed texels is // greater than the source texel range @@ -169,20 +165,17 @@ void add_copy_channel_offset_node( int32_t src_channel_offset, int32_t dst_channel_offset, const ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - // Likely need to prepad these numbers. - std::vector in_sizes = t_in->sizes(); - std::vector out_sizes = t_out->sizes(); + const std::vector in_sizes = graph.sizes_of(in); + const std::vector out_sizes = graph.sizes_of(out); - VK_CHECK_COND(check_packed_dim_is(*t_in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(*t_out, WHCN::kChannelsDim)); + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); // NOTE: This function should be able to support 1d and 2d tensors when // range=1, src_offset=dst_offset=1. - VK_CHECK_COND(t_in->dim() >= 3, "Src dim should be at least 3"); - VK_CHECK_COND(t_out->dim() >= 3, "Dst dim should be at least 3"); + VK_CHECK_COND(graph.dim_of(in) >= 3, "Src dim should be at least 3"); + VK_CHECK_COND(graph.dim_of(out) >= 3, "Dst dim should be at least 3"); VK_CHECK_COND( dim_at(in_sizes) >= src_channel_offset + channel_range, @@ -212,7 +205,7 @@ void add_copy_channel_offset_node( std::string kernel_name = "copy_channel_offset"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); int32_t out_channels = dim_at(out_sizes); diff --git a/backends/vulkan/runtime/graph/ops/impl/Dequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/Dequantize.cpp index 61fd76145a4..a217734653d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Dequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Dequantize.cpp @@ -23,10 +23,11 @@ void resize_dequantize_node( const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - out->virtual_resize(in->sizes()); + const std::vector in_sizes = graph->sizes_of(in); + graph->virtual_resize(out, in_sizes); } utils::uvec3 dequantize_per_channel_local_wg_size( @@ -106,9 +107,35 @@ void add_dequantize_per_tensor_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(input)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } vkapi::ParamsBindList param_ubos; std::vector push_constants; @@ -168,9 +195,35 @@ void add_dequantize_per_token_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(input)); + + int quant_min_val, quant_max_val; - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } int num_tokens = static_cast(graph.sizes_of(scale)[0]); @@ -234,10 +287,37 @@ void add_dequantize_per_channel_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); int axis_val = static_cast(graph.get_int(axis)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(input)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } // Normalize axis and convert from NCHW to WHCN using utility functions const auto input_sizes = graph.sizes_of(input); @@ -319,9 +399,35 @@ void add_dequantize_block_wise_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(input)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } const auto input_sizes = graph.sizes_of(input); const auto block_size_list = graph.get_int_list(block_size); @@ -422,6 +528,18 @@ void dequantize_per_tensor_impl( graph.dtype_of(input) == vkapi::kChar || graph.dtype_of(input) == vkapi::kInt); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); @@ -469,6 +587,18 @@ void dequantize_per_token_impl( graph.dtype_of(input) == vkapi::kChar || graph.dtype_of(input) == vkapi::kInt); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); @@ -544,6 +674,18 @@ void dequantize_per_channel_impl( graph.dtype_of(input) == vkapi::kChar || graph.dtype_of(input) == vkapi::kInt); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); @@ -624,6 +766,18 @@ void dequantize_affine_impl( graph.dtype_of(input) == vkapi::kChar || graph.dtype_of(input) == vkapi::kInt); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); diff --git a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp index 85c80e01c27..475e7796b09 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -23,15 +24,16 @@ using utils::GPUMemoryLayout; using utils::StorageType; void check_embedding_args( - const api::vTensor& weight, - const api::vTensor& in, - const api::vTensor& out) { + ComputeGraph& graph, + const ValueRef weight, + const ValueRef in, + const ValueRef out) { // The packing logic may not be trivial here. Input and output are Channel // Packed, which is default for the Vulkan backend. However, weight vector is // height-packed instead of channel-packed for space reason. - VK_CHECK_COND(check_packed_dim_is(weight, WHCN::kHeightDim)); - VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim)); + VK_CHECK_COND(graph.packed_dim_of(weight) == WHCN::kHeightDim); + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } void add_embedding_node( @@ -39,31 +41,27 @@ void add_embedding_node( ValueRef weight, ValueRef in, ValueRef out) { - vTensorPtr t_weight = graph.get_tensor(weight); - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - - check_embedding_args(*t_weight, *t_in, *t_out); + check_embedding_args(graph, weight, in, out); std::string kernel_name = "embedding"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, {{out, vkapi::kWrite}, {{in, weight}, vkapi::kRead}}, { - t_out->sizes_ubo(), + graph.sizes_ubo(out), }, // Push Constants {}, // Specialization Constants - {t_out->hashed_layout(), - t_in->hashed_layout(), - t_weight->hashed_layout()}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(in), + graph.hashed_layout_of(weight)}, // Resize Args {}, // Resizing Logic diff --git a/backends/vulkan/runtime/graph/ops/impl/Expand.cpp b/backends/vulkan/runtime/graph/ops/impl/Expand.cpp new file mode 100644 index 00000000000..1623a26b2a1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/Expand.cpp @@ -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. + */ + +#include + +#include +#include +#include +#include +#include + +namespace vkcompute { + +void add_expand_buffer_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef size, + const ValueRef out) { + std::string kernel_name = "expand"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + vkapi::ParamsBindList param_buffers = { + graph.buffer_meta_ubo(out), + graph.buffer_meta_ubo(in), + }; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_global_wg_size, + default_pick_local_wg_size, + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + // Parameter buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {size}, + // Resizing Logic + nullptr)); +} + +void expand(ComputeGraph& graph, const std::vector& args) { + int idx = 0; + const ValueRef in = args.at(idx++); + const ValueRef size = args.at(idx++); + const ValueRef implicit = args.at(idx++); + (void)implicit; + const ValueRef out = args.at(idx++); + + if (graph.is_buffer_storage(out)) { + return add_expand_buffer_node(graph, in, size, out); + } + + VK_THROW("Expand operator only supports buffer storage"); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(aten.expand_copy.default, expand); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Flip.cpp b/backends/vulkan/runtime/graph/ops/impl/Flip.cpp index 04aac2484ac..52288734704 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Flip.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Flip.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -15,9 +16,24 @@ namespace vkcompute { -void check_flip_args(const api::vTensor& in, const api::vTensor& out) { - VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim)); +// Custom global workgroup size function for flip +utils::uvec3 flip_global_wg_size( + 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_global_wg_size(out); +} + +void check_flip_args( + ComputeGraph& graph, + const ValueRef in, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } void resize_flip_node( @@ -25,10 +41,10 @@ void resize_flip_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - out->virtual_resize(in->sizes()); + graph->virtual_resize(out, graph->sizes_of(in)); } utils::ivec4 create_whcn_bitmap( @@ -48,21 +64,19 @@ void add_flip_node( const ValueRef in, const std::vector& dim_list, const ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - check_flip_args(*t_in, *t_out); + check_flip_args(graph, in, out); - const auto dim_bitmap = create_whcn_bitmap(dim_list, t_in->dim()); + const auto dim_bitmap = create_whcn_bitmap(dim_list, graph.dim_of(in)); std::string kernel_name("flip"); kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + flip_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs { {out, vkapi::kWrite}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Full.cpp b/backends/vulkan/runtime/graph/ops/impl/Full.cpp index 3ed18445463..fe2676e91e0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Full.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Full.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -19,44 +20,42 @@ void resize_full_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); + const ValueRef out = args.at(0).refs.at(0); std::vector out_sizes; - if (graph->val_is_tensor(extra_args[0])) { - out_sizes = graph->get_tensor(extra_args[0])->sizes(); + if (graph->val_is_tensor(extra_args.at(0))) { + out_sizes = graph->sizes_of(extra_args.at(0)); } else { - out_sizes = *graph->get_int_list(extra_args[0]); + out_sizes = *graph->get_int_list(extra_args.at(0)); } - out->virtual_resize(out_sizes); + graph->virtual_resize(out, out_sizes); } -// size_or_in is IntListPtr when op is full and vTensorPtr if op is full_like void add_full_node( ComputeGraph& graph, const ValueRef size_or_in, const ValueRef fill_value, const ValueRef out) { float fill_value_val = graph.extract_scalar(fill_value); - vTensorPtr t_out = graph.get_tensor(out); std::string kernel_name("full"); kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}}, // Shader params buffers - {t_out->sizes_ubo(), graph.create_params_buffer(fill_value_val)}, + {graph.sizes_ubo(out), graph.create_params_buffer(fill_value_val)}, // Push Constants {}, // Specialization Constants - {SV(t_out->packed_dim())}, + {graph.packed_dim_of(out)}, // Resize Args {size_or_in}, // Resizing Logic diff --git a/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp b/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp index 0624020c872..5f39c16d405 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GridPriors.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -23,13 +24,13 @@ void resize_grid_priors_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(extra_args[0]); - std::vector in_sizes = in->sizes(); - int64_t height = in_sizes.at(in_sizes.size() - 2); - int64_t width = in_sizes.at(in_sizes.size() - 1); - std::vector sizes = {height * width, 2}; - out->virtual_resize(sizes); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = extra_args.at(0); + const std::vector in_sizes = graph->sizes_of(in); + const int64_t height = in_sizes.at(in_sizes.size() - 2); + const int64_t width = in_sizes.at(in_sizes.size() - 1); + const std::vector sizes = {height * width, 2}; + graph->virtual_resize(out, sizes); } void add_grid_priors_node( @@ -38,29 +39,27 @@ void add_grid_priors_node( const ValueRef& stride_ref, const ValueRef& offset_ref, const ValueRef& out) { - vTensorPtr t_out = graph.get_tensor(out); - vTensorPtr t_in = graph.get_tensor(in); - int32_t stride = graph.extract_scalar(stride_ref); - float offset = graph.extract_scalar(offset_ref); + const int32_t stride = graph.extract_scalar(stride_ref); + const float offset = graph.extract_scalar(offset_ref); std::string kernel_name = "grid_priors"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - GridPriorsParam param = {stride, offset}; - graph.execute_nodes().emplace_back(new DispatchNode( + const GridPriorsParam param = {stride, offset}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs { {out, vkapi::kWrite}, }, // Shader params buffers { - t_in->sizes_ubo(), - t_out->sizes_ubo(), + graph.sizes_ubo(in), + graph.sizes_ubo(out), graph.create_params_buffer(param), }, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp index 8d2a848b0c4..368b95c9d3b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GroupNorm.cpp @@ -17,14 +17,6 @@ namespace vkcompute { -std::vector calc_group_norm_mean_sizes( - api::vTensor& self, - const int64_t group) { - const std::vector& input_sizes = self.sizes(); - const int64_t N = input_sizes.at(0); - return {N, group}; -} - utils::uvec3 group_norm_local_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp index 8203829c50f..576711a86f1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -18,12 +19,13 @@ namespace vkcompute { void check_index_select_args( - const api::vTensor& in, - const api::vTensor& idx, - const api::vTensor& out) { - VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(idx, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim)); + ComputeGraph& graph, + const ValueRef in, + const ValueRef idx, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(idx) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } void add_index_select_channel_node( @@ -31,23 +33,19 @@ void add_index_select_channel_node( ValueRef in, ValueRef idx, ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_idx = graph.get_tensor(idx); - vTensorPtr t_out = graph.get_tensor(out); - - check_index_select_args(*t_in, *t_idx, *t_out); + check_index_select_args(graph, in, idx, out); std::string kernel_name = "index_select_channel"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, - {t_out->sizes_ubo(), t_in->sizes_ubo()}, + {graph.sizes_ubo(out), graph.sizes_ubo(in)}, // Push Constants {}, // Specialization Constants @@ -64,14 +62,16 @@ struct IndexSelectParams final { }; IndexSelectParams create_index_select_params( + ComputeGraph& graph, const int64_t dim_idx, - const api::vTensor& in) { + const ValueRef in) { if (dim_idx == kWidth4D) { return {0, 1}; } else if (dim_idx == kHeight4D) { return {1, 1}; } else if (dim_idx == kBatch4D) { - int64_t n_channels = dim_at(in.sizes(), kChannel4D); + 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)}; } else { @@ -85,25 +85,21 @@ void add_index_select_node( const int64_t dim_idx, ValueRef idx, ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_idx = graph.get_tensor(idx); - vTensorPtr t_out = graph.get_tensor(out); + check_index_select_args(graph, in, idx, out); - check_index_select_args(*t_in, *t_idx, *t_out); - - IndexSelectParams params = create_index_select_params(dim_idx, *t_in); + IndexSelectParams params = create_index_select_params(graph, dim_idx, in); std::string kernel_name = "index_select"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, - {t_out->sizes_ubo(), graph.create_params_buffer(params)}, + {graph.sizes_ubo(out), graph.create_params_buffer(params)}, // Push Constants {}, // Specialization Constants @@ -115,10 +111,12 @@ void add_index_select_node( } int64_t get_dim_idx(ComputeGraph& graph, ValueRef in, ValueRef dim_ref) { - vTensorPtr t_in = graph.get_tensor(in); int64_t dim = graph.extract_scalar(dim_ref); - dim = normalize(dim, t_in->dim()); - return normalize_to_dim_index(*t_in, dim); + const int64_t ndim = graph.dim_of(in); + dim = normalize(dim, ndim); + + // Convert to DimIndex - this replicates normalize_to_dim_index logic + return dim < 0 ? dim : dim - ndim; } void index_select(ComputeGraph& graph, const std::vector& args) { diff --git a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp index 14ed9c84a32..38d70271f4f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Linear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Linear.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -18,6 +19,70 @@ namespace vkcompute { +// Custom global workgroup size function for addmm_naive_texture +utils::uvec3 addmm_naive_texture_global_wg_size( + 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->logical_limits_of(out); +} + +// Custom global workgroup size function for addmm_naive_buffer +utils::uvec3 addmm_naive_buffer_global_wg_size( + 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->size_at(-1, out), + graph->size_at(-2, out), + graph->size_at(-3, out) * graph->size_at(-4, out)}; +} + +// Custom global workgroup size function for addmm_optimized +utils::uvec3 addmm_optimized_global_wg_size( + 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); + const ValueRef mat1 = args.at(1).refs.at(0); + + std::vector mat1_sizes = graph->sizes_of(mat1); + int mat1_dims = mat1_sizes.size(); + + utils::uvec3 global_size = graph->logical_limits_of(out); + + if (mat1_sizes.at(mat1_dims - 2) < 8) { + global_size = utils::divup_vec(global_size, {4, 2, 1}); + } else { + global_size = utils::divup_vec(global_size, {4, 4, 1}); + } + return global_size; +} + +// Custom local workgroup size function for addmm_optimized +utils::uvec3 addmm_optimized_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)args; + (void)resize_args; + return adaptive_work_group_size(global_workgroup_size); +} + void check_addmm_args( ComputeGraph& graph, const ValueRef self, @@ -54,29 +119,31 @@ void resize_addmm_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr mat1 = graph->get_tensor(args[1].refs[0]); - vTensorPtr mat2 = graph->get_tensor(args[1].refs[1]); - vTensorPtr self = graph->get_tensor(args[1].refs[2]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef mat1 = args.at(1).refs.at(0); + const ValueRef mat2 = args.at(1).refs.at(1); + + const bool mat2_is_transposed = graph->get_bool(extra_args.at(0)); - bool mat2_is_transposed = graph->get_bool(extra_args[0]); + const std::vector mat1_sizes = graph->sizes_of(mat1); + const std::vector mat2_sizes = graph->sizes_of(mat2); - const int out_cols = utils::val_at(-2, mat1->sizes()); - const int out_rows = mat2_is_transposed ? utils::val_at(-2, mat2->sizes()) - : utils::val_at(-1, mat2->sizes()); + const int out_cols = utils::val_at(-2, mat1_sizes); + const int out_rows = mat2_is_transposed ? utils::val_at(-2, mat2_sizes) + : utils::val_at(-1, mat2_sizes); std::vector new_out_sizes(3); - if (mat1->sizes().size() == 2) { + if (mat1_sizes.size() == 2) { new_out_sizes.resize(2); new_out_sizes.at(0) = out_cols; new_out_sizes.at(1) = out_rows; } else { - new_out_sizes.at(0) = mat1->sizes().at(0); + new_out_sizes.at(0) = mat1_sizes.at(0); new_out_sizes.at(1) = out_cols; new_out_sizes.at(2) = out_rows; } - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } struct Params final { @@ -107,11 +174,11 @@ void add_addmm_naive_texture_node( add_dtype_suffix(kernel_name, graph.dtype_of(out)); utils::uvec3 global_wg_size = graph.logical_limits_of(out); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), + addmm_naive_texture_global_wg_size, + pick_hw_square_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2, self}, vkapi::kRead}}, // Shader params buffers @@ -174,11 +241,11 @@ void add_addmm_naive_buffer_node( ? 1 : 0; - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - graph.create_local_wg_size(global_size), + addmm_naive_buffer_global_wg_size, + pick_hw_square_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2, self}, vkapi::kRead}}, // Shader params buffers @@ -248,31 +315,13 @@ void add_addmm_optimized_node( } else { kernel_name += "_tile_row_4"; } - add_dtype_suffix(kernel_name, graph.dtype_of(out)); - utils::uvec3 global_size = graph.logical_limits_of(out); - - // Each thread computes a W=(2/4) x H=4 x C=(1/4) output tile. Therefore, the - // total number of threads is W/(2 or 4) x H/4 x C/1. Since the out tensor is - // channels packed, C does not need to be divided by 4. The "identity" of each - // thread is the (x, y, z) coordinate of the output tile it is computing, and - // this identity can be used to compute the tensor index of the top left - // element in the tile, which will be [W=x*(2 or 4), H=y*4, C=z*(1 or 4), N=0] - if (mat1_sizes.at(mat1_dims - 2) < 8) { - // Use `logical_extents` instead of `image_extents` because the workgroup - // axes need to correspond to tensor dimensions. - global_size = utils::divup_vec(global_size, {4, 2, 1}); - } else { - global_size = utils::divup_vec(global_size, {4, 4, 1}); - } - utils::uvec3 local_size = adaptive_work_group_size(global_size); - - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - local_size, + addmm_optimized_global_wg_size, + addmm_optimized_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1_W_packed, mat2_packed, self}, vkapi::kRead}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp index 73a625f3adf..47ecf5f18d2 100644 --- a/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/MatMul.cpp @@ -39,22 +39,25 @@ void resize_matmul_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr mat1 = graph->get_tensor(args[1].refs[0]); - vTensorPtr mat2 = graph->get_tensor(args[1].refs[1]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef mat1 = args.at(1).refs.at(0); + const ValueRef mat2 = args.at(1).refs.at(1); + + bool mat2_is_transposed = graph->get_bool(resize_args.at(0)); - bool mat2_is_transposed = graph->get_bool(resize_args[0]); + const std::vector mat1_sizes = graph->sizes_of(mat1); + const std::vector mat2_sizes = graph->sizes_of(mat2); - const int out_cols = utils::val_at(-2, mat1->sizes()); - const int out_rows = mat2_is_transposed ? utils::val_at(-2, mat2->sizes()) - : utils::val_at(-1, mat2->sizes()); + const int out_cols = utils::val_at(-2, mat1_sizes); + const int out_rows = mat2_is_transposed ? utils::val_at(-2, mat2_sizes) + : utils::val_at(-1, mat2_sizes); - const int64_t out_dim = out->dim(); - std::vector new_out_sizes(mat1->sizes()); + const int64_t out_dim = graph->dim_of(out); + std::vector new_out_sizes(mat1_sizes); new_out_sizes.at(out_dim - 1) = out_rows; new_out_sizes.at(out_dim - 2) = out_cols; - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } /** @@ -99,7 +102,7 @@ void add_matmul_naive_buffer_node( graph, VK_KERNEL_FROM_STR(kernel_name), matmul_naive_buffer_global_wg_size, - default_pick_local_wg_size, + pick_hw_square_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, // Shader params buffers @@ -155,7 +158,7 @@ void add_matmul_naive_texture3d_node( graph, pick_matmul_naive_texture3d_shader, default_pick_global_wg_size, - default_pick_local_wg_size, + pick_hw_square_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, mat2}, vkapi::kRead}}, // Shader params buffers @@ -270,7 +273,7 @@ void add_matmul_optimized_node( graph, pick_matmul_optimized_shader, matmul_optimized_global_wg_size, - default_pick_local_wg_size, + pick_hw_square_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1_W_packed, mat2_packed}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp index 100d6e33931..8e15b56b208 100644 --- a/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -18,10 +19,10 @@ namespace vkcompute { std::vector calc_out_mean_sizes( - api::vTensor& self, + const std::vector& self_sizes, int64_t normalized_shape_dim) { - std::vector output_size = self.sizes(); - int64_t self_dim = self.sizes().size(); + std::vector output_size = self_sizes; + int64_t self_dim = self_sizes.size(); for (int64_t i = 0; i < normalized_shape_dim; ++i) { output_size.at(self_dim - i - 1) = 1; } @@ -32,20 +33,21 @@ void resize_native_layer_norm_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr mean = graph->get_tensor(args[0].refs[1]); - vTensorPtr rstd = graph->get_tensor(args[0].refs[2]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); - std::vector in_sizes = in->sizes(); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef mean = args.at(0).refs.at(1); + const ValueRef rstd = args.at(0).refs.at(2); + const ValueRef in = args.at(1).refs.at(0); + const std::vector in_sizes = graph->sizes_of(in); - const auto normalized_shape_dim = graph->get_int_list(extra_args[0])->size(); + const auto normalized_shape_dim = + graph->get_int_list(extra_args.at(0))->size(); - std::vector mean_size = - calc_out_mean_sizes(*in, normalized_shape_dim); + const std::vector mean_size = + calc_out_mean_sizes(in_sizes, normalized_shape_dim); - out->virtual_resize(in_sizes); - mean->virtual_resize(mean_size); - rstd->virtual_resize(mean_size); + graph->virtual_resize(out, in_sizes); + graph->virtual_resize(mean, mean_size); + graph->virtual_resize(rstd, mean_size); } void add_native_layer_norm_node( @@ -74,16 +76,17 @@ void add_native_layer_norm_node( ValueRef arg_bias = prepack_standard_like(graph, bias_data, in); const auto out_val = graph.get_value_list(out); - vTensorPtr t_out = graph.get_tensor(out_val->at(0)); - vTensorPtr t_mean = graph.get_tensor(out_val->at(1)); - vTensorPtr t_input = graph.get_tensor(in); + const ValueRef out_tensor = out_val->at(0); + const ValueRef mean_tensor = out_val->at(1); + const ValueRef rstd_tensor = out_val->at(2); + float epsilon = graph.extract_scalar(eps); - VK_CHECK_COND(check_same_packed_dim(*t_input, *t_out)); + VK_CHECK_COND(check_same_packed_dim(graph, in, out_tensor)); - std::vector in_sizes = t_input->sizes(); + const std::vector in_sizes = graph.sizes_of(in); - utils::uvec3 global_size = t_out->logical_limits(); + utils::uvec3 global_size = graph.logical_limits_of(out_tensor); utils::uvec3 local_size; // Since the shader sets shared memory scale factor > 1, if dispatch is @@ -100,28 +103,28 @@ void add_native_layer_norm_node( std::string kernel_name("native_layer_norm"); kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out_tensor)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - local_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs - {{{out_val->at(0), out_val->at(1), out_val->at(2)}, vkapi::kWrite}, + {{{out_tensor, mean_tensor, rstd_tensor}, vkapi::kWrite}, {{in, arg_weight, arg_bias}, vkapi::kRead}}, // Shader params buffers {}, // Push Constants { - graph.logical_limits_pc_of(out_val->at(0)), - graph.sizes_pc_of(out_val->at(0)), + graph.logical_limits_pc_of(out_tensor), + graph.sizes_pc_of(out_tensor), PushConstantDataInfo(&epsilon, sizeof(epsilon)), }, // Specialization Constants { - t_input->hashed_layout(), - t_out->hashed_layout(), + graph.hashed_layout_of(in), + graph.hashed_layout_of(out_tensor), }, // Resize Args {normalized_shape}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Pad.cpp b/backends/vulkan/runtime/graph/ops/impl/Pad.cpp index 8f3ba7532a9..d225af05633 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Pad.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Pad.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -41,17 +42,17 @@ void resize_constant_pad_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); - IntListPtr pad_vec = graph->get_int_list(extra_args[0]); - std::vector in_size = self->sizes(); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); + const IntListPtr pad_vec = graph->get_int_list(extra_args.at(0)); + std::vector in_size = graph->sizes_of(self); int dim = in_size.size() - 1; for (int i = 0; i < pad_vec->size(); i += 2) { in_size.at(dim) += pad_vec->at(i) + pad_vec->at(i + 1); dim--; } - out->virtual_resize(in_size); + graph->virtual_resize(out, in_size); } void add_constant_pad_nd_node( @@ -60,34 +61,32 @@ void add_constant_pad_nd_node( const ValueRef& pad, const ValueRef& fill_value, const ValueRef& out) { - float fill_value_val = graph.extract_scalar(fill_value); - IntListPtr pad_vec = graph.get_int_list(pad); - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); + const float fill_value_val = graph.extract_scalar(fill_value); + const IntListPtr pad_vec = graph.get_int_list(pad); std::string kernel_name = ""; - PadParam pad_param = creat_pad_param(*pad_vec); + const PadParam pad_param = creat_pad_param(*pad_vec); if (pad_vec->size() <= 4) { kernel_name = "pad_height_width"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); } else { kernel_name = "pad_channel"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); } - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers - {t_out->sizes_ubo(), - t_in->sizes_ubo(), + {graph.sizes_ubo(out), + graph.sizes_ubo(in), graph.create_params_buffer(pad_param), graph.create_params_buffer(fill_value_val)}, // Push Constants diff --git a/backends/vulkan/runtime/graph/ops/impl/Permute.cpp b/backends/vulkan/runtime/graph/ops/impl/Permute.cpp index 6e6a6fa3bf2..9ac4c963bc3 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Permute.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Permute.cpp @@ -129,37 +129,22 @@ void add_permute_node( std::vector push_constants; vkapi::SpecVarList spec_vars; - if (graph.is_buffer_storage(out)) { - param_buffers.append(graph.sizes_ubo(in)); - param_buffers.append(graph.strides_ubo(out)); - param_buffers.append(graph.numel_ubo(out)); - - // Buffer storage - use permute_buffer shader - push_constants = { - graph.strides_pc_of(in), - PushConstantDataInfo(&whcn_permute_dims, sizeof(whcn_permute_dims)), - }; - - spec_vars = {graph.hashed_layout_of(out), graph.hashed_layout_of(in)}; - } else { - // Texture storage - use permute_texture shader - const int32_t out_channels = dim_at(graph.sizes_of(out)); - const int32_t in_channels = dim_at(graph.sizes_of(in)); - - const int32_t packed_dim = graph.packed_dim_of(in); - ivec2 channel_info = {out_channels, in_channels}; - if (packed_dim == WHCN::kChannelsDim) { - channel_info[0] = utils::align_up_4(channel_info[0]); - channel_info[1] = utils::align_up_4(channel_info[1]); - } + const int32_t out_channels = dim_at(graph.sizes_of(out)); + const int32_t in_channels = dim_at(graph.sizes_of(in)); + + const int32_t packed_dim = graph.packed_dim_of(in); + ivec2 channel_info = {out_channels, in_channels}; + if (packed_dim == WHCN::kChannelsDim) { + channel_info[0] = utils::align_up_4(channel_info[0]); + channel_info[1] = utils::align_up_4(channel_info[1]); + } - push_constants = { - graph.sizes_pc_of(out), - graph.sizes_pc_of(in), - PushConstantDataInfo(&whcn_permute_dims, sizeof(whcn_permute_dims))}; + push_constants = { + graph.sizes_pc_of(out), + graph.sizes_pc_of(in), + PushConstantDataInfo(&whcn_permute_dims, sizeof(whcn_permute_dims))}; - spec_vars = {graph.hashed_layout_of(out), graph.hashed_layout_of(in)}; - } + spec_vars = {graph.hashed_layout_of(out), graph.hashed_layout_of(in)}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -179,8 +164,83 @@ void add_permute_node( resize_permute_node)); } +struct WHCNPermuteDims { + int32_t whcn_permute_dims[api::kTensorDimLimit]; + + void initialize(const std::vector& permute_dims) { + const int32_t permute_ndim = permute_dims.size(); + for (int32_t whcn_i = 0; whcn_i < permute_ndim; whcn_i++) { + const int32_t nchw_i = permute_ndim - 1 - whcn_i; + int64_t index_val = permute_dims.at(nchw_i); + if (index_val < 0) { + index_val += permute_ndim; + } + const int32_t permute_dim_whcn = permute_ndim - 1 - index_val; + whcn_permute_dims[whcn_i] = permute_dim_whcn; + } + for (int32_t whcn_i = permute_ndim; whcn_i < api::kTensorDimLimit; + whcn_i++) { + whcn_permute_dims[whcn_i] = whcn_i; + } + } +}; + +void add_permute_buffer_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef permute_dims, + const ValueRef out) { + check_args(graph, in, permute_dims, out); + + WHCNPermuteDims whcn_permute_dims; + // Convert the permute dims to WHCN dimension order, which is the standard in + // our compute shaders. The following transformations are applied. + // 1. Change dimension index values from NCHW order valueto WHCN order value + // 2. Extend the permute array to kTensorDimLimit + { + IntListPtr permute_dims_ptr = graph.get_int_list(permute_dims); + whcn_permute_dims.initialize(*permute_dims_ptr); + } + + std::string kernel_name = "permute"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + vkapi::ParamsBindList param_buffers = { + graph.buffer_meta_ubo(out), + graph.buffer_meta_ubo(in), + graph.create_params_buffer(whcn_permute_dims), + }; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_global_wg_size, + default_pick_local_wg_size, + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + // Parameter buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {permute_dims}, + // Resizing Logic + resize_permute_node)); +} + void permute(ComputeGraph& graph, const std::vector& args) { - return add_permute_node(graph, args[0], args[1], args[2]); + int idx = 0; + const ValueRef in = args.at(idx++); + const ValueRef permute_dims = args.at(idx++); + const ValueRef out = args.at(idx++); + + if (graph.is_buffer_storage(args[2])) { + return add_permute_buffer_node(graph, in, permute_dims, out); + } + return add_permute_node(graph, in, permute_dims, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp index e8afafa9a45..250fcdd5490 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Pool.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Pool.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -17,44 +18,48 @@ namespace vkcompute { -void check_pool2d_args(const api::vTensor& in, const api::vTensor& out) { - VK_CHECK_COND(check_packed_dim_is(in, WHCN::kChannelsDim)); - VK_CHECK_COND(check_packed_dim_is(out, WHCN::kChannelsDim)); +void check_pool2d_args( + ComputeGraph& graph, + const ValueRef in, + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == WHCN::kChannelsDim); + VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } void resize_pool2d_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - bool is_max_pool2d = extra_args[3] != kDummyValueRef; + bool is_max_pool2d = extra_args.at(3) != kDummyValueRef; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); - size_t ndim = self->sizes().size(); + const std::vector self_sizes = graph->sizes_of(self); + size_t ndim = self_sizes.size(); std::vector new_out_sizes(ndim); // Batch, Channel if (ndim == 4) { - new_out_sizes.at(ndim - 4) = self->sizes().at(ndim - 4); + new_out_sizes.at(ndim - 4) = self_sizes.at(ndim - 4); } - new_out_sizes.at(ndim - 3) = self->sizes().at(ndim - 3); + new_out_sizes.at(ndim - 3) = self_sizes.at(ndim - 3); // Height, Width const auto& new_out_sizes_hw = calc_out_sizes_hw( *graph, - self->sizes(), - extra_args[0], + self_sizes, + extra_args.at(0), /*kernel_size_only = */ true, - {extra_args[1], extra_args[2], extra_args[3], extra_args[4]}); + {extra_args.at(1), extra_args.at(2), extra_args.at(3), extra_args.at(4)}); new_out_sizes.at(ndim - 2) = new_out_sizes_hw.at(0); new_out_sizes.at(ndim - 1) = new_out_sizes_hw.at(1); - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); if (is_max_pool2d) { - vTensorPtr indices = graph->get_tensor(args[0].refs[1]); - indices->virtual_resize(new_out_sizes); + const ValueRef indices = args.at(0).refs.at(1); + graph->virtual_resize(indices, new_out_sizes); } } @@ -71,18 +76,24 @@ void add_max_pool2d_node( const ValueRef dilation, const ValueRef ceil_mode, const ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - - const auto out_val = graph.get_value_list(out); - vTensorPtr t_out = graph.get_tensor(out_val->at(0)); - - check_pool2d_args(*t_in, *t_out); + ValueRef out_tensor = out; + // Placeholder tensor to fill binding slot for indices tensor in case we are + // computing max_pool2d instead of max_pool2d_with_indices. + TmpTensor tmp_indices_tensor = + TmpTensor(&graph, {}, graph.dtype_of(in), graph.storage_type_of(in)); + ValueRef indices_tensor = tmp_indices_tensor.vref; + int32_t write_indices = 0; + if (graph.val_is_value_list(out)) { + const auto out_val = graph.get_value_list(out); + out_tensor = out_val->at(0); + indices_tensor = out_val->at(1); + write_indices = 1; + } - utils::uvec3 global_size = t_out->logical_limits(); - utils::uvec3 local_size = adaptive_work_group_size(global_size); + check_pool2d_args(graph, in, out_tensor); std::string kernel_name("max_pool2d"); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out_tensor)); Kernel2dParams kernel_params = create_kernel2d_params( graph, @@ -92,23 +103,23 @@ void add_max_pool2d_node( padding, dilation); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - local_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs - {{{out_val->at(0), out_val->at(1)}, vkapi::kWrite}, {in, vkapi::kRead}}, + {{{out_tensor, indices_tensor}, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers { - t_out->logical_limits_ubo(), - t_in->sizes_ubo(), + graph.logical_limits_ubo(out_tensor), + graph.sizes_ubo(in), graph.create_params_buffer(kernel_params), }, // Push Constants {}, // Specialization Constants - {}, + {write_indices}, // Resize Args {kernel_size, stride, padding, dilation, ceil_mode}, // Resizing Logic @@ -150,16 +161,10 @@ void add_avg_pool2d_node( const ValueRef count_include_pad, const ValueRef divisor_override, const ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - - check_pool2d_args(*t_in, *t_out); - - utils::uvec3 global_size = t_out->logical_limits(); - utils::uvec3 local_size = adaptive_work_group_size(global_size); + check_pool2d_args(graph, in, out); std::string kernel_name("avg_pool2d"); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); Kernel2dParams kernel_params = create_kernel2d_params(graph, kernel_size, stride, padding); @@ -167,16 +172,16 @@ void add_avg_pool2d_node( DivisorParams divisor_params = create_divisor_params(graph, divisor_override, count_include_pad); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_size, - local_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers - {t_out->logical_limits_ubo(), - t_in->sizes_ubo(), + {graph.logical_limits_ubo(out), + graph.sizes_ubo(in), graph.create_params_buffer(kernel_params), graph.create_params_buffer(divisor_params)}, // Push Constants @@ -209,6 +214,7 @@ void avg_pool2d(ComputeGraph& graph, const std::vector& args) { REGISTER_OPERATORS { VK_REGISTER_OP(aten.avg_pool2d.default, avg_pool2d); VK_REGISTER_OP(aten.max_pool2d_with_indices.default, max_pool2d); + VK_REGISTER_OP(aten.max_pool2d.default, max_pool2d); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Quantize.cpp b/backends/vulkan/runtime/graph/ops/impl/Quantize.cpp index 92719505a0f..88f77261f4f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Quantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Quantize.cpp @@ -23,10 +23,11 @@ void resize_quantize_node( const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - out->virtual_resize(in->sizes()); + const std::vector in_sizes = graph->sizes_of(in); + graph->virtual_resize(out, in_sizes); } utils::uvec3 quantize_per_channel_local_wg_size( @@ -107,9 +108,35 @@ void add_quantize_per_tensor_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(output)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } vkapi::ParamsBindList param_ubos; std::vector push_constants; @@ -169,9 +196,35 @@ void add_quantize_per_token_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(output)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } int num_tokens = static_cast(graph.sizes_of(scale)[0]); @@ -242,10 +295,37 @@ void add_quantize_per_channel_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); int axis_val = static_cast(graph.get_int(axis)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(output)); + + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } // Normalize axis and convert from NCHW to WHCN using utility functions const auto input_sizes = graph.sizes_of(input); @@ -335,9 +415,35 @@ void add_quantize_block_wise_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(input)); add_dtype_suffix(kernel_name, graph.dtype_of(output)); + add_dtype_suffix(kernel_name, graph.dtype_of(scale)); + add_dtype_suffix(kernel_name, graph.dtype_of(zero_point)); + + // Handle optional quant_min and quant_max parameters independently + auto bounds = get_dtype_bounds(graph.dtype_of(output)); - int quant_min_val = static_cast(graph.get_int(quant_min)); - int quant_max_val = static_cast(graph.get_int(quant_max)); + int quant_min_val, quant_max_val; + + // Handle quant_min + if (graph.val_is_none(quant_min)) { + quant_min_val = bounds.first; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_min), + "quant_min must be an integer, got type: ", + graph.get_val_type(quant_min)); + quant_min_val = static_cast(graph.get_int(quant_min)); + } + + // Handle quant_max + if (graph.val_is_none(quant_max)) { + quant_max_val = bounds.second; + } else { + VK_CHECK_COND( + graph.val_is_int(quant_max), + "quant_max must be an integer, got type: ", + graph.get_val_type(quant_max)); + quant_max_val = static_cast(graph.get_int(quant_max)); + } const auto input_sizes = graph.sizes_of(input); const auto block_size_list = graph.get_int_list(block_size); @@ -426,6 +532,8 @@ void quantize_per_tensor_impl( // Check tensor types VK_CHECK_COND(graph.val_is_tensor(input)); + VK_CHECK_COND(graph.val_is_tensor(scale)); + VK_CHECK_COND(graph.val_is_tensor(zero_point)); VK_CHECK_COND(graph.val_is_tensor(output)); // Verify input is a floating point type @@ -434,6 +542,18 @@ void quantize_per_tensor_impl( graph.dtype_of(input) == vkapi::kFloat || graph.dtype_of(input) == vkapi::kHalf); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + add_quantize_per_tensor_node( graph, input, scale, zero_point, quant_min, quant_max, output); } @@ -465,6 +585,18 @@ void quantize_per_token_impl( graph.dtype_of(input) == vkapi::kFloat || graph.dtype_of(input) == vkapi::kHalf); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); @@ -538,6 +670,18 @@ void quantize_per_channel_impl( graph.dtype_of(input) == vkapi::kFloat || graph.dtype_of(input) == vkapi::kHalf); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); @@ -616,6 +760,18 @@ void quantize_affine_impl( graph.dtype_of(input) == vkapi::kFloat || graph.dtype_of(input) == vkapi::kHalf); + // Get scale and zero point dtypes + vkapi::ScalarType scale_dtype = graph.dtype_of(scale); + vkapi::ScalarType zero_point_dtype = graph.dtype_of(zero_point); + + // Verify supported types for scale (fp32 only for now) + VK_CHECK_COND(scale_dtype == vkapi::kFloat); + + // Verify supported types for zero point (int32, int8, fp32) + VK_CHECK_COND( + zero_point_dtype == vkapi::kInt || zero_point_dtype == vkapi::kChar || + zero_point_dtype == vkapi::kFloat); + // Check that scale and zero_point have buffer storage and width packing VK_CHECK_COND(graph.is_buffer_storage(scale)); VK_CHECK_COND(graph.packed_dim_of(scale) == WHCN::kWidthDim); diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp new file mode 100644 index 00000000000..51f8138485e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedConvolution.cpp @@ -0,0 +1,695 @@ +/* + * 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 vkcompute { + +// +// Utility functions +// + +struct Conv2DParams { + utils::ivec2 kernel_size; + utils::ivec2 stride; + utils::ivec2 padding; + utils::ivec2 dilation; + int32_t groups; + int32_t out_channels_per_group; + int32_t in_channels_per_group; + int32_t logical_K_per_group; + int32_t K_per_group; + int32_t K4_per_group; + int32_t logical_K; + int32_t K; + int32_t K4; +}; + +Conv2DParams create_conv2d_params( + ComputeGraph& graph, + const ValueRef& conv_input, + const ValueRef& conv_output, + const ValueRef& kernel_size, + const ValueRef& stride, + const ValueRef& padding, + const ValueRef& dilation, + const ValueRef& groups) { + const auto kernel_size_list = graph.get_int_list(kernel_size); + const auto stride_list = graph.get_int_list(stride); + const auto padding_list = graph.get_int_list(padding); + const auto dilation_list = graph.get_int_list(dilation); + const int32_t groups_val = graph.get_int(groups); + + // Pre-compute input and output channels per group + + std::vector out_sizes = graph.sizes_of(conv_output); + const int32_t out_channels = utils::val_at(-3, out_sizes); + const int32_t out_channels_per_group = out_channels / groups_val; + + std::vector in_sizes = graph.sizes_of(conv_input); + const int32_t in_channels = utils::val_at(-3, in_sizes); + const int32_t in_channels_per_group = in_channels / groups_val; + + // Pre-compute the number of elements along the K dimension per group. This + // quantity is aligned to the next multiple of 4 to ensure data loads are + // aligned to texel boundaries. + + const int32_t logical_K_per_group = + kernel_size_list->at(0) * kernel_size_list->at(1) * in_channels_per_group; + const int32_t K_per_group = utils::align_up_4(logical_K_per_group); + const int32_t K4_per_group = K_per_group / 4; + + // Pre-compute the "theoretical" size of the K dim of the input im2col matrix, + // which represents the flattened convolution window used to compute an output + // element. This is used for bounds checking. + + const int32_t logical_K = + kernel_size_list->at(0) * kernel_size_list->at(1) * in_channels; + + const int32_t K = K_per_group * groups_val; + // Used for texel stride calculations + const int32_t K4 = K / 4; + + return Conv2DParams{ + // Swap the order from HW to WH + utils::make_ivec2({kernel_size_list->at(1), kernel_size_list->at(0)}), + utils::make_ivec2({stride_list->at(1), stride_list->at(0)}), + utils::make_ivec2({padding_list->at(1), padding_list->at(0)}), + utils::make_ivec2({dilation_list->at(1), dilation_list->at(0)}), + groups_val, + out_channels_per_group, + in_channels_per_group, + logical_K_per_group, + K_per_group, + K4_per_group, + logical_K, + K, + K4, + }; +} + +std::vector calculate_input_im2col_sizes( + ComputeGraph* graph, + const ValueRef& input, + const ValueRef& output, + const ValueRef& kernel_size, + const ValueRef& groups) { + std::vector in_sizes = graph->sizes_of(input); + const int64_t in_channels = utils::val_at(-3, in_sizes); + + std::vector out_sizes = graph->sizes_of(output); + const int64_t batches = utils::val_at(-4, out_sizes); + const int64_t out_height = utils::val_at(-2, out_sizes); + const int64_t out_width = utils::val_at(-1, out_sizes); + + // Represents the number of channel groups + const int64_t groups_val = graph->extract_scalar(groups); + // No need to div_up because in_channels % groups_val = 0 + const int64_t in_channels_per_group = in_channels / groups_val; + + const auto kernel_size_list = graph->get_int_list(kernel_size); + + // Align to the next multiple of 4 to ensure that data loads align nicely with + // texel boundaries. We want to ensure that the first data element of each + // group is at the start of its texel. + const int64_t flattened_kernel_len = utils::align_up_4( + in_channels_per_group * kernel_size_list->at(0) * + kernel_size_list->at(1)); + + // K -> flattened convolution window (adjusted) + const int64_t K = flattened_kernel_len * groups_val; + // M -> number of elements in 2D output plane. This is aligned to the next + // multiple of 4 since the im2col shader operates on 4x4 blocks. + const int64_t M = utils::align_up_4(out_height * out_width * batches); + + return {M, K}; +} + +std::vector calculate_output_im2col_sizes( + ComputeGraph* graph, + const ValueRef& output) { + std::vector out_sizes = graph->sizes_of(output); + const int64_t batches = utils::val_at(-4, out_sizes); + const int64_t out_channels = utils::val_at(-3, out_sizes); + const int64_t out_height = utils::val_at(-2, out_sizes); + const int64_t out_width = utils::val_at(-1, out_sizes); + + // N -> output channels + const int64_t N = out_channels; + // M -> number of elements in 2D output plane + const int64_t M = out_height * out_width * batches; + + return {M, N}; +} + +// +// Shader dispatch utilities +// + +utils::uvec3 im2col_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef input_image = args.at(1).refs.at(0); + const ValueRef output_image = resize_args.at(0); + const ValueRef kernel_size = resize_args.at(1); + const ValueRef groups = resize_args.at(2); + + std::vector im2col_sizes = calculate_input_im2col_sizes( + graph, input_image, output_image, kernel_size, groups); + const uint32_t K = utils::safe_downcast(im2col_sizes[1]); + const uint32_t M = utils::safe_downcast(im2col_sizes[0]); + + // 1 output tile is 4x4 elements + const uint32_t K4 = utils::div_up(K, 4u); + const uint32_t M4 = utils::div_up(M, 4u); + + return {K4, M4, 1}; +} + +utils::uvec3 col2im_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef output = args.at(0).refs.at(0); + + std::vector im2col_sizes = + calculate_output_im2col_sizes(graph, output); + const uint32_t N = utils::safe_downcast(im2col_sizes[1]); + const uint32_t M = utils::safe_downcast(im2col_sizes[0]); + + // 1 output tile is 4x4 elements + const uint32_t N4 = utils::div_up(N, 4u); + const uint32_t M4 = utils::div_up(M, 4u); + + return {N4, M4, 1}; +} + +// +// Dispatch nodes +// + +void add_input_im2col_node( + ComputeGraph& graph, + const ValueRef input_image, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef output_image, + const ValueRef input_im2col) { + Conv2DParams conv_params = create_conv2d_params( + graph, + input_image, + output_image, + kernel_size, + stride, + padding, + dilation, + groups); + + std::string kernel_name = "im2col"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_im2col)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_image)); + add_dtype_suffix(kernel_name, graph.dtype_of(output_image)); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(input_im2col), + graph.sizes_ubo(input_image), + graph.sizes_ubo(output_image), + graph.create_params_buffer(conv_params)}; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + im2col_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{input_im2col, vkapi::kWrite}, {input_image, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize args + {output_image, kernel_size, groups}, + // Resizing Logic + nullptr)); +} + +void add_quantize_and_pack_im2col_node( + ComputeGraph& graph, + const ValueRef input_image, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef output_image, + const ValueRef input_int_im2col) { + Conv2DParams conv_params = create_conv2d_params( + graph, + input_image, + output_image, + kernel_size, + stride, + padding, + dilation, + groups); + + float inv_scale = 1.0f / graph.extract_scalar(input_scale); + int32_t zp = graph.extract_scalar(input_zp); + + // Get shader for quantized conv2d linear tiled + std::string kernel_name = "quantize_and_pack_im2col"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_int_im2col)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_image)); + add_dtype_suffix(kernel_name, graph.dtype_of(output_image)); + + vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(input_int_im2col), + graph.sizes_ubo(input_image), + graph.sizes_ubo(output_image), + graph.create_params_buffer(conv_params)}; + + std::vector push_constants = { + PushConstantDataInfo(&inv_scale, sizeof(inv_scale)), + PushConstantDataInfo(&zp, sizeof(zp)), + }; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + im2col_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{input_int_im2col, vkapi::kWrite}, {input_image, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + {}, + // Resize args + {output_image, kernel_size, groups}, + // Resizing Logic + nullptr)); +} + +void add_conv2d_q8csw_linear_node( + ComputeGraph& graph, + const ValueRef input_im2col, + const ValueRef input_image, + const ValueRef packed_weight, + const ValueRef packed_weight_scales, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef output_image) { + Conv2DParams conv_params = create_conv2d_params( + graph, + input_image, + output_image, + kernel_size, + stride, + padding, + dilation, + groups); + + // One limitation of the current implementation is that for grouped convs, + // the number of output_image channels per group must be a multiple of 4. One + // loaded 4x4 weight tile must all belong to the same group. + if (conv_params.groups > 1) { + VK_CHECK_COND(conv_params.out_channels_per_group % 4 == 0); + } + + std::string kernel_name = "conv2d_q8csw_linear_tiled"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(output_image)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_im2col)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph.dtype_of(output_image)); + vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output_image), + graph.sizes_ubo(input_image), + graph.create_params_buffer(conv_params)}; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + col2im_global_wg_size, + quantized_linear_local_wg_size, + // Inputs and Outputs + {{output_image, vkapi::kWrite}, + {{input_im2col, packed_weight, packed_weight_scales, packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {apply_bias}, + // Resize args + {}, + // Resizing Logic + nullptr)); +} + +void add_conv2d_q8ta_q8csw_linear_node( + ComputeGraph& graph, + const ValueRef input_int_im2col, + const ValueRef input_image, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef output_image) { + Conv2DParams conv_params = create_conv2d_params( + graph, + input_image, + output_image, + kernel_size, + stride, + padding, + dilation, + groups); + + // One limitation of the current implementation is that for grouped convs, + // the number of output channels per group must be a multiple of 4. One loaded + // 4x4 weight tile must all belong to the same group. + if (conv_params.groups > 1) { + VK_CHECK_COND(conv_params.out_channels_per_group % 4 == 0); + } + + float scale = graph.extract_scalar(input_scale); + int32_t zp = graph.extract_scalar(input_zp); + + std::string kernel_name = "conv2d_q8ta_q8csw_linear_tiled"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(output_image)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(input_int_im2col)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph.dtype_of(output_image)); + vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output_image), + graph.sizes_ubo(input_image), + graph.create_params_buffer(conv_params)}; + + std::vector push_constants = { + PushConstantDataInfo(&scale, sizeof(scale)), + PushConstantDataInfo(&zp, sizeof(zp)), + }; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + col2im_global_wg_size, + quantized_linear_local_wg_size, + // Inputs and Outputs + {{output_image, vkapi::kWrite}, + {{input_int_im2col, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + {apply_bias}, + // Resize args + {weight_data}, + // Resizing Logic + nullptr)); +} + +// +// High level operator impl +// + +void quantized_conv2d_impl( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef input_image, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef weight_data, + const ValueRef weight_sums_data, + const ValueRef weight_scales_data, + const ValueRef bias_data, + const ValueRef kernel_size, + const ValueRef stride, + const ValueRef padding, + const ValueRef dilation, + const ValueRef groups, + const ValueRef output_image) { + VK_CHECK_COND(weight_quant_config.granularity == kPerChannel); + VK_CHECK_COND(weight_quant_config.nbits == 8); + VK_CHECK_COND(weight_quant_config.is_symmetric); + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + + // Create a dummy tensor to fill the binding slot of the bias tensor if it is + // not provided. This helps simplify dispatch logic and makes it so that + // fewer shader variants need to be generated. + TmpTensor dummy_bias( + &graph, + {}, + graph.dtype_of(output_image), + utils::kBuffer, + utils::kWidthPacked); + + ValueRef packed_bias = dummy_bias.vref; + if (!graph.val_is_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + } + + std::vector input_im2col_sizes = calculate_input_im2col_sizes( + &graph, input_image, output_image, kernel_size, groups); + + // Use weight only quantized conv2d if at least one is true: + // 1. Device does not support int8 dot product + // 2. Input is not quantized + if (!graph.can_use_int8_dot_product() || + input_quant_config.granularity == kNoQuantization) { + TmpTensor input_im2col( + &graph, + input_im2col_sizes, + vkapi::kFloat, + utils::kBuffer, + utils::kWidthPacked); + + add_input_im2col_node( + graph, + input_image, + kernel_size, + stride, + padding, + dilation, + groups, + output_image, + input_im2col); + + add_conv2d_q8csw_linear_node( + graph, + input_im2col, + input_image, + packed_weight, + packed_weight_scales, + bias_data, + packed_bias, + kernel_size, + stride, + padding, + dilation, + groups, + output_image); + return; + } else { + // Otherwise, use activation + weight quantized conv2d + VK_CHECK_COND(input_quant_config.granularity == kPerTensor); + VK_CHECK_COND(weight_quant_config.nbits == 8); + VK_CHECK_COND(!weight_quant_config.is_dynamic); + + ValueRef packed_weight_sums = prepack_standard( + graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); + + // Allocate quantized + packed im2col matrix for input + const int64_t num_blocks_M = utils::div_up_4(input_im2col_sizes.at(0)); + const int64_t num_blocks_K = utils::div_up_4(input_im2col_sizes.at(1)); + + TmpTensor input_int_im2col( + &graph, + {num_blocks_M, num_blocks_K * 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + add_quantize_and_pack_im2col_node( + graph, + input_image, + input_scale, + input_zp, + kernel_size, + stride, + padding, + dilation, + groups, + output_image, + input_int_im2col); + + add_conv2d_q8ta_q8csw_linear_node( + graph, + input_int_im2col, + input_image, + input_scale, + input_zp, + weight_data, + packed_weight, + packed_weight_sums, + packed_weight_scales, + bias_data, + packed_bias, + kernel_size, + stride, + padding, + dilation, + groups, + output_image); + return; + }; +} + +void conv2d_q8ta_q8csw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef input_image = args.at(idx++); + const ValueRef input_scale = args.at(idx++); + const ValueRef input_zp = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_sums_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef kernel_size = args.at(idx++); + const ValueRef stride = args.at(idx++); + const ValueRef padding = args.at(idx++); + const ValueRef dilation = args.at(idx++); + const ValueRef groups = args.at(idx++); + const ValueRef output_image = args.at(idx++); + + const int64_t K = graph.size_at(-1, weight_data); + + QuantizationConfig input_quant_config(8, kPerTensor, {}, false); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_conv2d_impl( + graph, + input_quant_config, + weight_quant_config, + input_image, + input_scale, + input_zp, + weight_data, + weight_sums_data, + weight_scales_data, + bias_data, + kernel_size, + stride, + padding, + dilation, + groups, + output_image); +} + +void conv2d_q8csw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef input_image = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef kernel_size = args.at(idx++); + const ValueRef stride = args.at(idx++); + const ValueRef padding = args.at(idx++); + const ValueRef dilation = args.at(idx++); + const ValueRef groups = args.at(idx++); + const ValueRef output_image = args.at(idx++); + + const int64_t K = graph.size_at(-1, weight_data); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_conv2d_impl( + graph, + input_quant_config, + weight_quant_config, + input_image, + kDummyValueRef, // input scale + kDummyValueRef, // input zero point + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + bias_data, + kernel_size, + stride, + padding, + dilation, + groups, + output_image); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.conv2d_q8ta_q8csw.default, conv2d_q8ta_q8csw); + VK_REGISTER_OP(et_vk.conv2d_q8csw.default, conv2d_q8csw); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp new file mode 100644 index 00000000000..7fbfcee5cb1 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -0,0 +1,1049 @@ +/* + * 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 vkcompute { + +// +// Shader dispatch utilities +// + +bool is_gemv(ComputeGraph* graph, const ValueRef& fp_input) { + return graph->size_at(-2, fp_input) == 1; +} + +void resize_linear_qw_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + (void)extra_args; + + ValueRef output = args.at(0).refs.at(0); + ValueRef fp_input = args.at(1).refs.at(0); + ValueRef weight_data = extra_args.at(1); + + std::vector mat1_sizes = graph->sizes_of(fp_input); + std::vector mat2_sizes = graph->sizes_of(weight_data); + + const int64_t out_cols = utils::val_at(-2, mat1_sizes); + const int64_t out_rows = utils::val_at(-2, mat2_sizes); + + std::vector new_out_sizes(3); + if (mat1_sizes.size() == 2) { + new_out_sizes.resize(2); + new_out_sizes.at(0) = out_cols; + new_out_sizes.at(1) = out_rows; + } else { + new_out_sizes.at(0) = mat1_sizes.at(0); + new_out_sizes.at(1) = out_cols; + new_out_sizes.at(2) = out_rows; + } + + graph->virtual_resize(output, new_out_sizes); +} + +utils::uvec3 quantized_linear_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + + std::vector out_sizes = graph->sizes_of(out); + // width + const uint32_t N = utils::val_at(-1, out_sizes); + // height + const uint32_t M = utils::val_at(-2, out_sizes); + + uint32_t N_per_tile = 4; + uint32_t M_per_tile = 4; + + // For 4-bit weights, each output tile contains 8 columns + if (shader.kernel_name.find("q4") != std::string::npos) { + N_per_tile = 8; + } + if (shader.kernel_name.find("coop") != std::string::npos) { + M_per_tile = 1; + } + + const uint32_t num_N_tiles = utils::div_up(N, N_per_tile); + const uint32_t num_M_tiles = utils::div_up(M, M_per_tile); + + // Otherwise, each output tile contains 4 columns and 4 rows + return {num_N_tiles, num_M_tiles, 1}; +} + +utils::uvec3 quantized_linear_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + const bool use_coop_algorithm = + shader.kernel_name.find("_coop") != std::string::npos; + + if (use_coop_algorithm) { + return {1, 1, 64}; + } else { + return pick_hw_square_wg_size( + graph, shader, global_workgroup_size, args, resize_args); + } +} + +std::tuple get_quantized_input_num_blocks( + ComputeGraph& graph, + const ValueRef input) { + std::vector input_sizes = graph.sizes_of(input); + const int64_t ndim = graph.dim_of(input); + + const int64_t M = input_sizes.at(ndim - 2); + const int64_t K = input_sizes.at(ndim - 1); + + const int64_t num_blocks_M = utils::div_up(M, int64_t(4)); + const int64_t num_blocks_K = utils::div_up(K, int64_t(4)); + + return std::make_tuple(num_blocks_M, num_blocks_K); +} + +utils::uvec3 quant_pack_input_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef input = args.at(1).refs.at(0); + int64_t num_blocks_M, num_blocks_K; + std::tie(num_blocks_M, num_blocks_K) = + get_quantized_input_num_blocks(*graph, input); + + return { + utils::safe_downcast(num_blocks_K), + utils::safe_downcast(num_blocks_M), + 1u}; +} + +vkapi::ShaderInfo pick_quantize_and_pack_input_with_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 group_size = resize_args.at(0); + + const int64_t group_size_val = graph->extract_scalar(group_size); + + std::string shader_name = "quantize_and_pack_linear_input_with_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)); + + return VK_KERNEL_FROM_STR(shader_name); +} + +utils::uvec3 pick_quantize_and_pack_input_with_sums_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef fp_input = args.at(1).refs.at(0); + // For gemv cases, skip the quantize and pack input step in favor of computing + // the quantized linear as a weight only quantized linear operation. The + // rationale for this is that gemv is a memory bound operation and may not + // necessarily benefit from quantizing the input and computing with integer + // accumulation. + if (is_gemv(graph, fp_input)) { + return {0u, 0u, 0u}; + } + + const ValueRef group_size = resize_args.at(0); + int64_t num_blocks_M, num_blocks_K; + std::tie(num_blocks_M, num_blocks_K) = + get_quantized_input_num_blocks(*graph, fp_input); + + const int64_t group_size_val = graph->extract_scalar(group_size); + const int64_t blocks_per_group = group_size_val / 4; + + const int64_t num_groups = num_blocks_K / blocks_per_group; + + return { + utils::safe_downcast(num_groups), + utils::safe_downcast(num_blocks_M), + 1u}; +} + +utils::uvec3 pick_quantize_and_pack_input_with_sums_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + + const ValueRef fp_input = args.at(1).refs.at(0); + // For gemv, skip the quantize input step since the quantized linear is + // computed as a weight only quantized linear operation. + if (is_gemv(graph, fp_input)) { + return {1u, 1u, 1u}; + } + + uint32_t groups_per_wg = 2u; + uint32_t workers_per_group = 32u; + + if (shader.kernel_name.find("o4w16") != std::string::npos) { + groups_per_wg = 4u; + workers_per_group = 16u; + } + + return {groups_per_wg, 1u, workers_per_group}; +} + +vkapi::ShaderInfo pick_linear_qw_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + + const ValueRef output = args.at(0).refs.at(0); + const ValueRef fp_input = args.at(1).refs.at(0); + const ValueRef packed_int_weight = args.at(1).refs.at(1); + + const bool weight_is_4bit = resize_args.at(0) != kDummyValueRef; + const bool is_gemv_case = is_gemv(graph, fp_input); + + std::string kernel_name = "linear_"; + if (weight_is_4bit) { + kernel_name += "q4gsw"; + } else { + kernel_name += "q8csw"; + } + + if (weight_is_4bit && is_gemv_case) { + kernel_name += "_coop"; + } else { + kernel_name += "_tiled"; + } + add_storage_type_suffix(kernel_name, graph->storage_type_of(output)); + add_storage_type_suffix( + kernel_name, graph->storage_type_of(packed_int_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(output)); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +vkapi::ShaderInfo pick_linear_dqa_qw_shader( + 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 fp_input = args.at(1).refs.at(0); + const ValueRef int_input = args.at(1).refs.at(1); + (void)int_input; + const ValueRef int_weight = args.at(1).refs.at(5); + + const bool weight_is_4bit = resize_args.at(0) != kDummyValueRef; + const bool is_gemv_case = is_gemv(graph, fp_input); + + std::string kernel_name = "linear_"; + if (weight_is_4bit) { + kernel_name += "dq8ca_q4gsw"; + } else { + kernel_name += "dq8ca_q8csw"; + } + + if (weight_is_4bit && is_gemv_case) { + kernel_name += "_coop"; + } else { + kernel_name += "_tiled"; + } + add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); + add_storage_type_suffix(kernel_name, graph->storage_type_of(int_weight)); + add_dtype_suffix(kernel_name, graph->dtype_of(out)); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +// +// Prepacking nodes +// + +ValueRef prepack_quantized_linear_weight( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef qmat2_data) { + VK_CHECK_COND( + weight_quant_config.nbits == 8 || weight_quant_config.nbits == 4); + + std::vector qmat2_orig_sizes = graph.sizes_of(qmat2_data); + const int64_t ndim = graph.dim_of(qmat2_data); + + int64_t qmat2_width = qmat2_orig_sizes.at(ndim - 1); + int64_t qmat2_height = qmat2_orig_sizes.at(ndim - 2); + + int64_t K; + int64_t N; + if (weight_quant_config.nbits == 4) { + // For 4-bit quantization, weight source data has shape [N, K/2]. Each byte + // contains 2 * 4-bit values. + K = qmat2_width * 2; + N = qmat2_height; + } else { + // For 8-bit quantization, the weight source data has shape [N, K] + K = qmat2_width; + N = qmat2_height; + } + + // Sanity check that assumptions are correct. Data loads along the innermost + // dimension must be well aligned along texel boundaries. + if (weight_quant_config.nbits == 4) { + VK_CHECK_COND(K % 8 == 0); + } else { + VK_CHECK_COND(K % 4 == 0); + } + + // The packing format packs the weight tensor into blocks of 4 columns (K) and + // 4 rows (N) + int64_t N_per_block = 4; + int64_t K_per_block = 4; + + // For 4 bit, quantization, the amount of information contained in one block + // can be doubled. Each block will contain data for 8 rows (N) instead of the + // usual 4. + if (weight_quant_config.nbits == 4) { + N_per_block = 8; + } + + // To figure out the size of the output tensor, determine the number of blocks + // along each dimension. + const int64_t num_blocks_K = utils::div_up(K, K_per_block); + const int64_t num_blocks_N = utils::div_up(N, N_per_block); + + // The blocks are arranged in a transposed manner, such that the transposed + // weight block is indexed like packed_weights[k4][n4] - this is to allow for + // optimal memory coalescing when computing GEMM. + int64_t output_height = num_blocks_K; + // The base dtype of the packed tensor is int32 (each int32 contains 4x 8bit + // values) and each block is represented as a ivec4. Therefore the width dim + // of the packed tensor is multiplied by 4. + int64_t output_width = num_blocks_N * 4; + + // For 4 bit quantization, The blocks are arranged without the transposition, + // such that a weight block is accessed like packed_weights[n8][k4]. This is + // an optimization targeted for LLMs, which need to compute GEMV as well as + // GEMM. This memory layout provides better performance for the co-operative + // algorithm used to compute GEMV, at the cost of slightly reducing GEMM + // performance. + if (weight_quant_config.nbits == 4) { + output_height = num_blocks_N; + output_width = num_blocks_K * 4; + } + + // Store the original sizes of the weight data to pass to the shader + utils::ivec2 orig_sizes = { + utils::safe_downcast(K), utils::safe_downcast(N)}; + + std::vector qmat2_sizes{output_height, output_width}; + + utils::StorageType storage_type = utils::kTexture2D; + 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; + } + + ValueRef qmat2 = graph.add_tensor( + qmat2_sizes, vkcompute::vkapi::kInt, storage_type, utils::kWidthPacked); + + utils::uvec3 global_wg_size; + if (weight_quant_config.nbits == 4) { + // For 4-bit quantization, each thread writes out two adjacent blocks + global_wg_size = { + utils::safe_downcast(utils::div_up(num_blocks_K, int64_t(2))), + utils::safe_downcast(num_blocks_N), + 1u}; + } else { + global_wg_size = { + utils::safe_downcast(num_blocks_N), + utils::safe_downcast(num_blocks_K), + 1u}; + } + + std::string kernel_name = weight_quant_config.nbits == 4 + ? "pack_q4_linear_weight" + : "pack_q8_linear_weight"; + add_storage_type_suffix(kernel_name, storage_type); + + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + global_wg_size, + graph.create_local_wg_size(global_wg_size), + // Inputs and Outputs + qmat2_data, + qmat2, + // UBOs + {}, + // Specialization Constants + {}, + // Push Constants + {graph.sizes_pc_of(qmat2), + PushConstantDataInfo(&orig_sizes, sizeof(utils::ivec2))})); + + return qmat2; +} + +// +// Dispatch nodes +// + +/* + * Shader dispatch for linear with quantized weight but fp activations. + */ +DynamicDispatchNode make_linear_qw_node( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_scales, + const ValueRef packed_weight_zeros, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output) { + // Only certain quantization types supported at the moment + VK_CHECK_COND( + weight_quant_config.granularity == kPerChannel || + weight_quant_config.granularity == kPerGroup); + VK_CHECK_COND(weight_quant_config.is_symmetric); + VK_CHECK_COND( + weight_quant_config.nbits == 8 || weight_quant_config.nbits == 4); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + int32_t K4_per_group = 0; + if (weight_quant_config.nbits == 4) { + int32_t group_size_val = graph.extract_scalar(group_size); + K4_per_group = utils::div_up(group_size_val, int32_t(4)); + } + + const ValueRef is_4bit_flag = + weight_quant_config.nbits == 4 ? group_size : kDummyValueRef; + + return DynamicDispatchNode( + graph, + pick_linear_qw_shader, + quantized_linear_global_wg_size, + quantized_linear_local_wg_size, + // Inputs and Outputs + {{output, vkapi::kWrite}, + {{fp_input, packed_weight, packed_weight_scales, packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {apply_bias, K4_per_group}, + // Resize args + {is_4bit_flag, weight_data}, + // Resizing Logic + resize_linear_qw_node); +} + +DynamicDispatchNode make_quantize_and_pack_linear_input_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const ValueRef fp_input, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef packed_int_input, + const ValueRef group_size) { + // Only certain quantization types supported at the moment + VK_CHECK_COND(input_quant_config.granularity == kPerTensor); + + int64_t num_blocks_M, num_blocks_K; + std::tie(num_blocks_M, num_blocks_K) = + get_quantized_input_num_blocks(graph, fp_input); + + float inv_scale = 1.0f / graph.extract_scalar(input_scale_data); + int32_t zp = graph.extract_scalar(input_zp_data); + + std::string shader_name = "quantize_and_pack_linear_input_per_tensor"; + 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)); + + vkapi::ParamsBindList param_buffers = {graph.sizes_ubo(fp_input)}; + + std::vector push_constants = { + PushConstantDataInfo(&inv_scale, sizeof(inv_scale)), + PushConstantDataInfo(&zp, sizeof(zp)), + }; + + return DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(shader_name), + quant_pack_input_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{packed_int_input, vkapi::kWrite}, {fp_input, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + {}, + // Resize args + {}); +} + +DynamicDispatchNode make_quantize_and_pack_linear_input_with_sums_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const ValueRef fp_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scales, + const ValueRef packed_input_zps, + const ValueRef packed_int_input, + const ValueRef group_size) { + // Only certain quantization types supported at the moment + VK_CHECK_COND(input_quant_config.granularity == kPerChannel); + + int64_t num_blocks_M, num_blocks_K; + std::tie(num_blocks_M, num_blocks_K) = + get_quantized_input_num_blocks(graph, fp_input); + + vkapi::ParamsBindList param_buffers = {graph.sizes_ubo(fp_input)}; + + 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)); + + return DynamicDispatchNode( + graph, + pick_quantize_and_pack_input_with_sums_shader, + pick_quantize_and_pack_input_with_sums_global_wg_size, + pick_quantize_and_pack_input_with_sums_local_wg_size, + // Inputs and Outputs + {{{packed_int_input, int_input_sums}, vkapi::kWrite}, + {{fp_input, packed_input_scales, packed_input_zps}, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {blocks_per_group}, + // Resize args + {group_size}); +} + +DynamicDispatchNode make_linear_qa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output) { + VK_CHECK_COND(input_quant_config.granularity == kPerTensor); + VK_CHECK_COND(input_quant_config.nbits == 8); + VK_CHECK_COND(weight_quant_config.granularity == kPerChannel); + VK_CHECK_COND(weight_quant_config.is_symmetric); + VK_CHECK_COND(weight_quant_config.nbits == 8); + + float scale = graph.extract_scalar(input_scale_data); + int32_t zp = graph.extract_scalar(input_zp_data); + + // Get shader for quantized linear + std::string kernel_name = "linear_q8ta_q8csw_tiled"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(output)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); + add_dtype_suffix(kernel_name, graph.dtype_of(output)); + vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output), graph.sizes_ubo(packed_int_input)}; + + std::vector push_constants = { + PushConstantDataInfo(&scale, sizeof(scale)), + PushConstantDataInfo(&zp, sizeof(zp)), + }; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + // Add the compute node + return DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + quantized_linear_global_wg_size, + quantized_linear_local_wg_size, + // Inputs and Outputs + {{output, vkapi::kWrite}, + {{packed_int_input, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + {apply_bias}, + // Resize args + {fp_input}, + // Resizing Logic + nullptr); +} + +DynamicDispatchNode make_linear_dqa_qw_node( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef packed_int_input, + const ValueRef int_input_sums, + const ValueRef packed_input_scale, + const ValueRef packed_input_zp, + const ValueRef input_scale_data, + const ValueRef input_zp_data, + const ValueRef weight_data, + const ValueRef packed_weight, + const ValueRef packed_weight_sums, + const ValueRef packed_weight_scales, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef packed_bias, + const ValueRef output) { + VK_CHECK_COND(input_quant_config.granularity == kPerChannel); + VK_CHECK_COND(input_quant_config.nbits == 8); + VK_CHECK_COND(input_quant_config.is_dynamic); + + VK_CHECK_COND(weight_quant_config.granularity == kPerGroup); + VK_CHECK_COND(weight_quant_config.is_symmetric); + VK_CHECK_COND(weight_quant_config.nbits == 4); + + vkapi::ParamsBindList param_buffers = { + graph.sizes_ubo(output), graph.sizes_ubo(fp_input)}; + + uint32_t apply_bias = 1; + if (graph.val_is_none(bias_data)) { + apply_bias = 0; + } + + int32_t K4_per_group = 0; + if (weight_quant_config.nbits == 4) { + int32_t group_size_val = graph.extract_scalar(group_size); + K4_per_group = utils::div_up(group_size_val, int32_t(4)); + } + + const ValueRef is_4bit_flag = + weight_quant_config.nbits == 4 ? group_size : kDummyValueRef; + + // Add the compute node + return DynamicDispatchNode( + graph, + pick_linear_dqa_qw_shader, + quantized_linear_global_wg_size, + quantized_linear_local_wg_size, + // Inputs and Outputs + {{output, vkapi::kWrite}, + {{fp_input, + packed_int_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + packed_bias}, + vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + {}, + // Specialization Constants + {apply_bias, K4_per_group}, + // Resize args + {is_4bit_flag, weight_data}, + // Resizing Logic + resize_linear_qw_node); +} + +// +// High level operator impl +// + +void quantized_linear_impl( + ComputeGraph& graph, + const QuantizationConfig& input_quant_config, + const QuantizationConfig& weight_quant_config, + const ValueRef fp_input, + const ValueRef input_scale, + const ValueRef input_zp, + const ValueRef weight_data, + const ValueRef weight_sums_data, + const ValueRef weight_scales_data, + const ValueRef weight_zeros_data, + const ValueRef group_size, + const ValueRef bias_data, + const ValueRef output) { + std::vector input_sizes = graph.sizes_of(fp_input); + std::vector weight_sizes = graph.sizes_of(weight_data); + + const int64_t K = utils::val_at(-1, input_sizes); + // K (input channels) must be a multiple of 4 to ensure that reading a group + // of 4 input channels from the input tensor will be aligned on a texel + // boundary. + VK_CHECK_COND(K % 4 == 0); + + // Prepack weight data + + const ValueRef packed_weight = + prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + const ValueRef packed_weight_scales = prepack_standard( + graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); + // Weight affine quant not supported at the moment + const ValueRef packed_weight_zeros = kDummyValueRef; + + // Prepack bias data + + // Create a dummy tensor to fill the binding slot of the bias tensor if it is + // not provided. This helps simplify dispatch logic and makes it so that + // fewer shdaer variants need to be generated. + TmpTensor dummy_bias( + &graph, {}, graph.dtype_of(output), utils::kBuffer, utils::kWidthPacked); + + ValueRef packed_bias = dummy_bias.vref; + if (graph.val_is_not_none(bias_data)) { + packed_bias = + prepack_standard(graph, bias_data, utils::kBuffer, utils::kWidthPacked); + } + + // Use weight only quantized linear if at least one is true: + // 1. Device does not support int8 dot product + // 2. Input is not quantized + if (!graph.can_use_int8_dot_product() || + input_quant_config.granularity == kNoQuantization) { + DynamicDispatchNode linear_qw_node(make_linear_qw_node( + graph, + weight_quant_config, + fp_input, + weight_data, + packed_weight, + packed_weight_scales, + packed_weight_zeros, + group_size, + bias_data, + packed_bias, + output)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode(linear_qw_node)); + return; + } + // Otherwise, use input and weight quantized linear computed with integer + // accumulation + + // Input scale/zero point only used for activation & weight quantized linear + ValueRef packed_input_scale = input_scale; + ValueRef packed_input_zp = input_zp; + if (graph.val_is_tref(input_scale)) { + VK_CHECK_COND(graph.val_is_tref(packed_input_zp)); + packed_input_scale = prepack_standard( + graph, input_scale, utils::kTexture3D, utils::kWidthPacked); + packed_input_zp = prepack_standard( + graph, input_zp, utils::kTexture3D, utils::kWidthPacked); + } + + // Pre-computed per quant group weight sums are needed for int accumulation, + // but not for weight only + const ValueRef packed_weight_sums = prepack_standard( + graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); + + // Allocate temporary tensor to store quantized and packed input + + int64_t num_blocks_M, num_blocks_K; + std::tie(num_blocks_M, num_blocks_K) = + get_quantized_input_num_blocks(graph, fp_input); + + const int64_t int_input_height = num_blocks_M; + const int64_t int_input_width = num_blocks_K * 4; + + TmpTensor packed_int_input( + &graph, + {int_input_height, int_input_width}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + + // Non dynamically quantized input case + if (!input_quant_config.is_dynamic) { + DynamicDispatchNode quantize_and_pack_linear_node( + make_quantize_and_pack_linear_input_node( + graph, + input_quant_config, + fp_input, + packed_input_scale, + packed_input_zp, + input_scale, + input_zp, + packed_int_input, + group_size)); + + graph.execute_nodes().emplace_back( + new DynamicDispatchNode(quantize_and_pack_linear_node)); + + DynamicDispatchNode linear_qa_qw_node(make_linear_qa_qw_node( + graph, + input_quant_config, + weight_quant_config, + fp_input, + packed_int_input, + packed_input_scale, + packed_input_zp, + input_scale, + input_zp, + weight_data, + packed_weight, + packed_weight_sums, + packed_weight_scales, + group_size, + bias_data, + packed_bias, + output)); + + graph.execute_nodes().emplace_back( + new DynamicDispatchNode(linear_qa_qw_node)); + + return; + } + + // Otherwise, input is dynamically quantized. Currently only per group 4-bit + // quantized weights is supported for this mode. + VK_CHECK_COND(weight_quant_config.nbits == 4); + + int64_t num_groups = 1; + if (weight_quant_config.granularity == kPerGroup) { + num_groups = graph.size_at(-2, weight_scales_data); + } + + TmpTensor int_input_sums( + &graph, + {num_groups, K}, + graph.dtype_of(output), + utils::kBuffer, + utils::kWidthPacked); + + DynamicDispatchNode quantize_and_pack_input_with_sums_node( + make_quantize_and_pack_linear_input_with_sums_node( + graph, + input_quant_config, + fp_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + packed_int_input, + group_size)); + + graph.execute_nodes().emplace_back( + new DynamicDispatchNode(quantize_and_pack_input_with_sums_node)); + + DynamicDispatchNode linear_dqa_qw_node(make_linear_dqa_qw_node( + graph, + input_quant_config, + weight_quant_config, + fp_input, + packed_int_input, + int_input_sums, + packed_input_scale, + packed_input_zp, + input_scale, + input_zp, + weight_data, + packed_weight, + packed_weight_sums, + packed_weight_scales, + group_size, + bias_data, + packed_bias, + output)); + + graph.execute_nodes().emplace_back( + new DynamicDispatchNode(linear_dqa_qw_node)); +} + +void linear_q8ta_q8csw(ComputeGraph& graph, const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef input_scale = args.at(idx++); + const ValueRef input_zp = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_sums_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(8, kPerTensor, {}, false); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + input_scale, + input_zp, + weight_data, + weight_sums_data, + weight_scales_data, + kDummyValueRef, // weight_zeros_data + kDummyValueRef, // group_size + bias_data, + output); +} + +void linear_q8csw(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 bias_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 + bias_data, + output); +} + +void linear_q4gsw(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 group_size = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t group_size_val = graph.extract_scalar(group_size); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + 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 + group_size, // group size + bias_data, + output); +} + +void linear_dq8ca_q4gsw( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef input_scale = args.at(idx++); + const ValueRef input_zp = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_sums_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef group_size = args.at(idx++); + const ValueRef bias_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t group_size_val = graph.extract_scalar(group_size); + + QuantizationConfig input_quant_config(8, kPerChannel, {}, false, true); + QuantizationConfig weight_quant_config(4, kPerGroup, {group_size_val}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + input_scale, + input_zp, + weight_data, + weight_sums_data, + weight_scales_data, + kDummyValueRef, // weight_zeros_data + group_size, // group_size + bias_data, + output); +} + +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(et_vk.linear_q4gsw.default, linear_q4gsw); + VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h new file mode 100644 index 00000000000..7b62c98390d --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h @@ -0,0 +1,29 @@ +/* + * 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 vkcompute { + +utils::uvec3 quantized_linear_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args); + +ValueRef prepack_quantized_linear_weight( + ComputeGraph& graph, + const QuantizationConfig& weight_quant_config, + const ValueRef qmat2_data); + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp index 07502a7a107..89c9e847724 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -15,6 +16,99 @@ namespace vkcompute { +// Custom global workgroup size function for linear_qcs8w +utils::uvec3 linear_qcs8w_global_wg_size( + 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 {static_cast(graph->numel_of(out)), 1, 1}; +} + +// Custom local workgroup size function for linear_qcs8w +utils::uvec3 linear_qcs8w_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64, 1, 1}; +} + +// Custom global workgroup size function for linear_qcsnw_tiled +utils::uvec3 linear_qcsnw_tiled_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + const ValueRef mat1 = args.at(1).refs.at(0); + + // Determine quantization bits from shader name + int quant_nbits = 8; + if (shader.kernel_name.find("qcs4w") != std::string::npos) { + quant_nbits = 4; + } + + std::vector mat1_sizes = graph->sizes_of(mat1); + const int64_t M = utils::val_at(-2, mat1_sizes); + uint32_t out_tile_nrows = 4; + if (M % 6 == 0) { + out_tile_nrows = 2; + } else if (M % 4 == 0) { + out_tile_nrows = 4; + } else if (M % 1 == 0) { + out_tile_nrows = 1; + } else { + out_tile_nrows = 4; + } + + // Number of output texels in the output tile + uint32_t out_tile_ntxcols = 1; + if (quant_nbits == 4) { + out_tile_ntxcols = 2; + } + + utils::uvec3 out_limits = graph->logical_limits_of(out); + uint32_t global_wg_x = utils::div_up(out_limits[0], out_tile_ntxcols); + return { + global_wg_x * (utils::div_up(out_limits[1], out_tile_nrows)), + 1, + out_limits[2]}; +} + +// Custom local workgroup size function for linear_qcsnw_tiled +utils::uvec3 linear_qcsnw_tiled_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + + // Check if using cooperative algorithm from shader name + bool use_coop_algorithm = + shader.kernel_name.find("_coop") != std::string::npos; + + if (use_coop_algorithm) { + return {8, 1, 8}; + } else { + return {64, 1, 1}; + } +} + void check_linear_qcsnw_args( const ComputeGraph& graph, const int quant_nbits, @@ -55,30 +149,33 @@ void resize_linear_qcsnw_node( const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr mat1 = graph->get_tensor(args[1].refs[0]); - vTensorPtr qmat2 = graph->get_tensor(args[1].refs[1]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef mat1 = args.at(1).refs.at(0); + const ValueRef qmat2 = args.at(1).refs.at(1); + + const std::vector mat1_sizes = graph->sizes_of(mat1); + const std::vector qmat2_sizes = graph->sizes_of(qmat2); - const int out_cols = utils::val_at(-2, mat1->sizes()); - int out_rows = utils::val_at(-1, qmat2->sizes()); + const int out_cols = utils::val_at(-2, mat1_sizes); + int out_rows = utils::val_at(-1, qmat2_sizes); // Byte dtype suggests 4-bit quantization in which case the weight tensor is // packed with 2 values per byte. - if (qmat2->dtype() == vkapi::kByte) { + if (graph->dtype_of(qmat2) == vkapi::kByte) { out_rows *= 2; } std::vector new_out_sizes(3); - if (mat1->sizes().size() == 2) { + if (mat1_sizes.size() == 2) { new_out_sizes.resize(2); new_out_sizes.at(0) = out_cols; new_out_sizes.at(1) = out_rows; } else { - new_out_sizes.at(0) = mat1->sizes().at(0); + new_out_sizes.at(0) = mat1_sizes.at(0); new_out_sizes.at(1) = out_cols; new_out_sizes.at(2) = out_rows; } - out->virtual_resize(new_out_sizes); + graph->virtual_resize(out, new_out_sizes); } void add_linear_qcs8w_node( @@ -135,11 +232,11 @@ void add_linear_qcs8w_node( static_cast(graph.numel_of(out_W_packed)), 1, 1}; const utils::uvec3 local_wg{64, 1, 1}; - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg, - local_wg, + linear_qcs8w_global_wg_size, + linear_qcs8w_local_wg_size, // Inputs and Outputs {{out_W_packed, vkapi::MemoryAccessType::WRITE}, {{mat1_W_packed, q_mat2, scales}, vkapi::MemoryAccessType::READ}}, @@ -244,11 +341,11 @@ void add_linear_qcsnw_tiled_node( local_wg_size = {8, 1, 8}; } - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - local_wg_size, + linear_qcsnw_tiled_global_wg_size, + linear_qcsnw_tiled_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{mat1, q_mat2, scales}, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQGANW.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQGANW.cpp deleted file mode 100644 index 8c7c6b0cdf9..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQGANW.cpp +++ /dev/null @@ -1,224 +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. - */ - -#include - -#include -#include - -#include -#include - -namespace vkcompute { - -void check_linear_qga4w_args( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat2_data, - const ValueRef group_size, - const ValueRef scales_and_zeros, - const ValueRef out) { - VK_CHECK_COND(graph.val_is_tensor(mat1)); - VK_CHECK_COND(graph.val_is_tref(mat2_data)); - VK_CHECK_COND(graph.val_is_tref(scales_and_zeros)); - - VK_CHECK_COND(graph.dim_of(mat1) <= 3); - VK_CHECK_COND(graph.dim_of(mat2_data) == 2); - VK_CHECK_COND(graph.dim_of(scales_and_zeros) == 3); - - VK_CHECK_COND(graph.size_at(-3, mat1) == 1); - const int K = graph.size_at(-1, mat1); - VK_CHECK_COND(graph.size_at(-1, mat2_data) * 2 == K); - - const int group_size_val = graph.extract_scalar(group_size); - VK_CHECK_COND(K % group_size_val == 0); - // Due to the way weight packing works, group size needs to be a multiple of 8 - VK_CHECK_COND(group_size_val % 8 == 0); - - VK_CHECK_COND(graph.has_standard_axis_map(mat1)); - VK_CHECK_COND(graph.has_standard_axis_map(out)); -} - -void resize_linear_qga4w_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - (void)extra_args; - - ValueRef out = args.at(0).refs.at(0); - ValueRef mat1 = args.at(1).refs.at(0); - ValueRef mat2_data = extra_args.at(0); - - std::vector mat1_sizes = graph->sizes_of(mat1); - std::vector mat2_sizes = graph->sizes_of(mat2_data); - - const int64_t out_cols = utils::val_at(-2, mat1_sizes); - const int64_t out_rows = utils::val_at(-2, mat2_sizes); - - std::vector new_out_sizes(3); - if (mat1_sizes.size() == 2) { - new_out_sizes.resize(2); - new_out_sizes.at(0) = out_cols; - new_out_sizes.at(1) = out_rows; - } else { - new_out_sizes.at(0) = mat1_sizes.at(0); - new_out_sizes.at(1) = out_cols; - new_out_sizes.at(2) = out_rows; - } - - graph->virtual_resize(out, new_out_sizes); -} - -/** - * Determines if the cooperative algorithm should be used based on input tensor - * dimensions. Apply the coop algorithm for gemv cases, i.e. mat1 is avector as - * as opposed to a matrix. - */ -bool should_use_coop_algorithm(ComputeGraph* graph, const ValueRef& mat1) { - return graph->size_at(-2, mat1) == 1; -} - -vkapi::ShaderInfo pick_linear_qga4w_shader( - 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 mat1 = args.at(1).refs.at(0); - const ValueRef mat2 = args.at(1).refs.at(1); - - const bool use_coop_algorithm = should_use_coop_algorithm(graph, mat1); - - std::string kernel_name = "linear_qga4w"; - if (use_coop_algorithm) { - kernel_name += "_coop"; - } else { - kernel_name += "_tiled"; - } - add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(mat1)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(mat2)); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - - return VK_KERNEL_FROM_STR(kernel_name); -} - -utils::uvec3 linear_qga4w_global_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - if (!use_coop_algorithm) { - // Constructing the global workgroup size for the tiled algorithm - utils::uvec3 global_wg_size = graph->logical_limits_of(out); - // Each shader thread computes a 4 high x 8 wide tile of the output matrix, - // which is equivalent to 4 x 2 texels. Since the output tensor must be - // width packed, div-up the "texel-width" of the output by 2 and the height - // of the output tensor by 4 to obtain the number of tiles that need to be - // computed. - global_wg_size[0] = utils::div_up(global_wg_size[0], uint32_t(2)); - global_wg_size[1] = utils::div_up(global_wg_size[1], uint32_t(4)); - return global_wg_size; - } - - uint32_t output_channels = graph->size_at(-1, out); - uint32_t batch_size = graph->size_at(-2, out); - - // Constructing the global workgroup size of the co-operative algorithm. The - // local work group size is 64, and each local work group co-operates to - // compute 8 output channels of the output. Therefore, a total of - // (output_channels / 8 x 64) threads should be launched, assuming a batch - // size of 1. - return {64, utils::div_up(output_channels, 8u), batch_size}; -} - -utils::uvec3 linear_qga4w_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)args; - (void)resize_args; - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - if (use_coop_algorithm) { - return {64, 1, 1}; - } else { - return graph->create_local_wg_size(global_workgroup_size); - } -} - -void add_linear_qga4w_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat2_data, - const ValueRef group_size, - const ValueRef scales_and_zeros_data, - const ValueRef out) { - check_linear_qga4w_args( - graph, mat1, mat2_data, group_size, scales_and_zeros_data, out); - - const uint32_t group_size_val = graph.extract_scalar(group_size); - - ValueRef mat2 = - prepack_int4_linear_weight_transposed_block_4x8(graph, mat2_data); - - ValueRef scales_and_zeros = prepack_standard( - graph, scales_and_zeros_data, utils::kBuffer, utils::kWidthPacked); - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - pick_linear_qga4w_shader, - linear_qga4w_global_wg_size, - linear_qga4w_local_wg_size, - // Inputs and Outputs - {{out, vkapi::kWrite}, {{mat1, mat2, scales_and_zeros}, vkapi::kRead}}, - // Shader params buffers - {}, - // Push Constants - {graph.sizes_pc_of(out), - graph.sizes_pc_of(mat1), - graph.sizes_pc_of(mat2)}, - // Specialization Constants - {SV(group_size_val)}, - // Resize Args - {mat2_data}, - // Resizing Logic - resize_linear_qga4w_node)); -} - -void linear_weight_int4( - ComputeGraph& graph, - const std::vector& args) { - return add_linear_qga4w_node( - graph, - args[0], // mat1 - args[1], // mat2 - args[2], // group_size - args[3], // scales_and_zeros - // There is an unused variable inner_k_tiles which is used to call - // _convert_weight_to_int4pack in the AOT custom op, which is why the 4th - // argument is skipped. - args[5] // out - ); -} - -REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.linear_weight_int4.default, linear_weight_int4); -} - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear_QTA8A_QGA4W.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear_QTA8A_QGA4W.cpp deleted file mode 100644 index 728d38c3e2d..00000000000 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear_QTA8A_QGA4W.cpp +++ /dev/null @@ -1,267 +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. - */ - -#include - -#include -#include - -#include -#include - -namespace vkcompute { - -void check_linear_qta8a_qga4w_args( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat1_scale, - const ValueRef mat1_zero_point, - const ValueRef mat2_data, - const ValueRef group_size, - const ValueRef weight_scales, - const ValueRef weight_zeros, - const ValueRef out) { - VK_CHECK_COND(graph.val_is_tensor(mat1)); - VK_CHECK_COND(graph.val_is_tensor(mat1_scale)); - VK_CHECK_COND(graph.val_is_tensor(mat1_zero_point)); - VK_CHECK_COND(graph.val_is_tref(mat2_data)); - VK_CHECK_COND(graph.val_is_tref(weight_scales)); - VK_CHECK_COND(graph.val_is_tref(weight_zeros)); - - VK_CHECK_COND(graph.dim_of(mat1) <= 3); - VK_CHECK_COND(graph.dim_of(mat2_data) == 2); - VK_CHECK_COND(graph.dim_of(weight_scales) == 2); - VK_CHECK_COND(graph.dim_of(weight_zeros) == 2); - - VK_CHECK_COND(graph.size_at(-3, mat1) == 1); - const int K = graph.size_at(-1, mat1); - VK_CHECK_COND(graph.size_at(-1, mat2_data) * 2 == K); - - const int group_size_val = graph.extract_scalar(group_size); - VK_CHECK_COND(K % group_size_val == 0); - // Due to the way weight packing works, group size needs to be a multiple of 8 - VK_CHECK_COND(group_size_val % 8 == 0); - - VK_CHECK_COND(graph.has_standard_axis_map(mat1)); - VK_CHECK_COND(graph.has_standard_axis_map(out)); - - // Check that scale and zero_point tensors are buffer storage with width - // packing - VK_CHECK_COND(graph.is_buffer_storage(mat1_scale)); - VK_CHECK_COND(graph.packed_dim_of(mat1_scale) == WHCN::kWidthDim); - VK_CHECK_COND(graph.is_buffer_storage(mat1_zero_point)); - VK_CHECK_COND(graph.packed_dim_of(mat1_zero_point) == WHCN::kWidthDim); - - // Calculate number of tokens for input - int64_t input_num_tokens = 1; - const auto mat1_sizes = graph.sizes_of(mat1); - for (size_t i = 0; i < mat1_sizes.size() - 1; i++) { - input_num_tokens *= mat1_sizes[i]; - } - - // Verify scale and zero_point tensor sizes match number of tokens - const auto mat1_scale_sizes = graph.sizes_of(mat1_scale); - const auto mat1_zero_point_sizes = graph.sizes_of(mat1_zero_point); - - VK_CHECK_COND( - utils::val_at(-1, mat1_scale_sizes) == input_num_tokens); - VK_CHECK_COND( - utils::val_at(-1, mat1_zero_point_sizes) == input_num_tokens); - - // Verify weight scales and zeros have the same shape - const auto weight_scales_sizes = graph.sizes_of(weight_scales); - const auto weight_zeros_sizes = graph.sizes_of(weight_zeros); - VK_CHECK_COND(weight_scales_sizes == weight_zeros_sizes); -} - -void resize_linear_qta8a_qga4w_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - (void)extra_args; - - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr mat1 = graph->get_tensor(args[1].refs[0]); - vTensorPtr mat2 = graph->get_tensor(args[1].refs[1]); - - const int64_t out_cols = utils::val_at(-2, mat1->sizes()); - const int64_t out_rows = utils::val_at(-1, mat2->sizes()) * 2; - - std::vector new_out_sizes(3); - if (mat1->sizes().size() == 2) { - new_out_sizes.resize(2); - new_out_sizes.at(0) = out_cols; - new_out_sizes.at(1) = out_rows; - } else { - new_out_sizes.at(0) = mat1->sizes().at(0); - new_out_sizes.at(1) = out_cols; - new_out_sizes.at(2) = out_rows; - } - - out->virtual_resize(new_out_sizes); -} - -/** - * Determines if the cooperative algorithm should be used based on input tensor - * dimensions. Apply the coop algorithm for vectors (GEMV cases), tiled for - * matrices (GEMM cases). - */ -bool should_use_coop_algorithm_qta8a_qga4w( - ComputeGraph* graph, - const ValueRef& mat1) { - const uint32_t M = graph->size_at(-2, mat1); - // Use coop algorithm for vectors (GEMV), tiled for larger matrices (GEMM) - return M == 1; -} - -vkapi::ShaderInfo pick_linear_qta8a_qga4w_shader( - 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 mat1 = args.at(1).refs.at(0); - const ValueRef mat2 = args.at(1).refs.at(1); - - const bool use_coop_algorithm = - should_use_coop_algorithm_qta8a_qga4w(graph, mat1); - - std::string kernel_name = "linear_qta8a_qga4w"; - if (use_coop_algorithm) { - kernel_name += "_coop"; - } else { - kernel_name += "_tiled"; - } - add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(mat1)); - add_storage_type_suffix(kernel_name, graph->storage_type_of(mat2)); - add_dtype_suffix(kernel_name, graph->dtype_of(out)); - - return VK_KERNEL_FROM_STR(kernel_name); -} - -utils::uvec3 linear_qta8a_qga4w_global_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - // C = 1, H = 2, W = 3 - // global_wg_size = {round_up(C / 2f), round_up(H / 3f), W} --> (2W, 1H, 0C) - // --> {1, 1, 3} global - - utils::uvec3 global_wg_size = graph->logical_limits_of(out); - global_wg_size[0] = utils::div_up(global_wg_size[0], uint32_t(2)); - if (!use_coop_algorithm) { // GEMM - TILED - global_wg_size[1] = utils::div_up(global_wg_size[1], uint32_t(3)); - } - - return global_wg_size; -} - -utils::uvec3 linear_qta8a_qga4w_local_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const utils::uvec3& global_workgroup_size, - const std::vector& args, - const std::vector& resize_args) { - (void)args; - (void)resize_args; - - const bool use_coop_algorithm = - shader.kernel_name.find("_coop") != std::string::npos; - - utils::uvec3 local_wg_size; - if (use_coop_algorithm) { // GEMV - COOP - local_wg_size = {8, 1, 8}; - } else { // GEMM - TILED - local_wg_size = graph->create_local_wg_size(global_workgroup_size); - } - - return local_wg_size; -} - -void add_linear_qta8a_qga4w_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef mat1_scale, - const ValueRef mat1_zero_point, - const ValueRef mat2_data, - const ValueRef group_size, - const ValueRef weight_scales_data, - const ValueRef weight_zeros_data, - const ValueRef out) { - check_linear_qta8a_qga4w_args( - graph, - mat1, - mat1_scale, - mat1_zero_point, - mat2_data, - group_size, - weight_scales_data, - weight_zeros_data, - out); - const uint32_t group_size_val = graph.extract_scalar(group_size); - - ValueRef mat2 = - prepack_int4_linear_weight_transposed_interleaved(graph, mat2_data); - ValueRef weight_scales = prepack_standard( - graph, weight_scales_data, utils::kBuffer, utils::kWidthPacked); - ValueRef weight_zeros = prepack_standard( - graph, weight_zeros_data, utils::kBuffer, utils::kWidthPacked); - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - pick_linear_qta8a_qga4w_shader, - linear_qta8a_qga4w_global_wg_size, - linear_qta8a_qga4w_local_wg_size, - // Inputs and Outputs - {{out, vkapi::kWrite}, - {{mat1, mat2, weight_scales, weight_zeros, mat1_scale, mat1_zero_point}, - vkapi::kRead}}, - // Shader params buffers - {}, - // Push Constants - {graph.sizes_pc_of(out), - graph.sizes_pc_of(mat1), - graph.sizes_pc_of(mat2)}, - // Specialization Constants - {SV(group_size_val)}, - // Resize Args - {}, - // Resizing Logic - resize_linear_qta8a_qga4w_node)); -} - -void linear_qta8a_qga4w( - ComputeGraph& graph, - const std::vector& args) { - return add_linear_qta8a_qga4w_node( - graph, - args[0], // quantized input (char tensor) - args[1], // input_scale (float buffer tensor) - args[2], // input_zero_point (int buffer tensor) - args[3], // quantized weights (4-bit packed, byte) - args[4], // group_size (int) - args[5], // weight_scales (float tensor) - args[6], // weight_zeros (int tensor) - args[7] // float output tensor - ); -} - -REGISTER_OPERATORS { - VK_REGISTER_OP(et_vk.linear_qta8a_qga4w.default, linear_qta8a_qga4w); -} - -} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp index c0fd442ec50..6ad1d7f371d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp @@ -22,14 +22,34 @@ void resize_reduce_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - int32_t reduce_dim_nchw = graph->extract_scalar(resize_args.at(0)); + const int32_t reduce_dim_nchw = + graph->extract_scalar(resize_args.at(0)); - std::vector new_sizes = in->sizes(); + std::vector new_sizes = graph->sizes_of(in); new_sizes.at(normalize(reduce_dim_nchw, new_sizes.size())) = 1; - out->virtual_resize(new_sizes); + graph->virtual_resize(out, new_sizes); +} + +void resize_reduce2d_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); + + // Extract the dimensions to reduce over + const std::vector dims_list = + graph->extract_int_or_symint_list(resize_args.at(0)); + int32_t reduce_dim1_nchw = dims_list[0]; + int32_t reduce_dim2_nchw = dims_list[1]; + + std::vector new_sizes = graph->sizes_of(in); + new_sizes.at(normalize(reduce_dim1_nchw, new_sizes.size())) = 1; + new_sizes.at(normalize(reduce_dim2_nchw, new_sizes.size())) = 1; + graph->virtual_resize(out, new_sizes); } utils::uvec3 reduce_global_wg_size( @@ -137,15 +157,101 @@ void add_reduce_node( resize_reduce_node)); } +void add_reduce2d_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef dims_ref, + const ValueRef out, + const std::string& op_name) { + VK_CHECK_COND( + !graph.is_buffer_storage(in) && !graph.is_buffer_storage(out), + "Vulkan reduction only supports texture storage"); + + const int64_t ndim = graph.dim_of(in); + + // Extract the two dimensions to reduce over + const std::vector dims_list = + graph.extract_int_or_symint_list(dims_ref); + VK_CHECK_COND( + dims_list.size() == 2, "reduce2d requires exactly 2 dimensions"); + + int32_t reduce_dim1 = normalize(dims_list[0], ndim); + int32_t reduce_dim2 = normalize(dims_list[1], ndim); + + // Convert to WHCN format + reduce_dim1 = nchw_dim_to_whcn_dim(reduce_dim1, ndim); + reduce_dim2 = nchw_dim_to_whcn_dim(reduce_dim2, ndim); + + // Check that none of the reduction dims are packed + VK_CHECK_COND(graph.packed_dim_of(in) != reduce_dim1); + VK_CHECK_COND(graph.packed_dim_of(in) != reduce_dim2); + VK_CHECK_COND(graph.packed_dim_of(out) != reduce_dim1); + VK_CHECK_COND(graph.packed_dim_of(out) != reduce_dim2); + + // Check that the concat dim is not one of the reduction dims + if (graph.dim_of(in) == 4 && graph.size_at(0, in) > 1) { + VK_CHECK_COND(graph.concat_dim_of(in) != reduce_dim1); + VK_CHECK_COND(graph.concat_dim_of(in) != reduce_dim2); + VK_CHECK_COND(graph.concat_dim_of(out) != reduce_dim1); + VK_CHECK_COND(graph.concat_dim_of(out) != reduce_dim2); + } + + std::string kernel_name = op_name + "2d"; // Add "2d" suffix + kernel_name.reserve(kShaderNameReserve); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + // Calculate group_dim for specialization constants (use remaining dimension) + int32_t group_dim = 0; + for (int i = 0; i < 3; i++) { + if (i != reduce_dim1 && i != reduce_dim2) { + group_dim = i; + break; + } + } + + const ValueRef reduce_dim1_whcn_ref = + graph.get_or_add_value_for_int(reduce_dim1); + 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); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + reduce_global_wg_size, + reduce_local_wg_size, + // Inputs and Outputs + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + // Shader params buffers + {graph.logical_limits_ubo(in), graph.sizes_ubo(in)}, + // Push Constants + {}, + // Specialization Constants + {graph.packed_dim_of(out), reduce_dim1, reduce_dim2, group_dim}, + // Resize Args + {dims_ref, + reduce_dim1_whcn_ref, + reduce_dim2_whcn_ref, + group_dim_whcn_ref}, + // Resizing Logic + resize_reduce2d_node)); +} + #define DEFINE_REDUCE_FN(op_name, out_arg_idx) \ void op_name(ComputeGraph& graph, const std::vector& args) { \ const std::vector dims_list = \ graph.extract_int_or_symint_list(args[1]); \ - VK_CHECK_COND(dims_list.size() == 1); \ - const int64_t dim_val = dims_list.at(0); \ - const ValueRef dim_ref = graph.get_or_add_value_for_int(dim_val); \ - return add_reduce_node( \ - graph, args[0], dim_ref, args[out_arg_idx], #op_name); \ + if (dims_list.size() == 1) { \ + const int64_t dim_val = dims_list.at(0); \ + const ValueRef dim_ref = graph.get_or_add_value_for_int(dim_val); \ + return add_reduce_node( \ + graph, args[0], dim_ref, args[out_arg_idx], #op_name); \ + } \ + if (dims_list.size() == 2) { \ + return add_reduce2d_node( \ + graph, args[0], args[1], args[out_arg_idx], #op_name); \ + } \ + VK_CHECK_COND(false, "Only 1 or 2 dimensions supported"); \ } DEFINE_REDUCE_FN(sum, 4) diff --git a/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp b/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp index f472e4dad0d..72c1637a2c9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Repeat.cpp @@ -8,6 +8,7 @@ #include +#include #include #include #include @@ -20,39 +21,43 @@ namespace vkcompute { namespace { void check_args( - const api::vTensor& in, + ComputeGraph& graph, + const ValueRef in, const std::vector& repeats, - const api::vTensor& out) { - VK_CHECK_COND(check_same_packed_dim(in, out)); + const ValueRef out) { + VK_CHECK_COND(graph.packed_dim_of(in) == graph.packed_dim_of(out)); - VK_CHECK_COND(in.storage_type() == out.storage_type()); - if (in.storage_type() == utils::kTexture2D) { - VK_CHECK_COND(in.dim() <= 2); + VK_CHECK_COND(graph.storage_type_of(in) == graph.storage_type_of(out)); + if (graph.storage_type_of(in) == utils::kTexture2D) { + VK_CHECK_COND(graph.dim_of(in) <= 2); } - int64_t in_dim = in.dim(); + const int64_t in_dim = graph.dim_of(in); VK_CHECK_COND( in_dim <= repeats.size(), "Input tensor dim size must be not greater than the repeat argument's size"); + const std::vector in_sizes = graph.sizes_of(in); + const std::vector out_sizes = graph.sizes_of(out); + VK_CHECK_COND( - dim_at(in.sizes()) * dim_at(repeats) == - dim_at(out.sizes()), + dim_at(in_sizes) * dim_at(repeats) == + dim_at(out_sizes), "Output's width doesn't match input's width * repeat count"); VK_CHECK_COND( - dim_at(in.sizes()) * dim_at(repeats) == - dim_at(out.sizes()), + dim_at(in_sizes) * dim_at(repeats) == + dim_at(out_sizes), "Output's height doesn't match input's height * repeat count"); VK_CHECK_COND( - dim_at(in.sizes()) * dim_at(repeats) == - dim_at(out.sizes()), + dim_at(in_sizes) * dim_at(repeats) == + dim_at(out_sizes), "Output's channel doesn't match input's channel * repeat count"); VK_CHECK_COND( - dim_at(in.sizes()) * dim_at(repeats) == - dim_at(out.sizes()), + dim_at(in_sizes) * dim_at(repeats) == + dim_at(out_sizes), "Output's batch doesn't match input's batch * repeat count"); } @@ -65,15 +70,14 @@ void add_repeat_node( ValueRef out) { const std::vector repeats = *(graph.get_int_list(repeats_ref)); - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - check_args(*t_in, repeats, *t_out); + check_args(graph, in, repeats, out); + const std::vector in_sizes = graph.sizes_of(in); const utils::ivec4 src_dims{ - dim_at(t_in->sizes()), - dim_at(t_in->sizes()), - dim_at(t_in->sizes()), - dim_at(t_in->sizes())}; + dim_at(in_sizes), + dim_at(in_sizes), + dim_at(in_sizes), + dim_at(in_sizes)}; const utils::ivec4 dst_repeats{ dim_at(repeats), dim_at(repeats), @@ -82,22 +86,22 @@ void add_repeat_node( std::string kernel_name = "repeat"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); // A copy of range with the last element set to batch size of the input tensor - const utils::ivec3 wg_size = t_out->logical_limits(); + const utils::ivec3 wg_size = graph.logical_limits_of(out); const auto shader = VK_KERNEL_FROM_STR(kernel_name); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - wg_size, - graph.create_local_wg_size(wg_size), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs { - {out, vkapi::MemoryAccessType::WRITE}, - {in, vkapi::MemoryAccessType::READ}, + {out, vkapi::kWrite}, + {in, vkapi::kRead}, }, // Parameter buffers {}, diff --git a/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp b/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp index 5bfadf43160..221d0d23f51 100644 --- a/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/RepeatInterleave.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -20,17 +21,17 @@ void resize_repeat_interleave_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - const int64_t nrepeats = graph->extract_scalar(extra_args[0]); - int64_t repeat_dim = graph->extract_scalar(extra_args[1]); + const int64_t nrepeats = graph->extract_scalar(extra_args.at(0)); + int64_t repeat_dim = graph->extract_scalar(extra_args.at(1)); - std::vector new_sizes = in->sizes(); + std::vector new_sizes = graph->sizes_of(in); repeat_dim = normalize(repeat_dim, new_sizes.size()); new_sizes.at(repeat_dim) *= nrepeats; - out->virtual_resize(new_sizes); + graph->virtual_resize(out, new_sizes); } void add_repeat_interleave_node( @@ -49,16 +50,11 @@ void add_repeat_interleave_node( std::string kernel_name = "repeat_interleave"; add_dtype_suffix(kernel_name, graph.dtype_of(out)); - const utils::uvec3 global_wg_size = graph.logical_limits_of(in); - const utils::uvec3 local_wg_size = graph.create_local_wg_size(global_wg_size); - - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - // Shader VK_KERNEL_FROM_STR(kernel_name), - // Workgroup sizes - global_wg_size, - local_wg_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::MemoryAccessType::WRITE}, {in, vkapi::MemoryAccessType::READ}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index 6057f1e183a..8edaebd11ff 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -24,6 +24,58 @@ namespace vkcompute { +bool is_single_token(ComputeGraph* graph, const ValueRef& q_projected) { + return graph->size_at(-3, q_projected) == 1; +} + +// +// Resize functions +// + +void resize_compute_attn_weights_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef attn_weights = args.at(0).refs.at(0); + const ValueRef q_projected = args.at(1).refs.at(0); + const ValueRef input_pos_symint = resize_args.at(0); + + const uint32_t num_q_heads = graph->size_at(-2, q_projected); + const uint32_t seq_len = graph->size_at(-3, q_projected); + + const int32_t input_pos_val = graph->read_symint(input_pos_symint); + + const uint32_t context_len = seq_len + input_pos_val; + + std::vector out_sizes = { + 1, // batch + num_q_heads, + seq_len, + utils::align_up_4(context_len)}; + + graph->virtual_resize(attn_weights, out_sizes); +} + +void resize_sdpa_attn_weights_softmax_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef attn_weights_softmax = args.at(0).refs.at(0); + const ValueRef attn_weights = args.at(1).refs.at(0); + + graph->virtual_resize(attn_weights_softmax, graph->sizes_of(attn_weights)); +} + +void resize_sdpa_compute_out_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef q_projected = resize_args.at(0); + + graph->virtual_resize(out, graph->sizes_of(q_projected)); +} + void resize_sdpa_out( ComputeGraph* graph, const std::vector& args, @@ -33,203 +85,210 @@ void resize_sdpa_out( int arg_idx = 0; const ValueRef q_projected = extra_args[arg_idx++]; const ValueRef out = extra_args[arg_idx++]; - graph->get_tensor(out)->virtual_resize(graph->sizes_of(q_projected)); + graph->virtual_resize(out, graph->sizes_of(q_projected)); } -void resize_flash_attention_out( +// +// Shader dispatch pick functions +// + +utils::uvec3 kv_cache_update_global_wg_size( ComputeGraph* graph, + const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { + (void)shader; (void)resize_args; - // Find the output tensor in the args - it's the first tensor in the first - // ArgGroup - const ValueRef out = args.at(0).refs.at(0); - // Find the query tensor - it's the first tensor in the second ArgGroup - const ValueRef q_projected = args.at(1).refs.at(0); + const ValueRef projected = args.at(1).refs.at(0); + + const uint32_t head_dim_size = graph->size_at(-1, projected); + const uint32_t num_heads = graph->size_at(-2, projected); + const uint32_t seq_len = graph->size_at(-3, projected); - // Resize output to match query dimensions - graph->get_tensor(out)->virtual_resize(graph->sizes_of(q_projected)); + return {utils::div_up_4(head_dim_size), seq_len, num_heads}; } -// Flash Attention implementation using single compute shader -utils::uvec3 flash_attention_global_wg_size( +utils::uvec3 attn_weight_scale_and_mask_global_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { (void)shader; + (void)resize_args; - const ValueRef q_projected = resize_args.at(0); - const ValueRef block_size_r = resize_args.at(1); + const ValueRef attn_weight = args.at(0).refs.at(0); + + if (graph->is_buffer_storage(attn_weight)) { + return { + graph->size_at(-1, attn_weight), + graph->size_at(-2, attn_weight), + graph->size_at(-3, attn_weight), + }; + } else { + return graph->logical_limits_of(attn_weight); + } +} + +vkapi::ShaderInfo pick_sdpa_compute_attn_weights_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef q_projected = args.at(1).refs.at(0); + const ValueRef k_cache = args.at(1).refs.at(1); + + const bool is_gemv = is_single_token(graph, q_projected); - // Get tensor dimensions - PyTorch format is [B, N, H, D] - // But Vulkan uses negative indexing: -4=B, -3=N, -2=H, -1=D - const int32_t B = graph->size_at(-4, q_projected); // batch - const int32_t N = graph->size_at(-3, q_projected); // sequence length - const int32_t H = graph->size_at(-2, q_projected); // num heads - const int32_t Br = - static_cast(graph->extract_scalar(block_size_r)); + std::string shader_name = "sdpa_compute_attn_weights"; + if (is_gemv) { + shader_name += "_coop"; + } else { + shader_name += "_tiled"; + } - // Calculate number of row blocks - const int32_t Tr = (N + Br - 1) / Br; + add_storage_type_suffix(shader_name, graph->storage_type_of(q_projected)); + add_storage_type_suffix(shader_name, graph->storage_type_of(k_cache)); + add_dtype_suffix(shader_name, graph->dtype_of(q_projected)); - // Dispatch size: (B * H * Tr, 1, 1) - return {static_cast(B * H * Tr), 1, 1}; + return VK_KERNEL_FROM_STR(shader_name); } -void flash_attention_impl( - ComputeGraph& graph, - const std::vector& args) { - int arg_idx = 0; - const ValueRef q_projected = args[arg_idx++]; - const ValueRef k_cache = args[arg_idx++]; - const ValueRef v_cache = args[arg_idx++]; - const ValueRef input_pos_symint = args[arg_idx++]; - const ValueRef attn_mask = args[arg_idx++]; - const ValueRef dropout_p = args[arg_idx++]; - const ValueRef is_causal = args[arg_idx++]; - const ValueRef scale = args[arg_idx++]; +utils::uvec3 pick_sdpa_compute_attn_weights_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef q_projected = args.at(1).refs.at(0); + const ValueRef input_pos_symint = resize_args.at(0); - const ValueRef out = args[arg_idx++]; + const uint32_t num_q_heads = graph->size_at(-2, q_projected); + const uint32_t seq_len = graph->size_at(-3, q_projected); - // Extract input_pos value for causal masking - const int32_t input_pos_val = graph.read_symint(input_pos_symint); + const int32_t input_pos_val = graph->read_symint(input_pos_symint); - const ValueRef k_cache_tensor = k_cache; - const ValueRef v_cache_tensor = v_cache; + const uint32_t context_len = seq_len + input_pos_val; - // Validation checks - re-enable with correct indexing - VK_CHECK_COND(graph.size_at(-4, q_projected) == 1); // batch size = 1 - VK_CHECK_COND(graph.size_at(-4, k_cache_tensor) == 1); - VK_CHECK_COND(graph.size_at(-4, v_cache_tensor) == 1); - VK_CHECK_COND( - graph.sizes_of(k_cache_tensor) == graph.sizes_of(v_cache_tensor)); - VK_CHECK_COND( - graph.size_at(-1, q_projected) == - graph.size_at(-1, k_cache_tensor)); // head_dim must match - VK_CHECK_COND( - graph.val_is_none(dropout_p) || - graph.extract_scalar(dropout_p) == 0); - VK_CHECK_COND(graph.val_is_none(scale)); - VK_CHECK_COND( - graph.val_is_none(is_causal) || graph.extract_scalar(is_causal)); - VK_CHECK_COND(graph.val_is_none(attn_mask)); + const uint32_t N4 = utils::div_up_4(context_len); + const uint32_t M4 = utils::div_up_4(seq_len); - // Ensure all tensors use buffer storage for Flash Attention - VK_CHECK_COND(graph.is_buffer_storage(q_projected)); - VK_CHECK_COND(graph.is_buffer_storage(k_cache_tensor)); - VK_CHECK_COND(graph.is_buffer_storage(v_cache_tensor)); - VK_CHECK_COND(graph.is_buffer_storage(out)); + return {N4, M4, num_q_heads}; +} - // Calculate scale factor - const int32_t head_dim_size = graph.size_at(-1, q_projected); - const float scale_val = 1.0f / std::sqrt(static_cast(head_dim_size)); +utils::uvec3 pick_sdpa_compute_attn_weights_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + const bool use_coop_algorithm = + shader.kernel_name.find("_coop") != std::string::npos; - // Get number of heads for multi-query attention support - const int32_t num_heads = graph.size_at(-2, q_projected); - const int32_t num_kv_heads = graph.size_at(-2, k_cache_tensor); + if (use_coop_algorithm) { + return {1, 64, 1}; + } else { + return pick_hw_square_wg_size( + graph, shader, global_workgroup_size, args, resize_args); + } +} - const int32_t block_size_r = 32; // Row block size - const int32_t block_size_c = 32; // Column block size +utils::uvec3 pick_sdpa_attn_weights_softmax_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef q_projected = resize_args.at(0); - // l and m have shape [B, H, N] - std::vector lm_sizes = { - graph.size_at(-4, q_projected), // B (batch) - graph.size_at(-2, q_projected), // H (num heads) - graph.size_at(-3, q_projected) // N (sequence length) - }; + const uint32_t num_q_heads = graph->size_at(-2, q_projected); + const uint32_t seq_len = graph->size_at(-3, q_projected); - // t_l stores row-wise normalization sums for softmax computation - // t_m stores row-wise maximum values for numerical stability in softmax - TmpTensor t_l(&graph, lm_sizes, vkapi::kFloat); - TmpTensor t_m(&graph, lm_sizes, vkapi::kFloat); + return {1, seq_len, num_q_heads}; +} - std::string kernel_name = "flash_attention"; - add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); - add_dtype_suffix(kernel_name, graph.dtype_of(out)); +utils::uvec3 pick_sdpa_attn_weights_softmax_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + return {64, 1, 1}; +} - // Set up parameter buffers - vkapi::ParamsBindList param_ubos = { - graph.sizes_ubo(q_projected), // Q_sizes - graph.sizes_ubo(k_cache_tensor), // K_sizes - graph.sizes_ubo(v_cache_tensor), // V_sizes - graph.sizes_ubo(out), // O_sizes - graph.sizes_ubo(t_l), // l_sizes (3D) - graph.sizes_ubo(t_m), // m_sizes (3D) - graph.create_params_buffer(scale_val), // scale - graph.create_params_buffer(block_size_r), // block_size_r - graph.create_params_buffer(block_size_c), // block_size_c - graph.create_params_buffer(input_pos_val), // input_pos - graph.create_params_buffer(num_heads), // num_heads - graph.create_params_buffer(num_kv_heads) // num_kv_heads - }; - - // Create block size references for dispatch calculation - const ValueRef block_size_r_ref = - graph.add_scalar(static_cast(block_size_r)); +vkapi::ShaderInfo pick_sdpa_compute_out_shader( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef v_cache = args.at(1).refs.at(1); - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - flash_attention_global_wg_size, - default_pick_local_wg_size, - // Inputs and Outputs - { - {{out, t_l, t_m}, vkapi::kReadWrite}, - {{q_projected, k_cache_tensor, v_cache_tensor}, vkapi::kRead}, - }, - // Shader param buffers - param_ubos, - // Push Constants - {}, - // Specialization Constants - {}, - // Resize Args - {q_projected, block_size_r_ref}, - // Resizing Logic - resize_flash_attention_out)); + const ValueRef q_projected = resize_args.at(0); + + const bool is_gemv = is_single_token(graph, q_projected); + + std::string shader_name = "sdpa_compute_out"; + if (is_gemv) { + shader_name += "_coop"; + } else { + shader_name += "_tiled"; + } + + add_storage_type_suffix(shader_name, graph->storage_type_of(out)); + add_storage_type_suffix(shader_name, graph->storage_type_of(v_cache)); + add_dtype_suffix(shader_name, graph->dtype_of(out)); + + return VK_KERNEL_FROM_STR(shader_name); } -utils::uvec3 kv_cache_update_global_wg_size( +utils::uvec3 pick_sdpa_compute_out_global_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, const std::vector& args, const std::vector& resize_args) { - (void)shader; - (void)resize_args; + const ValueRef q_projected = resize_args.at(0); - const ValueRef cache = args.at(0).refs.at(0); - const ValueRef projected = args.at(1).refs.at(0); + const uint32_t head_dim = graph->size_at(-1, q_projected); + const uint32_t num_q_heads = graph->size_at(-2, q_projected); + const uint32_t seq_len = graph->size_at(-3, q_projected); + + const uint32_t N4 = utils::div_up_4(head_dim); + const uint32_t M4 = utils::div_up_4(seq_len); + + return {N4, M4, num_q_heads}; +} - if (graph->is_buffer_storage(cache)) { - return graph->create_global_wg_size(projected); +utils::uvec3 pick_sdpa_compute_out_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + const bool use_coop_algorithm = + shader.kernel_name.find("_coop") != std::string::npos; + + if (use_coop_algorithm) { + return {1, 64, 1}; } else { - return graph->logical_limits_of(projected); + return pick_hw_square_wg_size( + graph, shader, global_workgroup_size, args, resize_args); } } -void add_kv_cache_update_node( +// +// Dispatch nodes +// + +void add_sdpa_kv_cache_update_node( ComputeGraph& graph, const ValueRef input_pos_symint, const ValueRef projected, const ValueRef cache) { - std::string kernel_name("kv_cache_update"); + std::string kernel_name("sdpa_kv_cache_update"); add_storage_type_suffix(kernel_name, graph.storage_type_of(projected)); add_dtype_suffix(kernel_name, graph.dtype_of(projected)); - vkapi::ParamsBindList param_ubos; - - if (graph.is_buffer_storage(cache)) { - param_ubos = { - graph.numel_ubo(projected), - graph.strides_ubo(cache), - graph.get_or_create_int_param_buffer(input_pos_symint)}; - } else { - param_ubos = { - graph.logical_limits_ubo(projected), - graph.get_or_create_int_param_buffer(input_pos_symint)}; - } + vkapi::ParamsBindList param_ubos = { + graph.sizes_ubo(cache), + graph.sizes_ubo(projected), + graph.get_or_create_int_param_buffer(input_pos_symint)}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -245,121 +304,113 @@ void add_kv_cache_update_node( // Specialization Constants {}, // Resize Args - {}, + {input_pos_symint}, // Resizing Logic nullptr)); } -utils::uvec3 attn_weight_scale_and_mask_global_wg_size( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - - const ValueRef attn_weight = args.at(0).refs.at(0); - - if (graph->is_buffer_storage(attn_weight)) { - return { - graph->size_at(-1, attn_weight), - graph->size_at(-2, attn_weight), - graph->size_at(-3, attn_weight), - }; - } else { - return graph->logical_limits_of(attn_weight); - } -} - -void add_attn_weight_scale_and_mask_node( +void add_sdpa_compute_attn_weights_node( ComputeGraph& graph, - const ValueRef input_pos_symint, const ValueRef q_projected, - const ValueRef attn_weight) { - std::string kernel_name("sdpa_attn_weight_scale_and_mask"); - add_storage_type_suffix(kernel_name, graph.storage_type_of(attn_weight)); - add_dtype_suffix(kernel_name, graph.dtype_of(attn_weight)); - + const ValueRef k_cache, + const ValueRef input_pos_symint, + const ValueRef attn_weights) { const int32_t head_dim_size = graph.size_at(-1, q_projected); const float scale_val = 1.0f / std::sqrt(static_cast(head_dim_size)); - vkapi::ParamsBindList param_ubos; - - if (graph.is_buffer_storage(attn_weight)) { - param_ubos = { - graph.sizes_ubo(attn_weight), - graph.strides_ubo(attn_weight), - graph.create_params_buffer(scale_val)}; - } else { - param_ubos = { - graph.logical_limits_ubo(attn_weight), - graph.get_or_create_int_param_buffer(input_pos_symint), - graph.create_params_buffer(scale_val)}; - } + vkapi::ParamsBindList param_ubos = { + graph.sizes_ubo(q_projected), + graph.sizes_ubo(k_cache), + graph.get_or_create_int_param_buffer(input_pos_symint)}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - VK_KERNEL_FROM_STR(kernel_name), - attn_weight_scale_and_mask_global_wg_size, - default_pick_local_wg_size, + pick_sdpa_compute_attn_weights_shader, + pick_sdpa_compute_attn_weights_global_wg_size, + pick_sdpa_compute_attn_weights_local_wg_size, // Inputs and Outputs - {{attn_weight, vkapi::kReadWrite}}, + {{attn_weights, vkapi::kWrite}, {{q_projected, k_cache}, vkapi::kRead}}, // Shader param buffers param_ubos, // Push Constants {}, // Specialization Constants - {}, + {scale_val}, // Resize Args - {}, + {input_pos_symint}, // Resizing Logic - nullptr)); + resize_compute_attn_weights_node)); } -std::vector get_cache_slice_sizes( +void add_sdpa_attn_weights_softmax_node( ComputeGraph& graph, - ValueRef cache, - ValueRef input_pos_symint, - ValueRef q_projected) { - std::vector slice_sizes = graph.sizes_of(cache); - - // Cache slicing will always be in the channels dim - const int32_t input_pos_val = graph.read_symint(input_pos_symint); - const int64_t q_seq_len = graph.size_at(1, q_projected); - slice_sizes.at(1) = input_pos_val + q_seq_len; - return slice_sizes; -} + const ValueRef attn_weights, + const ValueRef q_projected, + const ValueRef input_pos_symint, + const ValueRef attn_weights_softmax) { + std::string shader_name = "sdpa_attn_weights_softmax"; + add_storage_type_suffix( + shader_name, graph.storage_type_of(attn_weights_softmax)); + add_dtype_suffix(shader_name, graph.dtype_of(attn_weights_softmax)); -void resize_cache_slice_view_node( - ComputeGraph* graph, - const std::vector& args, - const std::vector& extra_args) { - (void)args; - std::vector slice_sizes = get_cache_slice_sizes( - *graph, extra_args[0], extra_args[1], extra_args[2]); + vkapi::ParamsBindList param_ubos = { + graph.sizes_ubo(q_projected), + graph.get_or_create_int_param_buffer(input_pos_symint)}; - graph->get_tensor(extra_args[3])->virtual_resize(slice_sizes); + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(shader_name), + pick_sdpa_attn_weights_softmax_global_wg_size, + pick_sdpa_attn_weights_softmax_local_wg_size, + // Inputs and Outputs + {{attn_weights_softmax, vkapi::kWrite}, {attn_weights, vkapi::kRead}}, + // Shader param buffers + param_ubos, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {q_projected, input_pos_symint}, + // Resizing Logic + resize_sdpa_attn_weights_softmax_node)); } -void add_cache_slice_view_node( +void add_sdpa_compute_out_node( ComputeGraph& graph, - ValueRef cache, - ValueRef input_pos_symint, - ValueRef q_projected, - ValueRef cache_sliced, - const int64_t max_seq_len) { - std::vector slice_sizes = - get_cache_slice_sizes(graph, cache, input_pos_symint, q_projected); - // Initialize the slice to the maximum possible size to start - slice_sizes.at(1) = max_seq_len; - - graph.get_tensor(cache_sliced)->virtual_resize(slice_sizes); - - graph.execute_nodes().emplace_back(new ExecuteNode( - resize_cache_slice_view_node, - {cache, input_pos_symint, q_projected, cache_sliced})); + const ValueRef attn_weights_softmax, + const ValueRef v_cache, + const ValueRef q_projected, + const ValueRef input_pos_symint, + const ValueRef out) { + vkapi::ParamsBindList param_ubos = { + graph.sizes_ubo(q_projected), + graph.sizes_ubo(v_cache), + graph.get_or_create_int_param_buffer(input_pos_symint)}; + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_sdpa_compute_out_shader, + pick_sdpa_compute_out_global_wg_size, + pick_sdpa_compute_out_local_wg_size, + // Inputs and Outputs + {{out, vkapi::kWrite}, {{attn_weights_softmax, v_cache}, vkapi::kRead}}, + // Shader param buffers + param_ubos, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {q_projected, input_pos_symint}, + // Resizing Logic + resize_sdpa_compute_out_node)); } +// +// High level operator impl +// + void update_cache_impl(ComputeGraph& graph, const std::vector& args) { int arg_idx = 0; const ValueRef value = args[arg_idx++]; @@ -377,7 +428,7 @@ void update_cache_impl(ComputeGraph& graph, const std::vector& args) { VK_CHECK_COND( graph.size_at(-2, value) == graph.size_at(-2, cache)); - add_kv_cache_update_node(graph, input_pos_symint, value, cache); + add_sdpa_kv_cache_update_node(graph, input_pos_symint, value, cache); } void sdpa_impl(ComputeGraph& graph, const std::vector& args) { @@ -418,105 +469,39 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { graph.val_is_none(is_causal) || graph.extract_scalar(is_causal)); VK_CHECK_COND(graph.val_is_none(attn_mask)); - const int32_t max_seq_len = graph.size_at(1, k_cache); + const int64_t num_q_heads = graph.size_at(-2, q_projected); + const int64_t max_seq_len = graph.size_at(-3, q_projected); - // Slice caches from 0 to input_pos + sequence_len - const ValueRef k_cache_sliced = graph.add_tensor_view(k_cache); - const ValueRef v_cache_sliced = graph.add_tensor_view(v_cache); - add_cache_slice_view_node( - graph, - k_cache, - input_pos_symint, - q_projected, - k_cache_sliced, - max_seq_len); - add_cache_slice_view_node( - graph, - v_cache, - input_pos_symint, - q_projected, - v_cache_sliced, - max_seq_len); - - // Scalar values for various dims - const ValueRef channels = graph.add_scalar(1); - const ValueRef height = graph.add_scalar(2); - const ValueRef width = graph.add_scalar(3); - - // Repeat interleave - const int64_t num_heads = graph.size_at(2, q_projected); - const int64_t num_kv_heads = graph.size_at(2, k_cache); - - const ValueRef num_repeats = - graph.add_scalar(num_heads / num_kv_heads); - - std::vector cache_slice_repeated_sizes(graph.sizes_of(q_projected)); - cache_slice_repeated_sizes.at(1) = max_seq_len; - - TmpTensor k_cache_sliced_repeated( - &graph, cache_slice_repeated_sizes, graph.dtype_of(k_cache_sliced)); - TmpTensor v_cache_sliced_repeated( - &graph, cache_slice_repeated_sizes, graph.dtype_of(v_cache_sliced)); - - add_repeat_interleave_node( - graph, k_cache_sliced, num_repeats, height, k_cache_sliced_repeated); - add_repeat_interleave_node( - graph, v_cache_sliced, num_repeats, height, v_cache_sliced_repeated); - - // Transpose sequence and head dims - const ValueRef q_transposed = graph.add_tensor_view(q_projected); - const ValueRef k_transposed = graph.add_tensor_view(k_cache_sliced_repeated); - const ValueRef v_transposed = graph.add_tensor_view(v_cache_sliced_repeated); - - add_transpose_view_node(graph, q_projected, channels, height, q_transposed); - add_transpose_view_node( - graph, k_cache_sliced_repeated, channels, height, k_transposed); - add_transpose_view_node( - graph, v_cache_sliced_repeated, channels, height, v_transposed); - - // Transpose K again to prepare for matmul - const ValueRef k_transposed_2 = graph.add_tensor_view(k_transposed); - add_transpose_view_node(graph, k_transposed, height, width, k_transposed_2); - - // Initialize attn_weight to the maximum possible size - std::vector attn_weight_full_sizes = graph.sizes_of(q_transposed); - attn_weight_full_sizes.at(2) = max_seq_len; - attn_weight_full_sizes.at(3) = max_seq_len; - TmpTensor attn_weight( - &graph, attn_weight_full_sizes, graph.dtype_of(q_transposed)); - - // Resize attn_weight to the correct dim - std::vector attn_weight_sizes = attn_weight_full_sizes; - attn_weight_sizes.at(2) = graph.size_at(2, q_transposed); - attn_weight_sizes.at(3) = graph.size_at(2, k_transposed); - graph.get_tensor(attn_weight)->virtual_resize(attn_weight_sizes); - - // Calculate attention weight, which is a matmul of Q and K - const ValueRef mat2_is_transposed = graph.add_scalar(false); - add_matmul_node( - graph, q_transposed, k_transposed_2, attn_weight, mat2_is_transposed); - - // Apply scale and mask to the attention weight - add_attn_weight_scale_and_mask_node( - graph, input_pos_symint, q_projected, attn_weight); - - TmpTensor attn_weight_softmax( - &graph, attn_weight_full_sizes, graph.dtype_of(q_transposed)); - graph.get_tensor(attn_weight_softmax)->virtual_resize(attn_weight_sizes); - add_softmax_node(graph, attn_weight, width, attn_weight_softmax, false); - - // Calculate final output - const ValueRef out_transposed = graph.add_tensor_view(out); - add_transpose_view_node(graph, out, channels, height, out_transposed); - add_matmul_node( - graph, - attn_weight_softmax, - v_transposed, - out_transposed, - mat2_is_transposed); + const int64_t max_context_len = graph.size_at(-3, k_cache); + + std::vector attn_weight_full_sizes = { + 1, // batch + num_q_heads, + max_seq_len, + max_context_len}; + + TmpTensor attn_weights( + &graph, + attn_weight_full_sizes, + graph.dtype_of(q_projected), + graph.storage_type_of(q_projected), + utils::kWidthPacked); + + TmpTensor attn_weights_softmax( + &graph, + attn_weight_full_sizes, + graph.dtype_of(q_projected), + graph.storage_type_of(q_projected), + utils::kWidthPacked); + + add_sdpa_compute_attn_weights_node( + graph, q_projected, k_cache, input_pos_symint, attn_weights); + + add_sdpa_attn_weights_softmax_node( + graph, attn_weights, q_projected, input_pos_symint, attn_weights_softmax); - graph.execute_nodes().emplace_back( - new ExecuteNode(resize_sdpa_out, {q_projected, out})); + add_sdpa_compute_out_node( + graph, attn_weights_softmax, v_cache, q_projected, input_pos_symint, out); } void sdpa_with_kv_cache_impl( @@ -540,10 +525,10 @@ void sdpa_with_kv_cache_impl( (void)sequence_len; - const ValueRef k_cache = - prepack_standard_like(graph, k_cache_data, q_projected); - const ValueRef v_cache = - prepack_standard_like(graph, v_cache_data, q_projected); + const ValueRef k_cache = prepack_standard( + graph, k_cache_data, utils::kTexture3D, utils::kWidthPacked); + const ValueRef v_cache = prepack_standard( + graph, v_cache_data, utils::kTexture3D, utils::kWidthPacked); update_cache_impl(graph, {k_projected, k_cache, input_pos_symint, -1}); update_cache_impl(graph, {v_projected, v_cache, input_pos_symint, -1}); @@ -565,7 +550,6 @@ REGISTER_OPERATORS { VK_REGISTER_OP(sdpa_with_kv_cache.default, sdpa_with_kv_cache_impl); VK_REGISTER_OP(update_cache.default, update_cache_impl); VK_REGISTER_OP(llama.custom_sdpa.default, sdpa_impl); - VK_REGISTER_OP(llama.flash_attention.default, flash_attention_impl); } } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp index e37ef66434b..5e645e29e3d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Softmax.cpp @@ -67,11 +67,11 @@ void resize_softmax_node( const std::vector& args, const std::vector& resize_args) { (void)resize_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - std::vector in_sizes = in->sizes(); - out->virtual_resize(in_sizes); + const std::vector in_sizes = graph->sizes_of(in); + graph->virtual_resize(out, in_sizes); } void add_softmax_node( diff --git a/backends/vulkan/runtime/graph/ops/impl/Split.cpp b/backends/vulkan/runtime/graph/ops/impl/Split.cpp index 8002dadc538..f87af08ee69 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Split.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Split.cpp @@ -23,23 +23,22 @@ void add_split_with_sizes_default_node( const std::vector& split_sizes, int64_t dim, ValueRef out_list_ref) { - vTensorPtr t_in = graph.get_tensor(in); + const ValueListPtr out_list = graph.get_value_list(out_list_ref); - ValueListPtr out_list = graph.get_value_list(out_list_ref); - - DimIndex dim_index = normalize_to_dim_index(*t_in, dim); + const int64_t input_ndim = graph.dim_of(in); + const DimIndex dim_index = dim < 0 ? static_cast(dim) + : static_cast(dim - input_ndim); VK_CHECK_COND(out_list->size() == split_sizes.size()); for (int split_idx = 0; split_idx < split_sizes.size(); split_idx++) { - int64_t split_size = split_sizes[split_idx]; - ValueRef out_ref = (*out_list)[split_idx]; + const int64_t split_size = split_sizes.at(split_idx); + const ValueRef out_ref = out_list->at(split_idx); - vTensorPtr t_out = graph.get_tensor(out_ref); - VK_CHECK_COND(dim_at(*t_out, dim_index) == split_size); + VK_CHECK_COND(dim_at(graph.sizes_of(out_ref), dim_index) == split_size); } - const auto packed_dim = t_in->packed_dim(); + const auto packed_dim = graph.packed_dim_of(in); const auto packed_dim_index = static_cast(kWidth4D - packed_dim); // Index of dimension to be concatenated in (w, h, c * b) coordinate system @@ -53,15 +52,14 @@ void add_split_with_sizes_default_node( // if splitting channels if (is_splitting_channel) { // set source offset w as channel size of the input tensor - src_offset[3] = dim_at(t_in->sizes(), kChannel4D); + src_offset[3] = dim_at(graph.sizes_of(in), kChannel4D); } for (ValueRef out_ref : *out_list) { // Doesn't need to use split_size since we have already verified that the // output tensor's size matches with the split_size. - vTensorPtr t_out = graph.get_tensor(out_ref); - const auto out_channel_size = dim_at(t_out->sizes(), kChannel4D); - utils::ivec3 range = t_out->logical_limits(); + const auto out_channel_size = dim_at(graph.sizes_of(out_ref), kChannel4D); + const utils::ivec3 range = graph.logical_limits_of(out_ref); if (dim_index == packed_dim_index) { // if splitting channels, use add_copy_channel_offset_node function as @@ -79,7 +77,8 @@ void add_split_with_sizes_default_node( dst_offset[3] = is_splitting_channel ? out_channel_size : 0; add_copy_packed_dim_offset_node( graph, in, range, src_offset, dst_offset, out_ref); - src_offset[dim_xyz_index] += dim_at(t_out->sizes(), packed_dim_index); + src_offset[dim_xyz_index] += + dim_at(graph.sizes_of(out_ref), packed_dim_index); } } else { // set destination offset w as channel size of the output tensor if @@ -117,13 +116,14 @@ void add_split_tensor_node( ValueRef split_size_ref, ValueRef dim_ref, ValueRef out) { - int64_t split_size = graph.extract_scalar(split_size_ref); - int64_t dim = graph.extract_scalar(dim_ref); - - vTensorPtr t_in = graph.get_tensor(in); - DimIndex dim_index = normalize_to_dim_index(*t_in, dim); - int64_t size = dim_at(*t_in, dim_index); - std::vector split_sizes(size / split_size, split_size); + const int64_t split_size = graph.extract_scalar(split_size_ref); + const int64_t dim = graph.extract_scalar(dim_ref); + + const int64_t input_ndim = graph.dim_of(in); + const DimIndex dim_index = dim < 0 ? static_cast(dim) + : static_cast(dim - input_ndim); + const int64_t size = dim_at(graph.sizes_of(in), dim_index); + const std::vector split_sizes(size / split_size, split_size); add_split_with_sizes_default_node(graph, in, split_sizes, dim, out); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp index 249f5e7fa6b..13801b45cc7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -55,8 +56,49 @@ void add_squeeze_copy_dims_node( } } +void resize_squeeze_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef dims_ref = extra_args.at(0); + + const IntListPtr dims = graph->get_int_list(dims_ref); + + std::vector out_sizes = graph->sizes_of(in); + + // Remove the dimensions specified in dims if their size is 1 + for (int64_t dim : *dims) { + if (dim >= 0 && dim < static_cast(out_sizes.size()) && + out_sizes[dim] == 1) { + out_sizes.erase(out_sizes.begin() + dim); + // After erasing, all subsequent dims shift left by one + // So we need to decrement all subsequent dims in dims + for (auto& d : *dims) { + if (d > dim) { + --d; + } + } + } + } + + graph->virtual_resize(out, out_sizes); +} + void squeeze_copy_dims(ComputeGraph& graph, const std::vector& args) { - return add_squeeze_copy_dims_node(graph, args[0], args[1], args[2]); + int idx = 0; + const ValueRef in = args.at(idx++); + const ValueRef dims = args.at(idx++); + const ValueRef out = args.at(idx++); + + std::vector resize_args = {dims}; + + if (graph.is_buffer_storage(in)) { + return add_view_copy_buffer_node( + graph, in, out, resize_args, resize_squeeze_node); + } + return add_squeeze_copy_dims_node(graph, in, dims, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp index bfaad716059..40de9b59e81 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp @@ -27,15 +27,18 @@ void add_staging_to_tensor_node( VK_CHECK_COND(graph.val_is_staging(in_staging)); vkapi::ShaderInfo shader = get_nchw_to_tensor_shader( - *graph.get_tensor(out_tensor), graph.int8_buffers_enabled()); + graph, + out_tensor, + graph.dtype_of(in_staging), + graph.int8_buffers_enabled()); - std::vector pcs; + vkapi::ParamsBindList param_buffers = {}; if (graph.is_buffer_storage(out_tensor)) { - pcs = { - graph.sizes_pc_of(out_tensor), - graph.strides_pc_of(out_tensor), - graph.numel_pc_of(out_tensor)}; - } else { + param_buffers.append(graph.buffer_meta_ubo(out_tensor)); + } + + std::vector pcs; + if (graph.is_texture_storage(out_tensor)) { pcs = {graph.sizes_pc_of(out_tensor)}; } @@ -47,7 +50,7 @@ void add_staging_to_tensor_node( // Input and Outputs {{out_tensor, vkapi::kWrite}, {in_staging, vkapi::kRead}}, // Parameter Buffers - {}, + param_buffers, // Push Constants pcs, // Specialization Constants @@ -66,16 +69,6 @@ bool is_bitw8_shader(const vkapi::ShaderInfo& shader) { return shader_prefix_str == kBitw8PrefixStr; } -vkapi::ShaderInfo get_tensor_to_staging_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - const ValueRef in_tensor = args.at(1).refs.at(0); - return get_tensor_to_nchw_shader( - *graph->get_tensor(in_tensor), graph->int8_buffers_enabled()); -} - utils::uvec3 tensor_to_staging_global_wg_size( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -111,15 +104,18 @@ void add_tensor_to_staging_node( VK_CHECK_COND(graph.val_is_staging(out_staging)); vkapi::ShaderInfo shader = get_tensor_to_nchw_shader( - *graph.get_tensor(in_tensor), graph.int8_buffers_enabled()); + graph, + in_tensor, + graph.dtype_of(out_staging), + graph.int8_buffers_enabled()); - std::vector pcs; + vkapi::ParamsBindList param_buffers = {}; if (graph.is_buffer_storage(in_tensor)) { - pcs = { - graph.sizes_pc_of(in_tensor), - graph.strides_pc_of(in_tensor), - graph.numel_pc_of(in_tensor)}; - } else { + param_buffers.append(graph.buffer_meta_ubo(in_tensor)); + } + + std::vector pcs; + if (graph.is_texture_storage(in_tensor)) { pcs = {graph.sizes_pc_of(in_tensor)}; } @@ -135,7 +131,7 @@ void add_tensor_to_staging_node( // Input and Outputs {{out_staging, vkapi::kWrite}, {in_tensor, vkapi::kRead}}, // Parameter Buffers - {}, + param_buffers, // Push Constants pcs, // Specialization Constants @@ -152,7 +148,12 @@ void add_prepack_standard_node( const ValueRef tensor, const bool transpose_hw = false) { vkapi::ShaderInfo shader = get_nchw_to_tensor_shader( - *graph.get_tensor(tensor), graph.int8_buffers_enabled()); + graph, tensor, graph.dtype_of(tensor_data), graph.int8_buffers_enabled()); + + vkapi::ParamsBindList param_buffers = {}; + if (graph.is_buffer_storage(tensor)) { + param_buffers.append(graph.buffer_meta_ubo(tensor)); + } std::vector pcs; if (graph.is_buffer_storage(tensor)) { @@ -175,7 +176,7 @@ void add_prepack_standard_node( tensor_data, tensor, // Parameter Buffers - {}, + param_buffers, // Specialization Constants {graph.hashed_layout_of(tensor), transpose_hw_spec}, pcs)); @@ -322,75 +323,6 @@ ValueRef prepack_int4_linear_weight_transposed_interleaved( return qmat2; } -ValueRef prepack_int4_linear_weight_transposed_block_4x8( - ComputeGraph& graph, - const ValueRef qmat2_data) { - std::vector qmat2_orig_sizes = graph.sizes_of(qmat2_data); - const int64_t ndim = graph.dim_of(qmat2_data); - - const int64_t K_div2 = qmat2_orig_sizes.at(ndim - 1); // Input is [N, K/2] - const int64_t N = qmat2_orig_sizes.at(ndim - 2); - // Logical K dimension. Each value in the tensor is a uint8 that contains 2 - // packed 4-bit values. - const int64_t K = K_div2 * 2; - - // This packing format partitions the weight tensor into 4 wide x 8 high - // blocks. To figure out the size of the output tensor, determine the number - // of blocks along the width and height dims. - const int64_t num_blocks_K = utils::div_up(K, int64_t(4)); - const int64_t num_blocks_N = utils::div_up(N, int64_t(8)); - // Each transposed block is 8 wide x 4 high. In terms of 8-bit values, the - // block is 4 wide x 4 high. To maximize memory loading efficiency, the packed - // weight tensor will use a base data type of uint32_t; in terms of uint32_t, - // each block is 1 wide x 4 high. However, each block is also flattened as it - // is stored, so that the whole block can be loaded at once. As a result, the - // stored block will be 4 wide x 1 high. - const int64_t output_width = num_blocks_K * 4; - const int64_t output_height = num_blocks_N; - - // Store the original sizes of the tensor to pass to the shader - utils::ivec2 orig_sizes{ - utils::safe_downcast(K), utils::safe_downcast(N)}; - - std::vector qmat2_sizes{output_height, output_width}; - - utils::StorageType storage_type = utils::kTexture2D; - 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; - } - - ValueRef qmat2 = graph.add_tensor( - qmat2_sizes, vkcompute::vkapi::kUInt, storage_type, utils::kWidthPacked); - - // Global workgroup size: each thread writes out two adjacent blocks - utils::uvec3 global_wg_size{ - utils::div_up(utils::safe_downcast(num_blocks_K), uint32_t(2)), - utils::safe_downcast(num_blocks_N), - 1u}; - - std::string kernel_name = "pack_int4_linear_weight_transposed_block_4x8"; - add_storage_type_suffix(kernel_name, storage_type); - - graph.prepack_nodes().emplace_back(new PrepackNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - graph.create_local_wg_size(global_wg_size), - // Inputs and Outputs - qmat2_data, - qmat2, - // UBOs - {}, - // Specialization Constants - {}, - // Push Constants - {graph.sizes_pc_of(qmat2), - PushConstantDataInfo(&orig_sizes, sizeof(utils::ivec2))})); - - return qmat2; -} - void prepack_op(ComputeGraph& graph, const std::vector& args) { return add_prepack_standard_node(graph, args[0], args[1]); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Staging.h b/backends/vulkan/runtime/graph/ops/impl/Staging.h index 0b1568ca139..5f5cdd1eda0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Staging.h +++ b/backends/vulkan/runtime/graph/ops/impl/Staging.h @@ -95,8 +95,4 @@ ValueRef prepack_int4_linear_weight_transposed_interleaved( ComputeGraph& graph, const ValueRef qmat2_data); -ValueRef prepack_int4_linear_weight_transposed_block_4x8( - ComputeGraph& graph, - const ValueRef qmat2_data); - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Tan.cpp b/backends/vulkan/runtime/graph/ops/impl/Tan.cpp index 89c4a4d408f..687b3923354 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Tan.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Tan.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -20,10 +21,11 @@ void resize_tan_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); - out->virtual_resize(self->sizes()); + const std::vector self_sizes = graph->sizes_of(self); + graph->virtual_resize(out, self_sizes); } void add_tan_node(ComputeGraph& graph, const ValueRef in, const ValueRef out) { @@ -34,11 +36,11 @@ void add_tan_node(ComputeGraph& graph, const ValueRef in, const ValueRef out) { vkapi::ParamsBindList ubos({}); ubos.append({graph.logical_limits_ubo(out)}); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/ToCopy.cpp b/backends/vulkan/runtime/graph/ops/impl/ToCopy.cpp index d1145a925d4..b7e0218823a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ToCopy.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ToCopy.cpp @@ -19,10 +19,10 @@ void resize_to_copy_op_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); - out->virtual_resize(self->sizes()); + graph->virtual_resize(out, graph->sizes_of(self)); } void add_to_copy_node(ComputeGraph& graph, ValueRef in, ValueRef out) { diff --git a/backends/vulkan/runtime/graph/ops/impl/Transpose.cpp b/backends/vulkan/runtime/graph/ops/impl/Transpose.cpp index 8501d085bc8..b797536d817 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Transpose.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Transpose.cpp @@ -23,16 +23,16 @@ void resize_transpose_view_node( const std::vector& args, const std::vector& extra_args) { (void)args; - vTensorPtr out = graph->get_tensor(extra_args[0]); - vTensorPtr in = graph->get_tensor(extra_args[1]); + const ValueRef out = extra_args.at(0); + const ValueRef in = extra_args.at(1); - const int64_t dim0 = graph->extract_scalar(extra_args[2]); - const int64_t dim1 = graph->extract_scalar(extra_args[3]); + const int64_t dim0 = graph->extract_scalar(extra_args.at(2)); + const int64_t dim1 = graph->extract_scalar(extra_args.at(3)); - std::vector new_sizes = in->sizes(); + std::vector new_sizes = graph->sizes_of(in); // Transpose the resized input sizes std::iter_swap(new_sizes.begin() + dim0, new_sizes.begin() + dim1); - out->virtual_resize(new_sizes); + graph->virtual_resize(out, new_sizes); } void check_transpose_view_args( @@ -62,9 +62,8 @@ void add_transpose_view_node( const int64_t dim1 = graph.extract_scalar(dim1_ref); check_transpose_view_args(graph, input_ref, dim0, dim1, out_ref); - const vTensorPtr in = graph.get_tensor(input_ref); - graph.get_tensor(out_ref)->virtual_clone(*in); - graph.get_tensor(out_ref)->virtual_transpose(dim0, dim1); + graph.virtual_clone(out_ref, input_ref); + graph.virtual_transpose(out_ref, dim0, dim1); graph.execute_nodes().emplace_back(new ExecuteNode( resize_transpose_view_node, {out_ref, input_ref, dim0_ref, dim1_ref})); diff --git a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp index 085e8559980..9830a8e8784 100644 --- a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp @@ -26,10 +26,11 @@ void resize_unary_op_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); - out->virtual_resize(self->sizes()); + const std::vector self_sizes = graph->sizes_of(self); + graph->virtual_resize(out, self_sizes); } void add_unary_op_node( diff --git a/backends/vulkan/runtime/graph/ops/impl/Unsqueeze.cpp b/backends/vulkan/runtime/graph/ops/impl/Unsqueeze.cpp index c4de5d88f30..0a98f6d8f43 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Unsqueeze.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Unsqueeze.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -45,8 +46,42 @@ void add_unsqueeze_node( add_permute_node(graph, in, permute_dims_ref, out); } +void resize_unsqueeze_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& extra_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef dims_ref = extra_args.at(0); + + const IntListPtr dims = graph->get_int_list(dims_ref); + + std::vector out_sizes = graph->sizes_of(in); + + // Insert singleton dimensions at the specified positions + for (auto dim : *dims) { + int64_t d = dim; + if (d < 0) { + d += static_cast(out_sizes.size()) + 1; + } + out_sizes.insert(out_sizes.begin() + d, 1); + } + + graph->virtual_resize(out, out_sizes); +} + void unsqueeze(ComputeGraph& graph, const std::vector& args) { - return add_unsqueeze_node(graph, args[0], args[1], args[2]); + int idx = 0; + const ValueRef in = args.at(idx++); + const ValueRef dims = args.at(idx++); + const ValueRef out = args.at(idx++); + + std::vector resize_args = {dims}; + if (graph.is_buffer_storage(in)) { + return add_view_copy_buffer_node( + graph, in, out, resize_args, resize_unsqueeze_node); + } + return add_unsqueeze_node(graph, in, dims, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp index d098ed94c7f..6662ae367c5 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Upsample.cpp @@ -8,6 +8,7 @@ #include +#include #include #include @@ -22,12 +23,12 @@ void resize_upsample_nearest2d_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr self = graph->get_tensor(args[1].refs[0]); - std::vector out_sizes = self->sizes(); // NCHW + const ValueRef out = args.at(0).refs.at(0); + const ValueRef self = args.at(1).refs.at(0); + std::vector out_sizes = graph->sizes_of(self); // NCHW - const ValueRef output_sizes = extra_args[0]; // HW - const ValueRef scale_factors = extra_args[1]; // HW + const ValueRef output_sizes = extra_args.at(0); // HW + const ValueRef scale_factors = extra_args.at(1); // HW if (!graph->val_is_none(output_sizes)) { IntListPtr output_size_ref = graph->get_int_list(output_sizes); out_sizes.at(2) = output_size_ref->at(0); @@ -38,7 +39,7 @@ void resize_upsample_nearest2d_node( out_sizes.at(3) *= scales->at(1); } - out->virtual_resize(out_sizes); + graph->virtual_resize(out, out_sizes); } void add_upsample_nearest2d_node( @@ -114,11 +115,11 @@ void add_upsample_nearest2d_node( } add_dtype_suffix(kernel_name, graph.dtype_of(out)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - graph.create_global_wg_size(out), - graph.create_local_wg_size(out), + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::MemoryAccessType::WRITE}, {in, vkapi::MemoryAccessType::READ}}, diff --git a/backends/vulkan/runtime/graph/ops/impl/Var.cpp b/backends/vulkan/runtime/graph/ops/impl/Var.cpp index 41fdc41e982..d8fd367f18a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Var.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Var.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include @@ -14,21 +15,109 @@ namespace vkcompute { using namespace utils; +// Custom global workgroup size function for var_buffer +utils::uvec3 var_buffer_global_wg_size( + 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->size_at(-1, out), + graph->size_at(-2, out), + graph->size_at(-3, out) * graph->size_at(-4, out)}; +} + +// Custom local workgroup size function for var_buffer +utils::uvec3 var_buffer_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)global_workgroup_size; + const ValueRef in = args.at(1).refs.at(0); + const int dim = resize_args.at(0); + + const int64_t ndim = graph->dim_of(in); + int32_t reduce_dim = normalize(dim, ndim); + reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); + + const uint32_t nworkers_per_group = 4; + utils::uvec3 local_wg_size{1, 1, 1}; + local_wg_size[reduce_dim] = nworkers_per_group; + return local_wg_size; +} + +// Custom global workgroup size function for var_texture +utils::uvec3 var_texture_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const int dim = resize_args.at(0); + + const int64_t ndim = graph->dim_of(in); + int32_t reduce_dim = normalize(dim, ndim); + reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); + + utils::uvec3 global_wg_size = graph->logical_limits_of(out); + global_wg_size[reduce_dim] = 1; + return global_wg_size; +} + +// Custom local workgroup size function for var_texture +utils::uvec3 var_texture_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + const ValueRef in = args.at(1).refs.at(0); + const int dim = resize_args.at(0); + + const int64_t ndim = graph->dim_of(in); + int32_t reduce_dim = normalize(dim, ndim); + reduce_dim = nchw_dim_to_whcn_dim(reduce_dim, ndim); + + const uint32_t nworkers_per_group = 4; + const uint32_t ngroups = 4; + + utils::uvec3 local_wg_size{1, 1, 1}; + local_wg_size[reduce_dim] = nworkers_per_group; + const int other_dim_1 = (reduce_dim + 1) % 3; + const int other_dim_2 = (reduce_dim + 2) % 3; + if (global_workgroup_size[other_dim_1] > global_workgroup_size[other_dim_2]) { + local_wg_size[other_dim_1] = ngroups; + } else { + local_wg_size[other_dim_2] = ngroups; + } + return local_wg_size; +} + void resize_var_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - int dim = extra_args[0]; + const int dim = extra_args.at(0); - std::vector new_sizes = in->sizes(); + std::vector new_sizes = graph->sizes_of(in); if (!new_sizes.empty()) { new_sizes.at(normalize(dim, new_sizes.size())) = 1; } - out->virtual_resize(new_sizes); + + graph->virtual_resize(out, new_sizes); } void add_var_buffer_node( @@ -67,11 +156,11 @@ void add_var_buffer_node( int32_t unbiased_int = static_cast(unbiased); push_constants.emplace_back(&unbiased_int, sizeof(unbiased_int)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - local_wg_size, + var_buffer_global_wg_size, + var_buffer_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers @@ -142,12 +231,11 @@ void add_var_texture_node( int32_t unbiased_int = static_cast(unbiased); push_constants.emplace_back(&unbiased_int, sizeof(unbiased_int)); - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - // shader_descriptor, VK_KERNEL_FROM_STR(kernel_name), - global_wg_size, - local_wg_size, + var_texture_global_wg_size, + var_texture_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, // Shader params buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/View.cpp b/backends/vulkan/runtime/graph/ops/impl/View.cpp index 9dbe79faebb..8701a6246b0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/View.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/View.cpp @@ -44,15 +44,19 @@ void resize_view_node( ComputeGraph* graph, const std::vector& args, const std::vector& extra_args) { - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); - if (extra_args[0] == kDummyValueRef || graph->val_is_none(extra_args[0])) { - out->virtual_resize(in->sizes()); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + if (extra_args.at(0) == kDummyValueRef || + graph->val_is_none(extra_args.at(0))) { + const std::vector in_sizes = graph->sizes_of(in); + graph->virtual_resize(out, in_sizes); } else { std::vector view_sizes = - graph->extract_int_or_symint_list(extra_args[0]); - std::vector out_sizes = compute_out_sizes(in->sizes(), view_sizes); - out->virtual_resize(out_sizes); + graph->extract_int_or_symint_list(extra_args.at(0)); + const std::vector in_sizes = graph->sizes_of(in); + const std::vector out_sizes = + compute_out_sizes(in_sizes, view_sizes); + graph->virtual_resize(out, out_sizes); } } @@ -61,12 +65,9 @@ void add_view_node( ValueRef in, ValueRef sizes, ValueRef out) { - vTensorPtr t_in = graph.get_tensor(in); - vTensorPtr t_out = graph.get_tensor(out); - std::string kernel_name = "view"; kernel_name.reserve(kShaderNameReserve); - add_dtype_suffix(kernel_name, *t_out); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -81,15 +82,54 @@ void add_view_node( // Push Constants {{graph.sizes_pc_of(out), graph.sizes_pc_of(in)}}, // Specialization Constants - {SV(t_in->packed_dim()), SV(t_out->packed_dim())}, + {graph.packed_dim_of(in), graph.packed_dim_of(out)}, // Resize Args {sizes}, // Resizing Logic resize_view_node)); } +void add_view_copy_buffer_node( + ComputeGraph& graph, + ValueRef in, + ValueRef out, + const std::vector& resize_args, + const ExecuteNode::ResizeFunction& resize_fn) { + std::string kernel_name = "view_buffer"; + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_global_wg_size, + default_pick_local_wg_size, + // Inputs and Outputs + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + // Parameter Buffers + {graph.buffer_meta_ubo(out), graph.buffer_meta_ubo(in)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + resize_args, + // Resizing Logic + resize_fn)); +} + void view(ComputeGraph& graph, const std::vector& args) { - return add_view_node(graph, args[0], args[1], args[2]); + int idx = 0; + const ValueRef in = args.at(idx++); + const ValueRef sizes = args.at(idx++); + const ValueRef out = args.at(idx++); + + std::vector resize_args = {sizes}; + + if (graph.is_buffer_storage(out)) { + return add_view_copy_buffer_node( + graph, in, out, resize_args, resize_view_node); + } + return add_view_node(graph, in, sizes, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/View.h b/backends/vulkan/runtime/graph/ops/impl/View.h index a2038d184c3..7a7a8d57742 100644 --- a/backends/vulkan/runtime/graph/ops/impl/View.h +++ b/backends/vulkan/runtime/graph/ops/impl/View.h @@ -12,6 +12,18 @@ namespace vkcompute { +/* + * Dispatches the view_copy compute shader. This can be used to implement ops + * that preserve the "contiguous" indexes of elements between the input and + * output such as view_copy, squeeze_copy, unsqueeze_copy, etc. + */ +void add_view_copy_buffer_node( + ComputeGraph& graph, + ValueRef in, + ValueRef out, + const std::vector& resize_args, + const ExecuteNode::ResizeFunction& resize_fn); + void add_view_node( ComputeGraph& graph, ValueRef in, diff --git a/backends/vulkan/runtime/graph/ops/impl/Where.cpp b/backends/vulkan/runtime/graph/ops/impl/Where.cpp index ea610b1fe74..c1c482d9967 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Where.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Where.cpp @@ -10,6 +10,7 @@ #include +#include #include namespace vkcompute { @@ -19,11 +20,11 @@ void resize_where_node( const std::vector& args, const std::vector& extra_args) { (void)extra_args; - vTensorPtr out = graph->get_tensor(args[0].refs[0]); - vTensorPtr in = graph->get_tensor(args[1].refs[0]); + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); - std::vector in_sizes = in->sizes(); - out->virtual_resize(in_sizes); + const std::vector in_sizes = graph->sizes_of(in); + graph->virtual_resize(out, in_sizes); } void add_where_texture_node( @@ -37,16 +38,11 @@ void add_where_texture_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); - const utils::uvec3 global_wg_size = graph.create_global_wg_size(out); - const utils::uvec3 local_wg_size = graph.create_local_wg_size(global_wg_size); - - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - // Shader VK_KERNEL_FROM_STR(kernel_name), - // Workgroup sizes - global_wg_size, - local_wg_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{cond, self, other}, vkapi::kRead}}, // Parameter buffers @@ -72,9 +68,6 @@ void add_where_buffer_node( add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); - const utils::uvec3 global_wg_size = graph.create_global_wg_size(out); - const utils::uvec3 local_wg_size = graph.create_local_wg_size(global_wg_size); - vkapi::ParamsBindList ubos = { graph.numel_ubo(out), graph.strides_ubo(out), @@ -82,13 +75,11 @@ void add_where_buffer_node( graph.strides_ubo(self), graph.strides_ubo(other)}; - graph.execute_nodes().emplace_back(new DispatchNode( + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - // Shader VK_KERNEL_FROM_STR(kernel_name), - // Workgroup sizes - global_wg_size, - local_wg_size, + default_pick_global_wg_size, + default_pick_local_wg_size, // Inputs and Outputs {{out, vkapi::kWrite}, {{cond, self, other}, vkapi::kRead}}, // Parameter buffers diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h index 4bd8e9b900b..5ed07dece38 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h +++ b/backends/vulkan/runtime/graph/ops/impl/utils/DimUtils.h @@ -31,11 +31,6 @@ constexpr DimIndex kHeight4D = DimIndex::DIM_2ND_LAST; constexpr DimIndex kChannel4D = DimIndex::DIM_3RD_LAST; constexpr DimIndex kBatch4D = DimIndex::DIM_4TH_LAST; -inline DimIndex normalize_to_dim_index(const api::vTensor& v_in, int32_t dim) { - return dim < 0 ? static_cast(dim) - : static_cast(dim - v_in.dim()); -} - /* * Semantic dimension names for a 1D tensor */ @@ -83,15 +78,6 @@ int32_t dim_at(const std::vector& sizes) { return dim_at(sizes, DI); } -template -int32_t dim_at(const api::vTensor& v_in) { - return dim_at(v_in.sizes(), DI); -} - -inline int32_t dim_at(const api::vTensor& v_in, DimIndex dim_index) { - return dim_at(v_in.sizes(), dim_index); -} - inline std::ostream& operator<<(std::ostream& os, DimIndex dim_index) { switch (dim_index) { case kWidth4D: diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/QuantizationConfig.h b/backends/vulkan/runtime/graph/ops/impl/utils/QuantizationConfig.h new file mode 100644 index 00000000000..4bc8c7c3bfc --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/utils/QuantizationConfig.h @@ -0,0 +1,58 @@ +/* + * 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 { + +enum class QuantizationGranularity { + PerChannel, + PerTensor, + PerGroup, + NoQuantization, +}; + +static constexpr QuantizationGranularity kPerChannel = + QuantizationGranularity::PerChannel; +static constexpr QuantizationGranularity kPerTensor = + QuantizationGranularity::PerTensor; +static constexpr QuantizationGranularity kPerGroup = + QuantizationGranularity::PerGroup; +static constexpr QuantizationGranularity kNoQuantization = + QuantizationGranularity::NoQuantization; + +struct QuantizationConfig { + int nbits; + QuantizationGranularity granularity; + std::vector granularity_sizes; + bool is_symmetric; + bool is_dynamic; + + QuantizationConfig() + : nbits(8), + granularity(kPerTensor), + granularity_sizes(), + is_symmetric(true), + is_dynamic(false) {} + + QuantizationConfig( + int nbits_, + QuantizationGranularity granularity_, + const std::vector& granularity_sizes_, + bool is_symmetric_ = true, + bool is_dynamic_ = false) + : nbits(nbits_), + granularity(granularity_), + granularity_sizes(granularity_sizes_), + is_symmetric(is_symmetric_), + is_dynamic(is_dynamic_) {} +}; + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h index 8e10c4e2bfa..270bdd1cd6b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h +++ b/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h @@ -28,4 +28,22 @@ T extract_scalar(const Value& value) { VK_THROW("Cannot extract scalar from Value with type ", value.type()); } +// Helper function to get default quant_min and quant_max based on dtype +// This matches the logic in _get_and_check_qmin_qmax from quant_primitives.py +inline std::pair get_dtype_bounds(vkapi::ScalarType dtype) { + switch (dtype) { + case vkapi::kByte: // uint8 + return {0, 255}; + case vkapi::kChar: // int8 + return {-128, 127}; + case vkapi::kShort: // int16 + return {-(1 << 15), (1 << 15) - 1}; + case vkapi::kInt: // int32 + return {-(1LL << 31), (1LL << 31) - 1}; + default: + // For unsupported types, throw an error instead of assuming int8 + VK_THROW("Unsupported dtype for quantization bounds: ", dtype); + } +} + } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp index 2bcf2a3842f..a52572289a4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp @@ -15,15 +15,14 @@ namespace vkcompute { // std::vector calculate_broadcasted_output_size( - const api::vTensor& t1, - const api::vTensor& t2) { - std::vector out_sizes( - std::max(t1.sizes().size(), t2.sizes().size())); + const std::vector& sizes1, + const std::vector& sizes2) { + std::vector out_sizes(std::max(sizes1.size(), sizes2.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) = - std::max(utils::val_at(i, t1.sizes()), utils::val_at(i, t2.sizes())); + std::max(utils::val_at(i, sizes1), utils::val_at(i, sizes2)); } return out_sizes; @@ -33,30 +32,6 @@ std::vector calculate_broadcasted_output_size( // Tensor property checking functions // -bool check_ndim_is(const api::vTensor& t, size_t ndim) { - return t.sizes().size() == ndim; -} - -bool check_same_sizes_at( - const api::vTensor& t1, - const int64_t d1, - const api::vTensor& t2, - const int64_t d2) { - return utils::val_at(d1, t1.sizes()) == utils::val_at(d2, t2.sizes()); -} - -bool check_packed_dim_is(const api::vTensor& t, const int32_t packed_dim) { - return t.packed_dim() == packed_dim; -} - -bool check_same_ndim(const api::vTensor& t1, const api::vTensor& t2) { - return t1.sizes().size() == t2.sizes().size(); -} - -bool check_same_packed_dim(const api::vTensor& t1, const api::vTensor& t2) { - return t1.packed_dim() == t2.packed_dim(); -} - bool check_same_packed_dim( ComputeGraph& graph, const ValueRef in, @@ -64,42 +39,38 @@ bool check_same_packed_dim( return graph.packed_dim_of(in) == graph.packed_dim_of(out); } -bool check_same_packed_dim( - const api::vTensor& t1, - const api::vTensor& t2, - const api::vTensor& t3) { - if (t1.packed_dim() != t2.packed_dim()) { - return false; - } - return (t1.packed_dim() == t3.packed_dim()); -} - // // Broadcast flag functions // bool is_packed_dim_broadcasted( - const api::vTensor& sndr, - const api::vTensor& rcvr) { + ComputeGraph& graph, + const ValueRef sndr, + const ValueRef rcvr) { // We assume that the tensors are broadcastable. If values aren't equal at // some index, then the value of rcvr is 1 and hence should be broadcasted. - switch (sndr.packed_dim()) { + const std::vector sndr_sizes = graph.sizes_of(sndr); + const std::vector rcvr_sizes = graph.sizes_of(rcvr); + + switch (graph.packed_dim_of(sndr)) { case WHCN::kChannelsDim: - return utils::val_at(-3, sndr.sizes()) > utils::val_at(-3, rcvr.sizes()); + return utils::val_at(-3, sndr_sizes) > utils::val_at(-3, rcvr_sizes); case WHCN::kHeightDim: - return utils::val_at(-2, sndr.sizes()) > utils::val_at(-2, rcvr.sizes()); + return utils::val_at(-2, sndr_sizes) > utils::val_at(-2, rcvr_sizes); case WHCN::kWidthDim: - return utils::val_at(-1, sndr.sizes()) > utils::val_at(-1, rcvr.sizes()); + return utils::val_at(-1, sndr_sizes) > utils::val_at(-1, rcvr_sizes); default: VK_THROW("Invalid packed dim"); } } utils::ivec2 create_broadcast_params( - const api::vTensor& t1, - const api::vTensor& t2) { + ComputeGraph& graph, + const ValueRef t1, + const ValueRef t2) { return utils::make_ivec2( - {is_packed_dim_broadcasted(t2, t1), is_packed_dim_broadcasted(t1, t2)}); + {is_packed_dim_broadcasted(graph, t2, t1), + is_packed_dim_broadcasted(graph, t1, t2)}); } // diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h index 3b61083069e..b62bf661995 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h @@ -18,44 +18,31 @@ namespace vkcompute { // std::vector calculate_broadcasted_output_size( - const api::vTensor& t1, - const api::vTensor& t2); + const std::vector& sizes1, + const std::vector& sizes2); // // Tensor property checking functions // -bool check_ndim_is(const api::vTensor& t, size_t ndim); - -bool check_same_ndim(const api::vTensor& t1, const api::vTensor& t2); - -bool check_same_sizes_at( - const api::vTensor& t1, - int64_t d1, - const api::vTensor& t2, - int64_t d2); - -bool check_packed_dim_is(const api::vTensor& t, const int32_t packed_dim); - -bool check_same_packed_dim(const api::vTensor& t1, const api::vTensor& t2); - bool check_same_packed_dim( ComputeGraph& graph, const ValueRef in, const ValueRef out); -bool check_same_packed_dim( - const api::vTensor& t1, - const api::vTensor& t2, - const api::vTensor& t3); - // // Broadcast flag functions // +bool is_packed_dim_broadcasted( + ComputeGraph& graph, + const ValueRef sndr, + const ValueRef rcvr); + utils::ivec2 create_broadcast_params( - const api::vTensor& t1, - const api::vTensor& t2); + ComputeGraph& graph, + const ValueRef t1, + const ValueRef t2); // // Work group size calculation functions diff --git a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp index b3a72e27c43..e829f355fe2 100644 --- a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.cpp @@ -10,23 +10,6 @@ namespace vkcompute { -void bind_tensor_to_descriptor_set( - api::vTensor& tensor, - vkapi::PipelineBarrier& pipeline_barrier, - const vkapi::MemoryAccessFlags accessType, - vkapi::DescriptorSet& descriptor_set, - const uint32_t idx) { - if (tensor.buffer()) { - vkapi::VulkanBuffer& buffer = tensor.buffer( - pipeline_barrier, vkapi::PipelineStage::COMPUTE, accessType); - descriptor_set.bind(idx, buffer); - } else { - vkapi::VulkanImage& image = tensor.image( - pipeline_barrier, vkapi::PipelineStage::COMPUTE, accessType); - descriptor_set.bind(idx, image); - } -} - uint32_t bind_values_to_descriptor_set( ComputeGraph* graph, const std::vector& args, @@ -36,19 +19,8 @@ uint32_t bind_values_to_descriptor_set( uint32_t idx = base_idx; for (auto& arg : args) { for (auto& ref : arg.refs) { - if (graph->val_is_tensor(ref)) { - bind_tensor_to_descriptor_set( - *(graph->get_tensor(ref)), - pipeline_barrier, - arg.access, - descriptor_set, - idx++); - } else if (graph->val_is_staging(ref)) { - bind_staging_to_descriptor_set( - *(graph->get_staging(ref)), descriptor_set, idx++); - } else { - VK_THROW("Unsupported type: ", graph->get_val_type(ref)); - } + graph->bind_value_to_descriptor_set( + ref, pipeline_barrier, arg.access, descriptor_set, idx++); } } return idx; diff --git a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h index 671a18f7e91..307bec154f3 100644 --- a/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h +++ b/backends/vulkan/runtime/graph/ops/utils/BindingUtils.h @@ -16,13 +16,6 @@ namespace vkcompute { // For objects in the graph // -void bind_tensor_to_descriptor_set( - api::vTensor& tensor, - vkapi::PipelineBarrier& pipeline_barrier, - const vkapi::MemoryAccessFlags accessType, - vkapi::DescriptorSet& descriptor_set, - const uint32_t idx); - uint32_t bind_values_to_descriptor_set( ComputeGraph* graph, const std::vector& args, diff --git a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp index 6388a8ad091..231e6d0c7f6 100644 --- a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.cpp @@ -26,12 +26,6 @@ void add_storage_type_suffix( } } -void add_storage_type_suffix( - std::string& kernel_name, - const api::vTensor& tensor) { - return add_storage_type_suffix(kernel_name, tensor.storage_type()); -} - void add_dtype_suffix(std::string& kernel_name, const vkapi::ScalarType dtype) { switch (dtype) { case vkapi::kDouble: @@ -75,23 +69,6 @@ void add_dtype_suffix(std::string& kernel_name, const vkapi::ScalarType dtype) { } } -void add_dtype_suffix(std::string& kernel_name, const api::vTensor& tensor) { - return add_dtype_suffix(kernel_name, tensor.dtype()); -} - -void add_ndim_suffix(std::string& kernel_name, const api::vTensor& tensor) { - switch (tensor.storage_type()) { - case utils::kTexture3D: - kernel_name += "_3d"; - break; - case utils::kTexture2D: - kernel_name += "_2d"; - break; - default: - break; - } -} - void add_packed_dim_suffix(std::string& kernel_name, const int32_t packed_dim) { switch (packed_dim) { case WHCN::kWidthDim: @@ -108,10 +85,4 @@ void add_packed_dim_suffix(std::string& kernel_name, const int32_t packed_dim) { } } -void add_packed_dim_suffix( - std::string& kernel_name, - const api::vTensor& tensor) { - return add_packed_dim_suffix(kernel_name, tensor.packed_dim()); -} - } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h index 10084054964..4a2fddb5cf2 100644 --- a/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h +++ b/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h @@ -19,19 +19,11 @@ constexpr size_t kShaderNameReserve = 64u; void add_storage_type_suffix( std::string& kernel_name, const utils::StorageType storage_type); -void add_storage_type_suffix( - std::string& kernel_name, - const api::vTensor& tensor); void add_dtype_suffix(std::string& kernel_name, const vkapi::ScalarType dtype); -void add_dtype_suffix(std::string& kernel_name, const api::vTensor& tensor); void add_ndim_suffix(std::string& kernel_name, const size_t ndim); -void add_ndim_suffix(std::string& kernel_name, const api::vTensor& tensor); void add_packed_dim_suffix(std::string& kernel_name, const int32_t packed_dim); -void add_packed_dim_suffix( - std::string& kernel_name, - const api::vTensor& tensor); } // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp index ea3ae0fa1c3..c2adca526fb 100644 --- a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.cpp @@ -21,29 +21,32 @@ bool is_bitw8(vkapi::ScalarType dtype) { } vkapi::ShaderInfo get_nchw_to_tensor_shader( - const api::vTensor& v_dst, + ComputeGraph& graph, + const ValueRef dst, + const vkapi::ScalarType staging_dtype, bool int8_buffer_enabled, bool push_constant_variant) { std::string kernel_name; kernel_name.reserve(kShaderNameReserve); - if (is_bitw8(v_dst.dtype()) && v_dst.storage_type() != utils::kBuffer && + const vkapi::ScalarType dst_dtype = graph.dtype_of(dst); + const utils::StorageType dst_storage_type = graph.storage_type_of(dst); + + if (is_bitw8(dst_dtype) && dst_storage_type != utils::kBuffer && !int8_buffer_enabled) { kernel_name = "nchw_to_bitw8_image_nobitw8buffer"; if (!push_constant_variant) { kernel_name += "_no_pc"; } - add_storage_type_suffix(kernel_name, v_dst); - add_dtype_suffix(kernel_name, v_dst); + add_storage_type_suffix(kernel_name, dst_storage_type); + add_dtype_suffix(kernel_name, dst_dtype); return VK_KERNEL_FROM_STR(kernel_name); } - if (v_dst.storage_type() == utils::kBuffer) { + if (dst_storage_type == utils::kBuffer) { kernel_name = "nchw_to_buffer"; - if (!push_constant_variant) { - kernel_name += "_no_pc"; - } - add_dtype_suffix(kernel_name, v_dst); + add_dtype_suffix(kernel_name, dst_dtype); + add_dtype_suffix(kernel_name, staging_dtype); return VK_KERNEL_FROM_STR(kernel_name); } @@ -51,36 +54,40 @@ vkapi::ShaderInfo get_nchw_to_tensor_shader( if (!push_constant_variant) { kernel_name += "_no_pc"; } - add_storage_type_suffix(kernel_name, v_dst); - add_dtype_suffix(kernel_name, v_dst); + add_storage_type_suffix(kernel_name, dst_storage_type); + add_dtype_suffix(kernel_name, dst_dtype); + add_dtype_suffix(kernel_name, staging_dtype); return VK_KERNEL_FROM_STR(kernel_name); } vkapi::ShaderInfo get_tensor_to_nchw_shader( - const api::vTensor& v_src, + ComputeGraph& graph, + const ValueRef src, + const vkapi::ScalarType staging_dtype, bool int8_buffer_enabled, bool push_constant_variant) { std::string kernel_name; kernel_name.reserve(kShaderNameReserve); - if (is_bitw8(v_src.dtype()) && v_src.storage_type() != utils::kBuffer && + const vkapi::ScalarType src_dtype = graph.dtype_of(src); + const utils::StorageType src_storage_type = graph.storage_type_of(src); + + if (is_bitw8(src_dtype) && src_storage_type != utils::kBuffer && !int8_buffer_enabled) { kernel_name = "bitw8_image_to_nchw_nobitw8buffer"; if (!push_constant_variant) { kernel_name += "_no_pc"; } - add_storage_type_suffix(kernel_name, v_src); - add_dtype_suffix(kernel_name, v_src); + add_storage_type_suffix(kernel_name, src_storage_type); + add_dtype_suffix(kernel_name, src_dtype); return VK_KERNEL_FROM_STR(kernel_name); } - if (v_src.storage_type() == utils::kBuffer) { + if (src_storage_type == utils::kBuffer) { kernel_name = "buffer_to_nchw"; - if (!push_constant_variant) { - kernel_name += "_no_pc"; - } - add_dtype_suffix(kernel_name, v_src); + add_dtype_suffix(kernel_name, src_dtype); + add_dtype_suffix(kernel_name, staging_dtype); return VK_KERNEL_FROM_STR(kernel_name); } @@ -88,8 +95,9 @@ vkapi::ShaderInfo get_tensor_to_nchw_shader( if (!push_constant_variant) { kernel_name += "_no_pc"; } - add_storage_type_suffix(kernel_name, v_src); - add_dtype_suffix(kernel_name, v_src); + add_storage_type_suffix(kernel_name, src_storage_type); + add_dtype_suffix(kernel_name, src_dtype); + add_dtype_suffix(kernel_name, staging_dtype); return VK_KERNEL_FROM_STR(kernel_name); } diff --git a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h index 9e6b61d6cd8..a4419de3932 100644 --- a/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h +++ b/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h @@ -13,11 +13,16 @@ namespace vkcompute { vkapi::ShaderInfo get_nchw_to_tensor_shader( - const api::vTensor& v_dst, + ComputeGraph& graph, + const ValueRef dst, + const vkapi::ScalarType staging_dtype, bool int8_buffer_enabled = true, bool push_constant_variant = true); + vkapi::ShaderInfo get_tensor_to_nchw_shader( - const api::vTensor& v_src, + ComputeGraph& graph, + const ValueRef src, + const vkapi::ScalarType staging_dtype, bool int8_buffer_enabled = true, bool push_constant_variant = true); diff --git a/backends/vulkan/runtime/utils/VecUtils.h b/backends/vulkan/runtime/utils/VecUtils.h index 6d2e8c63bb9..d84eb54d2b9 100644 --- a/backends/vulkan/runtime/utils/VecUtils.h +++ b/backends/vulkan/runtime/utils/VecUtils.h @@ -275,6 +275,19 @@ struct vec final { VK_CHECK_COND(i >= 0 && i < N, "Index out of bounds!"); return data[i]; } + + bool operator==(const vec& other) const { + for (uint32_t i = 0; i < N; ++i) { + if (data[i] != other.data[i]) { + return false; + } + } + return true; + } + + bool operator!=(const vec& other) const { + return !(*this == other); + } }; } // namespace detail @@ -527,6 +540,16 @@ class WorkgroupSize final { inline constexpr uint32_t operator[](const int idx) const { return (val >> (11 * idx)) & 0x7ffu; } + + // Equality operator + bool operator==(const WorkgroupSize& other) const { + return val == other.val; + } + + // Inequality operator (optional, for completeness) + bool operator!=(const WorkgroupSize& other) const { + return !(*this == other); + } }; } // namespace utils diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index e08491c656b..aa76b202882 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -11,6 +11,7 @@ #include #include +#include namespace vkcompute { namespace vkapi { @@ -109,6 +110,9 @@ VkDevice create_logical_device( #ifdef VK_KHR_shader_float16_int8 VK_KHR_SHADER_FLOAT16_INT8_EXTENSION_NAME, #endif /* VK_KHR_shader_float16_int8 */ +#ifdef VK_KHR_shader_integer_dot_product + VK_KHR_SHADER_INTEGER_DOT_PRODUCT_EXTENSION_NAME, +#endif /* VK_KHR_shader_integer_dot_product */ #if defined(VK_KHR_pipeline_executable_properties) && defined(VULKAN_DEBUG) VK_KHR_PIPELINE_EXECUTABLE_PROPERTIES_EXTENSION_NAME, #endif /* VK_KHR_pipeline_executable_properties */ @@ -160,6 +164,14 @@ VkDevice create_logical_device( extension_list_top = &shader_float16_int8_types; #endif /* VK_KHR_shader_float16_int8 */ +#ifdef VK_KHR_shader_integer_dot_product + VkPhysicalDeviceShaderIntegerDotProductFeaturesKHR + shader_int_dot_product_features{ + physical_device.shader_int_dot_product_features}; + shader_int_dot_product_features.pNext = extension_list_top; + extension_list_top = &shader_int_dot_product_features; +#endif /* VK_KHR_shader_integer_dot_product */ + device_create_info.pNext = extension_list_top; VkDevice handle = nullptr; @@ -401,6 +413,112 @@ std::string Adapter::stringize() const { #endif /* VK_KHR_shader_float16_int8 */ ss << " }" << std::endl; + ss << " Shader 64bit Features {" << std::endl; + PRINT_BOOL(physical_device_.supports_int64_shader_types, shaderInt64) + PRINT_BOOL(physical_device_.supports_float64_shader_types, shaderFloat64) + ss << " }" << std::endl; + +#ifdef VK_KHR_shader_integer_dot_product + ss << " Shader Integer Dot Product Features {" << std::endl; + PRINT_PROP( + physical_device_.shader_int_dot_product_features, + shaderIntegerDotProduct); + ss << " }" << std::endl; + + ss << " Shader Integer Dot Product Properties {" << std::endl; + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct8BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct8BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct8BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct4x8BitPackedUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct4x8BitPackedSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct4x8BitPackedMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct16BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct16BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct16BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct32BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct32BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct32BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct64BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct64BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProduct64BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating8BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating8BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating8BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating4x8BitPackedUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating4x8BitPackedSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating4x8BitPackedMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating16BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating16BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating16BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating32BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating32BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating32BitMixedSignednessAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating64BitUnsignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating64BitSignedAccelerated); + PRINT_PROP( + physical_device_.shader_int_dot_product_properties, + integerDotProductAccumulatingSaturating64BitMixedSignednessAccelerated); + ss << " }" << std::endl; +#endif /* VK_KHR_shader_integer_dot_product */ + const VkPhysicalDeviceMemoryProperties& mem_props = physical_device_.memory_properties; diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index aa4c659c6d8..65d0977b533 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -212,10 +212,27 @@ class Adapter final { #endif /* VK_KHR_shader_float16_int8 */ } + inline bool supports_int8_dot_product() { +#ifdef VK_KHR_shader_integer_dot_product + return physical_device_.shader_int_dot_product_features + .shaderIntegerDotProduct == VK_TRUE; +#else + return false; +#endif /* VK_KHR_shader_integer_dot_product */ + } + inline bool supports_int16_shader_types() { return physical_device_.supports_int16_shader_types; } + inline bool supports_int64_shader_types() { + return physical_device_.supports_int64_shader_types; + } + + inline bool supports_float64_shader_types() { + return physical_device_.supports_float64_shader_types; + } + inline bool has_full_float16_buffers_support() { return supports_16bit_storage_buffers() && supports_float16_shader_types(); } diff --git a/backends/vulkan/runtime/vk_api/Device.cpp b/backends/vulkan/runtime/vk_api/Device.cpp index b9e3b444db2..7a3a825f5ec 100644 --- a/backends/vulkan/runtime/vk_api/Device.cpp +++ b/backends/vulkan/runtime/vk_api/Device.cpp @@ -16,6 +16,7 @@ #include #include #include +#include namespace vkcompute { namespace vkapi { @@ -36,9 +37,17 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle) shader_float16_int8_types{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT16_INT8_FEATURES_KHR}, #endif /* VK_KHR_shader_float16_int8 */ +#ifdef VK_KHR_shader_integer_dot_product + shader_int_dot_product_features{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_INTEGER_DOT_PRODUCT_FEATURES_KHR}, + shader_int_dot_product_properties{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_INTEGER_DOT_PRODUCT_PROPERTIES_KHR}, +#endif queue_families{}, num_compute_queues(0), supports_int16_shader_types(false), + supports_int64_shader_types(false), + supports_float64_shader_types(false), has_unified_memory(false), has_timestamps(false), timestamp_period(0), @@ -77,6 +86,13 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle) extension_list_top = &shader_float16_int8_types; #endif /* VK_KHR_shader_float16_int8 */ +#ifdef VK_KHR_shader_integer_dot_product + shader_int_dot_product_features.pNext = extension_list_top; + extension_list_top = &shader_int_dot_product_features; + shader_int_dot_product_properties.pNext = extension_list_top; + extension_list_top = &shader_int_dot_product_properties; +#endif /* VK_KHR_shader_integer_dot_product */ + features2.pNext = extension_list_top; vkGetPhysicalDeviceFeatures2(handle, &features2); @@ -84,6 +100,12 @@ PhysicalDevice::PhysicalDevice(VkPhysicalDevice physical_device_handle) if (features2.features.shaderInt16 == VK_TRUE) { supports_int16_shader_types = true; } + if (features2.features.shaderInt64 == VK_TRUE) { + supports_int64_shader_types = true; + } + if (features2.features.shaderFloat64 == VK_TRUE) { + supports_float64_shader_types = true; + } // Check if there are any memory types have both the HOST_VISIBLE and the // DEVICE_LOCAL property flags diff --git a/backends/vulkan/runtime/vk_api/Device.h b/backends/vulkan/runtime/vk_api/Device.h index 3fdfcc04a49..917df514c4b 100644 --- a/backends/vulkan/runtime/vk_api/Device.h +++ b/backends/vulkan/runtime/vk_api/Device.h @@ -12,7 +12,7 @@ #include -#include +#include #include namespace vkcompute { @@ -44,6 +44,12 @@ struct PhysicalDevice final { #ifdef VK_KHR_shader_float16_int8 VkPhysicalDeviceShaderFloat16Int8Features shader_float16_int8_types; #endif /* VK_KHR_shader_float16_int8 */ +#ifdef VK_KHR_shader_integer_dot_product + VkPhysicalDeviceShaderIntegerDotProductFeatures + shader_int_dot_product_features; + VkPhysicalDeviceShaderIntegerDotProductProperties + shader_int_dot_product_properties; +#endif /* VK_KHR_shader_integer_dot_product */ // Available GPU queues std::vector queue_families; @@ -51,6 +57,8 @@ struct PhysicalDevice final { // Metadata uint32_t num_compute_queues; bool supports_int16_shader_types; + bool supports_int64_shader_types; + bool supports_float64_shader_types; bool has_unified_memory; bool has_timestamps; float timestamp_period; diff --git a/backends/vulkan/runtime/vk_api/Exception.cpp b/backends/vulkan/runtime/vk_api/Exception.cpp index d26fbd8cb22..d3efa81e52a 100644 --- a/backends/vulkan/runtime/vk_api/Exception.cpp +++ b/backends/vulkan/runtime/vk_api/Exception.cpp @@ -92,6 +92,15 @@ std::ostream& operator<<(std::ostream& out, const VulkanExtension result) { case VulkanExtension::INT8_STORAGE: out << "VK_KHR_8bit_storage"; break; + case VulkanExtension::INTEGER_DOT_PRODUCT: + out << "VK_KHR_shader_integer_dot_product"; + break; + case VulkanExtension::SHADER_INT64: + out << "shaderInt64"; + break; + case VulkanExtension::SHADER_FLOAT64: + out << "shaderFloat64"; + break; } return out; } diff --git a/backends/vulkan/runtime/vk_api/Exception.h b/backends/vulkan/runtime/vk_api/Exception.h index a65afb1bcc5..aa1ef1f2526 100644 --- a/backends/vulkan/runtime/vk_api/Exception.h +++ b/backends/vulkan/runtime/vk_api/Exception.h @@ -82,6 +82,9 @@ enum class VulkanExtension : uint8_t { SHADER_INT16, INT16_STORAGE, INT8_STORAGE, + INTEGER_DOT_PRODUCT, + SHADER_INT64, + SHADER_FLOAT64, }; class ShaderNotSupportedError : public std::exception { diff --git a/backends/vulkan/runtime/vk_api/QueryPool.cpp b/backends/vulkan/runtime/vk_api/QueryPool.cpp index 2f6d433b887..e8b3ca55206 100644 --- a/backends/vulkan/runtime/vk_api/QueryPool.cpp +++ b/backends/vulkan/runtime/vk_api/QueryPool.cpp @@ -209,7 +209,7 @@ std::string QueryPool::generate_string_report() { std::stringstream ss; - int kernel_name_w = 40; + int kernel_name_w = 120; int global_size_w = 25; int local_size_w = 25; int duration_w = 25; diff --git a/backends/vulkan/runtime/vk_api/Shader.cpp b/backends/vulkan/runtime/vk_api/Shader.cpp index 458b1f83956..c932d0a264b 100644 --- a/backends/vulkan/runtime/vk_api/Shader.cpp +++ b/backends/vulkan/runtime/vk_api/Shader.cpp @@ -31,7 +31,10 @@ ShaderInfo::ShaderInfo( const utils::uvec3 tile_size, const bool requires_shader_int16_ext, const bool requires_16bit_storage_ext, - const bool requires_8bit_storage_ext) + const bool requires_8bit_storage_ext, + const bool requires_integer_dot_product_ext, + const bool requires_shader_int64_ext, + const bool requires_shader_float64_ext) : src_code{ spirv_bin, size, @@ -41,7 +44,10 @@ ShaderInfo::ShaderInfo( out_tile_size(tile_size), requires_shader_int16(requires_shader_int16_ext), requires_16bit_storage(requires_16bit_storage_ext), - requires_8bit_storage(requires_8bit_storage_ext) { + requires_8bit_storage(requires_8bit_storage_ext), + requires_integer_dot_product(requires_integer_dot_product_ext), + requires_shader_int64(requires_shader_int64_ext), + requires_shader_float64(requires_shader_float64_ext) { } bool operator==(const ShaderInfo& _1, const ShaderInfo& _2) { diff --git a/backends/vulkan/runtime/vk_api/Shader.h b/backends/vulkan/runtime/vk_api/Shader.h index 7d0fa7b7476..6311710f02b 100644 --- a/backends/vulkan/runtime/vk_api/Shader.h +++ b/backends/vulkan/runtime/vk_api/Shader.h @@ -65,6 +65,9 @@ struct ShaderInfo final { bool requires_shader_int16 = false; bool requires_16bit_storage = false; bool requires_8bit_storage = false; + bool requires_integer_dot_product = false; + bool requires_shader_int64 = false; + bool requires_shader_float64 = false; explicit ShaderInfo(); @@ -76,7 +79,10 @@ struct ShaderInfo final { const utils::uvec3 tile_size, const bool requires_shader_int16_ext, const bool requires_16bit_storage_ext, - const bool requires_8bit_storage_ext); + const bool requires_8bit_storage_ext, + const bool requires_integer_dot_product_ext, + const bool requires_shader_int64_ext, + const bool requires_shader_float64_ext); operator bool() const { return src_code.bin != nullptr; diff --git a/backends/vulkan/runtime/vk_api/memory/Buffer.cpp b/backends/vulkan/runtime/vk_api/memory/Buffer.cpp index 4f58e07b146..f10e40abdbb 100644 --- a/backends/vulkan/runtime/vk_api/memory/Buffer.cpp +++ b/backends/vulkan/runtime/vk_api/memory/Buffer.cpp @@ -20,6 +20,7 @@ VulkanBuffer::VulkanBuffer() allocator_(VK_NULL_HANDLE), memory_{}, owns_memory_(false), + memory_bundled_(false), is_copy_(false), handle_(VK_NULL_HANDLE) {} @@ -33,6 +34,7 @@ VulkanBuffer::VulkanBuffer( allocator_(vma_allocator), memory_{}, owns_memory_(allocate_memory), + memory_bundled_(allocate_memory), is_copy_(false), handle_(VK_NULL_HANDLE) { // If the buffer size is 0, allocate a buffer with a size of 1 byte. This is @@ -77,6 +79,7 @@ VulkanBuffer::VulkanBuffer( allocator_(other.allocator_), memory_(other.memory_), owns_memory_(false), + memory_bundled_(false), is_copy_(true), handle_(other.handle_) { // TODO: set the offset and range appropriately @@ -91,6 +94,7 @@ VulkanBuffer::VulkanBuffer(VulkanBuffer&& other) noexcept allocator_(other.allocator_), memory_(std::move(other.memory_)), owns_memory_(other.owns_memory_), + memory_bundled_(other.memory_bundled_), is_copy_(other.is_copy_), handle_(other.handle_) { other.handle_ = VK_NULL_HANDLE; @@ -99,16 +103,19 @@ VulkanBuffer::VulkanBuffer(VulkanBuffer&& other) noexcept VulkanBuffer& VulkanBuffer::operator=(VulkanBuffer&& other) noexcept { VkBuffer tmp_buffer = handle_; bool tmp_owns_memory = owns_memory_; + bool tmp_memory_bundled = memory_bundled_; buffer_properties_ = other.buffer_properties_; allocator_ = other.allocator_; memory_ = std::move(other.memory_); owns_memory_ = other.owns_memory_; + memory_bundled_ = other.memory_bundled_; is_copy_ = other.is_copy_; handle_ = other.handle_; other.handle_ = tmp_buffer; other.owns_memory_ = tmp_owns_memory; + other.memory_bundled_ = tmp_memory_bundled; return *this; } @@ -119,14 +126,22 @@ VulkanBuffer::~VulkanBuffer() { // ownership of the underlying resource. if (handle_ != VK_NULL_HANDLE && !is_copy_) { if (owns_memory_) { - vmaDestroyBuffer(allocator_, handle_, memory_.allocation); + if (memory_bundled_) { + vmaDestroyBuffer(allocator_, handle_, memory_.allocation); + // Prevent the underlying memory allocation from being freed; it was + // freed by vmaDestroyImage + memory_.allocation = VK_NULL_HANDLE; + } else { + vkDestroyBuffer(this->device(), handle_, nullptr); + // Allow underlying memory allocation to be freed by the destructor of + // Allocation class + } } else { vkDestroyBuffer(this->device(), handle_, nullptr); + // Prevent the underlying memory allocation from being freed since this + // object doesn't own it + memory_.allocation = VK_NULL_HANDLE; } - // Prevent the underlying memory allocation from being freed; it was either - // freed by vmaDestroyBuffer, or this resource does not own the underlying - // memory - memory_.allocation = VK_NULL_HANDLE; } } @@ -136,6 +151,24 @@ VmaAllocationInfo VulkanBuffer::allocation_info() const { return info; } +void VulkanBuffer::bind_allocation_impl(const Allocation& memory) { + VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); + if (!is_copy_) { + VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_)); + } +} + +void VulkanBuffer::bind_allocation(const Allocation& memory) { + bind_allocation_impl(memory); + memory_.allocation = memory.allocation; +} + +void VulkanBuffer::acquire_allocation(Allocation&& memory) { + bind_allocation_impl(memory); + memory_ = std::move(memory); + owns_memory_ = true; +} + VkMemoryRequirements VulkanBuffer::get_memory_requirements() const { VkMemoryRequirements memory_requirements; vkGetBufferMemoryRequirements(this->device(), handle_, &memory_requirements); diff --git a/backends/vulkan/runtime/vk_api/memory/Buffer.h b/backends/vulkan/runtime/vk_api/memory/Buffer.h index e1b441397b4..582b537465d 100644 --- a/backends/vulkan/runtime/vk_api/memory/Buffer.h +++ b/backends/vulkan/runtime/vk_api/memory/Buffer.h @@ -100,6 +100,10 @@ class VulkanBuffer final { Allocation memory_; // Indicates whether the underlying memory is owned by this resource bool owns_memory_; + // Indicates whether the allocation for the buffer was created with the buffer + // via vmaCreateBuffer; if this is false, the memory is owned but was bound + // separately via vmaBindBufferMemory + bool memory_bundled_; // Indicates whether this VulkanBuffer was copied from another VulkanBuffer, // thus it does not have ownership of the underlying VKBuffer bool is_copy_; @@ -162,13 +166,21 @@ class VulkanBuffer final { return (handle_ == other.handle_) && is_copy_; } - inline void bind_allocation(const Allocation& memory) { - VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); - if (!is_copy_) { - VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_)); - } - memory_.allocation = memory.allocation; - } + private: + void bind_allocation_impl(const Allocation& memory); + + public: + /* + * Given a memory allocation, bind it to the underlying VkImage. The lifetime + * of the memory allocation is assumed to be managed externally. + */ + void bind_allocation(const Allocation& memory); + + /* + * Given a rvalue memory allocation, bind it to the underlying VkImage and + * also acquire ownership of the memory allocation. + */ + void acquire_allocation(Allocation&& memory); VkMemoryRequirements get_memory_requirements() const; diff --git a/backends/vulkan/runtime/vk_api/memory/Image.cpp b/backends/vulkan/runtime/vk_api/memory/Image.cpp index da6ff76bccd..cadeb779c83 100644 --- a/backends/vulkan/runtime/vk_api/memory/Image.cpp +++ b/backends/vulkan/runtime/vk_api/memory/Image.cpp @@ -99,6 +99,7 @@ VulkanImage::VulkanImage() allocator_(VK_NULL_HANDLE), memory_{}, owns_memory_(false), + memory_bundled_(false), owns_view_(false), is_copy_(false), handles_{ @@ -125,6 +126,7 @@ VulkanImage::VulkanImage( allocator_(vma_allocator), memory_{}, owns_memory_{allocate_memory}, + memory_bundled_(allocate_memory), owns_view_(false), is_copy_(false), handles_{ @@ -195,6 +197,7 @@ VulkanImage::VulkanImage( allocator_(VK_NULL_HANDLE), memory_{}, owns_memory_(false), + memory_bundled_(false), is_copy_(false), handles_{ image, @@ -224,6 +227,7 @@ VulkanImage::VulkanImage(VulkanImage&& other) noexcept allocator_(other.allocator_), memory_(std::move(other.memory_)), owns_memory_(other.owns_memory_), + memory_bundled_(other.memory_bundled_), owns_view_(other.owns_view_), is_copy_(other.is_copy_), handles_(other.handles_), @@ -232,12 +236,14 @@ VulkanImage::VulkanImage(VulkanImage&& other) noexcept other.handles_.image_view = VK_NULL_HANDLE; other.handles_.sampler = VK_NULL_HANDLE; other.owns_memory_ = false; + other.memory_bundled_ = false; } VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { VkImage tmp_image = handles_.image; VkImageView tmp_image_view = handles_.image_view; bool tmp_owns_memory = owns_memory_; + bool tmp_memory_bundled = memory_bundled_; device_ = other.device_; image_properties_ = other.image_properties_; @@ -246,6 +252,7 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { allocator_ = other.allocator_; memory_ = std::move(other.memory_); owns_memory_ = other.owns_memory_; + memory_bundled_ = other.memory_bundled_; is_copy_ = other.is_copy_; handles_ = other.handles_; layout_ = other.layout_; @@ -253,6 +260,7 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { other.handles_.image = tmp_image; other.handles_.image_view = tmp_image_view; other.owns_memory_ = tmp_owns_memory; + other.memory_bundled_ = tmp_memory_bundled; return *this; } @@ -271,14 +279,22 @@ VulkanImage::~VulkanImage() { if (handles_.image != VK_NULL_HANDLE) { if (owns_memory_) { - vmaDestroyImage(allocator_, handles_.image, memory_.allocation); + if (memory_bundled_) { + vmaDestroyImage(allocator_, handles_.image, memory_.allocation); + // Prevent the underlying memory allocation from being freed; it was + // freed by vmaDestroyImage + memory_.allocation = VK_NULL_HANDLE; + } else { + vkDestroyImage(this->device(), handles_.image, nullptr); + // Allow underlying memory allocation to be freed by the destructor of + // Allocation class + } } else { vkDestroyImage(this->device(), handles_.image, nullptr); + // Prevent the underlying memory allocation from being freed since this + // object doesn't own it + memory_.allocation = VK_NULL_HANDLE; } - // Prevent the underlying memory allocation from being freed; it was either - // freed by vmaDestroyImage, or this resource does not own the underlying - // memory - memory_.allocation = VK_NULL_HANDLE; } } @@ -319,6 +335,31 @@ void VulkanImage::create_image_view() { &(handles_.image_view))); } +void VulkanImage::bind_allocation_impl(const Allocation& memory) { + VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); + // To prevent multiple instances of binding the same VkImage to a memory + // block, do not actually bind memory if this VulkanImage is a copy. Assume + // that the original VulkanImage is responsible for binding the image. + if (!is_copy_) { + VK_CHECK(vmaBindImageMemory(allocator_, memory.allocation, handles_.image)); + } + + // Only create the image view if the image has been bound to memory + owns_view_ = true; + create_image_view(); +} + +void VulkanImage::bind_allocation(const Allocation& memory) { + bind_allocation_impl(memory); + memory_.allocation = memory.allocation; +} + +void VulkanImage::acquire_allocation(Allocation&& memory) { + bind_allocation_impl(memory); + memory_ = std::move(memory); + owns_memory_ = true; +} + VkMemoryRequirements VulkanImage::get_memory_requirements() const { VkMemoryRequirements memory_requirements; vkGetImageMemoryRequirements( diff --git a/backends/vulkan/runtime/vk_api/memory/Image.h b/backends/vulkan/runtime/vk_api/memory/Image.h index 5bbdaf06b47..db632c34378 100644 --- a/backends/vulkan/runtime/vk_api/memory/Image.h +++ b/backends/vulkan/runtime/vk_api/memory/Image.h @@ -156,6 +156,10 @@ class VulkanImage final { Allocation memory_; // Indicates whether the underlying memory is owned by this resource bool owns_memory_; + // Indicates whether the allocation for the image was created with the image + // via vmaCreateImage; if this is false, the memory is owned but was bound + // separately via vmaBindImageMemory + bool memory_bundled_; // In some cases, a VulkanImage may be a copy of another VulkanImage but still // own a unique view of the VkImage. bool owns_view_; @@ -242,21 +246,21 @@ class VulkanImage final { return (handles_.image == other.handles_.image) && is_copy_; } - inline void bind_allocation(const Allocation& memory) { - VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); - // To prevent multiple instances of binding the same VkImage to a memory - // block, do not actually bind memory if this VulkanImage is a copy. Assume - // that the original VulkanImage is responsible for binding the image. - if (!is_copy_) { - VK_CHECK( - vmaBindImageMemory(allocator_, memory.allocation, handles_.image)); - } - memory_.allocation = memory.allocation; - - // Only create the image view if the image has been bound to memory - owns_view_ = true; - create_image_view(); - } + private: + void bind_allocation_impl(const Allocation& memory); + + public: + /* + * Given a memory allocation, bind it to the underlying VkImage. The lifetime + * of the memory allocation is assumed to be managed externally. + */ + void bind_allocation(const Allocation& memory); + + /* + * Given a rvalue memory allocation, bind it to the underlying VkImage and + * also acquire ownership of the memory allocation. + */ + void acquire_allocation(Allocation&& memory); VkMemoryRequirements get_memory_requirements() const; diff --git a/backends/vulkan/serialization/schema.fbs b/backends/vulkan/serialization/schema.fbs index 99ba6a86594..4bc12208ce7 100644 --- a/backends/vulkan/serialization/schema.fbs +++ b/backends/vulkan/serialization/schema.fbs @@ -20,6 +20,7 @@ enum VkDataType : byte { FLOAT32 = 5, FLOAT64 = 6, INT64 = 7, + UNSET = 127, } // Describes what kind of GPU resource should be used to represent a tensor. The @@ -55,6 +56,9 @@ table VkTensor { storage_type:VkStorageType = DEFAULT_STORAGE; // Memory layout that should be used to represent this tensor memory_layout:VkMemoryLayout = DEFAULT_LAYOUT; + // dtype to use for staging buffer. This may be different from the tensor's datatype + // if force_fp16 is enabled to force all float tensors to be represented as fp16. + staging_datatype:VkDataType = UNSET; } table Null {} @@ -118,6 +122,7 @@ table VkValue { table VkBytes { offset:ulong; length:ulong; + named_key:string; } table VkGraph { diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index b74a7fb1f8e..43ea6c7ce30 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.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 ctypes +import hashlib import logging import operator from types import NoneType @@ -25,6 +27,7 @@ is_symint_node, TensorRepr, ) +from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.utils import DelegateMappingBuilder from executorch.exir.tensor import TensorSpec @@ -47,15 +50,18 @@ def __init__( program: ExportedProgram, delegate_mapping_builder: DelegateMappingBuilder, downcast_64_bit: bool = True, + force_fp16: bool = False, ) -> None: self.program = program self.delegate_mapping_builder = delegate_mapping_builder self.downcast_64_bit = downcast_64_bit + self.force_fp16 = force_fp16 self.chain = [] self.values = [] self.input_ids = [] self.output_ids = [] self.const_tensors = [] + self.named_data_store = NamedDataStore() # Mapping from Node to VkValue id self.node_to_value_ids = {} @@ -129,8 +135,42 @@ def get_param_tensor(self, node: Node) -> torch.Tensor: def maybe_add_constant_tensor(self, node: Node) -> int: constant_id = -1 if is_param_node(self.program, node): - constant_id = len(self.const_tensors) - self.const_tensors.append(self.get_param_tensor(node)) + tensor = self.get_param_tensor(node) + + effective_dtype = self.get_effective_dtype(tensor.dtype) + + # Convert the tensor dtype if needed + if tensor.dtype != effective_dtype: + tensor = tensor.to(effective_dtype) + + # Serialize tensor data to bytes + tensor = tensor.contiguous() + size = tensor.untyped_storage().nbytes() + + if size > 0: + array_type = ctypes.c_char * size + array = ctypes.cast( + tensor.untyped_storage().data_ptr(), + ctypes.POINTER(array_type), + ).contents + + # Generate SHA256 hash as the named key + tensor_bytes = bytes(array) + sha256_hash = hashlib.sha256(tensor_bytes) + named_key = sha256_hash.hexdigest() + + # Add to named data store with 16-byte alignment (matching XNNPACK) + self.named_data_store.add_named_data( + named_key, tensor_bytes, alignment=16 + ) + + # Create VkBytes entry with named_key and set offset to indicate named data usage + constant_id = len(self.const_tensors) + self.const_tensors.append((named_key, size)) + else: + # Handle empty tensors + constant_id = len(self.const_tensors) + self.const_tensors.append(None) return constant_id @@ -190,6 +230,29 @@ def create_symint_value(self) -> int: self.values.append(vk_graph_schema.VkValue(vk_graph_schema.SymInt(0))) return new_id + def get_effective_dtype(self, dtype: torch.dtype) -> torch.dtype: + if self.downcast_64_bit and dtype == torch.float64: + return torch.float32 + elif self.downcast_64_bit and dtype == torch.int64: + return torch.int32 + elif self.force_fp16 and dtype == torch.float32: + return torch.float16 + else: + return dtype + + def get_staging_dtype(self, dtype: torch.dtype) -> torch.dtype: + # Since 64 bit types are not guaranteed to be supported on all GPUs, + # the conversion between 32 bit and 64 bit types is handled on the CPU + # side. The conversion will occur when copying the staging buffer + # contents to/from ETensor data pointers, rather than in the shader to + # copy between GPU buffer/image to staging buffer. + if self.downcast_64_bit and dtype == torch.float64: + return torch.float32 + elif self.downcast_64_bit and dtype == torch.int64: + return torch.int32 + else: + return dtype + def create_tensor_value(self, spec: TensorSpec, constant_id: int = -1) -> int: # Negative id indicates that this tensor will have its own dedicated memory. mem_obj_id = -1 @@ -204,14 +267,16 @@ def create_tensor_value(self, spec: TensorSpec, constant_id: int = -1) -> int: storage_type = spec.etvk_node_repr.storage_type memory_layout = spec.etvk_node_repr.memory_layout - # Apply downcast logic before getting VK datatype - effective_dtype = spec.dtype - if self.downcast_64_bit and spec.dtype == torch.float64: - effective_dtype = torch.float32 - elif self.downcast_64_bit and spec.dtype == torch.int64: - effective_dtype = torch.int32 + effective_dtype = self.get_effective_dtype(spec.dtype) + # For constant tensors, the datatype of the original tensor will have been + # converted to the effective dtype. Otherwise, the type of the staging buffer + # for inputs/outputs should match the original tensor dtype. + staging_dtype = ( + effective_dtype if constant_id >= 0 else self.get_staging_dtype(spec.dtype) + ) datatype = self.get_vk_datatype(effective_dtype) + staging_datatype = self.get_vk_datatype(staging_dtype) new_id = len(self.values) self.values.append( @@ -223,6 +288,7 @@ def create_tensor_value(self, spec: TensorSpec, constant_id: int = -1) -> int: mem_obj_id=mem_obj_id, storage_type=storage_type, memory_layout=memory_layout, + staging_datatype=staging_datatype, ) ) ) diff --git a/backends/vulkan/serialization/vulkan_graph_schema.py b/backends/vulkan/serialization/vulkan_graph_schema.py index f845e5601a7..cf5326f40cf 100644 --- a/backends/vulkan/serialization/vulkan_graph_schema.py +++ b/backends/vulkan/serialization/vulkan_graph_schema.py @@ -31,6 +31,7 @@ class VkDataType(IntEnum): FLOAT32 = 5 FLOAT64 = 6 INT64 = 7 + UNSET = 127 class VkStorageType(IntEnum): @@ -61,6 +62,7 @@ class VkTensor: mem_obj_id: int storage_type: VkStorageType = VkStorageType.DEFAULT_STORAGE memory_layout: VkMemoryLayout = VkMemoryLayout.DEFAULT_LAYOUT + staging_datatype: VkDataType = VkDataType.UNSET @dataclass @@ -137,6 +139,7 @@ class VkValue: class VkBytes: offset: int length: int + named_key: str = "" @dataclass diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 2ceedf73d10..96f944560a8 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 Arm Limited and/or its affiliates. # # pyre-strict # @@ -7,14 +8,15 @@ # LICENSE file in the root directory of this source tree. import ctypes +import importlib.resources as _resources import json import os import tempfile - from dataclasses import dataclass from typing import ClassVar, List -import pkg_resources +import executorch.backends.vulkan.serialization as serialization_package + import torch from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( @@ -22,7 +24,6 @@ VkGraph, ) from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass - from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile @@ -32,7 +33,9 @@ def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, "schema.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, "schema.fbs") + ) json_path = os.path.join(d, "schema.json") with open(json_path, "wb") as json_file: json_file.write(vk_graph_json.encode("ascii")) @@ -48,7 +51,9 @@ def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, "schema.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, "schema.fbs") + ) bin_path = os.path.join(d, "schema.bin") with open(bin_path, "wb") as bin_file: @@ -191,10 +196,21 @@ def serialize_constant_tensors( current_offset = len(raw_bytes) for tensor in const_tensors: - if tensor.numel() == 0: + # The tensor data is stored in the named data map + if isinstance(tensor, tuple): + named_key, size = tensor + vk_graph.constants.append( + VkBytes( + offset=18446744073709551615, # UINT64_MAX to indicate named data + length=size, + named_key=named_key, + ) + ) + elif tensor is None or ( + isinstance(tensor, torch.Tensor) and tensor.numel() == 0 + ): vk_graph.constants.append(VkBytes(current_offset, 0)) - continue - else: + elif isinstance(tensor, torch.Tensor): array_type = ctypes.c_char * tensor.untyped_storage().nbytes() array = ctypes.cast( tensor.untyped_storage().data_ptr(), @@ -208,6 +224,8 @@ def serialize_constant_tensors( vk_graph.constants.append(VkBytes(current_offset, len(tensor_bytes))) current_offset += aligned_size(len(tensor_bytes)) + else: + raise ValueError(f"Unsupported constant tensor type: {type(tensor)}") def serialize_custom_shaders( diff --git a/backends/vulkan/targets.bzl b/backends/vulkan/targets.bzl index 590e76e1486..a9ba62b6f9f 100644 --- a/backends/vulkan/targets.bzl +++ b/backends/vulkan/targets.bzl @@ -263,6 +263,7 @@ def define_common_targets(is_fbcode = False): ], exported_deps = [ ":vulkan_graph_runtime_shaderlib{}".format(suffix), + "//executorch/runtime/backend:interface", ], define_static_target = True, # Static initialization is used to register operators to the global operator registry, @@ -303,8 +304,8 @@ def define_common_targets(is_fbcode = False): ":vulkan_graph_runtime{}".format(suffix), "//executorch/backends/vulkan/serialization:vk_delegate_schema", "//executorch/runtime/core:event_tracer", - "//executorch/runtime/backend:interface", "//executorch/runtime/core/exec_aten/util:tensor_util", + "//executorch/runtime/core:named_data_map", ], define_static_target = True, # VulkanBackend.cpp needs to compile with executor as whole @@ -329,7 +330,7 @@ def define_common_targets(is_fbcode = False): "//executorch/exir:tensor", "//executorch/exir/backend/canonical_partitioners:config_partitioner_lib", "//executorch/backends/vulkan/serialization:lib", - ] + ], ) runtime.python_library( @@ -344,6 +345,7 @@ def define_common_targets(is_fbcode = False): ], deps = [ "//caffe2:torch", + "//executorch/backends/vulkan/patterns:vulkan_patterns", ] ) @@ -385,6 +387,8 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/transforms:view_copy_to_squeeze_unsqueeze", "//executorch/backends/vulkan/_passes:vulkan_passes", "//executorch/backends/vulkan/serialization:lib", + "//executorch/backends/transforms:remove_getitem_op", + "//executorch/backends/xnnpack/_passes:xnnpack_passes", "//executorch/exir/backend:backend_details", ], ) diff --git a/backends/vulkan/test/CMakeLists.txt b/backends/vulkan/test/CMakeLists.txt index da25b6e88d1..e3bce1d8baf 100644 --- a/backends/vulkan/test/CMakeLists.txt +++ b/backends/vulkan/test/CMakeLists.txt @@ -35,10 +35,11 @@ if(TARGET vulkan_backend) set(PYTHON_EXECUTABLE python3) endif() - # Include this file to access executorch_target_link_options_shared_lib This is required - # to provide access to executorch_target_link_options_shared_lib which allows libraries - # to be linked with the --whole-archive flag. This is required for libraries - # that perform dynamic registration via static initialization. + # Include this file to access executorch_target_link_options_shared_lib This + # is required to provide access to executorch_target_link_options_shared_lib + # which allows libraries to be linked with the --whole-archive flag. This is + # required for libraries that perform dynamic registration via static + # initialization. include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) include(../cmake/ShaderLibrary.cmake) @@ -82,8 +83,8 @@ if(TARGET vulkan_backend) ) target_include_directories(vulkan_compute_api_test PRIVATE ${COMMON_INCLUDES}) target_link_libraries( - vulkan_compute_api_test PRIVATE GTest::gtest_main vulkan_backend executorch_core - test_shaderlib + vulkan_compute_api_test PRIVATE GTest::gtest_main vulkan_backend + executorch_core test_shaderlib ) target_compile_options(vulkan_compute_api_test PRIVATE ${VULKAN_CXX_FLAGS}) diff --git a/backends/vulkan/test/TARGETS b/backends/vulkan/test/TARGETS index ef429ff21fa..53fad86f90c 100644 --- a/backends/vulkan/test/TARGETS +++ b/backends/vulkan/test/TARGETS @@ -14,6 +14,7 @@ python_unittest( "//executorch/kernels/portable:custom_ops_generated_lib", ], deps = [ + ":test_utils", "//caffe2:torch", "//executorch/backends/transforms:convert_dtype_pass", "//executorch/backends/vulkan:vulkan_preprocess", @@ -68,3 +69,23 @@ runtime.python_library( "//executorch/backends/vulkan:vulkan_preprocess", ] ) + +runtime.python_library( + name = "test_utils", + srcs = [ + "utils.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/xnnpack:xnnpack_preprocess", + "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer", + "//executorch/backends/xnnpack/partition:xnnpack_partitioner", + "//executorch/devtools:lib", + "//executorch/devtools/bundled_program/serialize:lib", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", # @manual + "//executorch/extension/pytree:pylib", + ], +) diff --git a/backends/vulkan/test/custom_ops/CMakeLists.txt b/backends/vulkan/test/custom_ops/CMakeLists.txt new file mode 100644 index 00000000000..97b632338db --- /dev/null +++ b/backends/vulkan/test/custom_ops/CMakeLists.txt @@ -0,0 +1,98 @@ +# 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. + +cmake_minimum_required(VERSION 3.19) +project(prototyping_shaders) + +if(ANDROID) + set(CMAKE_FIND_ROOT_PATH_MODE_LIBRARY BOTH) + set(CMAKE_FIND_ROOT_PATH_MODE_PACKAGE BOTH) +endif() + +find_package(executorch CONFIG REQUIRED COMPONENTS vulkan_backend) + +# Compile settings + +set(VULKAN_CXX_FLAGS "-fexceptions") +list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_WRAPPER") +list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_VOLK") + +message(STATUS "VULKAN_CXX_FLAGS: ${VULKAN_CXX_FLAGS}") + +# Only build if Vulkan was compiled +if(TARGET vulkan_backend) + if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) + endif() + + if(NOT PYTHON_EXECUTABLE) + set(PYTHON_EXECUTABLE python3) + endif() + + # Include this file to access executorch_target_link_options_shared_lib + include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) + include(${EXECUTORCH_ROOT}/backends/vulkan/cmake/ShaderLibrary.cmake) + + # Third party include paths + set(VULKAN_THIRD_PARTY_PATH ${EXECUTORCH_ROOT}/backends/vulkan/third-party) + set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include) + set(VOLK_PATH ${VULKAN_THIRD_PARTY_PATH}/volk) + set(VMA_PATH ${VULKAN_THIRD_PARTY_PATH}/VulkanMemoryAllocator) + + set(COMMON_INCLUDES ${EXECUTORCH_ROOT}/.. ${VULKAN_HEADERS_PATH} ${VOLK_PATH} + ${VMA_PATH} + ) + + # Prototyping utility files + set(PROTOTYPING_UTILS_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}) + set(PROTOTYPING_UTILS_CPP ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp) + + # Prototyping shaders + message(STATUS "shader stuff") + set(PROTOTYPING_SHADERS_PATH ${CMAKE_CURRENT_SOURCE_DIR}/glsl) + gen_vulkan_shader_lib_cpp(${PROTOTYPING_SHADERS_PATH}) + vulkan_shader_lib(prototyping_shaderlib ${generated_spv_cpp}) + target_compile_options(prototyping_shaderlib PRIVATE ${VULKAN_CXX_FLAGS}) + message(STATUS "done shader stuff") + + # Operator implementations library + file(GLOB OPERATOR_IMPL_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/impl/*.cpp) + add_library(operator_implementations STATIC ${OPERATOR_IMPL_SOURCES}) + target_include_directories( + operator_implementations PRIVATE ${COMMON_INCLUDES} + ) + target_link_libraries( + operator_implementations PRIVATE vulkan_backend executorch_core + prototyping_shaderlib + ) + target_compile_options(operator_implementations PRIVATE ${VULKAN_CXX_FLAGS}) + set_property(TARGET operator_implementations PROPERTY CXX_STANDARD 17) + + executorch_target_link_options_shared_lib(vulkan_backend) + executorch_target_link_options_shared_lib(operator_implementations) + + # Function to create operator prototype binaries + function(add_operator_prototype OPERATOR_NAME) + set(TARGET_NAME ${OPERATOR_NAME}) + set(SOURCE_FILE ${CMAKE_CURRENT_SOURCE_DIR}/${OPERATOR_NAME}.cpp) + + add_executable(${TARGET_NAME} ${SOURCE_FILE} ${PROTOTYPING_UTILS_CPP}) + target_include_directories(${TARGET_NAME} PRIVATE ${COMMON_INCLUDES}) + target_link_libraries( + ${TARGET_NAME} PRIVATE vulkan_backend executorch_core + prototyping_shaderlib operator_implementations + ) + target_compile_options(${TARGET_NAME} PRIVATE ${VULKAN_CXX_FLAGS}) + set_property(TARGET ${TARGET_NAME} PROPERTY CXX_STANDARD 17) + endfunction() + + # Define operator prototypes + add_operator_prototype(add) + add_operator_prototype(q8csw_linear) + add_operator_prototype(q8csw_conv2d) + add_operator_prototype(q4gsw_linear) + add_operator_prototype(choose_qparams_per_row) +endif() diff --git a/backends/qualcomm/aot/ir/TARGETS b/backends/vulkan/test/custom_ops/TARGETS similarity index 63% rename from backends/qualcomm/aot/ir/TARGETS rename to backends/vulkan/test/custom_ops/TARGETS index 0a42614a385..e84397dc20e 100644 --- a/backends/qualcomm/aot/ir/TARGETS +++ b/backends/vulkan/test/custom_ops/TARGETS @@ -2,4 +2,4 @@ load(":targets.bzl", "define_common_targets") oncall("executorch") -define_common_targets() +define_common_targets(is_fbcode = True) diff --git a/backends/vulkan/test/custom_ops/add.cpp b/backends/vulkan/test/custom_ops/add.cpp new file mode 100644 index 00000000000..bc20246a7d1 --- /dev/null +++ b/backends/vulkan/test/custom_ops/add.cpp @@ -0,0 +1,165 @@ +// 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 "utils.h" + +using namespace executorch::vulkan::prototyping; + +// Generate test cases for add operation +std::vector generate_add_test_cases() { + std::vector test_cases; + + // Set the data generation type as a local variable + DataGenType data_gen_type = DataGenType::ONES; + + // Define different input size configurations + std::vector> size_configs = { + {1, 64, 64}, // Small square + {1, 128, 128}, // Medium square + {1, 256, 256}, // Large square + {1, 512, 512}, // Very large square + {1, 1, 1024}, // Wide tensor + {1, 1024, 1}, // Tall tensor + {32, 32, 32}, // 3D cube + {16, 128, 64}, // 3D rectangular + }; + + // Storage types to test + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + + // Data types to test + std::vector data_types = {vkapi::kFloat, vkapi::kHalf}; + + // Generate test cases for each combination + for (const auto& sizes : size_configs) { + for (const auto& storage_type : storage_types) { + for (const auto& data_type : data_types) { + TestCase test_case; + + // Create a descriptive name for the test case + std::string size_str = ""; + for (size_t i = 0; i < sizes.size(); ++i) { + size_str += std::to_string(sizes[i]); + if (i < sizes.size() - 1) + size_str += "x"; + } + + std::string storage_str = + (storage_type == utils::kTexture3D) ? "Texture3D" : "Buffer"; + std::string dtype_str = (data_type == vkapi::kFloat) ? "Float" : "Half"; + + // Add data generation type to the name for clarity + std::string test_name = + "Add_" + size_str + "_" + storage_str + "_" + dtype_str; + test_case.set_name(test_name); + + // Set the operator name for the test case + test_case.set_operator_name("etvk.add_prototype"); + + // Add two input tensors with the same size, type, storage, and data + // generation method + ValueSpec input_a( + sizes, data_type, storage_type, utils::kWidthPacked, data_gen_type); + ValueSpec input_b( + sizes, data_type, storage_type, utils::kWidthPacked, data_gen_type); + + // Add output tensor with the same size, type, and storage as inputs + // (output uses ZEROS by default) + ValueSpec output( + sizes, + data_type, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + + test_case.add_input_spec(input_a); + test_case.add_input_spec(input_b); + test_case.add_output_spec(output); + + test_cases.push_back(test_case); + } + } + } + + return test_cases; +} + +// Custom FLOP calculator for add operation +// Add operation performs 1 FLOP (addition) per element +int64_t add_flop_calculator(const TestCase& test_case) { + // Calculate total elements from the first input tensor + int64_t total_elements = 1; + if (!test_case.empty() && test_case.num_inputs() > 0 && + test_case.inputs()[0].is_tensor()) { + const auto& sizes = test_case.inputs()[0].get_tensor_sizes(); + for (int64_t size : sizes) { + total_elements *= size; + } + } + + // Add operation: 1 FLOP per element (one addition) + return total_elements; +} + +// Reference implementation for add operator +void add_reference_compute(TestCase& test_case) { + const ValueSpec& input_a = test_case.inputs().at(0); + const ValueSpec& input_b = test_case.inputs().at(1); + + ValueSpec& output = test_case.outputs().at(0); + + if (input_a.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Calculate number of elements + int64_t num_elements = input_a.numel(); + + auto& input_a_data = input_a.get_float_data(); + auto& input_b_data = input_b.get_float_data(); + + auto& ref_data = output.get_ref_float_data(); + ref_data.resize(num_elements); + for (int64_t i = 0; i < num_elements; ++i) { + ref_data[i] = input_a_data[i] + input_b_data[i]; + } +} + +int main(int argc, char* argv[]) { + set_print_output(false); // Disable output tensor printing + set_print_latencies(false); // Enable latency timing printing + set_use_gpu_timestamps(true); // Enable GPU timestamps + + print_performance_header(); + std::cout << "Add Operation Prototyping Framework" << std::endl; + print_separator(); + + // Initialize Vulkan context + try { + api::context()->initialize_querypool(); + } catch (const std::exception& e) { + std::cerr << "Failed to initialize Vulkan context: " << e.what() + << std::endl; + return 1; + } + + // Execute test cases using the new framework with custom FLOP calculator and + // reference compute + auto results = execute_test_cases( + generate_add_test_cases, + add_flop_calculator, + "Add", + 3, + 10, + add_reference_compute); + + return 0; +} diff --git a/backends/vulkan/test/custom_ops/build_and_run.sh b/backends/vulkan/test/custom_ops/build_and_run.sh new file mode 100755 index 00000000000..2b9ce576e0e --- /dev/null +++ b/backends/vulkan/test/custom_ops/build_and_run.sh @@ -0,0 +1,177 @@ +#!/bin/zsh + +set -eux + +# Check that we're in the executorch directory +current_dir=$(pwd) +if [[ ! "$current_dir" =~ executorch$ ]]; then + echo "Error: This script must be run from a directory ending in 'executorch'" + echo "Current directory: $current_dir" + exit 1 +fi + +# Function to configure and build main project +configure_and_build_main() { + local android_args="" + if [[ "$ANDROID_MODE" == "true" ]]; then + cmake . \ + -DCMAKE_INSTALL_PREFIX=$CMAKE_OUT_DIR \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ + -DANDROID_ABI=arm64-v8a \ + -DANDROID_PLATFORM=android-28 \ + -DGLSLC_PATH=$(which glslc) \ + -B$CMAKE_OUT_DIR + else + cmake . \ + -DCMAKE_INSTALL_PREFIX=$CMAKE_OUT_DIR \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DGLSLC_PATH=$(which glslc) \ + -B$CMAKE_OUT_DIR + fi + + cmake --build $CMAKE_OUT_DIR -j16 --target install + # -DCMAKE_CXX_FLAGS="-DVULKAN_DEBUG" \ +} + +# Function to build main project only +build_main() { + cmake --build $CMAKE_OUT_DIR -j16 --target install +} + +# Function to configure and build tests +configure_and_build_tests() { + # Check if glslc is installed + if ! command -v glslc >/dev/null 2>&1; then + echo "Error: glslc is not installed or not found in PATH." + exit 1 + fi + + local android_args="" + if [[ "$ANDROID_MODE" == "true" ]]; then + cmake backends/vulkan/test/custom_ops/ \ + -DCMAKE_INSTALL_PREFIX=$CMAKE_OUT_DIR \ + -DCMAKE_BUILD_TYPE=Debug \ + -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ + -DANDROID_ABI=arm64-v8a \ + -DANDROID_PLATFORM=android-28 \ + -DGLSLC_PATH=$(which glslc) \ + -B$CMAKE_OUT_DIR/backends/vulkan/test/custom_ops + else + cmake backends/vulkan/test/custom_ops/ \ + -DCMAKE_INSTALL_PREFIX=$CMAKE_OUT_DIR \ + -DCMAKE_BUILD_TYPE=Debug \ + -DGLSLC_PATH=$(which glslc) \ + -B$CMAKE_OUT_DIR/backends/vulkan/test/custom_ops + fi + + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all + +} + +build_tests() { + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all +} + +# Function to rebuild both main and tests +rebuild_both() { + build_main + build_tests +} + +# Function to clean and rebuild everything +clean_and_rebuild() { + rm -rf $CMAKE_OUT_DIR + configure_and_build_main + configure_and_build_tests +} + +# Function to execute binary if specified +execute_binary() { + local binary_name="$1" + if [[ -n "$binary_name" ]]; then + local binary_path="$CMAKE_OUT_DIR/backends/vulkan/test/custom_ops/$binary_name" + echo "Executing binary: $binary_path" + + if [[ "$ANDROID_MODE" == "true" ]]; then + if [[ -f "$binary_path" ]]; then + echo "Pushing binary to Android device..." + adb push "$binary_path" /data/local/tmp/ + echo "Executing binary on Android device..." + adb shell "cd /data/local/tmp && ./$binary_name" + else + echo "Error: Binary '$binary_path' not found" + exit 1 + fi + else + if [[ -f "$binary_path" && -x "$binary_path" ]]; then + "$binary_path" + else + echo "Error: Binary '$binary_path' not found or not executable" + exit 1 + fi + fi + fi +} + +# Parse command line arguments +BINARY_TO_EXECUTE="" +ANDROID_MODE=false +CMAKE_OUT_DIR="cmake-out" + +# Check for --android flag and adjust arguments accordingly +if [[ "$1" == "--android" ]]; then + ANDROID_MODE=true + CMAKE_OUT_DIR="cmake-android-out" + shift # Remove --android from arguments + echo "Android mode enabled. Using $CMAKE_OUT_DIR as build directory." +fi + +case "${1:-}" in + --rebuild|-r) + echo "Rebuilding both main project and tests..." + BINARY_TO_EXECUTE="${2:-}" + rebuild_both + execute_binary "$BINARY_TO_EXECUTE" + ;; + --rebuild1|-r1) + echo "Rebuilding main project only..." + BINARY_TO_EXECUTE="${2:-}" + build_main + execute_binary "$BINARY_TO_EXECUTE" + ;; + --rebuild2|-r2) + echo "Rebuilding tests only..." + BINARY_TO_EXECUTE="${2:-}" + build_tests + execute_binary "$BINARY_TO_EXECUTE" + ;; + --clean|-c) + echo "WARNING: This will delete the entire $CMAKE_OUT_DIR directory and rebuild everything." + echo -n "Are you sure you want to continue? (y/N): " + read -r response + if [[ "$response" =~ ^[Yy]$ ]]; then + echo "Cleaning and rebuilding everything..." + BINARY_TO_EXECUTE="${2:-}" + clean_and_rebuild + execute_binary "$BINARY_TO_EXECUTE" + else + echo "Clean operation cancelled." + exit 0 + fi + ;; + "") + echo "Running full configure and build..." + configure_and_build_main + configure_and_build_tests + ;; + *) + # If first argument doesn't match any build option, treat it as binary name + # and use default build behavior + echo "Running full configure and build..." + BINARY_TO_EXECUTE="$1" + configure_and_build_main + configure_and_build_tests + execute_binary "$BINARY_TO_EXECUTE" + ;; +esac diff --git a/backends/vulkan/test/custom_ops/choose_qparams_per_row.cpp b/backends/vulkan/test/custom_ops/choose_qparams_per_row.cpp new file mode 100644 index 00000000000..f984b2b2b40 --- /dev/null +++ b/backends/vulkan/test/custom_ops/choose_qparams_per_row.cpp @@ -0,0 +1,363 @@ +// 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 "utils.h" + +#include + +using namespace executorch::vulkan::prototyping; +using namespace vkcompute; + +static constexpr int64_t kRefDimSizeLimit = 2050; +static constexpr float SMALL_SCALE_THRESHOLD = 6.1e-5f; + +// ChooseQParams configuration struct +struct ChooseQParamsConfig { + int64_t num_channels; // Height dimension (number of channels) + int64_t channel_size; // Width dimension (size per channel) + int32_t quant_min = -128; + int32_t quant_max = 127; + std::string test_case_name = "placeholder"; + std::string op_name = "choose_qparams_per_row"; +}; + +// Utility function to create a test case from a ChooseQParamsConfig +TestCase create_test_case_from_config( + const ChooseQParamsConfig& config, + utils::StorageType storage_type, + vkapi::ScalarType input_dtype) { + TestCase test_case; + + // Create a descriptive name for the test case + std::string storage_str = + (storage_type == utils::kTexture3D) ? "Texture3D" : "Buffer"; + std::string dtype_str = (input_dtype == vkapi::kFloat) ? "Float" : "Half"; + + std::string test_name = + config.test_case_name + "_" + storage_str + "_" + dtype_str; + test_case.set_name(test_name); + + // Set the operator name for the test case + std::string operator_name = "etvk." + config.op_name + ".default"; + test_case.set_operator_name(operator_name); + + // Input tensor (float) - [num_channels, channel_size] + std::vector input_size = {config.num_channels, config.channel_size}; + ValueSpec input_tensor( + input_size, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + + if (debugging()) { + print_valuespec_data(input_tensor, "input_tensor"); + } + + // Quantization parameters + ValueSpec quant_min(config.quant_min); + ValueSpec quant_max(config.quant_max); + + // Output scale tensor (float) - [num_channels] + ValueSpec scale_out( + {config.num_channels}, + vkapi::kFloat, + utils::kTexture3D, // Always buffer as per requirement + utils::kWidthPacked, + DataGenType::ZEROS); + + // Output zero_point tensor (int8) - [num_channels] + ValueSpec zero_point_out( + {config.num_channels}, + vkapi::kChar, // int8 for quantized zero point + utils::kTexture3D, // Always buffer as per requirement + utils::kWidthPacked, + DataGenType::ZEROS); + + // Add all specs to test case + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(quant_min); + test_case.add_input_spec(quant_max); + test_case.add_output_spec(scale_out); + test_case.add_output_spec(zero_point_out); + + return test_case; +} + +// CPU reference implementation matching the behavior from op_choose_qparams.cpp +void calculate_scale_and_zero_point_reference( + float min_val, + float max_val, + int32_t qmin, + int32_t qmax, + float& scale, + int32_t& zero_point) { + // Extend the [min, max] interval to ensure that it contains 0 + min_val = std::min(min_val, 0.0f); + max_val = std::max(max_val, 0.0f); + + // Use double precision for intermediate computation but use single precision + // in final number to reflect the actual number used during quantization. + double scale_double = + (static_cast(max_val) - min_val) / (qmax - qmin); + + // If scale is 0 or too small so its reciprocal is infinity, we arbitrary + // adjust the scale to 0.1 . We want to avoid scale's reciprocal being + // infinity because some of fbgemm code pre-computes scale's reciprocal to do + // multiplication instead of division in the time critical part of code. + if (static_cast(scale_double) == 0.0f || + std::isinf(1.0f / static_cast(scale_double))) { + scale_double = 0.1; + } + + // Cut off small scale + if (scale_double < SMALL_SCALE_THRESHOLD) { + float org_scale = static_cast(scale_double); + scale_double = SMALL_SCALE_THRESHOLD; + // Adjust the min and max based on the new scale + if (min_val == 0.0f) { + max_val = SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else if (max_val == 0.0f) { + min_val = -SMALL_SCALE_THRESHOLD * (qmax - qmin); + } else { + float amplifier = SMALL_SCALE_THRESHOLD / org_scale; + min_val *= amplifier; + max_val *= amplifier; + } + } + + // Zero-point computation. + // First the initial floating-point computation. The zero-point can be + // determined from solving an affine equation for any known pair + // (real value, corresponding quantized value). + // We know two such pairs: (rmin, qmin) and (rmax, qmax). + // The arithmetic error on the zero point computed from either pair + // will be roughly machine_epsilon * (sum of absolute values of terms) + // so we want to use the variant that adds the smaller terms. + double zero_point_from_min = qmin - min_val / scale_double; + double zero_point_from_max = qmax - max_val / scale_double; + double zero_point_from_min_error = + std::abs(qmin) - std::abs(min_val / scale_double); + double zero_point_from_max_error = + std::abs(qmax) - std::abs(max_val / scale_double); + double initial_zero_point = + zero_point_from_min_error < zero_point_from_max_error + ? zero_point_from_min + : zero_point_from_max; + + // Now we need to nudge the zero point to be an integer + // (our zero points are integer, and this is motivated by the requirement + // to be able to represent the real value "0" exactly as a quantized value, + // which is required in multiple places, for example in Im2col with zero + // padding). + int32_t nudged_zero_point = 0; + if (initial_zero_point < qmin) { + nudged_zero_point = qmin; + } else if (initial_zero_point > qmax) { + nudged_zero_point = qmax; + } else { + nudged_zero_point = + static_cast(nearbyint(static_cast(initial_zero_point))); + } + + scale = static_cast(scale_double); + zero_point = nudged_zero_point; +} + +// Generate easy test cases for choose_qparams_per_channel operation (for +// debugging) +std::vector generate_choose_qparams_per_channel_easy_cases() { + std::vector test_cases; + + // Single simple configuration for debugging + int num_channels = 4; + int channel_size = 8; + + ChooseQParamsConfig config = { + num_channels, // num_channels + channel_size, // channel_size + -128, // quant_min + 127, // quant_max + "simple", // test_case_name + }; + + // Test with both storage types + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + std::vector float_types = {vkapi::kFloat}; + + // Generate test cases for each combination + for (const auto& storage_type : storage_types) { + for (const auto& input_dtype : float_types) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, input_dtype)); + } + } + + return test_cases; +} + +// Generate test cases for choose_qparams_per_channel operation +std::vector generate_choose_qparams_per_channel_test_cases() { + std::vector test_cases; + + std::vector configs = { + {4, 16}, + {8, 32}, + {16, 64}, + {32, 128}, + {64, 256}, + {128, 512}, + {1, 512}, + // Performance cases + {256, 1024}, + {512, 2048}, + {1, 2048}, + {1, 8096}, + }; + + // Test with different storage types + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + + for (auto config : configs) { + std::string prefix = (config.num_channels < kRefDimSizeLimit && + config.channel_size < kRefDimSizeLimit) + ? "correctness_" + : "performance_"; + std::string generated_test_case_name = prefix + + std::to_string(config.num_channels) + "_" + + std::to_string(config.channel_size); + + config.test_case_name = generated_test_case_name; + + for (const auto& storage_type : storage_types) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, vkapi::kFloat)); + } + } + + return test_cases; +} + +// Reference implementation for choose_qparams_per_channel +void choose_qparams_per_channel_reference_impl(TestCase& test_case) { + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& quant_min_spec = test_case.inputs()[idx++]; + const ValueSpec& quant_max_spec = test_case.inputs()[idx++]; + const ValueSpec& eps_spec = test_case.inputs()[idx++]; + const ValueSpec& dtype_spec = test_case.inputs()[idx++]; + (void)eps_spec; // Unused in reference implementation + (void)dtype_spec; // Unused in reference implementation + + // Extract output specifications + ValueSpec& scale_out_spec = test_case.outputs()[0]; + ValueSpec& zero_point_out_spec = test_case.outputs()[1]; + + // Get tensor dimensions + auto input_sizes = + input_spec.get_tensor_sizes(); // [num_channels, channel_size] + int64_t num_channels = input_sizes[0]; + int64_t channel_size = input_sizes[1]; + + // Skip for large tensors since computation time will be extremely slow + if (num_channels > kRefDimSizeLimit || channel_size > kRefDimSizeLimit) { + throw std::invalid_argument( + "One or more dimensions (num_channels, channel_size) exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + int32_t quant_min = quant_min_spec.get_int_value(); + int32_t quant_max = quant_max_spec.get_int_value(); + + // Prepare output data + auto& scale_ref_data = scale_out_spec.get_ref_float_data(); + auto& zero_point_ref_data = zero_point_out_spec.get_ref_int8_data(); + scale_ref_data.resize(num_channels); + zero_point_ref_data.resize(num_channels); + + // Process each channel + for (int64_t channel = 0; channel < num_channels; ++channel) { + // Find min and max for this channel + float min_val = std::numeric_limits::max(); + float max_val = std::numeric_limits::lowest(); + + for (int64_t i = 0; i < channel_size; ++i) { + int64_t input_idx = channel * channel_size + i; + float val = input_data[input_idx]; + min_val = std::min(min_val, val); + max_val = std::max(max_val, val); + } + + // Calculate scale and zero point for this channel + float scale; + int32_t zero_point; + calculate_scale_and_zero_point_reference( + min_val, max_val, quant_min, quant_max, scale, zero_point); + + // Store results (cast zero_point to int8) + scale_ref_data[channel] = scale; + zero_point_ref_data[channel] = static_cast(zero_point); + } +} + +void reference_impl(TestCase& test_case) { + choose_qparams_per_channel_reference_impl(test_case); +} + +int64_t choose_qparams_per_channel_flop_calculator(const TestCase& test_case) { + // Get input dimensions + const auto& input_sizes = test_case.inputs()[0].get_tensor_sizes(); + int64_t num_channels = input_sizes[0]; + int64_t channel_size = input_sizes[1]; + + // Calculate FLOPs for choose_qparams_per_channel operation + // Each channel requires: + // - Min/max finding: approximately 2 * channel_size comparisons + // - Scale calculation: ~5 operations (division, min/max operations) + // - Zero point calculation: ~10 operations (multiple arithmetic operations) + int64_t ops_per_channel = 2 * channel_size + 15; // Simplified estimate + + int64_t flop = num_channels * ops_per_channel; + + return flop; +} + +int main(int argc, char* argv[]) { + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout << "Choose QParams Per Channel Operation Prototyping Framework" + << std::endl; + print_separator(); + + ReferenceComputeFunc ref_fn = reference_impl; + + auto results = execute_test_cases( + generate_choose_qparams_per_channel_test_cases, + choose_qparams_per_channel_flop_calculator, + "ChooseQParamsPerChannel", + 0, + 10, + ref_fn); + + return 0; +} diff --git a/backends/vulkan/test/custom_ops/glsl/add.yaml b/backends/vulkan/test/custom_ops/glsl/add.yaml new file mode 100644 index 00000000000..dd479cafd31 --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/add.yaml @@ -0,0 +1,29 @@ +# 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. + +add_buffer: + parameter_names_with_default_values: + NDIM: 3 + DTYPE: float + PACKING: C_packed + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: add_buffer + +add_texture: + parameter_names_with_default_values: + NDIM: 3 + DTYPE: float + PACKING: C_packed + generate_variant_forall: + DTYPE: + - VALUE: half + - VALUE: float + shader_variants: + - NAME: add_texture3d diff --git a/backends/vulkan/test/custom_ops/glsl/add_buffer.glsl b/backends/vulkan/test/custom_ops/glsl/add_buffer.glsl new file mode 100644 index 00000000000..8a0ddc4dba7 --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/add_buffer.glsl @@ -0,0 +1,39 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define VEC4_T ${texel_type(DTYPE)} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type("buffer")} +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_other", DTYPE, "buffer")} + +layout(push_constant) uniform restrict Block { + int out_numel; +}; + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const int out_bufi = int(gl_GlobalInvocationID.x); + if (out_bufi >= out_numel) { + return; + } + + // Simple addition without broadcasting + t_out[out_bufi] = t_in[out_bufi] + t_other[out_bufi]; +} \ No newline at end of file diff --git a/backends/vulkan/test/custom_ops/glsl/add_texture.glsl b/backends/vulkan/test/custom_ops/glsl/add_texture.glsl new file mode 100644 index 00000000000..f64c8e25d71 --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/add_texture.glsl @@ -0,0 +1,40 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +#define VEC4_T ${texel_type(DTYPE)} + +${define_active_storage_type("texture3d")} +${define_required_extensions(DTYPE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")} +${layout_declare_tensor(B, "r", "t_other", DTYPE, "texture3d")} + +layout(push_constant) uniform restrict Block { + ivec4 out_sizes; +}; + +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", "0")} + +void main() { + const ivec3 pos = ivec3(gl_GlobalInvocationID); + + // Simple addition without broadcasting - same position for all tensors + VEC4_T in_texel = texelFetch(t_in, pos, 0); + VEC4_T other_texel = texelFetch(t_other, pos, 0); + + imageStore(t_out, pos, in_texel + other_texel); +} diff --git a/backends/vulkan/test/custom_ops/glsl/float_canvas.glsl b/backends/vulkan/test/custom_ops/glsl/float_canvas.glsl new file mode 100644 index 00000000000..f821fa3586f --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/float_canvas.glsl @@ -0,0 +1,34 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_active_storage_type("texture3d")} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", "float", "texture3d")} +${layout_declare_tensor(B, "r", "nchw_in", "uint", "buffer")} + +${layout_declare_ubo(B, "ivec3", "out_limits")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const ivec3 lpos = ivec3(gl_GlobalInvocationID); + + if (any(greaterThanEqual(lpos, out_limits))) { + return; + } + + // Placeholder: just copy input to output + vec4 in_texel = vec4(1.0f); + imageStore(t_out, lpos, in_texel); +} diff --git a/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_buffer.glsl b/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_buffer.glsl new file mode 100644 index 00000000000..c1d90fadf7e --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_buffer.glsl @@ -0,0 +1,45 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_active_storage_type("texture3d")} + +#extension GL_EXT_debug_printf : enable + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", "int", "texture3d")} +${layout_declare_tensor(B, "r", "nchw_in", "uint", "buffer")} + +${layout_declare_ubo(B, "ivec3", "out_limits")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const ivec3 lpos = ivec3(gl_GlobalInvocationID); + + if (any(greaterThanEqual(lpos, out_limits))) { + return; + } + + // Pack four 8-bit values equal to 1 into a single uint + int packed = (1 << 0) | (1 << 8) | (1 << 16) | (1 << 24); + + debugPrintfEXT( + "t_out[%i, %i] = %i\\n", + lpos.x, lpos.y, + packed); + + + // Placeholder: just copy input to output + ivec4 in_texel = ivec4(packed); + imageStore(t_out, lpos, in_texel); +} diff --git a/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_texture3d.glsl b/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_texture3d.glsl new file mode 100644 index 00000000000..be6717efdaa --- /dev/null +++ b/backends/vulkan/test/custom_ops/glsl/packed_int32_canvas_texture3d.glsl @@ -0,0 +1,45 @@ +/* + * 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. + */ + +#version 450 core + +#define PRECISION ${PRECISION} + +${define_active_storage_type("texture2d")} + +#extension GL_EXT_debug_printf : enable + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_out", "int", "texture3d")} +${layout_declare_tensor(B, "r", "nchw_in", "uint", "buffer")} + +${layout_declare_ubo(B, "ivec3", "out_limits")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +void main() { + const ivec3 lpos = ivec3(gl_GlobalInvocationID); + + if (any(greaterThanEqual(lpos, out_limits))) { + return; + } + + // Pack four 8-bit values equal to 1 into a single uint + int packed = (1 << 0) | (1 << 8) | (1 << 16) | (1 << 24); + + debugPrintfEXT( + "t_out[%i, %i] = %i\\n", + lpos.x, lpos.y, + packed); + + + // Placeholder: just copy input to output + ivec4 in_texel = ivec4(packed); + imageStore(t_out, lpos, in_texel); +} diff --git a/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp b/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp new file mode 100644 index 00000000000..dc35153baf0 --- /dev/null +++ b/backends/vulkan/test/custom_ops/impl/AddPrototype.cpp @@ -0,0 +1,109 @@ +/* + * 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 { + +// Shader selection function for add operations +vkapi::ShaderInfo pick_add_shader( + 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 in1 = args.at(1).refs.at(0); + + // Build shader name following the binary_op pattern + std::string kernel_name = "add"; + add_storage_type_suffix(kernel_name, graph->storage_type_of(out)); + add_dtype_suffix(kernel_name, graph->dtype_of(in1)); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +// Global workgroup size function for add operations +utils::uvec3 add_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + return default_pick_global_wg_size(graph, shader, args, resize_args); +} + +// Local workgroup size function for add operations +utils::uvec3 add_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + return default_pick_local_wg_size( + graph, shader, global_workgroup_size, args, resize_args); +} + +void add_prototype(ComputeGraph& graph, const std::vector& args) { + int idx = 0; + const ValueRef input_a = args.at(idx++); + const ValueRef input_b = args.at(idx++); + const ValueRef output = args.at(idx++); + + // Prepare parameter buffers (empty for add operation) + vkapi::ParamsBindList param_buffers; + + // Prepare push constants based on storage type + std::vector push_constants; + push_constants.reserve(graph.is_buffer_storage(output) ? 1 : 1); + + if (graph.is_buffer_storage(output)) { + // Buffer storage: pass numel as push constant + push_constants.emplace_back(graph.numel_pc_of(output)); + } else { + // Texture storage: pass sizes as push constant + push_constants.emplace_back(graph.sizes_pc_of(output)); + } + + // Prepare specialization constants + vkapi::SpecVarList spec_vars; + if (graph.is_buffer_storage(output)) { + spec_vars = { + graph.hashed_layout_of(output), + graph.hashed_layout_of(input_a), + graph.hashed_layout_of(input_b)}; + } else { + spec_vars = {graph.hashed_layout_of(output)}; + } + + // Add the compute node + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + pick_add_shader, + add_global_wg_size, + add_local_wg_size, + // Inputs and Outputs + {{output, vkapi::kWrite}, {{input_a, input_b}, vkapi::kRead}}, + // Shader params buffers + param_buffers, + // Push Constants + push_constants, + // Specialization Constants + spec_vars, + // Resize args + {}, + // Resizing Logic + nullptr)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(etvk.add_prototype, add_prototype); +} + +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/q4gsw_linear.cpp b/backends/vulkan/test/custom_ops/q4gsw_linear.cpp new file mode 100644 index 00000000000..59d9d694c2c --- /dev/null +++ b/backends/vulkan/test/custom_ops/q4gsw_linear.cpp @@ -0,0 +1,560 @@ +// 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 "utils.h" + +#include + +using namespace executorch::vulkan::prototyping; + +using namespace vkcompute; + +static constexpr int64_t kRefDimSizeLimit = 300; + +// Linear configuration struct +struct LinearConfig { + int64_t M; // Batch size / number of rows in input + int64_t K; // Input features / columns in input, rows in weight + int64_t N; // Output features / columns in weight + int64_t group_size; // Number of input channels per quantization group + bool has_bias = false; + std::string test_case_name = "placeholder"; + std::string op_name = "linear_dq8ca_q4gsw"; +}; + +// Helper function to unpack 4-bit values from uint8 +std::pair unpack_4bit(uint8_t packed) { + // Extract lower 4 bits and upper 4 bits + int8_t lower = packed & 0x0F; + int8_t upper = (packed >> 4) & 0x0F; + + // Subtract 8 from unpacked 4-bit values + lower -= 8; + upper -= 8; + + return std::make_pair(lower, upper); +} + +// Utility function to create a test case from a LinearConfig +TestCase create_test_case_from_config( + const LinearConfig& config, + utils::StorageType storage_type, + vkapi::ScalarType input_dtype) { + TestCase test_case; + + // Create a descriptive name for the test case + std::string storage_str = + (storage_type == utils::kTexture3D) ? "Texture3D" : "Buffer"; + std::string dtype_str = (input_dtype == vkapi::kFloat) ? "Float" : "Half"; + + std::string test_name = + config.test_case_name + "_" + storage_str + "_" + dtype_str; + test_case.set_name(test_name); + + // Set the operator name for the test case + std::string operator_name = "et_vk." + config.op_name + ".default"; + test_case.set_operator_name(operator_name); + + // Derive sizes from M, K, N + std::vector input_size = {config.M, config.K}; + // Input tensor (float/half) - [M, K] + ValueSpec input_tensor( + input_size, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDINT); + + if (debugging()) { + print_valuespec_data(input_tensor, "input_tensor"); + } + + // For activation+weight quantized linear (linear_dq8ca_q4gsw) + // Input scale and zero point as per-input channel tensors + ValueSpec input_scale( + {1, config.M}, // Per-input channel tensor + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM_SCALES); + input_scale.set_constant(true); + + ValueSpec input_zero_point( + {1, config.M}, // Per-input channel tensor + vkapi::kChar, + storage_type, + utils::kWidthPacked, + DataGenType::RANDINT); + input_zero_point.set_constant(true); + + // For 4-bit weights, packed size is [N, K/2] since 2 weights per byte + std::vector weight_size = {config.N, config.K / 2}; + // Quantized weight tensor (uint8, packed 4-bit) - [N, K/2] + ValueSpec quantized_weight( + weight_size, + vkapi::kByte, // uint8 for packed 4-bit quantized weights + storage_type, + utils::kWidthPacked, + DataGenType::RANDINT4); + quantized_weight.set_constant(true); + quantized_weight.set_int4(true); + + if (debugging()) { + print_valuespec_data(quantized_weight, "weight_tensor"); + } + + // Weight quantization scales (float/half, per-group) + // For group symmetric quantization: [K/group_size, N] + // Each group of input features has scales for all output features + std::vector weight_scales_size = { + config.K / config.group_size, config.N}; + ValueSpec weight_scales( + weight_scales_size, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM_SCALES); + weight_scales.set_constant(true); + + // Pre-computed per-group weight sums for zero point adjustment + // This is needed for activation+weight quantized operations + // Size: [K/group_size, N] - one sum per group per output feature + ValueSpec weight_sums( + weight_scales_size, // Same size as weight_scales + vkapi::kInt, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + weight_sums.set_constant(true); + + // Compute weight_sums data based on quantized weights + int64_t num_groups = config.K / config.group_size; + compute_weight_sums_4bit_grouped( + weight_sums, quantized_weight, num_groups, config.N, config.group_size); + + // Group size parameter + ValueSpec group_size_spec(static_cast(config.group_size)); + + // Bias (optional, float/half) - [N] + ValueSpec bias( + {config.N}, // Per output feature + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + bias.set_constant(true); + if (!config.has_bias) { + bias.set_none(true); + } + + // Output tensor (float/half) - [M, N] + ValueSpec output( + {config.M, config.N}, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + + // Add all specs to test case based on operator type + if (config.op_name.find("dq8ca") != std::string::npos) { + // For activation+weight quantized linear (linear_dq8ca_q4gsw) + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(input_scale); + test_case.add_input_spec(input_zero_point); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_sums); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(group_size_spec); + test_case.add_input_spec(bias); + test_case.add_output_spec(output); + } else { + // For weight-only quantized linear (linear_q4gsw) + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(group_size_spec); + test_case.add_input_spec(bias); + test_case.add_output_spec(output); + } + + return test_case; +} + +// Generate easy test cases for quantized linear operation (for debugging) +std::vector generate_quantized_linear_easy_cases() { + std::vector test_cases; + + // Single simple configuration for debugging + int M = 8; + int K = 16; + int N = 16; + int group_size = 8; + + LinearConfig config = { + M, // Batch size + K, // Input features + N, // Output features + group_size, // Group size + true, // has_bias + "simple", // test_case_name + }; + + // Test with both storage types and data types for completeness + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + std::vector float_types = {vkapi::kFloat}; + + // Generate test cases for each combination + for (const auto& storage_type : storage_types) { + for (const auto& input_dtype : float_types) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, input_dtype)); + } + } + + return test_cases; +} + +// Generate test cases for quantized linear operation +std::vector generate_quantized_linear_test_cases() { + std::vector test_cases; + + std::vector configs = { + // // Gemv test cases + // {1, 128, 64, 32}, + // {1, 256, 128, 64}, + // Gemm + {4, 64, 32, 16}, + {4, 128, 64, 32}, + {4, 256, 128, 64}, + {32, 64, 32, 16}, + {32, 128, 64, 32}, + {32, 256, 128, 64}, + // No bias tests + {32, 128, 64, 32, false}, + {32, 256, 128, 64, false}, + // Performance test cases + {1, 2048, 2048, 128}, + {128, 2048, 2048, 128}, + {256, 2048, 2048, 128}, + {1024, 2048, 2048, 128}, + }; + + // Test with different storage types and data types + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + + for (auto config : configs) { + std::string prefix = + (config.M < kRefDimSizeLimit && config.K < kRefDimSizeLimit && + config.N < kRefDimSizeLimit) + ? "correctness_" + : "performance_"; + std::string generated_test_case_name = prefix + std::to_string(config.M) + + "_" + std::to_string(config.K) + "_" + std::to_string(config.N) + "_g" + + std::to_string(config.group_size); + if (!config.has_bias) { + generated_test_case_name += "_no_bias"; + } + + config.test_case_name = generated_test_case_name; + + for (const auto& storage_type : storage_types) { + // Test both activation+weight quantized and weight only quantized, but + // only if the current device supports int8 dot product + if (vkcompute::api::context() + ->adapter_ptr() + ->supports_int8_dot_product()) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, vkapi::kFloat)); + } + + LinearConfig wo_quant_config = config; + wo_quant_config.op_name = "linear_q4gsw"; + test_cases.push_back(create_test_case_from_config( + wo_quant_config, storage_type, vkapi::kFloat)); + } + } + + return test_cases; +} + +// Reference implementation for 4-bit group symmetric weight quantized linear +void linear_q4gsw_reference_impl(TestCase& test_case) { + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& group_size_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [batch_size, in_features] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [in_features, out_features/2] + auto output_sizes = + output_spec.get_tensor_sizes(); // [batch_size, out_features] + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = output_sizes[1]; + int64_t group_size = group_size_spec.get_int_value(); + + // Skip for large tensors since computation time will be extremely slow + if (batch_size > kRefDimSizeLimit || in_features > kRefDimSizeLimit || + out_features > kRefDimSizeLimit) { + throw std::invalid_argument( + "One or more dimensions (batch_size, in_features, out_features) exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + auto& weight_data = weight_spec.get_uint8_data(); + auto& weight_scales_data = weight_scales_spec.get_float_data(); + auto& bias_data = bias_spec.get_float_data(); + + // Calculate number of output elements + int64_t num_output_elements = batch_size * out_features; + + auto& ref_data = output_spec.get_ref_float_data(); + ref_data.resize(num_output_elements); + + // Perform quantized linear transformation (matrix multiplication) + for (int64_t b = 0; b < batch_size; ++b) { + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + float sum = 0.0f; + + // Matrix multiplication: output[b][out_f] = sum(input[b][in_f] * + // weight[out_f][in_f]) + for (int64_t in_f = 0; in_f < in_features; ++in_f) { + // Get input value + int64_t input_idx = b * in_features + in_f; + float input_val = input_data[input_idx]; + + // Get weight value and dequantize (4-bit group symmetric quantization) + int64_t group_idx = in_f / group_size; + int64_t scales_idx = group_idx * out_features + out_f; + + // Get packed weight value - weight matrix is [N, K/2] + int64_t weight_idx = (out_f) * (in_features / 2) + (in_f / 2); + uint8_t packed_weight = weight_data[weight_idx]; + + // Unpack 4-bit weight + auto unpacked = unpack_4bit(packed_weight); + int8_t weight_4bit = (in_f % 2 == 0) ? unpacked.first : unpacked.second; + + // Dequantize weight using group symmetric quantization (no zero point) + float weight_scale = weight_scales_data[scales_idx]; + float dequant_weight = static_cast(weight_4bit) * weight_scale; + + sum += input_val * dequant_weight; + } + + // Add bias and store result + if (!bias_spec.is_none()) { + sum += bias_data[out_f]; + } + int64_t output_idx = b * out_features + out_f; + ref_data[output_idx] = sum; + } + } +} + +// Reference implementation for activation+weight quantized linear (dq8ca_q4gsw) +void linear_dq8ca_q4gsw_reference_impl(TestCase& test_case) { + // Extract input specifications + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& input_scale_spec = test_case.inputs()[idx++]; + const ValueSpec& input_zeros_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_sums_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& group_size_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [batch_size, in_features] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [out_features, in_features/2] + auto output_sizes = + output_spec.get_tensor_sizes(); // [batch_size, out_features] + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = output_sizes[1]; + int64_t group_size = group_size_spec.get_int_value(); + + // Skip for large tensors since computation time will be extremely slow + if (batch_size > kRefDimSizeLimit || in_features > kRefDimSizeLimit || + out_features > kRefDimSizeLimit) { + throw std::invalid_argument( + "One or more dimensions (batch_size, in_features, out_features) exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + auto& input_scale_data = + input_scale_spec.get_float_data(); // Per-input channel tensor + auto& input_zero_point_data = + input_zeros_spec.get_int8_data(); // Per-input channel tensor + + auto& weight_data = weight_spec.get_uint8_data(); + auto& weight_sums_data = weight_sums_spec.get_int32_data(); + (void)weight_sums_data; // Unused for now + auto& weight_scales_data = weight_scales_spec.get_float_data(); + auto& bias_data = bias_spec.get_float_data(); + + // Calculate number of output elements + int64_t num_output_elements = batch_size * out_features; + + auto& ref_data = output_spec.get_ref_float_data(); + ref_data.resize(num_output_elements); + + // Perform quantized linear transformation (matrix multiplication) with + // integer accumulation + for (int64_t b = 0; b < batch_size; ++b) { + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + int32_t int_sum = 0; + (void)int_sum; + int32_t weight_sum = 0; // Track weight sum on the fly for each group + (void)weight_sum; + + // For group symmetric quantization, compute with proper grouping for + // accurate reference + float float_result = 0.0f; + + for (int64_t in_f = 0; in_f < in_features; ++in_f) { + // Get input value and quantize to int8 using per-input channel + // parameters + int64_t input_idx = b * in_features + in_f; + + // Use per-input channel scale and zero point - index by batch dimension + float input_scale = input_scale_data[b]; // {1, M} -> index by batch + int8_t input_zero_point = + input_zero_point_data[b]; // {1, M} -> index by batch + + float quant_input_f = + std::round(input_data[input_idx] / input_scale) + input_zero_point; + quant_input_f = std::min(std::max(quant_input_f, -128.0f), 127.0f); + int8_t quantized_input = static_cast(quant_input_f); + + // Get quantized weight and its scale + int64_t weight_idx = out_f * (in_features / 2) + (in_f / 2); + uint8_t packed_weight = weight_data[weight_idx]; + auto unpacked = unpack_4bit(packed_weight); + int8_t quantized_weight = + (in_f % 2 == 0) ? unpacked.first : unpacked.second; + + // Get the appropriate scale for this group + int64_t group_idx = in_f / group_size; + int64_t scales_idx = group_idx * out_features + out_f; + float weight_scale = weight_scales_data[scales_idx]; + + // Compute the contribution with proper scaling + float contribution = + static_cast(quantized_input - input_zero_point) * + static_cast(quantized_weight) * input_scale * weight_scale; + + float_result += contribution; + } + + // Add bias and store result + if (!bias_spec.is_none()) { + float_result += bias_data[out_f]; + } + int64_t output_idx = b * out_features + out_f; + ref_data[output_idx] = float_result; + } + } +} + +void reference_impl(TestCase& test_case) { + if (test_case.operator_name().find("dq8ca") != std::string::npos) { + linear_dq8ca_q4gsw_reference_impl(test_case); + } else { + linear_q4gsw_reference_impl(test_case); + } +} + +int64_t quantized_linear_flop_calculator(const TestCase& test_case) { + int input_idx = 0; + int weight_idx = 1; + if (test_case.operator_name().find("dq8ca") != std::string::npos) { + input_idx = 0; + weight_idx = 3; // Weight comes after input, input_scale, input_zero_point + } + + // Get input and weight dimensions + const auto& input_sizes = test_case.inputs()[0].get_tensor_sizes(); + const auto& output_sizes = test_case.outputs()[0].get_tensor_sizes(); + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = output_sizes[1]; + + // Calculate FLOPs for quantized linear operation + // Each output element requires: + // - in_features multiply-accumulate operations + // - Additional operations for quantization/dequantization + int64_t output_elements = batch_size * out_features; + int64_t ops_per_output = in_features; + + // Add quantization overhead (approximate) + // - Unpack 4-bit weight: 1 op per weight element used + // - Dequantize weight: 1 op per weight element used + // - Add bias: 1 op per output element + // - For activation+weight quantized: add input quantization ops + int64_t quantization_ops = ops_per_output * 2 + 1; // Simplified estimate + + int64_t flop = output_elements * (ops_per_output + quantization_ops); + + return flop; +} + +int main(int argc, char* argv[]) { + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout + << "4-bit Group Symmetric Weight Quantized Linear Operation Prototyping Framework" + << std::endl; + print_separator(); + + ReferenceComputeFunc ref_fn = reference_impl; + + // Execute easy test cases using the new framework with custom FLOP calculator + auto results = execute_test_cases( + generate_quantized_linear_test_cases, + quantized_linear_flop_calculator, + "QuantizedLinearQ4GSW", + 10, + 10, + ref_fn); + + return 0; +} diff --git a/backends/vulkan/test/custom_ops/q8csw_conv2d.cpp b/backends/vulkan/test/custom_ops/q8csw_conv2d.cpp new file mode 100644 index 00000000000..d566e5b2646 --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8csw_conv2d.cpp @@ -0,0 +1,785 @@ +// 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 "utils.h" + +#include + +using namespace executorch::vulkan::prototyping; + +using namespace vkcompute; + +static constexpr int64_t kRefDimSizeLimit = 100; + +// Component structs for better readability +struct KernelSize { + int32_t h; + int32_t w; + + KernelSize(int32_t height, int32_t width) : h(height), w(width) {} +}; + +struct Stride { + int32_t h; + int32_t w; + + Stride(int32_t height, int32_t width) : h(height), w(width) {} +}; + +struct Padding { + int32_t h; + int32_t w; + + Padding(int32_t height, int32_t width) : h(height), w(width) {} +}; + +struct Dilation { + int32_t h; + int32_t w; + + Dilation(int32_t height = 1, int32_t width = 1) : h(height), w(width) {} +}; + +struct OutInChannels { + int32_t out; + int32_t in; + + OutInChannels(int32_t out_channels, int32_t in_channels) + : out(out_channels), in(in_channels) {} +}; + +struct InputSize2D { + int32_t h; + int32_t w; + + InputSize2D(int32_t height, int32_t width) : h(height), w(width) {} +}; + +// Conv2d configuration struct +struct Conv2dConfig { + OutInChannels channels; + InputSize2D input_size; + KernelSize kernel; + Stride stride; + Padding padding; + Dilation dilation; + int32_t groups; // Number of groups for grouped convolution + std::string test_case_name = "placeholder"; + std::string op_name = "conv2d_q8ta_q8csw"; + + // Calculate output dimensions + int64_t get_output_height() const { + return (input_size.h + 2 * padding.h - dilation.h * (kernel.h - 1) - 1) / + stride.h + + 1; + } + + int64_t get_output_width() const { + return (input_size.w + 2 * padding.w - dilation.w * (kernel.w - 1) - 1) / + stride.w + + 1; + } +}; + +// Utility function to create a test case from a Conv2dConfig +TestCase create_test_case_from_config( + const Conv2dConfig& config, + utils::StorageType storage_type, + vkapi::ScalarType input_dtype) { + TestCase test_case; + + // Create a descriptive name for the test case + std::string storage_str = + (storage_type == utils::kTexture3D) ? "Texture3D" : "Buffer"; + std::string dtype_str = (input_dtype == vkapi::kFloat) ? "Float" : "Half"; + + std::string test_name = + config.test_case_name + "_" + storage_str + "_" + dtype_str; + test_case.set_name(test_name); + + // Set the operator name for the test case + std::string operator_name = "et_vk." + config.op_name + ".default"; + test_case.set_operator_name(operator_name); + + // Calculate output dimensions + int64_t H_out = config.get_output_height(); + int64_t W_out = config.get_output_width(); + + // Input tensor (float/half) - [1, C_in, H_in, W_in] (batch size always 1) + std::vector input_size = { + 1, config.channels.in, config.input_size.h, config.input_size.w}; + + ValueSpec input_tensor( + input_size, + input_dtype, + storage_type, + utils::kChannelsPacked, + DataGenType::RANDOM); + + if (debugging()) { + print_valuespec_data(input_tensor, "input_tensor"); + } + + float input_scale_val = 0.07f; + ValueSpec input_scale(input_scale_val); + + int32_t input_zero_point_val = -3; + ValueSpec input_zero_point(input_zero_point_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 + const int64_t in_channels_per_group = config.channels.in / config.groups; + const int64_t in_features = utils::align_up_4( + in_channels_per_group * config.kernel.h * config.kernel.w); + std::vector weight_size = {config.channels.out, in_features}; + ValueSpec quantized_weight( + weight_size, + vkapi::kChar, // int8 for quantized weights + storage_type, + utils::kWidthPacked, + DataGenType::RANDINT8); + quantized_weight.set_constant(true); + + if (debugging()) { + print_valuespec_data(quantized_weight, "weight_tensor"); + } + + // Weight quantization scales (float/half, per-channel) + ValueSpec weight_scales( + {config.channels.out}, // Per output channel + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM_SCALES); + weight_scales.set_constant(true); + + ValueSpec weight_sums( + {config.channels.out}, // Per output channel + vkapi::kInt, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + weight_sums.set_constant(true); + + // Compute weight_sums data based on quantized weights + compute_weight_sums( + weight_sums, quantized_weight, config.channels.out, in_features); + + // Bias (optional, float/half) - [C_out] + ValueSpec bias( + {config.channels.out}, // Per output channel + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + bias.set_constant(true); + + // Stride and padding parameters + ValueSpec stride({config.stride.h, config.stride.w}); + ValueSpec padding({config.padding.h, config.padding.w}); + + // Dilation and groups parameters + ValueSpec dilation({config.dilation.h, config.dilation.w}); + ValueSpec groups(config.groups); + + // Kernel size parameters + ValueSpec kernel_size({config.kernel.h, config.kernel.w}); + + // Output tensor (float/half) - [1, C_out, H_out, W_out] (batch size always 1) + ValueSpec output( + {1, config.channels.out, H_out, W_out}, + input_dtype, + storage_type, + utils::kChannelsPacked, + DataGenType::ZEROS); + + // Add all specs to test case + if (config.op_name.find("q8ta") != std::string::npos) { + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(input_scale); + test_case.add_input_spec(input_zero_point); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_sums); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(bias); + test_case.add_input_spec(kernel_size); + test_case.add_input_spec(stride); + test_case.add_input_spec(padding); + test_case.add_input_spec(dilation); + test_case.add_input_spec(groups); + } else { + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(bias); + test_case.add_input_spec(kernel_size); + test_case.add_input_spec(stride); + test_case.add_input_spec(padding); + test_case.add_input_spec(dilation); + test_case.add_input_spec(groups); + } + + test_case.add_output_spec(output); + + return test_case; +} + +// Generate easy test cases for quantized conv2d operation (for debugging) +std::vector generate_quantized_conv2d_easy_cases() { + std::vector test_cases; + + // Single simple configuration for debugging + Conv2dConfig config = { + OutInChannels(8, 3), // channels (out, in) + InputSize2D(8, 8), // input_size (h, w) + KernelSize(3, 3), // kernel + Stride(1, 1), // stride + Padding(0, 0), // padding + Dilation(1, 1), // dilation + 1, // groups + }; + + // Test with both storage types and data types for completeness + std::vector storage_types = {utils::kTexture3D}; + std::vector float_types = {vkapi::kFloat}; + + // Generate test cases for each combination + for (const auto& storage_type : storage_types) { + for (const auto& input_dtype : float_types) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, input_dtype)); + } + } + + return test_cases; +} + +// Generate test cases for quantized conv2d operation +std::vector generate_quantized_conv2d_test_cases() { + std::vector test_cases; + + std::vector configs = { + {OutInChannels(32, 3), + InputSize2D(64, 64), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1}, + {OutInChannels(32, 16), + InputSize2D(32, 32), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + {OutInChannels(64, 32), + InputSize2D(16, 16), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + // One output channel case + {OutInChannels(1, 32), + InputSize2D(55, 55), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + + // Stride 2 convolutions + {OutInChannels(32, 3), + InputSize2D(64, 64), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1}, + {OutInChannels(64, 32), + InputSize2D(32, 32), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1}, + // Different kernel sizes + {OutInChannels(32, 16), + InputSize2D(28, 28), + KernelSize(5, 5), + Stride(1, 1), + Padding(2, 2), + Dilation(1, 1), + 1}, + {OutInChannels(64, 32), + InputSize2D(14, 14), + KernelSize(7, 7), + Stride(1, 1), + Padding(3, 3), + Dilation(1, 1), + 1}, + + // Dilated convolutions + {OutInChannels(32, 16), + InputSize2D(32, 32), + KernelSize(3, 3), + Stride(1, 1), + Padding(2, 2), + Dilation(2, 2), + 1}, + {OutInChannels(64, 32), + InputSize2D(16, 16), + KernelSize(3, 3), + Stride(1, 1), + Padding(3, 3), + Dilation(3, 3), + 1}, + + // Grouped convolutions + {OutInChannels(32, 32), + InputSize2D(32, 32), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 4}, + {OutInChannels(64, 64), + InputSize2D(16, 16), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 8}, + // Performance test cases + {OutInChannels(256, 128), + InputSize2D(128, 128), + KernelSize(1, 1), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 8}, + {OutInChannels(128, 64), + InputSize2D(128, 128), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}}; + + // Test with different storage types and data types + std::vector storage_types = {utils::kTexture3D}; + + // Generate test cases for each combination + for (auto& config : configs) { + for (const auto& storage_type : storage_types) { + // Generate test case name programmatically + bool is_performance = config.channels.out > kRefDimSizeLimit || + config.channels.in > kRefDimSizeLimit || + config.input_size.h > kRefDimSizeLimit || + config.input_size.w > kRefDimSizeLimit; + std::string prefix = is_performance ? "performance_" : "correctness_"; + std::string suffix = std::to_string(config.channels.out) + "/" + + std::to_string(config.channels.in) + "_" + + std::to_string(config.input_size.h) + "/" + + std::to_string(config.input_size.w) + "_" + + std::to_string(config.kernel.h) + "/" + + std::to_string(config.kernel.w); + + config.test_case_name = prefix + suffix; + // The default operator tested is activation + weight quantized conv2d; + // however, only test this if the int8 dot product extension is supported + if (vkcompute::api::context() + ->adapter_ptr() + ->supports_int8_dot_product()) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, vkapi::kFloat)); + } + + Conv2dConfig wo_quant_config = config; + wo_quant_config.op_name = "conv2d_q8csw"; + test_cases.push_back(create_test_case_from_config( + wo_quant_config, storage_type, vkapi::kFloat)); + } + } + + return test_cases; +} + +// Reference implementation for weight only quantized conv2d (fp accumulation) +void conv2d_q8csw_reference_impl(TestCase& test_case) { + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + const ValueSpec& kernel_size_spec = test_case.inputs()[idx++]; + const ValueSpec& stride_spec = test_case.inputs()[idx++]; + const ValueSpec& padding_spec = test_case.inputs()[idx++]; + const ValueSpec& dilation_spec = test_case.inputs()[idx++]; + const ValueSpec& groups_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [N, C_in, H_in, W_in] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [C_out, C_in_per_group * K_h * K_w] + auto output_sizes = + output_spec.get_tensor_sizes(); // [N, C_out, H_out, W_out] + + int64_t N = input_sizes[0]; + int64_t C_in = input_sizes[1]; + int64_t H_in = input_sizes[2]; + int64_t W_in = input_sizes[3]; + int64_t C_out = output_sizes[1]; + int64_t H_out = output_sizes[2]; + int64_t W_out = output_sizes[3]; + + // Get kernel dimensions from kernel_size ValueSpec + auto kernel_size_data = kernel_size_spec.get_int32_data(); + int64_t K_h = kernel_size_data[0]; + int64_t K_w = kernel_size_data[1]; + + // Get stride, padding, dilation, and groups + auto stride_data = stride_spec.get_int32_data(); + auto padding_data = padding_spec.get_int32_data(); + auto dilation_data = dilation_spec.get_int32_data(); + int64_t stride_h = stride_data[0]; + int64_t stride_w = stride_data[1]; + int64_t pad_h = padding_data[0]; + int64_t pad_w = padding_data[1]; + int64_t dilation_h = dilation_data[0]; + 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) { + throw std::invalid_argument( + "One or more dimensions exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + 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(); + + // Calculate channels per group for grouped convolution + int64_t C_in_per_group = C_in / groups; + int64_t C_out_per_group = C_out / groups; + + // Calculate number of output elements + int64_t num_output_elements = N * C_out * H_out * W_out; + + 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); + + // Perform weight-only quantized conv2d operation (fp accumulation) + for (int64_t n = 0; n < N; ++n) { + for (int64_t out_c = 0; out_c < C_out; ++out_c) { + for (int64_t out_h = 0; out_h < H_out; ++out_h) { + for (int64_t out_w = 0; out_w < W_out; ++out_w) { + float sum = 0.0f; + + // Determine which group this output channel belongs to + int64_t group_idx = out_c / C_out_per_group; + int64_t in_c_start = group_idx * C_in_per_group; + int64_t in_c_end = (group_idx + 1) * C_in_per_group; + + // Convolution operation with dilation support and grouped convolution + for (int64_t in_c = in_c_start; in_c < in_c_end; ++in_c) { + for (int64_t kh = 0; kh < K_h; ++kh) { + for (int64_t kw = 0; kw < K_w; ++kw) { + // Calculate input position with dilation + int64_t in_h = out_h * stride_h - pad_h + kh * dilation_h; + int64_t in_w = out_w * stride_w - pad_w + kw * dilation_w; + + // Check bounds (zero padding) + if (in_h >= 0 && in_h < H_in && in_w >= 0 && in_w < W_in) { + // Get input value (keep as float) + int64_t input_idx = n * (C_in * H_in * W_in) + + in_c * (H_in * W_in) + in_h * W_in + in_w; + float input_val = input_data[input_idx]; + + // Get weight value and dequantize + // Weight layout: [C_out, C_in_per_group * K_h * K_w] + int64_t weight_idx = out_c * in_features + + (kh * (K_w * C_in_per_group) + kw * C_in_per_group + + (in_c % C_in_per_group)); + float dequant_weight = + (static_cast(weight_data[weight_idx])) * + weight_scales_data[out_c]; + + sum += input_val * dequant_weight; + } + } + } + } + + // Add bias and store result + sum += bias_data[out_c]; + int64_t output_idx = n * (C_out * H_out * W_out) + + out_c * (H_out * W_out) + out_h * W_out + out_w; + ref_data[output_idx] = sum; + } + } + } + } +} + +// Reference implementation for activation and weight quantized conv2d (int +// accumulation) +void conv2d_q8ta_q8csw_reference_impl(TestCase& test_case) { + // Extract input specifications + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& input_scale_spec = test_case.inputs()[idx++]; + const ValueSpec& input_zeros_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_sums_spec = test_case.inputs()[idx++]; + (void)weight_sums_spec; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + const ValueSpec& kernel_size_spec = test_case.inputs()[idx++]; + const ValueSpec& stride_spec = test_case.inputs()[idx++]; + const ValueSpec& padding_spec = test_case.inputs()[idx++]; + const ValueSpec& dilation_spec = test_case.inputs()[idx++]; + const ValueSpec& groups_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [N, C_in, H_in, W_in] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [C_out, C_in_per_group * K_h * K_w] + auto output_sizes = + output_spec.get_tensor_sizes(); // [N, C_out, H_out, W_out] + + int64_t N = input_sizes[0]; + int64_t C_in = input_sizes[1]; + int64_t H_in = input_sizes[2]; + int64_t W_in = input_sizes[3]; + int64_t C_out = output_sizes[1]; + int64_t H_out = output_sizes[2]; + int64_t W_out = output_sizes[3]; + + // Get kernel dimensions from kernel_size ValueSpec + auto kernel_size_data = kernel_size_spec.get_int32_data(); + int64_t K_h = kernel_size_data[0]; + int64_t K_w = kernel_size_data[1]; + + // Get stride, padding, dilation, and groups + auto stride_data = stride_spec.get_int32_data(); + auto padding_data = padding_spec.get_int32_data(); + auto dilation_data = dilation_spec.get_int32_data(); + int64_t stride_h = stride_data[0]; + int64_t stride_w = stride_data[1]; + int64_t pad_h = padding_data[0]; + int64_t pad_w = padding_data[1]; + int64_t dilation_h = dilation_data[0]; + 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) { + throw std::invalid_argument( + "One or more dimensions exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + const float input_scale = input_scale_spec.get_float_value(); + const int32_t input_zero_point = input_zeros_spec.get_int_value(); + + 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(); + + // Calculate channels per group for grouped convolution + int64_t C_in_per_group = C_in / groups; + int64_t C_out_per_group = C_out / groups; + + // Calculate number of output elements + int64_t num_output_elements = N * C_out * H_out * W_out; + + 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); + + // Perform activation and weight quantized conv2d operation (int accumulation) + for (int64_t n = 0; n < N; ++n) { + for (int64_t out_c = 0; out_c < C_out; ++out_c) { + for (int64_t out_h = 0; out_h < H_out; ++out_h) { + for (int64_t out_w = 0; out_w < W_out; ++out_w) { + int32_t int_sum = 0; + int32_t weight_sum = 0; // Track weight sum on the fly + + // Determine which group this output channel belongs to + int64_t group_idx = out_c / C_out_per_group; + int64_t in_c_start = group_idx * C_in_per_group; + int64_t in_c_end = (group_idx + 1) * C_in_per_group; + + // Convolution operation with integer accumulation + for (int64_t in_c = in_c_start; in_c < in_c_end; ++in_c) { + for (int64_t kh = 0; kh < K_h; ++kh) { + for (int64_t kw = 0; kw < K_w; ++kw) { + // Calculate input position with dilation + int64_t in_h = out_h * stride_h - pad_h + kh * dilation_h; + int64_t in_w = out_w * stride_w - pad_w + kw * dilation_w; + + // Check bounds (zero padding) + if (in_h >= 0 && in_h < H_in && in_w >= 0 && in_w < W_in) { + // Get input value and quantize to int8 + int64_t input_idx = n * (C_in * H_in * W_in) + + in_c * (H_in * W_in) + in_h * W_in + in_w; + + float quant_input_f = + std::round(input_data[input_idx] / input_scale) + + input_zero_point; + quant_input_f = + std::min(std::max(quant_input_f, -128.0f), 127.0f); + int8_t quantized_input = static_cast(quant_input_f); + + // Get quantized weight (already int8) + // Weight layout: [C_out, C_in_per_group * K_h * K_w] + int64_t weight_idx = out_c * in_features + + (kh * (K_w * C_in_per_group) + kw * C_in_per_group + + (in_c % C_in_per_group)); + int8_t quantized_weight = weight_data[weight_idx]; + + // Integer multiplication and accumulation + int_sum += static_cast(quantized_input) * + static_cast(quantized_weight); + + // Track weight sum for this output channel on the fly + weight_sum += static_cast(quantized_weight); + } else { + // For zero padding, we still need to account for the weight + // in weight_sum when input is effectively 0 (but quantized 0 + // is input_zero_point) + int64_t weight_idx = out_c * in_features + + (kh * (K_w * C_in_per_group) + kw * C_in_per_group + + (in_c % C_in_per_group)); + int8_t quantized_weight = weight_data[weight_idx]; + + // Add contribution from zero-padded input (quantized zero = + // input_zero_point) + int_sum += static_cast(input_zero_point) * + static_cast(quantized_weight); + + // Track weight sum for this output channel on the fly + weight_sum += static_cast(quantized_weight); + } + } + } + } + + // Convert accumulated integer result to float and apply scales + // Final result = (int_sum - zero_point_correction) * input_scale * + // weight_scale + bias zero_point_correction = input_zero_point * + // sum_of_weights_for_this_output_channel + int32_t zero_point_correction = input_zero_point * weight_sum; + int32_t accum_adjusted = int_sum - zero_point_correction; + float float_result = + accum_adjusted * input_scale * weight_scales_data[out_c]; + + // Add bias and store result + float_result += bias_data[out_c]; + int64_t output_idx = n * (C_out * H_out * W_out) + + out_c * (H_out * W_out) + out_h * W_out + out_w; + ref_data[output_idx] = float_result; + } + } + } + } +} + +void reference_impl(TestCase& test_case) { + if (test_case.operator_name().find("q8ta") != std::string::npos) { + conv2d_q8ta_q8csw_reference_impl(test_case); + } else { + conv2d_q8csw_reference_impl(test_case); + } +} + +// Custom FLOP calculator for quantized conv2d operation +int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { + int kernel_idx = 4; + if (test_case.operator_name().find("q8ta") != std::string::npos) { + kernel_idx = 7; + } + // Get input and weight dimensions + const auto& input_sizes = test_case.inputs()[0].get_tensor_sizes(); + const auto& output_sizes = test_case.outputs()[0].get_tensor_sizes(); + + const auto& kernel_sizes = test_case.inputs()[kernel_idx].get_int32_data(); + + int64_t N = input_sizes[0]; + int64_t C_in = input_sizes[1]; + int64_t C_out = output_sizes[1]; + int64_t K_h = kernel_sizes[0]; + int64_t K_w = kernel_sizes[1]; + int64_t H_out = output_sizes[2]; + int64_t W_out = output_sizes[3]; + + // Calculate FLOPs for quantized conv2d operation + // Each output element requires: + // - C_in * K_h * K_w multiply-accumulate operations + // - Additional operations for quantization/dequantization + int64_t output_elements = N * C_out * H_out * W_out; + int64_t ops_per_output = C_in * K_h * K_w; + + int64_t flop = output_elements * (ops_per_output); + + return flop; +} + +int main(int argc, char* argv[]) { + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout << "Quantized Conv2d Operation Prototyping Framework" << std::endl; + print_separator(); + + ReferenceComputeFunc ref_fn = reference_impl; + + // Execute test cases using the new framework with custom FLOP calculator + auto results = execute_test_cases( + generate_quantized_conv2d_test_cases, + quantized_conv2d_flop_calculator, + "QuantizedConv2d", + 0, + 10, + ref_fn); + + return 0; +} diff --git a/backends/vulkan/test/custom_ops/q8csw_linear.cpp b/backends/vulkan/test/custom_ops/q8csw_linear.cpp new file mode 100644 index 00000000000..23973426fcc --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8csw_linear.cpp @@ -0,0 +1,479 @@ +// 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 "utils.h" + +#include + +using namespace executorch::vulkan::prototyping; + +using namespace vkcompute; + +static constexpr int64_t kRefDimSizeLimit = 300; + +// Linear configuration struct +struct LinearConfig { + int64_t M; // Batch size / number of rows in input + int64_t K; // Input features / columns in input, rows in weight + int64_t N; // Output features / columns in weight + bool has_bias = true; + std::string test_case_name = "placeholder"; + std::string op_name = "linear_q8ta_q8csw"; +}; + +// Utility function to create a test case from a LinearConfig +TestCase create_test_case_from_config( + const LinearConfig& config, + utils::StorageType storage_type, + vkapi::ScalarType input_dtype) { + TestCase test_case; + + // Create a descriptive name for the test case + std::string storage_str = + (storage_type == utils::kTexture3D) ? "Texture3D" : "Buffer"; + std::string dtype_str = (input_dtype == vkapi::kFloat) ? "Float" : "Half"; + + std::string test_name = + config.test_case_name + "_" + storage_str + "_" + dtype_str; + test_case.set_name(test_name); + + // Set the operator name for the test case + std::string operator_name = "et_vk." + config.op_name + ".default"; + test_case.set_operator_name(operator_name); + + // Derive sizes from M, K, N + std::vector input_size = {config.M, config.K}; + std::vector weight_size = {config.N, config.K}; + + // Input tensor (float/half) - [M, K] + ValueSpec input_tensor( + input_size, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + + if (debugging()) { + print_valuespec_data(input_tensor, "input_tensor"); + } + + float input_scale_val = 0.008f; + ValueSpec input_scale(input_scale_val); + + int32_t input_zero_point_val = -2; + ValueSpec input_zero_point(input_zero_point_val); + + // Quantized weight tensor (int8) - [K, N] + ValueSpec quantized_weight( + weight_size, + vkapi::kChar, // int8 for quantized weights + storage_type, + utils::kWidthPacked, + DataGenType::RANDINT8); + quantized_weight.set_constant(true); + + if (debugging()) { + print_valuespec_data(quantized_weight, "weight_tensor"); + } + + // Weight quantization scales (float/half, per-channel) + ValueSpec weight_scales( + {config.N}, // Per output feature + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM_SCALES); + weight_scales.set_constant(true); + + ValueSpec weight_sums( + {config.N}, // Per output features + vkapi::kInt, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + weight_sums.set_constant(true); + + // Compute weight_sums data based on quantized weights + int64_t in_features = config.K; + int64_t out_features = config.N; + compute_weight_sums(weight_sums, quantized_weight, out_features, in_features); + + // Bias (optional, float/half) - [N] + ValueSpec bias( + {config.N}, // Per output feature + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + bias.set_constant(true); + if (!config.has_bias) { + bias.set_none(true); + } + + // Output tensor (float/half) - [M, N] + ValueSpec output( + {config.M, config.N}, + input_dtype, + storage_type, + utils::kWidthPacked, + DataGenType::ZEROS); + + // Add all specs to test case + if (config.op_name.find("q8ta") != std::string::npos) { + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(input_scale); + test_case.add_input_spec(input_zero_point); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_sums); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(bias); + test_case.add_output_spec(output); + } else { + test_case.add_input_spec(input_tensor); + test_case.add_input_spec(quantized_weight); + test_case.add_input_spec(weight_scales); + test_case.add_input_spec(bias); + test_case.add_output_spec(output); + } + + return test_case; +} + +// Generate easy test cases for quantized linear operation (for debugging) +std::vector generate_quantized_linear_easy_cases() { + std::vector test_cases; + + // Single simple configuration for debugging + int M = 4; + int K = 4; + int N = 4; + + LinearConfig config = { + M, // Batch size + K, // Input features + N, // Output features + true, // has_bias + "simple", // test_case_name + }; + + // Test with both storage types and data types for completeness + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + std::vector float_types = {vkapi::kFloat}; + + // Generate test cases for each combination + for (const auto& storage_type : storage_types) { + for (const auto& input_dtype : float_types) { + test_cases.push_back( + create_test_case_from_config(config, storage_type, input_dtype)); + } + } + + return test_cases; +} + +// Generate test cases for quantized linear operation +std::vector generate_quantized_linear_test_cases() { + std::vector test_cases; + + std::vector configs = { + {4, 64, 32}, + {4, 128, 64}, + {4, 256, 128}, + {32, 64, 32}, + {32, 128, 64}, + {32, 256, 128}, + // No bias tests + {32, 128, 64, false}, + {32, 256, 128, false}, + {256, 2048, 2048}, + {512, 2048, 2048}, + {1024, 2048, 2048}, + }; + + // Test with different storage types and data types + std::vector storage_types = { + utils::kTexture3D, utils::kBuffer}; + + for (auto config : configs) { + std::string prefix = + (config.M < kRefDimSizeLimit && config.K < kRefDimSizeLimit && + config.N < kRefDimSizeLimit) + ? "correctness_" + : "performance_"; + std::string generated_test_case_name = prefix + std::to_string(config.M) + + "_" + std::to_string(config.K) + "_" + std::to_string(config.N); + if (!config.has_bias) { + generated_test_case_name += "_no_bias"; + } + + config.test_case_name = generated_test_case_name; + + for (const auto& storage_type : storage_types) { + if (vkcompute::api::context() + ->adapter_ptr() + ->supports_int8_dot_product()) { + // Test both activation+weight quantized and weight only quantized + test_cases.push_back( + create_test_case_from_config(config, storage_type, vkapi::kFloat)); + } + + LinearConfig wo_quant_config = config; + wo_quant_config.op_name = "linear_q8csw"; + test_cases.push_back(create_test_case_from_config( + wo_quant_config, storage_type, vkapi::kFloat)); + } + } + + return test_cases; +} + +// Reference implementation for weight only quantized linear +void linear_q8csw_reference_impl(TestCase& test_case) { + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [batch_size, in_features] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [out_features, in_features] + auto output_sizes = + output_spec.get_tensor_sizes(); // [batch_size, out_features] + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = weight_sizes[0]; + + // Skip for large tensors since computation time will be extremely slow + if (batch_size > kRefDimSizeLimit || in_features > kRefDimSizeLimit || + out_features > kRefDimSizeLimit) { + throw std::invalid_argument( + "One or more dimensions (batch_size, in_features, out_features) exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + + 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(); + + // Calculate number of output elements + int64_t num_output_elements = batch_size * out_features; + + auto& ref_data = output_spec.get_ref_float_data(); + ref_data.resize(num_output_elements); + + // Perform quantized linear transformation (matrix multiplication) + for (int64_t b = 0; b < batch_size; ++b) { + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + float sum = 0.0f; + + // Matrix multiplication: output[b][out_f] = sum(input[b][in_f] * + // weight[out_f][in_f]) + for (int64_t in_f = 0; in_f < in_features; ++in_f) { + // Get input value and dequantize + int64_t input_idx = b * in_features + in_f; + float input_val = input_data[input_idx]; + + // Get weight value and dequantize + int64_t weight_idx = out_f * in_features + in_f; + float dequant_weight = (static_cast(weight_data[weight_idx])) * + weight_scales_data[out_f]; + + sum += input_val * dequant_weight; + } + + // Add bias and store result + if (!bias_spec.is_none()) { + sum += bias_data[out_f]; + } + int64_t output_idx = b * out_features + out_f; + ref_data[output_idx] = sum; + } + } +} + +void linear_q8ta_q8csw_reference_impl(TestCase& test_case) { + // Extract input specifications + int32_t idx = 0; + const ValueSpec& input_spec = test_case.inputs()[idx++]; + const ValueSpec& input_scale_spec = test_case.inputs()[idx++]; + const ValueSpec& input_zeros_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_spec = test_case.inputs()[idx++]; + const ValueSpec& weight_sums_spec = test_case.inputs()[idx++]; + (void)weight_sums_spec; + const ValueSpec& weight_scales_spec = test_case.inputs()[idx++]; + const ValueSpec& bias_spec = test_case.inputs()[idx++]; + + // Extract output specification (mutable reference) + ValueSpec& output_spec = test_case.outputs()[0]; + + // Get tensor dimensions + auto input_sizes = input_spec.get_tensor_sizes(); // [batch_size, in_features] + auto weight_sizes = + weight_spec.get_tensor_sizes(); // [out_features, in_features] + auto output_sizes = + output_spec.get_tensor_sizes(); // [batch_size, out_features] + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = weight_sizes[0]; + + // Skip for large tensors since computation time will be extremely slow + if (batch_size > kRefDimSizeLimit || in_features > kRefDimSizeLimit || + out_features > kRefDimSizeLimit) { + throw std::invalid_argument( + "One or more dimensions (batch_size, in_features, out_features) exceed the allowed limit for reference implementation."); + } + + if (input_spec.dtype != vkapi::kFloat) { + throw std::invalid_argument("Unsupported dtype"); + } + + // Get raw data pointers + auto& input_data = input_spec.get_float_data(); + const float input_scale = input_scale_spec.get_float_value(); + const int32_t input_zero_point = input_zeros_spec.get_int_value(); + + 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(); + + // Calculate number of output elements + int64_t num_output_elements = batch_size * out_features; + + auto& ref_data = output_spec.get_ref_float_data(); + ref_data.resize(num_output_elements); + + // Perform quantized linear transformation (matrix multiplication) with + // integer accumulation + for (int64_t b = 0; b < batch_size; ++b) { + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + int32_t int_sum = 0; + int32_t weight_sum = 0; // Track weight sum on the fly + + // Matrix multiplication with integer accumulation: + // int_sum = sum(quantized_input[b][in_f] * quantized_weight[out_f][in_f]) + for (int64_t in_f = 0; in_f < in_features; ++in_f) { + // Get input value and quantize to int8 + int64_t input_idx = b * in_features + in_f; + + float quant_input_f = + std::round(input_data[input_idx] / input_scale) + input_zero_point; + quant_input_f = std::min(std::max(quant_input_f, -128.0f), 127.0f); + int8_t quantized_input = static_cast(quant_input_f); + + // Get quantized weight (already int8) + int64_t weight_idx = out_f * in_features + in_f; + int8_t quantized_weight = weight_data[weight_idx]; + + // Integer multiplication and accumulation + int_sum += static_cast(quantized_input) * + static_cast(quantized_weight); + + // Track weight sum for this output channel on the fly + weight_sum += static_cast(quantized_weight); + } + + // Convert accumulated integer result to float and apply scales + // Final result = (int_sum - zero_point_correction) * input_scale * + // weight_scale + bias zero_point_correction = input_zero_point * + // sum_of_weights_for_this_output_channel + int32_t zero_point_correction = input_zero_point * weight_sum; + int32_t accum_adjusted = int_sum - zero_point_correction; + + float float_result = + accum_adjusted * input_scale * weight_scales_data[out_f]; + + // Add bias and store result + if (!bias_spec.is_none()) { + float_result += bias_data[out_f]; + } + int64_t output_idx = b * out_features + out_f; + ref_data[output_idx] = float_result; + } + } +} + +void reference_impl(TestCase& test_case) { + if (test_case.operator_name().find("q8ta") != std::string::npos) { + linear_q8ta_q8csw_reference_impl(test_case); + } else { + linear_q8csw_reference_impl(test_case); + } +} + +int64_t quantized_linear_flop_calculator(const TestCase& test_case) { + int input_idx = 0; + int weight_idx = 1; + if (test_case.operator_name().find("q8ta") != std::string::npos) { + input_idx = 0; + weight_idx = 3; + } + + // Get input and weight dimensions + const auto& input_sizes = test_case.inputs()[input_idx].get_tensor_sizes(); + const auto& weight_sizes = test_case.inputs()[weight_idx].get_tensor_sizes(); + + int64_t batch_size = input_sizes[0]; + int64_t in_features = input_sizes[1]; + int64_t out_features = weight_sizes[0]; + + // Calculate FLOPs for quantized linear operation + // Each output element requires: + // - in_features multiply-accumulate operations + // - Additional operations for quantization/dequantization + int64_t output_elements = batch_size * out_features; + int64_t ops_per_output = in_features; + + // Add quantization overhead (approximate) + // - Dequantize input: 1 op per input element used + // - Dequantize weight: 1 op per weight element used + // - Add bias: 1 op per output element + int64_t quantization_ops = ops_per_output + 1; // Simplified estimate + + int64_t flop = output_elements * (ops_per_output + quantization_ops); + + return flop; +} + +int main(int argc, char* argv[]) { + set_debugging(false); + set_print_output(false); + set_print_latencies(false); + set_use_gpu_timestamps(true); + + print_performance_header(); + std::cout << "Quantized Linear Operation Prototyping Framework" << std::endl; + print_separator(); + + ReferenceComputeFunc ref_fn = reference_impl; + + auto results = execute_test_cases( + generate_quantized_linear_test_cases, + quantized_linear_flop_calculator, + "QuantizedLinear", + 0, + 10, + ref_fn); + + return 0; +} diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl new file mode 100644 index 00000000000..3162857c2d3 --- /dev/null +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -0,0 +1,99 @@ +load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load( + "@fbsource//xplat/executorch/backends/vulkan:targets.bzl", + "get_platforms", + "vulkan_spv_shader_lib", +) + +def define_custom_op_test_binary(custom_op_name, extra_deps = [], src_file = None): + deps_list = [ + ":prototyping_utils", + ":operator_implementations", + ":custom_ops_shaderlib", + "//executorch/backends/vulkan:vulkan_graph_runtime", + runtime.external_dep_location("libtorch"), + ] + extra_deps + + src_file_str = src_file if src_file else "{}.cpp".format(custom_op_name) + + runtime.cxx_binary( + name = custom_op_name, + srcs = [ + src_file_str, + ], + platforms = get_platforms(), + define_static_target = False, + deps = deps_list, + ) + +def define_common_targets(is_fbcode = False): + if is_fbcode: + return + + # Shader library from GLSL files + runtime.filegroup( + name = "custom_ops_shaders", + srcs = native.glob([ + "glsl/*.glsl", + "glsl/*.yaml", + ]), + visibility = [ + "//executorch/backends/vulkan/test/custom_ops/...", + "@EXECUTORCH_CLIENTS", + ], + ) + + vulkan_spv_shader_lib( + name = "custom_ops_shaderlib", + spv_filegroups = { + ":custom_ops_shaders": "glsl", + }, + is_fbcode = is_fbcode, + ) + + # Prototyping utilities library + runtime.cxx_library( + name = "prototyping_utils", + srcs = [ + "utils.cpp", + ], + headers = [ + "utils.h", + ], + exported_headers = [ + "utils.h", + ], + platforms = get_platforms(), + deps = [ + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + visibility = [ + "//executorch/backends/vulkan/test/custom_ops/...", + "@EXECUTORCH_CLIENTS", + ], + ) + + # Operator implementations library + runtime.cxx_library( + name = "operator_implementations", + srcs = native.glob([ + "impl/*.cpp", + ]), + platforms = get_platforms(), + deps = [ + "//executorch/backends/vulkan:vulkan_graph_runtime", + ":custom_ops_shaderlib", + ], + visibility = [ + "//executorch/backends/vulkan/test/custom_ops/...", + "@EXECUTORCH_CLIENTS", + ], + link_whole = True, + ) + + define_custom_op_test_binary("add") + define_custom_op_test_binary("q8csw_linear") + define_custom_op_test_binary("q8csw_conv2d") + define_custom_op_test_binary("choose_qparams_per_row") + define_custom_op_test_binary("q4gsw_linear") diff --git a/backends/vulkan/test/custom_ops/utils.cpp b/backends/vulkan/test/custom_ops/utils.cpp new file mode 100644 index 00000000000..2aa827a4d5a --- /dev/null +++ b/backends/vulkan/test/custom_ops/utils.cpp @@ -0,0 +1,1776 @@ +// 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 "utils.h" +#include +#include +#include +#include + +#include + +namespace executorch { +namespace vulkan { +namespace prototyping { + +int get_seed() { + static int seed = 42; + return seed++; +} + +// Forward declarations for data generation utilities +void generate_random_float_data( + std::vector& data, + float min_val = -1.0f, + float max_val = 1.0f); +void generate_random_int_data( + std::vector& data, + int min_val = -10, + int max_val = 10); +void generate_randint_float_data( + std::vector& data, + int min_val = -10, + int max_val = 10); +void generate_randint_half_data( + std::vector& data, + int min_val = -10, + int max_val = 10); +void generate_random_int8_data( + std::vector& data, + int8_t min_val = -10, + int8_t max_val = 10); +void generate_random_uint8_data( + std::vector& data, + uint8_t min_val = 0, + uint8_t max_val = 255); +void generate_random_2xint4_data(std::vector& data); +void generate_random_2xint4_data(std::vector& data); +void generate_random_int4_data( + std::vector& data, + int8_t min_val = -8, + int8_t max_val = 7); +void generate_ones_data(std::vector& data); +void generate_zeros_data(std::vector& data); + +// Output and latency printing utilities +namespace { +static int print_output_enabled = 0; +static int print_latencies_enabled = 0; +static int gpu_timestamps_enabled = 0; +static int debugging_enabled = 0; +} // namespace + +bool print_output() { + return print_output_enabled > 0; +} + +void set_print_output(bool print_output) { + print_output_enabled = print_output ? 1 : 0; +} + +bool print_latencies() { + return print_latencies_enabled > 0; +} + +void set_print_latencies(bool print_latencies) { + print_latencies_enabled = print_latencies ? 1 : 0; +} + +bool use_gpu_timestamps() { + return gpu_timestamps_enabled > 0; +} + +void set_use_gpu_timestamps(bool use_timestamps) { + gpu_timestamps_enabled = use_timestamps ? 1 : 0; +} + +bool debugging() { + return debugging_enabled > 0; +} + +void set_debugging(bool enable_debugging) { + debugging_enabled = enable_debugging ? 1 : 0; +} + +// ValueSpec implementation +void ValueSpec::generate_tensor_data() { + if (spec_type != SpecType::Tensor) { + return; + } + + int64_t num_elements = numel(); + + switch (dtype) { + case vkapi::kFloat: { + float_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + generate_random_float_data(float_data); + } else if (data_gen_type == DataGenType::RANDOM_SCALES) { + generate_random_float_data(float_data, 0.005, 0.015); + } else if (data_gen_type == DataGenType::RANDINT) { + generate_randint_float_data(float_data); + } else if (data_gen_type == DataGenType::RANDINT8) { + generate_randint_float_data(float_data, -128, 127); + } else if (data_gen_type == DataGenType::RANDINT4) { + generate_randint_float_data(float_data, -8, 7); + } else if (data_gen_type == DataGenType::ONES) { + generate_ones_data(float_data); + } else if (data_gen_type == DataGenType::ZEROS) { + generate_zeros_data(float_data); + } else { + generate_zeros_data(float_data); + } + break; + } + case vkapi::kHalf: { + half_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + // Generate random float data first, then convert to half + std::vector temp_data(num_elements); + generate_random_float_data(temp_data); + for (size_t i = 0; i < temp_data.size(); ++i) { + // Simple conversion to uint16_t representation of half + half_data[i] = static_cast(temp_data[i] * 32767.0f); + } + } else if (data_gen_type == DataGenType::RANDINT) { + generate_randint_half_data(half_data); + } else if (data_gen_type == DataGenType::RANDINT8) { + generate_randint_half_data(half_data, -128, 127); + } else if (data_gen_type == DataGenType::RANDINT4) { + generate_randint_half_data(half_data, -8, 7); + } else if (data_gen_type == DataGenType::ONES) { + std::fill( + half_data.begin(), + half_data.end(), + static_cast(32767)); // 1.0 in half + } else if (data_gen_type == DataGenType::ZEROS) { + std::fill( + half_data.begin(), + half_data.end(), + static_cast(0)); // 0.0 in half + } else { + std::fill( + half_data.begin(), + half_data.end(), + static_cast(0)); // 0.0 in half + } + break; + } + case vkapi::kInt: { + int32_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + generate_random_int_data(int32_data); + } else if (data_gen_type == DataGenType::RANDINT) { + generate_random_int_data( + int32_data); // For int type, RANDINT is same as RANDOM + } else if (data_gen_type == DataGenType::RANDINT8) { + generate_random_int_data(int32_data, -128, 127); + } else if (data_gen_type == DataGenType::RANDINT4) { + generate_random_int_data(int32_data, -8, 7); + } else if (data_gen_type == DataGenType::ONES) { + std::fill(int32_data.begin(), int32_data.end(), 1); + } else if (data_gen_type == DataGenType::ZEROS) { + std::fill(int32_data.begin(), int32_data.end(), 0); + } else { + std::fill(int32_data.begin(), int32_data.end(), 0); + } + break; + } + case vkapi::kChar: { + int8_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + generate_random_int8_data(int8_data); + } else if (data_gen_type == DataGenType::RANDINT) { + generate_random_int8_data(int8_data); + } else if (data_gen_type == DataGenType::RANDINT8) { + generate_random_int8_data(int8_data, -128, 127); + } else if (data_gen_type == DataGenType::RANDINT4) { + generate_random_2xint4_data(int8_data); + } else if (data_gen_type == DataGenType::ONES) { + std::fill(int8_data.begin(), int8_data.end(), 1); + } else if (data_gen_type == DataGenType::ONES_INT4) { + int8_t packed_data = (1 << 4) | 1; + std::fill(int8_data.begin(), int8_data.end(), packed_data); + } else if (data_gen_type == DataGenType::ZEROS) { + std::fill(int8_data.begin(), int8_data.end(), 0); + } else { + std::fill(int8_data.begin(), int8_data.end(), 0); + } + break; + } + case vkapi::kByte: { + uint8_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + generate_random_uint8_data(uint8_data); + } else if (data_gen_type == DataGenType::RANDINT) { + generate_random_uint8_data(uint8_data); + } else if (data_gen_type == DataGenType::RANDINT8) { + generate_random_uint8_data(uint8_data, 0, 255); + } else if (data_gen_type == DataGenType::RANDINT4) { + generate_random_2xint4_data(uint8_data); + } else if (data_gen_type == DataGenType::ONES) { + std::fill(uint8_data.begin(), uint8_data.end(), 1); + } else if (data_gen_type == DataGenType::ONES_INT4) { + uint8_t packed_data = (9 << 4) | 9; + std::fill(uint8_data.begin(), uint8_data.end(), packed_data); + } else if (data_gen_type == DataGenType::ZEROS) { + std::fill(uint8_data.begin(), uint8_data.end(), 0); + } else { + std::fill(uint8_data.begin(), uint8_data.end(), 0); + } + break; + } + default: + // Default to float + float_data.resize(num_elements); + if (data_gen_type == DataGenType::RANDOM) { + generate_random_float_data(float_data); + } else if (data_gen_type == DataGenType::RANDINT) { + generate_randint_float_data(float_data); + } else if (data_gen_type == DataGenType::ONES) { + generate_ones_data(float_data); + } else if (data_gen_type == DataGenType::ZEROS) { + generate_zeros_data(float_data); + } else { + generate_zeros_data(float_data); + } + break; + } +} + +int64_t ValueSpec::numel() const { + if (spec_type == SpecType::Int || spec_type == SpecType::Float || + spec_type == SpecType::Bool) { + return 1; + } else if (spec_type == SpecType::IntList) { + return sizes.empty() ? 0 : sizes[0]; + } else { // Tensor + int64_t total = 1; + for (int64_t size : sizes) { + total *= size; + } + return total; + } +} + +size_t ValueSpec::nbytes() const { + size_t element_size = 0; + switch (dtype) { + case vkapi::kFloat: + element_size = sizeof(float); + break; + case vkapi::kHalf: + element_size = sizeof(uint16_t); + break; + case vkapi::kInt: + element_size = sizeof(int32_t); + break; + case vkapi::kChar: + element_size = sizeof(int8_t); + break; + case vkapi::kByte: + element_size = sizeof(uint8_t); + break; + default: + element_size = sizeof(float); // Default fallback + break; + } + return numel() * element_size; +} + +std::string ValueSpec::to_string() const { + std::string result = "ValueSpec("; + + switch (spec_type) { + case SpecType::Tensor: + result += "type=Tensor, sizes=["; + break; + case SpecType::IntList: + result += "type=IntList, count="; + result += std::to_string(sizes.empty() ? 0 : sizes[0]); + result += ", data_gen="; + result += (data_gen_type == DataGenType::FIXED) ? "FIXED" : "RANDOM"; + result += ")"; + return result; + case SpecType::Int: + result += "type=Int, value="; + result += std::to_string(get_int_value()); + result += ", data_gen="; + result += (data_gen_type == DataGenType::FIXED) ? "FIXED" : "RANDOM"; + result += ")"; + return result; + case SpecType::Float: + result += "type=Float, value="; + result += std::to_string(get_float_value()); + result += ", data_gen="; + result += (data_gen_type == DataGenType::FIXED) ? "FIXED" : "RANDOM"; + result += ")"; + return result; + case SpecType::Bool: + result += "type=Bool, value="; + result += get_bool_value() ? "true" : "false"; + result += ", data_gen="; + result += (data_gen_type == DataGenType::FIXED) ? "FIXED" : "RANDOM"; + result += ")"; + return result; + } + + for (size_t i = 0; i < sizes.size(); ++i) { + result += std::to_string(sizes[i]); + if (i < sizes.size() - 1) + result += ", "; + } + result += "]"; + + if (spec_type == SpecType::Tensor) { + result += ", dtype="; + switch (dtype) { + case vkapi::kFloat: + result += "float"; + break; + case vkapi::kHalf: + result += "half"; + break; + case vkapi::kInt: + result += "int32"; + break; + case vkapi::kChar: + result += "int8"; + break; + case vkapi::kByte: + result += "uint8"; + break; + default: + result += "unknown"; + break; + } + + result += ", memory_layout="; + switch (memory_layout) { + case utils::kWidthPacked: + result += "WidthPacked"; + break; + case utils::kHeightPacked: + result += "HeightPacked"; + break; + case utils::kChannelsPacked: + result += "ChannelsPacked"; + break; + default: + result += "unknown"; + break; + } + + result += ", storage_type="; + switch (storage_type) { + case utils::kTexture3D: + result += "Texture3D"; + break; + case utils::kBuffer: + result += "Buffer"; + break; + default: + result += "unknown"; + break; + } + } + + result += ", data_gen="; + switch (data_gen_type) { + case DataGenType::FIXED: + result += "FIXED"; + break; + case DataGenType::RANDOM: + result += "RANDOM"; + break; + case DataGenType::RANDINT: + result += "RANDINT"; + break; + case DataGenType::RANDINT8: + result += "RANDINT8"; + break; + case DataGenType::RANDINT4: + result += "RANDINT4"; + break; + case DataGenType::ONES: + result += "ONES"; + break; + case DataGenType::ZEROS: + result += "ZEROS"; + break; + default: + result += "unknown"; + break; + } + result += ")"; + return result; +} + +// Additional ValueSpec methods +void ValueSpec::resize_data(size_t new_size) { + switch (dtype) { + case vkapi::kFloat: + float_data.resize(new_size); + break; + case vkapi::kHalf: + half_data.resize(new_size); + break; + case vkapi::kInt: + int32_data.resize(new_size); + break; + case vkapi::kChar: + int8_data.resize(new_size); + break; + case vkapi::kByte: + uint8_data.resize(new_size); + break; + default: + float_data.resize(new_size); + break; + } +} + +void* ValueSpec::get_mutable_data_ptr() { + switch (dtype) { + case vkapi::kFloat: + return float_data.data(); + case vkapi::kHalf: + return half_data.data(); + case vkapi::kInt: + return int32_data.data(); + case vkapi::kChar: + return int8_data.data(); + case vkapi::kByte: + return uint8_data.data(); + default: + return float_data.data(); + } +} + +float ValueSpec::get_element(size_t index) const { + if (index >= static_cast(numel())) { + return 0.0f; + } + + switch (dtype) { + case vkapi::kFloat: + return index < float_data.size() ? float_data[index] : 0.0f; + case vkapi::kHalf: + return index < half_data.size() ? (half_data[index] / 32767.0f) : 0.0f; + case vkapi::kInt: + return index < int32_data.size() ? static_cast(int32_data[index]) + : 0.0f; + case vkapi::kChar: + return index < int8_data.size() ? static_cast(int8_data[index]) + : 0.0f; + case vkapi::kByte: + return index < uint8_data.size() ? static_cast(uint8_data[index]) + : 0.0f; + default: + return 0.0f; + } +} + +const void* ValueSpec::get_data_ptr() const { + switch (dtype) { + case vkapi::kFloat: + return float_data.data(); + case vkapi::kHalf: + return half_data.data(); + case vkapi::kInt: + return int32_data.data(); + case vkapi::kChar: + return int8_data.data(); + case vkapi::kByte: + return uint8_data.data(); + default: + throw std::runtime_error("Unsupported data type for get_data_ptr"); + } +} + +void generate_random_float_data( + std::vector& data, + float min_val, + float max_val) { + std::mt19937 gen(get_seed()); + std::uniform_real_distribution dis(min_val, max_val); + for (auto& val : data) { + val = dis(gen); + } +} + +void generate_random_int_data( + std::vector& data, + int min_val, + int max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = dis(gen); + } +} + +void generate_randint_float_data( + std::vector& data, + int min_val, + int max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = static_cast(dis(gen)); + } +} + +void generate_randint_half_data( + std::vector& data, + int min_val, + int max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = static_cast(std::abs(dis(gen)) % 65536); + } +} + +void generate_ones_data(std::vector& data) { + std::fill(data.begin(), data.end(), 1.0f); +} + +void generate_random_int8_data( + std::vector& data, + int8_t min_val, + int8_t max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = static_cast(dis(gen)); + } +} + +void generate_random_uint8_data( + std::vector& data, + uint8_t min_val, + uint8_t max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = static_cast(dis(gen)); + } +} + +void generate_random_int4_data( + std::vector& data, + int8_t min_val, + int8_t max_val) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(min_val, max_val); + for (auto& val : data) { + val = static_cast(dis(gen)); + } +} + +void generate_random_2xint4_data(std::vector& data) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(-8, 7); // Signed 4-bit range + for (auto& val : data) { + // Generate two separate 4-bit values + int8_t lower_4bits = static_cast(dis(gen)) & 0x0F; + int8_t upper_4bits = static_cast(dis(gen)) & 0x0F; + // Pack them into a single 8-bit value + val = (upper_4bits << 4) | lower_4bits; + } +} + +void generate_random_2xint4_data(std::vector& data) { + std::mt19937 gen(get_seed()); + std::uniform_int_distribution dis(0, 15); // Unsigned 4-bit range + for (auto& val : data) { + // Generate two separate 4-bit values + uint8_t lower_4bits = static_cast(dis(gen)) & 0x0F; + uint8_t upper_4bits = static_cast(dis(gen)) & 0x0F; + // Pack them into a single 8-bit value + val = (upper_4bits << 4) | lower_4bits; + } +} + +void generate_zeros_data(std::vector& data) { + std::fill(data.begin(), data.end(), 0.0f); +} + +// Correctness checking against reference data +bool ValueSpec::validate_against_reference( + float abs_tolerance, + float rel_tolerance) const { + // Only validate float tensors as specified in requirements + if (dtype != vkapi::kFloat || !is_tensor()) { + return true; // Skip validation for non-float or non-tensor types + } + + const auto& computed_data = get_float_data(); + const auto& reference_data = get_ref_float_data(); + + // Skip validation if no reference data is available + if (reference_data.empty()) { + return true; + } + + // Check if sizes match + if (computed_data.size() != reference_data.size()) { + if (debugging()) { + std::cout << "Size mismatch: computed=" << computed_data.size() + << ", reference=" << reference_data.size() << std::endl; + } + return false; + } + + // Element-wise comparison with both absolute and relative tolerance + for (size_t i = 0; i < computed_data.size(); ++i) { + float diff = std::abs(computed_data[i] - reference_data[i]); + float abs_ref = std::abs(reference_data[i]); + + // Check if either absolute or relative tolerance condition is satisfied + bool abs_tolerance_ok = diff <= abs_tolerance; + bool rel_tolerance_ok = diff <= rel_tolerance * abs_ref; + + if (!abs_tolerance_ok && !rel_tolerance_ok) { + std::cout << "Mismatch at element " << i + << ": computed=" << computed_data[i] + << ", reference=" << reference_data[i] << ", diff=" << diff + << ", abs_tolerance=" << abs_tolerance + << ", rel_tolerance=" << rel_tolerance + << ", rel_threshold=" << (rel_tolerance * abs_ref) << std::endl; + return false; + } + } + + if (debugging()) { + std::cout << "Correctness validation PASSED" << std::endl; + } + return true; +} + +// Helper function to collect GPU timing from querypool +float collect_gpu_timing_us(ComputeGraph& graph) { + graph.context()->querypool().extract_results(); + const auto results = graph.context()->querypool().get_shader_timestamp_data(); + if (!results.empty()) { + // Sum durations of all shaders that don't contain nchw_to or to_nchw + float total_duration_us = 0.0f; + for (const auto& shader_result : results) { + if (shader_result.kernel_name.find("nchw_to") == std::string::npos && + shader_result.kernel_name.find("to_nchw") == std::string::npos) { + // Calculate duration from start and end times, convert from ns to μs + uint64_t duration_ns = + shader_result.end_time_ns - shader_result.start_time_ns; + total_duration_us += static_cast(duration_ns) / 1000.0f; + } + } + return total_duration_us; + } + return 0.0f; +} + +// BenchmarkResult implementation +void BenchmarkResult::add_iter_timing(float time_us) { + iter_timings.push_back(time_us); +} + +float BenchmarkResult::get_avg_time_us() const { + if (iter_timings.empty()) { + return 0.0f; + } + + float sum = 0.0f; + for (float timing : iter_timings) { + sum += timing; + } + return sum / iter_timings.size(); +} + +float BenchmarkResult::get_min_time_us() const { + if (iter_timings.empty()) { + return 0.0f; + } + + return *std::min_element(iter_timings.begin(), iter_timings.end()); +} + +float BenchmarkResult::get_max_time_us() const { + if (iter_timings.empty()) { + return 0.0f; + } + + return *std::max_element(iter_timings.begin(), iter_timings.end()); +} + +float BenchmarkResult::get_std_dev_us() const { + if (iter_timings.size() <= 1) { + return 0.0f; + } + + float mean = get_avg_time_us(); + float sum_sq_diff = 0.0f; + + for (float timing : iter_timings) { + float diff = timing - mean; + sum_sq_diff += diff * diff; + } + + return std::sqrt(sum_sq_diff / (iter_timings.size() - 1)); +} + +void BenchmarkResult::print_summary( + int case_number, + const std::string& size_info, + float total_gflops) const { + static constexpr int OPERATOR_NAME_WIDTH = 50; + static constexpr int KERNEL_NAME_WIDTH = 70; + static constexpr int SIZE_INFO_WIDTH = 20; + static constexpr int TIMING_WIDTH = 20; + static constexpr int GFLOPS_WIDTH = 20; + static constexpr int CORRECTNESS_WIDTH = 10; + + std::string correctness_str; + switch (correctness_status_) { + case CorrectnessStatus::SKIPPED: + correctness_str = "SKIPPED"; + break; + case CorrectnessStatus::PASSED: + correctness_str = "PASSED"; + break; + case CorrectnessStatus::FAILED: + correctness_str = "FAILED"; + break; + } + + std::cout << std::left << std::setw(OPERATOR_NAME_WIDTH) + << get_operator_name() << " " << std::left + << std::setw(KERNEL_NAME_WIDTH) << get_kernel_name() << std::right + << " " << std::setw(SIZE_INFO_WIDTH) << size_info + << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) + << get_avg_time_us() << " μs " << std::setw(GFLOPS_WIDTH) + << std::fixed << std::setprecision(3) << total_gflops << " GFLOP/s " + << std::setw(CORRECTNESS_WIDTH) << correctness_str << std::endl; +} + +// TestResult implementation +void TestResult::add_result(const BenchmarkResult& result) { + results_.push_back(result); +} + +void TestResult::add_result(BenchmarkResult&& result) { + results_.push_back(std::move(result)); +} + +void TestResult::print_summary() const { + static constexpr int CASE_WIDTH = 80; + static constexpr int KERNEL_NAME_WIDTH = 20; + static constexpr int TIMING_WIDTH = 12; + static constexpr int PASS_WIDTH = 8; + + if (results_.empty()) { + std::cout << "No results to display" << std::endl; + return; + } + + std::cout << "\n=== " << operation_name_ + << " Performance Summary ===" << std::endl; + print_separator(); + + std::cout << std::left << std::setw(CASE_WIDTH) << "Case" << std::left + << std::setw(KERNEL_NAME_WIDTH) << "Kernel Name" << std::left + << std::setw(TIMING_WIDTH) << "Avg (μs)" << std::left + << std::setw(TIMING_WIDTH) << "Min (μs)" << std::left + << std::setw(TIMING_WIDTH) << "Max (μs)" << std::left + << std::setw(TIMING_WIDTH) << "Std Dev" << std::left + << std::setw(PASS_WIDTH) << "Pass" << std::endl; + print_separator(); + + for (size_t i = 0; i < results_.size(); ++i) { + const auto& result = results_[i]; + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + std::cout << std::left << std::setw(CASE_WIDTH) << i + 1 << std::left + << std::setw(KERNEL_NAME_WIDTH) + << result.get_kernel_name().substr(0, KERNEL_NAME_WIDTH - 1) + << std::left << std::setw(TIMING_WIDTH) << std::fixed + << std::setprecision(3) << result.get_avg_time_us() << std::left + << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) + << result.get_min_time_us() << std::left + << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) + << result.get_max_time_us() << std::left + << std::setw(TIMING_WIDTH) << std::fixed << std::setprecision(3) + << result.get_std_dev_us() << std::left << std::setw(PASS_WIDTH) + << (vulkan_execute_succeeded ? "✓" : "✗") << std::endl; + } + + print_separator(); + std::cout << "Total cases: " << results_.size() + << ", Passed: " << get_passed_count() + << ", Failed: " << get_failed_count() << std::endl; + std::cout << "Overall GFLOP/s: " << std::fixed << std::setprecision(3) + << gflops_ << std::endl; + std::cout << "Overall correctness: " + << (correctness_passed_ ? "PASSED" : "FAILED") << std::endl; +} + +void TestResult::print_detailed_results() const { + if (results_.empty()) { + std::cout << "No results to display" << std::endl; + return; + } + + std::cout << "\n=== " << operation_name_ + << " Detailed Results ===" << std::endl; + + for (size_t i = 0; i < results_.size(); ++i) { + const auto& result = results_[i]; + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + std::cout << "\nCase " << i + 1 << ": " << result.get_kernel_name() + << std::endl; + std::cout << " Iterations: " << result.get_num_iterations() << std::endl; + std::cout << " Average: " << std::fixed << std::setprecision(3) + << result.get_avg_time_us() << " μs" << std::endl; + std::cout << " Min: " << std::fixed << std::setprecision(3) + << result.get_min_time_us() << " μs" << std::endl; + std::cout << " Max: " << std::fixed << std::setprecision(3) + << result.get_max_time_us() << " μs" << std::endl; + std::cout << " Std Dev: " << std::fixed << std::setprecision(3) + << result.get_std_dev_us() << " μs" << std::endl; + std::cout << " Correctness: " + << (vulkan_execute_succeeded ? "PASSED" : "FAILED") << std::endl; + + if (result.get_num_iterations() > 0) { + std::cout << " Individual timings (μs): "; + const auto& timings = result.get_iter_timings(); + for (size_t j = 0; j < std::min(size_t(10), timings.size()); ++j) { + std::cout << std::fixed << std::setprecision(1) << timings[j]; + if (j < std::min(size_t(10), timings.size()) - 1) + std::cout << ", "; + } + if (timings.size() > 10) { + std::cout << " ... (" << (timings.size() - 10) << " more)"; + } + std::cout << std::endl; + } + } + + std::cout << "\nOverall Results:" << std::endl; + std::cout << " Total GFLOP/s: " << std::fixed << std::setprecision(3) + << gflops_ << std::endl; + std::cout << " Overall correctness: " + << (correctness_passed_ ? "PASSED" : "FAILED") << std::endl; +} + +void TestResult::print_statistics() const { + if (results_.empty()) { + std::cout << "No results to display statistics for" << std::endl; + return; + } + + std::cout << "\n=== " << operation_name_ << " Statistics ===" << std::endl; + std::cout << "Total test cases: " << results_.size() << std::endl; + std::cout << "Passed: " << get_passed_count() << std::endl; + std::cout << "Failed: " << get_failed_count() << std::endl; + std::cout << "Success rate: " << std::fixed << std::setprecision(1) + << (100.0f * get_passed_count() / results_.size()) << "%" + << std::endl; + + if (get_passed_count() > 0) { + std::cout << "Total average time: " << std::fixed << std::setprecision(3) + << get_total_avg_time_us() << " μs" << std::endl; + std::cout << "Total GFLOP/s: " << std::fixed << std::setprecision(3) + << get_total_gflops() << std::endl; + + const auto* fastest = get_fastest_result(); + const auto* slowest = get_slowest_result(); + const auto* highest_gflops = get_highest_gflops_result(); + + if (fastest) { + std::cout << "Fastest case: " << fastest->get_kernel_name() << " (" + << std::fixed << std::setprecision(3) + << fastest->get_avg_time_us() << " μs)" << std::endl; + } + + if (slowest) { + std::cout << "Slowest case: " << slowest->get_kernel_name() << " (" + << std::fixed << std::setprecision(3) + << slowest->get_avg_time_us() << " μs)" << std::endl; + } + + if (highest_gflops) { + std::cout << "Best performing case: " << highest_gflops->get_kernel_name() + << " (" << std::fixed << std::setprecision(3) + << highest_gflops->get_avg_time_us() << " μs)" << std::endl; + } + } +} + +void TestResult::print_brief_summary() const { + print_separator(); + std::cout << "Summary Statistics:" << std::endl; + + if (get_passed_count() > 0) { + std::cout << "Average execution time: " << std::fixed + << std::setprecision(3) << get_total_avg_time_us() << " μs" + << std::endl; + std::cout << "Total throughput: " << std::fixed << std::setprecision(3) + << get_gflops() << " GFLOP/s" << std::endl; + std::cout << "Successful test cases: " << get_passed_count() << "/" + << size() << std::endl; + std::cout << "Overall correctness: " + << (get_correctness_passed() ? "PASSED" : "FAILED") << std::endl; + } else { + std::cout << "No successful test cases to report" << std::endl; + } +} + +float TestResult::get_total_avg_time_us() const { + if (results_.empty()) { + return 0.0f; + } + + float sum = 0.0f; + size_t count = 0; + + for (const auto& result : results_) { + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + if (vulkan_execute_succeeded) { + sum += result.get_avg_time_us(); + count++; + } + } + + return count > 0 ? sum / count : 0.0f; +} + +float TestResult::get_total_gflops() const { + return gflops_; +} + +size_t TestResult::get_passed_count() const { + size_t count = 0; + for (const auto& result : results_) { + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + if (vulkan_execute_succeeded) { + count++; + } + } + return count; +} + +size_t TestResult::get_failed_count() const { + return results_.size() - get_passed_count(); +} + +const BenchmarkResult* TestResult::get_fastest_result() const { + const BenchmarkResult* fastest = nullptr; + + for (const auto& result : results_) { + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + if (vulkan_execute_succeeded) { + if (!fastest || result.get_avg_time_us() < fastest->get_avg_time_us()) { + fastest = &result; + } + } + } + + return fastest; +} + +const BenchmarkResult* TestResult::get_slowest_result() const { + const BenchmarkResult* slowest = nullptr; + + for (const auto& result : results_) { + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + if (vulkan_execute_succeeded) { + if (!slowest || result.get_avg_time_us() > slowest->get_avg_time_us()) { + slowest = &result; + } + } + } + + return slowest; +} + +const BenchmarkResult* TestResult::get_highest_gflops_result() const { + // Since GFLOPS is now a TestResult-level metric rather than per-case, + // this method returns the fastest result as a proxy for highest performance + return get_fastest_result(); +} + +// Default FLOP calculation function (assumes 1 FLOP per element) +int64_t default_flop_calculator(const TestCase& test_case) { + // Calculate total elements from the first input tensor + int64_t total_elements = 1; + if (!test_case.empty() && test_case.num_inputs() > 0 && + test_case.inputs()[0].is_tensor()) { + const auto& sizes = test_case.inputs()[0].get_tensor_sizes(); + for (int64_t size : sizes) { + total_elements *= size; + } + } + + // Assume 1 FLOP per element for basic operations + return total_elements; +} + +ComputeGraph setup_compute_graph(TestCase& test_case, std::string op_name) { + GraphConfig config; + config.enable_querypool = true; + ComputeGraph graph(config); + + std::vector input_values; + + // Process input ValueSpecs + for (size_t i = 0; i < test_case.num_inputs(); ++i) { + const ValueSpec& input_spec = test_case.inputs()[i]; + + if (input_spec.is_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())); + 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())); + input_values.push_back(input_value); + } else if (input_spec.is_bool()) { + 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 + const auto& int32_list = input_spec.get_int_list(); + std::vector int64_list; + int64_list.reserve(int32_list.size()); + for (int32_t val : int32_list) { + int64_list.push_back(static_cast(val)); + } + ValueRef input_value = graph.add_scalar_list(std::move(int64_list)); + input_values.push_back(input_value); + } else if (input_spec.is_constant()) { + 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( + input_spec.get_tensor_sizes(), + input_spec.dtype, + input_spec.storage_type, + input_spec.memory_layout); + input_values.push_back(input_io.value); + } + } + + std::vector output_values; + + // Process output ValueSpecs + for (size_t i = 0; i < test_case.num_outputs(); ++i) { + const ValueSpec& output_spec = test_case.outputs()[i]; + + if (!output_spec.is_tensor()) { + throw std::runtime_error("All output specifications must be tensors"); + } + + // Create output tensor + ValueRef output_value = graph.add_tensor( + output_spec.get_tensor_sizes(), + output_spec.dtype, + output_spec.storage_type, + output_spec.memory_layout); + + output_values.push_back(output_value); + } + + // Get the operator function and call it + auto opFn = VK_GET_OP_FN(op_name); + + // Create arguments vector for the operator function + std::vector op_args = input_values; + op_args.insert(op_args.end(), output_values.begin(), output_values.end()); + + opFn(graph, op_args); + + for (size_t i = 0; i < output_values.size(); ++i) { + graph.set_output_value(output_values[i]); + } + return graph; +} + +// Test execution utilities +BenchmarkResult +execute_test_case(TestCase& test_case, int warmup_runs, int benchmark_runs) { + BenchmarkResult result( + test_case.name().empty() ? "unnamed_test_case" : test_case.name()); + + // Initialize querypool if using GPU timestamps + if (use_gpu_timestamps()) { + api::context()->initialize_querypool(); + } + + // Create the compute graph for this test case using setup_compute_graph + ComputeGraph graph = + setup_compute_graph(test_case, test_case.operator_name()); + + // Prepare the graph + graph.prepare(); + graph.prepack(); + + // Copy input data into the graph's staging buffers + for (size_t i = 0; i < test_case.num_inputs(); ++i) { + const ValueSpec& input_spec = test_case.inputs()[i]; + if (input_spec.is_tensor() && i < graph.inputs().size()) { + // Skip copying data for constant tensors + if (input_spec.is_constant()) { + continue; + } + + const auto& input_ref = graph.inputs()[i]; + + // Get the appropriate data based on dtype + const void* data_ptr = nullptr; + size_t data_numel = input_spec.numel(); + + switch (input_spec.dtype) { + case vkapi::kFloat: + data_ptr = input_spec.get_float_data().data(); + break; + case vkapi::kHalf: + data_ptr = input_spec.get_half_data().data(); + break; + case vkapi::kInt: + data_ptr = input_spec.get_int32_data().data(); + break; + case vkapi::kChar: + data_ptr = input_spec.get_int8_data().data(); + break; + case vkapi::kByte: + data_ptr = input_spec.get_uint8_data().data(); + break; + default: + throw std::runtime_error("Unsupported data type for input tensor"); + } + + // Copy data into staging buffer + graph.copy_into_staging(input_ref.staging, data_ptr, data_numel); + } + } + + // Warmup runs + for (int run = 0; run < warmup_runs; ++run) { + graph.execute(); + } + + // Benchmark runs - collect individual iteration timings + float total_cpu_time_us = 0.0f; + float total_gpu_time_us = 0.0f; + + 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(); + auto cpu_end = std::chrono::high_resolution_clock::now(); + + auto cpu_duration = std::chrono::duration_cast( + cpu_end - cpu_start); + float cpu_time_us = static_cast(cpu_duration.count()); + total_cpu_time_us += cpu_time_us; + + // Collect GPU timing using helper function + float gpu_time_us = collect_gpu_timing_us(graph); + total_gpu_time_us += gpu_time_us; + + // Add the appropriate timing based on the flag + float iter_time_us = use_gpu_timestamps() ? gpu_time_us : cpu_time_us; + result.add_iter_timing(iter_time_us); + } + + // Calculate averages for display + float avg_cpu_time_us = total_cpu_time_us / benchmark_runs; + float avg_gpu_time_us = total_gpu_time_us / benchmark_runs; + + // Print both timings if latency printing is enabled + if (print_latencies()) { + if (use_gpu_timestamps()) { + graph.context()->querypool().print_results(); + } + std::cout << " CPU timing: " << std::fixed << std::setprecision(3) + << avg_cpu_time_us << " μs" << std::endl; + std::cout << " GPU timing: " << std::fixed << std::setprecision(3) + << avg_gpu_time_us << " μs" << std::endl; + std::cout << " Using " << (use_gpu_timestamps() ? "GPU" : "CPU") + << " timing for result" << std::endl; + } + + // Copy output data from the graph's staging buffers + for (size_t i = 0; i < test_case.num_outputs(); ++i) { + ValueSpec& output_spec = test_case.outputs()[i]; + + if (output_spec.is_tensor() && i < graph.outputs().size()) { + const auto& output_ref = graph.outputs()[i]; + + // Ensure output data vector is properly sized + size_t data_numel = output_spec.numel(); + output_spec.resize_data(data_numel); + + // Get mutable data pointer for the output + void* data_ptr = output_spec.get_mutable_data_ptr(); + + if (data_ptr != nullptr) { + // Copy data from staging buffer to output spec + graph.copy_from_staging(output_ref.staging, data_ptr, data_numel); + } + + // Print output tensor data if output printing is enabled + if (print_output()) { + std::string output_name = "Output[" + std::to_string(i) + "]"; + print_valuespec_data(output_spec, output_name); + } + } + } + + return result; +} + +TestResult execute_test_cases( + std::function()> test_case_generator, + FlopCalculatorFunc flop_calculator, + const std::string& operation_name, + int warmup_runs, + int benchmark_runs, + ReferenceComputeFunc reference_compute_func) { + TestResult results(operation_name); + + // Generate all test cases + std::vector test_cases = test_case_generator(); + + std::cout << "Executing " << test_cases.size() << " test cases for " + << operation_name << std::endl; + print_separator(); + + bool any_correctness_failed = false; + float total_gflops = 0.0f; + + for (size_t i = 0; i < test_cases.size(); ++i) { + TestCase& test_case = test_cases[i]; + + // Compute reference data if reference function is provided + bool skipped_reference_fn = true; + if (reference_compute_func) { + try { + reference_compute_func(test_case); + skipped_reference_fn = false; + } catch (const std::invalid_argument& e) { + if (debugging()) { + std::cout << "Compute reference skipped: " << e.what() << std::endl; + } + } + } + + // Execute single test case + BenchmarkResult result; + bool shader_not_supported = false; + try { + result = execute_test_case(test_case, warmup_runs, benchmark_runs); + result.set_operator_name(test_case.operator_name()); + } catch (const vkcompute::vkapi::ShaderNotSupportedError& e) { + result = BenchmarkResult( + test_case.name().empty() ? "unnamed_test_case" : test_case.name(), + test_case.operator_name()); + shader_not_supported = true; + } + + // Determine if this test case passed (has valid timing data) + bool vulkan_execute_succeeded = + result.get_num_iterations() > 0 && result.get_avg_time_us() > 0.0f; + + if (shader_not_supported) { + result.set_correctness_status(CorrectnessStatus::SKIPPED); + } else if (!vulkan_execute_succeeded) { + result.set_correctness_status(CorrectnessStatus::FAILED); + } else if (skipped_reference_fn) { + result.set_correctness_status(CorrectnessStatus::SKIPPED); + } else { + // Reference function provided and succeeded - validate outputs + bool correctness_passed = true; + + for (size_t output_idx = 0; output_idx < test_case.num_outputs(); + ++output_idx) { + const ValueSpec& output_spec = test_case.outputs()[output_idx]; + + if (!output_spec.validate_against_reference( + test_case.get_abs_tolerance(), test_case.get_rel_tolerance())) { + correctness_passed = false; + std::cout << " Correctness validation FAILED for test " + << result.get_kernel_name() << std::endl; + print_valuespec_data(output_spec, "vulkan output"); + print_valuespec_data(output_spec, "ref output", true); + + throw std::runtime_error("Correctness validation failed"); + } + } + + if (correctness_passed) { + result.set_correctness_status(CorrectnessStatus::PASSED); + } else { + any_correctness_failed = true; + result.set_correctness_status(CorrectnessStatus::FAILED); + } + } + + // Calculate GFLOPS for this test case using the provided FLOP calculator + float case_gflops = 0.0f; + if (vulkan_execute_succeeded) { + // Use the provided FLOP calculator to get total FLOPs for this test case + int64_t total_flops = flop_calculator(test_case); + float flops = static_cast(total_flops); + float avg_time_us = result.get_avg_time_us(); + if (avg_time_us > 0.0f && total_flops > 0) { + case_gflops = (flops / 1e9f) / (avg_time_us / 1e6f); + } + + total_gflops += case_gflops; + } else { + case_gflops = -1.0f; // Indicate failure + } + + // Calculate tensor info for display + std::string size_info = "["; + if (!test_case.empty() && test_case.num_inputs() > 0 && + test_case.inputs()[0].is_tensor()) { + const auto& sizes = test_case.inputs()[0].get_tensor_sizes(); + for (size_t j = 0; j < sizes.size(); ++j) { + size_info += std::to_string(sizes[j]); + if (j < sizes.size() - 1) + size_info += "x"; + } + } + size_info += "]"; + + // Print progress using the BenchmarkResult member function + result.print_summary(i + 1, size_info, case_gflops); + + // Add result to collection + results.add_result(std::move(result)); + } + + // Set the overall results on the TestResult + results.set_correctness_passed(!any_correctness_failed); + results.set_gflops(total_gflops); + + print_separator(); + std::cout << "Completed " << results.size() << " test cases" << std::endl; + + return results; +} + +// Convenience overload that uses the default FLOP calculator +TestResult execute_test_cases( + std::function()> test_case_generator, + const std::string& operation_name, + int warmup_runs, + int benchmark_runs, + ReferenceComputeFunc reference_compute_func) { + return execute_test_cases( + test_case_generator, + default_flop_calculator, + operation_name, + warmup_runs, + benchmark_runs, + reference_compute_func); +} + +// Utility functions for printing +void print_performance_header() { + std::cout << "\n=== Compute Shader Performance Benchmark ===" << std::endl; +} + +void print_separator() { + std::cout << std::string(70, '-') << std::endl; +} + +// ValueSpec data printing utilities +void print_valuespec_data( + const ValueSpec& spec, + const std::string& name, + const bool print_ref_data, + size_t max_elements, + int precision) { + std::cout << "\n" << name << " Data:" << std::endl; + std::cout << " Type: " << spec.to_string() << std::endl; + + if (!spec.is_tensor()) { + if (spec.is_int()) { + std::cout << " Value: " << spec.get_int_value() << std::endl; + } else if (spec.is_int_list()) { + const auto& int_list = spec.get_int_list(); + std::cout << " Values: ["; + size_t print_count = std::min(max_elements, int_list.size()); + for (size_t i = 0; i < print_count; ++i) { + std::cout << int_list[i]; + if (i < print_count - 1) + std::cout << ", "; + } + if (int_list.size() > max_elements) { + std::cout << ", ... (" << (int_list.size() - max_elements) << " more)"; + } + std::cout << "]" << std::endl; + } + return; + } + + // Print tensor data + size_t total_elements = spec.numel(); + size_t print_count = std::min(max_elements, total_elements); + + std::cout << " Total elements: " << total_elements << std::endl; + std::cout << " Data (first " << print_count << " elements): ["; + + std::cout << std::fixed << std::setprecision(precision); + + switch (spec.dtype) { + case vkapi::kFloat: { + auto data = spec.get_float_data().data(); + if (print_ref_data) { + data = spec.get_ref_float_data().data(); + } + for (size_t i = 0; i < print_count; ++i) { + std::cout << data[i]; + if (i < print_count - 1) + std::cout << ", "; + } + break; + } + case vkapi::kHalf: { + const auto& data = spec.get_half_data(); + for (size_t i = 0; i < print_count; ++i) { + // Convert uint16_t back to float for display + float value = data[i] / 32767.0f; + std::cout << value; + if (i < print_count - 1) + std::cout << ", "; + } + break; + } + case vkapi::kInt: { + const auto& data = spec.get_int32_data(); + for (size_t i = 0; i < print_count; ++i) { + std::cout << data[i]; + if (i < print_count - 1) + std::cout << ", "; + } + break; + } + case vkapi::kChar: { + const auto& data = spec.get_int8_data(); + if (spec.is_int4()) { + // Print each 4-bit value individually + size_t element_count = 0; + for (size_t i = 0; i < data.size() && element_count < print_count; + ++i) { + // Extract lower 4 bits (signed) + int8_t lower_4bits = data[i] & 0x0F; + if (lower_4bits > 7) + lower_4bits -= 16; // Convert to signed + std::cout << static_cast(lower_4bits); + element_count++; + + if (element_count < print_count) { + std::cout << ", "; + // Extract upper 4 bits (signed) + int8_t upper_4bits = (data[i] >> 4) & 0x0F; + if (upper_4bits > 7) + upper_4bits -= 16; // Convert to signed + std::cout << static_cast(upper_4bits); + element_count++; + + if (element_count < print_count) + std::cout << ", "; + } + } + } else { + for (size_t i = 0; i < print_count; ++i) { + std::cout << static_cast(data[i]); + if (i < print_count - 1) + std::cout << ", "; + } + } + break; + } + case vkapi::kByte: { + const auto& data = spec.get_uint8_data(); + if (spec.is_int4()) { + // Print each 4-bit value individually + size_t element_count = 0; + for (size_t i = 0; i < data.size() && element_count < print_count; + ++i) { + // Extract lower 4 bits + uint8_t lower_4bits = data[i] & 0x0F; + std::cout << static_cast(lower_4bits); + element_count++; + + if (element_count < print_count) { + std::cout << ", "; + // Extract upper 4 bits + uint8_t upper_4bits = (data[i] >> 4) & 0x0F; + std::cout << static_cast(upper_4bits); + element_count++; + + if (element_count < print_count) + std::cout << ", "; + } + } + } else { + for (size_t i = 0; i < print_count; ++i) { + std::cout << static_cast(data[i]); + if (i < print_count - 1) + std::cout << ", "; + } + } + break; + } + default: + std::cout << "unsupported data type"; + break; + } + + if (total_elements > max_elements) { + std::cout << ", ... (" << (total_elements - max_elements) << " more)"; + } + std::cout << "]" << std::endl; + + // Print some statistics for tensor data + if (total_elements > 0) { + float min_val = 0.0f, max_val = 0.0f, sum = 0.0f; + bool first = true; + + for (size_t i = 0; i < total_elements; ++i) { + float val = spec.get_element(i); + if (first) { + min_val = max_val = val; + first = false; + } else { + min_val = std::min(min_val, val); + max_val = std::max(max_val, val); + } + sum += val; + } + + float mean = sum / total_elements; + std::cout << " Statistics: min=" << std::setprecision(precision) << min_val + << ", max=" << max_val << ", mean=" << mean << ", sum=" << sum + << std::endl; + } +} + +ValueRef quantized_weights_canvas( + ComputeGraph& graph, + const ValueRef weight_ref) { + const auto original_sizes = graph.sizes_of(weight_ref); + + // Get the 2 highest values of original_sizes + std::vector sorted_sizes = original_sizes; + std::sort(sorted_sizes.begin(), sorted_sizes.end(), std::greater()); + int64_t largest1 = sorted_sizes.size() > 0 ? sorted_sizes[0] : 0; + int64_t largest2 = sorted_sizes.size() > 1 ? sorted_sizes[1] : 0; + + std::vector final_sizes = {1, largest1, largest1}; + + // Debug logging if debugging flag is set + if (debugging()) { + std::cout << "Debug: Creating quantized weights canvas tensor" << std::endl; + std::cout << "Debug: Original sizes: ["; + for (size_t i = 0; i < original_sizes.size(); ++i) { + std::cout << original_sizes[i]; + if (i < original_sizes.size() - 1) + std::cout << ", "; + } + std::cout << "]" << std::endl; + std::cout << "Debug: Canvas sizes: ["; + for (size_t i = 0; i < final_sizes.size(); ++i) { + std::cout << final_sizes[i]; + if (i < final_sizes.size() - 1) + std::cout << ", "; + } + std::cout << "]" << std::endl; + } + + ValueRef packed_weight = graph.add_tensor( + final_sizes, vkapi::kInt, utils::kTexture3D, utils::kWidthPacked); + + utils::uvec3 global_wg_size{ + utils::div_up(utils::safe_downcast(largest1), uint32_t(4)), + utils::safe_downcast(largest2), + utils::safe_downcast(std::min(largest1, int64_t(2048)))}; + + std::string kernel_name = "packed_int32_canvas"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(packed_weight)); + + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + graph.create_global_wg_size(packed_weight), + graph.create_local_wg_size(packed_weight), + weight_ref, + packed_weight, + // UBOs + {graph.logical_limits_ubo(packed_weight)}, + // Specialization constants + {}, + // Push Constants + {})); + + return packed_weight; +} + +ValueRef float_tensor_canvas(ComputeGraph& graph, const ValueRef weight_ref) { + const auto original_sizes = graph.sizes_of(weight_ref); + + // Get the 2 highest values of original_sizes + std::vector sorted_sizes = original_sizes; + std::sort(sorted_sizes.begin(), sorted_sizes.end(), std::greater()); + int64_t largest1 = sorted_sizes.size() > 0 ? sorted_sizes[0] : 0; + int64_t largest2 = sorted_sizes.size() > 1 ? sorted_sizes[1] : 0; + + std::vector final_sizes = {1, largest1, largest1}; + + // Debug logging if debugging flag is set + if (debugging()) { + std::cout << "Debug: Creating float tensor canvas" << std::endl; + std::cout << "Debug: Original sizes: ["; + for (size_t i = 0; i < original_sizes.size(); ++i) { + std::cout << original_sizes[i]; + if (i < original_sizes.size() - 1) + std::cout << ", "; + } + std::cout << "]" << std::endl; + std::cout << "Debug: Canvas sizes: ["; + for (size_t i = 0; i < final_sizes.size(); ++i) { + std::cout << final_sizes[i]; + if (i < final_sizes.size() - 1) + std::cout << ", "; + } + std::cout << "]" << std::endl; + } + + ValueRef packed_weight = graph.add_tensor( + final_sizes, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + + utils::uvec3 global_wg_size{ + utils::div_up(utils::safe_downcast(largest1), uint32_t(4)), + utils::safe_downcast(largest2), + utils::safe_downcast(std::min(largest1, int64_t(2048)))}; + + graph.prepack_nodes().emplace_back(new PrepackNode( + graph, + VK_KERNEL_FROM_STR("float_canvas"), + graph.create_global_wg_size(packed_weight), + graph.create_local_wg_size(packed_weight), + weight_ref, + packed_weight, + // UBOs + {graph.logical_limits_ubo(packed_weight)}, + // Specialization constants + {}, + // Push Constants + {})); + + return packed_weight; +} + +// Compute weight sums for quantized operations (linear and convolution) +void compute_weight_sums( + ValueSpec& weight_sums, + const ValueSpec& quantized_weight, + int64_t out_features, + int64_t elements_per_output_feature) { + auto& weight_sums_data = weight_sums.get_int32_data(); + auto& quantized_weight_data = quantized_weight.get_int8_data(); + + weight_sums_data.resize(out_features); + + // For each output feature, compute the sum of quantized weights + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + int32_t sum = 0; + for (int64_t elem = 0; elem < elements_per_output_feature; ++elem) { + // Weight indexing depends on the layout: + // For linear: [out_features, in_features] -> out_f * + // elements_per_output_feature + elem For conv2d: [C_out, C_in * K_h * + // K_w] -> out_f * elements_per_output_feature + elem + int64_t weight_idx = out_f * elements_per_output_feature + elem; + sum += static_cast(quantized_weight_data[weight_idx]); + } + weight_sums_data[out_f] = sum; + } +} + +// Helper function to unpack 4-bit values from uint8 (same as in +// q4gsw_linear.cpp) +std::pair unpack_4bit_utils(uint8_t packed) { + // Extract lower 4 bits and upper 4 bits + int8_t lower = packed & 0x0F; + int8_t upper = (packed >> 4) & 0x0F; + + // Subtract 8 from unpacked 4-bit values + lower -= 8; + upper -= 8; + + return std::make_pair(lower, upper); +} + +// Compute weight sums for 4-bit group symmetric quantized weights +void compute_weight_sums_4bit_grouped( + ValueSpec& weight_sums, + const ValueSpec& quantized_weight, + int64_t num_groups, + int64_t out_features, + int64_t group_size) { + auto& weight_sums_data = weight_sums.get_int32_data(); + auto& quantized_weight_data = quantized_weight.get_uint8_data(); + + // Resize to [num_groups, out_features] + weight_sums_data.resize(num_groups * out_features); + + // For each group and each output feature, compute the sum of quantized + // weights in that group + for (int64_t group_idx = 0; group_idx < num_groups; ++group_idx) { + for (int64_t out_f = 0; out_f < out_features; ++out_f) { + int32_t sum = 0; + + // Sum weights for this group and output feature + for (int64_t in_group = 0; in_group < group_size; ++in_group) { + int64_t in_f = group_idx * group_size + in_group; + + // Get packed weight value - weight matrix is [N, K/2] + int64_t weight_idx = + out_f * ((num_groups * group_size) / 2) + (in_f / 2); + uint8_t packed_weight = quantized_weight_data[weight_idx]; + + // Unpack 4-bit weight + auto unpacked = unpack_4bit_utils(packed_weight); + int8_t weight_4bit = (in_f % 2 == 0) ? unpacked.first : unpacked.second; + + sum += static_cast(weight_4bit); + } + + // Store sum for this group and output feature + int64_t sums_idx = group_idx * out_features + out_f; + weight_sums_data[sums_idx] = sum; + } + } +} + +} // namespace prototyping +} // namespace vulkan +} // namespace executorch diff --git a/backends/vulkan/test/custom_ops/utils.h b/backends/vulkan/test/custom_ops/utils.h new file mode 100644 index 00000000000..f1736f1d144 --- /dev/null +++ b/backends/vulkan/test/custom_ops/utils.h @@ -0,0 +1,669 @@ +// 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 + +namespace executorch { +namespace vulkan { +namespace prototyping { + +using namespace vkcompute; + +// +// Global configuration options +// + +bool print_output(); +void set_print_output(bool print_output); + +bool print_latencies(); +void set_print_latencies(bool print_latencies); + +bool use_gpu_timestamps(); +void set_use_gpu_timestamps(bool use_timestamps); + +bool debugging(); +void set_debugging(bool enable_debugging); + +// +// ValueSpec class +// + +enum class SpecType { Tensor, IntList, Int, Float, Bool }; + +// Data generation types +enum class DataGenType { + FIXED, + RANDOM, + RANDOM_SCALES, + RANDINT, + RANDINT8, + RANDINT4, + ONES, + ONES_INT4, + ZEROS +}; + +// Value specification struct +struct ValueSpec { + std::vector sizes; + vkapi::ScalarType dtype; + utils::GPUMemoryLayout memory_layout; + utils::StorageType storage_type; + SpecType spec_type; + DataGenType data_gen_type; + bool is_constant_tensor; + bool is_none_flag; + bool is_int4_tensor; + + 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::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, + utils::StorageType storage_type = utils::kTexture3D, + utils::GPUMemoryLayout memory_layout = utils::kWidthPacked) + : sizes(sizes), + dtype(dtype), + memory_layout(memory_layout), + storage_type(storage_type), + spec_type(SpecType::Tensor), + data_gen_type(DataGenType::ZEROS), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) { + generate_tensor_data(); + } + + // Constructor for tensor with custom data generation type + ValueSpec( + const std::vector& sizes, + vkapi::ScalarType dtype, + utils::StorageType storage_type, + utils::GPUMemoryLayout memory_layout, + DataGenType data_gen_type) + : sizes(sizes), + dtype(dtype), + memory_layout(memory_layout), + storage_type(storage_type), + spec_type(SpecType::Tensor), + data_gen_type(data_gen_type), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) { + generate_tensor_data(); + } + + // Constructor for single int + ValueSpec(int32_t value) + : sizes({1}), + dtype(vkapi::kInt), + memory_layout(utils::kWidthPacked), + storage_type(utils::kTexture3D), + spec_type(SpecType::Int), + data_gen_type(DataGenType::FIXED), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) { + int32_data.push_back(value); + } + + // Constructor for single float + ValueSpec(float value) + : sizes({1}), + dtype(vkapi::kFloat), + memory_layout(utils::kWidthPacked), + storage_type(utils::kTexture3D), + spec_type(SpecType::Float), + data_gen_type(DataGenType::FIXED), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) { + float_data.push_back(value); + } + + // Constructor for single bool + ValueSpec(bool value) + : sizes({1}), + dtype(vkapi::kInt), + memory_layout(utils::kWidthPacked), + storage_type(utils::kTexture3D), + spec_type(SpecType::Bool), + data_gen_type(DataGenType::FIXED), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) { + int32_data.push_back(value ? 1 : 0); + } + + // Constructor for int list + ValueSpec(const std::vector& values) + : sizes({static_cast(values.size())}), + dtype(vkapi::kInt), + memory_layout(utils::kWidthPacked), + storage_type(utils::kTexture3D), + spec_type(SpecType::IntList), + data_gen_type(DataGenType::FIXED), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false), + int32_data(values) {} + + // Default constructor + ValueSpec() + : dtype(vkapi::kFloat), + memory_layout(utils::kWidthPacked), + storage_type(utils::kTexture3D), + spec_type(SpecType::Tensor), + data_gen_type(DataGenType::ZEROS), + is_constant_tensor(false), + is_none_flag(false), + is_int4_tensor(false) {} + + int64_t numel() const; + size_t nbytes() const; + std::string to_string() const; + + bool is_tensor() const { + return spec_type == SpecType::Tensor; + } + bool is_int_list() const { + return spec_type == SpecType::IntList; + } + bool is_int() const { + return spec_type == SpecType::Int; + } + bool is_float() const { + return spec_type == SpecType::Float; + } + bool is_bool() const { + return spec_type == SpecType::Bool; + } + + int32_t get_int_value() const { + return int32_data.empty() ? 0 : int32_data[0]; + } + float get_float_value() const { + return float_data.empty() ? 0.0f : float_data[0]; + } + bool get_bool_value() const { + return int32_data.empty() ? false : (int32_data[0] != 0); + } + const std::vector& get_int_list() const { + return int32_data; + } + const std::vector& get_tensor_sizes() const { + return sizes; + } + + const std::vector& get_float_data() const { + return float_data; + } + const std::vector& get_int32_data() const { + return int32_data; + } + const std::vector& get_half_data() const { + return half_data; + } + const std::vector& get_int8_data() const { + return int8_data; + } + const std::vector& get_uint8_data() const { + return uint8_data; + } + + std::vector& get_float_data() { + return float_data; + } + std::vector& get_int32_data() { + return int32_data; + } + std::vector& get_half_data() { + return half_data; + } + std::vector& get_int8_data() { + return int8_data; + } + std::vector& get_uint8_data() { + return uint8_data; + } + + const std::vector& get_ref_float_data() const { + return ref_float_data; + } + const std::vector& get_ref_int32_data() const { + return ref_int32_data; + } + const std::vector& get_ref_half_data() const { + return ref_half_data; + } + const std::vector& get_ref_int8_data() const { + return ref_int8_data; + } + const std::vector& get_ref_uint8_data() const { + return ref_uint8_data; + } + + std::vector& get_ref_float_data() { + return ref_float_data; + } + std::vector& get_ref_int32_data() { + return ref_int32_data; + } + std::vector& get_ref_half_data() { + return ref_half_data; + } + std::vector& get_ref_int8_data() { + return ref_int8_data; + } + std::vector& get_ref_uint8_data() { + return ref_uint8_data; + } + + void resize_data(size_t new_size); + void* get_mutable_data_ptr(); + float get_element(size_t index) const; + + // Set/get constant flag + bool is_constant() const { + return is_constant_tensor; + } + void set_constant(bool is_constant) { + is_constant_tensor = is_constant; + } + + // Set/get none flag + bool is_none() const { + return is_none_flag; + } + + void set_none(bool is_none) { + is_none_flag = is_none; + } + + // Set/get int4 flag + bool is_int4() const { + return is_int4_tensor; + } + void set_int4(bool is_int4) { + is_int4_tensor = is_int4; + } + + const void* get_data_ptr() const; + + // Correctness checking against reference data + // Returns true if computed data matches reference data within tolerance + // Only validates float tensors as specified in requirements + bool validate_against_reference( + float abs_tolerance = 2e-3f, + float rel_tolerance = 1e-3f) const; + + private: + void generate_tensor_data(); +}; + +// +// TestCase +// + +class TestCase { + public: + TestCase() : abs_tolerance_(2e-3f), rel_tolerance_(1e-3f) {} + TestCase(const std::string& name) + : name_(name), abs_tolerance_(2e-3f), rel_tolerance_(1e-3f) {} + + void set_name(const std::string& name) { + name_ = name; + } + const std::string& name() const { + return name_; + } + + void set_operator_name(const std::string& op_name) { + operator_name_ = op_name; + } + const std::string& operator_name() const { + return operator_name_; + } + + // Tolerance settings + void set_abs_tolerance(float abs_tolerance) { + abs_tolerance_ = abs_tolerance; + } + float get_abs_tolerance() const { + return abs_tolerance_; + } + + void set_rel_tolerance(float rel_tolerance) { + rel_tolerance_ = rel_tolerance; + } + float get_rel_tolerance() const { + return rel_tolerance_; + } + + void add_input_spec(const ValueSpec& spec) { + inputs_.push_back(spec); + } + + const std::vector& inputs() const { + return inputs_; + } + + std::vector& inputs() { + return inputs_; + } + + size_t num_inputs() const { + return inputs_.size(); + } + + void add_output_spec(const ValueSpec& spec) { + outputs_.push_back(spec); + } + + const std::vector& outputs() const { + return outputs_; + } + + std::vector& outputs() { + return outputs_; + } + + size_t num_outputs() const { + return outputs_.size(); + } + + bool empty() const { + return inputs_.empty() && outputs_.empty(); + } + void clear() { + inputs_.clear(); + outputs_.clear(); + name_.clear(); + operator_name_.clear(); + abs_tolerance_ = 2e-3f; + rel_tolerance_ = 1e-3f; + } + + private: + std::string name_; + std::string operator_name_; + std::vector inputs_; + std::vector outputs_; + float abs_tolerance_; + float rel_tolerance_; +}; + +// +// BenchmarkResult +// + +enum class CorrectnessStatus { + SKIPPED, // No reference function provided + PASSED, // Reference function provided and validation passed + FAILED // Reference function provided but validation failed +}; + +class BenchmarkResult { + public: + BenchmarkResult() : correctness_status_(CorrectnessStatus::SKIPPED) {} + + BenchmarkResult(const std::string& name) + : kernel_name(name), correctness_status_(CorrectnessStatus::SKIPPED) {} + + BenchmarkResult( + const std::string& kernel_name, + const std::string& operator_name) + : kernel_name(kernel_name), + operator_name(operator_name), + correctness_status_(CorrectnessStatus::SKIPPED) {} + + // Add timing for a single iteration + void add_iter_timing(float time_us); + + // Getters + const std::string& get_kernel_name() const { + return kernel_name; + } + const std::string& get_operator_name() const { + return operator_name; + } + float get_avg_time_us() const; + size_t get_num_iterations() const { + return iter_timings.size(); + } + const std::vector& get_iter_timings() const { + return iter_timings; + } + CorrectnessStatus get_correctness_status() const { + return correctness_status_; + } + + // Setters + void set_kernel_name(const std::string& name) { + kernel_name = name; + } + void set_operator_name(const std::string& name) { + operator_name = name; + } + void set_correctness_status(CorrectnessStatus status) { + correctness_status_ = status; + } + + // Statistics + float get_min_time_us() const; + float get_max_time_us() const; + float get_std_dev_us() const; + + // Clear all timings + void clear_timings() { + iter_timings.clear(); + } + + // Print progress for this benchmark result + void print_summary( + int case_number, + const std::string& size_info, + float total_gflops) const; + + private: + std::string kernel_name; + std::string operator_name; + std::vector + iter_timings; // Individual iteration timings in microseconds + CorrectnessStatus correctness_status_; +}; + +// Test result collection and processing +class TestResult { + public: + TestResult() : gflops_(0.0f), correctness_passed_(true) {} + TestResult(const std::string& operation_name) + : operation_name_(operation_name), + gflops_(0.0f), + correctness_passed_(true) {} + + // Add a benchmark result + void add_result(const BenchmarkResult& result); + void add_result(BenchmarkResult&& result); + + // Getters + const std::string& get_operation_name() const { + return operation_name_; + } + float get_gflops() const { + return gflops_; + } + bool get_correctness_passed() const { + return correctness_passed_; + } + size_t size() const { + return results_.size(); + } + bool empty() const { + return results_.empty(); + } + + // Setters + void set_gflops(float gflops_val) { + gflops_ = gflops_val; + } + void set_correctness_passed(bool passed) { + correctness_passed_ = passed; + } + + // Access results + const BenchmarkResult& operator[](size_t index) const { + return results_[index]; + } + BenchmarkResult& operator[](size_t index) { + return results_[index]; + } + const std::vector& get_results() const { + return results_; + } + + // Iterator support + std::vector::iterator begin() { + return results_.begin(); + } + std::vector::iterator end() { + return results_.end(); + } + std::vector::const_iterator begin() const { + return results_.begin(); + } + std::vector::const_iterator end() const { + return results_.end(); + } + + // Processing and analysis + void print_summary() const; + void print_detailed_results() const; + void print_statistics() const; + void print_brief_summary() const; + + // Get aggregate statistics + float get_total_avg_time_us() const; + float get_total_gflops() const; + size_t get_passed_count() const; + size_t get_failed_count() const; + + // Find best/worst performing results + const BenchmarkResult* get_fastest_result() const; + const BenchmarkResult* get_slowest_result() const; + const BenchmarkResult* get_highest_gflops_result() const; + + // Clear all results + void clear() { + results_.clear(); + } + + // Set operation name + void set_operation_name(const std::string& name) { + operation_name_ = name; + } + + private: + std::string operation_name_; + std::vector results_; + float gflops_; + bool correctness_passed_; +}; + +// +// Test case execution +// + +using FlopCalculatorFunc = std::function; + +// Default FLOP calculation function (assumes 1 FLOP per element) +int64_t default_flop_calculator(const TestCase& test_case); + +using ReferenceComputeFunc = std::function; + +BenchmarkResult execute_test_case( + TestCase& test_case, + int warmup_runs = 3, + int benchmark_runs = 10); + +TestResult execute_test_cases( + std::function()> test_case_generator, + FlopCalculatorFunc flop_calculator, + const std::string& operation_name = "Operation", + int warmup_runs = 3, + int benchmark_runs = 10, + ReferenceComputeFunc reference_compute_func = nullptr); + +TestResult execute_test_cases( + std::function()> test_case_generator, + const std::string& operation_name = "Operation", + int warmup_runs = 3, + int benchmark_runs = 10, + ReferenceComputeFunc reference_compute_func = nullptr); + +// +// Print utilities +// + +void print_performance_header(); +void print_separator(); + +void print_valuespec_data( + const ValueSpec& spec, + const std::string& name = "ValueSpec", + const bool print_ref_data = false, + size_t max_elements = 20, + int precision = 6); + +ValueRef quantized_weights_canvas( + ComputeGraph& graph, + const ValueRef weight_ref); + +ValueRef float_tensor_canvas(ComputeGraph& graph, const ValueRef weight_ref); + +// Compute weight sums for quantized operations (linear and convolution) +void compute_weight_sums( + ValueSpec& weight_sums, + const ValueSpec& quantized_weight, + int64_t out_features, + int64_t elements_per_output_feature); + +// Compute weight sums for 4-bit group symmetric quantized weights +void compute_weight_sums_4bit_grouped( + ValueSpec& weight_sums, + const ValueSpec& quantized_weight, + int64_t num_groups, + int64_t out_features, + int64_t group_size); + +// Setup compute graph based on TestCase and operation name +ComputeGraph setup_compute_graph(TestCase& test_case, std::string op_name); + +} // namespace prototyping +} // namespace vulkan +} // namespace executorch diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index c19e818f63d..5e8991f8e50 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -29,10 +29,10 @@ if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) endif() -# Include this file to access executorch_target_link_options_shared_lib This is required to -# provide access to executorch_target_link_options_shared_lib which allows libraries to be -# linked with the --whole-archive flag. This is required for libraries that -# perform dynamic registration via static initialization. +# Include this file to access executorch_target_link_options_shared_lib This is +# required to provide access to executorch_target_link_options_shared_lib which +# allows libraries to be linked with the --whole-archive flag. This is required +# for libraries that perform dynamic registration via static initialization. include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) get_torch_base_path(TORCH_BASE_PATH) @@ -47,7 +47,7 @@ find_library(LIB_C10 c10 HINTS ${TORCH_INSTALL_PREFIX}/lib) # Third party include paths -set(VULKAN_THIRD_PARTY_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../../third-party) +set(VULKAN_THIRD_PARTY_PATH ${EXECUTORCH_ROOT}/backends/vulkan/third-party) set(GTEST_INCLUDE_PATH ${EXECUTORCH_ROOT}/third-party/googletest/googletest/include @@ -88,6 +88,12 @@ function(vulkan_op_test test_name test_src) 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}) + target_link_libraries( + test_utils PRIVATE vulkan_backend ${LIB_TORCH} ${LIB_TORCH_CPU} + ) + find_library( CUSTOM_OPS_LIB custom_ops_aot_lib HINTS ${CMAKE_INSTALL_PREFIX}/executorch/extension/llm/custom_ops @@ -95,7 +101,7 @@ if(TARGET vulkan_backend AND LIB_TORCH) if(CUSTOM_OPS_LIB) vulkan_op_test( vulkan_sdpa_test ${CMAKE_CURRENT_SOURCE_DIR}/sdpa_test.cpp - ${CUSTOM_OPS_LIB} + ${CUSTOM_OPS_LIB} test_utils ) else() message( @@ -104,10 +110,11 @@ if(TARGET vulkan_backend AND LIB_TORCH) endif() vulkan_op_test( vulkan_rope_test ${CMAKE_CURRENT_SOURCE_DIR}/rotary_embedding_test.cpp + test_utils ) vulkan_op_test( - vulkan_linear_weight_int4_test - ${CMAKE_CURRENT_SOURCE_DIR}/linear_weight_int4_test.cpp + quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp + test_utils ) # Only build generated op tests if a path to tags.yaml and diff --git a/backends/vulkan/test/op_tests/cases.py b/backends/vulkan/test/op_tests/cases.py index 22725a46100..8c5d0c4797b 100644 --- a/backends/vulkan/test/op_tests/cases.py +++ b/backends/vulkan/test/op_tests/cases.py @@ -55,16 +55,28 @@ def get_binary_elementwise_inputs(): ((3, 64, 1), (1, 64, 1)), ] ) - test_suite.layouts = [ - "utils::kWidthPacked", - "utils::kChannelsPacked", - ] test_suite.storage_types = [ "utils::kBuffer", "utils::kTexture3D", ] - return test_suite + highdim_test_suite = VkTestSuite( + [ + ((4, 5, 8, 1, 2, 1), (4, 5, 8, 1, 1, 1)), + ] + ) + highdim_test_suite.storage_types = [ + "utils::kBuffer", + ] + highdim_test_suite.test_name_suffix = "highdim" + + for suite in [test_suite, highdim_test_suite]: + suite.layouts = [ + "utils::kWidthPacked", + "utils::kChannelsPacked", + ] + + return [test_suite, highdim_test_suite] # Eq requires a different test generator so it was split from the other test case. @@ -258,11 +270,13 @@ def get_avg_pool2d_inputs(): return test_suite -@register_test_suite("aten.max_pool2d_with_indices.default") +@register_test_suite( + ["aten.max_pool2d_with_indices.default", "aten.max_pool2d.default"] +) def get_max_pool2d_inputs(): test_suite = VkTestSuite( [ - ((S, M1, M2), [2, 2], [1, 1], [0, 0], [1, 1]), + ((1, 7, 89, 77), [2, 2], [1, 1], [0, 0], [1, 1]), ] ) return test_suite @@ -297,6 +311,28 @@ def get_conv_inputs(): ) test_cases = [ + Test( + self=(1, 64, 256, 256), + weight=(64, 32, 3, 3), + bias=None, + stride=[1, 1], + padding=[1, 1], + dilation=[1, 1], + transposed=False, + output_padding=[0, 0], + groups=2, + ), + Test( + self=(1, 16, 3, 3), + weight=(16, 8, 3, 3), + bias=None, + stride=[1, 1], + padding=[1, 1], + dilation=[1, 1], + transposed=False, + output_padding=[0, 0], + groups=2, + ), Test( self=(1, 6, 40, 50), weight=(8, 6, 3, 3), @@ -877,7 +913,28 @@ def get_view_inputs(): "utils::kHeightPacked", "utils::kChannelsPacked", ] - return test_suite + + highdim_test_suite = VkTestSuite( + [ + ((1, 1, 3, 3, 3), (9, 3)), + ((2, 3, 4, 6, 5, 4), (6, 4, 6, 5, 4)), + ((2, 3, 3, 7, 8), (2, 3, 3, 8 * 7)), + ] + ) + highdim_test_suite.storage_types = [ + "utils::kBuffer", + ] + highdim_test_suite.test_name_suffix = "highdim" + highdim_test_suite.data_gen = "make_seq_tensor" + + for suite in [test_suite, highdim_test_suite]: + suite.layouts = [ + # "utils::kWidthPacked", + "utils::kHeightPacked", + "utils::kChannelsPacked", + ] + + return [test_suite, highdim_test_suite] @register_test_suite("aten.slice_copy.Tensor") @@ -1090,12 +1147,34 @@ def get_unsqueeze_inputs(): ((1, 10), -1), ] ) - test_suite.layouts = [ - "utils::kWidthPacked", - "utils::kChannelsPacked", + + highdim_test_suite = VkTestSuite( + [ + ((2, 3, 4, 5, 6), 0), + ((2, 3, 4, 5, 6), 1), + ((2, 3, 4, 5, 6), 5), + ((2, 3, 4, 5, 6), -1), + ((2, 3, 4, 5, 6), -2), + ((1, 2, 3, 4, 5), 0), + ((1, 2, 3, 4, 5), 3), + ((1, 2, 3, 4, 5), -1), + ((2, 3, 4, 5), 0), + ((1, 2, 3, 4), 1), + ] + ) + highdim_test_suite.storage_types = [ + "utils::kBuffer", ] - test_suite.data_gen = "make_seq_tensor" - return test_suite + highdim_test_suite.test_name_suffix = "highdim" + + for suite in [test_suite, highdim_test_suite]: + suite.layouts = [ + "utils::kWidthPacked", + "utils::kChannelsPacked", + ] + suite.data_gen = "make_seq_tensor" + + return [test_suite, highdim_test_suite] @register_test_suite("aten.clone.default") @@ -1115,11 +1194,28 @@ def get_clone_inputs(): ((XS,),), ] ) - test_suite.layouts = [ - "utils::kChannelsPacked", + + highdim_test_suite = VkTestSuite( + [ + ((2, 3, 4, 5, 6),), + ((2, 3, 4, 5, 1),), + ((1, 1, 3, 4, 5),), + ((2, 3, 4, 5, 6, 7),), + ((1, 2, 3, 4, 5, 6),), + ] + ) + highdim_test_suite.storage_types = [ + "utils::kBuffer", ] - test_suite.data_gen = "make_seq_tensor" - return test_suite + highdim_test_suite.test_name_suffix = "highdim" + + for suite in [test_suite, highdim_test_suite]: + suite.layouts = [ + "utils::kChannelsPacked", + ] + suite.data_gen = "make_seq_tensor" + + return [test_suite, highdim_test_suite] @register_test_suite("aten.repeat.default") @@ -1137,7 +1233,7 @@ def get_repeat_inputs(): "utils::kHeightPacked", "utils::kChannelsPacked", ] - test_suite_2d.storage_types = ["utils::kTexture2D"] + test_suite_2d.storage_types = ["utils::kTexture3D"] test_suite_2d.data_gen = "make_seq_tensor" test_suite_2d.dtypes = ["at::kFloat"] test_suite_2d.test_name_suffix = "2d" @@ -1233,66 +1329,81 @@ def get_repeat_interleave_inputs(): @register_test_suite("aten.cat.default") def get_cat_inputs(): # TensorList must be specified as list of tuples - test_suite = VkTestSuite( - [ - # Cat on Height - ([(M, M, 3, 5), (M, M, 0, 5)], 2), - ([(S1, S1, 3, 5), (S1, S1, 0, 5)], 2), - ([(M, M, 3, 5), (M, M, 4, 5)], 2), - ([(S1, S1, 3, 5), (S1, S1, 4, 5)], 2), - ([(M2, 3, 5), (M2, 4, 5)], 1), - ([(S1, 3, 5), (S1, 4, 5)], 1), - ([(3, 5), (4, 5)], 0), - ([(3, 5), (4, 5), (1, 5)], 0), - ( - [(3, 5)], - 0, - ), - # Cat on Width - ([(M, M, 5, 3), (M, M, 5, 4)], 3), - ([(S1, S1, 5, 3), (S1, S1, 5, 4)], 3), - ([(M, 5, 3), (M, 5, 4)], 2), - ([(S1, 5, 3), (S1, 5, 4)], 2), - ([(5, 0), (5, 4)], 1), - ([(5, 3), (5, 4)], 1), - ([(5, 3), (5, 4), (5, 1)], 1), - ( - [(5, 4)], - 1, - ), - ([(5,), (6,)], 0), - # Cat on Batch - ([(M, S1, 5, 4), (M1, S1, 5, 4)], 0), - ([(S, S1, 5, 4), (S1, S1, 5, 4)], 0), - ([(S, M, 5, 4), (S1, M, 5, 4)], 0), - ([(S, XS, 5, 4), (S1, XS, 5, 4)], 0), - ([(S, S2, 5, 4), (S1, S2, 5, 4)], 0), - ( - [ - (3, 1, 2, 5), - (3, 1, 2, 5), - (3, 1, 2, 5), - ], - 0, - ), - # Cat on Channel - ([(M, 5, 4), (0, 5, 4), (M1, 5, 4)], 0), - ([(S, 5, 4), (0, 5, 4), (S2, 5, 4)], 0), - ([(M, 5, 4), (M1, 5, 4), (M2, 5, 4)], 0), - ([(S, 5, 4), (S1, 5, 4), (S2, 5, 4)], 0), - ([(XS, 5, 4), (XS, 5, 4), (S2, 5, 4)], 0), - ([(XS, S, 5, 4), (XS, S1, 5, 4), (XS, S2, 5, 4)], 1), - ([(XS, XS, 5, 4), (XS, XS, 5, 4), (XS, S2, 5, 4)], 1), - ( - [ - (XS, 1, 2, 5), - (XS, 1, 2, 5), - (XS, 1, 2, 5), - ], - 1, - ), - ] - ) + suite_inputs = [ + # Cat on Height + ([(M, M, 3, 5), (M, M, 0, 5)], 2), + ([(S1, S1, 3, 5), (S1, S1, 0, 5)], 2), + ([(M, M, 3, 5), (M, M, 4, 5)], 2), + ([(S1, S1, 3, 5), (S1, S1, 4, 5)], 2), + ([(M2, 3, 5), (M2, 4, 5)], 1), + ([(S1, 3, 5), (S1, 4, 5)], 1), + ([(3, 5), (4, 5)], 0), + ([(3, 5), (4, 5), (1, 5)], 0), + ( + [(3, 5)], + 0, + ), + # Cat on Width + ([(M, M, 5, 3), (M, M, 5, 4)], 3), + ([(S1, S1, 5, 3), (S1, S1, 5, 4)], 3), + ([(M, 5, 3), (M, 5, 4)], 2), + ([(S1, 5, 3), (S1, 5, 4)], 2), + ([(5, 0), (5, 4)], 1), + ([(5, 3), (5, 4)], 1), + ([(5, 3), (5, 4), (5, 1)], 1), + ( + [(5, 4)], + 1, + ), + ([(5,), (6,)], 0), + # Cat on Batch + ([(M, S1, 5, 4), (M1, S1, 5, 4)], 0), + ([(S, S1, 5, 4), (S1, S1, 5, 4)], 0), + ([(S, M, 5, 4), (S1, M, 5, 4)], 0), + ([(S, XS, 5, 4), (S1, XS, 5, 4)], 0), + ([(S, S2, 5, 4), (S1, S2, 5, 4)], 0), + ( + [ + (3, 1, 2, 5), + (3, 1, 2, 5), + (3, 1, 2, 5), + ], + 0, + ), + # Cat on Channel + ([(M, 5, 4), (0, 5, 4), (M1, 5, 4)], 0), + ([(S, 5, 4), (0, 5, 4), (S2, 5, 4)], 0), + ([(M, 5, 4), (M1, 5, 4), (M2, 5, 4)], 0), + ([(S, 5, 4), (S1, 5, 4), (S2, 5, 4)], 0), + ([(XS, 5, 4), (XS, 5, 4), (S2, 5, 4)], 0), + ([(XS, S, 5, 4), (XS, S1, 5, 4), (XS, S2, 5, 4)], 1), + ([(XS, XS, 5, 4), (XS, XS, 5, 4), (XS, S2, 5, 4)], 1), + ( + [ + (XS, 1, 2, 5), + (XS, 1, 2, 5), + (XS, 1, 2, 5), + ], + 1, + ), + ] + + high_number_cat_inputs = [] + for num_input in [6, 9]: + odd_size = (3, 7, 29, 31) + even_size = (3, 8, 29, 32) + ones = (3, 1, 1, 1) + + for input_size in [odd_size, even_size, ones]: + input_sizes = [input_size] * num_input + # Test cat on height, width, and batch dim + high_number_cat_inputs.append((input_sizes, 3)) + high_number_cat_inputs.append((input_sizes, 2)) + high_number_cat_inputs.append((input_sizes, 1)) + high_number_cat_inputs.append((input_sizes, 0)) + + test_suite = VkTestSuite(suite_inputs + high_number_cat_inputs) + test_suite.layouts = [ "utils::kWidthPacked", "utils::kChannelsPacked", @@ -1724,7 +1835,31 @@ def get_squeeze_copy_dim_inputs(): ([1, M1, M1], 0), ] ) - return test_suite + + highdim_test_suite = VkTestSuite( + [ + ([1, 2, 3, 4, 5, 1], 0), + ([1, 2, 3, 4, 5, 1], 5), + ([1, 2, 3, 4, 5, 1], [0, 5]), + ([2, 1, 3, 1, 5, 6], 1), + ([2, 1, 3, 1, 5, 6], 3), + ([2, 1, 3, 1, 5, 6], [1, 3]), + ([1, 1, 3, 4, 5, 6], [0, 1]), + ([2, 3, 4, 1, 1, 6], [3, 4]), + ] + ) + highdim_test_suite.storage_types = [ + "utils::kBuffer", + ] + highdim_test_suite.test_name_suffix = "highdim" + + for suite in [test_suite, highdim_test_suite]: + suite.layouts = [ + "utils::kWidthPacked", + "utils::kChannelsPacked", + ] + + return [test_suite, highdim_test_suite] @register_test_suite("aten.flip.default") @@ -1747,6 +1882,48 @@ def get_flip_inputs(): return test_suite +@register_test_suite("aten.expand_copy.default") +def get_expand_inputs(): + test_suite = VkTestSuite( + [ + # Basic expansion cases + ((1,), [5]), + ((1, 1), [3, 4]), + ((1, 3), [2, 3]), + ((3, 1), [3, 4]), + ((1, 1, 1), [2, 3, 4]), + # Expand with same size (no-op) + ((3, 4), [3, 4]), + ((2, 3, 4), [2, 3, 4]), + # Expand with additional dimensions + ((3,), [2, 3]), + ((3, 4), [2, 3, 4]), + ((2, 3), [1, 2, 3]), + # Mixed expansion cases + ((1, 3, 1, 4), [2, 3, 5, 4]), + ((1, 1, 3, 1), [2, 4, 3, 5]), + # Larger tensor cases + ((1, S1), [M, S1]), + ((S2, 1), [S2, M1]), + ((1, 1, S), [S1, S2, S]), + ((1, S1, 1, S2), [M, S1, M1, S2]), + ] + ) + test_suite.storage_types = [ + "utils::kBuffer", + ] + test_suite.layouts = [ + "utils::kWidthPacked", + "utils::kChannelsPacked", + ] + test_suite.dtypes = [ + "at::kFloat", + "at::kHalf", + ] + test_suite.data_gen = "make_seq_tensor" + return test_suite + + @register_test_suite("aten.where.self") def get_where_inputs(): Test = namedtuple("Where", ["condition", "self", "other"]) diff --git a/backends/vulkan/test/op_tests/quantized_linear_test.cpp b/backends/vulkan/test/op_tests/quantized_linear_test.cpp index db95f4a793f..6e81820e735 100644 --- a/backends/vulkan/test/op_tests/quantized_linear_test.cpp +++ b/backends/vulkan/test/op_tests/quantized_linear_test.cpp @@ -33,43 +33,10 @@ class VulkanLinearQCS4WTest : public ::testing::Test { } }; -class VulkanLinearQTA8AQGA4WTest : public ::testing::Test { - public: - void SetUp() override { - if (!vkcompute::api::context() - ->adapter_ptr() - ->has_full_int8_buffers_support()) { - GTEST_SKIP(); - } - } - - void TearDown() override { - // Clean up any resources if needed - } -}; - // // Reference Implementations // -at::Tensor linear_qga4w_reference_impl( - const at::Tensor& x, - const at::Tensor& weights_4x2, - const int64_t groupsize, - const at::Tensor& scales_and_zeros, - const int64_t inner_k_tiles) { - const std::vector original_x_size(x.sizes().vec()); - const size_t ndim = original_x_size.size(); - const int64_t out_features = weights_4x2.size(0); - const at::Tensor x_flattened = x.reshape({-1, original_x_size[ndim - 1]}); - at::Tensor out = at::_weight_int4pack_mm_for_cpu( - x_flattened, weights_4x2, groupsize, scales_and_zeros); - std::vector out_shape( - original_x_size.begin(), original_x_size.end()); - out_shape.at(ndim - 1) = out_features; - return out.reshape(out_shape); -} - at::Tensor unpack_weights_4x2(const at::Tensor& weights_4x2) { std::vector weights_shape(weights_4x2.sizes().vec()); weights_shape[1] *= 2; @@ -94,41 +61,6 @@ at::Tensor unpack_weights_4x2(const at::Tensor& weights_4x2) { return weights_unpacked; } -at::Tensor dequantize_and_linear_qga4w( - const at::Tensor& x, - const at::Tensor& weights_4x2, - const int64_t groupsize, - const at::Tensor& scales_and_zeros, - const int64_t inner_k_tiles) { - std::vector weights_shape(weights_4x2.sizes().vec()); - weights_shape[1] *= 2; - - at::Tensor weights_dequantized = - at::empty(weights_shape, at::device(at::kCPU).dtype(at::kFloat)); - - const int64_t N = weights_dequantized.size(0); - const int64_t K = weights_dequantized.size(1); - - const int k_groups = K / groupsize; - for (int n = 0; n < N; n++) { - for (int k = 0; k < K; k += 2) { - const int group_idx = k / groupsize; - // const int scale_idx = k_groups * n + group_idx; - const uint8_t packed_val = weights_4x2[n][k / 2].item().to(); - const uint8_t second_val = packed_val & 0x0F; - const uint8_t first_val = (packed_val & 0xF0) >> 4; - - const float scale = scales_and_zeros[group_idx][n][0].item().to(); - const float zero = scales_and_zeros[group_idx][n][1].item().to(); - - weights_dequantized[n][k] = (float(first_val) - 8.0) * scale + zero; - weights_dequantized[n][k + 1] = (float(second_val) - 8.0) * scale + zero; - } - } - - return at::linear(x, weights_dequantized); -} - at::Tensor dequantize_and_linear_qcs4w( const at::Tensor& x, const at::Tensor& weights_4x2, @@ -179,197 +111,10 @@ at::Tensor linear_qcs4w_reference_impl( return out.reshape(out_shape); } -at::Tensor linear_qta8a_qga4w_quantized_matmul( - const at::Tensor& quantized_input, // [B, M, K] int8 quantized input - const at::Tensor& input_scale, // [B*M] per-token input scales - const at::Tensor& input_zero_point, // [B*M] per-token input zero points - const at::Tensor& weights_4x2, // [N, K/2] 4-bit packed weights - const int64_t group_size, // Group size for weight quantization - const at::Tensor& weight_scales, // [K/group_size, N] weight scales - const at::Tensor& weight_zeros) { // [K/group_size, N] weight zeros - - const int64_t B = quantized_input.size(0); - const int64_t M = quantized_input.size(1); - const int64_t K = quantized_input.size(2); - const int64_t N = weights_4x2.size(0); - - // Create output tensor for floating point results - at::Tensor float_output = - at::zeros({B, M, N}, at::device(at::kCPU).dtype(at::kFloat)); - - // Accessors for efficient access - auto input_accessor = quantized_input.accessor(); - auto output_accessor = float_output.accessor(); - auto weights_accessor = weights_4x2.accessor(); - auto weight_scales_accessor = weight_scales.accessor(); - auto weight_zeros_accessor = weight_zeros.accessor(); - auto input_scale_accessor = input_scale.accessor(); - auto input_zero_accessor = input_zero_point.accessor(); - - // Perform quantized matrix multiplication following quantization.md equation - // (5): result_real_value = lhs_scale * rhs_scale * Sum_over_k( - // (lhs_quantized_value[k] - lhs_zero_point) * - // (rhs_quantized_value[k] - rhs_zero_point) - // ) - for (int64_t b = 0; b < B; b++) { - for (int64_t m = 0; m < M; m++) { - const int64_t token_idx = b * M + m; - const float lhs_scale = - input_scale_accessor[token_idx]; // Per-token input scale - const int32_t lhs_zero_point = - input_zero_accessor[token_idx]; // Per-token input zero point - - for (int64_t n = 0; n < N; n++) { - float result_real_value = 0.0f; - - for (int64_t k = 0; k < K; k++) { - // Get per-group weight quantization parameters - const int64_t group_idx = k / group_size; - const float rhs_scale = - weight_scales_accessor[group_idx][n]; // Per-group weight scale - const int32_t rhs_zero_point = - weight_zeros_accessor[group_idx] - [n]; // Per-group weight zero point - - // Unpack the 4-bit weight for this position - const uint8_t packed_val = weights_accessor[n][k / 2]; - uint8_t weight_4bit; - if (k % 2 == 0) { - weight_4bit = (packed_val & 0xF0) >> 4; // First weight in pair - } else { - weight_4bit = packed_val & 0x0F; // Second weight in pair - } - - // Get quantized values - const int32_t lhs_quantized_value = - static_cast(input_accessor[b][m][k]); - // Convert 4-bit weight to signed: subtract 8 to get range [-8, 7] - const int32_t rhs_quantized_value = - static_cast(weight_4bit) - 8; - - // Apply proper quantization paradigm from quantization.md equation - // (3): real_value = scale * (quantized_value - zero_point) Following - // equation (5): result = lhs_scale * rhs_scale * - // (lhs_quantized - lhs_zero) * (rhs_quantized - rhs_zero) - const float lhs_diff = - static_cast(lhs_quantized_value - lhs_zero_point); - const float rhs_diff = - static_cast(rhs_quantized_value - rhs_zero_point); - - result_real_value += lhs_scale * rhs_scale * lhs_diff * rhs_diff; - } - - output_accessor[b][m][n] = result_real_value; - } - } - } - - return float_output; -} - -at::Tensor linear_qta8a_qga4w_4bit_dequant_impl( - const at::Tensor& quantized_input, - const at::Tensor& input_scale, - const at::Tensor& input_zero_point, - const at::Tensor& weights_4x2, - const int64_t group_size, - const at::Tensor& weight_scales, - const at::Tensor& weight_zeros) { - // Calculate number of input tokens - int64_t input_num_tokens = 1; - for (size_t i = 0; i < quantized_input.sizes().size() - 1; i++) { - input_num_tokens *= quantized_input.size(i); - } - - // Manually dequantize the char tensor using per-token quantization - at::Tensor x_float = at::zeros_like(quantized_input, at::kFloat); - - // Apply per-token dequantization - auto input_accessor = quantized_input.accessor(); - auto output_accessor = x_float.accessor(); - - for (int64_t token_idx = 0; token_idx < input_num_tokens; token_idx++) { - float scale_val = input_scale[token_idx].item(); - int zero_point_val = input_zero_point[token_idx].item(); - - // Calculate batch and sequence indices for this token - int64_t b = token_idx / quantized_input.size(1); - int64_t m = token_idx % quantized_input.size(1); - - // Apply dequantization for all features in this token - for (int64_t k = 0; k < quantized_input.size(-1); k++) { - float dequant_val = - (input_accessor[b][m][k] - zero_point_val) * scale_val; - output_accessor[b][m][k] = dequant_val; - } - } - - std::vector weights_shape(weights_4x2.sizes().vec()); - weights_shape[1] *= 2; - - at::Tensor weights_dequantized = - at::empty(weights_shape, at::device(at::kCPU).dtype(at::kFloat)); - - const int64_t N = weights_dequantized.size(0); - const int64_t K = weights_dequantized.size(1); - - for (int n = 0; n < N; n++) { - for (int k = 0; k < K; k += 2) { - const int group_idx = k / group_size; - const uint8_t packed_val = weights_4x2[n][k / 2].item().to(); - const uint8_t second_val = packed_val & 0x0F; - const uint8_t first_val = (packed_val & 0xF0) >> 4; - - const float scale = weight_scales[group_idx][n].item().to(); - const int zero = weight_zeros[group_idx][n].item().to(); - - weights_dequantized[n][k] = - ((float(first_val) - 8.0) - float(zero)) * scale; - weights_dequantized[n][k + 1] = - ((float(second_val) - 8.0) - float(zero)) * scale; - } - } - - at::Tensor linear_result = at::linear(x_float, weights_dequantized); - - return linear_result; -} - // // Test functions // -void test_reference_linear_qga4w( - const int B, - const int M, - const int K, - const int N, - const int group_size = 32, - const int inner_k_tiles = 8) { - assert(K % group_size == 0); - - at::Tensor x = at::rand({B, M, K}, at::device(at::kCPU).dtype(at::kFloat)); - at::Tensor weights_4x2 = - at::randint(0, 256, {N, K / 2}, at::device(at::kCPU).dtype(at::kByte)); - at::Tensor weights_int = unpack_weights_4x2(weights_4x2); - - const int k_groups = K / group_size; - at::Tensor scales_and_zeros = - at::rand({k_groups, N, 2}, at::device(at::kCPU).dtype(at::kFloat)); - - at::Tensor out = linear_qga4w_reference_impl( - x, - at::_convert_weight_to_int4pack_for_cpu(weights_int, group_size), - group_size, - scales_and_zeros, - inner_k_tiles); - - at::Tensor out_ref = dequantize_and_linear_qga4w( - x, weights_4x2, group_size, scales_and_zeros, inner_k_tiles); - - ASSERT_TRUE(at::allclose(out, out_ref)); -} - void test_reference_linear_qcs4w( const int B, const int M, @@ -389,118 +134,6 @@ void test_reference_linear_qcs4w( ASSERT_TRUE(at::allclose(out, out_ref)); } -void test_vulkan_linear_qga4w_impl( - const int B, - const int M, - const int K, - const int N, - const int group_size = 32, - const int inner_k_tiles = 8, - const vkcompute::utils::StorageType in_storage = - vkcompute::utils::kTexture3D, - const vkcompute::utils::StorageType out_storage = - vkcompute::utils::kTexture3D) { - assert(K % group_size == 0); - - at::Tensor x = at::rand({B, M, K}, at::device(at::kCPU).dtype(at::kFloat)); - at::Tensor weights_4x2 = - at::randint(0, 256, {N, K / 2}, at::device(at::kCPU).dtype(at::kByte)); - - const int k_groups = K / group_size; - at::Tensor scales_and_zeros = - at::rand({k_groups, N, 2}, at::device(at::kCPU).dtype(at::kFloat)); - - at::Tensor weights_int = unpack_weights_4x2(weights_4x2); - at::Tensor out_ref = linear_qga4w_reference_impl( - x, - at::_convert_weight_to_int4pack_for_cpu(weights_int, group_size), - group_size, - scales_and_zeros, - inner_k_tiles); - - // Build Vulkan graph - using namespace vkcompute; - - GraphConfig config; - config.set_storage_type_override(utils::kTexture3D); - ComputeGraph graph(config); - -#define MAKE_TENSORREF_FOR(x) \ - ValueRef r_##x = graph.add_tensorref( \ - x.sizes().vec(), \ - from_at_scalartype(x.scalar_type()), \ - x.const_data_ptr()); - - MAKE_TENSORREF_FOR(weights_4x2); - MAKE_TENSORREF_FOR(scales_and_zeros); - - IOValueRef r_x = graph.add_input_tensor( - x.sizes().vec(), from_at_scalartype(x.scalar_type()), in_storage); - - const ValueRef r_out = graph.add_tensor( - out_ref.sizes().vec(), - from_at_scalartype(out_ref.scalar_type()), - out_storage); - - VK_GET_OP_FN("et_vk.linear_weight_int4.default") - (graph, - {r_x.value, - r_weights_4x2, - graph.add_scalar(group_size), - r_scales_and_zeros, - kDummyValueRef, - r_out}); - - ValueRef staging_out = graph.set_output_tensor(r_out); - - graph.prepare(); - - graph.prepack(); - - // - // Run model - // - - graph.propagate_resize(); - graph.copy_into_staging(r_x.staging, x.const_data_ptr(), x.numel()); - - graph.execute(); - - at::Tensor vk_out = at::empty_like(out_ref); - graph.copy_from_staging( - staging_out, vk_out.mutable_data_ptr(), vk_out.numel()); - - ASSERT_TRUE(at::allclose(vk_out, out_ref, 1e-4, 1e-4)); -} - -void test_vulkan_linear_qga4w( - const int B, - const int M, - const int K, - const int N, - const int group_size = 32, - const int inner_k_tiles = 8) { - test_vulkan_linear_qga4w_impl( - B, - M, - K, - N, - group_size, - inner_k_tiles, - vkcompute::utils::kBuffer, - vkcompute::utils::kBuffer); - - test_vulkan_linear_qga4w_impl( - B, - M, - K, - N, - group_size, - inner_k_tiles, - vkcompute::utils::kTexture3D, - vkcompute::utils::kTexture3D); -} - void test_vulkan_linear_qcs4w_impl( const int B, const int M, @@ -579,210 +212,6 @@ void test_vulkan_linear_qcs4w( B, M, K, N, vkcompute::utils::kTexture3D, vkcompute::utils::kTexture3D); } -void test_vulkan_linear_qta8a_qga4w_impl( - const int B, - const int M, - const int K, - const int N, - const int group_size = 8, - const vkcompute::utils::StorageType in_storage = - vkcompute::utils::kTexture3D, - const vkcompute::utils::StorageType out_storage = - vkcompute::utils::kTexture3D) { - assert(K % group_size == 0); - - const int64_t input_num_tokens = B * M; - const int k_groups = K / group_size; - - at::Tensor input_scale = - at::rand({input_num_tokens}, at::device(at::kCPU).dtype(at::kFloat)); - at::Tensor input_zero_point = at::randint( - -10, 10, {input_num_tokens}, at::device(at::kCPU).dtype(at::kInt)); - - at::Tensor float_x = - at::rand({B, M, K}, at::device(at::kCPU).dtype(at::kFloat)); - - // Create a reference quantized tensor using per-token quantization - // Mimic per-token quantization using at::quantize_per_channel by reshaping - // [num_tokens, features] - at::Tensor float_x_reshaped = float_x.view({input_num_tokens, K}); - at::Tensor qx_ref_reshaped = at::quantize_per_channel( - float_x_reshaped, - input_scale.to(at::kDouble), - input_zero_point.to(at::kLong), - 0, // axis 0 for per-token (first dimension after reshape) - c10::ScalarType::QInt8); - - at::Tensor x = - at::int_repr(qx_ref_reshaped).view(float_x.sizes()).to(at::kChar); - - at::Tensor weights_4x2 = - at::randint(0, 256, {N, K / 2}, at::device(at::kCPU).dtype(at::kByte)); - at::Tensor weight_scales = - at::rand({k_groups, N}, at::device(at::kCPU).dtype(at::kFloat)); - at::Tensor weight_zeros = at::randint( - -128, 128, {k_groups, N}, at::device(at::kCPU).dtype(at::kInt)); - - at::Tensor out_ref = linear_qta8a_qga4w_4bit_dequant_impl( - x, - input_scale, - input_zero_point, - weights_4x2, - group_size, - weight_scales, - weight_zeros); - - // Build Vulkan graph - using namespace vkcompute; - - GraphConfig config; - config.set_storage_type_override(utils::kTexture3D); - ComputeGraph graph(config); - -#define MAKE_TENSORREF_FOR(x) \ - ValueRef r_##x = graph.add_tensorref( \ - x.sizes().vec(), \ - from_at_scalartype(x.scalar_type()), \ - x.const_data_ptr()); - - MAKE_TENSORREF_FOR(weights_4x2); - MAKE_TENSORREF_FOR(weight_scales); - MAKE_TENSORREF_FOR(weight_zeros); - - IOValueRef r_x = graph.add_input_tensor( - x.sizes().vec(), from_at_scalartype(x.scalar_type()), in_storage); - - IOValueRef r_input_scale = graph.add_input_tensor( - input_scale.sizes().vec(), - from_at_scalartype(input_scale.scalar_type()), - utils::kBuffer); - - IOValueRef r_input_zero_point = graph.add_input_tensor( - input_zero_point.sizes().vec(), - from_at_scalartype(input_zero_point.scalar_type()), - utils::kBuffer); - - const ValueRef r_out = graph.add_tensor( - out_ref.sizes().vec(), - from_at_scalartype(out_ref.scalar_type()), - out_storage); - - VK_GET_OP_FN("et_vk.linear_qta8a_qga4w.default") - (graph, - {r_x.value, - r_input_scale.value, - r_input_zero_point.value, - r_weights_4x2, - graph.add_scalar(group_size), - r_weight_scales, - r_weight_zeros, - r_out}); - - ValueRef staging_out = graph.set_output_tensor(r_out); - - graph.prepare(); - - graph.prepack(); - - // - // Run model - // - - graph.propagate_resize(); - graph.copy_into_staging(r_x.staging, x.const_data_ptr(), x.numel()); - graph.copy_into_staging( - r_input_scale.staging, input_scale.const_data_ptr(), input_scale.numel()); - graph.copy_into_staging( - r_input_zero_point.staging, - input_zero_point.const_data_ptr(), - input_zero_point.numel()); - - graph.execute(); - - at::Tensor vk_out = at::empty_like(out_ref); - graph.copy_from_staging( - staging_out, vk_out.mutable_data_ptr(), vk_out.numel()); - - // This is a reference implementation that uses the quantized - // matmul paradigm. It should follow closely with how the vulkan - // implementation works, and demonstrates reasonably close results. - at::Tensor qmm_ref = linear_qta8a_qga4w_quantized_matmul( - x, - input_scale, - input_zero_point, - weights_4x2, - group_size, - weight_scales, - weight_zeros); - - // For quantized int8 operations, allow for 1-unit differences due to rounding - bool is_close = at::allclose(vk_out, out_ref, 5e-3, 5e-3); - if (!is_close) { - std::cout << "qmm_ref: \n" << qmm_ref << std::endl; - std::cout << "out_ref: \n" << out_ref << std::endl; - std::cout << "vk_out: \n" << vk_out << std::endl; - } - - ASSERT_TRUE(is_close); -} - -void test_vulkan_linear_qta8a_qga4w( - const int B, - const int M, - const int K, - const int N, - const int group_size = 32) { - test_vulkan_linear_qta8a_qga4w_impl( - B, - M, - K, - N, - group_size, - vkcompute::utils::kBuffer, - vkcompute::utils::kBuffer); - - test_vulkan_linear_qta8a_qga4w_impl( - B, - M, - K, - N, - group_size, - vkcompute::utils::kTexture3D, - vkcompute::utils::kTexture3D); -} - -// Test linear_qga4w operator - -TEST(VulkanLinearQGA4WTest, test_reference_impl) { - test_reference_linear_qga4w( - /*B = */ 1, - /*M = */ 4, - /*K = */ 128, - /*N = */ 32); -} - -TEST(VulkanLinearQGA4WTest, test_vulkan_impl_small_m) { - test_vulkan_linear_qga4w( - /*B = */ 1, - /*M = */ 4, - /*K = */ 128, - /*N = */ 32); - - test_vulkan_linear_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 256, - /*N = */ 256); -} - -TEST(VulkanLinearQGA4WTest, test_vulkan_impl_gemm) { - test_vulkan_linear_qga4w( - /*B = */ 1, - /*M = */ 256, - /*K = */ 256, - /*N = */ 256); -} - // Test linear_qcs4w operator TEST_F(VulkanLinearQCS4WTest, test_reference_impl) { @@ -814,87 +243,3 @@ TEST_F(VulkanLinearQCS4WTest, test_vulkan_impl_gemm) { /*K = */ 32, /*N = */ 32); } - -// Test linear_qta8a_qga4w operator - -TEST_F( - VulkanLinearQTA8AQGA4WTest, - test_vulkan_linear_quant_gemm_custom_groupsize) { - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 2, - /*K = */ 8, - /*N = */ 8, - /*group_size = */ 8); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 2, - /*K = */ 16, - /*N = */ 8, - /*group_size = */ 8); -} - -TEST_F(VulkanLinearQTA8AQGA4WTest, test_vulkan_linear_quant_gemm) { - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 4, - /*K = */ 64, - /*N = */ 32); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 4, - /*K = */ 128, - /*N = */ 32); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 8, - /*K = */ 64, - /*N = */ 16); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 256, - /*K = */ 256, - /*N = */ 256); -} - -TEST_F( - VulkanLinearQTA8AQGA4WTest, - test_vulkan_linear_quant_gemv_custom_groupsize) { - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 8, - /*N = */ 8, - /*group_size = */ 8); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 16, - /*N = */ 8, - /*group_size = */ 8); -} - -TEST_F(VulkanLinearQTA8AQGA4WTest, test_vulkan_linear_quant_gemv) { - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 32, - /*N = */ 32); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 64, - /*N = */ 16); - - test_vulkan_linear_qta8a_qga4w( - /*B = */ 1, - /*M = */ 1, - /*K = */ 256, - /*N = */ 256); -} diff --git a/backends/vulkan/test/op_tests/sdpa_test.cpp b/backends/vulkan/test/op_tests/sdpa_test.cpp index 90a688047af..a94e68a53af 100644 --- a/backends/vulkan/test/op_tests/sdpa_test.cpp +++ b/backends/vulkan/test/op_tests/sdpa_test.cpp @@ -215,16 +215,13 @@ at::Tensor sdpa_reference_impl( void test_reference_sdpa( const int start_input_pos, const int sequence_len, - const int embedding_dim, + const int head_dim, const int num_heads, const int num_kv_heads, const int batch_size, const int max_seq_len, at::ScalarType dtype = at::kFloat) { - const int head_dim = embedding_dim / num_heads; - // K and V caches. Need an extra set for the reference implementation - at::Tensor k_cache = at::zeros( {batch_size, max_seq_len, num_kv_heads, head_dim}, at::device(at::kCPU).dtype(dtype)); @@ -265,19 +262,23 @@ void test_reference_sdpa( void test_vulkan_sdpa( const int start_input_pos, - const int base_sequence_len, - const int embedding_dim, + const std::vector& sequence_lens, + const int head_dim, const int num_heads, const int num_kv_heads, const int batch_size, - const int max_seq_len, - const bool dynamic_seq_len = true, + vkcompute::utils::StorageType storage_type, at::ScalarType dtype = at::kFloat) { - const int head_dim = embedding_dim / num_heads; + // compute the max sequence length + int max_seq_len = start_input_pos; + for (int i = 0; i < sequence_lens.size(); ++i) { + max_seq_len += sequence_lens[i]; + } + // Add some extra space to the max sequence length + max_seq_len += 128; - const int init_seq_len = dynamic_seq_len ? max_seq_len : base_sequence_len; + const int init_seq_len = max_seq_len; // K and V caches - at::Tensor k_cache = at::zeros( {batch_size, max_seq_len, num_kv_heads, head_dim}, at::device(at::kCPU).dtype(dtype)); @@ -300,7 +301,6 @@ void test_vulkan_sdpa( using namespace vkcompute; GraphConfig config; - config.set_storage_type_override(utils::kTexture3D); ComputeGraph graph(config); // "Data" variant for vulkan initialization @@ -319,7 +319,7 @@ void test_vulkan_sdpa( #define MAKE_INPUT_FOR(x) \ IOValueRef r_##x = graph.add_input_tensor( \ - x.sizes().vec(), from_at_scalartype(x.scalar_type())); + x.sizes().vec(), from_at_scalartype(x.scalar_type()), storage_type); MAKE_INPUT_FOR(q); MAKE_INPUT_FOR(k); @@ -328,7 +328,7 @@ void test_vulkan_sdpa( const ValueRef r_input_pos_symint = graph.add_symint(start_input_pos); const ValueRef r_out = graph.add_tensor( - out.sizes().vec(), from_at_scalartype(out.scalar_type())); + out.sizes().vec(), from_at_scalartype(out.scalar_type()), storage_type); VK_GET_OP_FN("sdpa_with_kv_cache.default") (graph, @@ -365,10 +365,10 @@ void test_vulkan_sdpa( graph.copy_from_staging( \ staging_##x, vk_##x.mutable_data_ptr(), vk_##x.numel()); - int seq_len = base_sequence_len; - for (int i = 0, input_pos = start_input_pos; - input_pos + seq_len < max_seq_len; - input_pos += seq_len, i++) { + torch::manual_seed(0); + + int input_pos = start_input_pos; + for (auto seq_len : sequence_lens) { q = at::rand( {batch_size, seq_len, num_heads, head_dim}, at::device(at::kCPU).dtype(dtype)); @@ -398,6 +398,46 @@ void test_vulkan_sdpa( const bool output_correct = at::allclose(reference_out, vk_out); if (!output_correct) { + // Print only differing tensor elements side by side for easier comparison + auto ref_flat = reference_out.flatten(); + auto vk_flat = vk_out.flatten(); + auto numel = ref_flat.numel(); + std::cout << "reference_out\tvk_out\tindex" << std::endl; + int first_diff_idx = -1; + auto sizes = reference_out.sizes(); + int d0 = sizes[0], d1 = sizes[1], d2 = sizes[2], d3 = sizes[3]; + for (int i = 0; i < numel; ++i) { + if (std::abs(ref_flat[i].item() - vk_flat[i].item()) > + 1e-4) { + // Compute 4-D index from flat index + int i0 = i / (d1 * d2 * d3); + int rem0 = i % (d1 * d2 * d3); + int i1 = rem0 / (d2 * d3); + int rem1 = rem0 % (d2 * d3); + int i2 = rem1 / d3; + int i3 = rem1 % d3; + std::cout << ref_flat[i].item() << "\t" << vk_flat[i].item() << "\t[" + << i0 << ", " << i1 << ", " << i2 << ", " << i3 << "]" + << std::endl; + if (first_diff_idx == -1) { + first_diff_idx = i; + } + break; + } + } + if (first_diff_idx != -1) { + // Compute 4-D index from flat index + int i0 = first_diff_idx / (d1 * d2 * d3); + int rem0 = first_diff_idx % (d1 * d2 * d3); + int i1 = rem0 / (d2 * d3); + int rem1 = rem0 % (d2 * d3); + int i2 = rem1 / d3; + int i3 = rem1 % d3; + std::cout << "First difference at flat index " << first_diff_idx + << " which is tensor index [" << i0 << ", " << i1 << ", " + << i2 << ", " << i3 << "]" << std::endl; + } + at::Tensor diffs = at::abs(reference_out - vk_out); std::cout << "Failed at input_pos " << input_pos << " with seq_len " @@ -414,513 +454,65 @@ void test_vulkan_sdpa( } ASSERT_TRUE(output_correct); - if (dynamic_seq_len) { - seq_len = base_sequence_len + (i % 3); - } + input_pos += seq_len; } } -TEST(VulkanSDPATest, test_sdpa_op_small_params) { - const int starting_input_pos = 0; - const int base_sequence_len = 3; - const int embedding_dim = 18; - const int num_heads = 6; - const int num_kv_heads = 2; - const int batch_size = 1; - const int max_seq_len = 7; - - test_vulkan_sdpa( - starting_input_pos, - base_sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len, - false); -} - -TEST(VulkanSDPATest, test_sdpa_op_small_params_dynamic) { - const int starting_input_pos = 0; - const int base_sequence_len = 3; - const int embedding_dim = 18; - const int num_heads = 6; - const int num_kv_heads = 2; - const int batch_size = 1; - const int max_seq_len = 12; - - test_vulkan_sdpa( - starting_input_pos, - base_sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_sdpa_op_llama3_params_dynamic) { - const int starting_input_pos = 0; - const int base_sequence_len = 3; - const int embedding_dim = 2048; - const int num_heads = 32; - const int num_kv_heads = 8; - const int batch_size = 1; - const int max_seq_len = 128; - - test_vulkan_sdpa( - starting_input_pos, - base_sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_reference_impl) { - const int starting_input_pos = 0; - const int base_sequence_len = 3; - const int embedding_dim = 2048; - const int num_heads = 32; - const int num_kv_heads = 8; - const int batch_size = 1; - const int max_seq_len = 128; - - test_reference_sdpa( - starting_input_pos, - base_sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -void test_vulkan_flash_attention( +void test_vulkan_sdpa( const int start_input_pos, - const int sequence_len, - const int embedding_dim, + const std::vector& sequence_lens, + const int head_dim, const int num_heads, const int num_kv_heads, const int batch_size, - const int max_seq_len, at::ScalarType dtype = at::kFloat) { - const int head_dim = embedding_dim / num_heads; - - at::Tensor k_cache = at::zeros( - {batch_size, max_seq_len, num_kv_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor v_cache = at::zeros_like(k_cache); - - at::Tensor q = at::rand( - {batch_size, sequence_len, num_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor k = at::rand( - {batch_size, sequence_len, num_kv_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor v = at::rand_like(k); - - // Get reference output using existing SDPA - at::Tensor reference_out = sdpa_reference_impl( - q, - k, - v, - k_cache, - v_cache, + // Test texture + test_vulkan_sdpa( start_input_pos, - sequence_len, - {}, - 0.0, - true, - {}); - - using namespace vkcompute; - - GraphConfig config; - config.set_storage_type_override( - utils::kBuffer); // Flash Attention requires buffer storage - ComputeGraph graph(config); - - // Create input references - IOValueRef r_q = graph.add_input_tensor( - q.sizes().vec(), from_at_scalartype(q.scalar_type())); - IOValueRef r_k = graph.add_input_tensor( - k.sizes().vec(), from_at_scalartype(k.scalar_type())); - IOValueRef r_v = graph.add_input_tensor( - v.sizes().vec(), from_at_scalartype(v.scalar_type())); - - // Create cache tensors (these would be updated by cache update operations in - // practice) - ValueRef r_k_cache = graph.add_tensorref( - k_cache.sizes().vec(), - from_at_scalartype(k_cache.scalar_type()), - k_cache.const_data_ptr()); - ValueRef r_v_cache = graph.add_tensorref( - v_cache.sizes().vec(), - from_at_scalartype(v_cache.scalar_type()), - v_cache.const_data_ptr()); - - const ValueRef r_input_pos_symint = graph.add_symint(start_input_pos); - const ValueRef r_out = - graph.add_tensor(q.sizes().vec(), from_at_scalartype(q.scalar_type())); - - // Call Flash Attention implementation - VK_GET_OP_FN("llama.flash_attention.default") - (graph, - { - r_q.value, - r_k.value, // Use actual K tensor, not cache - r_v.value, // Use actual V tensor, not cache - r_input_pos_symint, - kDummyValueRef, // attn_mask - kDummyValueRef, // dropout_p - kDummyValueRef, // is_causal - kDummyValueRef, // scale - r_out, - }); - - ValueRef staging_out = graph.set_output_tensor(r_out); - - graph.prepare(); - graph.encode_prepack(); - graph.prepack(); - - // Copy inputs and run - graph.copy_into_staging(r_q.staging, q.const_data_ptr(), q.numel()); - graph.copy_into_staging(r_k.staging, k.const_data_ptr(), k.numel()); - graph.copy_into_staging(r_v.staging, v.const_data_ptr(), v.numel()); - - graph.execute(); - - // Extract output - at::Tensor vk_out = at::zeros_like(q).contiguous(); - graph.copy_from_staging( - staging_out, vk_out.mutable_data_ptr(), vk_out.numel()); - - // Compare results - const bool output_correct = at::allclose(reference_out, vk_out, 1e-3, 1e-3); - - if (!output_correct) { - at::Tensor diffs = at::abs(reference_out - vk_out); - std::cout << "Flash Attention test failed!" << std::endl; - std::cout << "Maximum difference: " << at::max(diffs).item() << std::endl; - std::cout << "Maximum value observed: " - << at::max(at::abs(at::cat({reference_out, vk_out}, -1))).item() - << std::endl; - } - ASSERT_TRUE(output_correct); -} - -TEST(VulkanSDPATest, test_flash_attention_small_params) { - // TINY DEBUG PARAMETERS - easy to trace by hand - const int starting_input_pos = 0; - const int sequence_len = 2; // Very small sequence - const int embedding_dim = 4; // Very small embedding - const int num_heads = 2; // Just 2 heads - const int num_kv_heads = 2; // Match query heads (no multi-query complexity) - const int batch_size = 1; // Single batch - const int max_seq_len = 4; // Small cache - - test_vulkan_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, + sequence_lens, + head_dim, num_heads, num_kv_heads, batch_size, - max_seq_len); -} + vkcompute::utils::kTexture3D, + dtype); -TEST(VulkanSDPATest, test_flash_attention_multi_tile) { - // MULTI-TILE TEST - tests the tiling algorithm with multiple blocks - // With block_size_r=32, block_size_c=32 (from SDPA.cpp), and seq_len=48: - // - Tr = ceil(48/32) = 2 row tiles (blocks: 0-31, 32-47) - // - Tc = ceil(48/32) = 2 column tiles (blocks: 0-31, 32-47) - // - Total of 2x2 = 4 tile combinations to process per head - // - Memory usage: 48*2*16 = 1,536 elements per tensor (much more reasonable) - const int starting_input_pos = 0; - const int sequence_len = 48; // Moderate size to force multiple tiles - const int embedding_dim = 32; // head_dim = 32/2 = 16 per head - const int num_heads = 2; // 2 heads to keep manageable - const int num_kv_heads = 2; // Match query heads - const int batch_size = 1; // Single batch - const int max_seq_len = 64; // Reasonable cache size - - test_vulkan_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, + // Test buffer + test_vulkan_sdpa( + start_input_pos, + sequence_lens, + head_dim, num_heads, num_kv_heads, batch_size, - max_seq_len); + vkcompute::utils::kBuffer, + dtype); } -// Flash Attention tests corresponding to traditional SDPA tests - -TEST(VulkanSDPATest, test_flash_attention_op_small_params) { - // Corresponds to test_sdpa_op_small_params - const int starting_input_pos = 0; - const int sequence_len = 3; - const int embedding_dim = 18; - const int num_heads = 6; - const int num_kv_heads = 2; - const int batch_size = 1; - const int max_seq_len = 7; +TEST(VulkanSDPATest, test_sdpa_op_small_params) { + const int base_sequence_len = 3; + const int num_heads = 8; + const int head_dim = 4; + const int num_kv_heads = 4; - test_vulkan_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); + test_vulkan_sdpa( + 0, {3, 1, 1, 5, 1, 1, 2}, head_dim, num_heads, num_kv_heads, 1); } -TEST(VulkanSDPATest, test_flash_attention_op_small_params_dynamic) { - // Corresponds to test_sdpa_op_small_params_dynamic - // Note: Flash attention doesn't support dynamic sequence lengths in the same - // way as traditional SDPA, so we test with the base sequence length - const int starting_input_pos = 0; - const int sequence_len = 3; - const int embedding_dim = 18; +TEST(VulkanSDPATest, test_sdpa_op_small_params_dynamic) { + const int base_sequence_len = 3; + const int head_dim = 8; const int num_heads = 6; const int num_kv_heads = 2; - const int batch_size = 1; - const int max_seq_len = 12; - - test_vulkan_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_flash_attention_op_llama3_params) { - // Corresponds to test_sdpa_op_llama3_params_dynamic - // This is a large test that exercises the multi-tile algorithm extensively - const int starting_input_pos = 0; - const int sequence_len = 3; - const int embedding_dim = 2048; - const int num_heads = 32; - const int num_kv_heads = 8; - const int batch_size = 1; - const int max_seq_len = 128; - - test_vulkan_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_flash_attention_op_llama3_params_dynamic) { - // Corresponds to test_sdpa_op_llama3_params_dynamic - // Test with varying sequence lengths to ensure flash attention works with - // different sizes - const int starting_input_pos = 0; - const int embedding_dim = 2048; - const int num_heads = 32; - const int num_kv_heads = 8; - const int batch_size = 1; - const int max_seq_len = 128; - - // Test with different sequence lengths - std::vector sequence_lengths = {1, 3, 5, 7, 16, 32}; - - for (int seq_len : sequence_lengths) { - if (seq_len < max_seq_len) { - test_vulkan_flash_attention( - starting_input_pos, - seq_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); - } - } -} - -void test_reference_flash_attention( - const int start_input_pos, - const int sequence_len, - const int embedding_dim, - const int num_heads, - const int num_kv_heads, - const int batch_size, - const int max_seq_len, - at::ScalarType dtype = at::kFloat) { - const int head_dim = embedding_dim / num_heads; - - // For flash attention reference test, we test single-shot attention - // rather than iterative cache updates, since flash attention processes - // the entire sequence at once - - at::Tensor q = at::rand( - {batch_size, sequence_len, num_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor k = at::rand( - {batch_size, sequence_len, num_kv_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor v = at::rand_like(k); - - // Create empty caches for reference implementation - at::Tensor k_cache_ref = at::zeros( - {batch_size, max_seq_len, num_kv_heads, head_dim}, - at::device(at::kCPU).dtype(dtype)); - at::Tensor v_cache_ref = at::zeros_like(k_cache_ref); - - // Get reference implementation output - at::Tensor reference_out = sdpa_reference_impl( - q, - k, - v, - k_cache_ref, - v_cache_ref, - start_input_pos, - sequence_len, - {}, - 0.0, - true, - {}); - - // Build Vulkan Flash Attention graph - using namespace vkcompute; - GraphConfig config; - config.set_storage_type_override(utils::kBuffer); - ComputeGraph graph(config); - - IOValueRef r_q = graph.add_input_tensor( - q.sizes().vec(), from_at_scalartype(q.scalar_type())); - IOValueRef r_k = graph.add_input_tensor( - k.sizes().vec(), from_at_scalartype(k.scalar_type())); - IOValueRef r_v = graph.add_input_tensor( - v.sizes().vec(), from_at_scalartype(v.scalar_type())); - - // Create empty cache tensors for flash attention - at::Tensor k_cache_flash = at::zeros_like(k_cache_ref); - at::Tensor v_cache_flash = at::zeros_like(v_cache_ref); - - ValueRef r_k_cache = graph.add_tensorref( - k_cache_flash.sizes().vec(), - from_at_scalartype(k_cache_flash.scalar_type()), - k_cache_flash.const_data_ptr()); - ValueRef r_v_cache = graph.add_tensorref( - v_cache_flash.sizes().vec(), - from_at_scalartype(v_cache_flash.scalar_type()), - v_cache_flash.const_data_ptr()); - - const ValueRef r_input_pos_symint = graph.add_symint(start_input_pos); - const ValueRef r_out = - graph.add_tensor(q.sizes().vec(), from_at_scalartype(q.scalar_type())); - - VK_GET_OP_FN("llama.flash_attention.default") - (graph, - { - r_q.value, - r_k.value, - r_v.value, - r_input_pos_symint, - kDummyValueRef, // attn_mask - kDummyValueRef, // dropout_p - kDummyValueRef, // is_causal - kDummyValueRef, // scale - r_out, - }); - - ValueRef staging_out = graph.set_output_tensor(r_out); - - graph.prepare(); - graph.encode_prepack(); - graph.prepack(); - - graph.copy_into_staging(r_q.staging, q.const_data_ptr(), q.numel()); - graph.copy_into_staging(r_k.staging, k.const_data_ptr(), k.numel()); - graph.copy_into_staging(r_v.staging, v.const_data_ptr(), v.numel()); - - graph.execute(); - - at::Tensor flash_out = at::zeros_like(q).contiguous(); - graph.copy_from_staging( - staging_out, flash_out.mutable_data_ptr(), flash_out.numel()); - - // Compare flash attention output with reference implementation - const bool output_correct = - at::allclose(reference_out, flash_out, 1e-3, 1e-3); - - if (!output_correct) { - at::Tensor diffs = at::abs(reference_out - flash_out); - std::cout << "Flash Attention reference test failed" << std::endl; - std::cout << "Maximum difference: " << at::max(diffs).item() << std::endl; - std::cout - << "Maximum value observed: " - << at::max(at::abs(at::cat({reference_out, flash_out}, -1))).item() - << std::endl; - } - ASSERT_TRUE(output_correct); + test_vulkan_sdpa(0, {3, 1, 1, 5, 1, 1}, head_dim, num_heads, num_kv_heads, 1); } -TEST(VulkanSDPATest, test_flash_attention_reference_impl) { - const int starting_input_pos = 0; - const int sequence_len = 3; - const int embedding_dim = 2048; - const int num_heads = 32; +TEST(VulkanSDPATest, test_sdpa_op_llama3_params_dynamic) { + const int head_dim = 128; + const int num_heads = 24; const int num_kv_heads = 8; - const int batch_size = 1; - const int max_seq_len = 128; - test_reference_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_flash_attention_reference_impl_small) { - const int starting_input_pos = 0; - const int sequence_len = 2; - const int embedding_dim = 32; - const int num_heads = 4; - const int num_kv_heads = 2; - const int batch_size = 1; - const int max_seq_len = 16; - - test_reference_flash_attention( - starting_input_pos, - sequence_len, - embedding_dim, - num_heads, - num_kv_heads, - batch_size, - max_seq_len); -} - -TEST(VulkanSDPATest, test_flash_attention_edge_cases) { - // Test with single head (no multi-query complexity) - test_vulkan_flash_attention(0, 1, 8, 1, 1, 1, 4); - - // Test with equal heads (no multi-query complexity) - test_vulkan_flash_attention(0, 2, 16, 4, 4, 1, 8); - - // Test with large head dimension - test_vulkan_flash_attention(0, 2, 128, 2, 1, 1, 8); - - // Test with sequence length that exactly matches block size (32) - test_vulkan_flash_attention(0, 32, 64, 2, 1, 1, 64); - - // Test with sequence length slightly larger than block size - test_vulkan_flash_attention(0, 33, 64, 2, 1, 1, 64); + test_vulkan_sdpa( + 0, {111, 1, 1, 1, 57, 1, 1}, head_dim, num_heads, num_kv_heads, 1); } diff --git a/backends/vulkan/test/op_tests/utils/gen_computegraph.py b/backends/vulkan/test/op_tests/utils/gen_computegraph.py index 4fba14ca16e..490044340d6 100644 --- a/backends/vulkan/test/op_tests/utils/gen_computegraph.py +++ b/backends/vulkan/test/op_tests/utils/gen_computegraph.py @@ -549,15 +549,13 @@ def virtual_resize(self, ref: ValueRefList) -> str: return "" if ref.src_cpp_type == AT_TENSOR: - ret_str = f"{self.graph}{self.dot}get_tensor({ref.name}.value)" - ret_str += f"->virtual_resize({ref.src_cpp_name}.sizes().vec());\n" + ret_str = f"{self.graph}{self.dot}virtual_resize({ref.name}.value, " + ret_str += f"{ref.src_cpp_name}.sizes().vec());\n" elif ref.src_cpp_type == AT_TENSOR_LIST: ret_str = "" ret_str += f"for (int i=0; i < {ref.name}_io_value_refs.size(); i++) {{\n" - ret_str += ( - f" {self.graph}{self.dot}get_tensor({ref.name}_io_value_refs[i].value)" - ) - ret_str += f"->virtual_resize({ref.src_cpp_name}[i].sizes().vec());\n" + ret_str += f" {self.graph}{self.dot}virtual_resize({ref.name}_io_value_refs[i].value, " + ret_str += f"{ref.src_cpp_name}[i].sizes().vec());\n" ret_str += "}\n" else: raise AssertionError(f"{ref.src_cpp_type} not expected") diff --git a/backends/vulkan/test/scripts/test_model.sh b/backends/vulkan/test/scripts/test_model.sh new file mode 100755 index 00000000000..5f06d2c039b --- /dev/null +++ b/backends/vulkan/test/scripts/test_model.sh @@ -0,0 +1,180 @@ +#!/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. + +set -exu + +# Initialize variables +RUN_BUILD=false +RUN_CORRECTNESS_TEST=false +RUN_CLEAN=false +RUN_RECOMPILE=false +MODEL_NAME="" +OUTPUT_DIRECTORY="." + +# Parse arguments +SKIP_NEXT=false +for i in $(seq 1 $#); do + if [[ "$SKIP_NEXT" == true ]]; then + SKIP_NEXT=false + continue + fi + + arg="${!i}" + case $arg in + --build|-b) + RUN_BUILD=true + ;; + --clean|-c) + RUN_CLEAN=true + ;; + --recompile|-rc) + RUN_RECOMPILE=true + ;; + --output_directory|-o) + next_i=$((i + 1)) + if [[ $next_i -le $# ]]; then + OUTPUT_DIRECTORY="${!next_i}" + SKIP_NEXT=true + else + echo "Error: --output_directory|-o requires a value" + exit 1 + fi + ;; + --*|-*) + echo "Unknown argument: $arg" + exit 1 + ;; + *) + if [[ -z "$MODEL_NAME" ]]; then + MODEL_NAME="$arg" + else + echo "Multiple model names provided: $MODEL_NAME and $arg" + exit 1 + fi + ;; + esac +done + +# Determine execution mode based on parsed arguments +if [[ "$RUN_BUILD" == true ]] && [[ -z "$MODEL_NAME" ]]; then + # Build-only mode + RUN_CORRECTNESS_TEST=false +elif [[ "$RUN_BUILD" == true ]] && [[ -n "$MODEL_NAME" ]]; then + # Build and test mode + RUN_CORRECTNESS_TEST=true +elif [[ "$RUN_BUILD" == false ]] && [[ -n "$MODEL_NAME" ]]; then + # Test-only mode + RUN_CORRECTNESS_TEST=true +else + echo "Invalid argument combination. Usage:" + echo " $0 --build|-b [--clean|-c] [--recompile|-rc] [-o|--output_directory DIR] # Build-only mode" + echo " $0 model_name [--build|-b] [--clean|-c] [--recompile|-rc] [-o|--output_directory DIR] # Test mode or build+test mode" + exit 1 +fi + +if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then + PYTHON_EXECUTABLE=python3 +fi +which "${PYTHON_EXECUTABLE}" + +CMAKE_OUTPUT_DIR=cmake-out + +# Only set EXPORTED_MODEL if running correctness test +if [[ "${RUN_CORRECTNESS_TEST}" == true ]]; then + EXPORTED_MODEL=${MODEL_NAME}_vulkan +fi + + +clean_build_directory() { + echo "Cleaning build directory: ${CMAKE_OUTPUT_DIR}" + rm -rf ${CMAKE_OUTPUT_DIR} +} + +recompile() { + cmake --build cmake-out -j64 --target install +} + +build_core_libraries_and_devtools() { + echo "Building core libraries and devtools with comprehensive Vulkan support..." + + # Build core libraries with all required components + cmake . \ + -DCMAKE_INSTALL_PREFIX=cmake-out \ + -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \ + -DEXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=ON \ + -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \ + -DEXECUTORCH_BUILD_DEVTOOLS=ON \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DEXECUTORCH_BUILD_XNNPACK=ON \ + -Bcmake-out && \ + cmake --build cmake-out -j64 --target install + + # Build devtools example runner + cmake examples/devtools \ + -DCMAKE_INSTALL_PREFIX=cmake-out \ + -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -Bcmake-out/examples/devtools && \ + cmake --build cmake-out/examples/devtools -j16 --config Release +} + +run_example_runner() { + ./${CMAKE_OUTPUT_DIR}/examples/devtools/example_runner -bundled_program_path "${OUTPUT_DIRECTORY}/${EXPORTED_MODEL}.bpte" -output_verification +} + +test_bundled_model_with_vulkan() { + # Export model as bundled program with Vulkan backend + "${PYTHON_EXECUTABLE}" -m examples.vulkan.export --model_name="${MODEL_NAME}" --output_dir="${OUTPUT_DIRECTORY}" --bundled + + # Update exported model name for bundled program + EXPORTED_MODEL="${MODEL_NAME}_vulkan" + + # Verify the exported bundled model exists + if [[ ! -f "${OUTPUT_DIRECTORY}/${EXPORTED_MODEL}.bpte" ]]; then + echo "Error: Failed to export bundled model ${MODEL_NAME} with Vulkan backend" + exit 1 + fi + + # Note: Running bundled programs may require different executor runner + echo "Bundled program created successfully. Use appropriate bundled program runner to test." + + run_example_runner +} + + +# Main execution +if [[ "${RUN_BUILD}" == true ]]; then + if [[ "${RUN_CLEAN}" == true ]]; then + clean_build_directory + fi + build_core_libraries_and_devtools +fi + +if [[ "${RUN_RECOMPILE}" == true ]]; then + recompile +fi + +if [[ "${RUN_CORRECTNESS_TEST}" == true ]]; then + echo "Testing ${MODEL_NAME} with Vulkan backend..." + # Always use bundled program testing + test_bundled_model_with_vulkan + + # Check if test completed successfully + if [[ $? -eq 0 ]]; then + echo "Vulkan model test completed successfully!" + else + echo "Vulkan model test failed!" + exit 1 + fi +fi diff --git a/backends/vulkan/test/scripts/test_op.sh b/backends/vulkan/test/scripts/test_op.sh new file mode 100755 index 00000000000..1ec07b7f75f --- /dev/null +++ b/backends/vulkan/test/scripts/test_op.sh @@ -0,0 +1,266 @@ +#!/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. + +set -exu + +# Initialize variables +RUN_BUILD=false +RUN_CLEAN=false +RUN_CLEAN_TESTS=false +RUN_RECOMPILE=false +RUN_TESTS=false +TEST_BINARY="" +ATEN_OP="" + +# Parse arguments +SKIP_NEXT=false +if [[ $# -eq 0 ]]; then + # No arguments provided - run default test + TEST_BINARY="vulkan_op_correctness_tests" + RUN_TESTS=true +else + for i in $(seq 1 $#); do + if [[ "$SKIP_NEXT" == true ]]; then + SKIP_NEXT=false + continue + fi + + arg="${!i}" + case $arg in + --build|-b) + RUN_BUILD=true + ;; + --clean|-c) + RUN_CLEAN=true + RUN_BUILD=true + ;; + --clean_tests|-ct) + RUN_CLEAN_TESTS=true + ;; + --recompile|-rc) + RUN_RECOMPILE=true + ;; + --test|-t) + RUN_TESTS=true + ;; + --aten) + next_i=$((i + 1)) + if [[ $next_i -le $# ]]; then + ATEN_OP="${!next_i}" + TEST_BINARY="vulkan_op_correctness_tests" + RUN_TESTS=true + SKIP_NEXT=true + else + echo "Error: --aten requires an operator name" + exit 1 + fi + ;; + --*|-*) + echo "Unknown argument: $arg" + exit 1 + ;; + *) + if [[ -z "$TEST_BINARY" ]]; then + TEST_BINARY="$arg" + RUN_TESTS=true + else + echo "Multiple test binaries provided: $TEST_BINARY and $arg" + exit 1 + fi + ;; + esac + done +fi + +# Determine execution mode based on parsed arguments +if [[ "$RUN_BUILD" == true ]] && [[ -z "$TEST_BINARY" ]] && [[ "$RUN_TESTS" == false ]]; then + # Build-only mode + echo "Build-only mode" +elif [[ "$RUN_BUILD" == true ]] && [[ -n "$TEST_BINARY" ]]; then + # Build and test mode + echo "Build and test mode for: $TEST_BINARY" +elif [[ "$RUN_BUILD" == false ]] && [[ -n "$TEST_BINARY" ]]; then + # Test-only mode + echo "Test-only mode for: $TEST_BINARY" +elif [[ "$RUN_TESTS" == true ]] && [[ -z "$TEST_BINARY" ]]; then + # Run all available tests + echo "Running all available operator tests" +elif [[ $# -eq 0 ]]; then + # No arguments provided - run default test + TEST_BINARY="vulkan_op_correctness_tests" + RUN_TESTS=true + echo "No arguments provided, running default test: $TEST_BINARY" +else + echo "Invalid argument combination. Usage:" + echo " $0 # Run default vulkan_op_correctness_tests" + echo " $0 --build|-b [--clean|-c] [--clean_tests|-ct] [--recompile|-rc] # Build-only mode" + echo " $0 [test_binary_name] [--build|-b] [--clean|-c] [--clean_tests|-ct] [--recompile|-rc] # Test mode or build+test mode" + echo " $0 --test|-t [--build|-b] [--clean|-c] [--clean_tests|-ct] [--recompile|-rc] # Run all tests mode" + echo " $0 --aten [--build|-b] [--clean|-c] [--clean_tests|-ct] [--recompile|-rc] # Run specific ATen operator test" + echo " $0 --clean_tests|-ct # Clean and rebuild only operator tests" + echo "" + echo "Available test binaries:" + echo " - vulkan_op_correctness_tests" + echo " - vulkan_op_benchmarks" + echo " - compute_graph_op_tests" + echo " - sdpa_test" + exit 1 +fi + +if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then + PYTHON_EXECUTABLE=python3 +fi +which "${PYTHON_EXECUTABLE}" + +CMAKE_OUTPUT_DIR=cmake-out + +clean_build_directory() { + echo "Cleaning build directory: ${CMAKE_OUTPUT_DIR}" + rm -rf ${CMAKE_OUTPUT_DIR} +} + +clean_test_directory() { + echo "Cleaning test build directory: ${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests" + rm -rf ${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests +} + +build_core_libraries() { + cmake . \ + -DCMAKE_INSTALL_PREFIX=cmake-out \ + -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \ + -DEXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=ON \ + -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \ + -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ + -DEXECUTORCH_BUILD_KERNELS_LLM_AOT=ON \ + -DEXECUTORCH_BUILD_DEVTOOLS=ON \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DEXECUTORCH_BUILD_XNNPACK=ON \ + -DEXECUTORCH_BUILD_TESTS=ON \ + -Bcmake-out && \ + cmake --build cmake-out -j64 --target install +} + +build_operator_tests() { + echo "Building Vulkan operator tests..." + + # Prepare CMAKE arguments + CMAKE_ARGS=( + "backends/vulkan/test/op_tests" + "-DCMAKE_INSTALL_PREFIX=cmake-out" + "-DPYTHON_EXECUTABLE=$PYTHON_EXECUTABLE" + "-DCMAKE_CXX_STANDARD=17" + ) + + # Check if TORCH_OPS_YAML_PATH is set + if [[ -n "${TORCH_OPS_YAML_PATH:-}" ]]; then + # Verify that TORCH_OPS_YAML_PATH exists + if [[ ! -d "$TORCH_OPS_YAML_PATH" ]]; then + echo "Error: TORCH_OPS_YAML_PATH directory does not exist: $TORCH_OPS_YAML_PATH" + echo "Please set TORCH_OPS_YAML_PATH to a valid PyTorch native operations directory" + echo "Example: export TORCH_OPS_YAML_PATH=/path/to/pytorch/aten/src/ATen/native" + exit 1 + fi + + # Verify required YAML files exist + if [[ ! -f "$TORCH_OPS_YAML_PATH/native_functions.yaml" ]]; then + echo "Error: Required file not found: $TORCH_OPS_YAML_PATH/native_functions.yaml" + exit 1 + fi + + if [[ ! -f "$TORCH_OPS_YAML_PATH/tags.yaml" ]]; then + echo "Error: Required file not found: $TORCH_OPS_YAML_PATH/tags.yaml" + exit 1 + fi + + echo "Using TORCH_OPS_YAML_PATH: $TORCH_OPS_YAML_PATH" + CMAKE_ARGS+=("-DTORCH_OPS_YAML_PATH=$TORCH_OPS_YAML_PATH") + else + echo "WARNING: TORCH_OPS_YAML_PATH is not set. Building without PyTorch operator definitions." + echo "Some functionality may be limited. To enable full functionality, set TORCH_OPS_YAML_PATH to point to PyTorch's native operations directory." + echo "Example: export TORCH_OPS_YAML_PATH=/path/to/pytorch/aten/src/ATen/native" + fi + + # Build operator tests + cmake "${CMAKE_ARGS[@]}" \ + -Bcmake-out/backends/vulkan/test/op_tests && \ + cmake --build cmake-out/backends/vulkan/test/op_tests -j16 +} + +recompile() { + echo "Recompiling..." + cmake --build cmake-out -j64 --target install + cmake --build cmake-out/backends/vulkan/test/op_tests -j16 +} + +run_operator_test() { + local test_name="$1" + local test_binary_path="" + + case "$test_name" in + "aten") + test_binary_path="${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests/vulkan_op_correctness_tests" + ;; + *) + # Try to find the binary directly + test_binary_path="${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests/${test_name}" + ;; + esac + + if [[ -f "$test_binary_path" ]]; then + echo "Running test binary: $test_binary_path" + + # Add gtest filter if ATEN_OP is specified + if [[ -n "$ATEN_OP" ]]; then + echo "Filtering tests for ATen operator: $ATEN_OP" + "$test_binary_path" --gtest_filter="*${ATEN_OP}*" + else + "$test_binary_path" + fi + else + echo "Error: Test binary not found at $test_binary_path" + echo "Available binaries in ${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests/:" + ls -la "${CMAKE_OUTPUT_DIR}/backends/vulkan/test/op_tests/" 2>/dev/null || echo "Directory not found" + exit 1 + fi +} + +# Main execution +if [[ "${RUN_CLEAN_TESTS}" == true ]]; then + clean_test_directory + build_operator_tests +fi + +if [[ "${RUN_BUILD}" == true ]]; then + if [[ "${RUN_CLEAN}" == true ]]; then + clean_build_directory + fi + build_core_libraries + build_operator_tests +fi + +if [[ "${RUN_RECOMPILE}" == true ]]; then + recompile +fi + +if [[ "${RUN_TESTS}" == true ]]; then + run_operator_test "$TEST_BINARY" + + # Check if tests completed successfully + if [[ $? -eq 0 ]]; then + echo "Vulkan operator tests completed successfully!" + else + echo "Some Vulkan operator tests failed!" + exit 1 + fi +fi diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 4799a22882d..f8194f0b32c 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -10,6 +10,8 @@ import unittest from typing import Tuple +import executorch.backends.vulkan.test.utils as test_utils + import torch from executorch.backends.transforms.convert_dtype_pass import I64toI32 @@ -18,25 +20,35 @@ from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) + from executorch.exir import ( EdgeCompileConfig, EdgeProgramManager, ExecutorchProgramManager, + to_edge_transform_and_lower, +) +from executorch.extension.pybindings.portable_lib import ( # @manual + _load_for_executorch_from_buffer, ) -from torch.export import Dim, export, export_for_training, ExportedProgram +from executorch.extension.pytree import tree_flatten +from torch.export import Dim, export, ExportedProgram + +from torchao.quantization.granularity import PerGroup from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import Quantizer +from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ +from torchao.utils import unwrap_tensor_subclass -ctypes.CDLL("libvulkan.so.1") - - -from executorch.exir import to_edge_transform_and_lower -from executorch.extension.pybindings.portable_lib import ( # @manual - _load_for_executorch_from_buffer, -) -from executorch.extension.pytree import tree_flatten +try: + ctypes.CDLL("libvulkan.so.1") +except: + pass def lower_module( @@ -80,11 +92,11 @@ def quantize_and_lower_module( _skip_dim_order=False, # TODO(T182928844): Delegate dim order op to backend. ) - program = export_for_training( + program = export( model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True ).module() - program = prepare_pt2e(program, quantizer) # pyre-ignore + program = prepare_pt2e(program, quantizer) # Calibrate program(*sample_inputs) @@ -126,37 +138,47 @@ def assert_outputs_equal( # Multiple outputs executor always returns tuple, even if there is one output self.assertTrue(len(ref_output) == len(model_output)) if first_output_only: - self.assertTrue( - torch.allclose( - model_output[0], - ref_output[0], + result = torch.allclose( + model_output[0], + ref_output[0], + atol=atol, + rtol=rtol, + equal_nan=equal_nan, + ) + if not result: + test_utils.print_tensor_comparison_errors( + model_output[0], ref_output[0], atol, rtol + ) + self.assertTrue(result) + else: + for i in range(len(ref_output)): + result = torch.allclose( + model_output[i], + ref_output[i], atol=atol, rtol=rtol, equal_nan=equal_nan, ) - ) - else: - for i in range(len(ref_output)): - self.assertTrue( - torch.allclose( - model_output[i], - ref_output[i], - atol=atol, - rtol=rtol, - equal_nan=equal_nan, + if not result: + print(f"\n=== Output {i} comparison failed ===") + test_utils.print_tensor_comparison_errors( + model_output[i], ref_output[i], atol, rtol ) - ) + self.assertTrue(result) else: # If one output, eager returns tensor while executor tuple of size 1 - self.assertTrue( - torch.allclose( - model_output[0], - ref_output, - atol=atol, - rtol=rtol, - equal_nan=equal_nan, - ) + result = torch.allclose( + model_output[0], + ref_output, + atol=atol, + rtol=rtol, + equal_nan=equal_nan, ) + if not result: + test_utils.print_tensor_comparison_errors( + model_output[0], ref_output, atol, rtol + ) + self.assertTrue(result) def check_no_delegation(self, et_program: ExecutorchProgramManager): self.assertEqual( @@ -1774,20 +1796,6 @@ def forward(self, x): (torch.rand(size=[1, 5, 2, 3]),), ) - def test_vulkan_backend_high_dim_tensors_fail(self): - class UnsqueezeHigherDim(torch.nn.Module): - def __init__(self): - super().__init__() - - def forward(self, x): - return torch.unsqueeze(x, 2) - - self.lower_module_and_test_output( - UnsqueezeHigherDim(), - (torch.ones(size=[5, 4, 1, 2, 6]),), - expect_no_delegates=True, - ) - def test_vulkan_backend_large_linear_layer(self): class LinearModel(torch.nn.Module): def __init__(self, large_out_channels: int) -> None: @@ -2150,3 +2158,581 @@ def forward(self, a, b): ) self.lower_module_and_test_output(custom_complex_module, sample_inputs) + + def test_vulkan_backend_cat_width_dynamic_shapes(self): + class CatWidthModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x1, x2, x3, x4, x5, x6): + return torch.cat([x1, x2, x3, x4, x5, x6], dim=3) + + cat_width_module = CatWidthModule() + + # Create 6 tensors with different widths but same batch, channel, and height dimensions + sample_inputs = ( + torch.randn(size=(2, 3, 4, 5), dtype=torch.float32), # width=5 + torch.randn(size=(2, 3, 4, 3), dtype=torch.float32), # width=3 + torch.randn(size=(2, 3, 4, 7), dtype=torch.float32), # width=7 + torch.randn(size=(2, 3, 4, 2), dtype=torch.float32), # width=2 + torch.randn(size=(2, 3, 4, 4), dtype=torch.float32), # width=4 + torch.randn(size=(2, 3, 4, 6), dtype=torch.float32), # width=6 + ) + + # Define dynamic shapes for the width dimension (dim=3) for each input + width1 = Dim("width1", min=1, max=10) + width2 = Dim("width2", min=1, max=10) + width3 = Dim("width3", min=1, max=10) + width4 = Dim("width4", min=1, max=10) + width5 = Dim("width5", min=1, max=10) + width6 = Dim("width6", min=1, max=10) + + dynamic_shapes = { + "x1": {3: width1}, + "x2": {3: width2}, + "x3": {3: width3}, + "x4": {3: width4}, + "x5": {3: width5}, + "x6": {3: width6}, + } + + # Create test inputs with different width combinations + test_inputs = [ + ( + torch.randn(2, 3, 4, 2), # width=2 + torch.randn(2, 3, 4, 1), # width=1 + torch.randn(2, 3, 4, 3), # width=3 + torch.randn(2, 3, 4, 1), # width=1 + torch.randn(2, 3, 4, 2), # width=2 + torch.randn(2, 3, 4, 4), # width=4 + ), + ( + torch.randn(2, 3, 4, 8), # width=8 + torch.randn(2, 3, 4, 2), # width=2 + torch.randn(2, 3, 4, 1), # width=1 + torch.randn(2, 3, 4, 3), # width=3 + torch.randn(2, 3, 4, 5), # width=5 + torch.randn(2, 3, 4, 1), # width=1 + ), + ( + torch.randn(2, 3, 4, 1), # width=1 + torch.randn(2, 3, 4, 9), # width=9 + torch.randn(2, 3, 4, 2), # width=2 + torch.randn(2, 3, 4, 4), # width=4 + torch.randn(2, 3, 4, 1), # width=1 + torch.randn(2, 3, 4, 3), # width=3 + ), + ] + + self.lower_module_and_test_output( + cat_width_module, + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_cat_channels_dynamic_shapes(self): + class CatChannelsModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x1, x2, x3, x4, x5, x6): + return torch.cat([x1, x2, x3, x4, x5, x6], dim=1) + + cat_channels_module = CatChannelsModule() + + # Create 6 tensors with different channel counts but same batch, height, and width dimensions + sample_inputs = ( + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=4 + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=2 + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=6 + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=1 + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=3 + torch.randn(size=(2, 8, 8, 6), dtype=torch.float32), # channels=5 + ) + + # Define dynamic shapes for the channels dimension (dim=1) for each input + channels1 = Dim("channels1", min=1, max=8) + channels2 = Dim("channels2", min=1, max=8) + channels3 = Dim("channels3", min=1, max=8) + channels4 = Dim("channels4", min=1, max=8) + channels5 = Dim("channels5", min=1, max=8) + channels6 = Dim("channels6", min=1, max=8) + + dynamic_shapes = { + "x1": {1: channels1}, + "x2": {1: channels2}, + "x3": {1: channels3}, + "x4": {1: channels4}, + "x5": {1: channels5}, + "x6": {1: channels6}, + } + + # Create test inputs with different channel combinations + test_inputs = [ + ( + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 2, 8, 6), # channels=2 + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 3, 8, 6), # channels=3 + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 2, 8, 6), # channels=2 + ), + ( + torch.randn(2, 6, 8, 6), # channels=6 + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 3, 8, 6), # channels=3 + torch.randn(2, 2, 8, 6), # channels=2 + torch.randn(2, 4, 8, 6), # channels=4 + torch.randn(2, 1, 8, 6), # channels=1 + ), + ( + torch.randn(2, 2, 8, 6), # channels=2 + torch.randn(2, 7, 8, 6), # channels=7 + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 1, 8, 6), # channels=1 + torch.randn(2, 3, 8, 6), # channels=3 + torch.randn(2, 2, 8, 6), # channels=2 + ), + ] + + self.lower_module_and_test_output( + cat_channels_module, + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_high_dimensional_tensors(self): + class HighDimTensorModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, y): + # Unsqueeze inputs twice to create 5-dim tensors + x_5d = torch.unsqueeze(torch.unsqueeze(x, 0), 0) + y_5d = torch.unsqueeze(torch.unsqueeze(y, 0), 0) + # Add tensors together + result = x_5d + y_5d + return result + + high_dim_module = HighDimTensorModule() + # Create 2 4-dim inputs + sample_inputs = ( + torch.rand(size=(2, 3, 4, 5), dtype=torch.float32), + torch.rand(size=(2, 3, 4, 5), dtype=torch.float32), + ) + + self.lower_module_and_test_output(high_dim_module, sample_inputs) + + def test_vulkan_backend_torchao_wo_quantized_linear(self): + in_features = 1024 + out_features = 512 + bias = False + group_size = 64 + weight_bits = 4 + + class TorchAOQuantizedLinearModule(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + group_size: int = 64, + weight_bits: int = 4, + ): + super().__init__() + self.linear = torch.nn.Linear(in_features, out_features, bias=bias) + self.group_size = group_size + self.weight_bits = weight_bits + + if self.weight_bits == 4: + self.weight_dtype = torch.int4 + else: + self.weight_dtype = torch.int8 + + self.quant_granularity = PerGroup(self.group_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + def apply_quantization(self): + """Apply TorchAO weight-only quantization to the linear layer.""" + q_config = IntxWeightOnlyConfig( + weight_dtype=self.weight_dtype, + granularity=self.quant_granularity, + ) + quantize_(self, q_config) + unwrap_tensor_subclass(self) + return self + + # Test with GEMV pattern (batch_size=1, seq_len=1) + quantized_linear_module = TorchAOQuantizedLinearModule( + in_features=in_features, + out_features=out_features, + bias=bias, + group_size=group_size, + weight_bits=weight_bits, + ) + + # Apply quantization + quantized_linear_module = quantized_linear_module.apply_quantization() + + # Test with 2D input (GEMV pattern) + sample_inputs = (torch.randn(size=(1, in_features), dtype=torch.float32),) + + # Use higher tolerance since quantization introduces some error + self.lower_module_and_test_output( + quantized_linear_module, sample_inputs, atol=1e-2, rtol=1e-2 + ) + + # Test with GEMM pattern (batch_size > 1) + quantized_linear_module_gemm = TorchAOQuantizedLinearModule( + in_features=in_features, + out_features=out_features, + bias=bias, + group_size=group_size, + weight_bits=weight_bits, + ) + + # Apply quantization + quantized_linear_module_gemm = quantized_linear_module_gemm.apply_quantization() + + # Test with 3D input (GEMM pattern) + sample_inputs_gemm = ( + torch.randn(size=(1, 248, in_features), dtype=torch.float32), + ) + + # Use higher tolerance since quantization introduces some error + self.lower_module_and_test_output( + quantized_linear_module_gemm, sample_inputs_gemm, atol=1e-2, rtol=1e-2 + ) + + def test_vulkan_backend_xnnpack_pt2e_quantized_linear_sequence(self): + """ + Test a sequence of linear layers quantized with XNNPACK quantization config. + This test creates a module with multiple linear layers in sequence and applies + XNNPACK symmetric quantization to test the quantized model execution. + """ + + import executorch.backends.vulkan.test.utils as test_utils + + class LinearSequenceModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(128, 64, bias=False) + self.linear2 = torch.nn.Linear(64, 32, bias=False) + self.linear3 = torch.nn.Linear(32, 16, bias=False) + + MAX = 0.75 + MIN = -0.25 + self.linear1.weight.data = test_utils.random_uniform_tensor( + self.linear1.weight.shape, MIN, MAX + ) + self.linear2.weight.data = test_utils.random_uniform_tensor( + self.linear2.weight.shape, MIN, MAX + ) + self.linear3.weight.data = test_utils.random_uniform_tensor( + self.linear3.weight.shape, MIN, MAX + ) + + def forward(self, x): + x = self.linear1(x) + x = self.linear2(x) + x = self.linear3(x) + return x + + # Create the module + linear_sequence_module = LinearSequenceModule() + + M = 32 + # Create sample inputs + sample_inputs = ( + ( + test_utils.random_uniform_tensor( + (M, linear_sequence_module.linear1.in_features), + -0.25, + 0.75, + ) + ), + ) + + # Create XNNPACK quantizer with symmetric quantization config + quantizer = XNNPACKQuantizer() + operator_config = get_symmetric_quantization_config( + is_per_channel=True, + is_dynamic=False, + ) + quantizer.set_global(operator_config) + + # Test the quantized module using the existing quantize_and_lower_module function + # Use higher tolerance since quantization introduces some error + edge_program = quantize_and_lower_module( + linear_sequence_module, sample_inputs, quantizer + ) + + et_program = edge_program.to_executorch() + self.check_vk_delegation(et_program) + + self.run_delegated_model_and_check_output( + et_program, + linear_sequence_module, + sample_inputs, + atol=1e-2, + rtol=1e-1, + ) + + def test_vulkan_backend_xnnpack_pt2e_quantized_conv_sequence(self): + """ + Test a sequence of convolution layers quantized with PT2E quantization. + This test creates a module with multiple Conv2d layers in sequence and applies + XNNPACK symmetric quantization to test the quantized model execution. + Similar to the linear sequence test but using convolution layers. + """ + + import executorch.backends.vulkan.test.utils as test_utils + + class ConvSequenceModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d( + in_channels=3, + out_channels=16, + kernel_size=3, + padding=1, + bias=False, + ) + self.conv2 = torch.nn.Conv2d( + in_channels=16, + out_channels=32, + kernel_size=3, + padding=1, + bias=False, + ) + self.conv3 = torch.nn.Conv2d( + in_channels=32, + out_channels=64, + kernel_size=3, + padding=1, + bias=False, + ) + + MAX = 0.75 + MIN = -0.25 + self.conv1.weight.data = test_utils.random_uniform_tensor( + self.conv1.weight.shape, MIN, MAX + ) + self.conv2.weight.data = test_utils.random_uniform_tensor( + self.conv2.weight.shape, MIN, MAX + ) + self.conv3.weight.data = test_utils.random_uniform_tensor( + self.conv3.weight.shape, MIN, MAX + ) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + x = self.conv3(x) + return x + + # Create the module + conv_sequence_module = ConvSequenceModule() + + input_tensor = test_utils.random_uniform_tensor( + (1, 3, 32, 32), + -0.25, + 0.75, + ) + + # Create sample inputs + sample_inputs = (input_tensor,) + + # Create XNNPACK quantizer with symmetric quantization config + quantizer = XNNPACKQuantizer() + operator_config = get_symmetric_quantization_config( + is_per_channel=True, + is_dynamic=False, + ) + quantizer.set_global(operator_config) + + # Test the quantized module using the existing quantize_and_lower_module function + # Use higher tolerance since quantization introduces some error + edge_program = quantize_and_lower_module( + conv_sequence_module, sample_inputs, quantizer + ) + + et_program = edge_program.to_executorch() + self.check_vk_delegation(et_program) + + self.run_delegated_model_and_check_output( + et_program, + conv_sequence_module, + sample_inputs, + atol=1e-2, + rtol=1e-1, + ) + + def test_vulkan_backend_xnnpack_pt2e_quantized_conv_sequence_all_reduced(self): + """ + Test a sequence of convolution layers quantized with PT2E quantization. + This test creates a module with multiple Conv2d layers in sequence and applies + XNNPACK symmetric quantization to test the quantized model execution. + Similar to the linear sequence test but using convolution layers. + """ + + import executorch.backends.vulkan.test.utils as test_utils + + class ConvSequenceModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d( + in_channels=3, + out_channels=32, + kernel_size=3, + padding=1, + bias=False, + ) + self.conv2 = torch.nn.Conv2d( + in_channels=32, + out_channels=1, + kernel_size=3, + padding=1, + bias=False, + ) + + MAX = 0.75 + MIN = -0.25 + self.conv1.weight.data = test_utils.random_uniform_tensor( + self.conv1.weight.shape, MIN, MAX + ) + self.conv2.weight.data = test_utils.random_uniform_tensor( + self.conv2.weight.shape, MIN, MAX + ) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + # Create the module + conv_sequence_module = ConvSequenceModule() + + input_tensor = test_utils.random_uniform_tensor( + (1, 3, 32, 32), + -0.25, + 0.75, + ) + + # Create sample inputs + sample_inputs = (input_tensor,) + + # Create XNNPACK quantizer with symmetric quantization config + quantizer = XNNPACKQuantizer() + operator_config = get_symmetric_quantization_config( + is_per_channel=True, + is_dynamic=False, + ) + quantizer.set_global(operator_config) + + # Test the quantized module using the existing quantize_and_lower_module function + # Use higher tolerance since quantization introduces some error + edge_program = quantize_and_lower_module( + conv_sequence_module, sample_inputs, quantizer + ) + + et_program = edge_program.to_executorch() + self.check_vk_delegation(et_program) + + self.run_delegated_model_and_check_output( + et_program, + conv_sequence_module, + sample_inputs, + atol=1e-2, + rtol=1e-1, + ) + + @unittest.skip("Cannot run on swiftshader due to no 8-bit int support") + def test_vulkan_backend_torchao_8da4w_quantized_linear(self): + """ + Test TorchAO 8da4w quantization (int8 dynamic activation + int4 weight) with Vulkan backend. + This test uses the same quantization approach as the 8da4w qmode in quantize.py. + """ + in_features = 1024 + out_features = 512 + bias = False + group_size = 128 + + class TorchAO8da4wQuantizedLinearModule(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + group_size: int = 128, + ): + super().__init__() + self.linear = torch.nn.Linear(in_features, out_features, bias=bias) + self.group_size = group_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + def apply_8da4w_quantization(self): + """Apply TorchAO 8da4w quantization (int8 dynamic activation + int4 weight).""" + from torchao.quantization import ( + Int8DynamicActivationIntxWeightConfig, + quantize_, + ) + from torchao.quantization.granularity import PerGroup + from torchao.utils import unwrap_tensor_subclass + + quantize_( + self, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, granularity=PerGroup(self.group_size) + ), + ) + unwrap_tensor_subclass(self) + return self + + # Test with GEMV pattern (batch_size=1, seq_len=1) + quantized_linear_module = TorchAO8da4wQuantizedLinearModule( + in_features=in_features, + out_features=out_features, + bias=bias, + group_size=group_size, + ) + + # Apply 8da4w quantization + quantized_linear_module = quantized_linear_module.apply_8da4w_quantization() + + # Test with 2D input (GEMV pattern) + sample_inputs = (torch.randn(size=(1, in_features), dtype=torch.float32),) + + # Use higher tolerance since quantization introduces some error + self.lower_module_and_test_output( + quantized_linear_module, sample_inputs, atol=1e-2, rtol=1e-2 + ) + + # Test with GEMM pattern (batch_size > 1) + quantized_linear_module_gemm = TorchAO8da4wQuantizedLinearModule( + in_features=in_features, + out_features=out_features, + bias=bias, + group_size=group_size, + ) + + # Apply 8da4w quantization + quantized_linear_module_gemm = ( + quantized_linear_module_gemm.apply_8da4w_quantization() + ) + + # Test with 3D input (GEMM pattern) + sample_inputs_gemm = ( + torch.randn(size=(1, 248, in_features), dtype=torch.float32), + ) + + # Use higher tolerance since quantization introduces some error + self.lower_module_and_test_output( + quantized_linear_module_gemm, sample_inputs_gemm, atol=1e-2, rtol=1e-2 + ) diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index 4f54bc638ba..4a30ab6c2de 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -5,6 +5,7 @@ from executorch.backends.transforms.addmm_mm_to_linear import AddmmToLinearTransform from executorch.backends.vulkan._passes import FuseQuantizedOpsTransform +from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( get_symmetric_quantization_config, @@ -16,7 +17,6 @@ from executorch.exir.backend.canonical_partitioners.config_partitioner import ( format_target_name, ) -from torchao.quantization.linear_quant_modules import Int8DynActInt4WeightQuantizer from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import Quantizer @@ -57,7 +57,7 @@ def quantize_and_lower_module( _check_ir_validity=False, ) - program = torch.export.export_for_training( + program = torch.export.export( model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True ).module() @@ -94,7 +94,6 @@ def op_node_count(graph_module: torch.fx.GraphModule, canonical_op_name: str) -> class TestVulkanPasses(unittest.TestCase): - def test_fuse_int8pack_mm(self): K = 256 N = 256 @@ -155,55 +154,106 @@ def test_fuse_linear_qcs4w(self): self.assertEqual(op_node_count(gm, "linear_qcs4w.default"), 1) self.assertEqual(op_node_count(gm, "dequantize_per_channel.default"), 0) - def test_fuse_linear_qta8a_qga4w(self): - """Test fusion of dynamic activation + grouped weight quantized linear (QTA8A_QGA4W).""" - K = 256 - N = 256 - model = SingleLinearModule(K, N) - sample_inputs = model.get_sample_inputs() - - # Use source transform quantizer for dynamic activation + grouped weight quantization - quantizer = Int8DynActInt4WeightQuantizer( - groupsize=128, # Group size for 4-bit weights - padding_allowed=False, - precision=torch.float32, - scales_precision=torch.float32, - device=torch.device("cpu"), - ) - - # Apply source transform quantization - quantized_model = quantizer.quantize(model) - - # Export the quantized model + def test_fuse_rotary_emb(self): + """Test conversion of rotary embedding pattern to et_vk.apply_rotary_emb custom op.""" + + class RotaryEmbeddingModel(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ): + # This implementation matches the apply_rotary_emb function in rope.py + # Split into real and imaginary parts + xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1) + xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1) + + # Reshape frequencies for broadcasting + freqs_cos = self._reshape_for_broadcast(freqs_cos, xq_r) + freqs_sin = self._reshape_for_broadcast(freqs_sin, xq_r) + + # Apply rotary embedding + xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin + xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos + xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin + xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos + + # Recombine real and imaginary parts + xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3) + xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3) + + return xq_out.type_as(xq), xk_out.type_as(xk) + + def _reshape_for_broadcast(self, freqs_cis: torch.Tensor, x: torch.Tensor): + """Helper function to reshape frequencies for broadcasting""" + ndim = x.ndim + freqs_cis_ndim = freqs_cis.ndim + if freqs_cis_ndim == 3: + # freqs_cis: (seq_len, n_heads, head_dim // 2) + shape = [ + d if (i == ndim - 3 or i == ndim - 2 or i == ndim - 1) else 1 + for i, d in enumerate(x.shape) + ] + else: + # freqs_cis: (seq_len, head_dim // 2) + shape = [ + d if i == 1 or i == ndim - 1 else 1 + for i, d in enumerate(x.shape) + ] + return freqs_cis.view(shape) + + # Create sample inputs based on the test file + batch_size = 1 + seq_len = 5 + n_heads = 32 + n_kv_heads = 8 + head_dim = 2048 + + xq = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=torch.float) + xk = torch.randn(batch_size, seq_len, n_kv_heads, head_dim, dtype=torch.float) + freqs_cos = torch.randn(seq_len, head_dim // 2, dtype=torch.float) + freqs_sin = torch.randn(seq_len, head_dim // 2, dtype=torch.float) + + sample_inputs = (xq, xk, freqs_cos, freqs_sin) + + model = RotaryEmbeddingModel() + + # Export the model edge_compile_config = EdgeCompileConfig( _skip_dim_order=False, _check_ir_validity=False, ) - program = torch.export.export_for_training( - quantized_model, sample_inputs, strict=True - ).module() - - program = torch.export.export(program, sample_inputs) + program = torch.export.export(model, sample_inputs, strict=True) edge_manager = to_edge( program, compile_config=edge_compile_config, ) + # Apply the rotary embedding pass ep = edge_manager._edge_programs["forward"] - edge_manager.transform( - [ - AddmmToLinearTransform(), - FuseQuantizedOpsTransform(ep), - ] - ) + rotary_pass = FusePatternsPass(ep) + result = rotary_pass.call(ep.graph_module) - gm = ep.graph_module + # Verify that the pass was successful + self.assertTrue(result.modified) - # Check that the linear_qta8a_qga4w operator was created - self.assertEqual(op_node_count(gm, "linear_qta8a_qga4w.default"), 1) - # Check that the original quantization/dequantization nodes were removed - self.assertEqual(op_node_count(gm, "quantize_per_token.default"), 0) - self.assertEqual(op_node_count(gm, "dequantize_per_channel.default"), 0) - self.assertEqual(op_node_count(gm, "linear.default"), 0) + # Check that the custom op was created + gm = ep.graph_module + custom_op_count = 0 + for node in gm.graph.nodes: + if ( + node.op == "call_function" + and hasattr(node.target, "__name__") + and "apply_rotary_emb" in str(node.target) + ): + custom_op_count += 1 + + # We expect at least one custom op to be created + self.assertGreater(custom_op_count, 0) diff --git a/backends/vulkan/test/tester.py b/backends/vulkan/test/tester.py index def5aa05e5f..b2066a06ec0 100644 --- a/backends/vulkan/test/tester.py +++ b/backends/vulkan/test/tester.py @@ -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. -from typing import Any, List, Optional, Tuple +from typing import Any, List, Optional, Sequence, Tuple import executorch import executorch.backends.test.harness.stages as BaseStages @@ -13,8 +13,33 @@ from executorch.backends.test.harness import Tester as TesterBase from executorch.backends.test.harness.stages import StageType from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.quantizer.vulkan_quantizer import ( + get_symmetric_quantization_config as get_symmetric_quantization_config_vulkan, + VulkanQuantizer, +) from executorch.exir import EdgeCompileConfig from executorch.exir.backend.partitioner import Partitioner +from torchao.quantization.pt2e.quantizer import Quantizer + + +class Quantize(BaseStages.Quantize): + def __init__( + self, + quantizer: Optional[Quantizer] = None, + quantization_config: Any | None = None, + calibrate: bool = True, + calibration_samples: Optional[Sequence[Any]] = None, + is_qat: Optional[bool] = False, + ): + super().__init__( + quantizer=quantizer or VulkanQuantizer(), + quantization_config=( + quantization_config or get_symmetric_quantization_config_vulkan() + ), + calibrate=calibrate, + calibration_samples=calibration_samples, + is_qat=is_qat, + ) class Partition(BaseStages.Partition): diff --git a/backends/vulkan/test/utils.py b/backends/vulkan/test/utils.py new file mode 100644 index 00000000000..41c1d92bd00 --- /dev/null +++ b/backends/vulkan/test/utils.py @@ -0,0 +1,785 @@ +# 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 logging +from collections import OrderedDict +from copy import deepcopy + +from enum import auto, Enum +from typing import Any, List, Optional, Tuple + +import executorch.backends.vulkan.utils as utils + +import torch + +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner + +from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, +) +from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend +from executorch.devtools import BundledProgram +from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite +from executorch.devtools.bundled_program.serialize import ( + serialize_from_bundled_program_to_flatbuffer, +) +from executorch.exir import ExecutorchProgramManager, to_edge_transform_and_lower +from executorch.extension.pybindings.portable_lib import ( # @manual + _load_for_executorch_from_buffer, +) +from executorch.extension.pytree import tree_flatten +from torch.export import export + +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +class QuantizationMode(Enum): + """Enum to describe how a model should be quantized.""" + + NONE = auto() + INT8_STATIC_PER_CHANNEL = auto() + + +def get_exported_graph( + model, + sample_inputs, + dynamic_shapes=None, + qmode=QuantizationMode.NONE, +) -> torch.fx.GraphModule: + export_training_graph = export( + model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() + + if qmode == QuantizationMode.NONE: + return export_training_graph + + quantizer = XNNPACKQuantizer() + + operator_config = get_symmetric_quantization_config(is_per_channel=True) + quantizer.set_global(operator_config) + + prepared_graph = prepare_pt2e(export_training_graph, quantizer) + prepared_graph(*sample_inputs) + converted_graph = convert_pt2e(prepared_graph) + + return converted_graph + + +def random_uniform_tensor(shape, low=0.0, high=1.0, device=None, dtype=None): + if dtype is None: + dtype = torch.float32 + + return torch.empty(shape, device=device, dtype=dtype).uniform_(low, high) + + +def export_model_to_vulkan( + model, + sample_inputs, + dynamic_shapes=None, + operator_blocklist=None, + operator_allowlist=None, + nn_module_blocklist=None, + nn_module_allowlist=None, + qmode=QuantizationMode.NONE, +): + compile_options = {} + exported_graph = get_exported_graph(model, sample_inputs, qmode=qmode) + program = export( + exported_graph, + sample_inputs, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + + edge_program = to_edge_transform_and_lower( + program, + partitioner=[ + VulkanPartitioner( + compile_options, + operator_blocklist=operator_blocklist, + operator_allowlist=operator_allowlist, + nn_module_blocklist=nn_module_blocklist, + nn_module_allowlist=nn_module_allowlist, + ) + ], + transform_passes=None, + compile_config=None, + ) + + executorch_program = edge_program.to_executorch() + + # Check if the delegate ID matches VulkanBackend + if ( + executorch_program.executorch_program.execution_plan[0].delegates[0].id + != VulkanBackend.__name__ + ): + raise RuntimeError( + f"Expected delegate ID {VulkanBackend.__name__}, but got {executorch_program.executorch_program.execution_plan[0].delegates[0].id}" + ) + + return executorch_program + + +def export_model_to_xnnpack( + model, + sample_inputs, + dynamic_shapes=None, + operator_blocklist=None, + operator_allowlist=None, + nn_module_blocklist=None, + nn_module_allowlist=None, + qmode=QuantizationMode.NONE, +): + compile_options = {} + exported_graph = get_exported_graph(model, sample_inputs, qmode=qmode) + program = export( + exported_graph, + sample_inputs, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + + edge_program = to_edge_transform_and_lower( + program, + partitioner=[XnnpackPartitioner(compile_options)], + transform_passes=None, + compile_config=None, + ) + + executorch_program = edge_program.to_executorch() + + # Check if the delegate ID matches XnnpackBackend + if ( + executorch_program.executorch_program.execution_plan[0].delegates[0].id + != XnnpackBackend.__name__ + ): + raise RuntimeError( + f"Expected delegate ID {XnnpackBackend.__name__}, but got {executorch_program.executorch_program.execution_plan[0].delegates[0].id}" + ) + + return executorch_program + + +def print_tensor_comparison_errors( + tensor1, tensor2, atol=1e-03, rtol=1e-03, max_errors=10 +): + """ + Print the first max_errors tensor indexes that exceed the absolute/relative tolerance + and the error at each of those locations. + + Args: + tensor1: First tensor to compare + tensor2: Second tensor to compare + atol: Absolute tolerance + rtol: Relative tolerance + max_errors: Maximum number of errors to print (default: 10) + """ + # Handle lists/tuples of tensors + if isinstance(tensor1, (list, tuple)) and isinstance(tensor2, (list, tuple)): + if len(tensor1) != len(tensor2): + print(f"Tensor count mismatch: {len(tensor1)} vs {len(tensor2)}") + return + + for i, (t1, t2) in enumerate(zip(tensor1, tensor2)): + print(f"\n=== Tensor {i} comparison ===") + print_tensor_comparison_errors(t1, t2, atol, rtol, max_errors) + return + + # Handle single tensor comparison + if not isinstance(tensor1, torch.Tensor) or not isinstance(tensor2, torch.Tensor): + print("Error: Both inputs must be torch.Tensor objects") + return + + if tensor1.shape != tensor2.shape: + print(f"Shape mismatch: {tensor1.shape} vs {tensor2.shape}") + return + + # Calculate absolute and relative errors + abs_diff = torch.abs(tensor1 - tensor2) + rel_diff = abs_diff / ( + torch.abs(tensor2) + 1e-8 + ) # Add small epsilon to avoid division by zero + + # Find locations where tolerance is exceeded + tolerance_mask = (abs_diff > atol) & (rel_diff > rtol) + + if not tolerance_mask.any(): + print("All values are within tolerance") + return + + # Get indices where tolerance is exceeded + error_indices = torch.nonzero(tolerance_mask, as_tuple=False) + total_errors = error_indices.shape[0] + + print(f"Found {total_errors} values exceeding tolerance (atol={atol}, rtol={rtol})") + print(f"Showing first {min(max_errors, total_errors)} errors:") + print("Index -> tensor1_value, tensor2_value, abs_error, rel_error") + + # Print first max_errors locations + for i in range(min(max_errors, total_errors)): + idx = tuple(error_indices[i].tolist()) + val1 = tensor1[idx].item() + val2 = tensor2[idx].item() + abs_err = abs_diff[idx].item() + rel_err = rel_diff[idx].item() + + print( + f"{idx} -> {val1:.6f}, {val2:.6f}, abs_err={abs_err:.6f}, rel_err={rel_err:.6f}" + ) + + +def check_outputs_equal( + model_output, ref_output, atol=1e-03, rtol=1e-03, first_output_only=False +): + """ + Helper function that checks if model output and reference output are equal with some tolerance. + Returns True if equal, False otherwise. + """ + # Convert OrderedDict to list if needed + if isinstance(ref_output, OrderedDict): + ref_output = list(ref_output.values()) + + # Compare the result from executor and eager mode directly + if isinstance(ref_output, tuple) or isinstance(ref_output, list): + # Multiple outputs executor always returns tuple, even if there is one output + if len(ref_output) != len(model_output): + print_tensor_comparison_errors(model_output, ref_output, atol, rtol) + return False + if first_output_only: + result = torch.allclose( + model_output[0], ref_output[0], atol=atol, rtol=rtol + ) + if not result: + print_tensor_comparison_errors( + model_output[0], ref_output[0], atol, rtol + ) + return result + else: + for i in range(len(ref_output)): + if not torch.allclose( + model_output[i], ref_output[i], atol=atol, rtol=rtol + ): + print(f"\n=== Output {i} comparison failed ===") + print_tensor_comparison_errors( + model_output[i], ref_output[i], atol, rtol + ) + return False + return True + else: + # If one output, eager returns tensor while executor tuple of size 1 + result = torch.allclose(model_output[0], ref_output, atol=atol, rtol=rtol) + if not result: + print_tensor_comparison_errors(model_output[0], ref_output, atol, rtol) + return result + + +def run_and_check_output( + reference_model: torch.nn.Module, + executorch_program: ExecutorchProgramManager, + sample_inputs: Tuple[torch.Tensor], + atol=1e-03, + rtol=1e-01, + first_output_only=False, +) -> bool: + """ + Utility function that accepts an already lowered ExecuTorch program, executes it with + the provided sample input, and checks the output for correctness. + + Args: + executorch_program: Already lowered ExecutorchProgramManager + sample_inputs: Sample inputs to run the program with + reference_model: Reference model to generate reference outputs for comparison + atol: Absolute tolerance for output comparison + rtol: Relative tolerance for output comparison + first_output_only: Whether to compare only the first output + + Returns: + bool: True if outputs match within tolerance, False otherwise + """ + # Load the ExecutorTorch program + executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) + + # Flatten inputs for execution + inputs_flattened, _ = tree_flatten(sample_inputs) + + # Run the ExecutorTorch program + model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) + + # Generate reference outputs using the reference model + ref_output = reference_model(*sample_inputs) + + # Check if outputs are equal + return check_outputs_equal( + model_output, + ref_output, + atol=atol, + rtol=rtol, + first_output_only=first_output_only, + ) + + +def make_copy_of_inputs(sample_inputs: Tuple[Any]) -> Tuple[Any]: + sample_inputs_copy = [] + for input_val in sample_inputs: + if isinstance(input_val, torch.Tensor): + sample_inputs_copy.append(input_val.clone()) + else: + sample_inputs_copy.append(deepcopy(input_val)) + return tuple(sample_inputs_copy) + + +def lower_module_and_test_output( + model: torch.nn.Module, + sample_inputs: Tuple[torch.Tensor], + atol=1e-03, + rtol=1e-01, + dynamic_shapes=None, + test_inputs=None, + first_output_only=False, + operator_blocklist=None, + operator_allowlist=None, + nn_module_allowlist=None, + nn_module_blocklist=None, + xnnpack=False, +) -> bool: + """ + Helper testing function that takes a torch.nn.Module and lowers it to Vulkan with + the given sample inputs. It then runs the lowered module and compares its + outputs with the outputs of the eager module. + + Returns: + bool: True if all comparisons pass, False otherwise. + """ + # Export model to Vulkan using the helper function + if xnnpack: + executorch_program = export_model_to_xnnpack( + model, + make_copy_of_inputs(sample_inputs), + dynamic_shapes, + operator_blocklist, + operator_allowlist, + nn_module_blocklist, + nn_module_allowlist, + ) + else: + executorch_program = export_model_to_vulkan( + model, + make_copy_of_inputs(sample_inputs), + dynamic_shapes, + operator_blocklist=operator_blocklist, + operator_allowlist=operator_allowlist, + nn_module_blocklist=nn_module_blocklist, + nn_module_allowlist=nn_module_allowlist, + ) + + executorch_module = _load_for_executorch_from_buffer(executorch_program.buffer) + + inputs_flattened, _ = tree_flatten(sample_inputs) + + model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) + ref_output = model(*make_copy_of_inputs(sample_inputs)) + + if not check_outputs_equal( + model_output, + ref_output, + atol=atol, + rtol=rtol, + first_output_only=first_output_only, + ): + return False + + if test_inputs is not None: + for test_input in test_inputs: + test_inputs_flattened, _ = tree_flatten(test_input) + model_output = executorch_module.run_method( + "forward", tuple(test_inputs_flattened) + ) + ref_output = model(*test_input) + + if not check_outputs_equal( + model_output, + ref_output, + atol=atol, + rtol=rtol, + first_output_only=first_output_only, + ): + return False + + return True + + +def save_bundled_program( + model: torch.nn.Module, + sample_inputs: Tuple[torch.Tensor], + output_path: str, + method_name: str = "forward", + et_program: Optional[ExecutorchProgramManager] = None, + dynamic_shapes=None, +) -> str: + """ + Export a bundled .pte file containing the model and test cases. + + Args: + model: The PyTorch model to export + sample_inputs: Sample inputs for the model + output_path: Path where the bundled .pte file should be saved (should end with .bpte) + method_name: Name of the method to test (default: "forward") + et_program: Optional pre-exported ExecutorchProgramManager. If None, will export to Vulkan + dynamic_shapes: Optional dynamic shapes for export + + Returns: + str: Path to the saved bundled program file + """ + # If no ExecutorchProgramManager provided, export to Vulkan + if et_program is None: + et_program = export_model_to_vulkan(model, sample_inputs, dynamic_shapes) + + # Generate expected outputs by running the model + expected_outputs = [getattr(model, method_name)(*sample_inputs)] + + # Flatten sample inputs to match expected format + inputs_flattened, _ = tree_flatten(sample_inputs) + + # Create test suite with the sample inputs and expected outputs + test_suites = [ + MethodTestSuite( + method_name=method_name, + test_cases=[ + MethodTestCase( + inputs=inputs_flattened, + expected_outputs=expected_outputs, + ) + ], + ) + ] + + # Create bundled program + bp = BundledProgram(et_program, test_suites) + + # Serialize to flatbuffer + bp_buffer = serialize_from_bundled_program_to_flatbuffer(bp) + + # Ensure output path has correct extension + if not output_path.endswith(".bpte"): + output_path = output_path + ".bpte" + + # Write to file + with open(output_path, "wb") as file: + file.write(bp_buffer) + return output_path + + +def save_executorch_program( + executorch_program: ExecutorchProgramManager, + output_path: str, +) -> str: + """ + Save an ExecutorchProgramManager as a .pte file. + + Args: + executorch_program: The ExecutorchProgramManager to save + output_path: Path where the .pte file should be saved (should end with .pte) + + Returns: + str: Path to the saved .pte file + """ + # Ensure output path has correct extension + if not output_path.endswith(".pte"): + output_path = output_path + ".pte" + + # Write to file + with open(output_path, "wb") as file: + executorch_program.write_to_file(file) + + return output_path + + +def print_occurrences(edge_program, operator_list: List): + """ + Print the input/output information for all occurrences of specified operators in the edge program. + + Args: + edge_program: The edge program created by to_edge_transform_and_lower + operator_list: List of operators to search for in the graph + """ + logger = logging.getLogger("") + logger.setLevel(logging.INFO) + + logger.info( + f"Searching for occurrences of {len(operator_list)} operators in the graph..." + ) + + occurrence_count = 0 + + for node in edge_program.exported_program().graph.nodes: + if utils.is_torch_op_node(node): + target = node.target + # Handle auto_functionalized nodes + if node.target == torch.ops.higher_order.auto_functionalized: + first_arg = node.args[0] + if hasattr(first_arg, "name"): + target = first_arg.name() + elif hasattr(first_arg, "__name__"): + target = first_arg.__name__ + + # Check if this operator is in our list + if target in operator_list: + occurrence_count += 1 + logger.info(f"Occurrence {occurrence_count}: {node.format_node()}") + + # Get the node I/O string using the utils function + try: + io_str = utils.node_io_str(node) + logger.info(f" {io_str}") + except Exception as e: + logger.info(f" Error getting I/O string: {e}") + + if occurrence_count == 0: + logger.info("No occurrences of the specified operators found in the graph.") + else: + logger.info( + f"Found {occurrence_count} total occurrences of the specified operators." + ) + + +def op_ablation_test( # noqa: C901 + model: torch.nn.Module, + sample_inputs: Tuple[torch.Tensor], + atol=1e-03, + rtol=1e-01, + dynamic_shapes=None, + test_inputs=None, + first_output_only=False, +) -> dict: + """ + Fast binary search utility function to determine which operators work correctly when delegated to Vulkan. + + This function uses a binary search approach to efficiently find bad operators: + 1. Split operators into two halves (least frequent first, most frequent second) + 2. Test each half to see if it produces correct output + 3. Add good halves to known_good_ops and recursively search bad halves + 4. Continue until all operators are classified + + Args: + model: The PyTorch model to test + sample_inputs: Sample inputs for the model + atol: Absolute tolerance for output comparison + rtol: Relative tolerance for output comparison + dynamic_shapes: Optional dynamic shapes for export + test_inputs: Optional additional test inputs + first_output_only: Whether to compare only the first output + + Returns: + dict: Dictionary with keys: + - 'good_operators': List of operators that work correctly + - 'bad_operators': List of operators that cause failures + - 'operator_frequencies': Dictionary mapping operators to their occurrence count + - 'all_operators': List of all unique operators found in the graph + - 'test_count': Number of tests performed + """ + logger = logging.getLogger("") + logger.setLevel(logging.INFO) + + logger.info("Starting fast binary search operator ablation test...") + + # Step 1: Export model to get edge_program and extract operators + export_training_graph = export(model, sample_inputs, strict=True).module() + program = export( + export_training_graph, + sample_inputs, + dynamic_shapes=dynamic_shapes, + strict=True, + ) + edge_program = to_edge_transform_and_lower( + program, + partitioner=[], # No partitioner to get the full graph + transform_passes=None, + compile_config=None, + ) + + # Step 2: Scan edge_program.graph_module to obtain unique operators and their frequencies + operator_frequencies = {} + for node in edge_program.exported_program().graph.nodes: + if utils.is_torch_op_node(node): + target = node.target + # Handle auto_functionalized nodes + if node.target == torch.ops.higher_order.auto_functionalized: + first_arg = node.args[0] + if hasattr(first_arg, "name"): + target = first_arg.name() + elif hasattr(first_arg, "__name__"): + target = first_arg.__name__ + + if target in operator_frequencies: + operator_frequencies[target] += 1 + else: + operator_frequencies[target] = 1 + + all_operators = list(operator_frequencies.keys()) + logger.info(f"Found {len(all_operators)} unique operators in the graph") + + # Sort operators by frequency (most frequent first for binary search) + operators_by_frequency = sorted( + all_operators, key=lambda op: operator_frequencies[op], reverse=True + ) + + logger.info("Operator frequencies (sorted by occurrence, most frequent first):") + for op in operators_by_frequency: + logger.info(f" {op.name()}: {operator_frequencies[op]} occurrences") + + # Global test counter + test_count = 0 + + def test_operator_set(ops_to_test: List, known_good_ops: List) -> bool: + """Test if a set of operators works correctly when combined with known good operators.""" + nonlocal test_count + test_count += 1 + + test_allowlist = known_good_ops + ops_to_test + logger.info( + f"Test {test_count}: Testing {len(ops_to_test)} operators with {len(known_good_ops)} known good" + ) + + try: + success = lower_module_and_test_output( + model=model, + sample_inputs=sample_inputs, + atol=atol, + rtol=rtol, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + first_output_only=first_output_only, + operator_allowlist=test_allowlist, + ) + logger.info(f" {'✓ PASS' if success else '✗ FAIL'}") + + # Log known good ops + logger.info(" Known good:") + for op in known_good_ops: + logger.info(f" * {op.name()}") + + # Log tested ops + logger.info(" Tested ops:") + for op in ops_to_test: + logger.info(f" * {op.name()}") + + return success + except Exception as e: + logger.info(f" ! Error: {e}") + return False + + def find_bad_operators( + ops_to_test: List, known_good_ops: List + ) -> Tuple[List, List]: + """ + Recursively find bad operators using binary search. + + Returns: + Tuple of (good_operators, bad_operators) from ops_to_test + """ + if not ops_to_test: + return [], [] + + if len(ops_to_test) == 1: + # Base case: single operator + op = ops_to_test[0] + if test_operator_set([op], known_good_ops): + logger.info(f" Single operator {op.name()} is GOOD") + return [op], [] + else: + logger.info(f" Single operator {op.name()} is BAD") + return [], [op] + + # Split ops_to_test into two halves + mid = len(ops_to_test) // 2 + first_half = ops_to_test[:mid] # Least frequent operators + second_half = ops_to_test[mid:] # Most frequent operators + + logger.info( + f"Splitting {len(ops_to_test)} operators: {len(first_half)} + {len(second_half)}" + ) + + # Log known good ops + logger.info(" Known good:") + for op in known_good_ops: + logger.info(f" * {op.name()}") + + # Log first half ops + logger.info(" First half ops:") + for op in first_half: + logger.info(f" * {op.name()}") + + # Log second half ops + logger.info(" Second half ops:") + for op in second_half: + logger.info(f" * {op.name()}") + + good_ops = [] + bad_ops = [] + + first_half_good = test_operator_set(first_half, known_good_ops) + if first_half_good: + logger.info( + f"First half ({len(first_half)} ops) is good - adding to known good" + ) + good_ops.extend(first_half) + known_good_ops.extend(first_half) + + second_half_good = test_operator_set(second_half, known_good_ops) + if second_half_good: + logger.info( + f"Second half ({len(second_half)} ops) is good - adding to known good" + ) + good_ops.extend(second_half) + + if not first_half_good: + logger.info(f"First half ({len(first_half)} ops) is bad - recursing") + sub_good, sub_bad = find_bad_operators(first_half, known_good_ops) + good_ops.extend(sub_good) + bad_ops.extend(sub_bad) + known_good_ops.extend(sub_good) + if not second_half_good: + logger.info(f"Second half ({len(second_half)} ops) is bad - recursing") + sub_good, sub_bad = find_bad_operators(second_half, known_good_ops) + good_ops.extend(sub_good) + bad_ops.extend(sub_bad) + + return good_ops, bad_ops + + # Start the binary search + logger.info( + f"\n=== Starting binary search on {len(operators_by_frequency)} operators ===" + ) + good_operators, bad_operators = find_bad_operators(operators_by_frequency, []) + + # Summary of results + logger.info(f"\n=== Binary search complete after {test_count} tests ===") + logger.info(f"Good operators ({len(good_operators)}):") + for op in good_operators: + logger.info(f" ✓ {op.name()} (frequency: {operator_frequencies[op]})") + + logger.info(f"Bad operators ({len(bad_operators)}):") + for op in bad_operators: + logger.info(f" ✗ {op.name()} (frequency: {operator_frequencies[op]})") + + print_occurrences(edge_program, bad_operators) + + efficiency_gain = len(all_operators) - test_count + logger.info( + f"Efficiency: {test_count} tests instead of {len(all_operators)} (saved {efficiency_gain} tests)" + ) + + return { + "good_operators": good_operators, + "bad_operators": bad_operators, + "operator_frequencies": operator_frequencies, + "all_operators": all_operators, + "test_count": test_count, + } diff --git a/backends/vulkan/test/utils/test_utils.cpp b/backends/vulkan/test/utils/test_utils.cpp index faa0e7d0c47..f00dfa20976 100644 --- a/backends/vulkan/test/utils/test_utils.cpp +++ b/backends/vulkan/test/utils/test_utils.cpp @@ -14,9 +14,86 @@ #include #include +#include using namespace vkcompute; +bool is_bitw8(vkapi::ScalarType dtype) { + return dtype == vkapi::kByte || dtype == vkapi::kChar || + dtype == vkapi::kQInt8 || dtype == vkapi::kQUInt8; +} + +vkapi::ShaderInfo get_nchw_to_tensor_shader( + const api::vTensor& v_dst, + bool int8_buffer_enabled, + bool push_constant_variant) { + std::string kernel_name; + kernel_name.reserve(kShaderNameReserve); + + if (is_bitw8(v_dst.dtype()) && v_dst.storage_type() != utils::kBuffer && + !int8_buffer_enabled) { + kernel_name = "nchw_to_bitw8_image_nobitw8buffer"; + if (!push_constant_variant) { + kernel_name += "_no_pc"; + } + add_storage_type_suffix(kernel_name, v_dst.storage_type()); + add_dtype_suffix(kernel_name, v_dst.dtype()); + return VK_KERNEL_FROM_STR(kernel_name); + } + + if (v_dst.storage_type() == utils::kBuffer) { + kernel_name = "nchw_to_buffer"; + add_dtype_suffix(kernel_name, v_dst.dtype()); + add_dtype_suffix(kernel_name, v_dst.dtype()); + return VK_KERNEL_FROM_STR(kernel_name); + } + + kernel_name = "nchw_to_image"; + if (!push_constant_variant) { + kernel_name += "_no_pc"; + } + add_storage_type_suffix(kernel_name, v_dst.storage_type()); + add_dtype_suffix(kernel_name, v_dst.dtype()); + add_dtype_suffix(kernel_name, v_dst.dtype()); + + return VK_KERNEL_FROM_STR(kernel_name); +} + +vkapi::ShaderInfo get_tensor_to_nchw_shader( + const api::vTensor& v_src, + bool int8_buffer_enabled, + bool push_constant_variant) { + std::string kernel_name; + kernel_name.reserve(kShaderNameReserve); + + if (is_bitw8(v_src.dtype()) && v_src.storage_type() != utils::kBuffer && + !int8_buffer_enabled) { + kernel_name = "bitw8_image_to_nchw_nobitw8buffer"; + if (!push_constant_variant) { + kernel_name += "_no_pc"; + } + add_storage_type_suffix(kernel_name, v_src.storage_type()); + add_dtype_suffix(kernel_name, v_src.dtype()); + return VK_KERNEL_FROM_STR(kernel_name); + } + + if (v_src.storage_type() == utils::kBuffer) { + kernel_name = "buffer_to_nchw"; + add_dtype_suffix(kernel_name, v_src.dtype()); + add_dtype_suffix(kernel_name, v_src.dtype()); + return VK_KERNEL_FROM_STR(kernel_name); + } + + kernel_name = "image_to_nchw"; + if (!push_constant_variant) { + kernel_name += "_no_pc"; + } + add_storage_type_suffix(kernel_name, v_src.storage_type()); + add_dtype_suffix(kernel_name, v_src.dtype()); + add_dtype_suffix(kernel_name, v_src.dtype()); + + return VK_KERNEL_FROM_STR(kernel_name); +} // // Operator Recording Functions // @@ -41,9 +118,7 @@ void record_nchw_to_buffer_op( vkapi::PipelineStage::COMPUTE, vkapi::MemoryAccessType::WRITE), src_buffer, - v_dst.sizes_ubo(), - v_dst.strides_ubo(), - v_dst.numel_ubo()); + v_dst.buffer_meta_ubo()); } void record_buffer_to_nchw_op( @@ -61,9 +136,7 @@ void record_buffer_to_nchw_op( 0, dst_buffer, v_src.buffer(pipeline_barrier, vkapi::PipelineStage::COMPUTE), - v_src.sizes_ubo(), - v_src.strides_ubo(), - v_src.numel_ubo()); + v_src.buffer_meta_ubo()); } void record_nchw_to_image_op( @@ -121,8 +194,8 @@ void record_bitw8_image_to_nchw_nobitw8buffer_op( utils::uvec3 global_wg_size = {buffer_len, 1, 1}; std::string kernel_name = "bitw8_image_to_nchw_nobitw8buffer_no_pc"; - add_storage_type_suffix(kernel_name, v_src); - add_dtype_suffix(kernel_name, v_src); + add_storage_type_suffix(kernel_name, v_src.storage_type()); + add_dtype_suffix(kernel_name, v_src.dtype()); context->submit_compute_job( VK_KERNEL_FROM_STR(kernel_name), @@ -145,7 +218,7 @@ void record_binary_op( api::vTensor& v_in2, api::vTensor& v_dst) { std::string kernel_name = "binary_" + op_name + "_nobroadcast__test"; - add_dtype_suffix(kernel_name, v_dst); + add_dtype_suffix(kernel_name, v_dst.dtype()); vkapi::PipelineBarrier pipeline_barrier{}; vkapi::SpecVarList specialization_constants = {}; @@ -236,7 +309,7 @@ void record_scalar_add_buffer( vkapi::PipelineBarrier pipeline_barrier{}; vkapi::SpecVarList specialization_constants = {SV(offset)}; std::string kernel = "scalar_add_buffer"; - add_dtype_suffix(kernel, v_ten); + add_dtype_suffix(kernel, v_ten.dtype()); api::context()->submit_compute_job( VK_KERNEL_FROM_STR(kernel), pipeline_barrier, @@ -398,10 +471,9 @@ void fill_vtensor( const IOValueRef idx, float val, bool iota) { - vTensorPtr t = graph.get_tensor(idx.value); - std::vector data(t->numel()); - if (t->storage_type() != utils::kBuffer) { - data.resize(t->staging_buffer_numel()); + std::vector data(graph.numel_of(idx.value)); + if (graph.storage_type_of(idx.value) != utils::kBuffer) { + data.resize(graph.staging_buffer_numel_of(idx.value)); } if (iota) { std::iota(data.begin(), data.end(), val); @@ -489,13 +561,12 @@ void execute_graph_and_check_output( for (size_t i = 0; i < graph.outputs().size(); ++i) { IOValueRef out_ioval = graph.outputs().at(i); - vTensorPtr t_out = graph.get_tensor(out_ioval.value); - - std::vector output_data(t_out->staging_buffer_numel()); + std::vector output_data( + graph.staging_buffer_numel_of(out_ioval.value)); graph.copy_from_staging( out_ioval.staging, output_data.data(), output_data.size()); - for (size_t j = 0; j < t_out->numel(); ++j) { + for (size_t j = 0; j < graph.numel_of(out_ioval.value); ++j) { CHECK_VALUE(output_data, j, expected_outputs.at(i)); } } diff --git a/backends/vulkan/test/utils/test_utils.h b/backends/vulkan/test/utils/test_utils.h index 0f0d2647792..1fd40b6f815 100644 --- a/backends/vulkan/test/utils/test_utils.h +++ b/backends/vulkan/test/utils/test_utils.h @@ -214,9 +214,7 @@ inline int64_t get_buf_idx( vkcompute::ComputeGraph& graph, vkcompute::IOValueRef ref, const std::vector& tensor_coor) { - vkcompute::vTensorPtr vten_ptr = graph.get_tensor(ref.value); - - const std::vector& sizes = vten_ptr->sizes(); + const std::vector& sizes = graph.sizes_of(ref.value); int64_t c = vkcompute::dim_at(sizes); int64_t h = vkcompute::dim_at(sizes); diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 82df7e7d96f..a193d02da88 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -114,7 +114,7 @@ TEST_F(VulkanComputeAPITest, print_shader_executable_properties) { std::vector get_reference_strides( const std::vector& sizes, const utils::GPUMemoryLayout layout, - const bool unsqueezed = false) { + const bool flip_unsqueezed = false) { int64_t C = utils::val_at(-3, sizes); int64_t H = utils::val_at(-2, sizes); int64_t W = utils::val_at(-1, sizes); @@ -125,18 +125,20 @@ std::vector get_reference_strides( case utils::kWidthPacked: switch (sizes.size()) { case 1: - if (unsqueezed) - return {numel, numel, numel, 1}; + if (flip_unsqueezed) + return {1, numel, numel, numel}; return {1}; case 2: - if (unsqueezed) - return {numel, numel, W, 1}; + if (flip_unsqueezed) + return {1, W, numel, numel}; return {W, 1}; case 3: - if (unsqueezed) - return {numel, H * W, W, 1}; + if (flip_unsqueezed) + return {1, W, H * W, numel}; return {H * W, W, 1}; case 4: + if (flip_unsqueezed) + return {1, W, H * W, C * H * W}; return {C * H * W, H * W, W, 1}; default: return {}; @@ -145,18 +147,21 @@ std::vector get_reference_strides( case utils::kHeightPacked: switch (sizes.size()) { case 1: - if (unsqueezed) - return {numel, numel, numel, 1}; + if (flip_unsqueezed) + return {1, numel, numel, numel}; return {1}; case 2: - if (unsqueezed) - return {numel, numel, 1, H}; + if (flip_unsqueezed) + return {H, 1, numel, numel}; + return {1, H}; return {1, H}; case 3: - if (unsqueezed) - return {numel, H * W, 1, H}; + if (flip_unsqueezed) + return {H, 1, H * W, numel}; return {W * H, 1, H}; case 4: + if (flip_unsqueezed) + return {H, 1, W * H, C * W * H}; return {C * W * H, W * H, 1, H}; default: return {}; @@ -164,18 +169,20 @@ std::vector get_reference_strides( case utils::kChannelsPacked: switch (sizes.size()) { case 1: - if (unsqueezed) - return {numel, numel, numel, 1}; + if (flip_unsqueezed) + return {1, numel, numel, numel}; return {1}; case 2: - if (unsqueezed) - return {numel, numel, W, 1}; + if (flip_unsqueezed) + return {1, W, numel, numel}; return {W, 1}; case 3: - if (unsqueezed) - return {numel, 1, W * C, C}; + if (flip_unsqueezed) + return {C, W * C, 1, numel}; return {1, W * C, C}; case 4: + if (flip_unsqueezed) + return {C, W * C, 1, H * W * C}; return {H * W * C, 1, W * C, C}; default: return {}; @@ -184,6 +191,41 @@ std::vector get_reference_strides( return {}; } +/* + * Applies the following transformations to a tensor's dim_order vector: + * 1. Reverse the order of elements so that the fastest moving dimensions are + * first. + * 2. Convert NCHW dimension indices to WHCN indices, so that 0 represents the + * width dimension, 1 represents the height dimension, and 2 represents the + * channels dimension. + * 3. Unsqueeze the dim_order vector to the next multiple of 4. + */ +std::vector create_whcn_dim_order( + const std::vector& dim_order) { + size_t ndim = dim_order.size(); + std::vector whcn_order(ndim); + + // Convert from NCHW to WHCN index, and flip the dim order so that the fastest + // moving dimension is first. + // example: { 1, 2, 0} -> { 2, 0, 1} + // {height, width, channels} -> {channels, width, height} + for (size_t whcn_i = 0, nchw_i = (ndim - 1); whcn_i < ndim; + ++whcn_i, --nchw_i) { + whcn_order.at(whcn_i) = ndim - 1 - dim_order.at(nchw_i); + } + + // Unsqueeze to the next multiple of 4 + size_t ndim_up4 = utils::align_up_4(ndim); + whcn_order.resize(ndim_up4); + + // Append unsqueezed dimensions + for (size_t i = ndim; i < ndim_up4; ++i) { + whcn_order.at(i) = i; + } + + return whcn_order; +} + TEST_F(VulkanComputeAPITest, empty_init_shader_info_test) { vkapi::ShaderInfo empty_shader_info; EXPECT_FALSE(empty_shader_info); @@ -191,6 +233,20 @@ TEST_F(VulkanComputeAPITest, empty_init_shader_info_test) { EXPECT_TRUE(empty_shader_info.src_code.size == 0u); } +bool compare_vectors( + const std::vector& v32, + const std::vector& v64) { + if (v32.size() != v64.size()) { + return false; + } + for (size_t i = 0; i < v32.size(); ++i) { + if (static_cast(v32[i]) != v64[i]) { + return false; + } + } + return true; +} + TEST_F(VulkanComputeAPITest, calculate_dim_order_test) { // ndim, GPUMemoryLayout, expected dim order pairs std::vector>> test_cases = { @@ -238,17 +294,27 @@ TEST_F(VulkanComputeAPITest, calculate_tensor_strides_test) { std::vector dim_order = calculate_dim_order(sizes.size(), packed_dim); std::vector strides = calculate_strides(sizes, dim_order); + int64_t numel = utils::multiply_integers(sizes); + std::vector ref_strides = get_reference_strides(sizes, layout); ASSERT_TRUE(strides == ref_strides); - int64_t numel = utils::multiply_integers(sizes); std::vector unsqueezed_strides = - unsqueeze_strides(strides, numel); + flip_and_unsqueeze(strides, kTensorStrides, numel); + std::vector ref_unsqueezed_strides = get_reference_strides(sizes, layout, true); ASSERT_TRUE(unsqueezed_strides == ref_unsqueezed_strides); + std::vector whcn_dim_order = + flip_and_unsqueeze(dim_order, kTensorDimOrder, numel); + + std::vector ref_whcn_dim_order = + create_whcn_dim_order(dim_order); + + ASSERT_TRUE(whcn_dim_order == ref_whcn_dim_order); + // Create new vTensor and check that the strides are correct vTensor new_v_tensor( context(), @@ -498,7 +564,7 @@ TEST_F(VulkanComputeAPITest, update_params_between_submit) { vTensor a = CREATE_FLOAT_TEXTURE(sizes, /*allocate_memory = */ true); std::string kernel_name("fill_texture__test"); - add_dtype_suffix(kernel_name, a); + add_dtype_suffix(kernel_name, a.dtype()); struct Params final { utils::ivec3 size; @@ -1014,9 +1080,8 @@ TEST_F(VulkanComputeAPITest, texture_virtual_resize) { // Compute Graph Tests // -#define EXTRACT_TENSOR(name) \ - std::vector data_##name( \ - graph.get_tensor(name.value)->staging_buffer_numel()); \ +#define EXTRACT_TENSOR(name) \ + std::vector data_##name(graph.staging_buffer_numel_of(name.value)); \ graph.copy_from_staging(name.staging, data_##name.data(), data_##name.size()); // The purpose of this test is simply to track the size of various classes over @@ -1037,12 +1102,12 @@ TEST_F(VulkanComputeAPITest, print_object_sizes) { // Current known size on 64 bit system: 1040 B EXPECT_TRUE(sizeof(vTensor) < 1200); - // Current known size on 64 bit system: 48 B - EXPECT_TRUE(sizeof(Value) < 56); + // Current known size on 64 bit system: 80 B + EXPECT_TRUE(sizeof(Value) < 100); // Current known size on 64 bit system: 120 B EXPECT_TRUE(sizeof(StagingBuffer) < 500); - // Current known size on 64 bit system: 384 B - EXPECT_TRUE(sizeof(ComputeGraph) < 500); + // Current known size on 64 bit system: 608 B + EXPECT_TRUE(sizeof(ComputeGraph) < 700); // Current known size on 64 bit system: 248 B EXPECT_TRUE(sizeof(DispatchNode) < 500); } @@ -1177,6 +1242,7 @@ TEST(VulkanComputeGraphTest, test_zero_dim_tensor) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); // Run graph @@ -1193,7 +1259,7 @@ TEST(VulkanComputeGraphTest, test_zero_dim_tensor) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); ++i) { CHECK_VALUE(data_out, i, val_c); } } @@ -1219,6 +1285,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_buffer) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); // Run graph @@ -1233,7 +1300,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_buffer) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); ++i) { CHECK_VALUE(data_out, i, expected_val); } } @@ -1304,6 +1371,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); // Run graph @@ -1320,7 +1388,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); ++i) { CHECK_VALUE(data_out, i, val_c); } } @@ -1362,6 +1430,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); // Run graph @@ -1382,7 +1451,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); i++) { CHECK_VALUE(data_out, i, val_out); } } @@ -1445,7 +1514,7 @@ TEST(VulkanComputeGraphTest, test_simple_prepacked_graph) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); ++i) { CHECK_VALUE(data_out, i, val_out); } @@ -1520,6 +1589,7 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects_with_resize) { EXPECT_EQ(get_vma_allocation_count(), expected_vma_allocation_count); graph.prepare(); + graph.prepack(); // +3: shared memory allocations for tensors expected_vma_allocation_count += 3; @@ -1531,9 +1601,9 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects_with_resize) { {8, 44, 34}, {4, 13, 56}, {8, 12, 64}, {12, 55, 33}, {4, 54, 10}}; for (auto& new_sizes : new_sizes_list) { - graph.get_tensor(a.value)->virtual_resize(new_sizes); - graph.get_tensor(b.value)->virtual_resize(new_sizes); - graph.get_tensor(d.value)->virtual_resize(new_sizes); + graph.virtual_resize(a.value, new_sizes); + graph.virtual_resize(b.value, new_sizes); + graph.virtual_resize(d.value, new_sizes); graph.propagate_resize(); float val_a = new_sizes[1] + 4.0f; @@ -1551,7 +1621,7 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects_with_resize) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); i++) { + for (size_t i = 0; i < graph.numel_of(out.value); i++) { CHECK_VALUE(data_out, i, val_out); } } @@ -1566,7 +1636,7 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects_with_resize) { graph.propagate_resize(); // Check output shape - EXPECT_TRUE(graph.get_tensor(out.value)->sizes() == new_sizes); + EXPECT_TRUE(graph.sizes_of(out.value) == new_sizes); float val_a = new_sizes[1] + 6.0f; float val_b = new_sizes[2] + 2.5f; @@ -1583,7 +1653,7 @@ TEST(VulkanComputeGraphTest, test_simple_shared_objects_with_resize) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); i++) { + for (size_t i = 0; i < graph.numel_of(out.value); i++) { CHECK_VALUE(data_out, i, val_out); } } @@ -1660,6 +1730,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_tmp_tensors) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); // Run graph @@ -1681,7 +1752,7 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_tmp_tensors) { EXTRACT_TENSOR(out); // Sanity check that the values are correct - for (size_t i = 0; i < graph.get_tensor(out.value)->numel(); ++i) { + for (size_t i = 0; i < graph.numel_of(out.value); ++i) { CHECK_VALUE(data_out, i, val_out); } } @@ -1726,6 +1797,7 @@ TEST(VulkanComputeGraphTest, test_large_graph) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); auto build_end_time = std::chrono::system_clock::now(); @@ -1767,7 +1839,7 @@ TEST(VulkanComputeGraphTest, test_large_graph) { auto inference_time = std::chrono::duration_cast( inference_end_time - inference_start_time); - for (int i = 0; i < graph.get_tensor(out.value)->numel(); i++) { + for (int i = 0; i < graph.numel_of(out.value); i++) { CHECK_VALUE(data_out, i, val_e); } @@ -1802,6 +1874,7 @@ void test_clone( out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, a, 0.0f, /*iota = */ true); @@ -1886,6 +1959,7 @@ TEST(VulkanComputeGraphTest, test_etvk_copy_offset_node) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, a, 0.0f, /*iota = */ true); @@ -1949,6 +2023,7 @@ TEST(VulkanComputeGraphTest, DISABLED_test_etvk_copy_channel_offset_node) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, a, 0.0f, true); @@ -2039,6 +2114,7 @@ TEST( out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); float a_value = 1.0f; float b_value = 2.0f; @@ -2151,6 +2227,7 @@ TEST(VulkanComputeGraphTest, test_etvk_copy_offset_int_node) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, a, 0, /*iota = */ true); @@ -2214,6 +2291,7 @@ TEST(VulkanComputeGraphTest, DISABLED_test_etvk_copy_channel_offset_int_node) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, a, 0.0f, true); @@ -2273,6 +2351,7 @@ TEST(VulkanComputeGraphTest, test_view_change_packing) { out.staging = graph.set_output_tensor(out.value); graph.prepare(); + graph.prepack(); fill_vtensor(graph, in, 0.0, true); @@ -2282,7 +2361,7 @@ TEST(VulkanComputeGraphTest, test_view_change_packing) { // The extracted data is a flattened nchw buffer. Hence, should expect the // all elements inside the out array to match the index. - for (int i = 0; i < graph.get_tensor(out.value)->numel(); i++) { + for (int i = 0; i < graph.numel_of(out.value); i++) { CHECK_VALUE(data_out, i, i); } } @@ -2317,7 +2396,7 @@ void run_from_gpu_test( vTensor vten = vTensor(context(), sizes, dtype, storage_type, memory_layout); std::string kernel_name("idx_fill_texture"); - add_dtype_suffix(kernel_name, vten); + add_dtype_suffix(kernel_name, vten.dtype()); int32_t offset = -50; @@ -2431,10 +2510,9 @@ void compute_graph_round_trip_test( ValueRef r_staging_out = graph.set_output_tensor(r_tensor); graph.prepare(); + graph.prepack(); - vTensorPtr tensor = graph.get_tensor(r_tensor); - - std::vector data_in(tensor->numel()); + std::vector data_in(graph.numel_of(r_tensor)); for (int i = 0; i < data_in.size(); i++) { data_in[i] = T(i * -1); } @@ -2442,7 +2520,7 @@ void compute_graph_round_trip_test( graph.execute(); - std::vector data_out(tensor->staging_buffer_numel()); + std::vector data_out(graph.staging_buffer_numel_of(r_tensor)); graph.copy_from_staging(r_staging_out, data_out.data(), data_out.size()); for (int i = 0; i < data_in.size(); i++) { @@ -2623,7 +2701,6 @@ void test_mm( B, M, K, N, dtype, storage_type, memory_layout, mat2_data, prepack); graph.prepare(); - graph.prepack(); for (int i = 1; i < 4; i++) { @@ -2703,7 +2780,6 @@ void test_mm_with_resize_reencode( B, M, K, N, dtype, storage_type, memory_layout, mat2_data, false); graph.prepare(); - graph.prepack(); for (int i = 1; i < 4; i++) { @@ -2740,94 +2816,6 @@ TEST(VulkanComputeGraphOpsTest, test_graph_resize_reencode) { utils::kWidthPacked); } -void test_max_pool2d( - const std::vector& in_size, - const int64_t base_val, - std::vector& kernel) { - GraphConfig config; - ComputeGraph graph(config); - - // Build graph - - std::vector out_size(in_size); - int h = in_size.size() - 2; - int w = in_size.size() - 1; - out_size[h] = in_size[h] - kernel[0] + 1; - out_size[w] = in_size[w] - kernel[1] + 1; - - IOValueRef in_ioval = graph.add_input_tensor( - in_size, vkapi::kFloat, utils::GPUMemoryLayout::TENSOR_CHANNELS_PACKED); - IOValueRef out_ioval; - out_ioval.value = graph.add_tensor( - out_size, vkapi::kFloat, utils::GPUMemoryLayout::TENSOR_CHANNELS_PACKED); - IOValueRef idx_ioval; - idx_ioval.value = graph.add_tensor( - out_size, vkapi::kInt, utils::GPUMemoryLayout::TENSOR_CHANNELS_PACKED); - ValueRef out = graph.add_value_list({out_ioval.value, idx_ioval.value}); - - std::vector kernel_copy(kernel); - VK_GET_OP_FN("aten.max_pool2d_with_indices.default") - (graph, - {in_ioval.value, - graph.add_scalar_list(std::move(kernel)), - graph.add_scalar_list({1, 1}), - graph.add_scalar_list({0, 0}), - graph.add_scalar_list({1, 1}), - graph.add_scalar(false), - out}); - - out_ioval.staging = graph.set_output_tensor(out_ioval.value); - idx_ioval.staging = graph.set_output_tensor(idx_ioval.value); - - graph.prepare(); - - graph.prepack(); - - // Run graph - - fill_vtensor(graph, graph.inputs().at(0), base_val, /*iota = */ true); - - vTensorPtr t_in = graph.get_tensor(in_ioval.value); - std::vector input_data(t_in->staging_buffer_numel()); - graph.copy_from_staging( - in_ioval.staging, input_data.data(), input_data.size()); - - graph.execute(); - - vTensorPtr t_out = graph.get_tensor(out_ioval.value); - std::vector output_data(t_out->staging_buffer_numel()); - graph.copy_from_staging( - out_ioval.staging, output_data.data(), output_data.size()); - vTensorPtr t_idx = graph.get_tensor(idx_ioval.value); - std::vector index_data(t_idx->staging_buffer_numel()); - graph.copy_from_staging( - idx_ioval.staging, index_data.data(), index_data.size()); - - // Check results - - int h_offset = kernel_copy[0] - 1; - int w_offset = kernel_copy[1] - 1; - int h_out = utils::val_at(-2, t_out->sizes()); - int w_out = utils::val_at(-1, t_out->sizes()); - int w_in = utils::val_at(-1, t_in->sizes()); - for (size_t i = 0; i < h_out; ++i) { - for (size_t j = 0; j < w_out; ++j) { - size_t idx_out = i * w_out + j; - size_t idx_in = (i + h_offset) * w_in + (j + w_offset); - CHECK_VALUE(index_data, idx_out, idx_in); - CHECK_VALUE(output_data, idx_out, input_data[idx_in]); - } - } -} - -TEST(VulkanComputeGraphOpsTest, max_pool2d_smoke_test) { - std::vector kernel = {2, 3}; - test_max_pool2d( - /*in_size = */ {1, 4, 6}, - /*base_val = */ 10.0f, - kernel); -} - void test_grid_priors( std::vector input_sizes, std::vector output_sizes, @@ -2861,20 +2849,19 @@ void test_grid_priors( graph.prepack(); - vTensorPtr t_in = graph.get_tensor(in.value); - vTensorPtr t_out = graph.get_tensor(out.value); // Resize input graph.propagate_resize(); // run graph graph.execute(); - std::vector output_data(t_out->staging_buffer_numel()); + std::vector output_data(graph.staging_buffer_numel_of(out.value)); graph.copy_from_staging(out.staging, output_data.data(), output_data.size()); // check results - int h_out = utils::val_at(-2, t_out->sizes()); - int w_out = utils::val_at(-1, t_out->sizes()); + std::vector out_sizes = graph.sizes_of(out.value); + int h_out = utils::val_at(-2, out_sizes); + int w_out = utils::val_at(-1, out_sizes); for (size_t i = 0; i < h_out; ++i) { for (size_t j = 0; j < w_out; ++j) { size_t idx_out = i * w_out + j; @@ -3151,7 +3138,7 @@ void resize_dynamic_dispatch_node( std::vector out_sizes = graph->sizes_of(mat1); out_sizes.at(out_sizes.size() - 2) = 1; - graph->get_tensor(out)->virtual_resize(out_sizes); + graph->virtual_resize(out, out_sizes); } void add_dynamic_dispatch_test_node( @@ -3214,7 +3201,6 @@ void test_dynamic_dispatch(int M, int N) { ComputeGraph graph = build_dynamic_dispatch_test_graph(M, N); graph.prepare(); - graph.prepack(); for (int i = 1; i < 4; i++) { diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index fa45063a4d3..96f200eecbc 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -18,17 +18,21 @@ format_target_name, ) +from executorch.exir.dialects.edge._ops import EdgeOpOverload + from executorch.exir.tensor import TensorSpec -from torch._export.utils import is_buffer, is_param +from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param -from torch._subclasses.fake_tensor import FakeTensor +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorConverter from torch.export import ExportedProgram from torch.export.exported_program import InputKind from torch.export.graph_signature import TensorArgument +TorchOpType = Union[EdgeOpOverload, torch._ops.OpOverload, str] + _DQ_OPS = { "dequantize_per_tensor.tensor", "dequantize_per_tensor.default", @@ -54,6 +58,18 @@ MaybeNodeList = Union[torch.fx.Node, List[torch.fx.Node], Tuple[torch.fx.Node]] +def is_torch_op_node(node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + + if isinstance(node.target, EdgeOpOverload): + return True + if isinstance(node.target, torch._ops.OpOverload): + return True + + return False + + def is_dequant_node(node: torch.fx.Node) -> bool: if node.op != "call_function": return False @@ -68,6 +84,13 @@ def is_quant_node(node: torch.fx.Node) -> bool: return node_name in _Q_OPS +def is_choose_qparams_node(node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + node_name = format_target_name(node.target.__name__) # pyre-ignore + return "choose_qparams" in node_name + + def is_dequant_per_channel_node(node: torch.fx.Node) -> bool: if node.op != "call_function": return False @@ -75,6 +98,13 @@ def is_dequant_per_channel_node(node: torch.fx.Node) -> bool: return node_name == "dequantize_per_channel.default" +def is_view_copy_node(node: torch.fx.Node) -> bool: + if node.op != "call_function": + return False + node_name = format_target_name(node.target.__name__) # pyre-ignore + return "view_copy" in node_name + + def is_linear_node(node: torch.fx.Node) -> bool: if node.op != "call_function": return False @@ -236,6 +266,70 @@ def get_primary_arg_idx(self, node: torch.fx.Node) -> Optional[int]: return primary_arg_idx +def node_comes_from_any_nn_module_in_set( + node, + nn_module_typenames: Set[str], +) -> bool: + if isinstance(node, (list, tuple)): + return all( + node_comes_from_any_nn_module_in_set(n, nn_module_typenames) for n in node + ) + + if not isinstance(node, torch.fx.Node): + return False + + nn_module_stack = node.meta.get("nn_module_stack", None) + if nn_module_stack is None: + return False + + for _, packed in nn_module_stack.items(): + _, typename = packed + for partial_name in nn_module_typenames: + if partial_name in typename: + return True + + return False + + +def get_tensor_name(exp_prog: ExportedProgram, node: torch.fx.Node) -> str: + if node is None: + return "" + if is_param(exp_prog, node): + return exp_prog.graph_signature.inputs_to_parameters[node.name] + elif is_buffer(exp_prog, node): + return exp_prog.graph_signature.inputs_to_buffers[node.name] + elif is_lifted_tensor_constant(exp_prog, node): + return exp_prog.graph_signature.inputs_to_lifted_tensor_constants[node.name] + else: + assert isinstance(node.target, str) + return node.target + + return "" + + +def find_dequant_user(node: torch.fx.Node) -> Optional[torch.fx.Node]: + """ + Search the direct users of the given node and return the first one that is a + dequantization op. Returns None if no dequantization op is found. + """ + for user in node.users: + if is_dequant_node(user): + return user + return None + + +def find_quant_user(node: torch.fx.Node) -> Optional[torch.fx.Node]: + """ + Search the direct users of the given node and return the first one that is a + quantization op. Returns None if no quantization op is found. + """ + for user in node.users: + if is_quant_node(user): + return user + + return None + + ## ## Memory Layout, Storage Type Determination ## @@ -585,9 +679,9 @@ def make_filtered_tensor_repset( if extents_are_valid(extents, texture_limits): valid_texture_layouts.add(memory_layout) - # High dimensional tensors are currently not supported + # High dimensional tensors require buffer storage if len(tensor_val.shape) > 4: - return NO_STORAGE + return TensorRepSet(tensor_repset.valid_buffer_layouts, set()) # Bool tensors are currently not supported if tensor_val.dtype == torch.bool: @@ -607,6 +701,7 @@ def make_filtered_tensor_repset( CHANNELS_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_CHANNELS_PACKED}) ANY_TEXTURE = TensorRepSet(set(), all_memory_layouts) +ANY_BUFFER = TensorRepSet(all_memory_layouts, set()) ANY_STORAGE = TensorRepSet(all_memory_layouts, all_memory_layouts) NO_STORAGE = TensorRepSet(set(), set()) @@ -1028,11 +1123,146 @@ def get_node_repr(node) -> Union[TensorRepr, TensorReprList]: return get_node_spec_attr(node, "etvk_node_repr", False) +## +## Graph Pattern Matching +## + + +def maybe_skip_q_dq_arg_chain( + arg: torch.fx.node.Argument, +) -> Tuple[Optional[torch.fx.Node], Optional[torch.fx.Node], Optional[torch.fx.Node]]: + """ + Check if the given node argument is part of a Quantize/Dequantize chain produced by + the quant workflow. If so, return the source tensor that is the input to the Q/DQ + chain and the quantize/dequantize nodes in the chain. Otherwise, return the argument + as is and None, None + """ + if not isinstance(arg, torch.fx.Node): + return None, None, None + + # If the arg is a view copy node, check if the original node is a dequant node + if is_dequant_node(arg) or ( + is_view_copy_node(arg) and is_dequant_node(arg.args[0]) # pyre-ignore[6] + ): + dequant_node = arg + if is_view_copy_node(arg): + dequant_node = arg.args[0] + + quant_node = dequant_node.args[0] # pyre-ignore[16] + assert isinstance(quant_node, torch.fx.Node) + source_arg = quant_node.args[0] + assert isinstance(source_arg, torch.fx.Node) + assert isinstance(dequant_node, torch.fx.Node) + return source_arg, quant_node, dequant_node + else: + return arg, None, None + + +def trace_args_until_placeholder( + node: torch.fx.node.Argument, max_search_depth: int = 4 +) -> Tuple[Optional[torch.fx.Node], List[torch.fx.Node]]: + """ + Trace through node.args[0] of a given initial node until a placeholder node is found + then return it and the list of nodes traversed. If no placeholder node is found, + returns None and an empty list. + """ + cur_node = node + search_depth = 0 + + if not isinstance(cur_node, torch.fx.Node): + return None, [] + + traversed = [cur_node] + while cur_node.op != "placeholder" and search_depth < max_search_depth: + # Break if cur_node has no args + if len(cur_node.args) == 0: + break + + cur_node = cur_node.args[0] + if not isinstance(cur_node, torch.fx.Node): + break + traversed.append(cur_node) + search_depth += 1 + + if not isinstance(cur_node, torch.fx.Node): + return None, [] + if cur_node.op != "placeholder": + return None, [] + + assert isinstance(cur_node, torch.fx.Node) + return cur_node, traversed + + +def is_in_4bit_range(tensor: torch.Tensor) -> bool: + """ + Check if the given tensor is in the range of 4-bit quantization and is of integer type. + """ + if tensor.dtype not in (torch.int8, torch.uint8): + return False + + return tensor.min().item() >= -8 and tensor.max().item() <= 7 + + +def is_in_8bit_range(tensor: torch.Tensor) -> bool: + """ + Check if the given tensor is in the range of 4-bit quantization and is of integer type. + """ + if tensor.dtype not in (torch.int8, torch.uint8): + return False + + return tensor.min().item() >= -128 and tensor.max().item() <= 127 + + ## ## Misc ## +def get_tensor_val_str(tensor_val: FakeTensor) -> str: + return f"{tensor_val.dtype}: {tensor_val.shape}" + + +def get_node_val_str(node: torch.fx.Node) -> str: + if is_single_tensor_node(node): + assert isinstance(node.meta["val"], FakeTensor) + return get_tensor_val_str(node.meta["val"]) + elif is_tensor_collection_node(node): + assert isinstance(node.meta["val"], (list, tuple)) + return f"[{', '.join(get_tensor_val_str(t) for t in node.meta['val'])}]" + else: + if "val" not in node.meta: + return str(node) + return str(node.meta["val"]) + + +def get_arg_node_val_str(arg_node: Any) -> str: + if isinstance(arg_node, torch.fx.Node): + return get_node_val_str(arg_node) + elif isinstance(arg_node, (list, tuple)): + return f"[{', '.join(get_arg_node_val_str(n) for n in arg_node)}]" + else: + return str(arg_node) + + +def node_io_str(node: torch.fx.Node) -> str: + target = node.target + if isinstance(target, EdgeOpOverload): + assert isinstance(target, EdgeOpOverload) + target_name = target.__name__ + elif isinstance(target, torch._ops.OpOverload): + assert isinstance(target, torch._ops.OpOverload) + target_name = target.name() + else: + target_name = str(target) + + out_str = f"{get_node_val_str(node)} = {target_name}(" + for arg in node.args: + out_str += get_arg_node_val_str(arg) + ", " + + out_str += " ...)" + return out_str + + def update_program_state_dict( program: ExportedProgram, buffer_name: str, @@ -1058,3 +1288,39 @@ def update_program_state_dict( # Finally, overwrite the current tensor with updated tensor program.state_dict[target_name] = updated_tensor + + +def align_width_and_update_state_dict( + ep: ExportedProgram, + node: torch.fx.Node, + cur_tensor: torch.Tensor, + align_to: int = 4, + force_update: bool = False, +) -> torch.Tensor: + """ + Align the width of the given tensor to the given alignment value and update the + state dict of the program with the aligned tensor. + """ + added_padding = False + cur_width = cur_tensor.shape[-1] + # Only align the width of the tensor if it is not already aligned + if cur_width % align_to != 0: + num_padding = align_to - (cur_width % align_to) + # Align the width of the tensor to the given alignment value + aligned_tensor = torch.nn.functional.pad( + cur_tensor, (0, num_padding) + ).contiguous() + added_padding = True + else: + aligned_tensor = cur_tensor + + if added_padding or force_update: + update_program_state_dict(ep, node.name, aligned_tensor) + # FakeTensor needs to match updated tensor + cur_fake_tensor = node.meta["val"] + node.meta["val"] = FakeTensorConverter().from_real_tensor( + cur_fake_tensor.fake_mode, + aligned_tensor, + ) + + return aligned_tensor diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index a6d5737dbb8..95da66494e0 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -13,15 +13,13 @@ import executorch.backends.vulkan.utils as utils from executorch.backends.transforms.addmm_mm_to_linear import AddmmToLinearTransform -from executorch.backends.transforms.fuse_batch_norm_with_conv import ( - FuseBatchNormWithConvPass, -) from executorch.backends.transforms.fuse_conv_with_clamp import FuseClampPass from executorch.backends.transforms.fuse_view_copy import FuseViewCopyTransform from executorch.backends.transforms.view_copy_to_squeeze_unsqueeze import ( ViewCopyToSqueezeUnsqueezePass, ) from executorch.backends.vulkan._passes import ( + FoldQDQPass, FuseQuantizedOpsTransform, insert_prepack_nodes, RemoveLocalScalarDenseOpsTransform, @@ -29,6 +27,7 @@ SqueezeUnsqueezeInputs, TagMemoryMetaPass, ) +from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.backends.vulkan._passes.remove_asserts import RemoveAssertsTransform from executorch.backends.vulkan.serialization.vulkan_graph_builder import VkGraphBuilder @@ -39,6 +38,7 @@ from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( serialize_vulkan_graph, ) +from executorch.backends.xnnpack._passes import FuseBatchNormPass from executorch.exir.backend.backend_details import ( BackendDetails, @@ -112,6 +112,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "downcast_64_bit": options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + if spec.key == "force_fp16": + options[spec.key] = bool.from_bytes(spec.value, byteorder="little") + # Unhandled options are ignored return options @@ -145,6 +148,7 @@ def preprocess( # noqa: C901 "memory_layout_override", VkMemoryLayout.TENSOR_WIDTH_PACKED ) downcast_64_bit = compile_options.get("downcast_64_bit", True) + force_fp16 = compile_options.get("force_fp16", False) program = unsafe_remove_auto_functionalized_pass(program) @@ -154,13 +158,15 @@ def preprocess( # noqa: C901 program = apply_passes( program, [ + FusePatternsPass(program), RemoveRedundantOpsTransform(), AddmmToLinearTransform(), FuseQuantizedOpsTransform(program), + FoldQDQPass(program), SqueezeUnsqueezeInputs(), FuseViewCopyTransform(), ViewCopyToSqueezeUnsqueezePass(), - FuseBatchNormWithConvPass(program), + FuseBatchNormPass(program), FuseClampPass(), ], ) @@ -219,6 +225,7 @@ def preprocess( # noqa: C901 program, DelegateMappingBuilder(generated_identifiers=True), downcast_64_bit=downcast_64_bit, + force_fp16=force_fp16, ) vk_graph = graph_builder.build_graph() @@ -227,4 +234,5 @@ def preprocess( # noqa: C901 vk_graph, graph_builder.const_tensors, [] ), debug_handle_map=graph_builder.delegate_mapping_builder.get_delegate_mapping(), + data_store_output=graph_builder.named_data_store.get_named_data_store_output(), ) diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt index 51abb4f2356..33bf84b9066 100644 --- a/backends/xnnpack/CMakeLists.txt +++ b/backends/xnnpack/CMakeLists.txt @@ -59,10 +59,10 @@ foreach(fbs_file ${_xnnpack_schema__srcs}) ) endforeach() -if(WIN32) +if(CMAKE_HOST_SYSTEM_NAME STREQUAL "Windows") set(MV_COMMAND powershell -Command - "Move-Item -Path ${_xnnpack_flatbuffer__outputs} -Destination ${_xnnpack_schema__outputs}" + "Move-Item -Path ${_xnnpack_flatbuffer__outputs} -Destination ${_xnnpack_schema__outputs} -Force" ) else() set(MV_COMMAND mv ${_xnnpack_flatbuffer__outputs} ${_xnnpack_schema__outputs}) diff --git a/backends/xnnpack/TARGETS b/backends/xnnpack/TARGETS index 62a703bddb7..d5c6d6303d2 100644 --- a/backends/xnnpack/TARGETS +++ b/backends/xnnpack/TARGETS @@ -36,10 +36,7 @@ runtime.python_library( ], deps = [ ":xnnpack_preprocess", - "//executorch/export:lib", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", "//executorch/backends/xnnpack/utils:xnnpack_utils", - "//executorch/backends/xnnpack/recipes:xnnpack_recipe_provider", - "//executorch/backends/xnnpack/recipes:xnnpack_recipe_types", ], ) diff --git a/backends/xnnpack/__init__.py b/backends/xnnpack/__init__.py index 01b73101c86..b87dfab4f02 100644 --- a/backends/xnnpack/__init__.py +++ b/backends/xnnpack/__init__.py @@ -4,18 +4,11 @@ # 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 - # Exposed Partitioners in XNNPACK Package from .partition.xnnpack_partitioner import ( XnnpackDynamicallyQuantizedPartitioner, XnnpackPartitioner, ) -from .recipes.xnnpack_recipe_provider import XNNPACKRecipeProvider -from .recipes.xnnpack_recipe_types import XNNPackRecipeType - -# Auto-register XNNPACK recipe provider -recipe_registry.register_backend_recipe_provider(XNNPACKRecipeProvider()) # Exposed Configs in XNNPACK Package from .utils.configs import ( @@ -34,7 +27,6 @@ "XnnpackDynamicallyQuantizedPartitioner", "XnnpackPartitioner", "XnnpackBackend", - "XNNPackRecipeType", "capture_graph_for_xnnpack", "get_xnnpack_capture_config", "get_xnnpack_edge_compile_config", diff --git a/backends/xnnpack/_passes/TARGETS b/backends/xnnpack/_passes/TARGETS index 5a038383f20..6f7b13d8026 100644 --- a/backends/xnnpack/_passes/TARGETS +++ b/backends/xnnpack/_passes/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "xnnpack_passes", srcs = native.glob(["*.py"]), deps = [ diff --git a/backends/xnnpack/_passes/fuse_batch_norm.py b/backends/xnnpack/_passes/fuse_batch_norm.py index a83be194e66..a51920ed5ad 100644 --- a/backends/xnnpack/_passes/fuse_batch_norm.py +++ b/backends/xnnpack/_passes/fuse_batch_norm.py @@ -82,12 +82,16 @@ def call(self, graph_module: torch.fx.GraphModule): @staticmethod def can_fuse( - input_node: torch.fx.Node, bn: torch.fx.Node, program: ExportedProgram + input_node: torch.fx.Node, + bn: torch.fx.Node, + program: ExportedProgram, ) -> bool: """ Determine whether a BatchNorm node can be fused with the preceding convolution or linear node. """ + is_conv = input_node.target == exir_ops.edge.aten.convolution.default + # All users of the batch_norm node must be getitem ops. # batch_norm returns a 3-element tuple. # Each user must only access the first element of the tuple. @@ -110,6 +114,16 @@ def can_fuse( ].count(False): return False + # Check the rank of the convolutution input - only Conv1d and 2d are supported. + if is_conv: + conv_input = input_node.args[0] + if ( + not isinstance(conv_input, torch.fx.Node) + or "val" not in conv_input.meta + or len(conv_input.meta["val"].shape) not in (3, 4) + ): + return False + return True def _fuse_ops( diff --git a/backends/xnnpack/cmake/Dependencies.cmake b/backends/xnnpack/cmake/Dependencies.cmake index 60ab7db6c05..62c8e7e0288 100644 --- a/backends/xnnpack/cmake/Dependencies.cmake +++ b/backends/xnnpack/cmake/Dependencies.cmake @@ -35,24 +35,36 @@ set(XNNPACK_BUILD_TESTS set(XNNPACK_ENABLE_AVXVNNI OFF CACHE BOOL "" - ) -# Work around observed failure: https://github.com/pytorch/executorch/pull/10362#issuecomment-2906391232 +) +# Work around observed failure: +# https://github.com/pytorch/executorch/pull/10362#issuecomment-2906391232 set(XNNPACK_ENABLE_AVX512VNNIGFNI - OFF - CACHE BOOL "") - + OFF + CACHE BOOL "" +) +set(XNNPACK_ENABLE_ARM_SME2 + ON + CACHE BOOL "" +) if(EXECUTORCH_XNNPACK_ENABLE_KLEIDI) - set(XNNPACK_ENABLE_KLEIDIAI - ON - CACHE BOOL "" - ) + set(XNNPACK_ENABLE_KLEIDIAI + ON + CACHE BOOL "" + ) else() - set(XNNPACK_ENABLE_KLEIDIAI - OFF - CACHE BOOL "" - ) + set(XNNPACK_ENABLE_KLEIDIAI + OFF + CACHE BOOL "" + ) endif() +if(WIN32) + # These XNNPACK options don't currently build on Windows. + set_overridable_option(XNNPACK_ENABLE_AVX256SKX OFF) + set_overridable_option(XNNPACK_ENABLE_AVX256VNNI OFF) + set_overridable_option(XNNPACK_ENABLE_AVX256VNNIGFNI OFF) + set_overridable_option(XNNPACK_ENABLE_AVX512BF16 OFF) +endif() set(XNNPACK_BUILD_ALL_MICROKERNELS OFF diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index 90a9a3063e3..6a055c9413f 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -621,8 +621,12 @@ def get_serialized_buffer_index( ConstantDataOffset(offset=UINT64_MAX, size=size, named_key=named_key) ) - external_tag = tensor.meta.get("delegate_constant_tag", None) + custom_meta = tensor.meta.get("custom", None) + external_tag = ( + custom_meta.get("delegate_constant_tag", None) if custom_meta else None + ) if external_tag is not None: + external_tag = custom_meta.get("delegate_constant_tag", None) logging.info( f"Adding constant data with name {tensor.name}, key {named_key} and external_tag {external_tag} to named_data_store" ) diff --git a/backends/xnnpack/operators/op_permute.py b/backends/xnnpack/operators/op_permute.py index 4d62d457cd0..d4086ed68b7 100644 --- a/backends/xnnpack/operators/op_permute.py +++ b/backends/xnnpack/operators/op_permute.py @@ -44,14 +44,21 @@ def define_node( self.define_nodes_tensor_inputs_outputs(node, xnn_graph, vals_to_ids) # input - input_id = vals_to_ids[get_input_node(node, 0)] + input_node = get_input_node(node, 0) + input_id = vals_to_ids[input_node] # output output_id = vals_to_ids[node] # permutation + input_rank = input_node.meta["val"].dim() permute_order = cast(List[int], node.args[1]) + # Handle negative dimensions by converting them to positive indices + permute_order = [ + (dim + input_rank) if dim < 0 else dim for dim in permute_order + ] + # change permute order if under channels last is_channels_last = node.meta.get( ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index ddbe8edc42d..817f9d1cf50 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -10,6 +10,11 @@ from typing import List, Optional import torch +from executorch.backends.xnnpack.utils.quant_utils import ( + is_dequant, + is_qparam, + is_quant, +) from executorch.exir.backend.canonical_partitioners.config_partitioner import ( format_target_name, PartitionerConfig, @@ -223,9 +228,18 @@ def _check_node_has_valid_dtype(self, node): valid_dtypes = { torch.float32, torch.float16, - torch.int8, - torch.qint8, } + # Only allow int8 and quant dtypes for quant operations + if is_quant(node) or is_dequant(node) or is_qparam(node): + valid_dtypes.update( + { + torch.qint32, + torch.qint8, + torch.quint8, + torch.int8, + } + ) + if ( node.op != "placeholder" and node.op != "call_function" diff --git a/backends/xnnpack/partition/xnnpack_partitioner.py b/backends/xnnpack/partition/xnnpack_partitioner.py index e5532e17f36..44207e2247a 100644 --- a/backends/xnnpack/partition/xnnpack_partitioner.py +++ b/backends/xnnpack/partition/xnnpack_partitioner.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. +import inspect import itertools - import logging from typing import List, Optional, Type, Union @@ -65,6 +65,37 @@ def __init__( self.per_op_mode = per_op_mode super().__init__(delegation_spec, initialized_configs) + def _check_if_called_from_to_backend(self) -> bool: + """ + Check if the partition method is being called from the deprecated to_backend workflow. + Returns True if called from deprecated direct to_backend, False if called from to_edge_transform_and_lower. + """ + stack = inspect.stack() + + for frame_info in stack: + if frame_info.function == "to_edge_transform_and_lower": + return False + + for frame_info in stack: + if frame_info.function == "to_backend": + filename = frame_info.filename + if "program/_program.py" in filename: + return True + return False + + def partition(self, exported_program): + """ + Override partition to add deprecation warning when called from to_backend. + """ + # Check if we're being called from the deprecated to_backend workflow + if self._check_if_called_from_to_backend(): + logger.warning( + "\nDEPRECATION WARNING: You are using the deprecated 'to_edge() + to_backend()' workflow. " + "Please consider migrating to 'to_edge_transform_and_lower()' for better error handling and optimization. " + ) + + return super().partition(exported_program) + def generate_partitions(self, ep: ExportedProgram) -> List[Partition]: """ generate_partitions is different if partitioner is set to per_op_mode diff --git a/backends/xnnpack/quantizer/TARGETS b/backends/xnnpack/quantizer/TARGETS index 761d1ed3821..0a30773923d 100644 --- a/backends/xnnpack/quantizer/TARGETS +++ b/backends/xnnpack/quantizer/TARGETS @@ -1,6 +1,6 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -python_library( +runtime.python_library( name = "xnnpack_quantizer", srcs = ["xnnpack_quantizer.py"], deps = [ @@ -11,7 +11,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "xnnpack_quantizer_utils", srcs = ["xnnpack_quantizer_utils.py"], deps = [ diff --git a/backends/xnnpack/recipes/TARGETS b/backends/xnnpack/recipes/TARGETS index 60968a5085d..6b6c1ddfe82 100644 --- a/backends/xnnpack/recipes/TARGETS +++ b/backends/xnnpack/recipes/TARGETS @@ -2,6 +2,22 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +runtime.python_library( + name = "xnnpack_recipes", + srcs = [ + "__init__.py", + ], + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + deps = [ + "//executorch/export:recipe_registry", + ":xnnpack_recipe_provider", + ":xnnpack_recipe_types", + ], +) + runtime.python_library( name = "xnnpack_recipe_provider", srcs = [ @@ -30,6 +46,6 @@ runtime.python_library( "@EXECUTORCH_CLIENTS", ], deps = [ - "//executorch/export:lib", + "//executorch/export:recipe", ], ) diff --git a/backends/xnnpack/recipes/__init__.py b/backends/xnnpack/recipes/__init__.py new file mode 100644 index 00000000000..3fa8e9496e6 --- /dev/null +++ b/backends/xnnpack/recipes/__init__.py @@ -0,0 +1,19 @@ +# 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 .xnnpack_recipe_provider import XNNPACKRecipeProvider +from .xnnpack_recipe_types import XNNPackRecipeType + +# Auto-register XNNPACK recipe provider +recipe_registry.register_backend_recipe_provider(XNNPACKRecipeProvider()) + + +__all__ = [ + "XNNPACKRecipeProvider", + "XNNPackRecipeType", +] diff --git a/backends/xnnpack/recipes/xnnpack_recipe_provider.py b/backends/xnnpack/recipes/xnnpack_recipe_provider.py index 8fba58c12c3..2c80d528c45 100644 --- a/backends/xnnpack/recipes/xnnpack_recipe_provider.py +++ b/backends/xnnpack/recipes/xnnpack_recipe_provider.py @@ -6,6 +6,7 @@ # pyre-strict +import logging from typing import Any, Optional, Sequence import torch @@ -25,6 +26,7 @@ get_xnnpack_executorch_backend_config, ) from executorch.export import ( + AOQuantizationConfig, BackendRecipeProvider, ExportRecipe, LoweringRecipe, @@ -57,31 +59,37 @@ def create_recipe( if recipe_type == XNNPackRecipeType.FP32: return self._build_fp32_recipe(recipe_type) - elif recipe_type == XNNPackRecipeType.INT8_DYNAMIC_PER_CHANNEL: + elif recipe_type == XNNPackRecipeType.PT2E_INT8_DYNAMIC_PER_CHANNEL: return self._build_quantized_recipe( recipe_type, is_per_channel=True, is_dynamic=True ) - elif recipe_type == XNNPackRecipeType.INT8_STATIC_PER_CHANNEL: + elif recipe_type == XNNPackRecipeType.PT2E_INT8_STATIC_PER_CHANNEL: return self._build_quantized_recipe( recipe_type, is_per_channel=True, is_dynamic=False ) - elif recipe_type == XNNPackRecipeType.INT8_STATIC_PER_TENSOR: + elif recipe_type == XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR: return self._build_quantized_recipe( recipe_type, is_per_channel=False, is_dynamic=False ) - elif recipe_type == XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL: - return self._build_int8da_intx_weight_recipe( + elif ( + recipe_type + == XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL + ): + return self._build_torchao_quantized_recipe( recipe_type=recipe_type, is_per_channel=True, weight_dtype=torch.int4, ) - elif recipe_type == XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR: + elif ( + recipe_type + == XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR + ): group_size = kwargs.get("group_size", 32) - return self._build_int8da_intx_weight_recipe( + return self._build_torchao_quantized_recipe( recipe_type=recipe_type, is_per_channel=False, weight_dtype=torch.int4, @@ -132,7 +140,7 @@ def _build_quantized_recipe( executorch_backend_config=get_xnnpack_executorch_backend_config(), ) - def _build_int8da_intx_weight_recipe( + def _build_torchao_quantized_recipe( self, recipe_type: RecipeType, is_per_channel: bool = True, @@ -141,17 +149,21 @@ def _build_int8da_intx_weight_recipe( ) -> ExportRecipe: if is_per_channel: weight_granularity = PerAxis(axis=0) + assert weight_dtype == torch.int4 or weight_dtype == torch.int8 else: weight_granularity = PerGroup(group_size=group_size) + assert weight_dtype == torch.int4 - config = Int8DynamicActivationIntxWeightConfig( - weight_dtype=weight_dtype, - weight_granularity=weight_granularity, + config = AOQuantizationConfig( + Int8DynamicActivationIntxWeightConfig( + weight_dtype=weight_dtype, + weight_granularity=weight_granularity, + ) ) quant_recipe = QuantizationRecipe( quantizers=None, - ao_base_config=[config], + ao_quantization_configs=[config], ) return ExportRecipe( @@ -162,13 +174,16 @@ def _build_int8da_intx_weight_recipe( ) def _validate_recipe_kwargs(self, recipe_type: RecipeType, **kwargs: Any) -> None: - if recipe_type == XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR: + if ( + recipe_type + == XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR + ): expected_keys = {"group_size"} unexpected = set(kwargs.keys()) - expected_keys if unexpected: - raise ValueError( - f"Recipe '{recipe_type.value}' only accepts 'group_size' parameter. " - f"Unexpected parameters: {list(unexpected)}" + logging.warning( + f"XNNPACK recipe '{recipe_type.value}' ignoring unexpected parameters: {list(unexpected)}. " + f"Only 'group_size' is supported for this recipe." ) if "group_size" in kwargs: group_size = kwargs["group_size"] @@ -179,7 +194,7 @@ def _validate_recipe_kwargs(self, recipe_type: RecipeType, **kwargs: Any) -> Non elif kwargs: # All other recipes don't expect any kwargs unexpected = list(kwargs.keys()) - raise ValueError( - f"Recipe '{recipe_type.value}' does not accept any parameters. " - f"Unexpected parameters: {unexpected}" + logging.warning( + f"XNNPACK recipe '{recipe_type.value}' ignoring unexpected parameters: {unexpected}. " + f"This recipe does not accept any parameters." ) diff --git a/backends/xnnpack/recipes/xnnpack_recipe_types.py b/backends/xnnpack/recipes/xnnpack_recipe_types.py index 5675c3a5ffa..82296f912c2 100644 --- a/backends/xnnpack/recipes/xnnpack_recipe_types.py +++ b/backends/xnnpack/recipes/xnnpack_recipe_types.py @@ -12,20 +12,25 @@ class XNNPackRecipeType(RecipeType): """XNNPACK-specific recipe types""" - FP32 = "fp32" + FP32 = "xnnpack_fp32" + + ## PT2E-based quantization recipes # INT8 Dynamic Quantization - INT8_DYNAMIC_PER_CHANNEL = "int8_dynamic_per_channel" + PT2E_INT8_DYNAMIC_PER_CHANNEL = "xnnpack_pt2e_int8_dynamic_per_channel" + # INT8 Static Quantization, needs calibration dataset + PT2E_INT8_STATIC_PER_CHANNEL = "xnnpack_pt2e_int8_static_per_channel" + PT2E_INT8_STATIC_PER_TENSOR = "xnnpack_pt2e_int8_static_per_tensor" + + ## TorchAO-based quantization recipes # INT8 Dynamic Activations INT4 Weight Quantization, Axis = 0 - INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL = "int8da_int4w_per_channel" + TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL = ( + "xnnpack_torchao_int8da_int4w_per_channel" + ) # INT8 Dynamic Activations INT4 Weight Quantization, default group_size = 32 # can be overriden by group_size kwarg - INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR = "int8da_int4w_per_tensor" - # INT8 Static Activations INT4 Weight Quantization - INT8_STATIC_ACT_INT4_WEIGHT_PER_CHANNEL = "int8a_int4w_per_channel" - INT8_STATIC_ACT_INT4_WEIGHT_PER_TENSOR = "int8a_int44w_per_tensor" - # INT8 Static Quantization, needs calibration dataset - INT8_STATIC_PER_CHANNEL = "int8_static_per_channel" - INT8_STATIC_PER_TENSOR = "int8_static_per_tensor" + TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR = ( + "xnnpack_torchao_int8da_int4w_per_tensor" + ) @classmethod def get_backend_name(cls) -> str: diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index b886bba2857..78eaaf6d039 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -609,9 +609,6 @@ bool isQP8(const fb_xnnpack::XNNGraph* graph, const NodePtr node) { auto cvt_output_id = graph_node->output_id(); auto check_dtype = [graph](uint32_t id, DataType dtype) -> bool { - assert( - dtype == DataType::xnn_datatype_qdint8 || - dtype == DataType::xnn_datatype_qbint4); for (auto value : *graph->xvalues()) { if (value->xvalue_union_type() != fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) { @@ -631,6 +628,12 @@ bool isQP8(const fb_xnnpack::XNNGraph* graph, const NodePtr node) { return false; } + // XNNPACK dtypes which have qp8 support. + const std::vector supported_filter_dtypes = { + DataType::xnn_datatype_qbint4, + DataType::xnn_datatype_qcint4, + DataType::xnn_datatype_qcint8}; + // Find if the convert output is going to the right linear node. // Assuming if we can find one valid linear node, then we can use QP8 // for all the linear nodes consuming this convert output. @@ -638,9 +641,10 @@ bool isQP8(const fb_xnnpack::XNNGraph* graph, const NodePtr node) { if (node->xnode_union_type() == fb_xnnpack::XNodeUnion::XNNFullyConnected) { auto linear_node = node->xnode_union_as_XNNFullyConnected(); if (linear_node->input1_id() == cvt_output_id) { - if (check_dtype( - linear_node->filter_id(), DataType::xnn_datatype_qbint4)) { - return true; + for (auto supported_filter_dtype : supported_filter_dtypes) { + if (check_dtype(linear_node->filter_id(), supported_filter_dtype)) { + return true; + } } } } diff --git a/backends/xnnpack/runtime/XNNExecutor.cpp b/backends/xnnpack/runtime/XNNExecutor.cpp index 9802da5c06e..3b3b16dbb91 100644 --- a/backends/xnnpack/runtime/XNNExecutor.cpp +++ b/backends/xnnpack/runtime/XNNExecutor.cpp @@ -21,6 +21,7 @@ using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::is_contiguous_dim_order; using executorch::runtime::kTensorDimensionLimit; +using executorch::runtime::Span; /** * Initializes the XNNExecutor with the runtime and given number of @@ -69,7 +70,7 @@ ET_NODISCARD Error XNNExecutor::initialize( * runtime correspond to their index in the list of arg passed into * delegate->execute() */ -ET_NODISCARD Error XNNExecutor::prepare_args(EValue** args) { +ET_NODISCARD Error XNNExecutor::prepare_args(Span args) { ET_CHECK_OR_RETURN_ERROR( runtime_ != nullptr, Internal, @@ -196,7 +197,7 @@ ET_NODISCARD Error XNNExecutor::forward(BackendExecutionContext& context) { * XNNPACK gives the index tensor to us as int32, we need to convert it * back to int64 for ExecuTorch. */ -ET_NODISCARD Error XNNExecutor::resize_outputs(EValue** args) const { +ET_NODISCARD Error XNNExecutor::resize_outputs(Span args) const { size_t output_idx_start = input_ids_.size(); for (size_t i = output_idx_start; i < externals_.size(); ++i) { uint32_t ext_id = externals_[i].id; diff --git a/backends/xnnpack/runtime/XNNExecutor.h b/backends/xnnpack/runtime/XNNExecutor.h index 8131b6b8b2c..f7084a5dd88 100644 --- a/backends/xnnpack/runtime/XNNExecutor.h +++ b/backends/xnnpack/runtime/XNNExecutor.h @@ -69,7 +69,7 @@ class XNNExecutor { * any additional memory planning as needed */ ET_NODISCARD executorch::runtime::Error prepare_args( - executorch::runtime::EValue** args); + executorch::runtime::Span args); /** * Executes the graph using the args prepared at prepare_args(). @@ -83,7 +83,7 @@ class XNNExecutor { * Performs any post processing of outputs like tensor resizing */ ET_NODISCARD executorch::runtime::Error resize_outputs( - executorch::runtime::EValue** args) const; + executorch::runtime::Span args) const; friend class XNNCompiler; }; diff --git a/backends/xnnpack/runtime/XNNPACKBackend.cpp b/backends/xnnpack/runtime/XNNPACKBackend.cpp index 9e02d566d99..b05919ecf2b 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.cpp +++ b/backends/xnnpack/runtime/XNNPACKBackend.cpp @@ -33,6 +33,7 @@ using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::FreeableBuffer; using executorch::runtime::Result; +using executorch::runtime::Span; class XnnpackBackend final : public ::executorch::ET_RUNTIME_NAMESPACE::BackendInterface { @@ -126,7 +127,7 @@ class XnnpackBackend final Error execute( BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const override { + Span args) const override { auto executor = static_cast(handle); #ifdef ENABLE_XNNPACK_SHARED_WORKSPACE diff --git a/backends/xnnpack/serialization/xnnpack_graph_serialize.py b/backends/xnnpack/serialization/xnnpack_graph_serialize.py index 0fbd0ddc5ef..712ddccb452 100644 --- a/backends/xnnpack/serialization/xnnpack_graph_serialize.py +++ b/backends/xnnpack/serialization/xnnpack_graph_serialize.py @@ -1,22 +1,21 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 importlib.resources as _resources import json - import logging import os import tempfile - from dataclasses import dataclass, fields, is_dataclass from typing import ClassVar, Literal, Optional -import pkg_resources +import executorch.backends.xnnpack.serialization as serialization_package from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNNGraph from executorch.exir._serialize._dataclass import _DataclassEncoder - from executorch.exir._serialize._flatbuffer import _flatc_compile logger = logging.getLogger(__name__) @@ -317,7 +316,9 @@ def convert_to_flatbuffer(xnnpack_graph: XNNGraph) -> bytes: with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") with open(schema_path, "wb") as schema_file: - schema_file.write(pkg_resources.resource_string(__name__, "schema.fbs")) + schema_file.write( + _resources.read_binary(serialization_package, "schema.fbs") + ) json_path = os.path.join(d, "schema.json") with open(json_path, "wb") as json_file: json_file.write(xnnpack_graph_json.encode("ascii")) diff --git a/backends/xnnpack/targets.bzl b/backends/xnnpack/targets.bzl index aee5104b17a..0eab89a00f9 100644 --- a/backends/xnnpack/targets.bzl +++ b/backends/xnnpack/targets.bzl @@ -1,4 +1,5 @@ load("@fbsource//xplat/executorch/backends/xnnpack/third-party:third_party_libs.bzl", "third_party_dep") +load("@fbsource//xplat/executorch/build:build_variables.bzl", "XNNPACK_BACKEND_BUCK_SRCS") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime") def _get_preprocessor_flags(): @@ -37,10 +38,7 @@ def define_common_targets(): aten_suffix = "_aten" if aten_mode else "" runtime.cxx_library( name = "xnnpack_backend" + aten_suffix, - srcs = native.glob([ - "runtime/*.cpp", - "runtime/profiling/*.cpp", - ]), + srcs = XNNPACK_BACKEND_BUCK_SRCS, headers = native.glob([ "runtime/*.h", "runtime/profiling/*.h", diff --git a/backends/xnnpack/test/CMakeLists.txt b/backends/xnnpack/test/CMakeLists.txt index 12d0a6d45be..395fb01d189 100644 --- a/backends/xnnpack/test/CMakeLists.txt +++ b/backends/xnnpack/test/CMakeLists.txt @@ -17,9 +17,8 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) -set(_test_srcs - runtime/test_xnnexecutor.cpp - ${EXECUTORCH_ROOT}/extension/threadpool/test/threadpool_test.cpp +set(_test_srcs runtime/test_xnnexecutor.cpp + ${EXECUTORCH_ROOT}/extension/threadpool/test/threadpool_test.cpp ) et_cxx_test( diff --git a/backends/xnnpack/test/TARGETS b/backends/xnnpack/test/TARGETS index e024721b556..5f3581b6aeb 100644 --- a/backends/xnnpack/test/TARGETS +++ b/backends/xnnpack/test/TARGETS @@ -100,8 +100,12 @@ runtime.python_test( srcs = glob([ "recipes/*.py", ]), + env = { + "HTTP_PROXY": "http://fwdproxy:8080", + "HTTPS_PROXY": "http://fwdproxy:8080", + }, deps = [ - "//executorch/backends/xnnpack:xnnpack_delegate", + "//executorch/backends/xnnpack/recipes:xnnpack_recipes", "//executorch/export:lib", "//pytorch/vision:torchvision", # @manual "//executorch/backends/xnnpack/test/tester:tester", diff --git a/backends/xnnpack/test/ops/test_check_quant_params.py b/backends/xnnpack/test/ops/test_check_quant_params.py index 8be59aab50e..e077462ee2c 100644 --- a/backends/xnnpack/test/ops/test_check_quant_params.py +++ b/backends/xnnpack/test/ops/test_check_quant_params.py @@ -9,7 +9,7 @@ ) from executorch.backends.xnnpack.utils.utils import get_param_tensor from executorch.exir import to_edge_transform_and_lower -from torch.export import export_for_training +from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -52,7 +52,7 @@ def _test_check_quant_message(self, ep_modifier, expected_message): torch._dynamo.reset() mod = torch.nn.Linear(10, 10) quantizer = XNNPACKQuantizer() - captured = export_for_training(mod, (torch.randn(1, 10),), strict=True).module() + captured = export(mod, (torch.randn(1, 10),), strict=True).module() quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True)) prepared = prepare_pt2e(captured, quantizer) diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index 421e59c0b08..ac6fec25732 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -34,8 +34,9 @@ from torch.export.graph_signature import ExportGraphSignature, InputKind try: + from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import ( - int8_dynamic_activation_int4_weight, + Int8DynamicActivationIntxWeightConfig, quantize_, ) from torchao.utils import unwrap_tensor_subclass @@ -391,7 +392,12 @@ def _test_groupwise_dq_linear( """ Helper function to test groupwise dynamic quantized linear op with different configurations. """ - quantize_(mod, int8_dynamic_activation_int4_weight(group_size=group_size)) + quantize_( + mod, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(group_size) + ), + ) unwrap_tensor_subclass(mod) DynamicallyQuantizedPartitioner = XnnpackPartitioner( config_precisions=ConfigPrecisionType.DYNAMIC_QUANT, diff --git a/backends/xnnpack/test/ops/test_permute.py b/backends/xnnpack/test/ops/test_permute.py index 2991ba1773d..e0171f2178f 100644 --- a/backends/xnnpack/test/ops/test_permute.py +++ b/backends/xnnpack/test/ops/test_permute.py @@ -55,6 +55,20 @@ def test_fp32_permute(self): inputs = (torch.randn(1, 1, 4, 4),) self._test_permute(inputs) + def test_fp32_permute_negative_dim(self): + inputs = (torch.randn(1, 1, 4, 4),) + ( + Tester(self.Permute([0, -2, -1, 1]), inputs) + .export() + .check_count({"torch.ops.aten.permute.default": 1}) + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not(["executorch_exir_dialects_edge__ops_aten_permute_copy_default"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_fp32_permute_copy(self): inputs = (torch.randn(1, 1, 4, 4),) ( diff --git a/backends/xnnpack/test/passes/test_batch_norm_fusion.py b/backends/xnnpack/test/passes/test_batch_norm_fusion.py index a095fa236fe..4e4cd065fa5 100644 --- a/backends/xnnpack/test/passes/test_batch_norm_fusion.py +++ b/backends/xnnpack/test/passes/test_batch_norm_fusion.py @@ -55,6 +55,26 @@ def forward(self, x): y = y + y return self.bn(y) + class ModelConv3dBN(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + kernel_size: Tuple[int, int, int], + ): + super().__init__() + op = torch.nn.Conv3d + self.conv3d = op(in_features, out_features, kernel_size) + self.bn = torch.nn.BatchNorm3d(out_features) + self.forward(torch.randn(2, 2, 4, 4, 4) * 2 + 2) # update the BN stats + + def forward(self, x): + y = self.conv3d(x) + y = self.bn(y) + y = self.conv3d(y) + y = y + y + return self.bn(y) + def test_fp32_conv_batch_norm_fusion(self): for transpose in [False, True]: ( @@ -142,3 +162,18 @@ def forward(self, x): .to_edge_transform_and_lower() .check_count({self.bn_name: 1}) ) + + def test_fp32_conv3d_batch_norm_doesnt_partition(self): + """ + Conv3d is not currently supported by XNNPACK. We also don't support standalone + batch norms yet (i.e. batch norms that are not fused with a conv). As such, we don't + want to partition the standalone batch norm and then fail to lower. + """ + ( + Tester(self.ModelConv3dBN(2, 2, (2, 2, 2)), (torch.randn(2, 2, 4, 4, 4),)) + .export() + .dump_artifact() + .to_edge_transform_and_lower() + .check_count({self.bn_name: 2}) + .run_method_and_compare_outputs() + ) diff --git a/backends/xnnpack/test/quantizer/test_pt2e_quantization.py b/backends/xnnpack/test/quantizer/test_pt2e_quantization.py index 5d76ecd2d54..f0456a3604e 100644 --- a/backends/xnnpack/test/quantizer/test_pt2e_quantization.py +++ b/backends/xnnpack/test/quantizer/test_pt2e_quantization.py @@ -22,7 +22,7 @@ weight_observer_range_neg_127_to_127, ) from torch.ao.quantization.qconfig_mapping import QConfigMapping -from torch.export import export_for_training +from torch.export import export from torch.testing._internal.common_quantization import ( NodeSpec as ns, TestHelperModules, @@ -58,7 +58,7 @@ def _quantize(self, m, quantizer, example_inputs, is_qat: bool = False): # resetting dynamo cache torch._dynamo.reset() - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() if is_qat: m = prepare_qat_pt2e(m, quantizer) else: @@ -351,7 +351,7 @@ def test_disallow_eval_train(self) -> None: m.train() # After export: this is not OK - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() with self.assertRaises(NotImplementedError): m.eval() with self.assertRaises(NotImplementedError): @@ -405,7 +405,7 @@ def forward(self, x): m = M().train() example_inputs = (torch.randn(1, 3, 3, 3),) bn_train_op, bn_eval_op = self._get_bn_train_eval_ops() # pyre-ignore[23] - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() def _assert_ops_are_correct(m: torch.fx.GraphModule, train: bool) -> None: bn_op = bn_train_op if train else bn_eval_op @@ -474,7 +474,7 @@ def forward(self, x): quantizer.set_global(operator_config) example_inputs = (torch.randn(2, 2),) m = M().eval() - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() weight_meta = None for n in m.graph.nodes: # pyre-ignore[16] if ( @@ -503,7 +503,7 @@ def test_reentrant(self) -> None: quantizer = XNNPACKQuantizer().set_global( get_symmetric_quantization_config(is_per_channel=True, is_qat=True) ) - m.conv_bn_relu = export_for_training( # pyre-ignore[8] + m.conv_bn_relu = export( # pyre-ignore[8] m.conv_bn_relu, example_inputs, strict=True ).module() m.conv_bn_relu = prepare_qat_pt2e(m.conv_bn_relu, quantizer) # pyre-ignore[6,8] @@ -513,7 +513,7 @@ def test_reentrant(self) -> None: quantizer = XNNPACKQuantizer().set_module_type( torch.nn.Linear, get_symmetric_quantization_config(is_per_channel=False) ) - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() m = prepare_pt2e(m, quantizer) # pyre-ignore[6] m = convert_pt2e(m) @@ -575,7 +575,7 @@ def check_nn_module(node: torch.fx.Node) -> None: "ConvWithBNRelu" in node.meta["nn_module_stack"]["L__self__"][1] ) - m.conv_bn_relu = export_for_training( # pyre-ignore[8] + m.conv_bn_relu = export( # pyre-ignore[8] m.conv_bn_relu, example_inputs, strict=True ).module() for node in m.conv_bn_relu.graph.nodes: # pyre-ignore[16] @@ -591,7 +591,7 @@ def test_speed(self) -> None: def dynamic_quantize_pt2e(model, example_inputs) -> torch.fx.GraphModule: torch._dynamo.reset() - model = export_for_training(model, example_inputs, strict=True).module() + model = export(model, example_inputs, strict=True).module() # Per channel quantization for weight # Dynamic quantization for activation # Please read a detail: https://fburl.com/code/30zds51q @@ -648,7 +648,7 @@ def forward(self, x): example_inputs = (torch.randn(1, 3, 5, 5),) m = M() - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer().set_global( get_symmetric_quantization_config(), ) @@ -724,11 +724,10 @@ def test_save_load(self) -> None: class TestXNNPACKQuantizerNumericDebugger(PT2ENumericDebuggerTestCase): - def test_quantize_pt2e_preserve_handle(self): m = TestHelperModules.Conv2dThenConv1d() example_inputs = m.example_inputs() - ep = export_for_training(m, example_inputs, strict=True) + ep = export(m, example_inputs, strict=True) m = ep.module() quantizer = XNNPACKQuantizer().set_global( @@ -768,7 +767,7 @@ def test_quantize_pt2e_preserve_handle(self): def test_extract_results_from_loggers(self): m = TestHelperModules.Conv2dThenConv1d() example_inputs = m.example_inputs() - ep = export_for_training(m, example_inputs, strict=True) + ep = export(m, example_inputs, strict=True) m = ep.module() m_ref_logger = prepare_for_propagation_comparison(m) @@ -792,7 +791,7 @@ def test_extract_results_from_loggers(self): def test_extract_results_from_loggers_list_output(self): m = TestHelperModules.Conv2dWithSplit() example_inputs = m.example_inputs() - ep = export_for_training(m, example_inputs, strict=True) + ep = export(m, example_inputs, strict=True) m = ep.module() m_ref_logger = prepare_for_propagation_comparison(m) diff --git a/backends/xnnpack/test/quantizer/test_representation.py b/backends/xnnpack/test/quantizer/test_representation.py index 817f7f9e368..614bc8b83d6 100644 --- a/backends/xnnpack/test/quantizer/test_representation.py +++ b/backends/xnnpack/test/quantizer/test_representation.py @@ -8,7 +8,7 @@ XNNPACKQuantizer, ) from torch._higher_order_ops.out_dtype import out_dtype # noqa: F401 -from torch.export import export_for_training +from torch.export import export from torch.testing._internal.common_quantization import ( NodeSpec as ns, QuantizationTestCase, @@ -33,7 +33,7 @@ def _test_representation( ) -> None: # resetting dynamo cache torch._dynamo.reset() - model = export_for_training(model, example_inputs, strict=True).module() + model = export(model, example_inputs, strict=True).module() model_copy = copy.deepcopy(model) model = prepare_pt2e(model, quantizer) # pyre-ignore[6] diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 84b1a932a5b..1e1a473dd59 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -27,7 +27,7 @@ convert_to_reference_fx, prepare_fx, ) -from torch.export import export_for_training +from torch.export import export from torch.testing._internal.common_quantization import ( NodeSpec as ns, skip_if_no_torchvision, @@ -500,7 +500,7 @@ def forward(self, x): ) example_inputs = (torch.randn(2, 2),) m = M().eval() - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() m = prepare_pt2e(m, quantizer) # pyre-ignore[6] # Use a linear count instead of names because the names might change, but # the order should be the same. @@ -636,7 +636,7 @@ def test_propagate_annotation(self): example_inputs = (torch.randn(1, 3, 5, 5),) # program capture - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() m = prepare_pt2e(m, quantizer) m(*example_inputs) @@ -901,9 +901,7 @@ def forward(self, input_tensor, hidden_tensor): model_fx = _convert_to_reference_decomposed_fx(model_fx) with torchdynamo.config.patch(allow_rnn=True): - model_graph = export_for_training( - model_graph, example_inputs, strict=True - ).module() + model_graph = export(model_graph, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config( is_per_channel=False, is_dynamic=False @@ -963,9 +961,7 @@ def forward(self, input_tensor, hidden_tensor): model_fx = _convert_to_reference_decomposed_fx(model_fx) with torchdynamo.config.patch(allow_rnn=True): - model_graph = export_for_training( - model_graph, example_inputs, strict=True - ).module() + model_graph = export(model_graph, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config( is_per_channel=False, is_dynamic=False @@ -1173,7 +1169,7 @@ def test_resnet18(self): m = torchvision.models.resnet18().eval() m_copy = copy.deepcopy(m) # program capture - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config(is_per_channel=True) diff --git a/backends/xnnpack/test/recipes/test_xnnpack_recipes.py b/backends/xnnpack/test/recipes/test_xnnpack_recipes.py index 679743e42d3..e4a469418cd 100644 --- a/backends/xnnpack/test/recipes/test_xnnpack_recipes.py +++ b/backends/xnnpack/test/recipes/test_xnnpack_recipes.py @@ -6,7 +6,10 @@ # pyre-strict +import logging +import os import unittest +from typing import List, Optional, Tuple import torch from executorch.backends.xnnpack.recipes.xnnpack_recipe_provider import ( @@ -18,9 +21,17 @@ from executorch.examples.models.model_factory import EagerModelFactory from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS, QuantType from executorch.exir.schema import DelegateCall, Program -from executorch.export import export, ExportRecipe, recipe_registry -from torch import nn +from executorch.export import ( + export, + ExportRecipe, + ExportSession, + recipe_registry, + StageType, +) +from torch import nn, Tensor +from torch.testing import FileCheck from torch.testing._internal.common_quantization import TestHelperModules +from torchao.quantization.utils import compute_error class TestXnnpackRecipes(unittest.TestCase): @@ -38,6 +49,36 @@ def check_fully_delegated(self, program: Program) -> None: self.assertEqual(len(instructions), 1) self.assertIsInstance(instructions[0].instr_args, DelegateCall) + def _compare_eager_quantized_model_outputs( + self, + # pyre-ignore[11] + session: ExportSession, + example_inputs: List[Tuple[Tensor]], + atol: float, + ) -> None: + """Utility to compare eager quantized model output with session output after xnnpack lowering""" + torch_export_stage_output = session.get_stage_artifacts()[ + StageType.TORCH_EXPORT + ] + eager_quantized_model = torch_export_stage_output.data["forward"].module() + output = session.run_method("forward", example_inputs[0])[0] + expected = eager_quantized_model(*example_inputs[0]) + Tester._assert_outputs_equal(output, expected, atol=atol) + + def _compare_eager_unquantized_model_outputs( + self, + session: ExportSession, + eager_unquantized_model: nn.Module, + example_inputs: List[Tuple[Tensor]], + sqnr_threshold: int = 20, + ) -> None: + """Utility to compare eager unquantized model output with session output using SQNR""" + quantized_output = session.run_method("forward", example_inputs[0])[0] + original_output = eager_unquantized_model(*example_inputs[0]) + error = compute_error(original_output, quantized_output) + print(f"{self._testMethodName} - SQNR: {error} dB") + self.assertTrue(error > sqnr_threshold) + def test_basic_recipe(self) -> None: m_eager = TestHelperModules.TwoLinearModule().eval() example_inputs = [(torch.randn(9, 8),)] @@ -46,18 +87,13 @@ def test_basic_recipe(self) -> None: example_inputs=example_inputs, export_recipe=ExportRecipe.get_recipe(XNNPackRecipeType.FP32), ) - self.assertTrue( - torch.allclose( - session.run_method("forward", example_inputs[0])[0], - m_eager(*example_inputs[0]), - atol=1e-3, - ) - ) + self._compare_eager_quantized_model_outputs(session, example_inputs, 1e-3) self.check_fully_delegated(session.get_executorch_program()) + self._compare_eager_unquantized_model_outputs(session, m_eager, example_inputs) def test_int8_dynamic_quant_recipe(self) -> None: test_cases = [ - ExportRecipe.get_recipe(XNNPackRecipeType.INT8_DYNAMIC_PER_CHANNEL), + ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_DYNAMIC_PER_CHANNEL), ] for export_recipe in test_cases: @@ -70,19 +106,18 @@ def test_int8_dynamic_quant_recipe(self) -> None: example_inputs=example_inputs, export_recipe=export_recipe, ) - self.assertTrue( - torch.allclose( - session.run_method("forward", example_inputs[0])[0], - m_eager(*example_inputs[0]), - atol=1e-1, - ) + self._compare_eager_quantized_model_outputs( + session, example_inputs, 1e-1 ) self.check_fully_delegated(session.get_executorch_program()) + self._compare_eager_unquantized_model_outputs( + session, m_eager, example_inputs + ) def test_int8_static_quant_recipe(self) -> None: test_cases = [ - ExportRecipe.get_recipe(XNNPackRecipeType.INT8_STATIC_PER_CHANNEL), - ExportRecipe.get_recipe(XNNPackRecipeType.INT8_STATIC_PER_TENSOR), + ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_STATIC_PER_CHANNEL), + ExportRecipe.get_recipe(XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR), ] for export_recipe in test_cases: @@ -95,14 +130,13 @@ def test_int8_static_quant_recipe(self) -> None: example_inputs=example_inputs, export_recipe=export_recipe, ) - self.assertTrue( - torch.allclose( - session.run_method("forward", example_inputs[0])[0], - m_eager(*example_inputs[0]), - atol=1e-1, - ) + self._compare_eager_quantized_model_outputs( + session, example_inputs, 1e-2 ) self.check_fully_delegated(session.get_executorch_program()) + self._compare_eager_unquantized_model_outputs( + session, m_eager, example_inputs + ) def test_8a4w_recipe(self) -> None: class SimpleLinearModel(nn.Module): @@ -116,46 +150,45 @@ def forward(self, x) -> torch.Tensor: test_cases = [ ExportRecipe.get_recipe( - XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL, + XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_CHANNEL, ), ExportRecipe.get_recipe( - XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, + XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, group_size=32, ), ] for export_recipe in test_cases: with self.subTest(export_recipe=export_recipe): - model = SimpleLinearModel() + model = SimpleLinearModel().eval() example_inputs = [(torch.randn(1, 32),)] session = export( model=model, example_inputs=example_inputs, export_recipe=export_recipe, ) - self.assertTrue( - torch.allclose( - session.run_method("forward", example_inputs[0])[0], - model(*example_inputs[0]), - atol=1e-2, - ) - ) self.check_fully_delegated(session.get_executorch_program()) + self._compare_eager_quantized_model_outputs( + session, example_inputs, 1e-3 + ) def _get_recipe_for_quant_type(self, quant_type: QuantType) -> XNNPackRecipeType: # Map QuantType to corresponding recipe name. if quant_type == QuantType.STATIC_PER_CHANNEL: - return XNNPackRecipeType.INT8_STATIC_PER_CHANNEL + return XNNPackRecipeType.PT2E_INT8_STATIC_PER_CHANNEL elif quant_type == QuantType.DYNAMIC_PER_CHANNEL: - return XNNPackRecipeType.INT8_DYNAMIC_PER_CHANNEL + return XNNPackRecipeType.PT2E_INT8_DYNAMIC_PER_CHANNEL elif quant_type == QuantType.STATIC_PER_TENSOR: - return XNNPackRecipeType.INT8_STATIC_PER_TENSOR - elif quant_type == QuantType.NONE: - return XNNPackRecipeType.FP32 - else: - raise ValueError(f"Unsupported QuantType: {quant_type}") + return XNNPackRecipeType.PT2E_INT8_STATIC_PER_TENSOR + return XNNPackRecipeType.FP32 - def _test_model_with_factory(self, model_name: str) -> None: + def _test_model_with_factory( + self, + model_name: str, + tolerance: Optional[float] = None, + sqnr_threshold: Optional[float] = None, + ) -> None: + logging.info(f"Testing model {model_name}") if model_name not in MODEL_NAME_TO_MODEL: self.skipTest(f"Model {model_name} not found in MODEL_NAME_TO_MODEL") return @@ -182,40 +215,76 @@ def _test_model_with_factory(self, model_name: str) -> None: dynamic_shapes=dynamic_shapes, ) - # Verify outputs match - Tester._assert_outputs_equal( - session.run_method("forward", example_inputs)[0], - model(*example_inputs), - atol=1e-3, + all_artifacts = session.get_stage_artifacts() + quantized_model = all_artifacts[StageType.QUANTIZE].data["forward"] + + edge_program_manager = all_artifacts[StageType.TO_EDGE_TRANSFORM_AND_LOWER].data + lowered_module = edge_program_manager.exported_program().module() + + # Check if model got lowered to xnnpack backend + FileCheck().check("torch.ops.higher_order.executorch_call_delegate").run( + lowered_module.code ) - @unittest.skip("T187799178: Debugging Numerical Issues with Calibration") + if tolerance is not None: + quantized_output = quantized_model(*example_inputs) + lowered_output = lowered_module(*example_inputs) + if model_name == "dl3": + quantized_output = quantized_output["out"] + lowered_output = lowered_output["out"] + + # lowering error + try: + Tester._assert_outputs_equal( + lowered_output, quantized_output, atol=tolerance, rtol=tolerance + ) + except AssertionError as e: + raise AssertionError( + f"Model '{model_name}' lowering error check failed with tolerance {tolerance}" + ) from e + logging.info( + f"{self._testMethodName} - {model_name} - lowering error passed" + ) + + # verify sqnr between eager model and quantized model + if sqnr_threshold is not None: + original_output = model(*example_inputs) + quantized_output = quantized_model(*example_inputs) + # lowered_output = lowered_module(*example_inputs) + if model_name == "dl3": + original_output = original_output["out"] + quantized_output = quantized_output["out"] + error = compute_error(original_output, quantized_output) + logging.info(f"{self._testMethodName} - {model_name} - SQNR: {error} dB") + self.assertTrue( + error > sqnr_threshold, f"Model '{model_name}' SQNR check failed" + ) + def test_all_models_with_recipes(self) -> None: models_to_test = [ - "linear", - "add", - "add_mul", - "ic3", - "mv2", - "mv3", - "resnet18", - "resnet50", - "vit", - "w2l", - "llama2", + # Tuple format: (model_name, error tolerance, minimum sqnr) + ("linear", 1e-3, 20), + ("add", 1e-3, 20), + ("add_mul", 1e-3, 20), + ("dl3", 1e-3, 20), + ("ic3", None, None), + ("ic4", 1e-3, 20), + ("mv2", 1e-3, None), + ("mv3", 1e-3, None), + ("resnet18", 1e-3, 20), + ("resnet50", 1e-3, 20), + ("vit", 1e-1, 10), + ("w2l", 1e-3, 20), ] - for model_name in models_to_test: - with self.subTest(model=model_name): - self._test_model_with_factory(model_name) - - def test_validate_recipe_kwargs_fp32(self) -> None: - provider = XNNPACKRecipeProvider() - - with self.assertRaises(ValueError) as cm: - provider.create_recipe(XNNPackRecipeType.FP32, invalid_param=123) - - error_msg = str(cm.exception) - self.assertIn("Recipe 'fp32' does not accept any parameters", error_msg) + try: + for model_name, tolerance, sqnr in models_to_test: + with self.subTest(model=model_name): + with torch.no_grad(): + self._test_model_with_factory(model_name, tolerance, sqnr) + finally: + # Clean up dog.jpg file if it exists + if os.path.exists("dog.jpg"): + os.remove("dog.jpg") def test_validate_recipe_kwargs_int4_tensor_with_valid_group_size( self, @@ -224,12 +293,13 @@ def test_validate_recipe_kwargs_int4_tensor_with_valid_group_size( # Should not raise any exception recipe_w_default_group = provider.create_recipe( - XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR + XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR ) self.assertIsNotNone(recipe_w_default_group) recipe = provider.create_recipe( - XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, group_size=64 + XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, + group_size=64, ) self.assertIsNotNone(recipe) @@ -240,7 +310,7 @@ def test_validate_recipe_kwargs_int4_tensor_with_invalid_group_size( with self.assertRaises(ValueError) as cm: provider.create_recipe( - XNNPackRecipeType.INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, + XNNPackRecipeType.TORCHAO_INT8_DYNAMIC_ACT_INT4_WEIGHT_PER_TENSOR, group_size="32", # String instead of int ) diff --git a/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp b/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp index 342e3478e0f..85cac66c62d 100644 --- a/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp +++ b/backends/xnnpack/test/runtime/test_xnn_data_separation.cpp @@ -108,6 +108,21 @@ TEST_F(DataSeparationTest, TestE2E) { "forward", &mmm.get(), nullptr, linear_data_map_.get()); ASSERT_EQ(method.error(), Error::Ok); + // Set a dummy input. + int32_t sizes[1] = {3}; + uint8_t dim_order[1] = {0}; + int32_t strides[1] = {1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 1, + sizes, + nullptr, + dim_order, + strides); + auto input_err = method->set_input( + executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0); + ASSERT_EQ(input_err, Error::Ok); + // Can execute the method. Error err = method->execute(); ASSERT_EQ(err, Error::Ok); diff --git a/backends/xnnpack/test/runtime/test_xnnexecutor.cpp b/backends/xnnpack/test/runtime/test_xnnexecutor.cpp index 4ce1484dc6c..b2a56f6283d 100644 --- a/backends/xnnpack/test/runtime/test_xnnexecutor.cpp +++ b/backends/xnnpack/test/runtime/test_xnnexecutor.cpp @@ -14,6 +14,7 @@ using executorch::backends::xnnpack::delegate::XNNExecutor; using executorch::runtime::Error; using executorch::runtime::EValue; +using executorch::runtime::Span; using executorch::runtime::testing::TensorFactory; TEST(XNNExecutorTest, ArgumentWithTooManyDimensions) { @@ -90,6 +91,7 @@ TEST(XNNExecutorTest, ArgumentWithTooManyDimensions) { EValue input_ev(input_tensor); EValue output_ev(output_tensor); std::array args = {&input_ev, &output_ev}; + Span stack_args(args.data(), 2); // Check for invalid number of dimensions should fail without stack overflow. - EXPECT_EQ(executor.prepare_args(args.data()), Error::InvalidArgument); + EXPECT_EQ(executor.prepare_args(stack_args), Error::InvalidArgument); } diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py new file mode 100644 index 00000000000..8cd9eb92d56 --- /dev/null +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -0,0 +1,84 @@ +# 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 io +import logging +import unittest + +import torch +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.exir import to_edge, to_edge_transform_and_lower +from torch.export import export + + +class TestXnnpackPartitioner(unittest.TestCase): + """Test cases for XnnpackPartitioner functionality and deprecation warnings.""" + + class SimpleModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + return self.linear(x) + + def test_deprecation_warning_for_to_backend_workflow(self): + """ + Test that the deprecated to_edge + to_backend workflow shows a deprecation warning. + """ + model = self.SimpleModel() + x = torch.randn(1, 10) + + exported_model = export(model, (x,)) + + # Capture log output to check for deprecation warning + log_capture_string = io.StringIO() + ch = logging.StreamHandler(log_capture_string) + ch.setLevel(logging.WARNING) + + logger = logging.getLogger( + "executorch.backends.xnnpack.partition.xnnpack_partitioner" + ) + logger.addHandler(ch) + logger.setLevel(logging.WARNING) + + edge = to_edge(exported_model) + partitioner = XnnpackPartitioner() + + edge.to_backend(partitioner) + + log_contents = log_capture_string.getvalue() + self.assertIn("DEPRECATION WARNING", log_contents) + self.assertIn("to_edge() + to_backend()", log_contents) + self.assertIn("to_edge_transform_and_lower()", log_contents) + + def test_no_warning_for_to_edge_transform_and_lower_workflow(self): + """ + Test that the recommended to_edge_transform_and_lower workflow does NOT show a deprecation warning. + """ + + model = self.SimpleModel() + x = torch.randn(1, 10) + + exported_model = export(model, (x,)) + + # Capture log output to check for deprecation warning + log_capture_string = io.StringIO() + ch = logging.StreamHandler(log_capture_string) + ch.setLevel(logging.WARNING) + + logger = logging.getLogger( + "executorch.backends.xnnpack.partition.xnnpack_partitioner" + ) + logger.addHandler(ch) + logger.setLevel(logging.WARNING) + + partitioner = XnnpackPartitioner() + + to_edge_transform_and_lower(exported_model, partitioner=[partitioner]) + + log_contents = log_capture_string.getvalue() + self.assertNotIn("DEPRECATION WARNING", log_contents) diff --git a/backends/xnnpack/test/test_xnnpack_utils.py b/backends/xnnpack/test/test_xnnpack_utils.py index 1f9d8c47723..5a6c529b497 100644 --- a/backends/xnnpack/test/test_xnnpack_utils.py +++ b/backends/xnnpack/test/test_xnnpack_utils.py @@ -70,7 +70,7 @@ _convert_to_reference_decomposed_fx, prepare_fx, ) -from torch.export import export_for_training +from torch.export import export from torch.testing import FileCheck @@ -317,7 +317,7 @@ def quantize_and_test_model_with_quantizer( module.eval() # program capture - m = export_for_training(module, example_inputs, strict=True).module() + m = export(module, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config() diff --git a/backends/xnnpack/third-party/XNNPACK b/backends/xnnpack/third-party/XNNPACK index 52208356940..3131afead79 160000 --- a/backends/xnnpack/third-party/XNNPACK +++ b/backends/xnnpack/third-party/XNNPACK @@ -1 +1 @@ -Subproject commit 52208356940a7c7d3597cf386d500a0f776f7bd0 +Subproject commit 3131afead790c5c69a9aa12273dfc40399789ad7 diff --git a/backends/xnnpack/third-party/cpuinfo b/backends/xnnpack/third-party/cpuinfo index 33ed0be77d7..8a9210069b5 160000 --- a/backends/xnnpack/third-party/cpuinfo +++ b/backends/xnnpack/third-party/cpuinfo @@ -1 +1 @@ -Subproject commit 33ed0be77d7767d0e2010e2c3cf972ef36c7c307 +Subproject commit 8a9210069b5a37dd89ed118a783945502a30a4ae diff --git a/backends/xnnpack/utils/TARGETS b/backends/xnnpack/utils/TARGETS index 55615e11063..d018f294c17 100644 --- a/backends/xnnpack/utils/TARGETS +++ b/backends/xnnpack/utils/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "xnnpack_utils", srcs = glob(["*.py"]), deps = [ diff --git a/codegen/api/unboxing.py b/codegen/api/unboxing.py index d92ee8d557f..4e13246e5b1 100644 --- a/codegen/api/unboxing.py +++ b/codegen/api/unboxing.py @@ -34,7 +34,7 @@ class Unboxing: Takes a sequence of Bindings and unbox EValues to these Bindings. Return generated code that performs correct unboxing. A sample generated code: // aten::mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) - void mul_out(EValue** stack) { + void mul_out(Span stack) { EValue& self = *stack[0]; EValue& other = *stack[1]; EValue& out = *stack[2]; diff --git a/codegen/tools/CMakeLists.txt b/codegen/tools/CMakeLists.txt index 6690418dd6f..489a96aafb6 100644 --- a/codegen/tools/CMakeLists.txt +++ b/codegen/tools/CMakeLists.txt @@ -19,27 +19,16 @@ target_compile_definitions( # Include directories target_include_directories( - selective_build PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/../../.. + selective_build PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../.. ) # Compile options target_compile_options( - selective_build PUBLIC - -Wno-deprecated-declarations - -fPIC - -frtti - -fexceptions + selective_build PUBLIC -Wno-deprecated-declarations -fPIC -frtti -fexceptions ) # Link against required libraries -target_link_libraries( - selective_build PRIVATE - executorch_core - program_schema -) +target_link_libraries(selective_build PRIVATE executorch_core program_schema) # Install the module -install(TARGETS selective_build - LIBRARY DESTINATION executorch/codegen/tools -) +install(TARGETS selective_build LIBRARY DESTINATION executorch/codegen/tools) diff --git a/conftest.py b/conftest.py new file mode 100644 index 00000000000..6c9df86a1ce --- /dev/null +++ b/conftest.py @@ -0,0 +1,12 @@ +import sys + +collect_ignore_glob: list[str] = [] + +# Skip Apple tests on Windows. Note that some Core ML tests can run on Linux, as the AOT flow +# is available. Tests will manage this internally. However, the coremltools import is not available +# on Windows and causes collection to fail. The easiest way to manage this seems to be to just +# skip collection for this subdirectory on unsupported platforms. +if sys.platform == "win32": + collect_ignore_glob += [ + "backends/apple/**", + ] diff --git a/devtools/CMakeLists.txt b/devtools/CMakeLists.txt index 85492075b8c..a267232fe6d 100644 --- a/devtools/CMakeLists.txt +++ b/devtools/CMakeLists.txt @@ -5,7 +5,9 @@ # LICENSE file in the root directory of this source tree. # The include directory that will contain the generated schema headers. -set(DEVTOOLS_INCLUDE_DIR_NO_BUILD_INTERFACE ${CMAKE_BINARY_DIR}/devtools/include) +set(DEVTOOLS_INCLUDE_DIR_NO_BUILD_INTERFACE + ${CMAKE_BINARY_DIR}/devtools/include +) set(DEVTOOLS_INCLUDE_DIR $ ) diff --git a/devtools/TARGETS b/devtools/TARGETS index 06964b83876..2236702fbc7 100644 --- a/devtools/TARGETS +++ b/devtools/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = ["__init__.py"], deps = [ diff --git a/devtools/bundled_program/CMakeLists.txt b/devtools/bundled_program/CMakeLists.txt index 533a92a3e25..e9c5e0e424d 100644 --- a/devtools/bundled_program/CMakeLists.txt +++ b/devtools/bundled_program/CMakeLists.txt @@ -20,7 +20,10 @@ foreach(schema_file ${_schema_files}) ) endforeach() -file(MAKE_DIRECTORY ${DEVTOOLS_INCLUDE_DIR}/executorch/devtools/bundled_program) +file( + MAKE_DIRECTORY + ${DEVTOOLS_INCLUDE_DIR_NO_BUILD_INTERFACE}/executorch/devtools/bundled_program +) add_custom_command( OUTPUT ${_schema_outputs} COMMAND diff --git a/devtools/bundled_program/serialize/__init__.py b/devtools/bundled_program/serialize/__init__.py index 075436e9c11..ceba7670910 100644 --- a/devtools/bundled_program/serialize/__init__.py +++ b/devtools/bundled_program/serialize/__init__.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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. @@ -8,16 +9,15 @@ # TODO(T138924864): Refactor to unify the serialization for bundled program and executorch program. +import importlib.resources as _resources import json import os import tempfile import executorch.devtools.bundled_program.schema as bp_schema -# @manual=fbsource//third-party/pypi/setuptools:setuptools -import pkg_resources +import executorch.devtools.bundled_program.serialize as serialization_package from executorch.devtools.bundled_program.core import BundledProgram - from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile @@ -30,7 +30,7 @@ def write_schema(d: str, schema_name: str) -> None: schema_path = os.path.join(d, "{}.fbs".format(schema_name)) with open(schema_path, "wb") as schema_file: schema_file.write( - pkg_resources.resource_string(__name__, "{}.fbs".format(schema_name)) + _resources.read_binary(serialization_package, f"{schema_name}.fbs") ) diff --git a/devtools/bundled_program/test/test_end2end.py b/devtools/bundled_program/test/test_end2end.py index 7cee073be0e..3268a0df19a 100644 --- a/devtools/bundled_program/test/test_end2end.py +++ b/devtools/bundled_program/test/test_end2end.py @@ -5,21 +5,7 @@ # LICENSE file in the root directory of this source tree. # flake8: noqa: F401 -import functools -import inspect -import os -import random import unittest -from typing import Callable, Dict, Optional, Tuple, Type - -import executorch.exir as exir - -import executorch.exir.control_flow as control_flow - -# @manual=//executorch/extension/pytree:pybindings -import executorch.extension.pytree as pytree - -import torch from executorch.devtools.bundled_program.core import BundledProgram from executorch.devtools.bundled_program.serialize import ( @@ -35,8 +21,6 @@ try: from executorch.extension.pybindings.portable_lib import ( _load_bundled_program_from_buffer, - _load_for_executorch_from_buffer, - _load_for_executorch_from_bundled_program, ) kernel_mode = "lean" @@ -47,8 +31,6 @@ try: from executorch.extension.pybindings.aten_lib import ( # @manual=//executorch/extension/pybindings:aten_lib _load_bundled_program_from_buffer, - _load_for_executorch_from_buffer, - _load_for_executorch_from_bundled_program, ) assert kernel_mode is None @@ -75,19 +57,8 @@ def test_sample_model_e2e(self): bundled_program_buffer ) - executorch_module = _load_for_executorch_from_bundled_program( - executorch_bundled_program - ) - for method_name in eager_model.method_names: - executorch_module.load_bundled_input( - executorch_bundled_program, - method_name, - 0, - ) - executorch_module.plan_execute(method_name) - executorch_module.verify_result_with_bundled_expected_output( - executorch_bundled_program, + executorch_bundled_program.verify_result_with_bundled_expected_output( method_name, 0, ) diff --git a/devtools/bundled_program/util/TARGETS b/devtools/bundled_program/util/TARGETS index 7d019ce30fb..b41e1ffd775 100644 --- a/devtools/bundled_program/util/TARGETS +++ b/devtools/bundled_program/util/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "test_util", srcs = [ "test_util.py", diff --git a/devtools/debug_format/TARGETS b/devtools/debug_format/TARGETS index 46f63fa31b3..2cbf7cf8048 100644 --- a/devtools/debug_format/TARGETS +++ b/devtools/debug_format/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "base_schema", srcs = [ "base_schema.py", @@ -11,7 +11,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "et_schema", srcs = [ "et_schema.py", diff --git a/devtools/debug_format/et_schema.py b/devtools/debug_format/et_schema.py index bb15d70abc4..1a2ae14a09a 100644 --- a/devtools/debug_format/et_schema.py +++ b/devtools/debug_format/et_schema.py @@ -29,6 +29,11 @@ OperatorNode, ValueNode, ) + +from torch._higher_order_ops.auto_functionalize import ( + auto_functionalized, + auto_functionalized_v2, +) from torch._subclasses import FakeTensor @@ -121,6 +126,12 @@ def _parse_args( # noqa: C901 # pyre-ignore named_args = node.target._schema.arguments + if node.op == "call_function" and ( + node.target == auto_functionalized or node.target == auto_functionalized_v2 + ): + # for functioanlized HOPs, args for the corresponding functional op are stored in kwargs + args = tuple(kwargs.values()) + for index, arg in enumerate(args): if isinstance(arg, torch.fx.node.Node): if arg.target == exir.memory.alloc: diff --git a/devtools/etdump/CMakeLists.txt b/devtools/etdump/CMakeLists.txt index 040b100f940..ca4df1d2a82 100644 --- a/devtools/etdump/CMakeLists.txt +++ b/devtools/etdump/CMakeLists.txt @@ -4,24 +4,28 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -set( - _schema_files - etdump_schema_flatcc.fbs - scalar_type.fbs -) +set(_schema_files etdump_schema_flatcc.fbs scalar_type.fbs) set(_schema_outputs) foreach(schema_file ${_schema_files}) list(APPEND _etdump_schema__srcs "${CMAKE_CURRENT_SOURCE_DIR}/${schema_file}") string(REGEX REPLACE "[.]fbs$" "_reader.h" generated_reader "${schema_file}") - list(APPEND _schema_outputs "${DEVTOOLS_INCLUDE_DIR}/executorch/devtools/etdump/${generated_reader}") + list(APPEND _schema_outputs + "${DEVTOOLS_INCLUDE_DIR}/executorch/devtools/etdump/${generated_reader}" + ) - string(REGEX REPLACE "[.]fbs$" "_builder.h" generated_builder "${schema_file}") - list(APPEND _schema_outputs "${DEVTOOLS_INCLUDE_DIR}/executorch/devtools/etdump/${generated_builder}") + string(REGEX REPLACE "[.]fbs$" "_builder.h" generated_builder + "${schema_file}" + ) + list(APPEND _schema_outputs + "${DEVTOOLS_INCLUDE_DIR}/executorch/devtools/etdump/${generated_builder}" + ) endforeach() -file(MAKE_DIRECTORY ${DEVTOOLS_INCLUDE_DIR_NO_BUILD_INTERFACE}/executorch/devtools/etdump) +file(MAKE_DIRECTORY + ${DEVTOOLS_INCLUDE_DIR_NO_BUILD_INTERFACE}/executorch/devtools/etdump +) add_custom_command( OUTPUT ${_schema_outputs} COMMAND @@ -47,16 +51,13 @@ add_library( ) target_link_libraries( etdump - PUBLIC - flatccrt - PRIVATE - executorch + PUBLIC flatccrt + PRIVATE executorch ) target_include_directories( etdump - PUBLIC - ${DEVTOOLS_INCLUDE_DIR} - $ + PUBLIC ${DEVTOOLS_INCLUDE_DIR} + $ ) install( diff --git a/devtools/etdump/TARGETS b/devtools/etdump/TARGETS index 7dcc4c1e84b..b67802e64b3 100644 --- a/devtools/etdump/TARGETS +++ b/devtools/etdump/TARGETS @@ -29,6 +29,7 @@ runtime.python_library( }, visibility = [ "//executorch/devtools/...", + "//executorch/runtime/test/...", ], deps = [ "fbsource//third-party/pypi/setuptools:setuptools", diff --git a/devtools/etdump/serialize.py b/devtools/etdump/serialize.py index 4ed63bc385b..00fd1e70689 100644 --- a/devtools/etdump/serialize.py +++ b/devtools/etdump/serialize.py @@ -1,20 +1,20 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 importlib.resources as _resources import json import os import tempfile -import pkg_resources +import executorch.devtools.etdump as etdump_package from executorch.devtools.etdump.schema_flatcc import ETDumpFlatCC - from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass - from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile # The prefix of schema files used for etdump @@ -25,9 +25,7 @@ def _write_schema(d: str, schema_name: str) -> None: schema_path = os.path.join(d, "{}.fbs".format(schema_name)) with open(schema_path, "wb") as schema_file: - schema_file.write( - pkg_resources.resource_string(__name__, "{}.fbs".format(schema_name)) - ) + schema_file.write(_resources.read_binary(etdump_package, f"{schema_name}.fbs")) def _serialize_from_etdump_to_json(etdump: ETDumpFlatCC) -> str: diff --git a/devtools/etrecord/TARGETS b/devtools/etrecord/TARGETS index 09fc3212bf8..3bb70ef4564 100644 --- a/devtools/etrecord/TARGETS +++ b/devtools/etrecord/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "etrecord", srcs = [ "__init__.py", diff --git a/devtools/etrecord/_etrecord.py b/devtools/etrecord/_etrecord.py index 3b8a71279fd..a144d7e4eaf 100644 --- a/devtools/etrecord/_etrecord.py +++ b/devtools/etrecord/_etrecord.py @@ -52,6 +52,7 @@ class ETRecordReservedFileNames(StrEnum): ET_DIALECT_GRAPH_MODULE = "et_dialect_graph_module" DEBUG_HANDLE_MAP_NAME = "debug_handle_map" DELEGATE_MAP_NAME = "delegate_map" + INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME = "instruction_id_to_num_outs_map" REFERENCE_OUTPUTS = "reference_outputs" REPRESENTATIVE_INPUTS = "representative_inputs" @@ -67,29 +68,59 @@ def __init__( _delegate_map: Optional[ Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]] ] = None, + _instruction_id_to_num_outs_map: Optional[ + Dict[str, Dict[int, Union[int, List[int]]]] + ] = None, _reference_outputs: Optional[Dict[str, List[ProgramOutput]]] = None, - _representative_inputs: Optional[List[ProgramOutput]] = None, + _representative_inputs: Optional[List[ProgramInput]] = None, ): + """ + Please do not construct an ETRecord object directly. + + If you want to create an ETRecord for logging AOT information to further analysis, please mark `generate_etrecord` + as True in your export api, and get the ETRecord object from the `ExecutorchProgramManager`. + For exmaple: + ```python + exported_program = torch.export.export(model, inputs) + edge_program = to_edge_transform_and_lower(exported_program, generate_etrecord=True) + executorch_program = edge_program.to_executorch() + etrecord = executorch_program.get_etrecord() + ``` + + If user need to create an ETRecord manually, please use the `create_etrecord` function. + """ + self.exported_program = exported_program self.export_graph_id = export_graph_id self.edge_dialect_program = edge_dialect_program self.graph_map = graph_map self._debug_handle_map = _debug_handle_map self._delegate_map = _delegate_map + self._instruction_id_to_num_outs_map = _instruction_id_to_num_outs_map self._reference_outputs = _reference_outputs self._representative_inputs = _representative_inputs def save(self, path: Union[str, os.PathLike, BinaryIO, IO[bytes]]) -> None: """ - Serialize and save the ETRecord to the specified path. + Serialize and save the ETRecord to the specified path for use in Inspector. The ETRecord + should contains at least edge dialect program and executorch program information for further + analysis, otherwise it will raise an exception. Args: path: Path where the ETRecord file will be saved to. + + Raises: + RuntimeError: If the ETRecord does not contain essential information for Inpector. """ if isinstance(path, (str, os.PathLike)): # pyre-ignore[6]: In call `os.fspath`, for 1st positional argument, expected `str` but got `Union[PathLike[typing.Any], str]` path = os.fspath(path) + if not (self.edge_dialect_program and self._debug_handle_map): + raise RuntimeError( + "ETRecord must contain edge dialect program and executorch program to be saved" + ) + etrecord_zip = ZipFile(path, "w") try: @@ -146,6 +177,12 @@ def _save_metadata(self, etrecord_zip: ZipFile) -> None: json.dumps(self._delegate_map), ) + if self._instruction_id_to_num_outs_map is not None: + etrecord_zip.writestr( + ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME, + json.dumps(self._instruction_id_to_num_outs_map), + ) + if self._reference_outputs is not None: etrecord_zip.writestr( ETRecordReservedFileNames.REFERENCE_OUTPUTS, @@ -258,6 +295,7 @@ def add_executorch_program( if ( self._debug_handle_map is not None or self._delegate_map is not None + or self._instruction_id_to_num_outs_map is not None or self._reference_outputs is not None or self._representative_inputs is not None ): @@ -267,13 +305,18 @@ def add_executorch_program( ) # Process executorch program and extract data - debug_handle_map, delegate_map, reference_outputs, representative_inputs = ( - _process_executorch_program(executorch_program) - ) + ( + debug_handle_map, + delegate_map, + instruction_id_to_num_outs_map, + reference_outputs, + representative_inputs, + ) = _process_executorch_program(executorch_program) # Set the extracted data self._debug_handle_map = debug_handle_map self._delegate_map = delegate_map + self._instruction_id_to_num_outs_map = instruction_id_to_num_outs_map self._reference_outputs = reference_outputs self._representative_inputs = representative_inputs @@ -345,6 +388,56 @@ def add_edge_dialect_program( # Set the extracted data self.edge_dialect_program = processed_edge_dialect_program + def update_representative_inputs( + self, + representative_inputs: Union[List[ProgramInput], BundledProgram], + ) -> None: + """ + Update the representative inputs in the ETRecord. + + This method allows users to customize the representative inputs that will be + included when the ETRecord is saved. The representative inputs can be provided + directly as a list or extracted from a BundledProgram. + + Args: + representative_inputs: Either a list of ProgramInput objects or a BundledProgram + from which representative inputs will be extracted. + """ + if isinstance(representative_inputs, BundledProgram): + self._representative_inputs = _get_representative_inputs( + representative_inputs + ) + else: + self._representative_inputs = representative_inputs + + def update_reference_outputs( + self, + reference_outputs: Union[ + Dict[str, List[ProgramOutput]], List[ProgramOutput], BundledProgram + ], + ) -> None: + """ + Update the reference outputs in the ETRecord. + + This method allows users to customize the reference outputs that will be + included when the ETRecord is saved. The reference outputs can be provided + directly as a dictionary mapping method names to lists of outputs, as a + single list of outputs (which will be treated as {"forward": List[ProgramOutput]}), + or extracted from a BundledProgram. + + Args: + reference_outputs: Either a dictionary mapping method names to lists of + ProgramOutput objects, a single list of ProgramOutput objects (treated + as outputs for the "forward" method), or a BundledProgram from which + reference outputs will be extracted. + """ + if isinstance(reference_outputs, BundledProgram): + self._reference_outputs = _get_reference_outputs(reference_outputs) + elif isinstance(reference_outputs, list): + self._reference_outputs = {"forward": reference_outputs} + else: + self._reference_outputs = reference_outputs + def _get_reference_outputs( bundled_program: BundledProgram, @@ -517,7 +610,9 @@ def _process_executorch_program( executorch_program: Union[ ExecutorchProgram, ExecutorchProgramManager, BundledProgram ] -) -> tuple[Optional[Dict], Optional[Dict], Optional[Dict], Optional[List]]: +) -> tuple[ + Optional[Dict], Optional[Dict], Optional[Dict], Optional[Dict], Optional[List] +]: """Process executorch program and return debug maps and bundled program data.""" if isinstance(executorch_program, BundledProgram): reference_outputs = _get_reference_outputs(executorch_program) @@ -526,11 +621,30 @@ def _process_executorch_program( debug_handle_map = executorch_program.executorch_program.debug_handle_map # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `debug_handle_map` delegate_map = executorch_program.executorch_program.delegate_map - return debug_handle_map, delegate_map, reference_outputs, representative_inputs + # pyre-ignore[16]: Item `None` of `typing.Union[None, exir.program._program.ExecutorchProgram, exir.program._program.ExecutorchProgramManager]` has no attribute `instruction_id_to_num_outs_map` + instruction_id_to_num_outs_map = ( + executorch_program.executorch_program.instruction_id_to_num_outs_map + ) + return ( + debug_handle_map, + delegate_map, + instruction_id_to_num_outs_map, + reference_outputs, + representative_inputs, + ) else: debug_handle_map = executorch_program.debug_handle_map delegate_map = executorch_program.delegate_map - return debug_handle_map, delegate_map, None, None + instruction_id_to_num_outs_map = ( + executorch_program.instruction_id_to_num_outs_map + ) + return ( + debug_handle_map, + delegate_map, + instruction_id_to_num_outs_map, + None, + None, + ) def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 @@ -564,6 +678,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 graph_map: Dict[str, ExportedProgram] = {} debug_handle_map = None delegate_map = None + instruction_id_to_num_outs_map = None exported_program = None edge_dialect_program = None reference_outputs = None @@ -583,6 +698,12 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 delegate_map = json.loads( etrecord_zip.read(ETRecordReservedFileNames.DELEGATE_MAP_NAME) ) + elif entry == ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME: + instruction_id_to_num_outs_map = json.loads( + etrecord_zip.read( + ETRecordReservedFileNames.INSTRUCTION_ID_TO_NUM_OUTS_MAP_NAME + ) + ) elif entry == ETRecordReservedFileNames.ETRECORD_IDENTIFIER: continue elif entry == ETRecordReservedFileNames.EDGE_DIALECT_EXPORTED_PROGRAM: @@ -648,6 +769,7 @@ def parse_etrecord(etrecord_path: str) -> ETRecord: # noqa: C901 graph_map=graph_map, _debug_handle_map=debug_handle_map, _delegate_map=delegate_map, + _instruction_id_to_num_outs_map=instruction_id_to_num_outs_map, _reference_outputs=reference_outputs, _representative_inputs=representative_inputs, export_graph_id=export_graph_id, diff --git a/devtools/etrecord/tests/TARGETS b/devtools/etrecord/tests/TARGETS index fffa7f18341..f25f0464c9e 100644 --- a/devtools/etrecord/tests/TARGETS +++ b/devtools/etrecord/tests/TARGETS @@ -1,4 +1,4 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") @@ -7,16 +7,11 @@ python_unittest( name = "etrecord_test", srcs = ["etrecord_test.py"], deps = [ - "//caffe2:torch", - "//executorch/devtools/bundled_program:config", - "//executorch/devtools/bundled_program:core", - "//executorch/devtools/etrecord:etrecord", - "//executorch/exir:lib", - "//executorch/exir/tests:models", + ":etrecord_test_library" ], ) -python_library( +runtime.python_library( name = "etrecord_test_library", srcs = ["etrecord_test.py"], deps = [ @@ -26,5 +21,7 @@ python_library( "//executorch/devtools/etrecord:etrecord", "//executorch/exir:lib", "//executorch/exir/tests:models", + "//executorch/backends/xnnpack/partition:xnnpack_partitioner", + "//executorch/export:lib", ], ) diff --git a/devtools/etrecord/tests/etrecord_test.py b/devtools/etrecord/tests/etrecord_test.py index 25ea5a25e1f..535de2e9a56 100644 --- a/devtools/etrecord/tests/etrecord_test.py +++ b/devtools/etrecord/tests/etrecord_test.py @@ -10,10 +10,12 @@ import json import tempfile import unittest +from typing import List import executorch.exir.tests.models as models import torch from executorch import exir +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite from executorch.devtools.bundled_program.core import BundledProgram from executorch.devtools.etrecord import generate_etrecord, parse_etrecord @@ -23,13 +25,51 @@ ETRecord, ETRecordReservedFileNames, ) -from executorch.exir import EdgeCompileConfig, EdgeProgramManager, to_edge -from executorch.exir.program._program import to_edge_transform_and_lower +from executorch.exir import EdgeCompileConfig, EdgeProgramManager +from executorch.exir.program._program import to_edge, to_edge_transform_and_lower + +from executorch.export import export as etexport, ExportRecipe, StageType from torch.export import export # TODO : T154728484 Add test cases to cover multiple entry points class TestETRecord(unittest.TestCase): + def assert_representative_inputs_equal( + self, + expected_inputs: List, + actual_inputs: List, + msg: str = "Representative inputs do not match", + ) -> None: + """ + Utility function to compare representative inputs. + + This function handles the comparison of representative inputs, which are lists of tuples + containing tensors. It compares each input tuple element by element using torch.equal(). + + Args: + expected_inputs: List of expected input tuples + actual_inputs: List of actual input tuples + msg: Optional message to display on assertion failure + """ + self.assertEqual( + len(expected_inputs), + len(actual_inputs), + f"{msg}: Different number of input sets", + ) + + for i, (expected, actual) in enumerate(zip(expected_inputs, actual_inputs)): + self.assertEqual( + len(expected), + len(actual), + f"{msg}: Input set {i} has different number of tensors", + ) + + for j, (exp_tensor, act_tensor) in enumerate(zip(expected, actual)): + self.assertTrue( + torch.equal(exp_tensor, act_tensor), + f"{msg}: Tensor {j} in input set {i} does not match", + ) + def assert_etrecord_has_no_exported_program(self, etrecord: ETRecord) -> None: """Assert that ETRecord has no exported program data.""" self.assertIsNone(etrecord.exported_program) @@ -68,17 +108,16 @@ def assert_etrecord_saveable(self, etrecord: ETRecord) -> None: self.assertIsNotNone(etrecord._debug_handle_map) self.assertIsNotNone(etrecord._delegate_map) - def get_test_model(self): + def get_test_model(self, generate_etrecord=False): f = models.BasicSinMax() - captured_output = exir.capture(f, f.get_random_inputs(), exir.CaptureConfig()) - captured_output_copy = copy.deepcopy(captured_output) - edge_output = captured_output.to_edge( - # TODO(gasoon): Remove _use_edge_ops=False once serde is fully migrated to Edge ops - exir.EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=False) + aten_dialect = export(f, f.get_random_inputs(), strict=True) + edge_program: EdgeProgramManager = to_edge( + aten_dialect, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + generate_etrecord=generate_etrecord, ) - edge_output_copy = copy.deepcopy(edge_output) - et_output = edge_output.to_executorch() - return (captured_output_copy, edge_output_copy, et_output) + edge_program_copy = copy.deepcopy(edge_program) + return (aten_dialect, edge_program_copy, edge_program.to_executorch()) def get_test_model_with_bundled_program(self): f = models.BasicSinMax() @@ -96,26 +135,36 @@ def get_test_model_with_bundled_program(self): ], ) ] - captured_output = exir.capture(f, inputs[0], exir.CaptureConfig()) - captured_output_copy = copy.deepcopy(captured_output) - edge_output = captured_output.to_edge( - # TODO(gasoon): Remove _use_edge_ops=False once serde is fully migrated to Edge ops - exir.EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=False) - ) - edge_output_copy = copy.deepcopy(edge_output) - et_output = edge_output.to_executorch() - + aten_dialect, edge_program_copy, et_output = self.get_test_model() bundled_program = BundledProgram(et_output, method_test_suites) - return (captured_output_copy, edge_output_copy, bundled_program) + return (aten_dialect, edge_program_copy, bundled_program) - def get_test_model_with_manager(self): + def get_test_export_session(self, generate_etrecord=False, to_edge_flow=False): f = models.BasicSinMax() - aten_dialect = export(f, f.get_random_inputs(), strict=True) - edge_program: EdgeProgramManager = to_edge( - aten_dialect, compile_config=EdgeCompileConfig(_check_ir_validity=False) + example_inputs = [f.get_random_inputs()] + export_recipe = None + + if to_edge_flow: + export_recipe = ExportRecipe( + pipeline_stages=[ + StageType.TORCH_EXPORT, + StageType.TO_EDGE, + StageType.TO_BACKEND, + StageType.TO_EXECUTORCH, + ] + ) + else: + export_recipe = ExportRecipe() + + # Test with generate_etrecord=True + export_session = etexport( + model=f, + example_inputs=example_inputs, + export_recipe=export_recipe, + generate_etrecord=generate_etrecord, ) - edge_program_copy = copy.deepcopy(edge_program) - return (aten_dialect, edge_program_copy, edge_program.to_executorch()) + + return export_session # Serialized and deserialized graph modules are not completely the same, so we check # that they are close enough and match especially on the parameters we care about in the Developer Tools. @@ -160,16 +209,20 @@ def test_etrecord_generation(self): self.check_graph_closeness( etrecord.graph_map["aten_dialect_output/forward"], - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.check_graph_closeness( etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) self.assertEqual( etrecord._debug_handle_map, json.loads(json.dumps(et_output.debug_handle_map)), ) + self.assertEqual( + etrecord._instruction_id_to_num_outs_map, + json.loads(json.dumps(et_output.instruction_id_to_num_outs_map)), + ) def test_etrecord_generation_with_bundled_program(self): ( @@ -209,25 +262,6 @@ def test_etrecord_generation_with_bundled_program(self): ) ) - def test_etrecord_generation_with_manager(self): - captured_output, edge_output, et_output = self.get_test_model_with_manager() - with tempfile.TemporaryDirectory() as tmpdirname: - generate_etrecord( - tmpdirname + "/etrecord.bin", - edge_output, - et_output, - ) - - etrecord = parse_etrecord(tmpdirname + "/etrecord.bin") - self.check_graph_closeness( - etrecord.edge_dialect_program, - edge_output.exported_program().graph_module, - ) - self.assertEqual( - etrecord._debug_handle_map, - json.loads(json.dumps(et_output.debug_handle_map)), - ) - def test_etrecord_invalid_input(self): captured_output, edge_output, et_output = self.get_test_model() with tempfile.TemporaryDirectory() as tmpdirname: @@ -249,14 +283,14 @@ def test_etrecord_reserved_name(self): edge_output, et_output, extra_recorded_export_modules={ - reserved_name: captured_output.exported_program.graph_module + reserved_name: captured_output.graph_module }, ) def test_etrecord_generation_with_exported_program(self): """Test that exported program can be recorded and parsed back correctly.""" captured_output, edge_output, et_output = self.get_test_model() - original_exported_program = captured_output.exported_program + original_exported_program = captured_output expected_graph_id = id(original_exported_program.graph) with tempfile.TemporaryDirectory() as tmpdirname: @@ -281,7 +315,7 @@ def test_etrecord_generation_with_exported_program(self): # Validate other components are still present self.check_graph_closeness( etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) self.assertEqual( etrecord._debug_handle_map, @@ -371,6 +405,41 @@ def test_get_etrecord_from_executorch_program_manager(self): self.assertEqual(etrecord._debug_handle_map, et_manager.debug_handle_map) self.assertEqual(etrecord._delegate_map, et_manager.delegate_map) + def test_get_etrecord_from_executorch_program_manager_with_partitioner(self): + """Test getting ETRecord from ExecutorchProgramManager using get_etrecord() method.""" + f = models.BasicSinMax() + aten_program = export(f, f.get_random_inputs(), strict=True) + + # Generate edge manager with ETRecord + edge_manager = to_edge_transform_and_lower( + aten_program, + partitioner=[XnnpackPartitioner()], + generate_etrecord=True, + ) + + # Convert to executorch + et_manager = edge_manager.to_executorch() + + # Test get_etrecord method + etrecord = et_manager.get_etrecord() + self.assertIsNotNone(etrecord) + self.assert_etrecord_saveable(etrecord) + + # Verify the data matches the original input + self.check_graph_closeness( + etrecord.exported_program, + aten_program.graph_module, + ) + self.assertEqual( + etrecord.export_graph_id, + id(aten_program.graph), + ) + + # Verify the executorch program data matches + # ETRecord stores data directly (not JSON serialized), so compare with original data + self.assertEqual(etrecord._debug_handle_map, et_manager.debug_handle_map) + self.assertEqual(etrecord._delegate_map, et_manager.delegate_map) + def test_get_etrecord_from_executorch_program_manager_without_generation(self): """Test getting ETRecord from ExecutorchProgramManager when ETRecord was not generated.""" f = models.BasicSinMax() @@ -394,6 +463,82 @@ def test_get_etrecord_from_executorch_program_manager_without_generation(self): self.assertIn("ETRecord was not generated", str(context.exception)) + def test_to_edge_with_etrecord_generation(self): + """Test that to_edge generates ETRecord correctly.""" + aten_program, edge_manager, _ = self.get_test_model(generate_etrecord=True) + + # Verify that ETRecord was generated and attached + self.assertIsNotNone(edge_manager._etrecord) + etrecord = edge_manager._etrecord + self.assert_legal_etrecord_in_edge_program(etrecord) + + # Verify the exported program matches the input + self.check_graph_closeness( + etrecord.exported_program, + aten_program.graph_module, + ) + self.assertEqual( + etrecord.export_graph_id, + id(aten_program.graph), + ) + + # Verify the edge dialect program matches the edge manager + self.check_graph_closeness( + etrecord.edge_dialect_program, + edge_manager.exported_program().graph_module, + ) + + def test_to_edge_without_etrecord_generation(self): + """Test that to_edge works correctly without ETRecord generation.""" + # Test with generate_etrecord=False (default) + _, edge_manager, et_manager = self.get_test_model() + + # Verify that no ETRecord was generated + self.assertIsNone(edge_manager._etrecord) + + # Test get_etrecord method should raise RuntimeError + with self.assertRaises(RuntimeError): + et_manager.get_etrecord() + + def test_to_edge_etrecord_save_and_parse(self): + """Test that ETRecord generated by to_edge can be saved and parsed.""" + aten_program, _, et_manager = self.get_test_model(generate_etrecord=True) + + etrecord = et_manager.get_etrecord() + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_to_edge.bin" + + etrecord.save(etrecord_path) + + # Parse ETRecord back and verify + parsed_etrecord = parse_etrecord(etrecord_path) + + # Validate that all components are preserved + # Note: Skip graph structure comparison due to transformation differences + self.check_graph_closeness( + etrecord.exported_program, parsed_etrecord.exported_program + ) + self.check_graph_closeness( + etrecord.edge_dialect_program, parsed_etrecord.edge_dialect_program + ) + + # Validate executorch program data + self.assertEqual( + parsed_etrecord._debug_handle_map, + json.loads(json.dumps(et_manager.debug_handle_map)), + ) + self.assertEqual( + parsed_etrecord._delegate_map, + json.loads(json.dumps(et_manager.delegate_map)), + ) + + # Validate export graph id + self.assertEqual( + parsed_etrecord.export_graph_id, + id(aten_program.graph), + ) + def test_to_edge_transform_and_lower_etrecord_save_and_parse(self): """Test that ETRecord generated by to_edge_transform_and_lower can be saved and parsed.""" f = models.BasicSinMax() @@ -402,6 +547,7 @@ def test_to_edge_transform_and_lower_etrecord_save_and_parse(self): # Generate edge manager with ETRecord edge_manager = to_edge_transform_and_lower( aten_program, + partitioner=[XnnpackPartitioner()], generate_etrecord=True, ) @@ -447,13 +593,11 @@ def test_add_extra_export_modules(self): captured_output, edge_output, et_output = self.get_test_model() # Create an ETRecord instance with existing graph_map - initial_graph_map = { - "existing_module/forward": captured_output.exported_program - } + initial_graph_map = {"existing_module/forward": captured_output} etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), graph_map=initial_graph_map, _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, @@ -483,7 +627,7 @@ def test_add_extra_export_modules(self): # Verify the modules are correctly stored self.check_graph_closeness( etrecord.graph_map["existing_module/forward"], - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.check_graph_closeness( etrecord.graph_map["new_module/forward"], @@ -495,9 +639,9 @@ def test_add_extra_export_modules_reserved_name_validation(self): captured_output, edge_output, et_output = self.get_test_model() etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -505,21 +649,19 @@ def test_add_extra_export_modules_reserved_name_validation(self): # Test that reserved names are rejected for reserved_name in ETRecordReservedFileNames: with self.assertRaises(RuntimeError): - etrecord.add_extra_export_modules( - {reserved_name: captured_output.exported_program} - ) + etrecord.add_extra_export_modules({reserved_name: captured_output}) def test_etrecord_class_constructor_and_save(self): """Test that ETRecord class constructor and save method work correctly.""" captured_output, edge_output, et_output = self.get_test_model() - original_exported_program = captured_output.exported_program + original_exported_program = captured_output expected_graph_id = id(original_exported_program.graph) # Create ETRecord instance directly using constructor etrecord = ETRecord( exported_program=original_exported_program, export_graph_id=expected_graph_id, - edge_dialect_program=edge_output.exported_program, + edge_dialect_program=edge_output.exported_program(), graph_map={"test_module/forward": original_exported_program}, _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, @@ -544,7 +686,7 @@ def test_etrecord_class_constructor_and_save(self): self.assertIsNotNone(parsed_etrecord.edge_dialect_program) self.check_graph_closeness( parsed_etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) # Validate graph map @@ -582,9 +724,9 @@ def test_etrecord_class_with_bundled_program_data(self): # Create ETRecord instance with bundled program data etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=bundled_program.executorch_program.debug_handle_map, _delegate_map=bundled_program.executorch_program.delegate_map, _reference_outputs=reference_outputs, @@ -628,7 +770,7 @@ def test_etrecord_class_with_bundled_program_data(self): def test_etrecord_generation_with_exported_program_dict(self): """Test that exported program dictionary can be recorded and parsed back correctly.""" captured_output, edge_output, et_output = self.get_test_model() - original_exported_program = captured_output.exported_program + original_exported_program = captured_output exported_program_dict = {"forward": original_exported_program} expected_graph_id = id(original_exported_program.graph) @@ -654,7 +796,7 @@ def test_etrecord_generation_with_exported_program_dict(self): # Validate other components are still present self.check_graph_closeness( etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) self.assertEqual( etrecord._debug_handle_map, @@ -670,9 +812,9 @@ def test_add_executorch_program(self): # Create an ETRecord instance without executorch program data etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), ) # Verify initial state - no executorch program data @@ -684,14 +826,8 @@ def test_add_executorch_program(self): # Verify executorch program data is now present self.assertIsNotNone(etrecord._debug_handle_map) self.assertIsNotNone(etrecord._delegate_map) - self.assertEqual( - etrecord._debug_handle_map, - json.loads(json.dumps(et_output.debug_handle_map)), - ) - self.assertEqual( - etrecord._delegate_map, - json.loads(json.dumps(et_output.delegate_map)), - ) + self.assertEqual(etrecord._debug_handle_map, et_output.debug_handle_map) + self.assertEqual(etrecord._delegate_map, et_output.delegate_map) # For regular ExecutorchProgram, reference_outputs and representative_inputs should be None self.assertIsNone(etrecord._reference_outputs) self.assertIsNone(etrecord._representative_inputs) @@ -706,9 +842,9 @@ def test_add_executorch_program_with_bundled_program(self): # Create an ETRecord instance without executorch program data etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), ) # Verify initial state - no executorch program data @@ -757,9 +893,9 @@ def test_add_executorch_program_already_exists_exception(self): # Create an ETRecord instance with existing executorch program data etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -779,9 +915,9 @@ def test_add_executorch_program_partial_data_exists_exception(self): # Create an ETRecord instance with only debug_handle_map (partial data) etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, ) @@ -800,9 +936,9 @@ def test_add_executorch_program_and_save(self): # Create an ETRecord instance without executorch program data etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), ) # Add executorch program @@ -821,13 +957,13 @@ def test_add_executorch_program_and_save(self): self.assertIsNotNone(parsed_etrecord.exported_program) self.check_graph_closeness( parsed_etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertIsNotNone(parsed_etrecord.edge_dialect_program) self.check_graph_closeness( parsed_etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) # Validate executorch program data @@ -843,7 +979,7 @@ def test_add_executorch_program_and_save(self): # Validate export graph id self.assertEqual( parsed_etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) def test_add_exported_program(self): @@ -852,7 +988,7 @@ def test_add_exported_program(self): # Create an ETRecord instance without exported program etrecord = ETRecord( - edge_dialect_program=edge_output.exported_program, + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -861,18 +997,18 @@ def test_add_exported_program(self): self.assert_etrecord_has_no_exported_program(etrecord) # Add exported program - etrecord.add_exported_program(captured_output.exported_program) + etrecord.add_exported_program(captured_output) # Verify exported program is now present self.assertIsNotNone(etrecord.exported_program) self.assertIsNotNone(etrecord.export_graph_id) self.check_graph_closeness( etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertEqual( etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) def test_add_exported_program_with_dict(self): @@ -881,7 +1017,7 @@ def test_add_exported_program_with_dict(self): # Create an ETRecord instance without exported program etrecord = ETRecord( - edge_dialect_program=edge_output.exported_program, + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -891,7 +1027,7 @@ def test_add_exported_program_with_dict(self): self.assertIsNone(etrecord.export_graph_id) # Add exported program as dictionary - exported_program_dict = {"forward": captured_output.exported_program} + exported_program_dict = {"forward": captured_output} etrecord.add_exported_program(exported_program_dict) # Verify exported program is now present @@ -899,11 +1035,11 @@ def test_add_exported_program_with_dict(self): self.assertIsNotNone(etrecord.export_graph_id) self.check_graph_closeness( etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertEqual( etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) def test_add_exported_program_already_exists_exception(self): @@ -912,9 +1048,9 @@ def test_add_exported_program_already_exists_exception(self): # Create an ETRecord instance with existing exported program etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -940,15 +1076,15 @@ def test_add_exported_program_partial_data_exists_exception(self): # Create an ETRecord instance with only export_graph_id (partial data) etrecord = ETRecord( - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) # Verify that adding exported program raises RuntimeError even with partial data with self.assertRaises(RuntimeError) as context: - etrecord.add_exported_program(captured_output.exported_program) + etrecord.add_exported_program(captured_output) self.assertIn( "Exported program already exists in the ETRecord", @@ -961,7 +1097,7 @@ def test_add_exported_program_with_none(self): # Create an ETRecord instance without exported program etrecord = ETRecord( - edge_dialect_program=edge_output.exported_program, + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -981,13 +1117,13 @@ def test_add_exported_program_and_save(self): # Create an ETRecord instance without exported program etrecord = ETRecord( - edge_dialect_program=edge_output.exported_program, + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) # Add exported program - etrecord.add_exported_program(captured_output.exported_program) + etrecord.add_exported_program(captured_output) with tempfile.TemporaryDirectory() as tmpdirname: etrecord_path = tmpdirname + "/etrecord_with_added_exported_program.bin" @@ -1002,19 +1138,19 @@ def test_add_exported_program_and_save(self): self.assertIsNotNone(parsed_etrecord.exported_program) self.check_graph_closeness( parsed_etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertIsNotNone(parsed_etrecord.edge_dialect_program) self.check_graph_closeness( parsed_etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) # Validate export graph id self.assertEqual( parsed_etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) def test_add_edge_dialect_program(self): @@ -1023,8 +1159,8 @@ def test_add_edge_dialect_program(self): # Create an ETRecord instance without edge dialect program etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), + exported_program=captured_output, + export_graph_id=id(captured_output.graph), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -1039,37 +1175,7 @@ def test_add_edge_dialect_program(self): self.assertIsNotNone(etrecord.edge_dialect_program) self.check_graph_closeness( etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, - ) - - def test_add_edge_dialect_program_with_exir_exported_program(self): - """Test add_edge_dialect_program with ExirExportedProgram.""" - captured_output, edge_output, et_output = self.get_test_model() - - # Create an ETRecord instance without edge dialect program - etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - _debug_handle_map=et_output.debug_handle_map, - _delegate_map=et_output.delegate_map, - ) - - # Verify initial state - no edge dialect program - self.assertIsNone(etrecord.edge_dialect_program) - - # Create ExirExportedProgram from captured output - exir_exported_program = captured_output.to_edge( - exir.EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=False) - ) - - # Add edge dialect program using ExirExportedProgram - etrecord.add_edge_dialect_program(exir_exported_program) - - # Verify edge dialect program is now present - self.assertIsNotNone(etrecord.edge_dialect_program) - self.check_graph_closeness( - etrecord.edge_dialect_program, - exir_exported_program.exported_program.graph_module, + edge_output.exported_program().graph_module, ) def test_add_edge_dialect_program_already_exists_exception(self): @@ -1078,9 +1184,9 @@ def test_add_edge_dialect_program_already_exists_exception(self): # Create an ETRecord instance with existing edge dialect program etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), - edge_dialect_program=edge_output.exported_program, + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -1109,8 +1215,8 @@ def test_add_edge_dialect_program_and_save(self): # Create an ETRecord instance without edge dialect program etrecord = ETRecord( - exported_program=captured_output.exported_program, - export_graph_id=id(captured_output.exported_program.graph), + exported_program=captured_output, + export_graph_id=id(captured_output.graph), _debug_handle_map=et_output.debug_handle_map, _delegate_map=et_output.delegate_map, ) @@ -1131,19 +1237,19 @@ def test_add_edge_dialect_program_and_save(self): self.assertIsNotNone(parsed_etrecord.exported_program) self.check_graph_closeness( parsed_etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertIsNotNone(parsed_etrecord.edge_dialect_program) self.check_graph_closeness( parsed_etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) # Validate export graph id self.assertEqual( parsed_etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) def test_add_all_programs_sequentially(self): @@ -1157,7 +1263,7 @@ def test_add_all_programs_sequentially(self): self.assert_etrecord_is_empty(etrecord) # Add exported program - etrecord.add_exported_program(captured_output.exported_program) + etrecord.add_exported_program(captured_output) # Add edge dialect program etrecord.add_edge_dialect_program(edge_output) @@ -1175,24 +1281,18 @@ def test_add_all_programs_sequentially(self): # Verify the data matches expected values self.check_graph_closeness( etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.check_graph_closeness( etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) self.assertEqual( etrecord.export_graph_id, - id(captured_output.exported_program.graph), - ) - self.assertEqual( - etrecord._debug_handle_map, - json.loads(json.dumps(et_output.debug_handle_map)), - ) - self.assertEqual( - etrecord._delegate_map, - json.loads(json.dumps(et_output.delegate_map)), + id(captured_output.graph), ) + self.assertEqual(etrecord._debug_handle_map, et_output.debug_handle_map) + self.assertEqual(etrecord._delegate_map, et_output.delegate_map) # Test that the complete ETRecord can be saved and parsed with tempfile.TemporaryDirectory() as tmpdirname: @@ -1208,19 +1308,19 @@ def test_add_all_programs_sequentially(self): self.assertIsNotNone(parsed_etrecord.exported_program) self.check_graph_closeness( parsed_etrecord.exported_program, - captured_output.exported_program.graph_module, + captured_output.graph_module, ) self.assertIsNotNone(parsed_etrecord.edge_dialect_program) self.check_graph_closeness( parsed_etrecord.edge_dialect_program, - edge_output.exported_program.graph_module, + edge_output.exported_program().graph_module, ) # Validate all metadata self.assertEqual( parsed_etrecord.export_graph_id, - id(captured_output.exported_program.graph), + id(captured_output.graph), ) self.assertEqual( parsed_etrecord._debug_handle_map, @@ -1230,3 +1330,419 @@ def test_add_all_programs_sequentially(self): parsed_etrecord._delegate_map, json.loads(json.dumps(et_output.delegate_map)), ) + + def test_executorch_export_with_etrecord_generation(self): + """Test that executorch.export generates ETRecord correctly when generate_etrecord=True.""" + # Verify that ETRecord was generated and can be retrieved + export_session = self.get_test_export_session(generate_etrecord=True) + etrecord = export_session.get_etrecord() + self.assertIsNotNone(etrecord) + self.assert_etrecord_saveable(etrecord) + + # Verify the executorch program data matches + et_manager = export_session.get_executorch_program_manager() + self.assertEqual(etrecord._debug_handle_map, et_manager.debug_handle_map) + self.assertEqual(etrecord._delegate_map, et_manager.delegate_map) + + def test_executorch_export_without_etrecord_generation(self): + """Test that executorch.export works correctly without ETRecord generation.""" + # Test with generate_etrecord=False (default) + export_session = self.get_test_export_session(generate_etrecord=False) + + # Verify that no ETRecord was generated + with self.assertRaises(RuntimeError) as context: + export_session.get_etrecord() + + self.assertIn("ETRecord was not generated", str(context.exception)) + + # Verify that the export session still works correctly + self.assertIsNotNone(export_session.get_executorch_program_manager()) + self.assertTrue(len(export_session.get_pte_buffer()) > 0) + + def test_executorch_export_etrecord_save_and_parse(self): + """Test that ETRecord generated by executorch.export can be saved and parsed.""" + export_session = self.get_test_export_session(generate_etrecord=True) + + etrecord = export_session.get_etrecord() + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_export.bin" + + etrecord.save(etrecord_path) + + # Parse ETRecord back and verify + parsed_etrecord = parse_etrecord(etrecord_path) + + # Validate that all components are preserved + self.assertIsNotNone(parsed_etrecord.exported_program) + self.assertIsNotNone(parsed_etrecord.edge_dialect_program) + + # Validate executorch program data + et_manager = export_session.get_executorch_program_manager() + self.assertEqual( + parsed_etrecord._debug_handle_map, + json.loads(json.dumps(et_manager.debug_handle_map)), + ) + self.assertEqual( + parsed_etrecord._delegate_map, + json.loads(json.dumps(et_manager.delegate_map)), + ) + + # Validate export graph id is preserved + self.assertIsNotNone(parsed_etrecord.export_graph_id) + + def test_executorch_export_with_to_edge_flow(self): + """Test executorch.export with TO_EDGE flow and ETRecord generation.""" + export_session = self.get_test_export_session( + generate_etrecord=True, + to_edge_flow=True, + ) + + # Verify that ETRecord was generated + etrecord = export_session.get_etrecord() + self.assertIsNotNone(etrecord) + self.assert_etrecord_saveable(etrecord) + + def test_executorch_export_etrecord_with_to_edge_flow_save_and_parse(self): + """Test that ETRecord generated by executorch.export can be saved and parsed.""" + export_session = self.get_test_export_session( + generate_etrecord=True, + to_edge_flow=True, + ) + + etrecord = export_session.get_etrecord() + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_export.bin" + + etrecord.save(etrecord_path) + + # Parse ETRecord back and verify + parsed_etrecord = parse_etrecord(etrecord_path) + + # Validate that all components are preserved + self.assertIsNotNone(parsed_etrecord.exported_program) + self.assertIsNotNone(parsed_etrecord.edge_dialect_program) + + # Validate executorch program data + et_manager = export_session.get_executorch_program_manager() + self.assertEqual( + parsed_etrecord._debug_handle_map, + json.loads(json.dumps(et_manager.debug_handle_map)), + ) + self.assertEqual( + parsed_etrecord._delegate_map, + json.loads(json.dumps(et_manager.delegate_map)), + ) + + # Validate export graph id is preserved + self.assertIsNotNone(parsed_etrecord.export_graph_id) + + def test_update_representative_inputs_with_list(self): + """Test update_representative_inputs with a list of ProgramInput objects.""" + captured_output, edge_output, et_output = self.get_test_model() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=et_output.debug_handle_map, + _delegate_map=et_output.delegate_map, + ) + + # Verify initial state - no representative inputs + self.assertIsNone(etrecord._representative_inputs) + + # Create custom representative inputs + f = models.BasicSinMax() + custom_inputs = [f.get_random_inputs() for _ in range(3)] + + # Update representative inputs + etrecord.update_representative_inputs(custom_inputs) + + # Verify representative inputs are now set + self.assertIsNotNone(etrecord._representative_inputs) + self.assertEqual(len(etrecord._representative_inputs), 3) + + # Compare the inputs using utility function + self.assert_representative_inputs_equal( + custom_inputs, + etrecord._representative_inputs, + "Custom inputs do not match ETRecord representative inputs", + ) + + def test_update_representative_inputs_with_bundled_program(self): + """Test update_representative_inputs with a BundledProgram.""" + ( + captured_output, + edge_output, + bundled_program, + ) = self.get_test_model_with_bundled_program() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=bundled_program.executorch_program.debug_handle_map, + _delegate_map=bundled_program.executorch_program.delegate_map, + ) + + # Verify initial state - no representative inputs + self.assertIsNone(etrecord._representative_inputs) + + # Update representative inputs using bundled program + etrecord.update_representative_inputs(bundled_program) + + # Verify representative inputs are now set + self.assertIsNotNone(etrecord._representative_inputs) + + # Compare with expected inputs from bundled program using utility function + expected_inputs = _get_representative_inputs(bundled_program) + self.assert_representative_inputs_equal( + expected_inputs, + etrecord._representative_inputs, + "Bundled program inputs do not match ETRecord representative inputs", + ) + + def test_update_representative_inputs_overwrite_existing(self): + """Test that update_representative_inputs overwrites existing inputs.""" + ( + captured_output, + edge_output, + bundled_program, + ) = self.get_test_model_with_bundled_program() + + # Create an ETRecord instance with existing representative inputs + initial_inputs = _get_representative_inputs(bundled_program) + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=bundled_program.executorch_program.debug_handle_map, + _delegate_map=bundled_program.executorch_program.delegate_map, + _representative_inputs=initial_inputs, + ) + + # Verify initial inputs are set + self.assertIsNotNone(etrecord._representative_inputs) + + # Create new custom inputs + f = models.BasicSinMax() + new_inputs = [f.get_random_inputs() for _ in range(2)] + + # Update representative inputs with new inputs + etrecord.update_representative_inputs(new_inputs) + + # Verify inputs are updated using utility function + self.assertEqual(len(etrecord._representative_inputs), 2) + self.assert_representative_inputs_equal( + new_inputs, + etrecord._representative_inputs, + "New inputs do not match ETRecord representative inputs after overwrite", + ) + + def test_update_reference_outputs_with_dict(self): + """Test update_reference_outputs with a dictionary of outputs.""" + captured_output, edge_output, et_output = self.get_test_model() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=et_output.debug_handle_map, + _delegate_map=et_output.delegate_map, + ) + + # Verify initial state - no reference outputs + self.assertIsNone(etrecord._reference_outputs) + + # Create custom reference outputs + f = models.BasicSinMax() + inputs = [f.get_random_inputs() for _ in range(2)] + custom_outputs = { + "forward": [f.forward(*inp) for inp in inputs], + "custom_method": [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])], + } + + # Update reference outputs + etrecord.update_reference_outputs(custom_outputs) + + # Verify reference outputs are now set + self.assertIsNotNone(etrecord._reference_outputs) + self.assertIn("forward", etrecord._reference_outputs) + self.assertIn("custom_method", etrecord._reference_outputs) + + # Compare the outputs + self.assertEqual(len(etrecord._reference_outputs["forward"]), 2) + self.assertEqual(len(etrecord._reference_outputs["custom_method"]), 2) + + for expected, actual in zip( + custom_outputs["forward"], etrecord._reference_outputs["forward"] + ): + self.assertTrue(torch.equal(expected[0], actual[0])) + + for expected, actual in zip( + custom_outputs["custom_method"], + etrecord._reference_outputs["custom_method"], + ): + self.assertTrue(torch.equal(expected, actual)) + + def test_update_reference_outputs_with_list(self): + """Test update_reference_outputs with a single list of outputs.""" + captured_output, edge_output, et_output = self.get_test_model() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=et_output.debug_handle_map, + _delegate_map=et_output.delegate_map, + ) + + # Verify initial state - no reference outputs + self.assertIsNone(etrecord._reference_outputs) + + # Create custom reference outputs as a single list + f = models.BasicSinMax() + inputs = [f.get_random_inputs() for _ in range(2)] + custom_outputs_list = [f.forward(*inp) for inp in inputs] + + # Update reference outputs with a single list + etrecord.update_reference_outputs(custom_outputs_list) + + # Verify reference outputs are now set and treated as "forward" method + self.assertIsNotNone(etrecord._reference_outputs) + self.assertIn("forward", etrecord._reference_outputs) + self.assertEqual(len(etrecord._reference_outputs["forward"]), 2) + + # Compare the outputs + for expected, actual in zip( + custom_outputs_list, etrecord._reference_outputs["forward"] + ): + self.assertTrue(torch.equal(expected[0], actual[0])) + + def test_update_reference_outputs_with_bundled_program(self): + """Test update_reference_outputs with a BundledProgram.""" + ( + captured_output, + edge_output, + bundled_program, + ) = self.get_test_model_with_bundled_program() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=bundled_program.executorch_program.debug_handle_map, + _delegate_map=bundled_program.executorch_program.delegate_map, + ) + + # Verify initial state - no reference outputs + self.assertIsNone(etrecord._reference_outputs) + + # Update reference outputs using bundled program + etrecord.update_reference_outputs(bundled_program) + + # Verify reference outputs are now set + self.assertIsNotNone(etrecord._reference_outputs) + self.assertIn("forward", etrecord._reference_outputs) + + # Compare with expected outputs from bundled program + expected_outputs = _get_reference_outputs(bundled_program) + self.assertTrue( + torch.equal( + etrecord._reference_outputs["forward"][0][0], + expected_outputs["forward"][0][0], + ) + ) + self.assertTrue( + torch.equal( + etrecord._reference_outputs["forward"][1][0], + expected_outputs["forward"][1][0], + ) + ) + + def test_update_apis_and_save_parse(self): + """Test that ETRecord with updated inputs/outputs can be saved and parsed correctly.""" + captured_output, edge_output, et_output = self.get_test_model() + + # Create an ETRecord instance + etrecord = ETRecord( + exported_program=captured_output, + export_graph_id=id(captured_output.graph), + edge_dialect_program=edge_output.exported_program(), + _debug_handle_map=et_output.debug_handle_map, + _delegate_map=et_output.delegate_map, + ) + + # Create custom inputs and outputs + f = models.BasicSinMax() + custom_inputs = [f.get_random_inputs() for _ in range(2)] + custom_outputs = { + "forward": [f.forward(*inp) for inp in custom_inputs], + } + + # Update both inputs and outputs + etrecord.update_representative_inputs(custom_inputs) + etrecord.update_reference_outputs(custom_outputs) + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_with_custom_data.bin" + + # Save the ETRecord + etrecord.save(etrecord_path) + + # Parse ETRecord back and verify + parsed_etrecord = parse_etrecord(etrecord_path) + + # Verify representative inputs are preserved using utility function + self.assertIsNotNone(parsed_etrecord._representative_inputs) + self.assertEqual(len(parsed_etrecord._representative_inputs), 2) + self.assert_representative_inputs_equal( + custom_inputs, + parsed_etrecord._representative_inputs, + "Custom inputs do not match parsed ETRecord representative inputs", + ) + + # Verify reference outputs are preserved + self.assertIsNotNone(parsed_etrecord._reference_outputs) + self.assertIn("forward", parsed_etrecord._reference_outputs) + self.assertEqual(len(parsed_etrecord._reference_outputs["forward"]), 2) + for expected, actual in zip( + custom_outputs["forward"], parsed_etrecord._reference_outputs["forward"] + ): + self.assertTrue(torch.equal(expected[0], actual[0])) + + def test_save_missing_essential_info(self): + def expected_runtime_error(etrecord, etrecord_path): + with self.assertRaises(RuntimeError) as context: + etrecord.save(etrecord_path) + + self.assertIn( + "ETRecord must contain edge dialect program and executorch program to be saved", + str(context.exception), + ) + + """Test that save raises RuntimeError when essential info is missing.""" + _, edge_output, et_output = self.get_test_model() + + etrecord = ETRecord() + + with tempfile.TemporaryDirectory() as tmpdirname: + etrecord_path = tmpdirname + "/etrecord_no_edge.bin" + + expected_runtime_error(etrecord, etrecord_path) + etrecord.add_edge_dialect_program(edge_output) + + # Should raise runtime error due to missing executorch program related info + expected_runtime_error(etrecord, etrecord_path) + + etrecord.add_executorch_program(et_output) + + # All essential components are now present, so save should succeed + etrecord.save(etrecord_path) diff --git a/devtools/inspector/TARGETS b/devtools/inspector/TARGETS index d32698f784f..8834dfbb6ba 100644 --- a/devtools/inspector/TARGETS +++ b/devtools/inspector/TARGETS @@ -1,9 +1,9 @@ load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "inspector", srcs = [ "_inspector.py", @@ -33,7 +33,7 @@ python_binary( ], ) -python_library( +runtime.python_library( name = "inspector_utils", srcs = [ "_inspector_utils.py", @@ -50,7 +50,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "intermediate_output_capturer", srcs = [ "_intermediate_output_capturer.py", @@ -60,7 +60,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "lib", srcs = ["__init__.py"], deps = [ diff --git a/devtools/inspector/_inspector.py b/devtools/inspector/_inspector.py index 17a7451aadf..6d046d8f2e8 100644 --- a/devtools/inspector/_inspector.py +++ b/devtools/inspector/_inspector.py @@ -317,6 +317,8 @@ class Event: op_type: List of op types corresponding to the event. delegate_debug_identifier: Supplemental identifier used in combination with instruction id. debug_handles: Debug handles in the model graph to which this event is correlated. + num_outputs: Indicates the number of outputs generated by the node. + Right now only used for call_delegate nodes that output more than one tensor. stack_trace: A dictionary mapping the name of each associated op to its stack trace. module_hierarchy: A dictionary mapping the name of each associated op to its module hierarchy. is_delegated_op: Whether or not the event was delegated. @@ -337,6 +339,7 @@ class Event: op_types: List[str] = dataclasses.field(default_factory=list) delegate_debug_identifier: Optional[Union[int, str]] = None debug_handles: Optional[Union[int, Sequence[int]]] = None + num_outputs: int = 1 stack_traces: Dict[str, str] = dataclasses.field(default_factory=dict) module_hierarchy: Dict[str, Dict] = dataclasses.field(default_factory=dict) is_delegated_op: Optional[bool] = None @@ -928,6 +931,7 @@ def _gen_resolve_debug_handles( self, handle_map: Dict[str, List[int]], delegate_map: Optional[Dict[str, DelegateMetadata]] = None, + instruction_id_to_num_outs_map: Dict[int, int] = None, ): """ Given mappings from instruction id to debug handles, populate the @@ -945,6 +949,10 @@ def _gen_resolve_debug_handles( if (instruction_id := str(event._instruction_id)) not in handle_map: continue + num_outputs = 1 + if instruction_id_to_num_outs_map is not None: + num_outputs = instruction_id_to_num_outs_map.get(instruction_id, 1) + event.num_outputs = num_outputs # For non-delegated event, handles are found in handle_map if (delegate_debug_id := event.delegate_debug_identifier) is None: event.debug_handles = handle_map[instruction_id] @@ -1131,6 +1139,7 @@ def _consume_etrecord(self) -> None: if self._etrecord._delegate_map is not None else None ), + self._etrecord._instruction_id_to_num_outs_map[FORWARD], ) # (2) Event Metadata Association @@ -1170,7 +1179,7 @@ def _get_aot_intermediate_outputs_and_op_names( export_program = None - # Will use the exported program to extract intermediate output if and only if exported_program has been provided, and it is the greatest ancestor of the edge_dialect_program + # Will use the exported program to extract intermediate output if and only if exported_program has been provided, and it is one of the ancestors of the edge_dialect_program if self._etrecord.exported_program and propagate_back_debug_handle( self._etrecord.exported_program, self._etrecord.export_graph_id, @@ -1178,6 +1187,10 @@ def _get_aot_intermediate_outputs_and_op_names( ): export_program = self._etrecord.exported_program else: + log.warning( + "Either aten dialect exported program is not in ETRecord, or it is not one of the ancestors of current edge dialect program." + "Will fall back to use edge dialect program to extract intermediate output", + ) export_program = self._etrecord.edge_dialect_program graph_module = export_program.module() aot_debug_handle_to_op_name = get_aot_debug_handle_to_op_name_mapping( @@ -1192,7 +1205,7 @@ def _get_aot_intermediate_outputs_and_op_names( # TODO: Make it more extensible to further merge overlapping debug handles def _get_runtime_intermediate_outputs_and_op_names( self, - ) -> Tuple[Dict[DebugHandle, Any], Dict[DebugHandle, List[str]]]: + ) -> Tuple[Dict[DebugHandle, Tuple[Any, int]], Dict[DebugHandle, List[str]]]: """ Retrieve the runtime intermediate outputs(debug handles and intermediate values mappings) from the event blocks, along with the corresponding debug handles and op names mapping. @@ -1213,12 +1226,15 @@ def _get_runtime_intermediate_outputs_and_op_names( debug_handle = (debug_handle,) else: debug_handle = tuple(debug_handle) - current_entry = debug_handle_to_output.get(debug_handle, (-1, None)) + current_entry = debug_handle_to_output.get( + debug_handle, (-1, None, event.num_outputs) + ) # When event has same debug_handle, only keep the one with the largest instruction id if event._instruction_id > current_entry[0]: debug_handle_to_output[debug_handle] = ( event._instruction_id, event.debug_data, + event.num_outputs, ) # TODO: One debug handle can be associated with multiple op names debug_handle_to_op_names[debug_handle] = [event.name] @@ -1227,7 +1243,7 @@ def _get_runtime_intermediate_outputs_and_op_names( debug_handle_to_output ) return { - k: v[1] for k, v in debug_handle_to_output.items() + k: (v[1], v[2]) for k, v in debug_handle_to_output.items() }, debug_handle_to_op_names def to_dataframe( @@ -1392,7 +1408,9 @@ def calculate_numeric_gap(self, distance: str = "MSE"): """ Compares logged intermediate outputs from the exported graph (in ETRecord) with runtime outputs (in ETDump) using a user-specific numerical comparator. - To use this function, you must first generate the ETRecord using the `bundle_program`, + If the exported graph is not supported, the function will fall back to use edge dialect graph. + + To use this function, you must first generate the ETRecord with representative inputs, and then create the Inspector instance with the ETRecord and ETDump. The Inspector can then compare the intermediate outputs from the AOT and the runtime. diff --git a/devtools/inspector/_inspector_utils.py b/devtools/inspector/_inspector_utils.py index 2bda03b4873..ee7ebb2f5ea 100644 --- a/devtools/inspector/_inspector_utils.py +++ b/devtools/inspector/_inspector_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. +# copyright (c) meta platforms, inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the @@ -99,11 +99,14 @@ class NodeData: - source: A string indicating the origin of the node (either FROM_AOT or FROM_RUNTIME). - debug_handle: A tuple representing the unique identifier for the output. - output: The actual output data associated with the debug handle. + - num_outputs: Indicates the number of outputs generated by the node. + Right now only used for call_delegate nodes that output more than one tensor. """ source: NodeSource debug_handle: tuple[int] output: Any + num_outputs: int class NodeFilter: @@ -578,8 +581,8 @@ def _merge_runtime_debug_handles( def merge_runtime_overlapping_debug_handles( - runtime_intermediate_outputs: Dict[DebugHandle, Tuple[int, Any]] -) -> Dict[DebugHandle, Tuple[int, Any]]: + runtime_intermediate_outputs: Dict[DebugHandle, Tuple[int, Any, int]] +) -> Dict[DebugHandle, Tuple[int, Any, int]]: """ Merges runtimes with overlapping debug handles into a single key in the dict. @@ -592,12 +595,17 @@ def merge_runtime_overlapping_debug_handles( """ if len(runtime_intermediate_outputs) == 0: return {} - merged: Dict[DebugHandle, Tuple[int, Any]] = {} + merged: Dict[DebugHandle, Tuple[int, Any, int]] = {} for debug_handle, ( instruction_id, debug_data, + num_outputs, ) in runtime_intermediate_outputs.items(): - curr_debug_handle, last_value = debug_handle, (instruction_id, debug_data) + curr_debug_handle, last_value = debug_handle, ( + instruction_id, + debug_data, + num_outputs, + ) # Collect any existing keys that overlap with the current key to_remove = [] for existing_debug_handle, existing_value in merged.items(): @@ -634,7 +642,9 @@ def _debug_handles_have_overlap( def _combine_aot_overlapped_intermediate_outputs( - aot_nodes: List[Tuple[DebugHandle, Any]], runtime_node: Tuple[DebugHandle, Any] + aot_nodes: List[Tuple[DebugHandle, Any]], + runtime_node: Tuple[DebugHandle, Any, int], + negative_index: int, ) -> Tuple[DebugHandle, Any]: """ Ensure the AOT combined debug_handles are the same as the runtime debug_handles (order ignored), @@ -642,7 +652,7 @@ def _combine_aot_overlapped_intermediate_outputs( """ # Map AOT single element debug_handles to outputs aot_map = dict(aot_nodes) - runtime_debug_handle, _ = runtime_node + runtime_debug_handle, _, _ = runtime_node # Combine all AOT debug_handles into a list aot_combined_debug_handle = [t[0] for t in aot_map.keys()] @@ -652,14 +662,14 @@ def _combine_aot_overlapped_intermediate_outputs( return (-1,), None # Pick the last intermediate output - last_int = runtime_debug_handle[-1] + last_int = runtime_debug_handle[negative_index] key = (last_int,) return runtime_debug_handle, aot_map[key] def _create_debug_handle_overlap_graph( aot_intermediate_outputs: Dict[DebugHandle, Any], - runtime_intermediate_outputs: Dict[DebugHandle, Any], + runtime_intermediate_outputs: Dict[DebugHandle, Tuple[Any, int]], ) -> Tuple[List[NodeData], Dict[int, List[int]]]: """ Create a graph representing overlapping debug handles between AOT and runtime outputs. @@ -675,9 +685,14 @@ def _create_debug_handle_overlap_graph( """ nodes = [] for debug_handle, output in aot_intermediate_outputs.items(): - nodes.append(NodeData(NodeSource.AOT, debug_handle, output)) - for debug_handle, output in runtime_intermediate_outputs.items(): - nodes.append(NodeData(NodeSource.RUNTIME, debug_handle, output)) + # TODO: for aot outputs also derive the number of output tensors generated by the node + nodes.append(NodeData(NodeSource.AOT, debug_handle, output, 1)) + for debug_handle, value in runtime_intermediate_outputs.items(): + nodes.append( + NodeData( + NodeSource.RUNTIME, debug_handle, output=value[0], num_outputs=value[1] + ) + ) edges = {i: [] for i in range(len(nodes))} for i in range(len(nodes)): @@ -730,7 +745,7 @@ def dfs(node_id, component): def map_runtime_aot_intermediate_outputs( aot_intermediate_outputs: Dict[DebugHandle, Any], - runtime_intermediate_outputs: Dict[DebugHandle, Any], + runtime_intermediate_outputs: Dict[DebugHandle, Tuple[Any, int]], ) -> Dict[Tuple[DebugHandle, Any], Tuple[DebugHandle, Any]]: """ Map the runtime intermediate outputs to the AOT intermediate outputs @@ -757,7 +772,11 @@ def map_runtime_aot_intermediate_outputs( if nodes[node_id].source == NodeSource.AOT ] runtime_list = [ - (nodes[node_id].debug_handle, nodes[node_id].output) + ( + nodes[node_id].debug_handle, + nodes[node_id].output, + nodes[node_id].num_outputs, + ) for node_id in comp if nodes[node_id].source == NodeSource.RUNTIME ] @@ -772,50 +791,74 @@ def map_runtime_aot_intermediate_outputs( f"Expected only one runtime debug handle, but found {len(runtime_list)}: {runtime_list}" ) - runtime_debug_handle, runtime_intermediate_output = runtime_list[0] - - # Combine aot debug handles into a single key - aot_combined_debug_handle, aot_intermediate_output = ( - _combine_aot_overlapped_intermediate_outputs(aot_list, runtime_list[0]) + runtime_debug_handle, runtime_intermediate_output, num_outputs = ( + runtime_list[0] ) + # iterate through each of the output from runtime, + # get the corresponding debug handle + # and map it to the aot debug handle + # and create a dictionary that maps aot debug handle + aot output to + # runtime debug handle + runtime output + # Note this works only for delegate case for now. + for i in range(num_outputs): + + negative_index = -1 * (i + 1) + aot_mapped_runtime_intermediate_output = runtime_intermediate_output + # Combine aot debug handles into a single key + aot_combined_debug_handle, aot_intermediate_output = ( + _combine_aot_overlapped_intermediate_outputs( + aot_list, runtime_list[0], negative_index + ) + ) - if aot_combined_debug_handle == (-1,): - # Skip this mapping if the aot combined debug handle and runtime debug handle do not exact match. - continue + if aot_combined_debug_handle == (-1,): + # Skip this mapping if the aot combined debug handle and runtime debug handle do not exact match. + continue - if isinstance(aot_intermediate_output, Sequence): - if not isinstance(runtime_intermediate_output, Sequence): - raise TypeError( - "runtime intermediate output should be a sequence when aot intermediate output is a sequence" + if isinstance(aot_intermediate_output, Sequence): + if not isinstance(runtime_intermediate_output, Sequence): + raise TypeError( + "runtime intermediate output should be a sequence when aot intermediate output is a sequence" + ) + last_element = runtime_intermediate_output[negative_index] + # TODO: this (last_element = list) is never really the case because runtime never returns output as a list + # for delegate case. + if isinstance(last_element, list) and all( + isinstance(t, torch.Tensor) for t in last_element + ): + # If the last element is a list of tensors (delegate case) + aot_mapped_runtime_intermediate_output = last_element + elif isinstance(last_element, torch.Tensor): + # If the last element is a tensor, as is always the case for runtime. + # However, now we have a strange condition where aot_intermediate_output is a list of tensors + # while runtime_intermediate_output is a single tensor. So we should never really come here. + # TODO: fix this + aot_mapped_runtime_intermediate_output = ( + runtime_intermediate_output + ) + else: + raise ValueError( + "The last element of runtime argument list must be a tensor or a list of tensors when aot intermediate output is a sequence" + ) + # List can't be used as a key, so convert to tuple + aot_intermediate_output = tuple(aot_intermediate_output) + aot_mapped_runtime_intermediate_output = tuple( + aot_mapped_runtime_intermediate_output ) - last_element = runtime_intermediate_output[-1] - if isinstance(last_element, list) and all( - isinstance(t, torch.Tensor) for t in last_element - ): - # If the last element is a list of tensors (delegate case) - runtime_intermediate_output = last_element - elif isinstance(last_element, torch.Tensor): - # If the last element is a tensor (non-delegate case) - pass - else: - raise ValueError( - "The last element of runtime argument list must be a tensor or a list of tensors when aot intermediate output is a sequence" + + elif isinstance(runtime_intermediate_output, Sequence): + # delegate runtime call and AOT intermediate is not a sequence, just take the last element from runtime list + aot_mapped_runtime_intermediate_output = ( + runtime_intermediate_output[negative_index] ) - # List can't be used as a key, so convert to tuple - aot_intermediate_output = tuple(aot_intermediate_output) - runtime_intermediate_output = tuple(runtime_intermediate_output) - - elif isinstance(runtime_intermediate_output, Sequence): - # delegate runtime call and AOT intermediate is not a sequence, just take the last element from runtime list - runtime_intermediate_output = runtime_intermediate_output[-1] - - # Create a mapping between runtime and aot - aot_runtime_mapping[ - (aot_combined_debug_handle, aot_intermediate_output) - ] = ( - runtime_debug_handle, - runtime_intermediate_output, - ) + + # Create a mapping between runtime and aot + aot_runtime_mapping[ + (aot_combined_debug_handle, aot_intermediate_output) + ] = ( + runtime_debug_handle, + aot_mapped_runtime_intermediate_output, + ) return aot_runtime_mapping diff --git a/devtools/inspector/numerical_comparator/TARGETS b/devtools/inspector/numerical_comparator/TARGETS index ec190f5cc6d..eb2923d924d 100644 --- a/devtools/inspector/numerical_comparator/TARGETS +++ b/devtools/inspector/numerical_comparator/TARGETS @@ -1,15 +1,15 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "numerical_comparator_base", srcs = ["numerical_comparator_base.py"], deps = [], ) -python_library( +runtime.python_library( name = "l1_numerical_comparator", srcs = ["l1_numerical_comparator.py"], deps = [ @@ -18,7 +18,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "mse_numerical_comparator", srcs = ["mse_numerical_comparator.py"], deps = [ @@ -27,7 +27,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "snr_numerical_comparator", srcs = ["snr_numerical_comparator.py"], deps = [ @@ -36,7 +36,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "lib", srcs = ["__init__.py"], deps = [ diff --git a/devtools/inspector/tests/TARGETS b/devtools/inspector/tests/TARGETS index 250cc76fd31..048c8f6f791 100644 --- a/devtools/inspector/tests/TARGETS +++ b/devtools/inspector/tests/TARGETS @@ -1,11 +1,15 @@ load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbsource//tools/target_determinator/macros:ci.bzl", "ci") oncall("executorch") python_unittest( name = "inspector_test", srcs = ["inspector_test.py"], + labels = ci.labels( + ci.buckconfig("executorch.event_tracer_enabled", "true"), + ), deps = [ "//executorch/devtools:lib", "//executorch/devtools/debug_format:et_schema", @@ -78,7 +82,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "inspector_test_utils", srcs = [ "inspector_test_utils.py", diff --git a/devtools/inspector/tests/inspector_test.py b/devtools/inspector/tests/inspector_test.py index 37dc7921923..a3afed07ed8 100644 --- a/devtools/inspector/tests/inspector_test.py +++ b/devtools/inspector/tests/inspector_test.py @@ -7,6 +7,7 @@ # pyre-unsafe import copy +import os import random import statistics import tempfile @@ -21,7 +22,9 @@ import torch import torch.fx +import torch.utils._pytree as pytree +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner from executorch.devtools import generate_etrecord, parse_etrecord from executorch.devtools.debug_format.et_schema import OperatorNode from executorch.devtools.etdump.schema_flatcc import ProfileEvent @@ -52,6 +55,10 @@ EdgeProgramManager, ExecutorchProgramManager, to_edge, + to_edge_transform_and_lower, +) +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, ) from torch.export import export, ExportedProgram @@ -633,7 +640,9 @@ def test_get_runtime_intermediate_outputs_and_op_names(self): self.assertIn((4,), runtime_outputs) self.assertIn((4,), op_names) self.assertTrue( - torch.allclose(runtime_outputs[(4,)][0], torch.tensor([4.0, 5.0, 6.0])) + torch.allclose( + runtime_outputs[(4,)][0][0], torch.tensor([4.0, 5.0, 6.0]) + ) ) self.assertEqual(op_names[(4,)], ["op_3"]) @@ -641,8 +650,6 @@ def test_get_runtime_intermediate_outputs_and_op_names(self): for key in range(5, 9): self.assertIn((key,), runtime_outputs) self.assertIn((key,), op_names) - self.assertEqual(runtime_outputs[(key,)][0].size(0), RAW_DATA_SIZE) - self.assertEqual(op_names[(key,)], [f"op_{key-1}"]) def test_calculate_numeric_gap(self): # Create a context manager to patch functions called by Inspector.__init__ @@ -655,7 +662,6 @@ def test_calculate_numeric_gap(self): ), patch.object( _inspector, "gen_graphs_from_etrecord" ): - # Call the constructor of Inspector inspector_instance = Inspector( etdump_path=ETDUMP_PATH, @@ -668,8 +674,8 @@ def test_calculate_numeric_gap(self): } runtime_intermediate_outputs = { - (0,): torch.tensor([2.0, 1.0, 4.0]), - (1,): torch.tensor([3.0, 6.0, 5.0]), + (0,): ([torch.tensor([2.0, 1.0, 4.0])], 1), + (1,): ([torch.tensor([3.0, 6.0, 5.0])], 1), } aot_debug_handle_to_op_name = {(0,): "op_0", (1,): "op_1"} @@ -709,12 +715,121 @@ def test_calculate_numeric_gap(self): self.assertTrue( torch.allclose( row["runtime_intermediate_output"], - runtime_intermediate_outputs[key], + runtime_intermediate_outputs[key][0][0], ) ) # gap should equal 3.0 self.assertEqual(row["gap"][0], 3.0) + @unittest.skip("ci config values are not propagated") + def test_intermediate_tensor_comparison_with_torch_export(self): + """Test intermediate tensor comparison using torch.export.export and to_edge_transform_and_lower.""" + + class SimpleTestModel(torch.nn.Module): + """A simple test model for demonstration purposes.""" + + def __init__(self, hidden_size: int = 32, num_layers: int = 2): + super().__init__() + self.layers = torch.nn.ModuleList( + [ + torch.nn.Linear(hidden_size, hidden_size) + for _ in range(num_layers) + ] + ) + self.activation = torch.nn.ReLU() + self.output_layer = torch.nn.Linear(hidden_size, 10) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.activation(self.layers[0](x)) + y = self.activation(self.layers[1](x)) + return y, self.output_layer(x) + + # Create test model and inputs + model = SimpleTestModel(hidden_size=32, num_layers=2) + model.eval() + + # Create representative inputs (smaller for faster testing) + batch_size, seq_len, hidden_size = 1, 8, 32 + input_tensor = torch.randn(batch_size, seq_len, hidden_size) + example_inputs = (input_tensor,) + representative_inputs = [example_inputs] + + with tempfile.TemporaryDirectory() as tmp_dir: + model_path = os.path.join(tmp_dir, "model.pte") + etrecord_path = os.path.join(tmp_dir, "etrecord.bin") + + # Step 1: Export using torch.export.export + exported_program = torch.export.export(model, example_inputs) + self.assertIsNotNone(exported_program) + + # Step 2: Lower to XNNPACK with generate_etrecord=True + edge_compile_config = EdgeCompileConfig(_check_ir_validity=False) + edge_program_manager = to_edge_transform_and_lower( + exported_program, + partitioner=[XnnpackPartitioner()], + compile_config=edge_compile_config, + generate_etrecord=True, + ) + self.assertIsNotNone(edge_program_manager) + + # Step 3: Generate ETRecord from edge program manager + # Step 4: Convert to executorch and save as PTE + executorch_program = edge_program_manager.to_executorch() + et_record = executorch_program.get_etrecord() + self.assertIsNotNone(et_record) + + # Update with representative inputs + flattened_x = pytree.tree_flatten(representative_inputs[0])[0] + et_record.update_representative_inputs(flattened_x) + et_record.save(etrecord_path) + + with open(model_path, "wb") as f: + executorch_program.write_to_file(f) + + # Step 5: Test intermediate output comparison using pybind APIs + # Read the PTE file + with open(model_path, "rb") as f: + pte_buffer = f.read() + + etdump_path = os.path.join(tmp_dir, "etdump.etdp") + debug_buffer_path = os.path.join(tmp_dir, "debug_buffer.bin") + + # Load the PTE file with ETDump enabled using pybind API + executorch_module = _load_for_executorch_from_buffer( + pte_buffer, + enable_etdump=True, + debug_buffer_size=1024 * 1024, # 1MB for testing + ) + self.assertIsNotNone(executorch_module) + + # Run the model with the given input using pybind API + flattened_x = pytree.tree_flatten(representative_inputs[0])[0] + executorch_module.run_method("forward", tuple(flattened_x)) + + # Write the ETDump results to a file using pybind API + executorch_module.write_etdump_result_to_file( + etdump_path, debug_buffer_path + ) + + # Step 6: Use Inspector API to compare intermediate outputs + try: + inspector = Inspector( + etdump_path=etdump_path, + etrecord=etrecord_path, + debug_buffer_path=debug_buffer_path, + ) + except FileNotFoundError as e: + new_message = f"{e} You likely need to run the test with --config executorch.event_tracer_enabled=true" + raise RuntimeError(new_message) from e + self.assertIsNotNone(inspector) + + # Calculate numerical gap using SNR metric + df = inspector.calculate_numeric_gap("SNR") + + # Verify that we got some intermediate tensor comparisons + # The exact number will depend on the model structure and partitioning + self.assertEqual(len(df), 2) + def _gen_random_float_list(self) -> List[float]: return [random.uniform(0, 10) for _ in range(RAW_DATA_SIZE)] diff --git a/devtools/inspector/tests/inspector_utils_test.py b/devtools/inspector/tests/inspector_utils_test.py index ea8c0e653af..26fe38acfac 100644 --- a/devtools/inspector/tests/inspector_utils_test.py +++ b/devtools/inspector/tests/inspector_utils_test.py @@ -233,20 +233,20 @@ def test_compare_results_uint8(self): def test_merge_overlapping_debug_handles_basic(self): big_tensor = torch.rand(100, 100) intermediate_outputs = { - (1, 2, 3): (1, "val1"), - (2, 3, 4, 5): (2, "val2"), - (6, 7, 8): (3, "val3"), - (10, 11): (4, "val4"), - (11, 12): (5, big_tensor), + (1, 2, 3): (1, "val1", 1), + (2, 3, 4, 5): (2, "val2", 1), + (6, 7, 8): (3, "val3", 1), + (10, 11): (4, "val4", 1), + (11, 12): (5, big_tensor, 1), } # basic merge behavior intermediate_outputs = merge_runtime_overlapping_debug_handles( intermediate_outputs ) expected_intermediate_outputs = { - (1, 2, 3, 4, 5): (2, "val2"), - (6, 7, 8): (3, "val3"), - (10, 11, 12): (5, big_tensor), + (1, 2, 3, 4, 5): (2, "val2", 1), + (6, 7, 8): (3, "val3", 1), + (10, 11, 12): (5, big_tensor, 1), } self.assertEqual(intermediate_outputs, expected_intermediate_outputs) self.assertIs(expected_intermediate_outputs[(10, 11, 12)][1], big_tensor) @@ -258,11 +258,11 @@ def test_merge_overlapping_debug_handles_non_continuous(self): tensor4 = torch.randn(6, 7) tensor5 = torch.randn(8, 9) intermediate_outputs = { - (1, 10): (1, tensor1), - (2, 5): (2, tensor2), - (1, 7, 9): (3, tensor3), - (11, 13): (4, tensor4), - (11, 15): (5, tensor5), + (1, 10): (1, tensor1, 1), + (2, 5): (2, tensor2, 1), + (1, 7, 9): (3, tensor3, 1), + (11, 13): (4, tensor4, 1), + (11, 15): (5, tensor5, 1), } intermediate_outputs = merge_runtime_overlapping_debug_handles( intermediate_outputs @@ -280,22 +280,22 @@ def test_merge_overlapping_debug_handles_non_continuous(self): def test_merge_overlapping_debug_handles_edge_cases(self): intermediate_outputs = { - (9,): (1, "val1"), + (9,): (1, "val1", 1), ( 9, 9, 9, - ): (2, "val2"), + ): (2, "val2", 1), ( 9, 9, - ): (3, "val3"), + ): (3, "val3", 1), } intermediate_outputs = merge_runtime_overlapping_debug_handles( intermediate_outputs ) expected_intermediate_outputs = { - (9,): (3, "val3"), + (9,): (3, "val3", 1), } self.assertEqual(intermediate_outputs, expected_intermediate_outputs) @@ -312,7 +312,7 @@ def test_map_runtime_aot_intermediate_outputs_empty_inputs(self): def test_map_runtime_aot_intermediate_outputs_single_element_tuple(self): # Single element tuple aot_intermediate_outputs = {(0,): 100, (1,): 200, (2,): 300} - runtime_intermediate_outputs = {(0,): 150, (1,): 250, (2,): 350} + runtime_intermediate_outputs = {(0,): (150, 1), (1,): (250, 1), (2,): (350, 1)} actual = map_runtime_aot_intermediate_outputs( aot_intermediate_outputs, runtime_intermediate_outputs ) @@ -326,7 +326,7 @@ def test_map_runtime_aot_intermediate_outputs_single_element_tuple(self): def test_map_runtime_aot_intermediate_outputs_no_overlaps(self): # No overlaps between aot and runtime debug_handles aot_intermediate_outputs = {(0,): 100, (4,): 300} - runtime_intermediate_outputs = {(2, 3): 200, (8, 9): 300} + runtime_intermediate_outputs = {(2, 3): (200, 1), (8, 9): (300, 1)} actual = map_runtime_aot_intermediate_outputs( aot_intermediate_outputs, runtime_intermediate_outputs ) @@ -336,7 +336,7 @@ def test_map_runtime_aot_intermediate_outputs_no_overlaps(self): def test_map_runtime_aot_intermediate_outputs_partial_match(self): # Partial match between aot and runtime debug_handles will return empty aot_intermediate_outputs = {(2,): 100, (9,): 300} - runtime_intermediate_outputs = {(2, 3): 200, (8, 9): 300} + runtime_intermediate_outputs = {(2, 3): (200, 1), (8, 9): (300, 1)} actual = map_runtime_aot_intermediate_outputs( aot_intermediate_outputs, runtime_intermediate_outputs ) @@ -346,7 +346,7 @@ def test_map_runtime_aot_intermediate_outputs_partial_match(self): def test_map_runtime_aot_intermediate_outputs_multiple_aot_to_one_runtime(self): # Multiple aot debug_handles map to one runtime debug_handle aot_intermediate_outputs = {(0,): 100, (1,): 200, (2,): 300, (3,): 400} - runtime_intermediate_outputs = {(2, 3, 1): 250, (8, 9): 300} + runtime_intermediate_outputs = {(2, 3, 1): (250, 1), (8, 9): (300, 1)} actual = map_runtime_aot_intermediate_outputs( aot_intermediate_outputs, runtime_intermediate_outputs ) @@ -357,35 +357,45 @@ def test_map_runtime_aot_intermediate_outputs_delegated(self): # Currently, runtime_intermediate_output logs all delegate call arguments # Test that the map function correctly extracted out the delegated outputs aot_intermediate_outputs = { - (1,): torch.tensor([4, 1]), + (1,): torch.tensor([1, 2, 3]), (2,): torch.tensor([4, 5]), (3,): torch.tensor([10, 10, 13]), (4,): torch.tensor([10, 11, 12]), (5,): torch.tensor([13, 14, 15, 16, 21]), - (6,): torch.tensor([13, 14, 15, 16, 17]), + (6,): torch.tensor([2]), } runtime_intermediate_outputs = { - (1, 2): [torch.tensor([1, 2, 3]), torch.tensor([4, 5])], - (3, 4): [ - torch.tensor([6, 7, 8, 9]), - torch.tensor(1), - torch.tensor([10, 11, 12]), - ], - (5, 6): [ - torch.tensor([1]), - torch.tensor([2]), - torch.tensor([13, 14, 15, 16, 17]), - ], + (1, 2): ([torch.tensor([1, 2, 3]), torch.tensor([4, 5])], 2), + (3, 4): ( + [ + torch.tensor([10, 10, 13]), + torch.tensor([10, 11, 12]), + ], + 2, + ), + (5, 6): ( + [ + torch.tensor([13, 14, 15, 16, 21]), + torch.tensor([2]), + ], + 2, + ), } actual = map_runtime_aot_intermediate_outputs( aot_intermediate_outputs, runtime_intermediate_outputs ) expected = { + ((1, 2), torch.tensor([1, 2, 3])): ((1, 2), torch.tensor([1, 2, 3])), ((1, 2), torch.tensor([4, 5])): ((1, 2), torch.tensor([4, 5])), + ((3, 4), torch.tensor([10, 10, 13])): ((3, 4), torch.tensor([10, 10, 13])), ((3, 4), torch.tensor([10, 11, 12])): ((3, 4), torch.tensor([10, 11, 12])), - ((5, 6), torch.tensor([13, 14, 15, 16, 17])): ( + ((5, 6), torch.tensor([13, 14, 15, 16, 21])): ( (5, 6), - torch.tensor([13, 14, 15, 16, 17]), + torch.tensor([13, 14, 15, 16, 21]), + ), + ((5, 6), torch.tensor([2])): ( + (5, 6), + torch.tensor([2]), ), } self.assertEqual(len(actual), len(expected)) @@ -399,8 +409,10 @@ def test_map_runtime_aot_intermediate_outputs_delegated(self): act_runtime_key, act_runtime_value, ) in actual.items(): - if exp_aot_key == act_aot_key and torch.allclose( - exp_aot_value, act_aot_value + if ( + exp_aot_key == act_aot_key + and exp_aot_value.numel() == act_aot_value.numel() + and torch.allclose(exp_aot_value, act_aot_value) ): found = True self.assertEqual(exp_runtime_key, act_runtime_key) diff --git a/devtools/scripts/generate_profiling_csv.py b/devtools/scripts/generate_profiling_csv.py new file mode 100644 index 00000000000..71e0a4070f3 --- /dev/null +++ b/devtools/scripts/generate_profiling_csv.py @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright 2024-25 Arm Limited and/or its 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 executorch.devtools import Inspector + + +def generate_csv(etdump_path, output): + """ + Generate a CSV file from ETDump profiling data. + + Args: + etdump_path (str): Path to the ETDump file generated by executor_runner + output (str): Path for the output CSV file + """ + inspector = Inspector(etdump_path) + df = inspector.to_dataframe() + df.to_csv(output) + + +def main(): + """ + Main function to parse command line arguments and generate profiling CSV. + + Usage: + python generate_profiling_csv.py --etdump_path="my_etdump" --output="profiling.csv" + + Example: + python generate_profiling_csv.py --etdump_path="llama3_etdump" --output="op_profiling.csv" + """ + parser = argparse.ArgumentParser( + description="Generate profiling CSV from a model's etdump" + ) + parser.add_argument( + "--etdump_path", + type=str, + default="./model.etdump", + help="Path to the etdump file", + required=False, + ) + + parser.add_argument( + "--output", + type=str, + default="./model_profiling.csv", + help="Path to the output CSV file", + required=False, + ) + + args = parser.parse_args() + print(f"Generating CSV from {args.etdump_path}") + generate_csv(args.etdump_path, args.output) + print(f"Saved CSV to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/devtools/scripts/profile_model.sh b/devtools/scripts/profile_model.sh new file mode 100755 index 00000000000..8697c97cd02 --- /dev/null +++ b/devtools/scripts/profile_model.sh @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright 2024-25 Arm Limited and/or its 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. + +#!/bin/bash + +# ExecutorTorch Model Profiling Script +# +# This script automates the process of building executor_runner with profiling enabled, +# running model inference with ETDump collection, and generating CSV profiling reports. +# +# Usage: +# ./devtools/scripts/profile_model.sh [model_path] [etdump_path] +# +# Arguments: +# model_path - Path to the .pte model file (default: "my_model") +# etdump_path - Path for ETDump output file (default: "path_to_et_dump") +# +# Examples: +# ./devtools/scripts/profile_model.sh +# ./devtools/scripts/profile_model.sh llama3.pte llama3_etdump +# +# Note: This script must be run from the top-level executorch directory. + +set -e + +echo "Building executor_runner with profiling enabled..." + +cmake --preset profiling -B build-profiling -DCMAKE_BUILD_TYPE=Release +cmake --build build-profiling --target executor_runner + +echo "Build completed successfully!" + +MODEL_PATH=${1:-"my_model"} +ETDUMP_PATH=${2:-"path_to_et_dump"} + +echo "Running and profiling model: $MODEL_PATH" +echo "ETDump output path: $ETDUMP_PATH" + +./build-profiling/executor_runner --model_path="$MODEL_PATH" --etdump_path="$ETDUMP_PATH" + +echo "Profiling run completed!" + +echo "Generating profiling CSV..." +python devtools/scripts/generate_profiling_csv.py --etdump_path="$ETDUMP_PATH" --output="op_profiling.csv" + +echo "Profiling CSV generated: op_profiling.csv" +echo "Profiling workflow completed successfully!" diff --git a/devtools/size_analysis_tool/TARGETS b/devtools/size_analysis_tool/TARGETS index c365ba152d5..b37eee6a389 100644 --- a/devtools/size_analysis_tool/TARGETS +++ b/devtools/size_analysis_tool/TARGETS @@ -1,10 +1,10 @@ load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") -python_library( +runtime.python_library( name = "size_analysis_tool_lib", srcs = [ "size_analysis_tool.py", diff --git a/devtools/visualization/__init__.py b/devtools/visualization/__init__.py index df1d74c7fae..8e91d7ffdb2 100644 --- a/devtools/visualization/__init__.py +++ b/devtools/visualization/__init__.py @@ -9,4 +9,5 @@ SingletonModelExplorerServer, visualize, visualize_graph, + visualize_model_explorer, ) diff --git a/devtools/visualization/visualization_utils.py b/devtools/visualization/visualization_utils.py index d21d11082a3..b21a953f4d2 100644 --- a/devtools/visualization/visualization_utils.py +++ b/devtools/visualization/visualization_utils.py @@ -108,7 +108,7 @@ def visualize( **kwargs, ): """Wraps the visualize_from_config call from model_explorer. - For convenicence, figures out how to find the exported_program + For convenience, figures out how to find the exported_program from EdgeProgramManager and ExecutorchProgramManager for you. See https://github.com/google-ai-edge/model-explorer/wiki/4.-API-Guide#visualize-pytorch-models @@ -123,13 +123,22 @@ def visualize( ) if reuse_server: cur_config.set_reuse_server() - visualize_from_config( - cur_config, + visualize_model_explorer( + config=kwargs.pop("config", cur_config), no_open_in_browser=no_open_in_browser, **kwargs, ) +def visualize_model_explorer( + **kwargs, +): + """Wraps the visualize_from_config call from model_explorer.""" + visualize_from_config( + **kwargs, + ) + + def visualize_graph( graph_module: GraphModule, exported_program: ExportedProgram | EdgeProgramManager | ExecutorchProgramManager, diff --git a/docs/source/backend-arm-ethos-u-dedicated_sram.png b/docs/source/backend-arm-ethos-u-dedicated_sram.png new file mode 100644 index 00000000000..43b17266cf1 Binary files /dev/null and b/docs/source/backend-arm-ethos-u-dedicated_sram.png differ diff --git a/docs/source/backend-arm-ethos-u-shared_sram.png b/docs/source/backend-arm-ethos-u-shared_sram.png new file mode 100644 index 00000000000..d52d685d6e7 Binary files /dev/null and b/docs/source/backend-arm-ethos-u-shared_sram.png differ diff --git a/docs/source/backend-arm-ethos-u-sram_only.png b/docs/source/backend-arm-ethos-u-sram_only.png new file mode 100644 index 00000000000..bc09addbffa Binary files /dev/null and b/docs/source/backend-arm-ethos-u-sram_only.png differ diff --git a/docs/source/backend-delegates-xnnpack-reference.md b/docs/source/backend-delegates-xnnpack-reference.md index d38c5af60fa..cfb915aca59 100644 --- a/docs/source/backend-delegates-xnnpack-reference.md +++ b/docs/source/backend-delegates-xnnpack-reference.md @@ -106,9 +106,9 @@ quantizer.set_global(quantization_config) ### Quantizing your model with the XNNPACKQuantizer After configuring our quantizer, we are now ready to quantize our model ```python -from torch.export import export_for_training +from torch.export import export -exported_model = export_for_training(model_to_quantize, example_inputs).module() +exported_model = export(model_to_quantize, example_inputs).module() prepared_model = prepare_pt2e(exported_model, quantizer) print(prepared_model.graph) ``` diff --git a/docs/source/backends-arm-ethos-u.md b/docs/source/backends-arm-ethos-u.md index 71e3be105de..ae14cb9901f 100644 --- a/docs/source/backends-arm-ethos-u.md +++ b/docs/source/backends-arm-ethos-u.md @@ -50,14 +50,14 @@ compile_spec = ArmCompileSpecBuilder().ethosu_compile_spec( ).build() # Post training quantization -graph_module = torch.export.export_for_training(mobilenet_v2, example_inputs).module() +graph_module = torch.export.export(mobilenet_v2, example_inputs).module() quantizer = EthosUQuantizer(compile_spec) operator_config = get_symmetric_quantization_config(is_per_channel=False) quantizer.set_global(operator_config) graph_module = prepare_pt2e(graph_module, quantizer) graph_module(*example_inputs) graph_module = convert_pt2e(graph_module) -exported_program = torch.export.export_for_training(graph_module, example_inputs) +exported_program = torch.export.export(graph_module, example_inputs) # Lower the exported program to the Ethos-U backend and save pte file. edge_program_manager = to_edge_transform_and_lower( @@ -72,17 +72,116 @@ with open("mv2_arm_ethos_u55.pte", "wb") as file: edge_program_manager.write_to_file(file) ``` +### Ethos-U memory modes +The Ethos-U NPU provides two distinct memory interfaces: +- One interface for **low-latency, high-bandwidth memory** +Typically on-chip memory such as **SRAM**. +- One interface for **higher-latency, lower-bandwidth memory** +Typically external (off-chip) memory such as **Flash** or **DRAM**. + +On all Ethos-U NPUs(Ethos-U55, Ethos-U65, Ethos-U85), the low-latency interface is usually the SRAM of the SoC. +The external memory type depends on the SoC: +- On a low-power microcontorller, the external memory is usually Flash. +- On systems with Cortex-A and rich operating system, the external memory is typically DRAM. + +When running an inference, the Ethos-U compiler and Ethos-U driver make use of three logical memory regions: +- Ethos-U scratch buffer - a contiguous block of memory used by the NPU to store the intermediate tensors produced and consumed during inference. +- Neural Network - a contiguous block of memory holding constant data such as weights, biases, quantization parameters required to run an inference. +- Ethos-U fast scratch buffer - a contiguous block of memory, assumed to reside in on-chip memory in order to hide the higher latency/lower bandwidth of external memory. Only applicable for Ethos-U65 and Ethos-U85 on systems +with Cortex-A and the external memory is assumed to be DRAM. + +The placement of the scratch buffer and the Neural Network determine the memory mode to be used in the Ethos-U +compile specificiation. We support three different placements of the scratch buffer and the ML model. + +#### 1. Sram-Only Memory Mode +- Ethos-U scratch buffer resides in the SRAM. +- Neural Network resides in the SRAM. +- Ethos-U fast scratch buffer is not used. +- Characteristics: + - Provides the best performance since all the memory traffic passes via the low-latency/high-bandwidth memory. + - The performance uplift is especially noticeable on memory-bound workloads on the external interface. + - Available on Ethos-U55, Ethos-U65 and Ethos-U85. +- Limitations: + - Embedded SoCs often have limited SRAM and NNs are becoming larger. This memory mode may be unsuitable for a system running a big model relative to the amount of SRAM available on the SoC. +Below, you can see a visual representation of the placement of the two logical memory regions for the Sram Only configuration. + +![](backend-arm-ethos-u-sram_only.png) + +#### 2. Shared-Sram Memory Mode +- Ethos-U scratch buffer resides in the SRAM. +- Neural Network resides in the External memory. +- Ethos-U fast scratch buffer is not used. +- Characteristics: + - Intermediate tensors are stored in the SRAM, leveraging its low-latency and high-bandwidth. + - The Ethos-U compiler can prefetch weights from the external memory to the SRAM ahead of time so that when the NPU needs the data, it will already be avaialbe in the on-chip memory. + - In this mode, the external interface is Read-Only, the on-chip memory interface is Read/Write + - Shared-Sram offers great balance between performance and low SRAM usage. + - Available on Ethos-U55, Ethos-U65 and Ethos-U85. +- Limitations: + - You need to have enough space in the SRAM to hold the peak intermediate tensor. +Below, you can see a visual representation of the placement of the two logical memory regions for the Shared_Sram configuration. + +![](backend-arm-ethos-u-shared_sram.png) + +#### 3. Dedicated-Sram Memory Mode +- Ethos-U scratch buffer resides in the External memory. +- Neural Network resides in the External memory. +- Ethos-U fast scratch buffer resides in the on-chip memory. +- Characteristics: + - Used when the peak intermediate tensor is too big to fit into the on-chip memory. + - Enables silicon acceleration of large models. + - The NPU stores the results from the intermediate computations in the external memory. + - The dedicated SRAM acts as a software managed cache, improving performance by pre-fetching frequently accessed tensors to the on-chip memory. + - Available on Ethos-U65 and Ethos-U85. +- Limitations: + - The SRAM space must be dedicated exculisely to the Ethos-U(the host processor should not access it). + - Not available on Ethos-U55. +Below, you can see a visual representation of the placement of the two logical memory regions for the Shared_Sram configuration. + +![](backend-arm-ethos-u-dedicated_sram.png) + +Here is a table comparing the three memory modes: + +| Memory Mode | Ethos-U Scratch Buffer Placement | Neural Network Placement | When to Use | Trade-off | +|--------------------|----------------------------------|----------------------------|------------ |---------------------------------------------------------------------------| +| **SRAM-Only** | On-chip SRAM | On-chip SRAM | When the ML model, the Ethos-U scratch buffer and the wider software stack fit within the SRAM of the SoC | Limited by SRAM size; often not feasible for larger NNs | +| **Shared-SRAM** | On-chip SRAM | External memory (Flash/DRAM) | Most common mode on Cortex-M and Ethos-U systems; balances good performance and SRAM usage | Requires enough SRAM to hold the largest intermediate tensor | +| **Dedicated-SRAM** | External memory | External memory (Flash/DRAM) | Most common mode for Cortex-A and Ethos-U systems. For very large models where the peak intermediates cannot fit in SRAM | Need high-bandwidth external memory to deliver good performance | + + +The memory modes are defined within the [vela.ini file](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela/-/blob/main/ethosu/config_files/Arm/vela.ini?ref_type=heads). When you install +ExecuTorch for the Ethos-U backend, you automatically install the compiler containing the vela.ini file so you can directly create a compile specification with these memory modes. + +#### Interpreting the output from the Ethos-U compiler regarding the memory footprint +As part of the `to_edge_transform_and_lower` step, you will see a memory footprint information presented as: + +``` +Total SRAM used 2467.27 KiB +Total Off-chip Flash used 12.20 KiB +```` +The `Total SRAM used` indicates the peak SRAM utilization needed by the NPU in order to perform an inference. In the snippet above, the Ethos-U compiler requires 2467.27 KiB of SRAM in order to schedule the inference. +Therefore, from an application standpoint, you need to ensure you have at least 2467.27 KiB of SRAM on the SoC to run this model. The Ethos-U compiler provides a scheduling algorithm allowing to +lower the peak SRAM usage within reasonable limits, you need to add the `--optimise Size` or `--arena-cache-size` CLI options for to the compile spec. You can read more about the options of the +Ethos-U compiler in the documentation [here](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela/-/blob/main/OPTIONS.md#optimise). If the peak SRAM usage remains too high in +Shared Sram memory mode, you would need to us the Dedicated Sram mode in order to store the Neural Network and the Ethos-U scratch buffer in the external memory. +The main advantage of the Dedicated_Sram memory mode is that you can run large models and still benefit from the low-latency/high-bandwidth of the SRAM, used as a cache. + +It is important to highlight that when you specify a memory mode in the compile spec, in the runtime, the user is expected to place the scratch buffer and NN in the correct memory location. +In other words, when you specify for ex. Shared Sram memory mode, the runtime application logic should place the ethos-U scratch buffer in the on-chip memory and the NN in the external memory for optimal performance. + +You can see how we are doing this coupling between the memory mode and runtime application the [Ethos-U porting guide](../../examples/arm/ethos-u-porting-guide.md). + ### Partitioner API `EthosUPartitioner` tries to partition as much of the model as possible. It will never delegate unsupported operators, but a user can pass additional checks to the constructor to avoid partitioning additional operators. To do this, subclass `OperatorSupportBase` and implement the function `is_node_supported`. A few such checks exist in `executorch.exir.backend.operator_support`: - `DontPartition`: Don't partition operators based on operator type. - `DontPartitionModule`: Don't partition operators based on which python module the operator comes from. -- `DontPartitionName`: Don't partition opertors based on the operator name. +- `DontPartitionName`: Don't partition operators based on the operator name. ### Quantization -A fully integer model is required for using the Arm Ethos-U backend. As discussed above, you can quantize floating point models with the the `EthosUQuantizer`. Quantizers are backend specific, which means the `EthosUQuantizer` is configured to quantize models correctly for the target. +A fully integer model is required for using the Arm Ethos-U backend. As discussed above, you can quantize floating point models with the `EthosUQuantizer`. Quantizers are backend specific, which means the `EthosUQuantizer` is configured to quantize models correctly for the target. ## Runtime Integration @@ -95,4 +194,4 @@ Finally, run the elf file on FVP using the script `executorch/backends/arm/scripts/run_fvp.sh --elf=executorch/mv2_arm_ethos_u55/cmake-out/arm_executor_runner --target=ethos-u55-128`. ## See Also -- [Arm Ethos-U Backend Tutorial](tutorial-arm-ethos-u.md) +- [Arm Ethos-U Backend Tutorial](tutorial-arm.md) diff --git a/docs/source/backends-coreml.md b/docs/source/backends-coreml.md index dbf87e7d697..fe6748617a0 100644 --- a/docs/source/backends-coreml.md +++ b/docs/source/backends-coreml.md @@ -1,6 +1,6 @@ # Core ML Backend -Core ML delegate is the ExecuTorch solution to take advantage of Apple's [CoreML framework](https://developer.apple.com/documentation/coreml) for on-device ML. With CoreML, a model can run on CPU, GPU, and the Apple Neural Engine (ANE). +Core ML delegate is the ExecuTorch solution to take advantage of Apple's [Core ML framework](https://developer.apple.com/documentation/coreml) for on-device ML. With Core ML, a model can run on CPU, GPU, and the Apple Neural Engine (ANE). ## Features @@ -9,7 +9,7 @@ Core ML delegate is the ExecuTorch solution to take advantage of Apple's [CoreML ## Target Requirements -Below are the minimum OS requirements on various hardware for running a CoreML-delegated ExecuTorch model: +Below are the minimum OS requirements on various hardware for running a Core ML-delegated ExecuTorch model: - [macOS](https://developer.apple.com/macos) >= 13.0 - [iOS](https://developer.apple.com/ios/) >= 16.0 - [iPadOS](https://developer.apple.com/ipados/) >= 16.0 @@ -30,9 +30,9 @@ xcode-select --install ---- -## Using the CoreML Backend +## Using the Core ML Backend -To target the CoreML backend during the export and lowering process, pass an instance of the `CoreMLPartitioner` to `to_edge_transform_and_lower`. The example below demonstrates this process using the MobileNet V2 model from torchvision. +To target the Core ML backend during the export and lowering process, pass an instance of the `CoreMLPartitioner` to `to_edge_transform_and_lower`. The example below demonstrates this process using the MobileNet V2 model from torchvision. ```python import torch @@ -55,42 +55,140 @@ with open("mv2_coreml.pte", "wb") as file: ### Partitioner API -The CoreML partitioner API allows for configuration of the model delegation to CoreML. Passing a `CoreMLPartitioner` instance with no additional parameters will run as much of the model as possible on the CoreML backend with default settings. This is the most common use case. For advanced use cases, the partitioner exposes the following options via the [constructor](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/partition/coreml_partitioner.py#L60): +The Core ML partitioner API allows for configuration of the model delegation to Core ML. Passing a `CoreMLPartitioner` instance with no additional parameters will run as much of the model as possible on the Core ML backend with default settings. This is the most common use case. For advanced use cases, the partitioner exposes the following options via the [constructor](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/partition/coreml_partitioner.py#L60): - - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by CoreML. By default, all ops that CoreML supports will be delegated. See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation. -- `compile_specs`: A list of `CompileSpec`s for the CoreML backend. These control low-level details of CoreML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32). These are discussed more below. -- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [CoreML `MLState`](https://developer.apple.com/documentation/coreml/mlstate). If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the CoreML lowered module under the hood. Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0. + - `skip_ops_for_coreml_delegation`: Allows you to skip ops for delegation by Core ML. By default, all ops that Core ML supports will be delegated. See [here](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/test/test_coreml_partitioner.py#L42) for an example of skipping an op for delegation. +- `compile_specs`: A list of `CompileSpec`s for the Core ML backend. These control low-level details of Core ML delegation, such as the compute unit (CPU, GPU, ANE), the iOS deployment target, and the compute precision (FP16, FP32). These are discussed more below. +- `take_over_mutable_buffer`: A boolean that indicates whether PyTorch mutable buffers in stateful models should be converted to [Core ML `MLState`](https://developer.apple.com/documentation/coreml/mlstate). If set to `False`, mutable buffers in the PyTorch graph are converted to graph inputs and outputs to the Core ML lowered module under the hood. Generally, setting `take_over_mutable_buffer` to true will result in better performance, but using `MLState` requires iOS >= 18.0, macOS >= 15.0, and Xcode >= 16.0. +- `take_over_constant_data`: A boolean that indicates whether PyTorch constant data like model weights should be consumed by the Core ML delegate. If set to False, constant data is passed to the Core ML delegate as inputs. By deafault, take_over_constant_data=True. +- `lower_full_graph`: A boolean that indicates whether the entire graph must be lowered to Core ML. If set to True and Core ML does not support an op, an error is raised during lowering. If set to False and Core ML does not support an op, the op is executed on the CPU by ExecuTorch. Although setting `lower_full_graph`=False can allow a model to lower where it would otherwise fail, it can introduce performance overhead in the model when there are unsupported ops. You will see warnings about unsupported ops during lowering if there are any. By default, `lower_full_graph`=False. -#### CoreML CompileSpec + +#### Core ML CompileSpec A list of `CompileSpec`s is constructed with [`CoreMLBackend.generate_compile_specs`](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/compiler/coreml_preprocess.py#L210). Below are the available options: -- `compute_unit`: this controls the compute units (CPU, GPU, ANE) that are used by CoreML. The default value is `coremltools.ComputeUnit.ALL`. The available options from coremltools are: +- `compute_unit`: this controls the compute units (CPU, GPU, ANE) that are used by Core ML. The default value is `coremltools.ComputeUnit.ALL`. The available options from coremltools are: - `coremltools.ComputeUnit.ALL` (uses the CPU, GPU, and ANE) - `coremltools.ComputeUnit.CPU_ONLY` (uses the CPU only) - `coremltools.ComputeUnit.CPU_AND_GPU` (uses both the CPU and GPU, but not the ANE) - `coremltools.ComputeUnit.CPU_AND_NE` (uses both the CPU and ANE, but not the GPU) -- `minimum_deployment_target`: The minimum iOS deployment target (e.g., `coremltools.target.iOS18`). The default value is `coremltools.target.iOS15`. -- `compute_precision`: The compute precision used by CoreML (`coremltools.precision.FLOAT16` or `coremltools.precision.FLOAT32`). The default value is `coremltools.precision.FLOAT16`. Note that the compute precision is applied no matter what dtype is specified in the exported PyTorch model. For example, an FP32 PyTorch model will be converted to FP16 when delegating to the CoreML backend by default. Also note that the ANE only supports FP16 precision. -- `model_type`: Whether the model should be compiled to the CoreML [mlmodelc format](https://developer.apple.com/documentation/coreml/downloading-and-compiling-a-model-on-the-user-s-device) during .pte creation ([`CoreMLBackend.MODEL_TYPE.COMPILED_MODEL`](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/compiler/coreml_preprocess.py#L71)), or whether it should be compiled to mlmodelc on device ([`CoreMLBackend.MODEL_TYPE.MODEL`](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/compiler/coreml_preprocess.py#L70)). Using `CoreMLBackend.MODEL_TYPE.COMPILED_MODEL` and doing compilation ahead of time should improve the first time on-device model load time. +- `minimum_deployment_target`: The minimum iOS deployment target (e.g., `coremltools.target.iOS18`). By default, the smallest deployment target needed to deploy the model is selected. During export, you will see a warning about the "Core ML specification version" that was used for the model, which maps onto a deployment target as discussed [here](https://apple.github.io/coremltools/mlmodel/Format/Model.html#model). If you need to control the deployment target, please specify it explicitly. +- `compute_precision`: The compute precision used by Core ML (`coremltools.precision.FLOAT16` or `coremltools.precision.FLOAT32`). The default value is `coremltools.precision.FLOAT16`. Note that the compute precision is applied no matter what dtype is specified in the exported PyTorch model. For example, an FP32 PyTorch model will be converted to FP16 when delegating to the Core ML backend by default. Also note that the ANE only supports FP16 precision. +- `model_type`: Whether the model should be compiled to the Core ML [mlmodelc format](https://developer.apple.com/documentation/coreml/downloading-and-compiling-a-model-on-the-user-s-device) during .pte creation ([`CoreMLBackend.MODEL_TYPE.COMPILED_MODEL`](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/compiler/coreml_preprocess.py#L71)), or whether it should be compiled to mlmodelc on device ([`CoreMLBackend.MODEL_TYPE.MODEL`](https://github.com/pytorch/executorch/blob/14ff52ff89a89c074fc6c14d3f01683677783dcd/backends/apple/coreml/compiler/coreml_preprocess.py#L70)). Using `CoreMLBackend.MODEL_TYPE.COMPILED_MODEL` and doing compilation ahead of time should improve the first time on-device model load time. + +### Dynamic and Enumerated Shapes in Core ML Export + +When exporting an `ExportedProgram` to Core ML, **dynamic shapes** are mapped to [`RangeDim`](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#set-the-range-for-each-dimension). +This enables Core ML `.pte` files to accept inputs with varying dimensions at runtime. + +⚠️ **Note:** The Apple Neural Engine (ANE) does not support true dynamic shapes. If a model relies on `RangeDim`, Core ML will fall back to scheduling the model on the CPU or GPU instead of the ANE. + +--- + +#### Enumerated Shapes + +To enable limited flexibility on the ANE—and often achieve better performance overall—you can export models using **[enumerated shapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes)**. + +- Enumerated shapes are *not fully dynamic*. +- Instead, they define a **finite set of valid input shapes** that Core ML can select from at runtime. +- This approach allows some adaptability while still preserving ANE compatibility. + +--- + +#### Specifying Enumerated Shapes + +Unlike `RangeDim`, **enumerated shapes are not part of the `ExportedProgram` itself.** +They must be provided through a compile spec. + +For reference on how to do this, see: +- The annotated code snippet below, and +- The [end-to-end test in ExecuTorch](https://github.com/pytorch/executorch/blob/main/backends/apple/coreml/test/test_enumerated_shapes.py), which demonstrates how to specify enumerated shapes during export. + + +```python +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(10, 5) + self.linear2 = torch.nn.Linear(11, 5) + + def forward(self, x, y): + return self.linear1(x).sum() + self.linear2(y) + +model = Model() +example_inputs = ( + torch.randn((4, 6, 10)), + torch.randn((5, 11)), +) + +# Specify the enumerated shapes. Below we specify that: +# +# * x can take shape [1, 5, 10] and y can take shape [3, 11], or +# * x can take shape [4, 6, 10] and y can take shape [5, 11] +# +# Any other input shapes will result in a runtime error. +# +# Note that we must export x and y with dynamic shapes in the ExportedProgram +# because some of their dimensions are dynamic +enumerated_shapes = {"x": [[1, 5, 10], [4, 6, 10]], "y": [[3, 11], [5, 11]]} +dynamic_shapes = [ + { + 0: torch.export.Dim.AUTO(min=1, max=4), + 1: torch.export.Dim.AUTO(min=5, max=6), + }, + {0: torch.export.Dim.AUTO(min=3, max=5)}, +] +ep = torch.export.export( + model.eval(), example_inputs, dynamic_shapes=dynamic_shapes +) + +# If enumerated shapes are specified for multiple inputs, we must export +# for iOS18+ +compile_specs = CoreMLBackend.generate_compile_specs( + minimum_deployment_target=ct.target.iOS18 +) +compile_specs.append( + CoreMLBackend.generate_enumerated_shapes_compile_spec( + ep, + enumerated_shapes, + ) +) + +# When using an enumerated shape compile spec, you must specify lower_full_graph=True +# in the CoreMLPartitioner. We do not support using enumerated shapes +# for partially exported models +partitioner = CoreMLPartitioner( + compile_specs=compile_specs, lower_full_graph=True +) +delegated_program = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[partitioner], +) +et_prog = delegated_program.to_executorch() +``` + +### Backward compatibility + +Core ML supports backward compatibility via the `minimum_deployment_target` option. A model exported with a specific deployment target is guaranteed to work on all deployment targets >= the specified deployment target. For example, a model exported with `coremltools.target.iOS17` will work on iOS 17 or higher. ### Testing the Model -After generating the CoreML-delegated .pte, the model can be tested from Python using the ExecuTorch runtime Python bindings. This can be used to quickly check the model and evaluate numerical accuracy. See [Testing the Model](using-executorch-export.md#testing-the-model) for more information. +After generating the Core ML-delegated .pte, the model can be tested from Python using the ExecuTorch runtime Python bindings. This can be used to quickly check the model and evaluate numerical accuracy. See [Testing the Model](using-executorch-export.md#testing-the-model) for more information. ---- ### Quantization -To quantize a PyTorch model for the CoreML backend, use the `CoreMLQuantizer`. +To quantize a PyTorch model for the Core ML backend, use the `CoreMLQuantizer`. ### 8-bit Quantization using the PT2E Flow -Quantization with the CoreML backend requires exporting the model for iOS 17 or later. +Quantization with the Core ML backend requires exporting the model for iOS 17 or later. To perform 8-bit quantization with the PT2E flow, follow these steps: 1) Create a [`coremltools.optimize.torch.quantization.LinearQuantizerConfig`](https://apple.github.io/coremltools/source/coremltools.optimize.torch.quantization.html#coremltools.optimize.torch.quantization.LinearQuantizerConfig) and use to to create an instance of a `CoreMLQuantizer`. -2) Use `torch.export.export_for_training` to export a graph module that will be prepared for quantization. +2) Use `torch.export.export` to export a graph module that will be prepared for quantization. 3) Call `prepare_pt2e` to prepare the model for quantization. 4) Run the prepared model with representative samples to calibrate the quantizated tensor activation ranges. 5) Call `convert_pt2e` to quantize the model. @@ -126,7 +224,7 @@ static_8bit_config = ct.optimize.torch.quantization.LinearQuantizerConfig( quantizer = CoreMLQuantizer(static_8bit_config) # Step 2: Export the model for training -training_gm = torch.export.export_for_training(mobilenet_v2, sample_inputs).module() +training_gm = torch.export.export(mobilenet_v2, sample_inputs).module() # Step 3: Prepare the model for quantization prepared_model = prepare_pt2e(training_gm, quantizer) @@ -139,7 +237,7 @@ for calibration_sample in [torch.randn(1, 3, 224, 224)]: # Step 5: Convert the calibrated model to a quantized model quantized_model = convert_pt2e(prepared_model) -# Step 6: Export the quantized model to CoreML +# Step 6: Export the quantized model to Core ML et_program = to_edge_transform_and_lower( torch.export.export(quantized_model, sample_inputs), partitioner=[ @@ -173,6 +271,67 @@ Quantizing activations requires calibrating the model on representative data. A See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/main/tutorials_source/pt2e_quant_ptq.html) for more information. +### LLM quantization with quantize_ + +The Core ML backend also supports quantizing models with the [torchao](https://github.com/pytorch/ao) quantize_ API. This is most commonly used for LLMs, requiring more advanced quantization. Since quantize_ is not backend aware, it is important to use a config that is compatible with Core ML: + +* Quantize embedding/linear layers with IntxWeightOnlyConfig (with weight_dtype torch.int4 or torch.int8, using PerGroup or PerAxis granularity). Using 4-bit or PerGroup quantization requires exporting with minimum_deployment_target >= ct.target.iOS18. Using 8-bit quantization with per-axis granularity is supported on ct.target.IOS16+. See [Core ML `CompileSpec`](#coreml-compilespec) for more information on setting the deployment target. +* Quantize embedding/linear layers with CodebookWeightOnlyConfig (with dtype torch.uint1 through torch.uint8, using various block sizes). Quantizing with CodebookWeightOnlyConfig requires exporting with minimum_deployment_target >= ct.target.iOS18, see [Core ML `CompileSpec`](#coreml-compilespec) for more information on setting the deployment target. + +Below is an example that quantizes embeddings to 8-bits per-axis and linear layers to 4-bits using group_size=32 with affine quantization: + +```python +from torchao.quantization.granularity import PerGroup, PerAxis +from torchao.quantization.quant_api import ( + IntxWeightOnlyConfig, + quantize_, +) + +# Quantize embeddings with 8-bits, per channel +embedding_config = IntxWeightOnlyConfig( + weight_dtype=torch.int8, + granularity=PerAxis(0), +) +quantize_( + eager_model, + embedding_config, + lambda m, fqn: isinstance(m, torch.nn.Embedding), +) + +# Quantize linear layers with 4-bits, per-group +linear_config = IntxWeightOnlyConfig( + weight_dtype=torch.int4, + granularity=PerGroup(32), +) +quantize_( + eager_model, + linear_config, +) +``` + +Below is another example that uses codebook quantization to quantize both embeddings and linear layers to 3-bits. +In the coremltools documentation, this is called [palettization](https://apple.github.io/coremltools/docs-guides/source/opt-palettization-overview.html): + +``` +from torchao.quantization.quant_api import ( + quantize_, +) +from torchao.prototype.quantization.codebook_coreml import CodebookWeightOnlyConfig + +quant_config = CodebookWeightOnlyConfig( + dtype=torch.uint3, + # There is one LUT per 16 rows + block_size=[16, -1], +) + +quantize_( + eager_model, + quant_config, + lambda m, fqn: isinstance(m, torch.nn.Embedding) or isinstance(m, torch.nn.Linear), +) +``` + +Both of the above examples will export and lower to Core ML with the to_edge_transform_and_lower API. ---- @@ -180,7 +339,7 @@ See [PyTorch 2 Export Post Training Quantization](https://docs.pytorch.org/ao/ma To run the model on device, use the standard ExecuTorch runtime APIs. See [Running on Device](getting-started.md#running-on-device) for more information, including building the iOS frameworks. -When building from source, pass `-DEXECUTORCH_BUILD_COREML=ON` when configuring the CMake build to compile the CoreML backend. +When building from source, pass `-DEXECUTORCH_BUILD_COREML=ON` when configuring the CMake build to compile the Core ML backend. Due to the use of static initializers for registration, it may be necessary to use whole-archive to link against the `coremldelegate` target. This can typically be done by passing `"$"` to `target_link_libraries`. @@ -197,7 +356,7 @@ target_link_libraries( $) ``` -No additional steps are necessary to use the backend beyond linking the target. A CoreML-delegated .pte file will automatically run on the registered backend. +No additional steps are necessary to use the backend beyond linking the target. A Core ML-delegated .pte file will automatically run on the registered backend. --- @@ -205,7 +364,7 @@ No additional steps are necessary to use the backend beyond linking the target. ### Extracting the mlpackage -[CoreML *.mlpackage files](https://apple.github.io/coremltools/docs-guides/source/convert-to-ml-program.html#save-ml-programs-as-model-packages) can be extracted from a CoreML-delegated *.pte file. This can help with debugging and profiling for users who are more familiar with *.mlpackage files: +[Core ML *.mlpackage files](https://apple.github.io/coremltools/docs-guides/source/convert-to-ml-program.html#save-ml-programs-as-model-packages) can be extracted from a Core ML-delegated *.pte file. This can help with debugging and profiling for users who are more familiar with *.mlpackage files: ```bash python examples/apple/coreml/scripts/extract_coreml_models.py -m /path/to/model.pte ``` @@ -217,7 +376,7 @@ Note that if the ExecuTorch model has graph breaks, there may be multiple extrac ### During lowering 1. "ValueError: In op, of type [X], named [Y], the named input [Z] must have the same data type as the named input x. However, [Z] has dtype fp32 whereas x has dtype fp16." -This happens because the model is in FP16, but CoreML interprets some of the arguments as FP32, which leads to a type mismatch. The solution is to keep the PyTorch model in FP32. Note that the model will be still be converted to FP16 during lowering to CoreML unless specified otherwise in the compute_precision [CoreML `CompileSpec`](#coreml-compilespec). Also see the [related issue in coremltools](https://github.com/apple/coremltools/issues/2480). +This happens because the model is in FP16, but Core ML interprets some of the arguments as FP32, which leads to a type mismatch. The solution is to keep the PyTorch model in FP32. Note that the model will be still be converted to FP16 during lowering to Core ML unless specified otherwise in the compute_precision [Core ML `CompileSpec`](#coreml-compilespec). Also see the [related issue in coremltools](https://github.com/apple/coremltools/issues/2480). 2. coremltools/converters/mil/backend/mil/load.py", line 499, in export raise RuntimeError("BlobWriter not loaded") @@ -227,4 +386,4 @@ If you're using Python 3.13, try reducing your python version to Python 3.12. c ### At runtime 1. [ETCoreMLModelCompiler.mm:55] [Core ML] Failed to compile model, error = Error Domain=com.apple.mlassetio Code=1 "Failed to parse the model specification. Error: Unable to parse ML Program: at unknown location: Unknown opset 'CoreML7'." UserInfo={NSLocalizedDescription=Failed to par$ -This means the model requires the the CoreML opset 'CoreML7', which requires running the model on iOS >= 17 or macOS >= 14. +This means the model requires the the Core ML opset 'CoreML7', which requires running the model on iOS >= 17 or macOS >= 14. diff --git a/docs/source/backends-mps.md b/docs/source/backends-mps.md index 0dcf8b13c13..184bd88e3a7 100644 --- a/docs/source/backends-mps.md +++ b/docs/source/backends-mps.md @@ -15,8 +15,8 @@ The MPS backend device maps machine learning computational graphs and primitives * [Introduction to ExecuTorch](intro-how-it-works.md) * [Getting Started](getting-started.md) * [Building ExecuTorch with CMake](using-executorch-building-from-source.md) -* [ExecuTorch iOS Demo App](https://github.com/pytorch-labs/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) -* [ExecuTorch iOS LLaMA Demo App](llm/llama-demo-ios.md) +* [ExecuTorch iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) +* [ExecuTorch LLM iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) ::: :::: diff --git a/docs/source/backends-qualcomm.md b/docs/source/backends-qualcomm.md index fb7e9c40931..f427c7c7cea 100644 --- a/docs/source/backends-qualcomm.md +++ b/docs/source/backends-qualcomm.md @@ -431,11 +431,12 @@ For practical examples, see [`test_qnn_delegate.py`](https://github.com/pytorch/ #### Step 3: Configure Compile Specs During this step, you will need to specify the target SoC, data type, and other QNN compiler spec. ```python -from executorch.backends.qualcomm.compiler import ( +from executorch.backends.qualcomm.utils.utils import ( generate_qnn_executorch_compiler_spec, generate_htp_compiler_spec, + QcomChipset, + to_edge_transform_and_lower_to_qnn, ) -from executorch.backends.qualcomm.utils.utils import QcomChipset # HTP Compiler Configuration backend_options = generate_htp_compiler_spec( @@ -450,11 +451,6 @@ compile_spec = generate_qnn_executorch_compiler_spec( ``` #### Step 4: Lower and Export the Model ```python -from executorch.backends.qualcomm.partition.qnn_partitioner import ( - to_edge_transform_and_lower_to_qnn, -) -from executorch.exir import ExecutorchBackendConfig - # Lower to QNN backend delegated_program = to_edge_transform_and_lower_to_qnn( quantized_model if quantized else model, @@ -463,9 +459,7 @@ delegated_program = to_edge_transform_and_lower_to_qnn( ) # Export to ExecuTorch format -executorch_program = delegated_program.to_executorch( - config=ExecutorchBackendConfig(extract_delegate_segments=False) -) +executorch_program = delegated_program.to_executorch() # Save the compiled model model_name = "custom_model_qnn.pte" diff --git a/docs/source/backends-xnnpack.md b/docs/source/backends-xnnpack.md index b7fca261850..d1a120e69fa 100644 --- a/docs/source/backends-xnnpack.md +++ b/docs/source/backends-xnnpack.md @@ -79,7 +79,7 @@ Weight-only quantization is not currently supported on XNNPACK. To perform 8-bit quantization with the PT2E flow, perform the following steps prior to exporting the model: 1) Create an instance of the `XnnpackQuantizer` class. Set quantization parameters. -2) Use `torch.export.export_for_training` to prepare for quantization. +2) Use `torch.export.export` to prepare for quantization. 3) Call `prepare_pt2e` to prepare the model for quantization. 4) For static quantization, run the prepared model with representative samples to calibrate the quantizated tensor activation ranges. 5) Call `convert_pt2e` to quantize the model. @@ -103,7 +103,7 @@ qparams = get_symmetric_quantization_config(is_per_channel=True) # (1) quantizer = XNNPACKQuantizer() quantizer.set_global(qparams) -training_ep = torch.export.export_for_training(model, sample_inputs).module() # (2) +training_ep = torch.export.export(model, sample_inputs).module() # (2) prepared_model = prepare_pt2e(training_ep, quantizer) # (3) for cal_sample in [torch.randn(1, 3, 224, 224)]: # Replace with representative model inputs diff --git a/docs/source/bundled-io.md b/docs/source/bundled-io.md index 3e8accce80e..79897737268 100644 --- a/docs/source/bundled-io.md +++ b/docs/source/bundled-io.md @@ -96,7 +96,7 @@ from executorch.devtools.bundled_program.config import MethodTestCase, MethodTes from executorch.devtools.bundled_program.serialize import ( serialize_from_bundled_program_to_flatbuffer, ) -from torch.export import export, export_for_training +from torch.export import export # Step 1: ExecuTorch Program Export @@ -130,7 +130,7 @@ capture_input = ( # Export method's FX Graph. method_graph = export( - export_for_training(model, capture_input).module(), + export(model, capture_input).module(), capture_input, ) @@ -238,7 +238,7 @@ from executorch.exir import to_edge from executorch.devtools import BundledProgram from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite -from torch.export import export, export_for_training +from torch.export import export class Module(torch.nn.Module): @@ -262,7 +262,7 @@ inputs = (torch.ones(2, 2, dtype=torch.float), ) # Find each method of model needs to be traced my its name, export its FX Graph. method_graph = export( - export_for_training(model, inputs).module(), + export(model, inputs).module(), inputs, ) @@ -374,7 +374,7 @@ from executorch.exir import to_edge from executorch.devtools import BundledProgram from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite -from torch.export import export, export_for_training +from torch.export import export class Module(torch.nn.Module): @@ -398,7 +398,7 @@ inputs = (torch.ones(2, 2, dtype=torch.float),) # Find each method of model needs to be traced my its name, export its FX Graph. method_graph = export( - export_for_training(model, inputs).module(), + export(model, inputs).module(), inputs, ) diff --git a/docs/source/compiler-delegate-and-partitioner.md b/docs/source/compiler-delegate-and-partitioner.md index 4e1cb22e9d0..c633bb1fd12 100644 --- a/docs/source/compiler-delegate-and-partitioner.md +++ b/docs/source/compiler-delegate-and-partitioner.md @@ -99,7 +99,7 @@ ET_NODISCARD virtual Result init( ET_NODISCARD virtual Error execute( BackendExecutionContext& context, DelegateHandle* handle, - EValue** args); + Span args); // [optional] Runtime destroy. Destroy the resource held by the backend virtual void destroy(ET_UNUSED DelegateHandle* handle); diff --git a/docs/source/conf.py b/docs/source/conf.py index 7128e34ed8d..65845c03868 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -202,7 +202,6 @@ "export-overview": "using-executorch-export.html", "runtime-build-and-cross-compilation": "using-executorch-building-from-source.html", "tutorials/export-to-executorch-tutorial": "../using-executorch-export.html", - "running-a-model-cpp-tutorial": "using-executorch-cpp.html", "build-run-vulkan": "backends-vulkan.html", "executorch-arm-delegate-tutorial": "backends-arm-ethos-u.html", "build-run-coreml": "backends-coreml.html", diff --git a/docs/source/executorch-runtime-api-reference.rst b/docs/source/executorch-runtime-api-reference.rst index 2b4239271c1..8853e5444eb 100644 --- a/docs/source/executorch-runtime-api-reference.rst +++ b/docs/source/executorch-runtime-api-reference.rst @@ -4,7 +4,7 @@ Runtime API Reference The ExecuTorch C++ API provides an on-device execution framework for exported PyTorch models. For a tutorial style introduction to the runtime API, check out the -`runtime tutorial `__ and its `simplified `__ version. +`using executorch with cpp tutorial `__ and its `simplified `__ version. For detailed information on how APIs evolve and the deprecation process, please refer to the `ExecuTorch API Life Cycle and Deprecation Policy `__. diff --git a/docs/source/executorch_custom_versions.py b/docs/source/executorch_custom_versions.py index 29c48a337ea..590f21b10ec 100644 --- a/docs/source/executorch_custom_versions.py +++ b/docs/source/executorch_custom_versions.py @@ -7,6 +7,9 @@ """ Sphinx extension to replace ${executorch_version:TAG} with version numbers. +It also defines a special variable ${executorch_version} that is set to the value +of `EXECUTORCH_VERSION` defined in this file. + This custom extension pulls third-party version strings from files in the .ci/docker/ci_commit_pins directory, and uses them to expand specific strings in markdown files. @@ -24,10 +27,13 @@ "pytorch.txt", ] +EXECUTORCH_VERSION = "0.7.0" + variables: dict[str, str] = {} -def read_version_files(): +def populate_version_variable(): + variables["${executorch_version}"] = EXECUTORCH_VERSION cwd = os.getcwd() version_file_path = os.path.join(cwd, "..", ".ci", "docker", "ci_commit_pins") @@ -38,7 +44,7 @@ def read_version_files(): variables[var_name] = f.read().strip() -read_version_files() +populate_version_variable() def replace_variables(app, doctree, docname): diff --git a/docs/source/extension-module.md b/docs/source/extension-module.md index 24f16aa8a3a..29aa6712d37 100644 --- a/docs/source/extension-module.md +++ b/docs/source/extension-module.md @@ -2,7 +2,7 @@ **Author:** [Anthony Shoumikhin](https://github.com/shoumikhin) -In the [Running an ExecuTorch Model in C++ Tutorial](running-a-model-cpp-tutorial.md), we explored the lower-level ExecuTorch APIs for running an exported model. While these APIs offer zero overhead, great flexibility, and control, they can be verbose and complex for regular use. To simplify this and resemble PyTorch's eager mode in Python, we introduce the `Module` facade APIs over the regular ExecuTorch runtime APIs. The `Module` APIs provide the same flexibility but default to commonly used components like `DataLoader` and `MemoryAllocator`, hiding most intricate details. +In the [Detailed C++ Runtime APIs Tutorial](running-a-model-cpp-tutorial.md), we explored the lower-level ExecuTorch APIs for running an exported model. While these APIs offer zero overhead, great flexibility, and control, they can be verbose and complex for regular use. To simplify this and resemble PyTorch's eager mode in Python, we introduce the `Module` facade APIs over the regular ExecuTorch runtime APIs. The `Module` APIs provide the same flexibility but default to commonly used components like `DataLoader` and `MemoryAllocator`, hiding most intricate details. ## Example diff --git a/docs/source/getting-started.md b/docs/source/getting-started.md index be15e7d6ea2..d3d9662f5c3 100644 --- a/docs/source/getting-started.md +++ b/docs/source/getting-started.md @@ -101,7 +101,7 @@ print("Comparing against original PyTorch module") print(torch.allclose(output[0], eager_reference_output, rtol=1e-3, atol=1e-5)) ``` -For complete examples of exporting and running the model, please refer to our [examples GitHub repository](https://github.com/pytorch-labs/executorch-examples/tree/main/mv2/python). +For complete examples of exporting and running the model, please refer to our [examples GitHub repository](https://github.com/meta-pytorch/executorch-examples/tree/main/mv2/python). Additionally, if you work with Hugging Face models, the [*huggingface/optimum-executorch*](https://github.com/huggingface/optimum-executorch) library simplifies running these models end-to-end with ExecuTorch, using familiar Hugging Face APIs. Visit the repository for specific examples and supported models. @@ -124,7 +124,7 @@ To add the library to your app, add the following dependency to gradle build rul ``` # app/build.gradle.kts dependencies { - implementation("org.pytorch:executorch-android:0.6.0") + implementation("org.pytorch:executorch-android:${executorch_version}") } # See latest available versions in https://mvnrepository.com/artifact/org.pytorch/executorch-android @@ -147,7 +147,7 @@ EValue[] output = model.forward(input_evalue); float[] scores = output[0].toTensor().getDataAsFloatArray(); ``` -For a full example of running a model on Android, see the [DeepLabV3AndroidDemo](https://github.com/pytorch-labs/executorch-examples/tree/main/dl3/android/DeepLabV3Demo). For more information on Android development, including building from source, a full description of the Java APIs, and information on using ExecuTorch from Android native code, see [Using ExecuTorch on Android](using-executorch-android.md). +For a full example of running a model on Android, see the [DeepLabV3AndroidDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo). For more information on Android development, including building from source, a full description of the Java APIs, and information on using ExecuTorch from Android native code, see [Using ExecuTorch on Android](using-executorch-android.md). ### iOS @@ -214,7 +214,7 @@ if (result.ok()) { For more information on the C++ APIs, see [Running an ExecuTorch Model Using the Module Extension in C++](extension-module.md) and [Managing Tensor Memory in C++](extension-tensor.md). -For complete examples of building and running C++ application, please refer to our [examples GitHub repository](https://github.com/pytorch-labs/executorch-examples/tree/main/mv2/cpp). +For complete examples of building and running C++ application, please refer to our [examples GitHub repository](https://github.com/meta-pytorch/executorch-examples/tree/main/mv2/cpp).
diff --git a/docs/source/index.md b/docs/source/index.md index f0ec1d2c6b3..d0c9142cf4a 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -41,8 +41,8 @@ ExecuTorch provides support for: - [Quantization](quantization-overview) - [FAQs](using-executorch-faqs) #### Examples -- [Android Demo Apps](https://github.com/pytorch-labs/executorch-examples/tree/main/dl3/android/DeepLabV3Demo#executorch-android-demo-app) -- [iOS Demo Apps](https://github.com/pytorch-labs/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) +- [Android Demo Apps](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo#executorch-android-demo-app) +- [iOS Demo Apps](https://github.com/meta-pytorch/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) - [Hugging Face Models](https://github.com/huggingface/optimum-executorch/blob/main/README.md) #### Backends - [Overview](backends-overview) @@ -71,7 +71,7 @@ ExecuTorch provides support for: - [Overview](runtime-overview) - [Extension Module](extension-module) - [Extension Tensor](extension-tensor) -- [Running a Model (C++ Tutorial)](running-a-model-cpp-tutorial) +- [Detailed C++ Runtime APIs Tutorial](running-a-model-cpp-tutorial) - [Backend Delegate Implementation and Linking](runtime-backend-delegate-implementation-and-linking) - [Platform Abstraction Layer](runtime-platform-abstraction-layer) #### Portable C++ Programming @@ -147,8 +147,8 @@ using-executorch-faqs :hidden: Building an ExecuTorch Android Demo App -Building an ExecuTorch iOS Demo App -tutorial-arm-ethos-u.md +Building an ExecuTorch iOS Demo App +tutorial-arm.md ``` ```{toctree} diff --git a/docs/source/llm/export-custom-llm.md b/docs/source/llm/export-custom-llm.md index bbdf596d21b..57537ba31d8 100644 --- a/docs/source/llm/export-custom-llm.md +++ b/docs/source/llm/export-custom-llm.md @@ -41,7 +41,7 @@ import torch from executorch.exir import EdgeCompileConfig, to_edge from torch.nn.attention import sdpa_kernel, SDPBackend -from torch.export import export, export_for_training +from torch.export import export from model import GPT @@ -66,7 +66,7 @@ dynamic_shape = ( # Trace the model, converting it to a portable intermediate representation. # The torch.no_grad() call tells PyTorch to exclude training-specific logic. with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): - m = export_for_training(model, example_inputs, dynamic_shapes=dynamic_shape).module() + m = export(model, example_inputs, dynamic_shapes=dynamic_shape).module() traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape) # Convert the model into a runnable ExecuTorch program. @@ -125,7 +125,7 @@ from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower import torch from torch.export import export from torch.nn.attention import sdpa_kernel, SDPBackend -from torch.export import export_for_training +from torch.export import export from model import GPT @@ -152,7 +152,7 @@ dynamic_shape = ( # Trace the model, converting it to a portable intermediate representation. # The torch.no_grad() call tells PyTorch to exclude training-specific logic. with torch.nn.attention.sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): - m = export_for_training(model, example_inputs, dynamic_shapes=dynamic_shape).module() + m = export(model, example_inputs, dynamic_shapes=dynamic_shape).module() traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape) # Convert the model into a runnable ExecuTorch program. @@ -209,7 +209,7 @@ xnnpack_quant_config = get_symmetric_quantization_config( xnnpack_quantizer = XNNPACKQuantizer() xnnpack_quantizer.set_global(xnnpack_quant_config) -m = export_for_training(model, example_inputs).module() +m = export(model, example_inputs).module() # Annotate the model for quantization. This prepares the model for calibration. m = prepare_pt2e(m, xnnpack_quantizer) diff --git a/docs/source/llm/export-llm.md b/docs/source/llm/export-llm.md index 35f17a8aa72..462d9a51849 100644 --- a/docs/source/llm/export-llm.md +++ b/docs/source/llm/export-llm.md @@ -2,6 +2,16 @@ Instead of needing to manually write code to call torch.export(), use ExecuTorch's assortment of lowering APIs, or even interact with TorchAO quantize_ APIs for quantization, we have provided an out of box experience which performantly exports a selection of supported models to ExecuTorch. +## Prerequisites + +The LLM export functionality requires the `pytorch_tokenizers` package. If you encounter a `ModuleNotFoundError: No module named 'pytorch_tokenizers'` error, install it from the ExecutorTorch source code: + +```bash +pip install -e ./extension/llm/tokenizers/ +``` + +## Supported Models + As of this doc, the list of supported LLMs include the following: - Llama 2/3/3.1/3.2 - Qwen 2.5/3 diff --git a/docs/source/llm/getting-started.md b/docs/source/llm/getting-started.md index c75d5bbc3f5..849418342b6 100644 --- a/docs/source/llm/getting-started.md +++ b/docs/source/llm/getting-started.md @@ -23,4 +23,4 @@ Deploying LLMs to ExecuTorch can be boiled down to a two-step process: (1) expor - [Running with C++](run-with-c-plus-plus.md) - [Running on Android (XNNPack)](llama-demo-android.md) - [Running on Android (Qualcomm)](build-run-llama3-qualcomm-ai-engine-direct-backend.md) -- [Running on iOS](llama-demo-ios.md) +- [Running on iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) diff --git a/docs/source/llm/run-on-ios.md b/docs/source/llm/run-on-ios.md index 3348a03feb3..88ad94c38d3 100644 --- a/docs/source/llm/run-on-ios.md +++ b/docs/source/llm/run-on-ios.md @@ -24,7 +24,7 @@ import ExecuTorchLLM ### TextLLMRunner -The `ExecuTorchTextLLMRunner` class (bridged to Swift as `TextLLMRunner`) provides a simple Objective-C/Swift interface for loading a text-generation model, configuring its tokenizer with custom special tokens, generating token streams, and stopping execution. +The `ExecuTorchLLMTextRunner` class (bridged to Swift as `TextLLMRunner`) provides a simple Objective-C/Swift interface for loading a text-generation model, configuring its tokenizer with custom special tokens, generating token streams, and stopping execution. This API is experimental and subject to change. #### Initialization @@ -38,7 +38,7 @@ NSString *modelPath = [[NSBundle mainBundle] pathForResource:@"llama-3.2-ins NSString *tokenizerPath = [[NSBundle mainBundle] pathForResource:@"tokenizer" ofType:@"model"]; NSArray *specialTokens = @[ @"<|bos|>", @"<|eos|>" ]; -ExecuTorchTextLLMRunner *runner = [[ExecuTorchTextLLMRunner alloc] initWithModelPath:modelPath +ExecuTorchLLMTextRunner *runner = [[ExecuTorchLLMTextRunner alloc] initWithModelPath:modelPath tokenizerPath:tokenizerPath specialTokens:specialTokens]; ``` @@ -123,4 +123,4 @@ runner.stop() ## Demo -Get hands-on with our [LLaMA iOS Demo App](llama-demo-ios.md) to see the LLM runtime APIs in action. +Get hands-on with our [etLLM iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) to see the LLM runtime APIs in action. diff --git a/docs/source/llm/run-with-c-plus-plus.md b/docs/source/llm/run-with-c-plus-plus.md index 77c8990b42d..f987fcab2a5 100644 --- a/docs/source/llm/run-with-c-plus-plus.md +++ b/docs/source/llm/run-with-c-plus-plus.md @@ -251,7 +251,7 @@ Supported tokenizer formats include: 3. **TikToken**: BPE tokenizers 4. **Llama2c**: BPE tokenizers in the Llama2.c format -For custom tokenizers, you can find implementations in the [pytorch-labs/tokenizers](https://github.com/pytorch-labs/tokenizers) repository. +For custom tokenizers, you can find implementations in the [meta-pytorch/tokenizers](https://github.com/meta-pytorch/tokenizers) repository. ## Other APIs diff --git a/docs/source/pte-file-format.md b/docs/source/pte-file-format.md index 9f5757ed250..66b27fe8a4c 100644 --- a/docs/source/pte-file-format.md +++ b/docs/source/pte-file-format.md @@ -71,7 +71,10 @@ Optional extended header: | byte offset zero above. I.e., it includes these headers. | [24..31] uint64_t offset (from byte offset zero above) to the start of the | first segment, or zero if there are no segments. -| [31..?] Any zero-padding necessary to preserve the alignment of the data +| [32..39] uint64_t size of the segment data, ie. the size from the segment_base_offset +| to the end of the segments. Note, the last segment should not have any +| trailing padding. +| [40..?] Any zero-padding necessary to preserve the alignment of the data | that follows. End of optional extended header. ``` @@ -81,13 +84,16 @@ Example: Offset to flatbuffer root (0x38) | File magic ("ET??") | | Extended header magic ("eh??") - | | | Extended header size (0x18) + | | | Extended header size (0x20) vvvvvvvvvvv vvvvvvvvvvv vvvvvvvvvvv vvvvvvvvvvv -0x0000 38 00 00 00 45 54 3F 3F 65 68 3F 3F 18 00 00 00 +0x0000 38 00 00 00 45 54 3F 3F 65 68 3F 3F 20 00 00 00 0x0010 F0 02 00 00 00 00 00 00 00 10 00 00 00 00 00 00 +0x0020 20 ^^^^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^^^^ | Offset to segments (0x1000) Size of program flatbuffer data (0x2f0) + | + Segment data size (0x20) ``` ## Program data diff --git a/docs/source/running-a-model-cpp-tutorial.md b/docs/source/running-a-model-cpp-tutorial.md index 43692f49a1b..a12ef122bc8 100644 --- a/docs/source/running-a-model-cpp-tutorial.md +++ b/docs/source/running-a-model-cpp-tutorial.md @@ -1,8 +1,8 @@ -# Running an ExecuTorch Model in C++ Tutorial +# Detailed C++ Runtime APIs Tutorial **Author:** [Jacob Szwejbka](https://github.com/JacobSzwejbka) -In this tutorial, we will cover how to run an ExecuTorch model in C++ using the more detailed, lower-level APIs: prepare the `MemoryManager`, set inputs, execute the model, and retrieve outputs. However, if you’re looking for a simpler interface that works out of the box, consider trying the [Module Extension Tutorial](extension-module.md). +In this tutorial, we will cover how to run an ExecuTorch model in C++ using the more detailed, lower-level APIs: prepare the `MemoryManager`, set inputs, execute the model, and retrieve outputs. However, if you’re looking for a simpler interface that works out of the box, consider trying the [Module Extension Tutorial](extension-module.md) and [Using ExecuTorch with C++](using-executorch-cpp.md). For a high level overview of the ExecuTorch Runtime please see [Runtime Overview](runtime-overview.md), and for more in-depth documentation on each API please see the [Runtime API Reference](executorch-runtime-api-reference.rst). diff --git a/docs/source/runtime-overview.md b/docs/source/runtime-overview.md index b1aa3870dd6..96a618a2a41 100644 --- a/docs/source/runtime-overview.md +++ b/docs/source/runtime-overview.md @@ -155,6 +155,7 @@ However, please note: For more details about the ExecuTorch runtime, please see: +* [Using ExecuTorch with C++](using-executorch-cpp.md) * [Detailed Runtime APIs Tutorial](running-a-model-cpp-tutorial.md) * [Simplified Runtime APIs Tutorial](extension-module.md) * [Building from Source](using-executorch-building-from-source.md) diff --git a/docs/source/tutorial-arm-ethos-u.md b/docs/source/tutorial-arm.md similarity index 70% rename from docs/source/tutorial-arm-ethos-u.md rename to docs/source/tutorial-arm.md index 528bcd93452..0692b631154 100644 --- a/docs/source/tutorial-arm-ethos-u.md +++ b/docs/source/tutorial-arm.md @@ -1,5 +1,4 @@ - -# Arm Ethos-U Backend Tutorial +# Arm® Backend Tutorial ::::{grid} 2 @@ -13,17 +12,23 @@ :::{grid-item-card} What you will learn in this tutorial: :class-card: card-prerequisites -In this tutorial you will learn how to export a simple PyTorch model for ExecuTorch Arm Ethos-U backend delegate and run it on a Corstone FVP emulators. +In this tutorial you will learn how to export a simple PyTorch model for ExecuTorch Arm backends. ::: :::: ```{warning} -This ExecuTorch backend delegate is under active development. You may encounter some rough edges and features which may be documented or planned but not implemented. +This delegate is under active development, to get best results please use a recent version. +The TOSA and Ethos(tm) backend support is reasonably mature and used in production by some users. +The VGF backend support is in early development and you may encounter issues. +You may encounter some rough edges and features which may be documented or planned but not implemented, please refer to the in-tree documentation for the latest status of features. ``` ```{tip} -If you are already familiar with this delegate, you may want to jump directly to the examples source dir - [https://github.com/pytorch/executorch/tree/main/examples/arm](https://github.com/pytorch/executorch/tree/main/examples/arm) +If you are already familiar with this delegate, you may want to jump directly to the examples: +* [Examples in the ExecuTorch repository](https://github.com/pytorch/executorch/tree/main/examples/arm) +* [Compilation for Ethos-U](https://github.com/pytorch/executorch/blob/main/examples/arm/ethos_u_minimal_example.ipynb) +* [A commandline compiler for example models](https://github.com/pytorch/executorch/blob/main/examples/arm/aot_arm_compiler.py) ``` ## Prerequisites @@ -32,110 +37,64 @@ Let's make sure you have everything you need before you get started. ### Hardware -To successfully complete this tutorial, you will need a Linux-based host machine with Arm aarch64 or x86_64 processor architecture. +To successfully complete this tutorial, you will need a Linux or MacOS host machine with Arm aarch64 or x86_64 processor architecture. -The target device will be an embedded platform with an Arm Cortex-M CPUs and Ethos-U NPUs (ML processor). This tutorial will show you how to run PyTorch models on both. +The target device will be an emulated platform to enable development without a specific development board. This tutorial has guidance for both Ethos-U targets and VGF via the ML SDK for Vulkan®. -We will be using a [Fixed Virtual Platform (FVP)](https://www.arm.com/products/development-tools/simulation/fixed-virtual-platforms), simulating [Corstone-300](https://developer.arm.com/Processors/Corstone-300)(cs300) and [Corstone-320](https://developer.arm.com/Processors/Corstone-320)(cs320)systems. Since we will be using the FVP (think of it as virtual hardware), we won't be requiring any real embedded hardware for this tutorial. +For Ethos-U and Cortex-M, We will be using a [Fixed Virtual Platform (FVP)](https://www.arm.com/products/development-tools/simulation/fixed-virtual-platforms), simulating [Corstone-300](https://developer.arm.com/Processors/Corstone-300)(cs300) and [Corstone-320](https://developer.arm.com/Processors/Corstone-320)(cs320)systems. Since we will be using the FVP (think of it as virtual hardware), we won't be requiring any real embedded hardware for this tutorial. -### Software +For VGF we will be using the [ML SDK for Vulkan(R)](https://github.com/arm/ai-ml-sdk-for-vulkan/)) to emulate the program consumer. -First, you will need to install ExecuTorch. Please follow the recommended tutorials if you haven't already, to set up a working ExecuTorch development environment. +### Software -To generate software which can be run on an embedded platform (real or virtual), we will need a tool chain for cross-compilation and an Arm Ethos-U software development kit, including the Vela compiler for Ethos-U NPUs. +First, you will need to install ExecuTorch. Please follow the recommended tutorials if you haven't already, to set up a working ExecuTorch development environment. For the VGF backend it's recommended you [install from source](https://docs.pytorch.org/executorch/stable/using-executorch-building-from-source.html), or from a [nightly](https://download.pytorch.org/whl/nightly/executorch/). -In the following sections we will walk through the steps to download each of the dependencies listed above. +In addition to this, you need to install a number of SDK dependencies for generating Ethos-U command streams or VGF files. There are scripts which automate this, which are found in the main [ExecuTorch repository](https://github.com/pytorch/executorch/tree/main/examples/arm/). ## Set Up the Developer Environment -In this section, we will do a one-time setup, like downloading and installing necessary software, for the platform support files needed to run ExecuTorch programs in this tutorial. +In this section, we will do a one-time setup of the platform support files needed to run ExecuTorch programs in this tutorial. It is recommended to run the script in a conda or venv environment. -For that we will use the `examples/arm/setup.sh` script to pull each item in an automated fashion. It is recommended to run the script in a conda environment. +With a checkout of the ExecuTorch repository, we will use the `examples/arm/setup.sh` script to pull each item in an automated fashion. + +For Ethos-U run: ```bash -examples/arm/setup.sh --i-agree-to-the-contained-eula +./examples/arm/setup.sh --i-agree-to-the-contained-eula ``` -Upon successful execution, you can directly go to [the next step](#convert-the-pytorch-model-to-the-pte-file). - -As mentioned before, we currently support only Linux based platforms with x86_64 or aarch64 processor architecture. Let’s make sure we are indeed on a supported platform. +For VGF run: ```bash -uname -s -# Linux - -uname -m -# x86_64 or aarch64 +./examples/arm/setup.sh --i-agree-to-the-contained-eula --disable-ethos-u-deps --enable-mlsdk-deps ``` +It is possible to install both sets of dependencies if you omit the disable options. -Next we will walk through the steps performed by the `setup.sh` script to better understand the development setup. - -### Download and Set Up the Corstone-300 and Corstone-320 FVP -Fixed Virtual Platforms (FVPs) are pre-configured, functionally accurate simulations of popular system configurations. Here in this tutorial, we are interested in Corstone-300 and Corstone-320 systems. We can download this from the Arm website. +### Notes: -```{note} - By downloading and running the FVP software, you will be agreeing to the FVP [End-user license agreement (EULA)](https://developer.arm.com/downloads/-/arm-ecosystem-fvps/eula). +```{warning} +The `setup.sh` script has generated a `setup_path.sh` script that you need to source whenever you restart your shell. ``` -To download, we can either download `Corstone-300 Ecosystem FVP` and `Corstone-320 Ecosystem FVP`from [here](https://developer.arm.com/downloads/-/arm-ecosystem-fvps). or `setup.sh` script does that for you under `setup_fvp` function. - -### Download and Install the Arm GNU AArch32 Bare-Metal Toolchain - -Similar to the FVP, we would also need a tool-chain to cross-compile ExecuTorch runtime, executor-runner bare-metal application, as well as the rest of the bare-metal stack for Cortex-M55/M85 CPU available on the Corstone-300/Corstone-320 platform. - -These toolchains are available [here](https://developer.arm.com/downloads/-/arm-gnu-toolchain-downloads). We will be using GCC 13.3.rel1 targeting `arm-none-eabi` here for our tutorial. Just like FVP, `setup.sh` script will down the toolchain for you. See `setup_toolchain` function. - -### Setup the Arm Ethos-U Software Development - -This git repository is the root directory for all Arm Ethos-U software. It is to help us download required repositories and place them in a tree structure. See `setup_ethos_u` function of the setup script for more details. - -Once this is done, you should have a working FVP simulator, a functioning toolchain for cross compilation, and the Ethos-U software development setup ready for the bare-metal developement. - -### Install the Vela Compiler -Once this is done, the script will finish the setup by installing the Vela compiler for you, details are in `setup_vela` function. +i.e. run +`source executorch/examples/arm/ethos-u-scratch/setup_path.sh` -### Install the TOSA reference model -This is the last step of the setup process, using `setup_tosa_reference_model` function `setup.sh` script will install TOSA reference model for you. -At the end of the setup, if everything goes well, your top level devlopement dir might look something like this, +To confirm your environment is set up correctly and will enable you to generate .pte's for your target: +For Ethos-U run: ```bash -. -├── arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi # for x86-64 hosts -├── arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi.tar.xz -├── ethos-u -│   ├── core_platform -│   ├── core_software -│   ├── fetch_externals.py -│ └── [...] -├── FVP-corstone300 -│ ├── FVP_Corstone_SSE-300.sh -│ └── [...] -├── FVP-corstone320 -│ ├── FVP_Corstone_SSE-320.sh -│ └── [...] -├── FVP_corstone300.tgz -├── FVP_corstone320.tgz -└── setup_path.sh +# Check for Vela, which converts TOSA to Ethos-U command streams. +which vela ``` -### Notes: - -The `setup.sh` script has generated a `setup_path.sh` script that you need to source everytime you restart you shell. - -e.g. run -`source executorch/examples/arm/ethos-u-scratch/setup_path.sh` - -As `setup.sh` will download and setup the needed Arm toolchain make sure it is used by calling - -`which arm-none-eabi-gcc` - -It should show `arm-none-eabi-gcc` in the `executorch` project and not anything in `/usr/bin` something like: +For VGF run: +```bash +# Check for model-converter, which converts TOSA to ML-SDK VGF format. +which model-converter +``` -`/examples/arm/ethos-u-scratch/arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi/bin/arm-none-eabi-gcc` -or -`/examples/arm/ethos-u-scratch/arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi/bin/arm-none-eabi-gcc` +To ensure there's no environment pollution you should confirm these binaries reside within your executorch checkout, under the examples/arm tree. Other versions may present compatibility issues, so this should be corrected by modifying your environment variables such as ${PATH} appropriately. -If not you might need to uninstall `arm-none-eabi-gcc` or make sure its picked after the one in the project in your $PATH env varable. ## Convert the PyTorch Model to the `.pte` File @@ -242,27 +201,50 @@ graph_module_edge.exported_program = to_backend( Similar to the non-delegate flow, the same script will server as a helper utility to help generate the `.pte` file. Notice the `--delegate` option to enable the `to_backend` call. +For Ethos targets: ```bash python3 -m examples.arm.aot_arm_compiler --model_name="add" --delegate +# This targets the default of ethos-u55-128, see --help for further targets # should produce ./add_arm_delegate_ethos-u55-128.pte ``` -### Delegated Quantized Workflow -Generating the `.pte` file can be done using the aot_arm_compiler: +For basic post-training quantization: ```bash python3 -m examples.arm.aot_arm_compiler --model_name="mv2" --delegate --quantize +# This targets the default of ethos-u55-128, see --help for further targets # should produce ./mv2_arm_delegate_ethos-u55-128.pte ``` + +For VGF targets: +```bash +python3 -m examples.arm.aot_arm_compiler --model_name="add" --target=vgf --delegate +# should produce ./add_arm_delegate_vgf.pte +``` + +For basic post-training quantization: +```bash +python3 -m examples.arm.aot_arm_compiler --model_name="mv2" --target=vgf --delegate --quantize +# should produce ./mv2_arm_delegate_vgf.pte +``` + +To capture intermediates such as VGF for lower level integration, invoke with the "-i" option: +```bash +python3 -m examples.arm.aot_arm_compiler --model_name="mv2" --target=vgf --delegate --quantize -i ./mv2_output +# should produce ./mv2_arm_delegate_vgf.pte and intermediates in ./mv2_out/ +``` +
-At the end of this, you should have three different `.pte` files. +At the end of this, you should have a number of different `.pte` files. -- The first one contains the [SoftmaxModule](#softmaxmodule), without any backend delegates. -- The second one contains the [AddModule](#addmodule), with Arm Ethos-U backend delegate enabled. -- The third one contains the [quantized MV2Model](#mv2module), with the Arm Ethos-U backend delegate enabled as well. +- the SoftmaxModule, without any backend delegates. +- the AddModule, targeting the Arm Ethos-U backend. +- the Quantized MV2Model, targeting the Arm Ethos-U backend. +- the AddModule, targeting the VGF backend. +- the Quantized MV2Model, targeting the VGF backend. -Now let's try to run these `.pte` files on a Corstone-300 and Corstone-320 platforms in a bare-metal environment. +Now let's try to run these `.pte` files on a target. ## Getting a Bare-Metal Executable @@ -300,17 +282,13 @@ To run a `.pte` file with the Arm backend delegate call instructions, you will n - `libexecutorch_delegate_ethos_u.a` -These libraries are generated by the `backends/arm/scripts/build_executorch.sh` and `backends/arm/scripts/build_portable_kernels.sh` scripts called from the `run.sh` script. - -The `--portable_kernels` flag can be used to set the build flag `EXECUTORCH_SELECT_OPS_LIST` when running `backends/arm/scripts/build_portable_kernels.sh` that will decide the number of portable operators included in the build and are available at runtime. It must match with `.pte` file's requirements, otherwise you will get `Missing Operator` error at runtime. - -For example, there in the command line above, to run SoftmaxModule, you only included the softmax CPU operator. Similarly, to run AddModule in a non-delegated manner you will need add op and so on. As you might have already realized, for the delegated operators, which will be executed by the Arm backend delegate, you do not need to include those operators in this list. This is only for *non-delegated* operators. +These libraries are generated by the `backends/arm/scripts/build_executorch.sh` script called from the `run.sh` script. ### Building the executor_runner Bare-Metal Application The SDK dir is the same one prepared [earlier](#setup-the-arm-ethos-u-software-development). And, you will be passing the `.pte` file (any one of them) generated above. -Note, you have to generate a new `executor-runner` binary if you want to change the model or the `.pte` file. This constraint is from the constrained bare-metal runtime environment you have for Corstone-300/Corstone-320 platforms. +Note, you have to generate a new `executor-runner` binary if you want to change the model or the `.pte` file. This constraint is from the constrained bare-metal runtime environment you have for Corstone-300/Corstone-320 platforms. The build also generates a kernel registration library for the relevant operators which could not be delegated to the EthosU, see the [Kernel Library Selective Build documentation](https://docs.pytorch.org/executorch/stable/kernel-library-selective-build.html). This step is executed by the build_executor_runner.sh script, which is invoked from the run.sh in the backends/arm/scripts folder. @@ -434,6 +412,40 @@ I [executorch:arm_executor_runner.cpp:179] The `run.sh` script provides various options to select a particular FVP target, use desired models, select portable kernels and can be explored using the `--help` argument ``` +## Running on the VGF backend with the standard executor_runner for Linux + +Follow typical [Building ExecuTorch with CMake](using-executorch-building-from-source.md) flow to build the linux target, ensuring that the VGF delegate is enabled. + +```bash +-DEXECUTORCH_BUILD_VGF=ON +``` + +A full example buld line is: +``` +cmake bash \ + -DCMAKE_INSTALL_PREFIX=cmake-out \ + -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_XNNPACK=OFF \ + -DEXECUTORCH_BUILD_VULKAN=ON \ + -DEXECUTORCH_BUILD_VGF=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \ + -DPYTHON_EXECUTABLE=python \ + -Bcmake-out . +cmake --build cmake-out -j25 --target install --config Release +``` + +You can then invoke the executor runner on the host machine, which will use the VGF delegate, and requires the vulkan layer drivers we installed with setup.sh. + +```bash +./cmake-out/executor_runner -model_path add_arm_delegate_vgf.pte +``` + + ## Takeaways In this tutorial you have learnt how to use the ExecuTorch software to both export a standard model from PyTorch and to run it on the compact and fully functioned ExecuTorch runtime, enabling a smooth path for offloading models from PyTorch to Arm based platforms. diff --git a/docs/source/tutorial-xnnpack-delegate-lowering.md b/docs/source/tutorial-xnnpack-delegate-lowering.md index 04cab007f65..bccd4e4add3 100644 --- a/docs/source/tutorial-xnnpack-delegate-lowering.md +++ b/docs/source/tutorial-xnnpack-delegate-lowering.md @@ -77,13 +77,13 @@ After lowering to the XNNPACK Program, we can then prepare it for executorch and The XNNPACK delegate can also execute symmetrically quantized models. To understand the quantization flow and learn how to quantize models, refer to [Custom Quantization](quantization-custom-quantization.md) note. For the sake of this tutorial, we will leverage the `quantize()` python helper function conveniently added to the `executorch/executorch/examples` folder. ```python -from torch.export import export_for_training +from torch.export import export from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower mobilenet_v2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval() sample_inputs = (torch.randn(1, 3, 224, 224), ) -mobilenet_v2 = export_for_training(mobilenet_v2, sample_inputs).module() # 2-stage export for quantization path +mobilenet_v2 = export(mobilenet_v2, sample_inputs).module() # 2-stage export for quantization path from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( @@ -110,7 +110,7 @@ def quantize(model, example_inputs): quantized_mobilenetv2 = quantize(mobilenet_v2, sample_inputs) ``` -Quantization requires a two stage export. First we use the `export_for_training` API to capture the model before giving it to `quantize` utility function. After performing the quantization step, we can now leverage the XNNPACK delegate to lower the quantized exported model graph. From here, the procedure is the same as for the non-quantized model lowering to XNNPACK. +Quantization requires a two stage export. First we use the `export` API to capture the model before giving it to `quantize` utility function. After performing the quantization step, we can now leverage the XNNPACK delegate to lower the quantized exported model graph. From here, the procedure is the same as for the non-quantized model lowering to XNNPACK. ```python # Continued from earlier... diff --git a/docs/source/tutorials_source/export-to-executorch-tutorial.py b/docs/source/tutorials_source/export-to-executorch-tutorial.py index 2ca6a207d17..4dd2818fabe 100644 --- a/docs/source/tutorials_source/export-to-executorch-tutorial.py +++ b/docs/source/tutorials_source/export-to-executorch-tutorial.py @@ -173,8 +173,8 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # ----------------------- # # To quantize a model, we first need to capture the graph with -# ``torch.export.export_for_training``, perform quantization, and then -# call ``torch.export``. ``torch.export.export_for_training`` returns a +# ``torch.export.export``, perform quantization, and then +# call ``torch.export``. ``torch.export.export`` returns a # graph which contains ATen operators which are Autograd safe, meaning they are # safe for eager-mode training, which is needed for quantization. We will call # the graph at this level, the ``Pre-Autograd ATen Dialect`` graph. @@ -187,12 +187,10 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # will annotate the nodes in the graph with information needed to quantize the # model properly for a specific backend. -from torch.export import export_for_training +from torch.export import export example_args = (torch.randn(1, 3, 256, 256),) -pre_autograd_aten_dialect = export_for_training( - SimpleConv(), example_args, strict=True -).module() +pre_autograd_aten_dialect = export(SimpleConv(), example_args, strict=True).module() print("Pre-Autograd ATen Dialect Graph") print(pre_autograd_aten_dialect) @@ -543,7 +541,7 @@ def forward(self, a, x, b): # Here is an example for an entire end-to-end workflow: import torch -from torch.export import export, export_for_training, ExportedProgram +from torch.export import export, export, ExportedProgram class M(torch.nn.Module): @@ -557,7 +555,7 @@ def forward(self, x): example_args = (torch.randn(3, 4),) -pre_autograd_aten_dialect = export_for_training(M(), example_args, strict=True).module() +pre_autograd_aten_dialect = export(M(), example_args, strict=True).module() # Optionally do quantization: # pre_autograd_aten_dialect = convert_pt2e(prepare_pt2e(pre_autograd_aten_dialect, CustomBackendQuantizer)) aten_dialect = export(pre_autograd_aten_dialect, example_args, strict=True) diff --git a/docs/source/using-executorch-android.md b/docs/source/using-executorch-android.md index 8ac179d325d..23513302063 100644 --- a/docs/source/using-executorch-android.md +++ b/docs/source/using-executorch-android.md @@ -28,13 +28,13 @@ The AAR library can be used for generic Android device with arm64-v8a or x86_64 ExecuTorch is available on [Maven Central](https://mvnrepository.com/artifact/org.pytorch/executorch-android). -Simply add the target [`org.pytorch:executorch-android:0.6.0-rc1`](https://repo.maven.apache.org/maven2/org/pytorch/executorch-android/0.6.0-rc1/) to your Android app dependency (build.gradle), and build your app. +Simply add the target [`org.pytorch:executorch-android:${executorch_version}`](https://repo.maven.apache.org/maven2/org/pytorch/executorch-android/${executorch_version}/) to your Android app dependency (build.gradle), and build your app. For example: ``` # app/build.gradle.kts dependencies { - implementation("org.pytorch:executorch-android:0.6.0-rc1") + implementation("org.pytorch:executorch-android:${executorch_version}") } ``` @@ -53,7 +53,8 @@ You can also directly specify an AAR file in the app. We upload pre-built AAR to | Version | AAR | SHASUMS | | ------- | --- | ------- | -| [v0.6.0-rc1](https://github.com/pytorch/executorch/releases/tag/v0.6.0-rc1) | [executorch.aar](https://ossci-android.s3.amazonaws.com/executorch/release/v0.6.0-rc1/executorch.aar) | [executorch.aar.sha256sums](https://ossci-android.s3.amazonaws.com/executorch/release/v0.6.0-rc1/executorch.aar.sha256sums) | +| [${executorch_version}](https://github.com/pytorch/executorch/releases/tag/${executorch_version}) | [executorch.aar](https://ossci-android.s3.amazonaws.com/executorch/release/${executorch_version}/executorch.aar) | [executorch.aar.sha256sums](https://ossci-android.s3.amazonaws.com/executorch/release/${executorch_version}/executorch.aar.sha256sums) | +| [v0.6.0](https://github.com/pytorch/executorch/releases/tag/v0.6.0) | [executorch.aar](https://ossci-android.s3.amazonaws.com/executorch/release/v0.6.0/executorch.aar) | [executorch.aar.sha256sums](https://ossci-android.s3.amazonaws.com/executorch/release/v0.6.0/executorch.aar.sha256sums) | | [v0.5.0](https://github.com/pytorch/executorch/releases/tag/v0.5.0) | [executorch.aar](https://ossci-android.s3.amazonaws.com/executorch/release/v0.5.0-rc3/executorch.aar) | [executorch.aar.sha256sums](https://ossci-android.s3.amazonaws.com/executorch/release/v0.5.0-rc3/executorch.aar.sha256sums) | ### Snapshots from main branch @@ -90,7 +91,7 @@ implementation("com.facebook.fbjni:fbjni:0.5.1") In your app working directory, such as executorch/examples/demo-apps/android/LlamaDemo, ``` mkdir -p app/libs -curl https://ossci-android.s3.amazonaws.com/executorch/release/v0.6.0-rc1/executorch.aar -o app/libs/executorch.aar +curl https://ossci-android.s3.amazonaws.com/executorch/release/${executorch_version}/executorch.aar -o app/libs/executorch.aar ``` And include it in gradle: @@ -200,7 +201,7 @@ adb push extension/module/test/resources/add.pte /data/local/tmp/ This example loads an ExecuTorch module, prepares input data, runs inference, and processes the output data. -Please use [DeepLabV3AndroidDemo](https://github.com/pytorch-labs/executorch-examples/tree/main/dl3/android/DeepLabV3Demo) +Please use [DeepLabV3AndroidDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo) and [LlamaDemo](https://github.com/pytorch/executorch/tree/main/examples/demo-apps/android/LlamaDemo) for the code examples using ExecuTorch AAR package. diff --git a/docs/source/using-executorch-building-from-source.md b/docs/source/using-executorch-building-from-source.md index 973e9c5f55b..d48f9d26db7 100644 --- a/docs/source/using-executorch-building-from-source.md +++ b/docs/source/using-executorch-building-from-source.md @@ -7,91 +7,72 @@ like Make, Ninja or Xcode. For information, see [cmake-generators(7)](https://cm ## System Requirements ### Operating System -We've tested these instructions on the following systems, although they should -also work in similar environments. - - -Linux (x86_64) -- CentOS 8+ -- Ubuntu 20.04.6 LTS+ -- RHEL 8+ - -macOS (x86_64/ARM64) -- Big Sur (11.0)+ - -Windows (x86_64) -- Windows Subsystem for Linux (WSL) with any of the Linux options - -### Software +ExecuTorch is tested on the following systems, although it should also work in similar environments. + + * Linux (x86_64) + * CentOS 8+ + * Ubuntu 20.04.6 LTS+ + * RHEL 8+ + * macOS (x86_64/ARM64) + * Big Sur (11.0)+ + * Windows (x86_64) + * Windows Subsystem for Linux (WSL) with any of the Linux options + * Windows 10+ with Visual Studio 2022+ (experimental) + +### Software Requirements * `conda` or another virtual environment manager - - We recommend `conda` as it provides cross-language + - `conda` is recommended as it provides cross-language support and integrates smoothly with `pip` (Python's built-in package manager) - Otherwise, Python's built-in virtual environment manager `python venv` is a good alternative. * `g++` version 7 or higher, `clang++` version 5 or higher, or another C++17-compatible toolchain. * `python` version 3.10-3.12 +* `Xcode Command Line Tools` (macOS only) * `ccache` (optional) - A compiler cache that speeds up recompilation +Additional dependencies will be installed automatically when running the [Python installation](#building-the-python-package). Note that the cross-compilable core runtime code supports a wider range of toolchains, down to C++17. See the [Runtime Overview](runtime-overview.md) for portability details. ## Environment Setup - -### Clone ExecuTorch - + Clone the ExecuTorch repository from GitHub and create a conda environment as follows. Venv can be used in place on conda. ```bash - # Clone the ExecuTorch repo from GitHub - git clone -b viable/strict https://github.com/pytorch/executorch.git && cd executorch + git clone -b viable/strict https://github.com/pytorch/executorch.git + cd executorch + conda create -yn executorch python=3.10.0 + conda activate executorch ``` -### Create a Virtual Environment - -Create and activate a Python virtual environment: - ```bash - python3 -m venv .venv && source .venv/bin/activate && pip install --upgrade pip - ``` +
-Or alternatively, [install conda on your machine](https://conda.io/projects/conda/en/latest/user-guide/install/index.html). Then, create a Conda environment named "executorch". - ```bash - conda create -yn executorch python=3.10.0 && conda activate executorch - ``` +## Building the Python package + To build and install the ExecuTorch Python components, used for PTE creation and Python runtime bindings, run the following command. + This will install the ExecuTorch python package and its dependencies into the active Python environment. -## Install ExecuTorch pip package from source ```bash - # Install ExecuTorch pip package and its dependencies, as well as - # development tools like CMake, and backend support for XNNPACK and CoreML. - # If developing on a Mac, make sure to install the Xcode Command Line Tools first. - # Intel-based macOS systems require building PyTorch from source (see below) + # Install ExecuTorch pip package and its dependencies. ./install_executorch.sh ``` - See the [PyTorch instructions](https://github.com/pytorch/pytorch#installation) on how to build PyTorch from source. + The `install_executorch.sh` script supports the following flags: - Use the [`--use-pt-pinned-commit` flag](../../install_executorch.py) to install ExecuTorch with an existing PyTorch build: + * `--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. + * `--use-pt-pinned-commit`: Install the pinned PyTorch commit. When not specified, the latest PyTorch nightly build is installed. - ```bash - ./install_executorch.sh --use-pt-pinned-commit - ``` - - For Intel-based macOS systems, use the [`--use-pt-pinned-commit --minimal` flags](../../install_executorch.py): - ```bash - ./install_executorch.sh --use-pt-pinned-commit --minimal - ``` + 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). - Notice that only XNNPACK and CoreML backends are supported by default. You can enable additional backends or disable default backends by setting the corresponding CMake flags: + 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: ```bash # Enable the MPS backend CMAKE_ARGS="-DEXECUTORCH_BUILD_MPS=ON" ./install_executorch.sh ``` - ```bash - # Disable the XNNPACK backend - CMAKE_ARGS="-DEXECUTORCH_BUILD_XNNPACK=OFF" ./install_executorch.sh - ``` - - For development mode, run the command with `--editable`, which allows us to modify Python source code and see changes reflected immediately. + ### Editable Install + For development, include the `--editable` flag, which allows for local changes to ExecuTorch Python code to be reflected without a re-install. Note that when C++ files are modified, you will need to re-run the full installation to reflect the changes. ```bash ./install_executorch.sh --editable @@ -100,10 +81,8 @@ Or alternatively, [install conda on your machine](https://conda.io/projects/cond pip install -e . --no-build-isolation ``` - If C++ files are being modified, you will still have to reinstall ExecuTorch from source. - > **_WARNING:_** -> Some modules can't be imported directly in editable mode. This is a known [issue](https://github.com/pytorch/executorch/issues/9558) and we are actively working on a fix for this. To workaround this: +> Some modules can't be imported directly in editable mode. This is a known [issue](https://github.com/pytorch/executorch/issues/9558) and we are actively working on a fix for this. To work around this: > ```bash > # This will fail > python -c "from executorch.exir import CaptureConfig" @@ -129,31 +108,15 @@ Or alternatively, [install conda on your machine](https://conda.io/projects/cond > > The `--clean` command removes build artifacts, pip outputs, and also clears the ccache if it's installed, ensuring a completely fresh build environment. -## Build ExecuTorch C++ runtime from source +
-ExecuTorch's CMake build system covers the pieces of the runtime that are -likely to be useful to embedded systems users. - -- `libexecutorch.a`: The core of the ExecuTorch runtime. Does not contain any - operator/kernel definitions or backend definitions. -- `libportable_kernels.a`: The implementations of ATen-compatible operators, - following the signatures in `//kernels/portable/functions.yaml`. -- `libportable_kernels_bindings.a`: Generated code that registers the contents - of `libportable_kernels.a` with the runtime. - - NOTE: This must be linked into your application with a flag like - `-Wl,-force_load` or `-Wl,--whole-archive`. It contains load-time functions - that automatically register the kernels, but linkers will often prune those - functions by default because there are no direct calls to them. -- `executor_runner`: An example tool that runs a `.pte` program file using all - `1` values as inputs, and prints the outputs to stdout. It is linked with - `libportable_kernels.a`, so the program may use any of the operators it - implements. +## Building the C++ Runtime +The ExecuTorch C++ runtime is built using CMake. It can be compiled standalone to run examples, added as a CMake dependency, or cross-compiled for Android, iOS, or embedded platforms. -### Configure the CMake build +### Configuring -Follow these steps after cloning or pulling the upstream repo, since the build -dependencies may have changed. +Configuration should be done after cloning, pulling the upstream repo, or changing build options. Once this is done, you won't need to do it again until you pull from the upstream repo or modify any CMake-related files. ```bash # cd to the root of the executorch repo @@ -165,24 +128,79 @@ cd executorch (mkdir cmake-out && cd cmake-out && cmake ..) ``` -Once this is done, you don't need to do it again until you pull from the upstream repo again, or if you modify any CMake-related files. +### Building -### CMake build options +Build all targets with `cmake --build`. -The release build offers optimizations intended to improve performance and reduce binary size. It disables program verification and executorch logging, and adds optimizations flags. ```bash --DCMAKE_BUILD_TYPE=Release +# cd to the root of the executorch repo +cd executorch + +# Build using the configuration that you previously generated under the +# `cmake-out` directory. +# +# NOTE: The `-j` argument specifies how many jobs/processes to use when +# building, and tends to speed up the build significantly. It's typical to use +# "core count + 1" as the `-j` value. +cmake --build cmake-out -j9 ``` -To further optimize the release build for size, use both: +> **_TIP:_** For faster rebuilds, consider installing ccache (see [Compiler Cache section](#compiler-cache-ccache) above). On first builds, ccache populates its cache. Subsequent builds with the same compiler flags can be significantly faster. + +### Build Presets + +ExecuTorch provides fine-grained control over what is built, as described in [Build Options](#build-options). These options are grouped into CMake presets to cover common scenarios, while providing the ability to override individual options. Presets can be specified when configuring CMake by specifying `--preset [name]` when configuring. + +Preset values for common scenarios are listed below. Using a platform preset is recommended to avoid needing to specify many fine-grained build options. + + * `arm-baremetal` - Build for bare-metal ARM targets. + * `ios` - Build features and backends common for iOS targets. + * `macos` - Build features and backends common for Mac targets. + * `linux` - Build features and backends for Linux targets. + * `llm` - Build Large Language Model-specific features. + * `profiling` - Build the ExecuTorch runtime with profiling enabled. + * `zephyr` - Build for Zephyr RTOS. + ```bash --DCMAKE_BUILD_TYPE=Release \ --DEXECUTORCH_OPTIMIZE_SIZE=ON +# Configure the build with the ios preset. +cmake .. --preset ios ``` -#### Compiler Cache (ccache) +### CMake Targets and Libraries + +To link against the ExecuTorch framework from CMake, the following top-level targets are exposed: + + * `executorch::backends`: Contains all configured backends. + * `executorch::extensions`: Contains all configured extensions. + * `executorch::kernels`: Contains all configured kernel libraries. -ExecuTorch automatically detects and enables [ccache](https://ccache.dev/) if it's installed on your system. This significantly speeds up recompilation by caching previously compiled objects: +The backends, extensions, and kernels included in these targets are controlled by the various `EXECUTORCH_` CMake options specified by the build. Using these targets will automatically pull in the required dependencies to use the configured features. + +### Running an Example Model + +The example `executor_runner` binary can be used to run a model and sanity-check the build. Run the following commands to generate and run a simple model. +You should see the message "Model executed successfully" followed by the output values. + +``` bash +python -m examples.portable.scripts.export --model_name="add" +./cmake-out/executor_runner --model_path add.pte +``` + +``` +I 00:00:00.000526 executorch:executor_runner.cpp:82] Model file add.pte is loaded. +I 00:00:00.000595 executorch:executor_runner.cpp:91] Using method forward +I 00:00:00.000612 executorch:executor_runner.cpp:138] Setting up planned buffer 0, size 48. +I 00:00:00.000669 executorch:executor_runner.cpp:161] Method loaded. +I 00:00:00.000685 executorch:executor_runner.cpp:171] Inputs prepared. +I 00:00:00.000764 executorch:executor_runner.cpp:180] Model executed successfully. +I 00:00:00.000770 executorch:executor_runner.cpp:184] 1 outputs: +Output 0: tensor(sizes=[1], [2.]) +``` + + +### Compiler Cache (ccache) + +ExecuTorch automatically detects and enables [ccache](https://ccache.dev/) if it's installed. This significantly speeds up recompilation by caching previously compiled objects: - If ccache is detected, you'll see: `ccache found and enabled for faster builds` - If ccache is not installed, you'll see: `ccache not found, builds will not be cached` @@ -205,177 +223,223 @@ No additional configuration is needed - the build system will automatically use See [CMakeLists.txt](https://github.com/pytorch/executorch/blob/main/CMakeLists.txt) -### Build the runtime components +
-Build all targets with +## Build Options -```bash -# cd to the root of the executorch repo -cd executorch +CMake options can be used to for fine-grained control of build type, control which features are built, and configure functionality, such as logging. Options are typically specified during CMake configuration. Default values of each option are set by the active preset, but can be overridden by specifying the option when configuring. -# Build using the configuration that you previously generated under the -# `cmake-out` directory. -# -# NOTE: The `-j` argument specifies how many jobs/processes to use when -# building, and tends to speed up the build significantly. It's typical to use -# "core count + 1" as the `-j` value. -cmake --build cmake-out -j9 +Note that many build options require other options to be enabled. This may require enabling multiple options to enable a given feature. The CMake build output will provide an error message when a required option is not enabled. + +#### Build Type + +The CMake build is typically set to `Debug` or `Release`. For production use or profiling, release mode should be used to improve performance and reduce binary size. It disables program verification and executorch logging and adds optimizations flags. The `EXECUTORCH_OPTIMIZE_SIZE` flag can be used to further optimize for size with a small performance tradeoff. + +```bash +# Specify build type during CMake configuration +cmake .. -DCMAKE_BUILD_TYPE=Release ``` -> **_TIP:_** For faster rebuilds, consider installing ccache (see [Compiler Cache section](#compiler-cache-ccache) above). On first builds, ccache populates its cache. Subsequent builds with the same compiler flags can be significantly faster. +#### Backends -## Use an example binary `executor_runner` to execute a .pte file +Typically, each hardware backend exposes a CMake option to control whether the backend is built. See backend-specific documentation for more details. -First, generate a .pte file, either by exporting an example model or following -the instructions in [Model Export and Lowering](using-executorch-export.md). + * `EXECUTORCH_BUILD_CADENCE` - Build the Cadence DSP backend. + * `EXECUTORCH_BUILD_COREML` - Build the Apple CoreML backend. + * `EXECUTORCH_BUILD_CORTEX_M` - Build the ARM Cortex-M backend. + * `EXECUTORCH_BUILD_MPS` - Build the Apple Metal Performance Shader backend. + * `EXECUTORCH_BUILD_NEURON` - Build the MediaTek Neuron backend. + * `EXECUTORCH_BUILD_OPENVINO` - Build the Intel OpenVINO backend. + * `EXECUTORCH_BUILD_QNN` - Build the Qualcomm AI Engine backend. + * `EXECUTORCH_BUILD_VGF` - Build the ARM VGF backend. + * `EXECUTORCH_BUILD_VULKAN` - Build the Vulkan GPU backend. + * `EXECUTORCH_BUILD_XNNPACK` - Build the XNNPACK CPU backend. -To generate a simple model file, run the following command from the ExecuTorch directory. It -will create a file named "add.pte" in the current directory. -``` -python -m examples.portable.scripts.export --model_name="add" -``` -Then, pass it to the command line tool: ```bash -./cmake-out/executor_runner --model_path add.pte +# Build the XNNPACK and Vulkan backends. +cmake .. -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_VULKAN=ON ``` -You should see the message "Model executed successfully" followed -by the output values. +#### Extensions -``` -I 00:00:00.000526 executorch:executor_runner.cpp:82] Model file add.pte is loaded. -I 00:00:00.000595 executorch:executor_runner.cpp:91] Using method forward -I 00:00:00.000612 executorch:executor_runner.cpp:138] Setting up planned buffer 0, size 48. -I 00:00:00.000669 executorch:executor_runner.cpp:161] Method loaded. -I 00:00:00.000685 executorch:executor_runner.cpp:171] Inputs prepared. -I 00:00:00.000764 executorch:executor_runner.cpp:180] Model executed successfully. -I 00:00:00.000770 executorch:executor_runner.cpp:184] 1 outputs: -Output 0: tensor(sizes=[1], [2.]) -``` +ExecuTorch extensions provide optional functionality outside of the core runtime. As the core runtime is designed to run in constrained environments, these features are typically disabled by default. Extensions include higher-level APIs (Module and Tensor), multi-threading support (Threadpool), training, and more. -### CMake Targets + * `EXECUTORCH_BUILD_EXTENSION_APPLE` - Build the Apple extension. This provides Swift and Objective-C bindings, log routing, and platform integration with Mac and iOS. See [Using ExecuTorch on iOS](using-executorch-ios.md). + * `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER` - Build the data loader extension. Provides classes to load PTEs from files or buffers. + * `EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR` - Build the flat tensor extension. Provides functionality to load and save tensor data in .ptd format. + * `EXECUTORCH_BUILD_EXTENSION_LLM` - Build the Large Language Model extension. Provides LLM-specific functionality, such as tokenizer APIs. See [Working with LLMs](llm/getting-started.md). + * `EXECUTORCH_BUILD_EXTENSION_LLM_APPLE` - Build the Large Language Model Apple extensions. + * `EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER` - Build the Large Language Model runner extension. + * `EXECUTORCH_BUILD_EXTENSION_MODULE` - Build the Module API extension. See [High-Level APIs](using-executorch-cpp.md#high-level-apis). + * `EXECUTORCH_BUILD_EXTENSION_TENSOR` - Build the Tensor API extension. Provides convenience APIs for creating and managing tensors. See [High-Level APIs](using-executorch-cpp.md#high-level-apis) and [extension/tensor](https://github.com/pytorch/executorch/tree/main/extension/tensor). + * `EXECUTORCH_BUILD_EXTENSION_TRAINING` - Build the training extension. This is experimental. + * `EXECUTORCH_BUILD_EXTENSION_EVALUE_UTIL` - Build the EValue utility extension. Provides a method to print EValue objects. See [print_evalue.h](https://github.com/pytorch/executorch/blob/main/extension/evalue_util/print_evalue.h). + * `EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL` - Build the runner utility extension. Provides utility methods for running models, such as allocating input and output tensor memory and generating inputs. See [executor_runner.cpp](https://github.com/pytorch/executorch/blob/main/examples/portable/executor_runner/executor_runner.cpp) for example usage. -To link against the ExecuTorch framework from CMake, the following top-level targets are exposed: + ``` +# Enable the data loader extension. +cmake .. -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON + ``` - * `executorch::backends`: Contains all configured backends. - * `executorch::extensions`: Contains all configured extensions. - * `executorch::kernels`: Contains all configured kernel libraries. +#### Logging -The backends, extensions, and kernels included in these targets are controlled by the various `EXECUTORCH_` CMake options specified by the build. +Logging is enabled by default in debug builds and disabled in release. When enabled, the default log level is Info. Both log enable and level can be overriden with options. See [Logging](using-executorch-runtime-integration.md#logging). Disabling logging and decreasing log verbosity will reduce binary size by stripping unused strings from the build. -## Build ExecuTorch for Windows +* `EXECUTORCH_ENABLE_LOGGING` - Enable or disable framework log messages. +* `EXECUTORCH_LOG_LEVEL` - The minimum log level to emit. One of `debug`, `info`, `error`, or `fatal`. -This document outlines the current known working build instructions for building and validating ExecuTorch on a Windows machine. + ``` +# Enable logging at debug +cmake .. -DEXECUTORCH_ENABLE_LOGGING=ON -DEXECUTORCH_LOG_LEVEL=debug + ``` -This demo uses the -[MobileNet v2](https://pytorch.org/vision/main/models/mobilenetv2.html) model to classify images using the [XNNPACK](https://github.com/google/XNNPACK) backend. +#### Output Libraries -Note that all commands should be executed on Windows powershell in administrator mode. +To link against the runtime from outside of the CMake ecosystem, the runtime can be first built with CMake and then linked directly. A few of the relevant top-level targets are described below. Note that this is a more involved process than using CMake and is only recommended when using CMake is not viable. -### Pre-requisites +- `libexecutorch.a`: The core of the ExecuTorch runtime. Does not contain any + operator/kernel definitions or backend definitions. +- `libportable_kernels.a`: The implementations of ATen-compatible operators, + following the signatures in `//kernels/portable/functions.yaml`. +- `libportable_kernels_bindings.a`: Generated code that registers the contents + of `libportable_kernels.a` with the runtime. + - NOTE: This must be linked into your application with a flag like + `-Wl,-force_load` or `-Wl,--whole-archive`. It contains load-time functions + that automatically register the kernels, but linkers will often prune those + functions by default because there are no direct calls to them. + `libportable_kernels.a`, so the program may use any of the operators it + implements. -#### 1. Install Miniconda for Windows -Install miniconda for Windows from the [official website](https://docs.conda.io/en/latest/miniconda.html). +Backends typically introduce additional targets. See backend-specific documentation for more details. -#### 2. Install Git for Windows -Install Git for Windows from the [official website](https://git-scm.com/download/win). +
-#### 3. Install ClangCL for Windows -Install ClangCL for Windows from the [official website](https://learn.microsoft.com/en-us/cpp/build/clang-support-msbuild?view=msvc-170). +## Cross-Compiling for Android +### Pre-requisites +- Set up a Python environment and clone the ExecuTorch repository, as described in [Environment Setup](#environment-setup). +- Install the [Android SDK](https://developer.android.com/studio). Android Studio is recommended. +- Install the [Android NDK](https://developer.android.com/ndk). + - Option 1: Install via [Android Studio](https://developer.android.com/studio/projects/install-ndk). + - Option 2: Download from [NDK Downloads](https://developer.android.com/ndk/downloads). -### Create the Conda Environment -To check if conda is detected by the powershell prompt, try `conda list` or `conda --version` +### Building the AAR -If conda is not detected, you could run the powershell script for conda named `conda-hook.ps1`. -To verify that Conda is available in the in the powershell environment, run try `conda list` or `conda --version`. -If Conda is not available, run conda-hook.ps1 as follows: -```bash -$miniconda_dir\\shell\\condabin\\conda-hook.ps1 -``` -where `$miniconda_dir` is the directory where you installed miniconda -This is `“C:\Users\\AppData\Local”` by default. +With the NDK installed, the `build_android_library.sh` script will build the ExecuTorch Java AAR. This file contains the ExecuTorch Java bindings +and native code. See [Using the AAR File](using-executorch-android.md#using-aar-file) for usage. -#### Create and activate the conda environment: ```bash -conda create -yn et python=3.12 -conda activate et +export ANDROID_ABIS=arm64-v8a +export BUILD_AAR_DIR=aar-out +mkdir -p $BUILD_AAR_DIR +sh scripts/build_android_library.sh ``` -### Check Symlinks -Set the following environment variable to enable symlinks: -```bash -git config --global core.symlinks true -``` +### Building the Example Runner -### Set up ExecuTorch -Clone ExecuTorch from the [official GitHub repository](https://github.com/pytorch/executorch). +The native executor runner can be cross-compiled for android and deployed via ADB. This step is intended as +an example of CMake cross compilation and is not necessary for integration into an app. ```bash -git clone --recurse -submodules https://github.com/pytorch/executorch.git -``` +# Run the following lines from the `executorch/` folder +./install_executorch.sh --clean +mkdir cmake-android-out && cd cmake-android-out -### Run the Setup Script +# point -DCMAKE_TOOLCHAIN_FILE to the location where ndk is installed +cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a .. -Currently, there are a lot of components that are not buildable on Windows. The below instructions install a very minimal ExecuTorch which can be used as a sanity check. +cd .. +cmake --build cmake-android-out -j9 -#### Move into the `executorch` directory -```bash -cd executorch +adb shell mkdir -p /data/local/tmp/executorch +# push the binary to an Android device +adb push cmake-android-out/executor_runner /data/local/tmp/executorch +# push the model file +adb push add.pte /data/local/tmp/executorch + +adb shell "/data/local/tmp/executorch/executor_runner --model_path /data/local/tmp/executorch/add.pte" ``` -#### (Optional) Run a --clean script prior to running the .bat file. +
+ +## Cross-Compiling for iOS + +For iOS, we'll build [frameworks](https://developer.apple.com/documentation/xcode/creating-a-multi-platform-binary-framework-bundle) instead of static libraries. The frameworks contain the compiled ExecuTorch runtime and public headers. + +### Pre-requisites + +* Install Xcode from the +[Mac App Store](https://apps.apple.com/app/xcode/id497799835) and install +the Command Line Tools using the terminal. + ```bash -./install_executorch.bat --clean +xcode-select --install ``` -#### Run the setup script. -You could run the .bat file or the python script. +### Building + +1. Build the frameworks: + ```bash -./install_executorch.bat -# OR -# python install_executorch.py +./scripts/build_apple_frameworks.sh ``` -### Export MobileNet V2 +Run the above command with `--help` flag to learn more on how to build additional backends +(like [Core ML](backends-coreml.md), [MPS](backends-mps.md) or XNNPACK), etc. +Note that some backends may require additional dependencies and certain versions of Xcode and iOS. +See backend-specific documentation for more details. -Create the following script named export_mv2.py +2. Copy over the generated `.xcframework` bundles to your Xcode project, link them against +your targets and don't forget to add an extra linker flag `-all_load`. -```bash -from torchvision.models import mobilenet_v2 -from torchvision.models.mobilenetv2 import MobileNet_V2_Weights +Check out the [iOS Demo App](https://github.com/meta-pytorch/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) tutorial for more info. -mv2 = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT) # This is torch.nn.Module +
-import torch -from executorch.exir import to_edge -from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +## Building on Windows -model = mv2.eval() # turn into evaluation mode +ExecuTorch provides experimental support for native Windows builds. -example_inputs = (torch.randn((1, 3, 224, 224)),) # Necessary for exporting the model +> **_NOTE:_** All commands should be executed on Windows powershell in administrator mode. -exported_graph = torch.export.export(model, example_inputs) # Core Aten graph +### Environment Setup -edge = to_edge(exported_graph) # Edge Dialect +#### Pre-requisites -edge_delegated = edge.to_backend(XnnpackPartitioner()) # Parts of the graph are delegated to XNNPACK +1. Install miniconda for Windows from the [official website](https://docs.conda.io/en/latest/miniconda.html). +2. Install Git for Windows from the [official website](https://git-scm.com/download/win). +3. Install ClangCL for Windows from the [official website](https://learn.microsoft.com/en-us/cpp/build/clang-support-msbuild?view=msvc-170) or through a [Visual Studio](https://learn.microsoft.com/en-us/cpp/build/clang-support-msbuild?view=msvc-170) or [Visual Studio Code](https://code.visualstudio.com/docs/cpp/config-clang-mac) installation. -executorch_program = edge_delegated.to_executorch() # ExecuTorch program +#### Clone and Configure Environment -pte_path = "mv2_xnnpack.pte" +```bash +git config --global core.symlinks true +git clone --recurse -submodules https://github.com/pytorch/executorch.git +cd executorch +conda create -yn et python=3.12 +conda activate et +``` -with open(pte_path, "wb") as file: - executorch_program.write_to_file(file) # Serializing into .pte file +If Conda is not available, run conda-hook.ps1, where `$miniconda_dir` is the directory where miniconda is installed. +This is `“C:\Users\\AppData\Local”` by default. + +```bash +$miniconda_dir\\shell\\condabin\\conda-hook.ps1 ``` -#### Run the export script to create a `mv2_xnnpack.pte` file. +### Build the Python Package + +Run `install_executorch.bat` to build and install the ExecuTorch Python package and runtime bindings. ```bash -python .\\export_mv2.py +cd executorch +./install_executorch.bat ``` -### Build and Install C++ Libraries + Binaries +> **_NOTE_** Many components are not currently buildable on Windows. These instructions install a very minimal ExecuTorch which can be used as a sanity check. + +### Build the C++ Runtime + ```bash del -Recurse -Force cmake-out; ` cmake . ` @@ -395,103 +459,45 @@ cmake . ` -Bcmake-out; ` cmake --build cmake-out -j64 --target install --config Release ``` -where `$miniconda_dir` is the directory where you installed miniconda -This is `“C:\Users\\AppData\Local”` by default. - -### Run Mobilenet V2 model with XNNPACK delegation - -```bash -.\\cmake-out\\backends\\xnnpack\\Release\\xnn_executor_runner.exe --model_path=.\\mv2_xnnpack.pte -``` - -The expected output would print a tensor of size 1x1000, containing values of class scores. - -```bash -Output 0: tensor(sizes=[1, 1000], [ - -0.50986, 0.30064, 0.0953904, 0.147726, 0.231205, 0.338555, 0.206892, -0.0575775, … ]) -``` - -Congratulations! You've successfully set up ExecuTorch on your Windows device and ran a MobileNet V2 model. -Now, you can explore and enjoy the power of ExecuTorch on your own Windows device! - -## Cross compilation -Following are instruction on how to perform cross compilation for Android and iOS. +> **_NOTE_** `$miniconda_dir` is the directory where you installed miniconda. This is `“C:\Users\\AppData\Local”` by default. -### Android +### Running an Example Model -#### Building executor_runner shell binary -- Prerequisite: [Android NDK](https://developer.android.com/ndk), choose one of the following: - - Option 1: Download Android Studio by following the instructions to [install ndk](https://developer.android.com/studio/projects/install-ndk). - - Option 2: Download Android NDK directly from [here](https://developer.android.com/ndk/downloads). +To validate the installation by running a model, create a file named export_mv2.py. Then, run the powershell commands to export and run the model. +The expected output is a tensor of size 1x1000, containing class scores. -Assuming Android NDK is available, run: -```bash -# Run the following lines from the `executorch/` folder -./install_executorch.sh --clean -mkdir cmake-android-out && cd cmake-android-out - -# point -DCMAKE_TOOLCHAIN_FILE to the location where ndk is installed -cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a .. - -cd .. -cmake --build cmake-android-out -j9 - -adb shell mkdir -p /data/local/tmp/executorch -# push the binary to an Android device -adb push cmake-android-out/executor_runner /data/local/tmp/executorch -# push the model file -adb push add.pte /data/local/tmp/executorch +```py +# export_mv2.py +import torch +from executorch.exir import to_edge_transform_and_lower +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from torchvision.models import mobilenet_v2 +from torchvision.models.mobilenetv2 import MobileNet_V2_Weights -adb shell "/data/local/tmp/executorch/executor_runner --model_path /data/local/tmp/executorch/add.pte" -``` +mv2 = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT).eval() +example_inputs = (torch.randn((1, 3, 224, 224)),) -#### Building AAR for app integration from source -- Prerequisite: Android NDK from the previous section, and Android SDK (Android Studio is recommended). +program = to_edge_transform_and_lower( + torch.export.export(model, example_inputs) +).to_executorch() -Assuming Android NDK and SDK is available, run: -```bash -export ANDROID_ABIS=arm64-v8a -export BUILD_AAR_DIR=aar-out -mkdir -p $BUILD_AAR_DIR -sh scripts/build_android_library.sh +with open("mv2_xnnpack.pte", "wb") as file: + executorch_program.write_to_file(file) ``` -This script will build the AAR, which contains the Java API and its corresponding JNI library. Please see -[this documentation](using-executorch-android.md#using-aar-file) for usage. - -### iOS - -For iOS we'll build [frameworks](https://developer.apple.com/documentation/xcode/creating-a-multi-platform-binary-framework-bundle) instead of static libraries, that will also contain the public headers inside. - -1. Install Xcode from the -[Mac App Store](https://apps.apple.com/app/xcode/id497799835) and then install -the Command Line Tools using the terminal: - ```bash -xcode-select --install +python .\\export_mv2.py +.\\cmake-out\\backends\\xnnpack\\Release\\xnn_executor_runner.exe --model_path=.\\mv2_xnnpack.pte ``` -2. Build the frameworks: - ```bash -./scripts/build_apple_frameworks.sh +Output 0: tensor(sizes=[1, 1000], [ + -0.50986, 0.30064, 0.0953904, 0.147726, 0.231205, 0.338555, 0.206892, -0.0575775, … ]) ``` -Run the above command with `--help` flag to learn more on how to build additional backends -(like [Core ML](backends-coreml.md), [MPS](backends-mps.md) or XNNPACK), etc. -Note, some backends may require additional dependencies and certain versions of Xcode and iOS. - -3. Copy over the generated `.xcframework` bundles to your Xcode project, link them against -your targets and don't forget to add an extra linker flag `-all_load`. - -Check out the [iOS Demo App](https://github.com/pytorch-labs/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) tutorial for more info. - - -## Next steps - -You have successfully cross-compiled `executor_runner` binary to iOS and Android platforms. You can start exploring advanced features and capabilities. Here is a list of sections you might want to read next: +## Next Steps -* [Selective build](kernel-library-selective-build.md) to build the runtime that links to only kernels used by the program, which can provide significant binary size savings. -* Tutorials on building [Android](https://github.com/pytorch-labs/executorch-examples/tree/main/dl3/android/DeepLabV3Demo#executorch-android-demo-app) and [iOS](https://github.com/pytorch-labs/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) demo apps. +* [Selective Build](kernel-library-selective-build.md) to link only kernels used by the program. This can provide significant binary size savings. +* Tutorials on building [Android](https://github.com/meta-pytorch/executorch-examples/tree/main/dl3/android/DeepLabV3Demo#executorch-android-demo-app) and [iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/mv3/apple/ExecuTorchDemo) demo apps. * Tutorials on deploying applications to embedded devices such as [ARM Cortex-M/Ethos-U](backends-arm-ethos-u.md) and [XTensa HiFi DSP](backends-cadence.md). diff --git a/docs/source/using-executorch-cpp.md b/docs/source/using-executorch-cpp.md index b1227aec7b3..3736226bc06 100644 --- a/docs/source/using-executorch-cpp.md +++ b/docs/source/using-executorch-cpp.md @@ -32,11 +32,11 @@ if (result.ok()) { For more information on the Module class, see [Running an ExecuTorch Model Using the Module Extension in C++](extension-module.md). For information on high-level tensor APIs, see [Managing Tensor Memory in C++](extension-tensor.md). -For complete examples of building and running a C++ application using the Module API, refer to our [examples GitHub repository](https://github.com/pytorch-labs/executorch-examples/tree/main/mv2/cpp). +For complete examples of building and running a C++ application using the Module API, refer to our [examples GitHub repository](https://github.com/meta-pytorch/executorch-examples/tree/main/mv2/cpp). ## Low-Level APIs -Running a model using the low-level runtime APIs allows for a high-degree of control over memory allocation, placement, and loading. This allows for advanced use cases, such as placing allocations in specific memory banks or loading a model without a file system. For an end to end example using the low-level runtime APIs, see [Running an ExecuTorch Model in C++ Tutorial](running-a-model-cpp-tutorial.md). +Running a model using the low-level runtime APIs allows for a high-degree of control over memory allocation, placement, and loading. This allows for advanced use cases, such as placing allocations in specific memory banks or loading a model without a file system. For an end to end example using the low-level runtime APIs, see [Detailed C++ Runtime APIs Tutorial](running-a-model-cpp-tutorial.md). ## Building with CMake diff --git a/docs/source/using-executorch-export.md b/docs/source/using-executorch-export.md index da9cadf3ec2..2a887bb346d 100644 --- a/docs/source/using-executorch-export.md +++ b/docs/source/using-executorch-export.md @@ -129,14 +129,16 @@ To generate a `model.pte`, `model.ptd` pair with the weights inside `model.ptd`, ```python from executorch.exir.passes.external_constants_pass import ( - delegate_external_constants_pass, + delegate_external_constants_pass_unlifted, ) -partial_function = partial( - delegate_external_constants_pass, - ep=exported_program, +# Tag the unlifted ep.module(). +tagged_module = exported_program.module() +delegate_external_constants_pass_unlifted( + module=tagged_module, gen_tag_fn=lambda x: "model", # This is the filename the weights will be saved to. In this case, weights will be saved as "model.ptd" ) - +# Re-export to get the EP. +exported_program = export(tagged_module, inputs, dynamic_shapes=dynamic_shapes) executorch_program = to_edge_transform_and_lower( exported_program, transform_passes = [partial_function], @@ -192,7 +194,7 @@ method = program.load_method("forward") outputs = method.execute([input_tensor]) ``` -Pybindings currently does not support loading program and data. To run a model with PTE and PTD components, please use the [Extension Module](extension-module.md). There is also an E2E demo in [executorch-examples](https://github.com/pytorch-labs/executorch-examples/tree/main/program-data-separation). +Pybindings currently does not support loading program and data. To run a model with PTE and PTD components, please use the [Extension Module](extension-module.md). There is also an E2E demo in [executorch-examples](https://github.com/meta-pytorch/executorch-examples/tree/main/program-data-separation). For more information, see [Runtime API Reference](executorch-runtime-api-reference.md). diff --git a/docs/source/using-executorch-faqs.md b/docs/source/using-executorch-faqs.md index f639524d69c..d1bd0390569 100644 --- a/docs/source/using-executorch-faqs.md +++ b/docs/source/using-executorch-faqs.md @@ -14,6 +14,13 @@ sudo apt install python-dev ``` if you are using Ubuntu, or use an equivalent install command. +### ModuleNotFoundError: No module named 'pytorch_tokenizers' + +The `pytorch_tokenizers` package is required for LLM export functionality. Install it from the ExecutorTorch source code: +``` +pip install -e ./extension/llm/tokenizers/ +``` + ## Export ### Missing out variants: { _ } diff --git a/docs/source/using-executorch-ios.md b/docs/source/using-executorch-ios.md index 263f58a7dd0..3e12f174177 100644 --- a/docs/source/using-executorch-ios.md +++ b/docs/source/using-executorch-ios.md @@ -14,6 +14,7 @@ The ExecuTorch Runtime for iOS and macOS (ARM64) is distributed as a collection * `kernels_llm` - Custom kernels for LLMs * `kernels_optimized` - Accelerated generic CPU kernels * `kernels_quantized` - Quantized kernels +* `kernels_torchao` - Quantized CPU kernels from torchao Link your binary with the ExecuTorch runtime and any backends or kernels used by the exported ML model. It is recommended to link the core runtime to the components that use ExecuTorch directly, and link kernels and backends against the main app target. @@ -243,7 +244,7 @@ let imageBuffer: UnsafeMutableRawPointer = ... // Existing image buffer let inputTensor = Tensor(&imageBuffer, shape: [1, 3, 224, 224]) // Execute the 'forward' method with the given input tensor and get an output tensor back. -let outputTensor: Tensor = try module.forward(inputTensor)! +let outputTensor = try Tensor(module.forward(inputTensor)) // Copy the tensor data into logits array for easier access. let logits = outputTensor.scalars() @@ -711,7 +712,10 @@ Inputs can be any type conforming to `ValueConvertible` (like `Tensor`, `Int`, ` - `forward(_:)`: A convenient shortcut for executing the common "forward" method. The API provides overloads for single inputs, multiple inputs, or no inputs. -Outputs are always returned as an array of `Value`. + +Outputs are returned in two ways: +- As an array of `Value`s, letting you inspect and cast results yourself. +- As your expected type. The generic overloads decode the result directly into your desired Swift type (such as a single `Tensor`, an array, or any custom type conforming to the `ValueSequenceConstructible` protocol). If the output doesn’t match the expected type (e.g. multiple Values returned when a single object is expected, or a tensor data type mismatch), an invalid type error is thrown. Objective-C: @@ -777,6 +781,10 @@ do { let logits = try outputTensor.scalars() print("First 5 logits: \(logits.prefix(5))") } + + // Try casting the outputs to a single typed object. + let tensorOutput = try Tensor(module.forward(inputTensor1, inputTensor2)) + let logits = tensorOutput.scalars() } catch { print("Execution failed: \(error)") } diff --git a/examples/README.md b/examples/README.md index 3dbdacaac68..3af92f8ef90 100644 --- a/examples/README.md +++ b/examples/README.md @@ -21,6 +21,7 @@ examples | └── mps # Contains end-to-end demos of MPS backend ├── arm # Contains demos of the Arm TOSA and Ethos-U NPU flows ├── qualcomm # Contains demos of Qualcomm QNN backend +�├── samsung # Contains demos of Samsung Exynos backend ├── cadence # Contains demos of exporting and running a simple model on Xtensa DSPs ├── third-party # Third-party libraries required for working on the demos └── README.md # This file @@ -33,7 +34,7 @@ A user's journey may commence by exploring the demos located in the [`portable/` ## Demos Apps -Explore mobile apps with ExecuTorch models integrated and deployable on [Android](demo-apps/android) and [iOS](demo-apps/apple_ios). This provides end-to-end instructions on how to export Llama models, load on device, build the app, and run it on device. +Explore mobile apps with ExecuTorch models integrated and deployable on [Android](demo-apps/android) and [iOS](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple). This provides end-to-end instructions on how to export Llama models, load on device, build the app, and run it on device. For specific details related to models and backend, you can explore the various subsections. @@ -69,6 +70,10 @@ The [`arm`](arm) directory contains scripts to help you run a PyTorch model on a You will find demos of [ExecuTorch QNN Backend](qualcomm) in the [`qualcomm`](qualcomm) directory. +### Exynos Backend + +You will find demos of [ExecuTorch Exynos Backend](samsung) in the [`samsung`](samsung) directory. + ### Cadence HiFi4 DSP The [`Cadence`](cadence) directory hosts a demo that showcases the process of exporting and executing a model on Xtensa Hifi4 DSP. You can utilize [this tutorial](../docs/source/backends-cadence.md) to guide you in configuring the demo and running it. diff --git a/examples/apple/coreml/scripts/build_executor_runner.sh b/examples/apple/coreml/scripts/build_executor_runner.sh index 00ee9821201..4da897bda96 100755 --- a/examples/apple/coreml/scripts/build_executor_runner.sh +++ b/examples/apple/coreml/scripts/build_executor_runner.sh @@ -93,7 +93,7 @@ find "$CMAKE_BUILD_DIR_PATH/" -name 'libcoreml_inmemoryfs.a' -exec cp -f "{}" "$ find "$CMAKE_BUILD_DIR_PATH/" -name 'libcoremldelegate.a' -exec cp -f "{}" "$LIBRARIES_DIR_PATH/libcoremldelegate.a" \; find "$CMAKE_BUILD_DIR_PATH/" -name 'libportable_ops_lib.a' -exec cp -f "{}" "$LIBRARIES_DIR_PATH/libportable_ops_lib.a" \; find "$CMAKE_BUILD_DIR_PATH/" -name 'libportable_kernels.a' -exec cp -f "{}" "$LIBRARIES_DIR_PATH/libportable_kernels.a" \; -cp -f "$CMAKE_BUILD_DIR_PATH/third-party/flatcc_external_project/lib/libflatccrt.a" "$LIBRARIES_DIR_PATH/libflatccrt.a" +cp -f "$CMAKE_BUILD_DIR_PATH/third-party/flatcc_ep/lib/libflatccrt.a" "$LIBRARIES_DIR_PATH/libflatccrt.a" # Build the runner echo "ExecuTorch: Building runner" diff --git a/examples/apple/coreml/scripts/export.py b/examples/apple/coreml/scripts/export.py index 0b5f64d13c2..e7756fa49ae 100644 --- a/examples/apple/coreml/scripts/export.py +++ b/examples/apple/coreml/scripts/export.py @@ -223,7 +223,6 @@ def main(): pte_base_name = get_pte_base_name(args) if args.use_partitioner: model = model.eval() - assert not args.generate_etrecord, "ETRecord is not supported with partitioner" ep = torch.export.export( model, args=example_args, @@ -234,9 +233,12 @@ def main(): delegated_program = exir.to_edge_transform_and_lower( ep, partitioner=[CoreMLPartitioner(compile_specs=compile_specs)], + generate_etrecord=args.generate_etrecord, ) exec_program = delegated_program.to_executorch() save_pte_program(exec_program, pte_base_name) + if args.generate_etrecord: + exec_program.get_etrecord().save(f"{pte_base_name}_coreml_etrecord.bin") if args.run_with_pybindings: run_with_pybindings( executorch_program=exec_program, diff --git a/examples/apple/mps/CMakeLists.txt b/examples/apple/mps/CMakeLists.txt index 3f61cedec8e..8a562dd206b 100644 --- a/examples/apple/mps/CMakeLists.txt +++ b/examples/apple/mps/CMakeLists.txt @@ -76,16 +76,10 @@ if(NOT CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") ) # - # The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. + # The `__srcs` lists are defined by executorch_load_build_variables. # - set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../../executorch_srcs.cmake" - ) - - extract_sources(${EXECUTORCH_SRCS_FILE}) - + executorch_load_build_variables() set(_mps_schema_headers ${CMAKE_BINARY_DIR}/../../../schema/include/) - include(${EXECUTORCH_SRCS_FILE}) target_include_directories( bundled_program INTERFACE ${CMAKE_CURRENT_BINARY_DIR}/../../../devtools/include diff --git a/examples/apple/mps/scripts/build_mps_executor_runner.sh b/examples/apple/mps/scripts/build_mps_executor_runner.sh index 625bc08a663..5d4e087d19e 100755 --- a/examples/apple/mps/scripts/build_mps_executor_runner.sh +++ b/examples/apple/mps/scripts/build_mps_executor_runner.sh @@ -38,7 +38,7 @@ done rm -rf "$OUTPUT" -cmake -DBUCK2="$BUCK" \ +cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE="$MODE" \ -DEXECUTORCH_BUILD_DEVTOOLS=ON \ diff --git a/examples/apple/mps/scripts/mps_example.py b/examples/apple/mps/scripts/mps_example.py index 42ea79435ed..46e0f6af242 100644 --- a/examples/apple/mps/scripts/mps_example.py +++ b/examples/apple/mps/scripts/mps_example.py @@ -170,9 +170,7 @@ def parse_args(): # pre-autograd export. eventually this will become torch.export with torch.no_grad(): - model = torch.export.export_for_training( - model, example_inputs, strict=True - ).module() + model = torch.export.export(model, example_inputs, strict=True).module() edge: EdgeProgramManager = export_to_edge( model, example_inputs, diff --git a/examples/apple/mps/test_mps.sh b/examples/apple/mps/test_mps.sh index bca28628473..2d0507fcf56 100755 --- a/examples/apple/mps/test_mps.sh +++ b/examples/apple/mps/test_mps.sh @@ -15,7 +15,7 @@ cmake_install_executorch_devtools_lib() { echo "Installing libexecutorch.a, libportable_kernels.a, libetdump.a, libbundled_program.a" rm -rf cmake-out - retry cmake -DBUCK2="$BUCK" \ + retry cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_DEVTOOLS=ON \ @@ -56,11 +56,5 @@ then PYTHON_EXECUTABLE=python3 fi -if [[ -z $BUCK ]]; -then - BUCK=buck2 -fi - - cmake_install_executorch_devtools_lib test_cmake_mps diff --git a/examples/arm/CMakeLists.txt b/examples/arm/CMakeLists.txt deleted file mode 100644 index 58466faeca5..00000000000 --- a/examples/arm/CMakeLists.txt +++ /dev/null @@ -1,66 +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. - -# Kernel library for portable kernels. Please this file formatted by running: -# ~~~ -# cmake-format -i CMakeLists.txt -# ~~~ - -cmake_minimum_required(VERSION 3.19) -project(arm_example) - -# Option to register op list -option(EXECUTORCH_SELECT_OPS_LIST "Register the following list of ops" OFF) - -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) -# Source root directory for executorch. -if(NOT EXECUTORCH_ROOT) - set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) -endif() - -include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) - -if(NOT PYTHON_EXECUTABLE) - resolve_python_executable() -endif() - -set(_common_compile_options -Wno-deprecated-declarations -fPIC) - -add_compile_options("-Wall" "-Werror") - -# Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/..) - -find_package(executorch CONFIG REQUIRED HINTS ${CMAKE_INSTALL_PREFIX}) -target_include_directories(executorch INTERFACE ${_common_include_directories}) - -include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) -include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) - -# Generate C++ bindings to register kernels into both PyTorch (for AOT) and -# Executorch (for runtime). Here select all ops in functions.yaml -gen_selected_ops( - LIB_NAME - "arm_portable_ops_lib" - OPS_SCHEMA_YAML - "" - ROOT_OPS - "${EXECUTORCH_SELECT_OPS_LIST}" - INCLUDE_ALL_OPS - "" -) -generate_bindings_for_kernels( - LIB_NAME "arm_portable_ops_lib" FUNCTIONS_YAML - ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml -) -gen_operators_lib( - LIB_NAME "arm_portable_ops_lib" KERNEL_LIBS portable_kernels DEPS executorch -) - -if(EXECUTORCH_ENABLE_EVENT_TRACER) - target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) - target_compile_options(portable_ops_lib INTERFACE -DET_EVENT_TRACER_ENABLED) -endif() diff --git a/examples/arm/README.md b/examples/arm/README.md index a326db70e64..9cce33bdade 100644 --- a/examples/arm/README.md +++ b/examples/arm/README.md @@ -1,36 +1,127 @@ -## ExecuTorch on ARM Cortex-M55 + Ethos-U55 +## ExecuTorch for Arm backends Ethos-U, VGF and Cortex-M -This dir contains scripts to help you prepare setup needed to run a PyTorch -model on an ARM Corstone-300 platform via ExecuTorch. Corstone-300 platform -contains the Cortex-M55 CPU and Ethos-U55 NPU. +This project contains scripts to help you setup and run a PyTorch +model on a Arm backend via ExecuTorch. This backend supports Ethos-U and VGF as +targets (using TOSA) but you can also use the Ethos-U example runner as an example +on Cortex-M if you do not delegate the model. + +The main scripts are `setup.sh`, `run.sh` and `aot_arm_compiler.py`. + +`setup.sh` will install the needed tools and with --root-dir +you can change the path to a scratch folder where it will download and generate build +artifacts. If supplied, you must also supply the same folder to run.sh with +--scratch-dir= If not supplied both script will use examples/arm/ethos-u-scratch + +`run.sh` can be used to build, run and test a model in an easy way and it will call cmake for you +and in cases you want to run a simulator it will start it also. The script will call `aot_arm_compiler.py` +to convert a model and include it in the build/run. + +Build and test artifacts are by default placed under the folder arm_test folder +this can be changed with --et_build_root= + +`aot_arm_compiler.py` is used to convert a Python model or a saved .pt model to a PTE file and is used by `run.sh` +and other test script but can also be used directly. + +If you prefer to use the ExecuTorch API, there is also the `ethos_u_minimal_example.ipynb` notebook example. +This shows the workflow if you prefer to integrate a python torch.export and ExecuTorch flow directly into your +model codebase. This is particularly useful if you want to perform more complex training, such as quantization +aware training using the ArmQuantizer. + +## Create a PTE file for Arm backends + +There is an easy to use example flow to compile your PyTorch model to a PTE file for the Arm backend called `aot_arm_compiler.py` +that you can use to generate PTE files, it can generate PTE files for the supported targets `-t` or even non delegated (Cortex-M) +using different memory modes and can both use a python file as input or just use the models from examples/models with `--model_input`. +It also supports generating Devtools artifacts like BundleIO BPTE files, and ETRecords. Run it with `--help` to check its capabilities. + +You point out the model to convert with `--model_name=` It supports running a model from examples/models or models +from a python file if you just specify `ModelUnderTest` and `ModelInput` in it. + +``` +$ python3 -m examples.arm.aot_arm_compiler --help +``` + +This is how you generate a BundleIO BPTE of a simple add example + +``` +$ python3 -m examples.arm.aot_arm_compiler --model_name=examples/arm/example_modules/add.py --target=ethos-u55-128 --bundleio +``` + +The example model used has added two extra variables that is picked up to make this work. + +`ModelUnderTest` should be a `torch.nn.module` instance. + +`ModelInputs` should be a tuple of inputs to the forward function. + + +You can also use the models from example/models directly by just using the short name e.g. + +``` +$ python3 -m examples.arm.aot_arm_compiler --model_name=mv2 --target=ethos-u55-64 +``` + + +The `aot_arm_compiler.py` is called from the scripts below so you don't need to, but it can be useful to do by hand in some cases. + + +## ExecuTorch on Arm Ethos-U55/U65 and U85 + +This example code will help you get going with the Corstone™-300/320 platforms and +run on the FVP and can be used a a starting guide in your porting to your board/HW We will start from a PyTorch model in python, export it, convert it to a `.pte` file - A binary format adopted by ExecuTorch. Then we will take the `.pte` model file and embed that with a baremetal application executor_runner. We will then take the executor_runner file, which contains not only the `.pte` binary but also necessary software components to run standalone on a baremetal system. -Lastly, we will run the executor_runner binary on a Corstone-300 FVP Simulator -platform. +The build flow will pick up the non delegated ops from the generated PTE file and +add CPU implementation of them. +Lastly, we will run the executor_runner binary on a Corstone™-300/320 FVP Simulator platform. + ### Example workflow -There are two main scripts, setup.sh and run.sh. Each takes one optional, -positional argument. It is a path to a scratch dir to download and generate -build artifacts. If supplied, the same argument must be supplied to both the scripts. +Below is example workflow to build an application for Ethos-U55/85. The script below requires an internet connection: -To run these scripts. On a Linux system, in a terminal, with a working internet connection, ``` # Step [1] - setup necessary tools $ cd -$ executorch/examples/arm/setup.sh --i-agree-to-the-contained-eula [optional-scratch-dir] +$ ./examples/arm/setup.sh --i-agree-to-the-contained-eula + +# Step [2] - Setup path to tools, The `setup.sh` script has generated a script that you need to source every time you restart you shell. +$ source examples/arm/ethos-u-scratch/setup_path.sh + +# Step [3] - build and run ExecuTorch and executor_runner baremetal example application +# on a Corstone(TM)-320 FVP to run a simple PyTorch model from a file. +$ ./examples/arm/run.sh --model_name=examples/arm/example_modules/add.py --target=ethos-u85-128 +``` + +The argument `--model_name=` is passed to `aot_arm_compiler.py` so you can use it in the same way +e.g. you can also use the models from example/models directly in the same way as above. + +``` +$ ./examples/arm/run.sh --model_name=mv2 --target=ethos-u55-64 +``` + +The runner will by default set all inputs to "1" and you are supposed to add/change the code +handling the input for your hardware target to give the model proper input, maybe from your camera +or mic hardware. + +While testing you can use the --bundleio flag to use the input from the python model file and +generate a .bpte instead of a .pte file. This will embed the input example data and reference output +in the bpte file/data, which is used to verify the model's output. You can also use --etdump to generate +an ETRecord and a ETDump trace files from your target (they are printed as base64 strings in the serial log). -# Step [2] - Setup Patch to tools, The `setup.sh` script has generated a script that you need to source everytime you restart you shell. -$ source executorch/examples/arm/ethos-u-scratch/setup_path.sh +Just keep in mind that CPU cycles are NOT accurate on the FVP simulator and it can not be used for +performance measurements, so you need to run on FPGA or actual ASIC to get good results from --etdump. +As a note the printed NPU cycle numbers are still usable and closer to real values if the timing +adaptor is setup correctly. -# Step [3] - build + run ExecuTorch and executor_runner baremetal application -# suited for Corstone FVP's to run a simple PyTorch model. -$ executorch/examples/arm/run.sh --model_name=mv2 --target=ethos-u85-128 [--scratch-dir=same-optional-scratch-dir-as-before] ``` +# Build + run with BundleIO and ETDump +$ ./examples/arm/run.sh --model_name=lstm --target=ethos-u85-128 --bundleio --etdump +``` + ### Ethos-U minimal example @@ -42,6 +133,19 @@ pip install jupyter jupyter notebook ethos_u_minimal_example.ipynb ``` +## ExecuTorch on ARM Cortex-M + +For Cortex-M you run the script without delegating e.g `--no_delegate` as the build flow already supports picking up +the non delegated ops from the generated PTE file and add CPU implementation of them this will work out of the box in +most cases. + +To run mobilenet_v2 on the Cortex-M55 only, without using the Ethos-U try this: + +``` +$ ./examples/arm/run.sh --model_name=mv2 --target=ethos-u55-128 --no_delegate +``` + + ### Online Tutorial We also have a [tutorial](https://pytorch.org/executorch/main/backends-arm-ethos-u) explaining the steps performed in these diff --git a/examples/arm/aot_arm_compiler.py b/examples/arm/aot_arm_compiler.py index 5f3eb60c44f..d7e1b64e3ca 100644 --- a/examples/arm/aot_arm_compiler.py +++ b/examples/arm/aot_arm_compiler.py @@ -8,6 +8,7 @@ # Example script for exporting simple models to flatbuffer import argparse +import copy import json import logging import os @@ -19,7 +20,6 @@ from examples.devtools.scripts.export_bundled_program import save_bundled_program from executorch.backends.arm.arm_backend import ( ArmCompileSpecBuilder, - get_tosa_spec, is_ethosu, is_tosa, is_vgf, @@ -31,20 +31,27 @@ TOSAQuantizer, VgfQuantizer, ) -from executorch.backends.arm.tosa_partitioner import TOSAPartitioner -from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.backends.arm.tosa.specification import get_tosa_spec from executorch.backends.arm.util.arm_model_evaluator import ( GenericModelEvaluator, MobileNetV2Evaluator, ) -from executorch.backends.arm.vgf_partitioner import VgfPartitioner +from executorch.backends.arm.vgf import VgfPartitioner # To use Cortex-M backend +from executorch.backends.cortex_m.passes.quantized_op_fusion_pass import ( + QuantizedOpFusionPass, +) + from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import ( ReplaceQuantNodesPass, ) + +from executorch.devtools import generate_etrecord from executorch.devtools.backend_debug import get_delegation_info from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite @@ -341,7 +348,6 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): "ethos-u85-1024", "ethos-u85-2048", "vgf", - "TOSA-0.80+BI", "TOSA-1.0+INT", "TOSA-1.0+FP", ] @@ -393,7 +399,7 @@ def get_compile_spec( try: tosa_spec = TosaSpecification.create_from_string(target) except: - tosa_spec = TosaSpecification.create_from_string("TOSA-0.80+BI") + tosa_spec = TosaSpecification.create_from_string("TOSA-1.0+INT") spec_builder = ArmCompileSpecBuilder().tosa_compile_spec(tosa_spec) elif "ethos-u" in target: spec_builder = ArmCompileSpecBuilder().ethosu_compile_spec( @@ -506,6 +512,13 @@ def get_args(): default=False, help="Flag for producing BundleIO bpte file with input/output test/ref data.", ) + parser.add_argument( + "--etrecord", + action="store_true", + required=False, + default=False, + help="Flag for producing a etrecord file.", + ) parser.add_argument( "-t", "--target", @@ -588,6 +601,11 @@ def get_args(): action="store_false", help="Disable strict checking while exporting models.", ) + parser.add_argument( + "--enable_qdq_fusion_pass", + action="store_true", + help="Enable the QuantizedOpFusionPass fusion step", + ) args = parser.parse_args() if args.evaluate and ( @@ -703,7 +721,7 @@ def quantize_model(args, model: torch.nn.Module, example_inputs, compile_spec): args.evaluate_config, ) # Wrap quantized model back into an exported_program - exported_program = torch.export.export_for_training( + exported_program = torch.export.export( model_int8, example_inputs, strict=args.strict_export ) @@ -779,11 +797,21 @@ def to_edge_no_delegate(exported_program, args, model: torch.nn.Module, example_ return model_int8, edge -def transform_for_cortex_m_backend(edge): +def transform_for_cortex_m_backend(edge, args): # Let's make sure we are using optimized Cortex M backend # NB: If we can't find and replace ops those are expected to be replaced, # bad things will happen at runtime, like "missing operator" errors! - edge = edge.transform([ReplaceQuantNodesPass()]) + + # Instantiate the mandatory ReplaceQuantNodesPass + passes = [ReplaceQuantNodesPass()] + + # Conditionally add the QuantizedOpFusionPass + if args.enable_qdq_fusion_pass: + passes.append(QuantizedOpFusionPass()) + + # Apply the passes + edge = edge.transform(passes) + return edge @@ -796,9 +824,9 @@ def transform_for_cortex_m_backend(edge): ) model = original_model.eval() - # export_for_training under the assumption we quantize, the exported form also works + # export under the assumption we quantize, the exported form also works # in to_edge if we don't quantize - exported_program = torch.export.export_for_training( + exported_program = torch.export.export( model, example_inputs, strict=args.strict_export ) model = exported_program.module() @@ -818,11 +846,14 @@ def transform_for_cortex_m_backend(edge): exported_program, args, model, example_inputs ) - # Transform so we can use ops from the Cortex M backend - edge = transform_for_cortex_m_backend(edge) + if args.target != "vgf": + # Transform so we can use ops from the Cortex M backend + edge = transform_for_cortex_m_backend(edge, args) dump_delegation_info(edge, args.intermediates) + edge_program_manager_copy = copy.deepcopy(edge) + try: exec_prog = edge.to_executorch( config=ExecutorchBackendConfig(extract_delegate_segments=False) @@ -844,9 +875,9 @@ def transform_for_cortex_m_backend(edge): ) if args.bundleio: - output_name = f"{output_name}.bpte" + output_file_name = f"{output_name}.bpte" else: - output_name = f"{output_name}.pte" + output_file_name = f"{output_name}.pte" if args.output is not None: if args.output.endswith(".pte") or args.output.endswith(".bpte"): @@ -859,19 +890,25 @@ def transform_for_cortex_m_backend(edge): raise RuntimeError( f"When not using --bundleio a .bpte file should not be use as --output {args.output}" ) - output_name = args.output + output_file_name = args.output else: # --output is a folder - output_name = os.path.join(args.output, output_name) + output_file_name = os.path.join(args.output, output_file_name) + + if args.bundleio or args.etrecord: + etrecord_file_name = os.path.splitext(output_file_name)[0] + "_etrecord.bin" + # Generate ETRecord + generate_etrecord(etrecord_file_name, edge_program_manager_copy, exec_prog) + print(f"ETRecord saved as {etrecord_file_name}") if args.bundleio: # Realize the quantization impact on numerics when generating reference output reference_model = original_model if not model_int8 else model_int8 - save_bpte_program(exec_prog, reference_model, output_name) - print(f"Bundle PTE file saved as {output_name}") + save_bpte_program(exec_prog, reference_model, output_file_name) + print(f"Bundle PTE file saved as {output_file_name}") else: - save_pte_program(exec_prog, output_name) - print(f"PTE file saved as {output_name}") + save_pte_program(exec_prog, output_file_name) + print(f"PTE file saved as {output_file_name}") if args.evaluate: evaluate_model( diff --git a/examples/arm/ethos-u-porting-guide.md b/examples/arm/ethos-u-porting-guide.md new file mode 100644 index 00000000000..a5ff0c50c6c --- /dev/null +++ b/examples/arm/ethos-u-porting-guide.md @@ -0,0 +1,128 @@ +# Introduction + +As you have seen in in the [Arm® Ethos™-U NPU backend tutorial](../../docs/source/backends-arm-ethos-u.md), ExecuTorch has two distinct parts: +- Ahead-of-time(AoT) compile flow +- Ethos-U on-device runtime + +In the porting guide, we guide you through the main steps to port your SoC with an Ethos-U to the ExecuTorch Ethos-U backend in order to leverage the ExecuTorch Ethos-U enablement. We assume +you are familiar with the concepts introduced in `backends-arm-ethos-u.md`, you have already generated a pte in the AoT flow and want to deploy the ML model on device. +Fundamentally, there are two big approaches you can take in porting a SoC with an Ethos-U NPU towards the ExecuTorch runtime. +- You can use the enablement we have done in ExecuTorch for the Arm® Corstone™-300(Arm® Cortex®-M55 and Arm® Ethos™-U55 reference design) and +Arm® Corstone™-320(Arm® Cortex®-M85 and Arm® Ethos™-U85 reference design) and migrate from the Corstone platform towards a new platform. +- If the SoC comes with an SDK that is not based on ExecuTorch, you can replace the runtime SDKs with the corresponding APIs from ExecuTorch runtime. + +It is important to understand that irrespective of whether the SoC comes with or without SDK, there is +[a single Ethos-U driver](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-driver) that any SoC with Ethos-U relies on. For that reason, there will be +overlap between the two approaches when it comes to the enablement of the Ethos-U NPU. + +## Functioning of a system with an Ethos-U +A system with an Ethos-U and Cortex-M or Arm® Cortex®-A functions in the following way: +- The CPU(Cortex-M or Cortex-A) dispatches an inference job to the Ethos-U NPU. The inference job is in the form of command stream - a sequence of instructions that the NPU +executes. The command stream is generated as part of the `to_edge_transform_and_lower` AoT compile stage. The command stream is embedded within the pte file and at runtime, +the pte file is stored in the memory of the SoC. +- The Ethos-U NPU autonomously reads the command stream from memory. When the NPU finishes processing the command stream, it raises an interrupt to the CPU to signal +that the inference job is complete. The CPU executes the interrupt handler and resumes its normal execution. + +### Ethos-U memory regions +In order to allow this functioning, the Ethos-U driver defines the following regions that the NPU hardware will access: +- Ethos-U scratch buffer - a contiguous block of memory used by the NPU to store the intermediate tensors produced and consumed during inference. Applicable for any Ethos-U NPU. +- Neural Network - a contiguous block of memory holding constant data such as weights, biases, quantization parameters required to run an inference. Applicable for +any Ethos-U NPU. +- Ethos-U fast scratch buffer - a contiguous block of memory for the case when the Ethos-U scratch buffer and Neural Network are both in the external memory. +Applicable only for Ethos-U65 and Ethos-U85 in Dedicated_Sram memory mode. + +### Ethos-U driver +The key function of the Ethos-U driver enabling the interaction with the NPU is +(ethosu_invoke_v3)[https://github.com/pytorch/executorch/blob/main/backends/arm/runtime/EthosUBackend.cpp#L324]. The `ethosu_invoke_v3` function takes as input a driver handle, +a pointer towards the command stream and the size of the command stream, base address as well as the size of the base addresses. For a system with Cortex-M, there is a 1:1 +mapping between base pointer and region, so we will pass three base pointers and each base pointer will correspond to one region. Then, as part of the compilation stage +in `to_edge_transform_and_lower`, the Ethos-U compiler will generate command stream taking into account the three regions. In other words, at runtime, the ethos-U driver +knows the address of the command stream, its size, as well as the address for the locations in memory needed to store the intermediate tensors. +The Ethos-U driver will pass these address to the NPU and the NPU will issue memory requests to the on-chip or external memories in order to access the necessary data +(e.g. read weights, store an intermediate result into the scratch buffer, etc). The `backends/arm/runtime/EthosUBackend.cpp` already integrates the Ethos-U driver and it +already supports the three memory modes of the Ethos-U. Therefore, you should reuse `backends/arm/runtime/EthosUBackend.cpp` as is, without modifications. +The key question for any porting is how to initialize the NPU and +make sure it works. Let's analyse these questions in the following section. + +**Note:** Interrupts work differently between Cortex-M and Cortex-A and a system with Cortex-A will use more base pointers and won't have a 1:1 mapping between Ethos-U +driver base pointer and Ethos-U region. The `backends/arm/runtime/EthosUBackend.cpp` is for a system with Cortex-M. Going forward, we assume we have a system with Cortex-M, similar to the Corstone platforms. + +## NPU initialization +In order to initialize the NPU hardware, the software needs to provide correct information about: +- The base address of the Ethos-U NPU on the memory map of the SoC. +- The interrupt assignment for the Ethos-U. You also need to provide interrupt priority. + +In the `executorch/examples/arm/executor_runner/arm_executor_runner.cpp` sample application, we inherit the Corstone-300/Corstone-320 NPU initialization done in the [core-platform project](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-platform). +We include the core-platform project as a dependency in the `backends/arm/scripts/corstone_utils.cmake` script. Note that in [corstone_utils.sh](https://github.com/pytorch/executorch/blob/main/backends/arm/scripts/corstone_utils.cmake#L69) +depending on whether we target Ethos-U55 or Ethos-U85, we include the corresponding target from core-platform. Then, inside core-platform, the NPU base address and interrupt assignment are defined in the +[target.cpp](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-platform/-/blob/main/targets/corstone-320/target.cpp?ref_type=heads#L44) as per the memory map of the Corstone-300/Corstone-320. +It's worth mentioning that the code in core-platform(code we reuse in the `examples/arm/executor_runner/arm_executor_runner.cpp`) also calls the `ethosu_init` +function to initialize the NPU. The `ethosu_init` function is [defined in the Ethos-U driver](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-driver/-/blob/main/src/ethosu_driver.c?ref_type=heads#L409). The ethos-u driver itself is +included within the core-platform CMake. In other words, to initialize the NPU in the ExecuTorch executor runner application, we reuse the Ethos-U initialization that has been done in the core-platform project. +Core-platform includes a [tutorial](https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-core-platform/-/blob/main/PORTING.md?ref_type=heads) on porting a new target. If you port your target to +core-platform, you can then easily reuse it in the ExecuTorch runtime. + +Also, as explained in the comments in `backends/arm/scripts/corstone_utils.cmake`, note that REGIONCFG register of the Ethos-U controls the memory(on-chip or external memory) used by the NPU to +access the the Ethos-U scratch buffer, ML model and Ethos-U fast scratch buffer. The REGIONCFG is defined in the Ethos-U driver and you need to configure it differently depending on the memory mode. +You can see in `backends/arm/scripts/corstone_utils.cmake` how we overwrite the register to the correct value depending on the memory mode. + +For the Corstone platforms, we make use of the Timing Adapters to model different memory latencies. The Timing Adapter is only applicable to the Corstone targets and should not be included for another SoC. + +## Corstone linker scripts + +The linker scripts point to the linker where to place various objects in memory when the application is loaded onto the target. +In the `arm_executor_runner.cpp` application, we reuse the linker scripts from the core-platform project. Note that the Global Offset Table(.got symbols) needs to be 16-byte aligned. The linker scripts are highly specific to the memory map of system. +For example on the Corstone-300, in order to allow us to build a lot of portable kernels, we relocate the portable kernels from the .text section, living in the ITCM, to a bigger memory. Also on the Corstone-300. the linker script defines two +load regions- rom_exec and rom_dram corresponding to loading the application in the ITCM and in the DDR. When you deploy the application, the boot loader copies the two binaries from the rom_exec/rom_dram regions to their physical address in memory - +a process known as scatter loader. Upon powering on the device, the very first instruction that the Cortex-M executes is the ResetHandler function. The ResetHandler is the first entry in the Vector Interrupt Table, and the location of the Vector Interrupt +Table is specified in the linker script(the `KEEP(*(.vectors))` symbol). The assembly boot code powering on the Cortex-M is itself defined in `core_software/cmsis_6/CMSIS/CoreValidation/Layer/Target/CM55S/RTE/Device/ARMCM55/startup_ARMCM55.c`. +The CMSIS start-up code for Cortex-M is added as part of the build system of the core-platform applications. + +## Coupling between the AoT compile specification memory mode, linker script and the application logic +It is important to note that when you specify a memory mode in the Python script to generate the pte file, in the runtime, the user is +expected to place the scratch buffer and NN in the correct memory location. + +For example, if you generate a pte file with compile specification for Shared Sram, the scratch buffer should be placed in the SRAM and the NN in the external memory in the runtime application code. You can see we are following +this approach in the `examples/arm/executor_runner/arm_executor_runner.cpp` example application. In the linker scripts for the application(`examples/arm/executor_runner/Corstone-320.ld` and +`examples/arm/executor_runner/Corstone-300.ld`) we check the value of `ETHOSU_ARENA` to determine whether the ethos-u scratch buffer is placed in the on-chip memory or in the external memory. In this +way, depending on the `ETHOSU_ARENA` parameter, the linker knows whether the symbol is to be placed in the .ddr or .sram.bss sections. The `ETHOSU_ARENA` parameter is set in the `backends/arm/scripts/corstone_utils.cmake` and +its value is derived based on the memory mode parameter that is passed to the `examples/arm/run.sh` shell script. Then, at link time, the .ddr section is always placed in the external memory and the .sram.bss is always placed in the SRAM. +Finally, note that in the `examples/arm/executor_runner/arm_executor_runner.cpp` application code, we place the buffers for the Ethos-u scratch and the neural network in the correct symbol from the linker script. For instance, +the Ethos-u scratch buffer corresponds to the the `.bss.tensor_arena` section in the linker script, In the application code, when we allocate memory for the Ethos-u scratch buffer, we place this array in the .bss.tensor_arena section in the memory map. + +``` +unsigned char __attribute__(( + section(".bss.tensor_arena"), + aligned(16))) temp_allocation_pool[temp_allocation_pool_size]; +``` +and the `.bss.tensor_arena` section is placed in the correct location in the memory map thanks to +the `ETHOSU_ARENA` parameter. + +There is a tight coupling between the memory mode for the Ethos-U and the placement of the ethos-u scratch buffer, +ethos-u-fast scratch buffer (only applicable for Dedicated_Sram) and the neural network in the memory map of the +SoC. The `arm_executor_runner.cpp` application built with the `examples/arm/run.sh` shell script and corresponding linker scripts are aimed to serve as +example implementation for the correct placement of the various objects in memory. + +It's also worth mentioning that in the AoT Python flow, by default the input to the pte file is in FP32. Therefore, the pte file contains a Quantize node, an Ethos-U custom delegate and a Dequantize node. +Sometimes, you may want to feed quantized input to the Ethos-U custom delegate straight away, for example if you have a camera input outputting RGB data in (u)int8. You can apply the `QuantizeInputs` and +`QuantizeOutputs` passes in the AoT flow for that purpose. Here is a snippet showing how to achieve it: +``` +edge_program_manager = to_edge_transform_and_lower(...) +from executorch.exir.passes.quantize_io_pass import QuantizeInputs +from executorch.exir.passes.quantize_io_pass import QuantizeOutputs +# Apply the QuantizeInputs & QuantizeOutputs passes to input & output tensor 0 +edge_program_manager.transform(passes=[QuantizeInputs(edge_program_manager, [0]), + QuantizeOutputs(edge_program_manager, [0])]) +# Convert edge program to executorch +executorch_program_manager = edge_program_manager.to_executorch( + config=ExecutorchBackendConfig(extract_delegate_segments=False) + ) +``` +If you apply the `QuantizeInputs` pass in the AoT flow, when you populate the input tensor in the runtime application logic, you need to use int8 numbers and not FP32. In the `examples/arm/executor_runner/arm_executor_runner.cpp` application, you can see +how we populate the input tensor depending on its data type. + +## Conclusion +The ExecuTorch project already provides an Arm Ethos-U backend in `executorch/backends/arm/runtime/` that you can reuse as is. The key steps to bring up a new platform is to reuse the Ethos-U driver +and ensure that the NPU base address and interrupt assignment matches your SoC. The `examples/arm/executor_runner/arm_executor_runner.cpp` is an example application running on the Corstone platform. +For the `arm_executor_runner.cpp` application, we are relying on the NPU initialization done in the core-platform project and we integrate core-platform in the +`backends/arm/scripts/corstone_utils.cmake` script. Then, we inherit the core-platform integration of the Ethos-U driver and the CMSIS boot code for the Cortex-M core. \ No newline at end of file diff --git a/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake b/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake index 68fbf8985e9..45e786e4acf 100644 --- a/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake +++ b/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake @@ -97,7 +97,5 @@ add_compile_options( # -Wall -Wextra -Wcast-align -Wdouble-promotion -Wformat # -Wmissing-field-initializers -Wnull-dereference -Wredundant-decls -Wshadow # -Wswitch -Wswitch-default -Wunused -Wno-redundant-decls - -Wno-error=deprecated-declarations - -Wno-error=shift-count-overflow - -Wno-psabi + -Wno-error=deprecated-declarations -Wno-error=shift-count-overflow -Wno-psabi ) diff --git a/examples/arm/ethos-u-setup/core_platform/0001-Add-got-section-to-the-DDR.patch b/examples/arm/ethos-u-setup/core_platform/0001-Add-got-section-to-the-DDR.patch deleted file mode 100644 index f2088f3c933..00000000000 --- a/examples/arm/ethos-u-setup/core_platform/0001-Add-got-section-to-the-DDR.patch +++ /dev/null @@ -1,25 +0,0 @@ -From e637571ca767671d8114542d85bca7965e0a4251 Mon Sep 17 00:00:00 2001 -From: Per Held -Date: Fri, 25 Apr 2025 13:25:29 +0200 -Subject: [PATCH 1/2] Add got section to the DDR - ---- - targets/corstone-300/platform.ld | 2 +- - 1 file changed, 1 insertion(+), 1 deletion(-) - -diff --git a/targets/corstone-300/platform.ld b/targets/corstone-300/platform.ld -index d586b97..b746aa0 100644 ---- a/targets/corstone-300/platform.ld -+++ b/targets/corstone-300/platform.ld -@@ -281,7 +281,7 @@ SECTIONS - #endif - * (expected_output_data_sec) - * (sec_command_stream, sec_weight_data, sec_input_data) -- -+ * (.got*) - * (ethosu_core_in_queue) - * (ethosu_core_out_queue) - . = ALIGN(4); --- -2.43.0 - diff --git a/examples/arm/ethos-u-setup/core_platform/0002-Move-input_data_sec-to-NOLOAD-area.patch b/examples/arm/ethos-u-setup/core_platform/0002-Move-input_data_sec-to-NOLOAD-area.patch deleted file mode 100644 index e9f1c332b42..00000000000 --- a/examples/arm/ethos-u-setup/core_platform/0002-Move-input_data_sec-to-NOLOAD-area.patch +++ /dev/null @@ -1,71 +0,0 @@ -From 42a16a7e9c73e79e55ee25534e3bbc39f169af62 Mon Sep 17 00:00:00 2001 -From: Per Held -Date: Mon, 28 Apr 2025 10:56:09 +0200 -Subject: [PATCH 2/2] Move input_data_sec to NOLOAD area - ---- - targets/corstone-300/platform.ld | 10 ++++++++-- - targets/corstone-320/platform.ld | 8 ++++++-- - 2 files changed, 14 insertions(+), 4 deletions(-) - -diff --git a/targets/corstone-300/platform.ld b/targets/corstone-300/platform.ld -index b746aa0..5043be2 100644 ---- a/targets/corstone-300/platform.ld -+++ b/targets/corstone-300/platform.ld -@@ -273,19 +273,25 @@ SECTIONS - *(.bss.tensor_arena) - #endif - -- . = ALIGN(4); -- *(input_data_sec) - . = ALIGN(16); - #if (ETHOSU_MODEL == 1) - *(network_model_sec) - #endif - * (expected_output_data_sec) -+ . = ALIGN(16); - * (sec_command_stream, sec_weight_data, sec_input_data) - * (.got*) - * (ethosu_core_in_queue) - * (ethosu_core_out_queue) - . = ALIGN(4); - } > DDR :rom_dram -+ .ddr_noload (NOLOAD) : -+ { -+ . = ALIGN(16); -+ *(input_data_sec) -+ . = ALIGN(16); -+ } > DDR :null -+ - - __eddr_data = ALIGN(4); - .sram.data : -diff --git a/targets/corstone-320/platform.ld b/targets/corstone-320/platform.ld -index 1f4f521..8c5e402 100644 ---- a/targets/corstone-320/platform.ld -+++ b/targets/corstone-320/platform.ld -@@ -268,8 +268,6 @@ SECTIONS - *(network_model_sec) - #endif - -- . = ALIGN(4); -- *(input_data_sec) - *(expected_output_data_sec) - *(output_data_sec) - -@@ -279,6 +277,12 @@ SECTIONS - __etext = .; - } > DDR :rom_dram - -+ .ddr_noload (NOLOAD) : -+ { -+ . = ALIGN(16); -+ *(input_data_sec) -+ } > DDR :null -+ - .bss : - { - . = ALIGN(4); --- -2.43.0 - diff --git a/examples/arm/ethos-u-setup/core_platform/0003-Move-the-portable-kernels-to-the-BRAM.patch b/examples/arm/ethos-u-setup/core_platform/0003-Move-the-portable-kernels-to-the-BRAM.patch deleted file mode 100644 index ccb27b83711..00000000000 --- a/examples/arm/ethos-u-setup/core_platform/0003-Move-the-portable-kernels-to-the-BRAM.patch +++ /dev/null @@ -1,42 +0,0 @@ -From 81f5bd9092bc25c343d8d85b692698c6d961d0bd Mon Sep 17 00:00:00 2001 -From: George Gekov -Date: Mon, 28 Jul 2025 15:23:50 +0100 -Subject: [PATCH] Move the portable kernels to the BRAM - -On the Corstone-300, we have 512KB of ITCM and by default, -the .text section lives in the ITCM. However, as we build more -portable kernels, we sometimes overflow and the .text section -no longer fits in the ITCM. This patch moves the portable kernels -to the BRAM as we have 1MB of BRAM ---- - targets/corstone-300/platform.ld | 7 ++++++- - 1 file changed, 6 insertions(+), 1 deletion(-) - -diff --git a/targets/corstone-300/platform.ld b/targets/corstone-300/platform.ld -index 5043be2..399e9f7 100644 ---- a/targets/corstone-300/platform.ld -+++ b/targets/corstone-300/platform.ld -@@ -135,7 +135,11 @@ SECTIONS - { - _vectors = .; - KEEP(*(.vectors)) -- *(.text*) -+ *(EXCLUDE_FILE( -+ *op_*.cpp.obj -+ ) -+ -+ .text*) - - KEEP(*(.init)) - KEEP(*(.fini)) -@@ -299,6 +303,7 @@ SECTIONS - __sram_data_start__ = .; - *(.sram.data) - . = ALIGN(4); -+ *op_*.cpp.obj (*.text*) - __sram_data_end__ = .; - } > BRAM AT >DDR :rom_dram - --- -2.39.5 (Apple Git-154) - diff --git a/examples/arm/ethos_u_minimal_example.ipynb b/examples/arm/ethos_u_minimal_example.ipynb index fd9dcfdf338..e63a7d37e58 100644 --- a/examples/arm/ethos_u_minimal_example.ipynb +++ b/examples/arm/ethos_u_minimal_example.ipynb @@ -18,13 +18,12 @@ "source": [ "# Ethos-U delegate flow example\n", "\n", - "This guide demonstrates the full flow for running a module on Arm Ethos-U using ExecuTorch. \n", + "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", - "3. Add Arm cross-compilation toolchain and simulators to PATH using `examples/arm/ethos-u-scratch/setup_path.sh` \n", + "2. Install Arm cross-compilation toolchain and simulators using `./examples/arm/setup.sh --i-agree-to-the-contained-eula`\n", "\n", "With all commands executed from the base `executorch` folder.\n", "\n", @@ -58,7 +57,7 @@ "\n", "model = Add()\n", "model = model.eval()\n", - "exported_program = torch.export.export_for_training(model, example_inputs)\n", + "exported_program = torch.export.export(model, example_inputs)\n", "graph_module = exported_program.module()\n", "\n", "_ = graph_module.print_readable()" @@ -70,7 +69,9 @@ "source": [ "To run on Ethos-U the `graph_module` must be quantized using the `arm_quantizer`. Quantization can be done in multiple ways and it can be customized for different parts of the graph; shown here is the recommended path for the EthosUBackend. Quantization also requires calibrating the module with example inputs.\n", "\n", - "Again printing the module, it can be seen that the quantization wraps the node in quantization/dequantization nodes which contain the computed quanitzation parameters." + "Again printing the module, it can be seen that the quantization wraps the node in quantization/dequantization nodes which contain the computed quanitzation parameters.\n", + "\n", + "With the default passes for the Arm Ethos-U backend, assuming the model lowers fully to the Ethos-U, the exported program is composed of a Quantize node, Ethos-U custom delegate and a Dequantize node. In some circumstances, you may want to feed quantized input to the Neural Network straight away, e.g. if you have a camera sensor outputting (u)int8 data and keep all the arithmetic of the application in the int8 domain. For these cases, you can apply the `exir/passes/quantize_io_pass.py`. See the unit test in `backends/arm/test/passes/test_ioquantization_pass.py`for an example how to feed quantized inputs and obtain quantized outputs.\n" ] }, { @@ -86,13 +87,11 @@ ")\n", "from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e\n", "\n", - "target = \"ethos-u55-128\"\n", - "\n", "# Create a compilation spec describing the target for configuring the quantizer\n", "# Some args are used by the Arm Vela graph compiler later in the example. Refer to Arm Vela documentation for an\n", "# explanation of its flags: https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela/-/blob/main/OPTIONS.md\n", "spec_builder = ArmCompileSpecBuilder().ethosu_compile_spec(\n", - " target,\n", + " target=\"ethos-u55-128\",\n", " system_config=\"Ethos_U55_High_End_Embedded\",\n", " memory_mode=\"Shared_Sram\",\n", " extra_flags=\"--output-format=raw --debug-force-regor\"\n", @@ -112,7 +111,7 @@ "_ = quantized_graph_module.print_readable()\n", "\n", "# Create a new exported program using the quantized_graph_module\n", - "quantized_exported_program = torch.export.export_for_training(quantized_graph_module, example_inputs)" + "quantized_exported_program = torch.export.export(quantized_graph_module, example_inputs)" ] }, { @@ -137,7 +136,6 @@ "metadata": {}, "outputs": [], "source": [ - "import os\n", "from executorch.backends.arm.ethosu import EthosUPartitioner\n", "from executorch.exir import (\n", " EdgeCompileConfig,\n", @@ -163,15 +161,10 @@ " config=ExecutorchBackendConfig(extract_delegate_segments=False)\n", " )\n", "\n", - "executorch_program_manager.exported_program().module().print_readable()\n", + "_ = executorch_program_manager.exported_program().module().print_readable()\n", "\n", "# Save pte file\n", - "cwd_dir = os.getcwd()\n", - "pte_base_name = \"simple_example\"\n", - "pte_name = pte_base_name + \".pte\"\n", - "pte_path = os.path.join(cwd_dir, pte_name)\n", - "save_pte_program(executorch_program_manager, pte_name)\n", - "assert os.path.exists(pte_path), \"Build failed; no .pte-file found\"" + "save_pte_program(executorch_program_manager, \"ethos_u_minimal_example.pte\")" ] }, { @@ -180,10 +173,9 @@ "source": [ "## Build executor runtime\n", "\n", - "After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced .pte-file using the Arm cross-compilation toolchain. This is done in three steps:\n", - "1. Build the executorch library and EthosUDelegate.\n", - "2. Build any external kernels required. In this example this is not needed as the graph is fully delegated, but its included for completeness.\n", - "3. Build and link the `arm_executor_runner`." + "After the AOT compilation flow is done, the runtime can be cross compiled and linked to the produced .pte-file using the Arm cross-compilation toolchain. This is done in two steps:\n", + "1. Build and install the executorch libraries and EthosUDelegate.\n", + "2. Build and link the `arm_executor_runner` and generate kernel bindings for any non delegated ops." ] }, { @@ -192,25 +184,37 @@ "metadata": {}, "outputs": [], "source": [ - "import subprocess\n", - "\n", - "# Setup paths\n", - "et_dir = os.path.join(cwd_dir, \"..\", \"..\")\n", - "et_dir = os.path.abspath(et_dir)\n", - "script_dir = os.path.join(et_dir, \"backends\", \"arm\", \"scripts\")\n", - "\n", - "# Cross-compile executorch \n", - "subprocess.run(os.path.join(script_dir, \"build_executorch.sh\"), shell=True, cwd=et_dir)\n", - "\n", - "# Cross-compile portable kernels\n", - "subprocess.run(os.path.join(script_dir, \"build_portable_kernels.sh\"), shell=True, cwd=et_dir)\n", - "\n", - "# Cross-compile executorch runner\n", - "args = f\"--pte={pte_path} --target={target}\"\n", - "subprocess.run(os.path.join(script_dir, \"build_executor_runner.sh\") + \" \" + args, shell=True, cwd=et_dir)\n", - "\n", - "elf_path = os.path.join(cwd_dir, pte_base_name, \"cmake-out\", \"arm_executor_runner\")\n", - "assert os.path.exists(elf_path), \"Build failed; no .elf-file found\"" + "%%bash\n", + "# Ensure the arm-none-eabi-gcc toolchain and FVP:s are available on $PATH\n", + "source ethos-u-scratch/setup_path.sh\n", + "\n", + "# Build executorch libraries cross-compiled for arm baremetal to executorch/cmake-out-arm\n", + "cmake --preset arm-baremetal \\\n", + "-DCMAKE_BUILD_TYPE=Release \\\n", + "-B../../cmake-out-arm ../..\n", + "cmake --build ../../cmake-out-arm --target install -j$(nproc) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash \n", + "source ethos-u-scratch/setup_path.sh\n", + "\n", + "# Build example executor runner application to examples/arm/ethos_u_minimal_example\n", + "cmake -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \\\n", + " -DCMAKE_BUILD_TYPE=Release \\\n", + " -DET_PTE_FILE_PATH=ethos_u_minimal_example.pte \\\n", + " -DTARGET_CPU=cortex-m55 \\\n", + " -DETHOSU_TARGET_NPU_CONFIG=ethos-u55-128 \\\n", + " -DMEMORY_MODE=Shared_Sram \\\n", + " -DSYSTEM_CONFIG=Ethos_U55_High_End_Embedded \\\n", + " -Bethos_u_minimal_example \\\n", + " executor_runner\n", + "cmake --build ethos_u_minimal_example -j$(nproc) -- arm_executor_runner" ] }, { @@ -219,7 +223,7 @@ "source": [ "# Run on simulated model\n", "\n", - "We can finally use the `backends/arm/scripts/run_fvp.sh` utility script to run the .elf-file on simulated Arm hardware. This Script runs the model with an input of ones, so the expected result of the addition should be close to 2." + "We can finally use the `backends/arm/scripts/run_fvp.sh` utility script to run the .elf-file on simulated Arm hardware. The example application is by default built with an input of ones, so the expected result of the quantized addition should be close to 2." ] }, { @@ -228,14 +232,17 @@ "metadata": {}, "outputs": [], "source": [ - "args = f\"--elf={elf_path} --target={target}\"\n", - "subprocess.run(os.path.join(script_dir, \"run_fvp.sh\") + \" \" + args, shell=True, cwd=et_dir)" + "%%bash \n", + "source ethos-u-scratch/setup_path.sh\n", + "\n", + "# Run the example\n", + "../../backends/arm/scripts/run_fvp.sh --elf=ethos_u_minimal_example/arm_executor_runner --target=ethos-u55-128" ] } ], "metadata": { "kernelspec": { - "display_name": ".venv", + "display_name": ".venv (3.10.15)", "language": "python", "name": "python3" }, diff --git a/examples/arm/example_modules/README.md b/examples/arm/example_modules/README.md deleted file mode 100644 index 9a746114b98..00000000000 --- a/examples/arm/example_modules/README.md +++ /dev/null @@ -1,7 +0,0 @@ -# Example of an external model for the ARM AOT Compiler -Example of an external Python file to be used as a module by the `run.sh` (and the `aot_arm_compiler.py`) scripts in `examples/arm` directory. -Just pass the path of the `add.py` file as `--model_name`: - -`ModelUnderTest` should be a `torch.nn.module` instance. - -`ModelInputs` should be a tuple of inputs to the forward function. diff --git a/examples/arm/example_modules/add.py b/examples/arm/example_modules/add.py index 6942e97f807..d29206083f8 100644 --- a/examples/arm/example_modules/add.py +++ b/examples/arm/example_modules/add.py @@ -1,3 +1,18 @@ +# All rights reserved. +# Copyright 2023-2025 Arm Limited and/or its affiliates. +# +# Example of an external model for the Arm AOT Compiler +# +# Example of an external Python file to be used as a module by the `run.sh` +# (and the `aot_arm_compiler.py`) scripts in `examples/arm` directory. +# +# Just pass the path of the `add.py` file as `--model_name` +# +# These two variables are picked up by the `aot_arm_compiler.py` and used: +# `ModelUnderTest` should be a `torch.nn.module` instance. +# `ModelInputs` should be a tuple of inputs to the forward function. +# + import torch diff --git a/examples/arm/executor_runner/CMakeLists.txt b/examples/arm/executor_runner/CMakeLists.txt index 0db57e9d15a..ff6f73398c3 100644 --- a/examples/arm/executor_runner/CMakeLists.txt +++ b/examples/arm/executor_runner/CMakeLists.txt @@ -12,6 +12,11 @@ option( "Set ET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE to specify memory alloction pool size" OFF ) +option( + ET_MODEL_PTE_ADDR + "Place in memory that the PTE file is located/flashed, if set to OFF the PTE is built into the code as a big data area." + OFF +) option(ET_BUNDLE_IO "Set to compile in BundleIO support" OFF) option(ET_ATOL "Set atol to use for BundleIO testing" OFF) option(ET_RTOL "Set rtol to use for BundleIO testing" OFF) @@ -20,23 +25,32 @@ option(ET_DUMP_OUTPUT "Dump output in log" ON) option(FETCH_ETHOS_U_CONTENT "Fetch ethos_u dependencies instead of relying on pre-downloads" ON ) +set(ET_NUM_INFERENCES + "1" + CACHE STRING "Number of inferences to run" +) -if(NOT DEFINED ET_PTE_FILE_PATH AND NOT ${SEMIHOSTING}) +if(NOT DEFINED ET_MODEL_PTE_ADDR + AND NOT DEFINED ET_PTE_FILE_PATH + AND NOT DEFINED SEMIHOSTING +) message( FATAL_ERROR - "ET_PTE_FILE_PATH must specify a model .pte or .bpte, for bare metal systems the " - "model is built into the binary." + "You must specify one of " + "ET_MODEL_PTE_ADDR - .pte (or .bpte) on address, " + "ET_PTE_FILE_PATH - .pte (or .bpte) built into the binary or " + "SEMIHOSTING - pte and input via host filesystem (use by pytest)" ) endif() # Example ExecuTorch demo for bare metal Cortex-M based systems set(ET_DIR_PATH - "../../.." + "${CMAKE_CURRENT_SOURCE_DIR}/../../.." CACHE PATH "Path to ExecuTorch dir" ) set(ET_BUILD_DIR_PATH - "${ET_DIR_PATH}/cmake-out" - CACHE PATH "Path to ExecuTorch build dir" + "${ET_DIR_PATH}/cmake-out-arm" + CACHE PATH "Path to ExecuTorch build/install dir" ) set(ET_INCLUDE_PATH "${ET_DIR_PATH}/.." @@ -77,14 +91,7 @@ set(MEMORY_MODE message(STATUS "SYSTEM_CONFIG is ${SYSTEM_CONFIG}") message(STATUS "MEMORY_MODE is ${MEMORY_MODE}") - -get_filename_component(ET_BUILD_DIR_PATH ${ET_BUILD_DIR_PATH} REALPATH) -get_filename_component(ET_DIR_PATH ${ET_DIR_PATH} REALPATH) -get_filename_component(ET_INCLUDE_PATH ${ET_INCLUDE_PATH} REALPATH) -get_filename_component(ETHOS_SDK_PATH ${ETHOS_SDK_PATH} REALPATH) -if(NOT ${SEMIHOSTING}) - get_filename_component(ET_PTE_FILE_PATH ${ET_PTE_FILE_PATH} REALPATH) -endif() +message(STATUS "ET_NUM_INFERENCES is ${ET_NUM_INFERENCES}") # By default, use 2MB of temporary scratch buffer For Dedicated_Sram, use 64MB # for the temporary scratch buffer and 384KB for the fast scratch buffer(the @@ -115,27 +122,8 @@ find_package( executorch REQUIRED HINTS "${ET_BUILD_DIR_PATH}/lib/cmake/ExecuTorch" ) -add_library(arm_portable_ops_lib STATIC IMPORTED) -set_property( - TARGET arm_portable_ops_lib - PROPERTY IMPORTED_LOCATION - "${ET_BUILD_DIR_PATH}/examples/arm/libarm_portable_ops_lib.a" -) -add_library(cortex_m_ops_lib STATIC IMPORTED) -set_property( - TARGET cortex_m_ops_lib - PROPERTY IMPORTED_LOCATION - "${ET_BUILD_DIR_PATH}/backends/cortex_m/libcortex_m_ops_lib.a" -) -add_library(cortex_m_kernels STATIC IMPORTED) -set_property( - TARGET cortex_m_kernels - PROPERTY IMPORTED_LOCATION - "${ET_BUILD_DIR_PATH}/backends/cortex_m/libcortex_m_kernels.a" -) - # Convert pte to header -if(NOT ${SEMIHOSTING}) +if(NOT ${ET_MODEL_PTE_ADDR} AND NOT SEMIHOSTING) add_custom_target( gen_model_header DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/model_pte.h ) @@ -157,8 +145,47 @@ target_sources( arm_memory_allocator.cpp ) -# Include the target's bare-metal linker script -ethosu_eval_link_options(arm_executor_runner) +# Check for "U55" in SYSTEM_CONFIG +string(FIND "${SYSTEM_CONFIG}" "U55" U55_FOUND) + +# Check for "U85" in SYSTEM_CONFIG +string(FIND "${SYSTEM_CONFIG}" "U85" U85_FOUND) + +# Check if neither "U55" nor "U85" was found +if(U55_FOUND EQUAL -1 AND U85_FOUND EQUAL -1) + message( + FATAL_ERROR + "SYSTEM_CONFIG does not contain 'U55' or 'U85'. Configuration aborting." + ) +endif() + +# Proceed with specific actions if either is found +if(NOT U55_FOUND EQUAL -1) + message(STATUS "SYSTEM_CONFIG contains 'U55'.") + set(LINK_FILE_IN "${CMAKE_SOURCE_DIR}/Corstone-300.ld") +endif() + +if(NOT U85_FOUND EQUAL -1) + message(STATUS "SYSTEM_CONFIG contains 'U85'.") + set(LINK_FILE_IN "${CMAKE_SOURCE_DIR}/Corstone-320.ld") +endif() + +if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + set(LINK_FILE_EXT ld) + set(LINK_FILE_OPTION "-T") + set(COMPILER_PREPROCESSOR_OPTIONS -E -x c -P) +endif() + +get_filename_component(LINK_FILE_OUT_BASE ${LINK_FILE} NAME) +set(LINK_FILE_OUT + ${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT} +) + +execute_process( + COMMAND ${CMAKE_C_COMPILER} ${COMPILER_PREPROCESSOR_OPTIONS} -o + ${LINK_FILE_OUT} ${LINK_FILE_IN} +) +target_link_options(arm_executor_runner PRIVATE "-T" "${LINK_FILE_OUT}") set(arm_executor_runner_link) list( @@ -168,10 +195,9 @@ list( ethosu_target_init executorch quantized_ops_lib + cortex_m_ops_lib "-Wl,--whole-archive" executorch_delegate_ethos_u - cortex_m_ops_lib - arm_portable_ops_lib quantized_kernels cortex_m_kernels portable_kernels @@ -180,6 +206,90 @@ list( -Map=arm_executor_runner.map ) +# Prefer to generate kernel bindings from model file if possible, which is when +# 1. Not building for semihosting 2. Not building with bundleio If that is not +# the case, fallback to select_ops_list If the model file does not contain any +# aten ops, a workaround is currently needed to avoid crashing. +execute_process( + COMMAND + python "${ET_DIR_PATH}/codegen/tools/gen_oplist.py" + --model_file_path=${ET_PTE_FILE_PATH} + --output_path=${CMAKE_CURRENT_BINARY_DIR}/temp.yaml + OUTPUT_VARIABLE CMD_RESULT +) + +if(CMD_RESULT MATCHES "aten::" OR CMD_RESULT MATCHES "dim_order_ops::") + set(FOUND_OPS_IN_FILE "true") +else() + set(FOUND_OPS_IN_FILE "false") +endif() + +if(${SEMIHOSTING}) + set(EXECUTORCH_SELECT_OPS_MODEL "") + message( + "gen_oplist: Building with semihosting, no model is used to auto generate ops from will use EXECUTORCH_SELECT_OPS_LIST=${EXECUTORCH_SELECT_OPS_LIST}" + ) +elseif(${FOUND_OPS_IN_FILE}) + set(EXECUTORCH_SELECT_OPS_LIST "") + set(EXECUTORCH_SELECT_OPS_MODEL "${ET_PTE_FILE_PATH}") + message( + "gen_oplist: EXECUTORCH_SELECT_OPS_MODEL=${ET_PTE_FILE_PATH} is used to auto generate ops from" + ) +elseif(NOT ${FOUND_OPS_IN_FILE} AND ${ET_BUNDLE_IO}) + set(EXECUTORCH_SELECT_OPS_MODEL "") + message( + "gen_oplist: Building with ET_BUNDLE_IO and .bpte is not supported to auto generate ops from will use EXECUTORCH_SELECT_OPS_LIST=${EXECUTORCH_SELECT_OPS_LIST}" + ) +else() + set(EXECUTORCH_SELECT_OPS_LIST "") + set(EXECUTORCH_SELECT_OPS_MODEL "") + message( + "gen_oplist: No non delagated ops was found in ${ET_PTE_FILE_PATH} no ops added to build" + ) +endif() + +# Ensure that either executorch_select_ops_list or executorch_select_ops_model +# is set - otherwise assume no kernels needs to be registered +if(NOT ("${EXECUTORCH_SELECT_OPS_LIST}" STREQUAL "" + AND "${EXECUTORCH_SELECT_OPS_MODEL}" STREQUAL "") +) + set(EXECUTORCH_ROOT ${ET_DIR_PATH}) + include(${ET_DIR_PATH}/tools/cmake/Utils.cmake) + include(${ET_DIR_PATH}/tools/cmake/Codegen.cmake) + + gen_selected_ops( + LIB_NAME + "arm_portable_ops_lib" + OPS_SCHEMA_YAML + "" + ROOT_OPS + "${EXECUTORCH_SELECT_OPS_LIST}" + INCLUDE_ALL_OPS + "" + OPS_FROM_MODEL + "${EXECUTORCH_SELECT_OPS_MODEL}" + DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_ENABLE_DTYPE_SELECTIVE_BUILD}" + ) + + generate_bindings_for_kernels( + LIB_NAME "arm_portable_ops_lib" FUNCTIONS_YAML + ${ET_DIR_PATH}/kernels/portable/functions.yaml DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_ENABLE_DTYPE_SELECTIVE_BUILD}" + ) + gen_operators_lib( + LIB_NAME + "arm_portable_ops_lib" + KERNEL_LIBS + portable_kernels + DEPS + executorch + DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_ENABLE_DTYPE_SELECTIVE_BUILD}" + ) + list(APPEND arm_executor_runner_link arm_portable_ops_lib) +endif() + if(EXECUTORCH_ENABLE_EVENT_TRACER) target_compile_options(arm_executor_runner PUBLIC -DET_EVENT_TRACER_ENABLED) @@ -208,10 +318,12 @@ target_compile_definitions( arm_executor_runner PRIVATE C10_USING_CUSTOM_GENERATED_MACROS ) +if(NOT ${ET_MODEL_PTE_ADDR} AND NOT SEMIHOSTING) + add_dependencies(arm_executor_runner gen_model_header) +endif() + if(SEMIHOSTING) target_compile_definitions(arm_executor_runner PUBLIC SEMIHOSTING) -else() - add_dependencies(arm_executor_runner gen_model_header) endif() if(ET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE) @@ -235,6 +347,12 @@ if(DEFINED ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) ) endif() +if(ET_MODEL_PTE_ADDR) + target_compile_definitions( + arm_executor_runner PUBLIC -DET_MODEL_PTE_ADDR=${ET_MODEL_PTE_ADDR} + ) +endif() + if(ET_BUNDLE_IO) target_compile_definitions(arm_executor_runner PUBLIC -DET_BUNDLE_IO) endif() @@ -255,6 +373,12 @@ if(ET_DUMP_OUTPUT) target_compile_definitions(arm_executor_runner PUBLIC -DET_DUMP_OUTPUT) endif() +if(ET_NUM_INFERENCES) + target_compile_definitions( + arm_executor_runner PUBLIC ET_NUM_INFERENCES=${ET_NUM_INFERENCES} + ) +endif() + # Fixup compilation of retarget.c if(SEMIHOSTING) # Remove this when MLBEDSW-8910 is closed. diff --git a/examples/arm/executor_runner/Corstone-300.ld b/examples/arm/executor_runner/Corstone-300.ld new file mode 100644 index 00000000000..f5b063a35c6 --- /dev/null +++ b/examples/arm/executor_runner/Corstone-300.ld @@ -0,0 +1,300 @@ +/* + * Copyright 2025 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 is a simplified linkerscript for the Corstone-300 memory system. + * This example has been modified to place certain sections in specific memory. + * Please refer to the Corstone SSE-300 Technical Reference Manual for + * further information. + * + * https://developer.arm.com/Processors/Corstone-300 + */ + +#ifndef ETHOSU_MODEL + /* default value - '1', for DRAM */ + #define ETHOSU_MODEL 1 +#endif + +#ifndef ETHOSU_ARENA + /* default value - '1', for DRAM */ + #define ETHOSU_ARENA 1 +#endif + +__STACK_SIZE = 0x00008000; +__HEAP_SIZE = 0x00008000; + +MEMORY +{ + ITCM (rx) : ORIGIN = 0x10000000, LENGTH = 0x00080000 + BRAM (rw) : ORIGIN = 0x11000000, LENGTH = 0x00100000 + DTCM (rw) : ORIGIN = 0x30000000, LENGTH = 0x00080000 + SRAM (rw) : ORIGIN = 0x31000000, LENGTH = 0x00200000 + QSPI (rw) : ORIGIN = 0x38000000, LENGTH = 0x00800000 + DDR (rwx) : ORIGIN = 0x70000000, LENGTH = 0x60000000 +} + +PHDRS +{ + rom_exec PT_LOAD; + rom_dram PT_LOAD; + null PT_NULL; +} + +/* Linker script to place sections and symbol values. Should be used together + * with other linker script that defines memory regions ITCM and RAM. + * It references following symbols, which must be defined in code: + * Reset_Handler : Entry of reset handler + * + * It defines following symbols, which code can use without definition: + * __exidx_start + * __exidx_end + * __copy_table_start__ + * __copy_table_end__ + * __zero_table_start__ + * __zero_table_end__ + * __etext + * __data_start__ + * __preinit_array_start + * __preinit_array_end + * __init_array_start + * __init_array_end + * __fini_array_start + * __fini_array_end + * __data_end__ + * __bss_start__ + * __bss_end__ + * __end__ + * end + * __HeapLimit + * __StackLimit + * __StackTop + * __stack + */ + +ENTRY(Reset_Handler) + +SECTIONS +{ + .text : + { + _vectors = .; + KEEP(*(.vectors)) + *(EXCLUDE_FILE( + *op_*.cpp.obj + ) + .text*) + KEEP(*(.init)) + KEEP(*(.fini)) + + /* .ctors */ + *crtbegin.o(.ctors) + *crtbegin?.o(.ctors) + *(EXCLUDE_FILE(*crtend?.o *crtend.o) .ctors) + *(SORT(.ctors.*)) + *(.ctors) + + /* .dtors */ + *crtbegin.o(.dtors) + *crtbegin?.o(.dtors) + *(EXCLUDE_FILE(*crtend?.o *crtend.o) .dtors) + *(SORT(.dtors.*)) + *(.dtors) + + KEEP(*(.eh_frame*)) + } > ITCM :rom_exec + + /* + * SG veneers: + * All SG veneers are placed in the special output section .gnu.sgstubs. Its start address + * must be set, either with the command line option '--section-start' or in a linker script, + * to indicate where to place these veneers in memory. + */ +/* + .gnu.sgstubs : + { + . = ALIGN(32); + } > ITCM :rom_exec +*/ + .ARM.extab : + { + *(.ARM.extab* .gnu.linkonce.armextab.*) + } > ITCM :rom_exec + + .ARM.exidx : + { + __exidx_start = .; + *(.ARM.exidx* .gnu.linkonce.armexidx.*) + __exidx_end = .; + } > ITCM :rom_exec + + .copy.table : + { + . = ALIGN(4); + __copy_table_start__ = .; + LONG (__etext) + LONG (__data_start__) + LONG ((__data_end__ - __data_start__) / 4) + + LONG (__eddr_data) + LONG (__sram_data_start__) + LONG ((__sram_data_end__ - __sram_data_start__) / 4) + + LONG (__eddr_data + (__sram_data_end__ - __sram_data_start__)) + LONG (__rodata_start__) + LONG ((__rodata_end__ - __rodata_start__) / 4) + + __copy_table_end__ = .; + } > ITCM :rom_exec + + .zero.table : + { + . = ALIGN(4); + __zero_table_start__ = .; + LONG (__bss_start__) + LONG ((__bss_end__ - __bss_start__) / 4) + __zero_table_end__ = .; + + /** + * Location counter can end up 2byte aligned with narrow Thumb code but + * __etext is assumed by startup code to be the LMA of a section in DTCM + * which must be 4byte aligned + */ + __etext = ALIGN (4); + + } > ITCM :rom_exec + + .data : AT(__etext) + { + __data_start__ = .; + *(vtable) + *(.data) + *(.data.*) + + . = ALIGN(4); + /* preinit data */ + PROVIDE_HIDDEN (__preinit_array_start = .); + KEEP(*(.preinit_array)) + PROVIDE_HIDDEN (__preinit_array_end = .); + + . = ALIGN(4); + /* init data */ + PROVIDE_HIDDEN (__init_array_start = .); + KEEP(*(SORT(.init_array.*))) + KEEP(*(.init_array)) + PROVIDE_HIDDEN (__init_array_end = .); + + . = ALIGN(4); + /* finit data */ + PROVIDE_HIDDEN (__fini_array_start = .); + KEEP(*(SORT(.fini_array.*))) + KEEP(*(.fini_array)) + PROVIDE_HIDDEN (__fini_array_end = .); + + KEEP(*(.jcr*)) + . = ALIGN(4); + /* All data end */ + __data_end__ = .; + } > DTCM :rom_exec + + .sram.bss : + { + . = ALIGN(16); +#if (ETHOSU_MODEL == 0) + * (network_model_sec) +#endif + +#if (ETHOSU_ARENA == 0) + . = ALIGN(32); + *(.bss.tensor_arena) +#endif + . = ALIGN(16); + *(.bss.ethosu_scratch); + *.(output_data_sec) + } > SRAM :null + + .ddr : + { +#if (ETHOSU_ARENA == 1) + . = ALIGN(32); + *(.bss.tensor_arena) +#endif + + . = ALIGN(4); + *(input_data_sec) + . = ALIGN(16); +#if (ETHOSU_MODEL == 1) + *(network_model_sec) +#endif + * (expected_output_data_sec) + . = ALIGN(16); + * (sec_command_stream, sec_weight_data, sec_input_data) + * (.got*) + * (ethosu_core_in_queue) + * (ethosu_core_out_queue) + . = ALIGN(4); + } > DDR :rom_dram + .ddr_noload (NOLOAD) : + { + . = ALIGN(16); + *(input_data_sec) + . = ALIGN(16); + } > DDR :null + __eddr_data = ALIGN(4); + .sram.data : + { + __sram_data_start__ = .; + *(.sram.data) + . = ALIGN(4); + *op_*.cpp.obj (*.text*) + __sram_data_end__ = .; + } > BRAM AT >DDR :rom_dram + + .rodata : + { + __rodata_start__ = .; + *(.rodata) + *(.rodata.*) + . = ALIGN(4); + __rodata_end__ = .; + } > DTCM AT >DDR :rom_dram + + .bss : + { + . = ALIGN(4); + __bss_start__ = .; + *(.bss) + *(.bss.*) + *(COMMON) + . = ALIGN(4); + __bss_end__ = .; + } > DTCM :null + + .heap (COPY) : + { + . = ALIGN(8); + __end__ = .; + PROVIDE(end = .); + . = . + __HEAP_SIZE; + . = ALIGN(8); + __HeapLimit = .; + } > DTCM :null + + .stack (ORIGIN(DTCM) + LENGTH(DTCM) - __STACK_SIZE) (COPY) : + { + . = ALIGN(8); + __StackLimit = .; + . = . + __STACK_SIZE; + . = ALIGN(8); + __StackTop = .; + } > DTCM :null + PROVIDE(__stack = __StackTop); + + __RAM_segment_used_end__ = .; + + /* Check if data + heap + stack exceeds DTCM limit */ + ASSERT(__StackLimit >= __HeapLimit, "region DTCM overflowed with stack") +} diff --git a/examples/arm/executor_runner/Corstone-320.ld b/examples/arm/executor_runner/Corstone-320.ld new file mode 100644 index 00000000000..62bb6240913 --- /dev/null +++ b/examples/arm/executor_runner/Corstone-320.ld @@ -0,0 +1,295 @@ +/* + * Copyright 2025 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 is a simplified linkerscript for the Corstone-300 memory system. + * This example has been modified to place certain sections in specific memory. + * Please refer to the Corstone SSE-300 Technical Reference Manual for + * further information. + * + * https://developer.arm.com/Processors/Corstone-320 + */ + +/* default value - '1', for DRAM */ +#ifndef ETHOSU_MODEL +#define ETHOSU_MODEL 1 +#endif + +/* default value - '1', for DRAM */ +#ifndef ETHOSU_ARENA +#define ETHOSU_ARENA 1 +#endif + +#ifndef STACK_SIZE +#define STACK_SIZE 0x8000 +#endif + +#ifndef HEAP_SIZE +#define HEAP_SIZE 0x10000 +#endif + +__STACK_SIZE = STACK_SIZE; +__HEAP_SIZE = HEAP_SIZE; + +MEMORY +{ + ITCM (rwx) : ORIGIN = 0x10000000, LENGTH = 0x00008000 + BROM (rx) : ORIGIN = 0x11000000, LENGTH = 0x00020000 + BRAM (rwx) : ORIGIN = 0x12000000, LENGTH = 0x00200000 + DTCM (rw) : ORIGIN = 0x30000000, LENGTH = 0x00008000 + SRAM (rw) : ORIGIN = 0x31000000, LENGTH = 0x00400000 + QSPI (rw) : ORIGIN = 0x38000000, LENGTH = 0x00800000 + DDR (rw) : ORIGIN = 0x70000000, LENGTH = 0x10000000 +} + +PHDRS +{ + rom_boot PT_LOAD; + rom_exec PT_LOAD; + rom_dram PT_LOAD; + null PT_NULL; +} + +/* Linker script to place sections and symbol values. Should be used together + * with other linker script that defines memory regions ITCM and RAM. + * It references following symbols, which must be defined in code: + * Reset_Handler : Entry of reset handler + * + * It defines following symbols, which code can use without definition: + * __exidx_start + * __exidx_end + * __copy_table_start__ + * __copy_table_end__ + * __zero_table_start__ + * __zero_table_end__ + * __etext + * __data_start__ + * __preinit_array_start + * __preinit_array_end + * __init_array_start + * __init_array_end + * __fini_array_start + * __fini_array_end + * __data_end__ + * __bss_start__ + * __bss_end__ + * __end__ + * end + * __HeapLimit + * __StackLimit + * __StackTop + * __stack + */ + +ENTRY(Reset_Handler) + +SECTIONS +{ + .text.vectors : + { + KEEP(*(.vectors)) + } > BROM :rom_boot + + /* + /* Vector table relocation to read write memory + * Alignment requirement from up to 496 interrupts, rounded to the closest + * power of two equals 512 (words), thus 2048 bytes. + */ + .data.vtable_rw (COPY): + { + . = ALIGN(0x800); + KEEP(*(.vtable_rw)) + } > ITCM :null + + .text : + { + *crt* (.text*) + *startup_ARMCM85.c.obj (.text*) + *system_ARMCM85.c.obj (.text*) + *target.cpp.obj (.text*) + + KEEP(*(.init)) + KEEP(*(.fini)) + + /* .ctors */ + *crtbegin.o(.ctors) + *crtbegin?.o(.ctors) + *(EXCLUDE_FILE(*crtend?.o *crtend.o) .ctors) + *(SORT(.ctors.*)) + *(.ctors) + + /* .dtors */ + *crtbegin.o(.dtors) + *crtbegin?.o(.dtors) + *(EXCLUDE_FILE(*crtend?.o *crtend.o) .dtors) + *(SORT(.dtors.*)) + *(.dtors) + + *(.text*) + + KEEP(*(.eh_frame*)) + } > BRAM :rom_exec + + .data : + { + . = ALIGN(4); + __data_start__ = .; + + *(vtable) + *(.data) + *(.data.*) + *(.rodata*) + + . = ALIGN(4); + __data_end__ = .; + } > BRAM :rom_exec + + .ARM.extab : + { + *(.ARM.extab* .gnu.linkonce.armextab.*) + } > BRAM :rom_exec + + .ARM.exidx : + { + __exidx_start = .; + *(.ARM.exidx* .gnu.linkonce.armexidx.*) + __exidx_end = .; + } > BRAM :rom_exec + + .copy.table : + { + . = ALIGN(4); + __copy_table_start__ = .; + + LONG (LOADADDR(.sram)) + LONG (ADDR(.sram)) + LONG (SIZEOF(.sram) / 4) + + __copy_table_end__ = .; + } > BRAM :rom_exec + + .zero.table : + { + . = ALIGN(4); + __zero_table_start__ = .; + + LONG (ADDR(.bss)) + LONG (SIZEOF(.bss) / 4) + + LONG (ADDR(.sram.bss)) + LONG (SIZEOF(.sram.bss) / 4) + + __zero_table_end__ = .; + } > BRAM :rom_exec + + .sram : AT(__etext) + { +#if (ETHOSU_MODEL == 0) + . = ALIGN(16); + *(network_model_sec) +#endif + + . = ALIGN(16); + *(.sram.data) + + . = ALIGN(4); + /* preinit data */ + PROVIDE_HIDDEN (__preinit_array_start = .); + KEEP(*(.preinit_array)) + PROVIDE_HIDDEN (__preinit_array_end = .); + + . = ALIGN(4); + /* init data */ + PROVIDE_HIDDEN (__init_array_start = .); + KEEP(*(SORT(.init_array.*))) + KEEP(*(.init_array)) + PROVIDE_HIDDEN (__init_array_end = .); + + . = ALIGN(4); + /* finit data */ + PROVIDE_HIDDEN (__fini_array_start = .); + KEEP(*(SORT(.fini_array.*))) + KEEP(*(.fini_array)) + PROVIDE_HIDDEN (__fini_array_end = .); + + KEEP(*(.jcr*)) + } > SRAM :rom_dram + + .sram.bss : + { +#if (ETHOSU_ARENA == 0) + . = ALIGN(32); + *(.bss.tensor_arena) +#endif + + . = ALIGN(16); + *(.bss.ethosu_scratch); + } > SRAM :null + + .ddr : + { +#if (ETHOSU_ARENA == 1) + . = ALIGN(32); + *(.bss.tensor_arena) +#endif + +#if (ETHOSU_MODEL == 1) + . = ALIGN(16); + *(network_model_sec) +#endif + + . = ALIGN(4); + *(input_data_sec) + *(expected_output_data_sec) + *(output_data_sec) + + *(ethosu_core_in_queue ethosu_core_out_queue) + + /* Place data for scatter loading here */ + __etext = .; + } > DDR :rom_dram + .ddr_noload (NOLOAD) : + { + . = ALIGN(16); + *(input_data_sec) + } > DDR :null + .bss : + { + . = ALIGN(4); + __bss_start__ = .; + + *(.bss) + *(.bss.*) + *(COMMON) + + . = ALIGN(4); + __bss_end__ = .; + } > BRAM :null + + .heap (ORIGIN(BRAM) + LENGTH(BRAM) - __HEAP_SIZE) (COPY) : + { + . = ALIGN(8); + __end__ = .; + PROVIDE(end = .); + . = . + __HEAP_SIZE; + . = ALIGN(8); + __HeapLimit = .; + } > BRAM :null + + .stack (ORIGIN(DTCM) + LENGTH(DTCM) - __STACK_SIZE) (COPY) : + { + . = ALIGN(8); + __StackLimit = .; + . = . + __STACK_SIZE; + . = ALIGN(8); + __StackTop = .; + } > DTCM :null + PROVIDE(__stack = __StackTop); + + /* Check if stack exceeds DTCM limit */ + ASSERT(LENGTH(DTCM) >= __STACK_SIZE, "region DTCM overflowed with stack") +} diff --git a/examples/arm/executor_runner/arm_executor_runner.cpp b/examples/arm/executor_runner/arm_executor_runner.cpp index 794c271154e..d56710e27ad 100644 --- a/examples/arm/executor_runner/arm_executor_runner.cpp +++ b/examples/arm/executor_runner/arm_executor_runner.cpp @@ -6,6 +6,64 @@ * LICENSE file in the root directory of this source tree. */ +/* This is an example executorch runner running on Arm Cortex-m and Ethos-U + * based hardware. This example tries to illustrate a few ways to use ExecuTorch + * and you can use it as is or remove the unneeded parts. Please use this code + * as inpiration. + * + * Some defines used to configure the code: + * + * ET_MODEL_PTE_ADDR - Where in memory/flash your PTE model data is, if + * not set the model is supposed to have been converted to + * a c-array named model_pte and put into model_pte.h + * this is placed in network_model_sec linker section + * that is controlled by your memory mode via the + * ETHOSU_MODEL cmake parameter. + * If SEMIHOSTING is define this is not used + * ET_DUMP_INPUT - Control if you want input to be dumped to the log. + * ET_DUMP_OUTPUT - Control if you want output to be dumped to the log. + * ET_BUNDLE_IO - Build in devtools BundelIO, this makes it possible to + * use bpte with bundled input and output refdata to + * compare output. + * See also ET_ATOL and ET_RTOL + * ET_ATOL - The atol used to compare the output and ref data when + * using ET_BUNDLE_IO + * ET_RTOL - The rtol used to compare the output and ref data when + * using ET_BUNDLE_IO + * ET_EVENT_TRACER_ENABLED - Build in devtools event trace code to generate + * ETDump and print it base64 coded of it in the logs + * so you can get it out of your embedded target. + * This can be used to benchmark where time is spent. + * If you run on Ethos-U the delegate/commandstream + * is run in one go, this means that per op + * measurements is not possible. + * Warning: CPU time meassurements is NOT possible in the FVP simulator and a + * real target or FPGA must be used. NPU number are roughly OK, and can be used + * as guidance if timeing adaptor values are set correctly. + * + * SEMIHOSTING - When using the FVP simulator it can be built to access your dev + * machines filesystem, this is used for testing models in + * unittest/pytest and a special version of the runner is built + * to read model and input as files and output is saved to the + * filesystem. The backends/arm/test/setup_testing.sh script will + * build this for you so you can use it from pytest to test with + * the FVP simulator. + * + * Memory areas used: + * You might want to configure this differently on your HW, like maybe all + * left over memory after code is linked. This needs to be big enough to fit + * and run your model. In our example using the FVP simulator we have much + * memory and set this quite high to be able to test larger models. + * Regarding heap/mallocs type of allocation from executorch, + * et_pal_allocate() is not implemented or needed. + * + * ET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE - Size of memory area + * used when setting up the model + * ET_ARM_BAREMETAL_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE - Size of memory area + * used when running + * inferences + */ + #include #include #include @@ -55,7 +113,17 @@ unsigned char __attribute__(( char* model_pte = nullptr; #else +#if defined(ET_MODEL_PTE_ADDR) +/** + * Set ET_MODEL_PTE_ADDR to the memory address where your PTE is placed + * e.g. if you for example flash it to 0x7000000 set + * -DET_MODEL_PTE_ADDR=0x7000000 You can run the Corstone FVP with the --data + * flag to place it on a address if you use the FVP. + */ +char* model_pte = reinterpret_cast(ET_MODEL_PTE_ADDR); + +#else /** * This header file is generated by the build process based on the .pte file * specified in the ET_PTE_FILE_PATH variable to the cmake build. @@ -66,13 +134,11 @@ char* model_pte = nullptr; * e.g. This includes the pte as a big chunk of data struct into this file */ #include "model_pte.h" - +#endif #endif using executorch::aten::ScalarType; using executorch::aten::Tensor; -using executorch::aten::TensorImpl; -using executorch::extension::BufferCleanup; using executorch::extension::BufferDataLoader; using executorch::runtime::Error; using executorch::runtime::EValue; @@ -94,6 +160,7 @@ using executorch::bundled_program::verify_method_outputs; #if defined(ET_EVENT_TRACER_ENABLED) using executorch::etdump::ETDumpGen; using executorch::etdump::ETDumpResult; +using torch::executor::etdump_result; #endif /** * The method_allocation_pool should be large enough to fit the setup, input @@ -130,6 +197,12 @@ const float et_rtol = 0.01; #endif +#if defined(ET_NUM_INFERENCES) +const int num_inferences = ET_NUM_INFERENCES; +#else +const int num_inferences = 1; +#endif + /** * The temp_allocation_pool is used for allocating temporary data during kernel * or delegate execution. This will be reset after each kernel or delegate call. @@ -431,7 +504,7 @@ struct RunnerContext { Box temp_allocator; Box> method; #if defined(ET_EVENT_TRACER_ENABLED) - Box etdump_gen; + Box etdump_gen; #endif #if defined(SEMIHOSTING) Box input_file_allocator; @@ -517,8 +590,9 @@ void runner_init( ET_LOG(Info, "Setting up planned buffer %zu, size %zu.", id, buffer_size); /* Move to it's own allocator when MemoryPlanner is in place. */ - uint8_t* buffer = - reinterpret_cast(ctx.method_allocator->allocate(buffer_size)); + /* Ethos-U driver requires 16 bit alignment. */ + uint8_t* buffer = reinterpret_cast( + ctx.method_allocator->allocate(buffer_size, 16UL)); ET_CHECK_MSG( buffer != nullptr, "Could not allocate memory for memory planned buffer size %zu", @@ -638,27 +712,24 @@ void runner_init( ET_LOG(Info, "Input prepared."); } -void run_model(RunnerContext& ctx) { - ET_LOG(Info, "Starting the model execution..."); - - StartMeasurements(); - // Run the model. - Error status = ctx.method.value()->execute(); - StopMeasurements(); - - ET_CHECK_MSG( - status == Error::Ok, - "Execution of method %s failed with status 0x%" PRIx32, - ctx.method_name, - status); -} - void log_mem_status(const RunnerContext& ctx) { size_t executor_memsize = ctx.method_allocator->used_size() - ctx.executor_membase; +#if defined(ET_MODEL_PTE_ADDR) + ET_LOG( + Info, + "model_pte_program_size: %lu bytes. (pte size unknown when not baked into elf)", + ctx.program_data_len); + ET_LOG( + Info, + "model_pte_loaded_size: %lu bytes. (pte size unknown when not baked into elf)", + ctx.pte_size); +#else ET_LOG(Info, "model_pte_program_size: %lu bytes.", ctx.program_data_len); ET_LOG(Info, "model_pte_loaded_size: %lu bytes.", ctx.pte_size); +#endif + #if defined(SEMIHOSTING) if (ctx.input_file_allocator->size() > 0) { ET_LOG( @@ -692,13 +763,7 @@ void log_mem_status(const RunnerContext& ctx) { ET_LOG(Info, "method_allocator_executor: %zu bytes", executor_memsize); } if (ctx.temp_allocator->size() > 0) { - ET_LOG( - Info, - "peak_temp_allocator: %zu / %zu free: %zu ( used: %zu %% ) ", - ctx.temp_allocator->peak_used(), - ctx.temp_allocator->size(), - ctx.temp_allocator->free_size(), - 100 * ctx.temp_allocator->peak_used() / ctx.temp_allocator->size()); + ET_LOG(Info, "temp_allocator: %zu", ctx.temp_allocator->size()); } } @@ -770,7 +835,7 @@ void write_etdump(RunnerContext& ctx) { if (result.buf != nullptr && result.size > 0) { // On a device with no file system we can't just write it out // to the file-system so we base64 encode it and dump it on the log. - int mode = 0; + int mode = base64_enc_modifier_padding | base64_dec_modifier_skipspace; size_t len = result.size; size_t encoded_len = base64_encoded_size(result.size, mode); uint8_t* encoded_buf = reinterpret_cast( @@ -781,7 +846,7 @@ void write_etdump(RunnerContext& ctx) { encoded_buf[encoded_len] = 0x00; // Ensure null termination ET_LOG(Info, "Writing etdump.bin [base64]"); printf( - "#---\necho \"%s\" | base64 -d >etdump.bin\npython3 -m devtools.inspector.inspector_cli --etdump_path etdump.bin --source_time_scale cycles --target_time_scale cycles\n#---\n", + "#[RUN THIS]\necho \"%s\" | base64 -d >etdump.bin\npython3 -m devtools.inspector.inspector_cli --etdump_path etdump.bin --source_time_scale cycles --target_time_scale cycles\n#[END]\n", encoded_buf); } else { ET_LOG( @@ -793,11 +858,11 @@ void write_etdump(RunnerContext& ctx) { #else // Dump the etdump data containing profiling/debugging data to the specified // file. - etdump_result result = etdump_gen.get_etdump_data(); + etdump_result result = ctx.etdump_gen->get_etdump_data(); if (result.buf != nullptr && result.size > 0) { // On a device with a file system we can just write it out // to the file-system. - char etdump_filename = "etdump.bin"; + const char* etdump_filename = "etdump.bin"; ET_LOG(Info, "Writing etdump to file: %s", etdump_filename); FILE* f = fopen(etdump_filename, "w+"); fwrite((uint8_t*)result.buf, 1, result.size, f); @@ -853,6 +918,37 @@ void verify_result(RunnerContext& ctx, const void* model_pte) { #endif } +void run_model(RunnerContext& ctx, const void* model_pte) { + Error status; + ET_LOG(Info, "Starting running %d inferences...", num_inferences); + int n = 0; + StartMeasurements(); + for (n = 0; n < num_inferences; n++) { + ET_LOG(Debug, "Running inference number %d", n); + // Run the model. + status = ctx.method.value()->execute(); + if (status != Error::Ok) { + break; + } + // Reset the temporary allocator holding the scratch buffer between + // inferences. We want to reuse the temp_allocator between inferences of the + // same Ethos-U custom delegate, not allocate memory with every new + // inference. + ctx.temp_allocator.reset(temp_allocation_pool_size, temp_allocation_pool); + } + StopMeasurements(n); + + ET_CHECK_MSG( + status == Error::Ok, + "Execution of method %s failed with status 0x%" PRIx32, + ctx.method_name, + status); + + ET_LOG(Info, "%d inferences finished", num_inferences); + print_outputs(ctx); + verify_result(ctx, model_pte); +} + } // namespace int main(int argc, const char* argv[]) { @@ -877,7 +973,15 @@ int main(int argc, const char* argv[]) { executorch::runtime::runtime_init(); std::vector> input_buffers; + +#if defined(ET_MODEL_PTE_ADDR) + // pte not in a known array but just on a memory/flash address + // As we dont know the size we pick something big enough + // Actual model is read from this area. + size_t pte_size = 0x10000000; +#else size_t pte_size = sizeof(model_pte); +#endif RunnerContext ctx; @@ -930,15 +1034,22 @@ int main(int argc, const char* argv[]) { } } #endif + + // Byte 4-7 is usually a nice magic number that could be good to print to make + // sure it's OK ETxx for PTE and BPxx for bundled pte where xx is a number. ET_LOG( - Info, "PTE in %p %c Size: %lu bytes", model_pte, model_pte[0], pte_size); + Info, + "PTE @ %p [----%c%c%c%c]", + model_pte, + model_pte[4], + model_pte[5], + model_pte[6], + model_pte[7]); runner_init(ctx, input_buffers, pte_size); - run_model(ctx); + run_model(ctx, model_pte); log_mem_status(ctx); - print_outputs(ctx); write_etdump(ctx); - verify_result(ctx, model_pte); ET_LOG(Info, "Program complete, exiting."); #if defined(SEMIHOSTING) diff --git a/examples/arm/executor_runner/arm_memory_allocator.cpp b/examples/arm/executor_runner/arm_memory_allocator.cpp index e22439a239d..6b627625ae1 100644 --- a/examples/arm/executor_runner/arm_memory_allocator.cpp +++ b/examples/arm/executor_runner/arm_memory_allocator.cpp @@ -7,7 +7,7 @@ #include "arm_memory_allocator.h" ArmMemoryAllocator::ArmMemoryAllocator(uint32_t size, uint8_t* base_address) - : MemoryAllocator(size, base_address), used_(0), peak_used_(0) {} + : MemoryAllocator(size, base_address), used_(0) {} void* ArmMemoryAllocator::allocate(size_t size, size_t alignment) { void* ret = executorch::runtime::MemoryAllocator::allocate(size, alignment); @@ -22,8 +22,6 @@ void* ArmMemoryAllocator::allocate(size_t size, size_t alignment) { } else { used_ = (used_ | (alignment - 1)) + 1 + size; } - if (used_ > peak_used_) - peak_used_ = used_; } return ret; } @@ -32,10 +30,6 @@ size_t ArmMemoryAllocator::used_size() const { return used_; } -size_t ArmMemoryAllocator::peak_used() const { - return peak_used_; -} - size_t ArmMemoryAllocator::free_size() const { return executorch::runtime::MemoryAllocator::size() - used_; } diff --git a/examples/arm/executor_runner/arm_memory_allocator.h b/examples/arm/executor_runner/arm_memory_allocator.h index f7e8939c655..1d7bbdecb4c 100644 --- a/examples/arm/executor_runner/arm_memory_allocator.h +++ b/examples/arm/executor_runner/arm_memory_allocator.h @@ -21,15 +21,10 @@ class ArmMemoryAllocator : public executorch::runtime::MemoryAllocator { // Returns the used size of the allocator's memory buffer. size_t used_size() const; - // Returns the peak memory usage of the allocator's memory buffer - // Peak usage is useful when doing multiple allocations & resets - size_t peak_used() const; - // Returns the free size of the allocator's memory buffer. size_t free_size() const; void reset(); private: size_t used_; - size_t peak_used_; }; diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp index 82ecc222c11..58a47105743 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.cpp +++ b/examples/arm/executor_runner/arm_perf_monitor.cpp @@ -4,8 +4,8 @@ * LICENSE file in the root directory of this source tree. */ +#include #include -#include #include "arm_perf_monitor.h" @@ -14,29 +14,31 @@ #include #include -static uint32_t ethosu_inference_count = 0; -static uint64_t ethosu_ArmCycleCountStart = 0; -static uint64_t ethosu_ArmBackendExecuteCycleCountStart = 0; -static uint64_t ethosu_ArmBackendExecuteCycleCount = 0; -static uint64_t ethosu_ArmWhenNPURunCycleCountStart = 0; -static uint64_t ethosu_ArmWhenNPURunCycleCount = 0; -static uint64_t ethosu_pmuCycleCount = 0; -static std::vector ethosu_pmuEventCounts( - ETHOSU_PMU_Get_NumEventCounters(), - 0); +namespace { #if defined(ETHOSU55) || defined(ETHOSU65) -static const uint32_t ethosu_pmuCountersUsed = 4; +const uint32_t ethosu_pmuCountersUsed = 4; #elif defined(ETHOSU85) -static const uint32_t ethosu_pmuCountersUsed = 5; +const uint32_t ethosu_pmuCountersUsed = 5; #else #error No NPU target defined #endif +uint32_t ethosu_delegation_count = 0; +uint64_t ethosu_ArmCycleCountStart = 0; +uint64_t ethosu_ArmBackendExecuteCycleCountStart = 0; +uint64_t ethosu_ArmBackendExecuteCycleCount = 0; +uint64_t ethosu_ArmWhenNPURunCycleCountStart = 0; +uint64_t ethosu_ArmWhenNPURunCycleCount = 0; +uint64_t ethosu_pmuCycleCount = 0; +std::array ethosu_pmuEventCounts = {0}; + // ethosu_pmuCountersUsed should match numbers of counters setup in // ethosu_inference_begin() and not be more then the HW supports static_assert(ETHOSU_PMU_NCOUNTERS >= ethosu_pmuCountersUsed); +} // namespace + extern "C" { // Callback invoked at start of NPU execution @@ -85,7 +87,7 @@ 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_inference_count++; + ethosu_delegation_count++; ethosu_pmuCycleCount += ETHOSU_PMU_Get_CCNTR(drv); for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { @@ -113,6 +115,7 @@ void EthosUBackend_execute_end() { } void StartMeasurements() { + ethosu_delegation_count = 0; ethosu_ArmBackendExecuteCycleCount = 0; ethosu_ArmWhenNPURunCycleCount = 0; ethosu_pmuCycleCount = 0; @@ -123,32 +126,43 @@ void StartMeasurements() { ethosu_ArmCycleCountStart = ARM_PMU_Get_CCNTR(); } -void StopMeasurements() { +void StopMeasurements(int num_inferences) { ARM_PMU_CNTR_Disable( PMU_CNTENCLR_CCNTR_ENABLE_Msk | PMU_CNTENCLR_CNT0_ENABLE_Msk | PMU_CNTENCLR_CNT1_ENABLE_Msk); uint32_t cycle_count = ARM_PMU_Get_CCNTR() - ethosu_ArmCycleCountStart; // Number of comand streams handled by the NPU - ET_LOG(Info, "NPU Inferences : %d", ethosu_inference_count); + ET_LOG(Info, "NPU Inferences : %d", num_inferences); + ET_LOG( + Info, + "NPU delegations: %d (%.2f per inference)", + ethosu_delegation_count, + (double)ethosu_delegation_count / num_inferences); ET_LOG(Info, "Profiler report, CPU cycles per operator:"); // This is number of CPU cycles for the ethos-u operator from start to finish // in the framework If there is more then one commandstream the time is added // together ET_LOG( Info, - "ethos-u : cycle_cnt : %d cycles", - ethosu_ArmBackendExecuteCycleCount); + "ethos-u : cycle_cnt : %d cycles (%.2f per inference)", + ethosu_ArmBackendExecuteCycleCount, + (double)ethosu_ArmBackendExecuteCycleCount / num_inferences); // 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 ET_LOG( Info, - "Operator(s) total: %d CPU cycles", - ethosu_ArmBackendExecuteCycleCount); + "Operator(s) total: %d CPU cycles (%.2f per inference)", + ethosu_ArmBackendExecuteCycleCount, + (double)ethosu_ArmBackendExecuteCycleCount / num_inferences); // Total CPU cycles used in the executorch method->execute() // Other delegates and no delegates are counted in this - ET_LOG(Info, "Inference runtime: %d CPU cycles total", cycle_count); + ET_LOG( + Info, + "Inference runtime: %d CPU cycles total (%.2f per inference)", + cycle_count, + (double)cycle_count / num_inferences); ET_LOG( Info, @@ -174,14 +188,24 @@ void StopMeasurements() { // If there is more then one commandstream the time is added together ET_LOG( Info, - "cpu_wait_for_npu_cntr : %" PRIu64 " CPU cycles", - ethosu_ArmWhenNPURunCycleCount); + "cpu_wait_for_npu_cntr : %" PRIu64 " CPU cycles (%.2f per inference)", + ethosu_ArmWhenNPURunCycleCount, + (double)ethosu_ArmWhenNPURunCycleCount / num_inferences); ET_LOG(Info, "Ethos-U PMU report:"); - ET_LOG(Info, "ethosu_pmu_cycle_cntr : %" PRIu64, ethosu_pmuCycleCount); + ET_LOG( + Info, + "ethosu_pmu_cycle_cntr : % " PRIu64 " (%.2f per inference)", + ethosu_pmuCycleCount, + (double)ethosu_pmuCycleCount / num_inferences); for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { - ET_LOG(Info, "ethosu_pmu_cntr%zd : %" PRIu64, i, ethosu_pmuEventCounts[i]); + ET_LOG( + Info, + "ethosu_pmu_cntr%zd : %" PRIu64 " (%.2f per inference)", + i, + ethosu_pmuEventCounts[i], + (double)ethosu_pmuEventCounts[i] / num_inferences); } #if defined(ETHOSU55) || defined(ETHOSU65) ET_LOG( @@ -199,6 +223,8 @@ void StopMeasurements() { #else void StartMeasurements() {} -void StopMeasurements() {} +void StopMeasurements(int num_inferences) { + (void)num_inferences; +} #endif diff --git a/examples/arm/executor_runner/arm_perf_monitor.h b/examples/arm/executor_runner/arm_perf_monitor.h index 3925a9a5713..afce6562654 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.h +++ b/examples/arm/executor_runner/arm_perf_monitor.h @@ -1,4 +1,4 @@ -/* Copyright 2024 Arm Limited and/or its affiliates. +/* Copyright 2024-2025 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. @@ -7,4 +7,4 @@ #pragma once void StartMeasurements(); -void StopMeasurements(); +void StopMeasurements(int num_inferences); diff --git a/examples/arm/run.sh b/examples/arm/run.sh index f2bc303b739..77dddfe6451 100755 --- a/examples/arm/run.sh +++ b/examples/arm/run.sh @@ -23,7 +23,6 @@ model_input="" aot_arm_compiler_flag_delegate="--delegate" aot_arm_compiler_flag_quantize="--quantize" aot_arm_compiler_flags="" -portable_kernels="aten::_softmax.out" target="ethos-u55-128" output_folder_set=false output_folder="." @@ -35,10 +34,13 @@ build_only=false system_config="" config="" memory_mode="" +pte_placement="elf" et_build_root="${et_root_dir}/arm_test" ethos_u_scratch_dir=${script_dir}/ethos-u-scratch scratch_dir_set=false toolchain=arm-none-eabi-gcc +select_ops_list="aten::_softmax.out" +qdq_fusion_op=false function help() { echo "Usage: $(basename $0) [options]" @@ -49,20 +51,26 @@ function help() { echo " --aot_arm_compiler_flags= Extra flags to pass to aot compiler" echo " --no_delegate Do not delegate the model (can't override builtin models)" echo " --no_quantize Do not quantize the model (can't override builtin models)" - echo " --portable_kernels= Comma separated list of portable (non delagated) kernels to include Default: ${portable_kernels}" + echo " --portable_kernels= TO BE DEPRECATED: Alias to select_ops_list." + echo " --select_ops_list= Comma separated list of portable (non delagated) kernels to include Default: ${select_ops_list}" + echo " NOTE: This is used when select_ops_model is not possible to use, e.g. for semihosting or bundleio." + echo " See https://docs.pytorch.org/executorch/stable/kernel-library-selective-build.html for more information." echo " --target= Target to build and run for Default: ${target}" echo " --output= Target build output folder Default: ${output_folder}" echo " --bundleio Create Bundled pte using Devtools BundelIO with Input/RefOutput included" echo " --etdump Adds Devtools etdump support to track timing, etdump area will be base64 encoded in the log" echo " --build_type= Build with Release, Debug or RelWithDebInfo, default is ${build_type}" echo " --extra_build_flags= Extra flags to pass to cmake like -DET_ARM_BAREMETAL_METHOD_ALLOCATOR_POOL_SIZE=60000 Default: none " - echo " --build_only Only build, don't run FVP" - echo " --system_config= System configuration to select from the Vela configuration file (see vela.ini). Default: Ethos_U55_High_End_Embedded for EthosU55 targets, Ethos_U85_SYS_DRAM_Mid for EthosU85 targets." + echo " --build_only Only build, don't run" + echo " --toolchain= Ethos-U: Toolchain can be specified (e.g. bare metal as arm-none-eabi-gcc or zephyr as arm-zephyr-eabi-gcc Default: ${toolchain}" + echo " --system_config= Ethos-U: System configuration to select from the Vela configuration file (see vela.ini). Default: Ethos_U55_High_End_Embedded for EthosU55 targets, Ethos_U85_SYS_DRAM_Mid for EthosU85 targets." echo " NOTE: If given, this option must match the given target. This option also sets timing adapter values customized for specific hardware, see ./executor_runner/CMakeLists.txt." - echo " --config= System configuration file that specifies system configurations (vela.ini)" - echo " --memory_mode= Memory mode to select from the Vela configuration file (see vela.ini), e.g. Shared_Sram/Sram_Only. Default: 'Shared_Sram' for Ethos-U55 targets, 'Sram_Only' for Ethos-U85 targets" + echo " --config= Ethos-U: System configuration file that specifies system configurations (vela.ini)" + echo " --memory_mode= Ethos-U: Memory mode to select from the Vela configuration file (see vela.ini), e.g. Shared_Sram/Sram_Only. Default: 'Shared_Sram' for Ethos-U55 targets, 'Sram_Only' for Ethos-U85 targets" + echo " --pte_placement= Ethos-U: Control if runtime has PTE baked into the elf or if its placed in memory outside of the elf, defaults to ${pte_placement}" echo " --et_build_root= Executorch build output root folder to use, defaults to ${et_build_root}" echo " --scratch-dir= Path to your Ethos-U scrach dir if you not using default ${ethos_u_scratch_dir}" + echo " --qdq_fusion_op Enable QDQ fusion op" exit 0 } @@ -74,34 +82,45 @@ for arg in "$@"; do --aot_arm_compiler_flags=*) aot_arm_compiler_flags="${arg#*=}";; --no_delegate) aot_arm_compiler_flag_delegate="" ;; --no_quantize) aot_arm_compiler_flag_quantize="" ;; - --portable_kernels=*) portable_kernels="${arg#*=}";; + --portable_kernels=*) select_ops_list="${arg#*=}" ; echo "WARNING: --portable_kernels is DEPRECATED use select_ops_list." ;; + --select_ops_list=*) select_ops_list="${arg#*=}";; --target=*) target="${arg#*=}";; - --toolchain=*) toolchain="${arg#*=}";; --output=*) output_folder="${arg#*=}" ; output_folder_set=true ;; --bundleio) bundleio=true ;; --etdump) build_with_etdump=true ;; --build_type=*) build_type="${arg#*=}";; --extra_build_flags=*) extra_build_flags="${arg#*=}";; --build_only) build_only=true ;; + --toolchain=*) toolchain="${arg#*=}";; --system_config=*) system_config="${arg#*=}";; --config=*) config="${arg#*=}";; --memory_mode=*) memory_mode="${arg#*=}";; + --pte_placement=*) pte_placement="${arg#*=}";; --et_build_root=*) et_build_root="${arg#*=}";; --scratch-dir=*) ethos_u_scratch_dir="${arg#*=}" ; scratch_dir_set=true ;; + --qdq_fusion_op) qdq_fusion_op=true;; *) ;; esac done +if ! [[ ${pte_placement} == "elf" ]]; then + if ! [[ "$pte_placement" =~ ^0x[0-9a-fA-F]{1,16}$ ]]; then + echo "ERROR: Placing the PTE in memory failed, address is larger then 64bit $pte_placement" + exit 1 + fi + echo "PTE is expected to be placed at $pte_placement instead of the elf." +fi + # Default Ethos-u tool folder override with --scratch-dir= ethos_u_scratch_dir=$(realpath ${ethos_u_scratch_dir}) setup_path_script=${ethos_u_scratch_dir}/setup_path.sh if [[ ${toolchain} == "arm-none-eabi-gcc" ]]; then toolchain_cmake=${et_root_dir}/examples/arm/ethos-u-setup/${toolchain}.cmake -elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then +elif [[ ${toolchain} == "arm-zephyr-eabi-gcc" ]]; then toolchain_cmake=${et_root_dir}/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake else - echo "Error: Invalid toolchain selection, provided: ${tolchain}" + echo "Error: Invalid toolchain selection, provided: ${toolchain}" echo " Valid options are {arm-none-eabi-gcc, arm-zephyr-eabi-gcc}" exit 1; fi @@ -179,9 +198,12 @@ fi cd $et_root_dir devtools_flag="" bundleio_flag="" +etrecord_flag="" et_dump_flag="" +qdq_fusion_op_flag="" if [ "$build_with_etdump" = true ] ; then et_dump_flag="--etdump" + etrecord_flag="--etrecord" fi if [ "$bundleio" = true ] ; then @@ -189,19 +211,22 @@ if [ "$bundleio" = true ] ; then bundleio_flag="--bundleio" fi +if [ "$qdq_fusion_op" = true ] ; then + qdq_fusion_op_flag="--enable_qdq_fusion_pass" +fi + backends/arm/scripts/build_executorch.sh --et_build_root="${et_build_root}" --build_type=$build_type $devtools_flag $et_dump_flag --toolchain="${toolchain}" -backends/arm/scripts/build_portable_kernels.sh --et_build_root="${et_build_root}" --build_type=$build_type --portable_kernels=$portable_kernels --toolchain="${toolchain}" if [[ -z "$model_name" ]]; then # the test models run, and whether to delegate test_model=( - "softmax" # 0 - "add" # 1 - "add3" # 2 - "qadd" # 3 - "qadd2" # 4 - "qops" # 5 - "mv2" # 6 + "softmax" # 0 + "add" # 1 + "add3" # 2 + "qadd" # 3 + "qadd2" # 4 + "qops" # 5 + "mv2" # 6 ) model_compiler_flags=( "" # 0 softmax @@ -236,12 +261,11 @@ for i in "${!test_model[@]}"; do # Name aligned with default aot_arm_compiler output model_filename=${model_short_name}_arm_delegate_${target} fi - elf_folder=${model_filename} if [ "$bundleio" = true ] ; then - model_filename=${model_filename}.bpte + model_filename_ext=${model_filename}.bpte else - model_filename=${model_filename}.pte + model_filename_ext=${model_filename}.pte fi if [ "$output_folder_set" = false ] ; then @@ -250,16 +274,16 @@ for i in "${!test_model[@]}"; do mkdir -p ${output_folder} output_folder=$(realpath ${output_folder}) - pte_file="${output_folder}/${model_filename}" + pte_file="${output_folder}/${model_filename_ext}" # Remove old pte files - rm -f "${output_folder}/${model_filename}" + rm -f "${output_folder}/${model_filename_ext}" if [ "$model_input_set" = true ]; then model_compiler_flags="${model_compiler_flags} --model_input=${model_input}" fi - ARM_AOT_CMD="python3 -m examples.arm.aot_arm_compiler --model_name=${model} --target=${target} ${model_compiler_flags} --intermediate=${output_folder} --output=${pte_file} --system_config=${system_config} --memory_mode=${memory_mode} $bundleio_flag --config=${config}" + ARM_AOT_CMD="python3 -m examples.arm.aot_arm_compiler --model_name=${model} --target=${target} ${model_compiler_flags} --intermediate=${output_folder} --output=${pte_file} --system_config=${system_config} --memory_mode=${memory_mode} $bundleio_flag ${etrecord_flag} --config=${config} $qdq_fusion_op_flag" echo "CALL ${ARM_AOT_CMD}" >&2 ${ARM_AOT_CMD} 1>&2 @@ -271,14 +295,34 @@ for i in "${!test_model[@]}"; do if [[ ${target} == *"TOSA"* ]]; then echo "Build for ${target} skip generating a .elf and running it" + elif [[ ${target} == *"vgf"* ]]; then + echo "Build and run for VKML, (target: ${target})" + set -x + backends/arm/scripts/build_executor_runner_vkml.sh --build_type=${build_type} \ + --extra_build_flags="${extra_build_flags}" \ + --output="${output_folder}" + if [ "$build_only" = false ] ; then + backends/arm/scripts/run_vkml.sh --model=${pte_file} --build_path=${output_folder} + fi + set +x + else + # Build the application, the pte is imported as a header/c array or the address specified by --pte_placement + model_data="" + pte_file_or_mem="${pte_file}" + elf_file="${output_folder}/${model_filename}/cmake-out/arm_executor_runner" + if ! [[ ${pte_placement} == "elf" ]]; then + # Place PTE in memory specified by pte_placement + pte_file_or_mem="${pte_placement}" + model_data="--data=${pte_file}@${pte_placement}" + elf_file="${et_build_root}/${target}_${pte_placement}/cmake-out/arm_executor_runner" + fi + set -x - # Rebuild the application as the pte is imported as a header/c array - backends/arm/scripts/build_executor_runner.sh --et_build_root="${et_build_root}" --pte="${pte_file}" --build_type=${build_type} --target=${target} --system_config=${system_config} --memory_mode=${memory_mode} ${bundleio_flag} ${et_dump_flag} --extra_build_flags="${extra_build_flags}" --ethosu_tools_dir="${ethos_u_scratch_dir}" --toolchain="${toolchain}" + backends/arm/scripts/build_executor_runner.sh --et_build_root="${et_build_root}" --pte="${pte_file_or_mem}" --build_type=${build_type} --target=${target} --system_config=${system_config} --memory_mode=${memory_mode} ${bundleio_flag} ${et_dump_flag} --extra_build_flags="${extra_build_flags}" --ethosu_tools_dir="${ethos_u_scratch_dir}" --toolchain="${toolchain}" --select_ops_list="${select_ops_list}" if [ "$build_only" = false ] ; then # Execute the executor_runner on FVP Simulator - elf_file="${output_folder}/${elf_folder}/cmake-out/arm_executor_runner" - backends/arm/scripts/run_fvp.sh --elf=${elf_file} --target=$target + backends/arm/scripts/run_fvp.sh --elf=${elf_file} ${model_data} --target=$target fi set +x fi diff --git a/examples/arm/run_mcu_models_fvp.sh b/examples/arm/run_mcu_models_fvp.sh new file mode 100755 index 00000000000..68d5ec03003 --- /dev/null +++ b/examples/arm/run_mcu_models_fvp.sh @@ -0,0 +1,308 @@ +#!/usr/bin/env bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# Copyright 2023-2025 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. + +# Prerequisite steps: (run the following commands before running this script) +# 1. Setup your environment for Arm FVP +# a. Setup Conda environment / venv +# b. ./install_executorch.sh --clean ; ./install_executorch.sh --editable; +# c. examples/arm/setup.sh --i-agree-to-the-contained-eula; +# d. source examples/arm/ethos-u-scratch/setup_path.sh +# 2. bash examples/selective_build/test_selective_build.sh cmake + +set -u + +# Valid targets for MCU model validation +VALID_TARGETS=( + "cortex-m55" + "cortex-m85" +) + +# Default models for MCU validation with portable kernels +DEFAULT_MODELS=(mv2 mv3 lstm) +# Available models (on FVP) +AVAILABLE_MODELS=(mv2 mv3 lstm) +# Add the following models if you want to enable them later (atm they are not working on FVP) +# edsr w2l ic3 ic4 resnet18 resnet50 + +# Variables +TARGET="" +MODELS=() +PASSED_MODELS=() +FAILED_MODELS=() + +# Function to validate target +validate_target() { + local target=$1 + for valid_target in "${VALID_TARGETS[@]}"; do + if [[ "$target" == "$valid_target" ]]; then + return 0 + fi + done + return 1 +} + +# Function to validate models +validate_models() { + local invalid_models=() + for model in "${MODELS[@]}"; do + if [[ ! " ${AVAILABLE_MODELS[*]} " =~ " $model " ]]; then + invalid_models+=("$model") + fi + done + + if [[ ${#invalid_models[@]} -gt 0 ]]; then + echo "❌ Error: Invalid model(s): ${invalid_models[*]}" + echo "Available models: ${AVAILABLE_MODELS[*]}" + return 1 + fi + return 0 +} + +cpu_to_ethos_target() { + local cpu=$1 + case $cpu in + cortex-m55) + echo "ethos-u55-128" + ;; + cortex-m85) + echo "ethos-u85-128" + ;; + *) + echo "Unknown CPU: $cpu" >&2 + return 1 + ;; + esac +} + +# Function to show usage +show_usage() { + echo "Usage: $0 --target= [--models=]" + echo "" + echo "MCU Model Validation without delegation" + echo "" + echo "Required arguments:" + echo " --target= Target platform for validation" + echo "" + echo "Optional arguments:" + echo " --models= Comma-separated list of models to test" + echo " (overrides default model list)" + echo "" + echo "Valid targets:" + printf ' %s\n' "${VALID_TARGETS[@]}" + echo "" + echo "Available models:" + printf ' %s\n' "${AVAILABLE_MODELS[@]}" + echo "" + echo "Examples:" + echo " $0 --target=ethos-u85-128" + echo " $0 --target=ethos-u55-128 --models=mv2,mv3,resnet18" + echo "" + echo "Default behavior:" + echo " - Uses all available models: ${DEFAULT_MODELS[*]}" + echo " - Runs with portable kernels (no delegation)" +} + +# Function to display summary +show_summary() { + local total_models=${#MODELS[@]} + + echo "" + echo "════════════════════════════════════════════════════════════════" + echo "🏁 MCU MODEL VALIDATION SUMMARY - TARGET: $TARGET" + echo "════════════════════════════════════════════════════════════════" + echo "" + + # Show individual results + for model in "${MODELS[@]}"; do + if [[ " ${PASSED_MODELS[*]} " =~ " $model " ]]; then + printf "%-12s : ✅ Passed\n" "$model" + elif [[ " ${FAILED_MODELS[*]} " =~ " $model " ]]; then + printf "%-12s : ❌ Failed\n" "$model" + else + printf "%-12s : ⏭️ Skipped\n" "$model" + fi + done + + echo "" + echo "────────────────────────────────────────────────────────────────" + + # Show statistics + local passed_count=${#PASSED_MODELS[@]} + local failed_count=${#FAILED_MODELS[@]} + local success_rate=$((passed_count * 100 / total_models)) + + echo "📊 STATISTICS:" + echo " Total Models : $total_models" + echo " ✅ Passed : $passed_count" + echo " ❌ Failed : $failed_count" + echo " 📈 Success Rate : $success_rate%" + echo "" + + # Show model selection info + if [[ ${#MODELS[@]} -eq ${#DEFAULT_MODELS[@]} ]] && [[ "${MODELS[*]}" == "${DEFAULT_MODELS[*]}" ]]; then + echo "📋 Model Selection: Default (all available models)" + else + echo "📋 Model Selection: Custom (${MODELS[*]})" + fi + echo "" + + # Overall result + if [[ $failed_count -eq 0 ]]; then + echo "🎉 OVERALL RESULT: ALL TESTS PASSED!" + echo "🔧 Mode: Portable Kernels (No Delegation)" + else + echo "⚠️ OVERALL RESULT: $failed_count/$total_models TESTS FAILED" + echo "🔧 Mode: Portable Kernels (No Delegation)" + echo "" + echo "🔍 Failed models: ${FAILED_MODELS[*]}" + fi + + echo "════════════════════════════════════════════════════════════════" + echo "" +} + +# Parse command line arguments +while [[ $# -gt 0 ]]; do + case $1 in + --target=*) + TARGET="${1#*=}" + shift + ;; + --models=*) + IFS=',' read -ra MODELS <<< "${1#*=}" + shift + ;; + -h|--help) + show_usage + exit 0 + ;; + *) + echo "❌ Error: Unknown argument '$1'" + echo "" + show_usage + exit 1 + ;; + esac +done + +# Check if target is provided +if [[ -z "$TARGET" ]]; then + echo "❌ Error: --target argument is required" + echo "" + show_usage + exit 1 +fi + +# Validate target +if ! validate_target "$TARGET"; then + echo "❌ Error: Invalid target '$TARGET'" + echo "" + show_usage + exit 1 +fi + +# Use default models if none specified +if [[ ${#MODELS[@]} -eq 0 ]]; then + MODELS=("${DEFAULT_MODELS[@]}") +fi + +# Validate models +if ! validate_models; then + exit 1 +fi + +# Remove duplicates from models array +IFS=" " read -r -a MODELS <<< "$(printf '%s\n' "${MODELS[@]}" | sort -u | tr '\n' ' ')" + +echo "🎯 MCU Model Validation - Target: $TARGET" +echo "📋 Processing models: ${MODELS[*]}" +echo "🔧 Mode: Portable Kernels (No Delegation)" +echo "" + +echo "🔨 Building ExecuteTorch libraries (one-time setup)..." +if ! backends/arm/scripts/build_executorch.sh; then + echo "❌ Failed to build ExecuteTorch libraries" + exit 1 +fi +echo "✅ ExecuteTorch libraries built successfully" +echo "" + +ETHOS_TARGET=$(cpu_to_ethos_target "$TARGET") +if [[ $? -ne 0 ]]; then + echo "Invalid CPU target: $TARGET" + exit 1 +fi +echo "Using ETHOS target: $ETHOS_TARGET" + +# Process each model +for model in "${MODELS[@]}"; do + echo "=== 🚀 Processing $model for $TARGET ===" + + # Track if this model succeeds + MODEL_SUCCESS=true + + # Step 1: Create directory + echo "📁 Creating directory arm_test/$model" + mkdir -p "arm_test/$model" + + # Step 2: AOT compilation (quantized, no delegation = portable kernels) + echo "⚙️ AOT compilation for $model" + if ! python3 -m examples.arm.aot_arm_compiler \ + -m "$model" \ + --target="$ETHOS_TARGET" \ + --quantize \ + --output="arm_test/$model"; then + echo "❌ AOT compilation failed for $model" + MODEL_SUCCESS=false + fi + + # Step 3: Build executor runner (only if AOT succeeded) + if [[ "$MODEL_SUCCESS" == true ]]; then + echo "🔨 Building executor runner for $model" + if ! backends/arm/scripts/build_executor_runner.sh \ + --pte="arm_test/$model/${model}_arm_${ETHOS_TARGET}.pte" \ + --target="$ETHOS_TARGET" \ + --output="arm_test/$model"; then + echo "❌ Executor runner build failed for $model" + MODEL_SUCCESS=false + fi + fi + + # Step 4: Run on FVP (only if build succeeded) + if [[ "$MODEL_SUCCESS" == true ]]; then + echo "🏃 Running $model on FVP with portable kernels" + if ! backends/arm/scripts/run_fvp.sh \ + --elf="arm_test/$model/arm_executor_runner" \ + --target="$ETHOS_TARGET"; then + echo "❌ FVP execution failed for $model" + MODEL_SUCCESS=false + fi + fi + + # Record result + if [[ "$MODEL_SUCCESS" == true ]]; then + echo "✅ $model completed successfully" + PASSED_MODELS+=("$model") + else + echo "❌ $model failed" + FAILED_MODELS+=("$model") + fi + + echo "" +done + +# Show comprehensive summary +show_summary + +# Exit with appropriate code for CI +if [[ ${#FAILED_MODELS[@]} -eq 0 ]]; then + exit 0 # Success +else + exit 1 # Failure +fi diff --git a/examples/arm/setup.sh b/examples/arm/setup.sh index 6ab59d9544b..2aa6590c64d 100755 --- a/examples/arm/setup.sh +++ b/examples/arm/setup.sh @@ -9,9 +9,9 @@ set -u -######## +######################## ### Hardcoded constants -######## +######################## script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) et_dir=$(realpath $script_dir/../..) ARCH="$(uname -m)" @@ -23,10 +23,10 @@ target_toolchain="" enable_fvps=1 enable_vela=1 enable_model_converter=0 # model-converter tool for VGF output -enable_vgf_lib=0 # vgf reader - runtime backend dependency +enable_vgf_lib=0 # vgf reader - runtime backend dependency enable_emulation_layer=0 # Vulkan layer driver - emulates Vulkan ML extensions -mlsdk_manifest_url="" - +enable_vulkan_sdk=0 # Download and export Vulkan SDK required by emulation layer +mlsdk_manifest_url="https://github.com/arm/ai-ml-sdk-manifest.git" # Figure out if setup.sh was called or sourced and save it into "is_script_sourced" (return 0 2>/dev/null) && is_script_sourced=1 || is_script_sourced=0 @@ -36,34 +36,6 @@ toolchain_url="" toolchain_dir="" toolchain_md5_checksum="" -if [[ "${ARCH}" == "x86_64" ]]; then - # FVPs - corstone300_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-300/FVP_Corstone_SSE-300_11.22_20_Linux64.tgz?rev=018659bd574f4e7b95fa647e7836ccf4&hash=22A79103C6FA5FFA7AFF3BE0447F3FF9" - corstone300_model_dir="Linux64_GCC-9.3" - corstone300_md5_checksum="98e93b949d0fbac977292d8668d34523" - - corstone320_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-320/FVP_Corstone_SSE-320_11.27_25_Linux64.tgz?rev=a507bffc219a4d5792f1192ab7002d89&hash=D9A824AA8227D2E679C9B9787FF4E8B6FBE3D7C6" - corstone320_model_dir="Linux64_GCC-9.3" - corstone320_md5_checksum="3deb3c68f9b2d145833f15374203514d" -elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then - # FVPs - corstone300_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-300/FVP_Corstone_SSE-300_11.22_20_Linux64_armv8l.tgz?rev=9cc6e9a32bb947ca9b21fa162144cb01&hash=7657A4CF27D42E892E3F08D452AAB073" - corstone300_model_dir="Linux64_armv8l_GCC-9.3" - corstone300_md5_checksum="cbbabbe39b07939cff7a3738e1492ef1" - - corstone320_url="https://developer.arm.com/-/media/Arm%20Developer%20Community/Downloads/OSS/FVP/Corstone-320/FVP_Corstone_SSE-320_11.27_25_Linux64_armv8l.tgz?rev=b6ebe0923cb84f739e017385fd3c333c&hash=8965C4B98E2FF7F792A099B08831FE3CB6120493" - corstone320_model_dir="Linux64_armv8l_GCC-9.3" - corstone320_md5_checksum="3889f1d80a6d9861ea4aa6f1c88dd0ae" -else - echo "[main] Error: only x86-64 & aarch64/arm64 architecture is supported for now!"; exit 1; -fi - -# Vela -vela_repo_url="https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela" -vela_rev="d37febc1715edf0d236c2ff555739a8a9aadcf9a" - -# MLSDK dependencies -mlsdk_manifest_dir="ml-sdk-for-vulkan-manifest" # List of supported options and their descriptions OPTION_LIST=( @@ -156,6 +128,10 @@ function check_options() { enable_emulation_layer=1 shift ;; + --enable-vulkan-sdk) + enable_vulkan_sdk=1 + shift + ;; --disable-ethos-u-deps) enable_baremetal_toolchain=0 enable_fvps=0 @@ -166,6 +142,7 @@ function check_options() { enable_model_converter=1 enable_vgf_lib=1 enable_emulation_layer=1 + enable_vulkan_sdk=1 shift ;; --mlsdk-manifest-url) @@ -199,195 +176,44 @@ function check_options() { function setup_root_dir() { mkdir -p ${root_dir} root_dir=$(realpath ${root_dir}) - setup_path_script="${root_dir}/setup_path.sh" -} - -function check_fvp_eula () { - # Mandatory user arg --i-agree-to-the-contained-eula - eula_acceptance_by_variable="${ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA:-False}" - - if [[ "${eula_acceptance}" -eq 0 ]]; then - if [[ ${eula_acceptance_by_variable} != "True" ]]; then - echo "Must pass argument '--i-agree-to-the-contained-eula' to agree to EULA associated with downloading the FVP." - echo "Alternativly set environment variable ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA=True." - echo "Exiting!" - exit 1 - else - echo "Arm EULA for FVP agreed to with ARM_FVP_INSTALL_I_AGREE_TO_THE_CONTAINED_EULA=True environment variable" - fi - fi -} - -function setup_fvp() { - # check EULA, forward argument - check_fvp_eula - - if [[ "${OS}" != "Linux" ]]; then - # Check if FVP is callable - if command -v FVP_Corstone_SSE-300_Ethos-U55 &> /dev/null; then - echo "[${FUNCNAME[0]}] Info: FVP for MacOS seem to be installed. Continuing..." - return 0 # If true exit gracefully and proceed with setup - else - echo "[${FUNCNAME[0]}] Warning: FVP only supported with Linux OS, skipping FVP setup..." - echo "[${FUNCNAME[0]}] Warning: For MacOS, using https://github.com/Arm-Examples/FVPs-on-Mac is recommended." - echo "[${FUNCNAME[0]}] Warning: Follow the instructions and make sure the path is set correctly." - return 1 # Throw error. User need to install FVP according to ^^^ - fi - fi - - # Download and install the Corstone 300 FVP simulator platform - fvps=("corstone300" "corstone320") - - for fvp in "${fvps[@]}"; do - cd "${root_dir}" - if [[ ! -e "FVP_${fvp}.tgz" ]]; then - echo "[${FUNCNAME[0]}] Downloading FVP ${fvp}..." - url_variable=${fvp}_url - fvp_url=${!url_variable} - curl --output "FVP_${fvp}.tgz" "${fvp_url}" - md5_variable=${fvp}_md5_checksum - fvp_md5_checksum=${!md5_variable} - verify_md5 ${fvp_md5_checksum} FVP_${fvp}.tgz || exit 1 - fi - - echo "[${FUNCNAME[0]}] Installing FVP ${fvp}..." - rm -rf FVP-${fvp} - mkdir -p FVP-${fvp} - cd FVP-${fvp} - tar xf ../FVP_${fvp}.tgz - - # Install the FVP - case ${fvp} in - corstone300) - ./FVP_Corstone_SSE-300.sh --i-agree-to-the-contained-eula --force --destination ./ --quiet --no-interactive - ;; - corstone320) - ./FVP_Corstone_SSE-320.sh --i-agree-to-the-contained-eula --force --destination ./ --quiet --no-interactive - ;; - *) - echo "[${FUNCNAME[0]}] Error: Unknown FVP model ${fvp}. Exiting." - exit 1 - ;; - esac - done + setup_path_script="${root_dir}/setup_path" } -function select_toolchain() { - if [[ "${ARCH}" == "x86_64" ]]; then - if [[ "${OS}" == "Linux" ]]; then - if [[ "${target_toolchain}" == "zephyr" ]]; then - # TODO can include support for zephyr toolchain for other host platforms later - toolchain_url="https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/toolchain_linux-x86_64_arm-zephyr-eabi.tar.xz" - toolchain_dir="arm-zephyr-eabi" - toolchain_md5_checksum="93128be0235cf5cf5f1ee561aa6eac5f" - else - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi" - toolchain_md5_checksum="0601a9588bc5b9c99ad2b56133b7f118" - fi - else - echo "[main] Error: only Linux is currently supported for x86-64 architecture now!"; exit 1; - fi - elif [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then - if [[ "${OS}" == "Darwin" ]]; then - if [[ "${target_toolchain}" == "zephyr" ]]; then - echo "[main] Error: only Linux OS is currently supported for aarch64 architecture targeting Zephyr now!"; exit 1; - else - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi" - toolchain_md5_checksum="f1c18320bb3121fa89dca11399273f4e" - fi - elif [[ "${OS}" == "Linux" ]]; then - if [[ "${target_toolchain}" == "zephyr" ]]; then - # eventually, this can be support by downloading the the toolchain from - # "https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v0.17.2/toolchain_linux-aarch64_arm-zephyr-eabi.tar.xz" - # but for now, we error if user tries to specify this - echo "[main] Error: currently target_toolchain zephyr is only support for x86-64 Linux host systems!"; exit 1; - else - toolchain_url="https://armkeil.blob.core.windows.net/developer/Files/downloads/gnu/13.3.rel1/binrel/arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi.tar.xz" - toolchain_dir="arm-gnu-toolchain-13.3.rel1-aarch64-arm-none-eabi" - toolchain_md5_checksum="303102d97b877ebbeb36b3158994b218" - fi - fi - else - echo "[main] Error: only x86-64 & aarch64/arm64 architecture is supported for now!"; exit 1; - fi - echo "[main] Info selected ${toolchain_dir} for ${ARCH} - ${OS} platform" -} - -function setup_toolchain() { - # Download and install the arm toolchain (default is arm-none-eabi) - # setting --target-toolchain to zephyr sets this to arm-zephyr-eabi - cd "${root_dir}" - if [[ ! -e "${toolchain_dir}.tar.xz" ]]; then - echo "[${FUNCNAME[0]}] Downloading ${toolchain_dir} toolchain ..." - curl --output "${toolchain_dir}.tar.xz" -L "${toolchain_url}" - verify_md5 ${toolchain_md5_checksum} "${toolchain_dir}.tar.xz" || exit 1 - fi - - echo "[${FUNCNAME[0]}] Installing ${toolchain_dir} toolchain ..." - rm -rf "${toolchain_dir}" - tar xf "${toolchain_dir}.tar.xz" -} - -function setup_vela() { - pip install ethos-u-vela@git+${vela_repo_url}@${vela_rev} +function setup_ethos_u_tools() { + CMAKE_POLICY_VERSION_MINIMUM=3.5 BUILD_PYBIND=1 pip install --no-dependencies -r $et_dir/backends/arm/requirements-arm-ethos-u.txt } function create_setup_path(){ cd "${root_dir}" - echo "" > "${setup_path_script}" + clear_setup_path if [[ "${enable_fvps}" -eq 1 ]]; then - fvps=("corstone300" "corstone320") - for fvp in "${fvps[@]}"; do - model_dir_variable=${fvp}_model_dir - fvp_model_dir=${!model_dir_variable} - fvp_bin_path="${root_dir}/FVP-${fvp}/models/${fvp_model_dir}" - echo "export PATH=\${PATH}:${fvp_bin_path}" >> ${setup_path_script} - done - - # Fixup for Corstone-320 python dependency - echo "export LD_LIBRARY_PATH=${root_dir}/FVP-corstone320/python/lib/" >> ${setup_path_script} - - echo "hash FVP_Corstone_SSE-300_Ethos-U55" >> ${setup_path_script} - echo "hash FVP_Corstone_SSE-300_Ethos-U65" >> ${setup_path_script} - echo "hash FVP_Corstone_SSE-320" >> ${setup_path_script} + setup_path_fvp fi if [[ "${enable_baremetal_toolchain}" -eq 1 ]]; then - toolchain_bin_path="$(cd ${toolchain_dir}/bin && pwd)" - echo "export PATH=\${PATH}:${toolchain_bin_path}" >> ${setup_path_script} + setup_path_toolchain + fi + + if [[ "${enable_vulkan_sdk}" -eq 1 ]]; then + setup_path_vulkan fi if [[ "${enable_model_converter}" -eq 1 ]]; then - cd "${root_dir}" - model_converter_bin_path="$(cd ${mlsdk_manifest_dir}/sw/model-converter/build && pwd)" - echo "export PATH=\${PATH}:${model_converter_bin_path}" >> ${setup_path_script} + setup_path_model_converter fi - # Add Path for vgf-lib and emulation-layer if [[ "${enable_vgf_lib}" -eq 1 ]]; then - cd "${root_dir}" - model_vgf_lib_bin_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/build && pwd)" - echo "export PATH=\${PATH}:${model_vgf_lib_bin_path}" >> ${setup_path_script} + setup_path_vgf_lib fi if [[ "${enable_emulation_layer}" -eq 1 ]]; then - cd "${root_dir}" - model_emulation_layer_bin_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/build && pwd)" - echo "export PATH=\${PATH}:${model_emulation_layer_bin_path}" >> ${setup_path_script} + setup_path_emulation_layer fi -} -function check_platform_support() { - # Make sure we are on a supported platform - if [[ "${ARCH}" != "x86_64" ]] && [[ "${ARCH}" != "aarch64" ]] \ - && [[ "${ARCH}" != "arm64" ]]; then - echo "[main] Error: only x86-64 & aarch64 architecture is supported for now!" - exit 1 - fi + echo "[main] Update path by running 'source ${setup_path_script}.sh'" + echo "[main] Or for fish shell use 'source ${setup_path_script}.fish'" } @@ -401,7 +227,15 @@ if [[ $is_script_sourced -eq 0 ]]; then check_options "$@" + # Import utils + source $et_dir/backends/arm/scripts/utils.sh + source $et_dir/backends/arm/scripts/fvp_utils.sh + source $et_dir/backends/arm/scripts/toolchain_utils.sh + source $et_dir/backends/arm/scripts/vulkan_utils.sh + + echo "[main]: Checking platform and os" check_platform_support + check_os_support cd "${script_dir}" @@ -415,57 +249,47 @@ if [[ $is_script_sourced -eq 0 ]]; then echo "enable-model-converter=${enable_model_converter}" echo "enable-vgf-lib=${enable_vgf_lib}" echo "enable-emulation-layer=${enable_emulation_layer}" + echo "enable-vulkan-sdk=${enable_vulkan_sdk}" echo "enable-vela=${enable_vela}" echo "mlsdk-manifest-url=${mlsdk_manifest_url}" - # Import utils - source $et_dir/backends/arm/scripts/utils.sh - - # Select appropriate toolchain - select_toolchain # Setup toolchain if [[ "${enable_baremetal_toolchain}" -eq 1 ]]; then + # Select appropriate toolchain + select_toolchain setup_toolchain fi # Setup FVP if [[ "${enable_fvps}" -eq 1 ]]; then + check_fvp_eula setup_fvp + install_fvp fi - - if [[ -z "$mlsdk_manifest_url" && "${enable_model_converter}" -eq 1 ]]; then - echo "Warning: mlsdk-manifest-url is not set, but model converter setup is not skipped." - echo " Please set the --mlsdk-manifest-url option to the correct URL." - echo " Skipping MLSDK model converter setup." - enable_model_converter=0 # Q: Can we assume if we enable mlsdk, we will always enable model converter - enable_vgf_lib=0 - enable_emulation_layer=0 + # Setup Vulkan SDK + if [[ "${enable_vulkan_sdk}" -eq 1 ]]; then + setup_vulkan_sdk fi - if [[ "${enable_model_converter}" -eq 1 ]]; then + if [[ "${enable_model_converter}" -eq 1 || \ + "${enable_vgf_lib}" -eq 1 || \ + "${enable_emulation_layer}" -eq 1 ]]; then source $et_dir/backends/arm/scripts/mlsdk_utils.sh -u "${mlsdk_manifest_url}" - setup_model_converter ${root_dir} ${mlsdk_manifest_dir} ${enable_vgf_lib} ${enable_emulation_layer} + setup_model_converter ${root_dir} ${mlsdk_manifest_dir} ${enable_model_converter} ${enable_vgf_lib} ${enable_emulation_layer} fi - # Create new setup_path script - if [[ "${enable_baremetal_toolchain}" -eq 1 || \ - "${enable_fvps}" -eq 1 || \ - "${enable_model_converter}" -eq 1 ]]; then - create_setup_path - fi + # Create the setup_path.sh used to create the PATH variable for shell + create_setup_path - # Setup the tosa_reference_model - $et_dir/backends/arm/scripts/install_reference_model.sh ${root_dir} + # Setup the tosa_reference_model and dependencies + CMAKE_POLICY_VERSION_MINIMUM=3.5 BUILD_PYBIND=1 pip install --no-dependencies -r $et_dir/backends/arm/requirements-arm-tosa.txt - # Setup vela and patch in codegen fixes if [[ "${enable_vela}" -eq 1 ]]; then - setup_vela + setup_ethos_u_tools fi - echo "[main] update path by doing 'source ${setup_path_script}'" - echo "[main] success!" exit 0 fi diff --git a/examples/arm/vgf_minimal_example.ipynb b/examples/arm/vgf_minimal_example.ipynb new file mode 100644 index 00000000000..35378817a7d --- /dev/null +++ b/examples/arm/vgf_minimal_example.ipynb @@ -0,0 +1,320 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Copyright 2025 Arm Limited and/or its affiliates.\n", + "#\n", + "# This source code is licensed under the BSD-style license found in the\n", + "# LICENSE file in the root directory of this source tree." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# VGF Backend flow example\n", + "\n", + "This guide demonstrates the full flow for lowering a module using the VGF backend using ExecuTorch. \n", + "Tested on Linux x86_64. 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 MLSDK and Tosa using `examples/arm/setup.sh --disable-ethos-u-deps --enable-mlsdk-deps` (For further guidance, refer to https://docs.pytorch.org/executorch/main/tutorial-arm.html)\n", + "3. Export vulkan environment variables and add MLSDK components to PATH and LD_LIBRARY_PATH using `examples/arm/ethos-u-scratch/setup_path.sh`\n", + "\n", + "With all commands executed from the base `executorch` folder.\n", + "\n", + "*Some scripts in this notebook produce long output logs: Configuring the 'Customizing Notebook Layout' settings to enable 'Output:scrolling' and setting 'Output:Text Line Limit' makes this more manageable*" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## AOT Flow\n", + "\n", + "The first step is creating the PyTorch module and exporting it. Exporting converts the python code in the module into a graph structure. The result is still runnable python code, which can be displayed by printing the `graph_module` of the exported program. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "class Add(torch.nn.Module):\n", + " def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n", + " return x + y\n", + "\n", + "example_inputs = (torch.ones(1,1,1,1),torch.ones(1,1,1,1))\n", + "\n", + "model = Add()\n", + "model = model.eval()\n", + "exported_program = torch.export.export_for_training(model, example_inputs)\n", + "graph_module = exported_program.module()\n", + "\n", + "_ = graph_module.print_readable()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# VGF backend supports both INT and FP targets. \n", + "\n", + "To lower the graph_module for FP targets using the VGF backend, we run it through the default FP lowering pipeline. \n", + "\n", + "FP lowering can be customized for different subgraphs; the sequence shown here is the recommended workflow for VGF.\n", + "Because we are staying in floating-point precision, no calibration with example inputs is required. \n", + "\n", + "If you print the module again, you will see that nodes are left in FP form (or annotated with any necessary casts) without any quantize/dequantize wrappers.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder\n", + "from executorch.backends.arm.tosa import ( \n", + " TosaSpecification,\n", + ")\n", + "\n", + "# Create a compilation spec describing the floating point target.\n", + "tosa_spec = TosaSpecification.create_from_string(\"TOSA-1.0+FP\")\n", + "\n", + "spec_builder = ArmCompileSpecBuilder().vgf_compile_spec(tosa_spec)\n", + "compile_spec = spec_builder.build()\n", + "\n", + "_ = graph_module.print_readable()\n", + "\n", + "# Create a new exported program using the graph_module\n", + "exported_program = torch.export.export_for_training(graph_module, example_inputs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To lower the graph_module for INT targets using the VGF backend, we apply the arm_quantizer. \n", + "\n", + "Quantization can be performed in various ways and tailored to different subgraphs; the sequence shown here represents the recommended workflow for VGF. \n", + "\n", + "This step also requires calibrating the module with representative inputs. \n", + "\n", + "If you print the module again, you’ll see that each node is now wrapped in quantization/dequantization nodes that embed the calculated quantization parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from executorch.backends.arm.quantizer import (\n", + " VgfQuantizer,\n", + " get_symmetric_quantization_config,\n", + ")\n", + "from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e\n", + "\n", + "# Create a compilation spec describing the target for configuring the quantizer\n", + "tosa_spec = TosaSpecification.create_from_string(\"TOSA-1.0+INT\")\n", + "\n", + "spec_builder = ArmCompileSpecBuilder().vgf_compile_spec(tosa_spec)\n", + "compile_spec = spec_builder.build()\n", + "\n", + "# Create and configure quantizer to use a symmetric quantization config globally on all nodes\n", + "quantizer = VgfQuantizer(compile_spec)\n", + "operator_config = get_symmetric_quantization_config(is_per_channel=False)\n", + "quantizer.set_global(operator_config)\n", + "\n", + "# Post training quantization\n", + "quantized_graph_module = prepare_pt2e(graph_module, quantizer)\n", + "quantized_graph_module(*example_inputs) # Calibrate the graph module with the example input\n", + "quantized_graph_module = convert_pt2e(quantized_graph_module)\n", + "\n", + "_ = quantized_graph_module.print_readable()\n", + "\n", + "# Create a new exported program using the quantized_graph_module\n", + "quantized_exported_program = torch.export.export_for_training(quantized_graph_module, example_inputs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# In the example below, we will make use of the quantized graph module.\n", + "\n", + "The lowering in the VGFBackend happens in five steps:\n", + "\n", + "1. **Lowering to core Aten operator set**: Transform module to use a subset of operators applicable to edge devices. \n", + "2. **Partitioning**: Find subgraphs that will be lowered by the VGF backend.\n", + "3. **Lowering to TOSA compatible operator set**: Perform transforms to make the VGF subgraph(s) compatible with TOSA \n", + "4. **Serialization to TOSA**: Compiles the graph module into a TOSA graph \n", + "5. **Compilation to VGF**: Compiles the FX GraphModule into a VGF representation using the model_converter and the previously created compile_spec. It also prints a network summary for each processed VGF partition.\n", + "\n", + "All of this happens behind the scenes in `to_edge_transform_and_lower`. Printing the graph module shows that what is left in the graph is two quantization nodes for `x` and `y` going into an `executorch_call_delegate` node, followed by a dequantization node." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "# Ensure the vulkan environment variables and MLSDK components are available on $PATH\n", + "source ethos-u-scratch/setup_path.sh" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from executorch.backends.arm.vgf import VgfPartitioner\n", + "from executorch.exir import (\n", + " EdgeCompileConfig,\n", + " ExecutorchBackendConfig,\n", + " to_edge_transform_and_lower,\n", + ")\n", + "from executorch.extension.export_util.utils import save_pte_program\n", + "\n", + "# Create partitioner from compile spec\n", + "partitioner = VgfPartitioner(compile_spec)\n", + "\n", + "# Lower the exported program to the VGF backend\n", + "edge_program_manager = to_edge_transform_and_lower(\n", + " quantized_exported_program,\n", + " partitioner=[partitioner],\n", + " compile_config=EdgeCompileConfig(\n", + " _check_ir_validity=False,\n", + " ),\n", + ")\n", + "\n", + "# Convert edge program to executorch\n", + "executorch_program_manager = edge_program_manager.to_executorch(\n", + " config=ExecutorchBackendConfig(extract_delegate_segments=False)\n", + ")\n", + "\n", + "executorch_program_manager.exported_program().module().print_readable()\n", + "\n", + "# Save pte file\n", + "cwd_dir = os.getcwd()\n", + "pte_base_name = \"simple_example\"\n", + "pte_name = pte_base_name + \".pte\"\n", + "pte_path = os.path.join(cwd_dir, pte_name)\n", + "save_pte_program(executorch_program_manager, pte_name)\n", + "assert os.path.exists(pte_path), \"Build failed; no .pte-file found\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Build executor runtime\n", + "\n", + "### Prerequisite\n", + "With our VGF inside our PTE we now need to setup the runtime. To do this we will use the previously built MLSDK dependencies, but we will also need to setup a Vulkan environment externally to Executorch.\n", + "Plese follow https://vulkan.lunarg.com/sdk/home in order to setup. \n", + "\n", + "\n", + "After the AOT compilation flow is done, we need to build the executor_runner target. For this example the generic version will be used.\n", + "To do this, please ensure the following commands are executed before moving onto the next step." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%bash\n", + "# Ensure the vulkan environment variables and MLSDK components are available on $PATH\n", + "source ethos-u-scratch/setup_path.sh\n", + "\n", + "# Compiled programs will appear in the executorch/cmake-out directory we create here.\n", + "# Build example executor runner application to examples/arm/vgf_minimal_example\n", + "cmake \\\n", + " -DCMAKE_INSTALL_PREFIX=cmake-out \\\n", + " -DCMAKE_BUILD_TYPE=Debug \\\n", + " -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \\\n", + " -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \\\n", + " -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \\\n", + " -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \\\n", + " -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \\\n", + " -DEXECUTORCH_BUILD_XNNPACK=OFF \\\n", + " -DEXECUTORCH_BUILD_VULKAN=ON \\\n", + " -DEXECUTORCH_BUILD_VGF=ON \\\n", + " -DEXECUTORCH_ENABLE_LOGGING=ON \\\n", + " -DPYTHON_EXECUTABLE=python \\\n", + " -B../../cmake-out-vkml ../..\n", + "\n", + "cmake --build ../../cmake-out-vkml --target executor_runner" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Run on VKML Emulator\n", + "\n", + "We can finally use the `backends/arm/scripts/run_vkml.sh` utility script to run the .pte end-to-end and proving out a backend’s kernel implementation. This Script runs the model with an input of ones, so the expected result of the addition should be close to 2." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "\n", + "# Setup paths\n", + "et_dir = os.path.join(cwd_dir, \"..\", \"..\")\n", + "et_dir = os.path.abspath(et_dir)\n", + "script_dir = os.path.join(et_dir, \"backends\", \"arm\", \"scripts\")\n", + "\n", + "args = f\"--model={pte_path}\"\n", + "subprocess.run(os.path.join(script_dir, \"run_vkml.sh\") + \" \" + args, shell=True, cwd=et_dir)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/cadence/models/babyllama.py b/examples/cadence/models/babyllama.py index 58a3035723f..1b576a1a3eb 100644 --- a/examples/cadence/models/babyllama.py +++ b/examples/cadence/models/babyllama.py @@ -12,7 +12,7 @@ import torch -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from executorch.examples.models.llama.llama_transformer import ModelArgs, Transformer @@ -35,7 +35,7 @@ def main() -> None: model = Transformer(args) example_inputs = (torch.randint(0, 10, [b, seq], dtype=torch.int64),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) if __name__ == "__main__": diff --git a/examples/cadence/models/mobilenet_v2.py b/examples/cadence/models/mobilenet_v2.py index 3c31af88de9..9c0791a70e5 100644 --- a/examples/cadence/models/mobilenet_v2.py +++ b/examples/cadence/models/mobilenet_v2.py @@ -13,7 +13,7 @@ from executorch.backends.cadence.aot.ops_registrations import * # noqa -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from torchvision.models import mobilenet_v2, MobileNet_V2_Weights @@ -27,4 +27,4 @@ model.eval() example_inputs = (torch.randn(1, 3, 64, 64),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/resnet18.py b/examples/cadence/models/resnet18.py index ebf0533bd84..c728d4dc714 100644 --- a/examples/cadence/models/resnet18.py +++ b/examples/cadence/models/resnet18.py @@ -13,7 +13,7 @@ from executorch.backends.cadence.aot.ops_registrations import * # noqa -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from torchvision.models import resnet18, ResNet18_Weights @@ -27,4 +27,4 @@ model.eval() example_inputs = (torch.randn(1, 3, 64, 64),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/resnet50.py b/examples/cadence/models/resnet50.py index cac73433d1d..d330e912074 100644 --- a/examples/cadence/models/resnet50.py +++ b/examples/cadence/models/resnet50.py @@ -13,7 +13,7 @@ from executorch.backends.cadence.aot.ops_registrations import * # noqa -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from torchvision.models import resnet50, ResNet50_Weights @@ -27,4 +27,4 @@ model.eval() example_inputs = (torch.randn(1, 3, 64, 64),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/rnnt_encoder.py b/examples/cadence/models/rnnt_encoder.py index d89f1ca42e7..8fe5f08b4c4 100644 --- a/examples/cadence/models/rnnt_encoder.py +++ b/examples/cadence/models/rnnt_encoder.py @@ -14,7 +14,7 @@ from typing import List, Optional, Tuple -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from torchaudio.prototype.models import ConvEmformer @@ -157,4 +157,4 @@ def infer( transcriber_lengths, ) - export_model(encoder, example_inputs) + export_and_run_model(encoder, example_inputs) diff --git a/examples/cadence/models/rnnt_joiner.py b/examples/cadence/models/rnnt_joiner.py index b6aaa7c91b7..fbb2ddb404e 100644 --- a/examples/cadence/models/rnnt_joiner.py +++ b/examples/cadence/models/rnnt_joiner.py @@ -14,7 +14,7 @@ from typing import Tuple -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -62,4 +62,4 @@ def forward( target_encodings, ) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/rnnt_predictor.py b/examples/cadence/models/rnnt_predictor.py index d83f593d3e3..90825f18074 100644 --- a/examples/cadence/models/rnnt_predictor.py +++ b/examples/cadence/models/rnnt_predictor.py @@ -14,7 +14,7 @@ from typing import Tuple -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -66,4 +66,4 @@ def forward( predictor_lengths, ) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/vision_transformer.py b/examples/cadence/models/vision_transformer.py index 79c9e3f1963..7acbedcb83b 100644 --- a/examples/cadence/models/vision_transformer.py +++ b/examples/cadence/models/vision_transformer.py @@ -13,7 +13,7 @@ import torch import torchvision -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -25,4 +25,4 @@ model = torchvision.models.vit_b_16() example_inputs = (torch.randn(1, 3, 224, 224),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/models/wav2vec2.py b/examples/cadence/models/wav2vec2.py index 5db9ea2a6d5..a9270d85899 100644 --- a/examples/cadence/models/wav2vec2.py +++ b/examples/cadence/models/wav2vec2.py @@ -12,7 +12,7 @@ import torch -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from torchaudio.models.wav2vec2.model import wav2vec2_model, Wav2Vec2Model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -58,7 +58,7 @@ def forward(self, x): audio_len = 1680 example_inputs = (torch.rand(1, audio_len),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) if __name__ == "__main__": diff --git a/examples/cadence/operators/test_add_op.py b/examples/cadence/operators/test_add_op.py index 31177741536..8089fe3a6b1 100644 --- a/examples/cadence/operators/test_add_op.py +++ b/examples/cadence/operators/test_add_op.py @@ -15,7 +15,7 @@ import torch import torch.nn as nn -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model class ATenOpTestCases(unittest.TestCase): @@ -64,7 +64,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor): Y = torch.randn(Yshape) model.eval() - export_model( + export_and_run_model( model, (X, Y), file_name=self._testMethodName, run_and_compare=False ) @@ -114,7 +114,7 @@ def forward(self, x: torch.Tensor, y: float): Y = 2.34 model.eval() - export_model( + export_and_run_model( model, (X, Y), file_name=self._testMethodName, run_and_compare=False ) diff --git a/examples/cadence/operators/test_g3_ops.py b/examples/cadence/operators/test_g3_ops.py index c33290416be..b096399c1a6 100644 --- a/examples/cadence/operators/test_g3_ops.py +++ b/examples/cadence/operators/test_g3_ops.py @@ -10,13 +10,13 @@ import torch import torch.nn as nn -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model class ATenOpTestCases(unittest.TestCase): def run_and_verify(self, model: nn.Module, inputs: Tuple[Any, ...]) -> None: model.eval() - export_model( + export_and_run_model( model, inputs, file_name=self._testMethodName, run_and_compare=False ) diff --git a/examples/cadence/operators/test_quantized_conv1d_op.py b/examples/cadence/operators/test_quantized_conv1d_op.py index 06ae03e6cfd..c178f7187f9 100644 --- a/examples/cadence/operators/test_quantized_conv1d_op.py +++ b/examples/cadence/operators/test_quantized_conv1d_op.py @@ -16,7 +16,7 @@ from executorch.backends.cadence.aot.ops_registrations import * # noqa -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -59,4 +59,4 @@ def forward(self, x: torch.Tensor): example_inputs = (torch.randn(cast(Sequence[int], shape)),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/operators/test_quantized_linear_op.py b/examples/cadence/operators/test_quantized_linear_op.py index c220bf62610..86ba3412b6e 100644 --- a/examples/cadence/operators/test_quantized_linear_op.py +++ b/examples/cadence/operators/test_quantized_linear_op.py @@ -14,7 +14,7 @@ from executorch.backends.cadence.aot.ops_registrations import * # noqa -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -41,4 +41,4 @@ def forward(self, x: torch.Tensor): example_inputs = (torch.ones(shape),) - export_model(model, example_inputs) + export_and_run_model(model, example_inputs) diff --git a/examples/cadence/operators/test_requantize_op.py b/examples/cadence/operators/test_requantize_op.py index 4e4528d330a..b87895b109d 100644 --- a/examples/cadence/operators/test_requantize_op.py +++ b/examples/cadence/operators/test_requantize_op.py @@ -18,7 +18,7 @@ import itertools -from executorch.backends.cadence.aot.export_example import export_model +from executorch.backends.cadence.aot.export_example import export_and_run_model from parameterized import parameterized @@ -98,7 +98,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # For example, if the ref output is [33, 50] and the real output is [33, 49], # the MSE loss is around 0.5, but the relative error is < 2%. So we set # the epsilon to a higher value. - export_model(model, inputs, eps_error=1.0) + export_and_run_model(model, inputs, eps_error=1.0) if __name__ == "__main__": diff --git a/examples/demo-apps/android/LlamaDemo/app/build.gradle.kts b/examples/demo-apps/android/LlamaDemo/app/build.gradle.kts index 893b1ee4784..19cfda847db 100644 --- a/examples/demo-apps/android/LlamaDemo/app/build.gradle.kts +++ b/examples/demo-apps/android/LlamaDemo/app/build.gradle.kts @@ -11,6 +11,8 @@ plugins { id("org.jetbrains.kotlin.android") } +val qnnVersion: String? = project.findProperty("qnnVersion") as? String + android { namespace = "com.example.executorchllamademo" compileSdk = 34 @@ -61,6 +63,9 @@ dependencies { implementation("com.google.android.material:material:1.12.0") implementation("androidx.activity:activity:1.9.0") implementation("org.json:json:20250107") + if (!qnnVersion.isNullOrEmpty()) { + implementation("com.qualcomm.qti:qnn-runtime:$qnnVersion") + } testImplementation("junit:junit:4.13.2") androidTestImplementation("androidx.test.ext:junit:1.1.5") androidTestImplementation("androidx.test.espresso:espresso-core:3.5.1") diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java index b26031d89a6..f995c5bc65a 100644 --- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java +++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java @@ -81,7 +81,6 @@ public class MainActivity extends AppCompatActivity implements Runnable, LlmCall private Runnable memoryUpdater; private boolean mThinkMode = false; private int promptID = 0; - private long startPos = 0; private static final int CONVERSATION_HISTORY_MESSAGE_LOOKBACK = 2; private Executor executor; @@ -178,7 +177,8 @@ private void setLocalModel(String modelPath, String tokenizerPath, float tempera if (mCurrentSettingsFields.getModelType() == ModelType.LLAVA_1_5) { ETLogging.getInstance().log("Llava start prefill prompt"); - startPos = mModule.prefillPrompt(PromptFormat.getLlavaPresetPrompt(), 0, 1, 0); + mModule.resetContext(); + mModule.prefillPrompt(PromptFormat.getLlavaPresetPrompt()); ETLogging.getInstance().log("Llava completes prefill prompt"); } } @@ -645,13 +645,11 @@ private void showMediaPreview(List uris) { ETLogging.getInstance().log("Starting runnable prefill image"); ETImage img = processedImageList.get(0); ETLogging.getInstance().log("Llava start prefill image"); - startPos = - mModule.prefillImages( - img.getInts(), - img.getWidth(), - img.getHeight(), - ModelUtils.VISION_MODEL_IMAGE_CHANNELS, - startPos); + mModule.prefillImages( + img.getInts(), + img.getWidth(), + img.getHeight(), + ModelUtils.VISION_MODEL_IMAGE_CHANNELS); }; executor.execute(runnable); } @@ -778,12 +776,8 @@ public void run() { mCurrentSettingsFields.getModelType(), mCurrentSettingsFields.getBackendType()) == ModelUtils.VISION_MODEL) { - mModule.generateFromPos( - finalPrompt, - ModelUtils.VISION_MODEL_SEQ_LEN, - startPos, - MainActivity.this, - false); + mModule.generate( + finalPrompt, ModelUtils.VISION_MODEL_SEQ_LEN, MainActivity.this, false); } else if (mCurrentSettingsFields.getModelType() == ModelType.LLAMA_GUARD_3) { String llamaGuardPromptForClassification = PromptFormat.getFormattedLlamaGuardPrompt(rawPrompt); diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelUtils.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelUtils.java index 32ed33cd302..cf7ab1756ce 100644 --- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelUtils.java +++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelUtils.java @@ -21,6 +21,9 @@ public class ModelUtils { // MediaTek static final int MEDIATEK_TEXT_MODEL = 3; + // QNN static llama + static final int QNN_TEXT_MODEL = 4; + public static int getModelCategory(ModelType modelType, BackendType backendType) { if (backendType.equals(BackendType.XNNPACK)) { switch (modelType) { @@ -35,6 +38,8 @@ public static int getModelCategory(ModelType modelType, BackendType backendType) } } else if (backendType.equals(BackendType.MEDIATEK)) { return MEDIATEK_TEXT_MODEL; + } else if (backendType.equals(BackendType.QUALCOMM)) { + return QNN_TEXT_MODEL; } return TEXT_MODEL; // default diff --git a/examples/demo-apps/android/LlamaDemo/docs/delegates/mediatek_README.md b/examples/demo-apps/android/LlamaDemo/docs/delegates/mediatek_README.md index 2ad87df0653..f72e1b0fbc7 100644 --- a/examples/demo-apps/android/LlamaDemo/docs/delegates/mediatek_README.md +++ b/examples/demo-apps/android/LlamaDemo/docs/delegates/mediatek_README.md @@ -54,7 +54,6 @@ zstd -cdq ".zst" > "/buck2" && chmod ### Set Environment Variables ``` -export BUCK2=path_to_buck/buck2 # Download BUCK2 and create BUCK2 executable export ANDROID_NDK=path_to_android_ndk export NEURON_BUFFER_ALLOCATOR_LIB=path_to_buffer_allocator/libneuron_buffer_allocator.so export NEURON_USDK_ADAPTER_LIB=path_to_usdk_adapter/libneuronusdk_adapter.mtk.so diff --git a/examples/demo-apps/android/LlamaDemo/docs/delegates/qualcomm_README.md b/examples/demo-apps/android/LlamaDemo/docs/delegates/qualcomm_README.md index 68aed7000c8..c1a93d02f93 100644 --- a/examples/demo-apps/android/LlamaDemo/docs/delegates/qualcomm_README.md +++ b/examples/demo-apps/android/LlamaDemo/docs/delegates/qualcomm_README.md @@ -212,6 +212,8 @@ sh examples/demo-apps/android/LlamaDemo/setup-with-qnn.sh This is running the shell script which configures the required core ExecuTorch, Llama2/3, and Android libraries, builds them into AAR, and copies it to the app. Note: If you are building the Android app mentioned in the next section on a separate machine (i.e. MacOS but building and exporting for QNN backend on Linux), make sure you copy the aar file generated from setup-with-qnn script to "examples/demo-apps/android/LlamaDemo/app/libs" before building the Android app. +6. Set up the correct QNN version in gradle rule +Currently, the gralde rule searches for the property `qnnVersion`. When this variable is defined, it will add QNN runtime library to the dependency. To use it, append the string `qnnVersion=` (ex. `qnnVersion=2.37.0`) to the end of the `gradle.properties` file. ## Run the Android Demo App diff --git a/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.jar b/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.jar index e708b1c023e..a4b76b9530d 100644 Binary files a/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.jar and b/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.jar differ diff --git a/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.properties b/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.properties index 2a7f77d2fbb..9355b415575 100644 --- a/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.properties +++ b/examples/demo-apps/android/LlamaDemo/gradle/wrapper/gradle-wrapper.properties @@ -1,6 +1,7 @@ -#Mon Sep 25 11:23:11 PDT 2023 distributionBase=GRADLE_USER_HOME distributionPath=wrapper/dists -distributionUrl=https\://services.gradle.org/distributions/gradle-8.0-bin.zip +distributionUrl=https\://services.gradle.org/distributions/gradle-8.10-bin.zip +networkTimeout=10000 +validateDistributionUrl=true zipStoreBase=GRADLE_USER_HOME zipStorePath=wrapper/dists diff --git a/examples/demo-apps/android/LlamaDemo/gradlew b/examples/demo-apps/android/LlamaDemo/gradlew index 4f906e0c811..f5feea6d6b1 100755 --- a/examples/demo-apps/android/LlamaDemo/gradlew +++ b/examples/demo-apps/android/LlamaDemo/gradlew @@ -1,7 +1,7 @@ -#!/usr/bin/env sh +#!/bin/sh # -# Copyright 2015 the original author or authors. +# Copyright © 2015-2021 the original authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,69 +15,104 @@ # See the License for the specific language governing permissions and # limitations under the License. # +# SPDX-License-Identifier: Apache-2.0 +# ############################################################################## -## -## Gradle start up script for UN*X -## +# +# Gradle start up script for POSIX generated by Gradle. +# +# Important for running: +# +# (1) You need a POSIX-compliant shell to run this script. If your /bin/sh is +# noncompliant, but you have some other compliant shell such as ksh or +# bash, then to run this script, type that shell name before the whole +# command line, like: +# +# ksh Gradle +# +# Busybox and similar reduced shells will NOT work, because this script +# requires all of these POSIX shell features: +# * functions; +# * expansions «$var», «${var}», «${var:-default}», «${var+SET}», +# «${var#prefix}», «${var%suffix}», and «$( cmd )»; +# * compound commands having a testable exit status, especially «case»; +# * various built-in commands including «command», «set», and «ulimit». +# +# Important for patching: +# +# (2) This script targets any POSIX shell, so it avoids extensions provided +# by Bash, Ksh, etc; in particular arrays are avoided. +# +# The "traditional" practice of packing multiple parameters into a +# space-separated string is a well documented source of bugs and security +# problems, so this is (mostly) avoided, by progressively accumulating +# options in "$@", and eventually passing that to Java. +# +# Where the inherited environment variables (DEFAULT_JVM_OPTS, JAVA_OPTS, +# and GRADLE_OPTS) rely on word-splitting, this is performed explicitly; +# see the in-line comments for details. +# +# There are tweaks for specific operating systems such as AIX, CygWin, +# Darwin, MinGW, and NonStop. +# +# (3) This script is generated from the Groovy template +# https://github.com/gradle/gradle/blob/HEAD/platforms/jvm/plugins-application/src/main/resources/org/gradle/api/internal/plugins/unixStartScript.txt +# within the Gradle project. +# +# You can find Gradle at https://github.com/gradle/gradle/. +# ############################################################################## # Attempt to set APP_HOME + # Resolve links: $0 may be a link -PRG="$0" -# Need this for relative symlinks. -while [ -h "$PRG" ] ; do - ls=`ls -ld "$PRG"` - link=`expr "$ls" : '.*-> \(.*\)$'` - if expr "$link" : '/.*' > /dev/null; then - PRG="$link" - else - PRG=`dirname "$PRG"`"/$link" - fi +app_path=$0 + +# Need this for daisy-chained symlinks. +while + APP_HOME=${app_path%"${app_path##*/}"} # leaves a trailing /; empty if no leading path + [ -h "$app_path" ] +do + ls=$( ls -ld "$app_path" ) + link=${ls#*' -> '} + case $link in #( + /*) app_path=$link ;; #( + *) app_path=$APP_HOME$link ;; + esac done -SAVED="`pwd`" -cd "`dirname \"$PRG\"`/" >/dev/null -APP_HOME="`pwd -P`" -cd "$SAVED" >/dev/null -APP_NAME="Gradle" -APP_BASE_NAME=`basename "$0"` - -# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. -DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' +# This is normally unused +# shellcheck disable=SC2034 +APP_BASE_NAME=${0##*/} +# Discard cd standard output in case $CDPATH is set (https://github.com/gradle/gradle/issues/25036) +APP_HOME=$( cd -P "${APP_HOME:-./}" > /dev/null && printf '%s +' "$PWD" ) || exit # Use the maximum available, or set MAX_FD != -1 to use that value. -MAX_FD="maximum" +MAX_FD=maximum warn () { echo "$*" -} +} >&2 die () { echo echo "$*" echo exit 1 -} +} >&2 # OS specific support (must be 'true' or 'false'). cygwin=false msys=false darwin=false nonstop=false -case "`uname`" in - CYGWIN* ) - cygwin=true - ;; - Darwin* ) - darwin=true - ;; - MINGW* ) - msys=true - ;; - NONSTOP* ) - nonstop=true - ;; +case "$( uname )" in #( + CYGWIN* ) cygwin=true ;; #( + Darwin* ) darwin=true ;; #( + MSYS* | MINGW* ) msys=true ;; #( + NONSTOP* ) nonstop=true ;; esac CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar @@ -87,9 +122,9 @@ CLASSPATH=$APP_HOME/gradle/wrapper/gradle-wrapper.jar if [ -n "$JAVA_HOME" ] ; then if [ -x "$JAVA_HOME/jre/sh/java" ] ; then # IBM's JDK on AIX uses strange locations for the executables - JAVACMD="$JAVA_HOME/jre/sh/java" + JAVACMD=$JAVA_HOME/jre/sh/java else - JAVACMD="$JAVA_HOME/bin/java" + JAVACMD=$JAVA_HOME/bin/java fi if [ ! -x "$JAVACMD" ] ; then die "ERROR: JAVA_HOME is set to an invalid directory: $JAVA_HOME @@ -98,88 +133,120 @@ Please set the JAVA_HOME variable in your environment to match the location of your Java installation." fi else - JAVACMD="java" - which java >/dev/null 2>&1 || die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. + JAVACMD=java + if ! command -v java >/dev/null 2>&1 + then + die "ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. Please set the JAVA_HOME variable in your environment to match the location of your Java installation." + fi fi # Increase the maximum file descriptors if we can. -if [ "$cygwin" = "false" -a "$darwin" = "false" -a "$nonstop" = "false" ] ; then - MAX_FD_LIMIT=`ulimit -H -n` - if [ $? -eq 0 ] ; then - if [ "$MAX_FD" = "maximum" -o "$MAX_FD" = "max" ] ; then - MAX_FD="$MAX_FD_LIMIT" - fi - ulimit -n $MAX_FD - if [ $? -ne 0 ] ; then - warn "Could not set maximum file descriptor limit: $MAX_FD" - fi - else - warn "Could not query maximum file descriptor limit: $MAX_FD_LIMIT" - fi +if ! "$cygwin" && ! "$darwin" && ! "$nonstop" ; then + case $MAX_FD in #( + max*) + # In POSIX sh, ulimit -H is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + MAX_FD=$( ulimit -H -n ) || + warn "Could not query maximum file descriptor limit" + esac + case $MAX_FD in #( + '' | soft) :;; #( + *) + # In POSIX sh, ulimit -n is undefined. That's why the result is checked to see if it worked. + # shellcheck disable=SC2039,SC3045 + ulimit -n "$MAX_FD" || + warn "Could not set maximum file descriptor limit to $MAX_FD" + esac fi -# For Darwin, add options to specify how the application appears in the dock -if $darwin; then - GRADLE_OPTS="$GRADLE_OPTS \"-Xdock:name=$APP_NAME\" \"-Xdock:icon=$APP_HOME/media/gradle.icns\"" -fi +# Collect all arguments for the java command, stacking in reverse order: +# * args from the command line +# * the main class name +# * -classpath +# * -D...appname settings +# * --module-path (only if needed) +# * DEFAULT_JVM_OPTS, JAVA_OPTS, and GRADLE_OPTS environment variables. # For Cygwin or MSYS, switch paths to Windows format before running java -if [ "$cygwin" = "true" -o "$msys" = "true" ] ; then - APP_HOME=`cygpath --path --mixed "$APP_HOME"` - CLASSPATH=`cygpath --path --mixed "$CLASSPATH"` - - JAVACMD=`cygpath --unix "$JAVACMD"` - - # We build the pattern for arguments to be converted via cygpath - ROOTDIRSRAW=`find -L / -maxdepth 1 -mindepth 1 -type d 2>/dev/null` - SEP="" - for dir in $ROOTDIRSRAW ; do - ROOTDIRS="$ROOTDIRS$SEP$dir" - SEP="|" - done - OURCYGPATTERN="(^($ROOTDIRS))" - # Add a user-defined pattern to the cygpath arguments - if [ "$GRADLE_CYGPATTERN" != "" ] ; then - OURCYGPATTERN="$OURCYGPATTERN|($GRADLE_CYGPATTERN)" - fi +if "$cygwin" || "$msys" ; then + APP_HOME=$( cygpath --path --mixed "$APP_HOME" ) + CLASSPATH=$( cygpath --path --mixed "$CLASSPATH" ) + + JAVACMD=$( cygpath --unix "$JAVACMD" ) + # Now convert the arguments - kludge to limit ourselves to /bin/sh - i=0 - for arg in "$@" ; do - CHECK=`echo "$arg"|egrep -c "$OURCYGPATTERN" -` - CHECK2=`echo "$arg"|egrep -c "^-"` ### Determine if an option - - if [ $CHECK -ne 0 ] && [ $CHECK2 -eq 0 ] ; then ### Added a condition - eval `echo args$i`=`cygpath --path --ignore --mixed "$arg"` - else - eval `echo args$i`="\"$arg\"" + for arg do + if + case $arg in #( + -*) false ;; # don't mess with options #( + /?*) t=${arg#/} t=/${t%%/*} # looks like a POSIX filepath + [ -e "$t" ] ;; #( + *) false ;; + esac + then + arg=$( cygpath --path --ignore --mixed "$arg" ) fi - i=`expr $i + 1` + # Roll the args list around exactly as many times as the number of + # args, so each arg winds up back in the position where it started, but + # possibly modified. + # + # NB: a `for` loop captures its iteration list before it begins, so + # changing the positional parameters here affects neither the number of + # iterations, nor the values presented in `arg`. + shift # remove old arg + set -- "$@" "$arg" # push replacement arg done - case $i in - 0) set -- ;; - 1) set -- "$args0" ;; - 2) set -- "$args0" "$args1" ;; - 3) set -- "$args0" "$args1" "$args2" ;; - 4) set -- "$args0" "$args1" "$args2" "$args3" ;; - 5) set -- "$args0" "$args1" "$args2" "$args3" "$args4" ;; - 6) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" ;; - 7) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" ;; - 8) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" ;; - 9) set -- "$args0" "$args1" "$args2" "$args3" "$args4" "$args5" "$args6" "$args7" "$args8" ;; - esac fi -# Escape application args -save () { - for i do printf %s\\n "$i" | sed "s/'/'\\\\''/g;1s/^/'/;\$s/\$/' \\\\/" ; done - echo " " -} -APP_ARGS=`save "$@"` -# Collect all arguments for the java command, following the shell quoting and substitution rules -eval set -- $DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS "\"-Dorg.gradle.appname=$APP_BASE_NAME\"" -classpath "\"$CLASSPATH\"" org.gradle.wrapper.GradleWrapperMain "$APP_ARGS" +# Add default JVM options here. You can also use JAVA_OPTS and GRADLE_OPTS to pass JVM options to this script. +DEFAULT_JVM_OPTS='"-Xmx64m" "-Xms64m"' + +# Collect all arguments for the java command: +# * DEFAULT_JVM_OPTS, JAVA_OPTS, JAVA_OPTS, and optsEnvironmentVar are not allowed to contain shell fragments, +# and any embedded shellness will be escaped. +# * For example: A user cannot expect ${Hostname} to be expanded, as it is an environment variable and will be +# treated as '${Hostname}' itself on the command line. + +set -- \ + "-Dorg.gradle.appname=$APP_BASE_NAME" \ + -classpath "$CLASSPATH" \ + org.gradle.wrapper.GradleWrapperMain \ + "$@" + +# Stop when "xargs" is not available. +if ! command -v xargs >/dev/null 2>&1 +then + die "xargs is not available" +fi + +# Use "xargs" to parse quoted args. +# +# With -n1 it outputs one arg per line, with the quotes and backslashes removed. +# +# In Bash we could simply go: +# +# readarray ARGS < <( xargs -n1 <<<"$var" ) && +# set -- "${ARGS[@]}" "$@" +# +# but POSIX shell has neither arrays nor command substitution, so instead we +# post-process each arg (as a line of input to sed) to backslash-escape any +# character that might be a shell metacharacter, then use eval to reverse +# that process (while maintaining the separation between arguments), and wrap +# the whole thing up as a single "set" statement. +# +# This will of course break if any of these variables contains a newline or +# an unmatched quote. +# + +eval "set -- $( + printf '%s\n' "$DEFAULT_JVM_OPTS $JAVA_OPTS $GRADLE_OPTS" | + xargs -n1 | + sed ' s~[^-[:alnum:]+,./:=@_]~\\&~g; ' | + tr '\n' ' ' + )" '"$@"' exec "$JAVACMD" "$@" diff --git a/examples/demo-apps/android/LlamaDemo/gradlew.bat b/examples/demo-apps/android/LlamaDemo/gradlew.bat index b4fb785a6d0..9b42019c791 100644 --- a/examples/demo-apps/android/LlamaDemo/gradlew.bat +++ b/examples/demo-apps/android/LlamaDemo/gradlew.bat @@ -1,9 +1,3 @@ -@REM Copyright (c) Meta Platforms, Inc. and affiliates. -@REM All rights reserved. -@REM -@REM This source code is licensed under the BSD-style license found in the -@REM LICENSE file in the root directory of this source tree. - @rem @rem Copyright 2015 the original author or authors. @rem @@ -19,8 +13,10 @@ @rem See the License for the specific language governing permissions and @rem limitations under the License. @rem +@rem SPDX-License-Identifier: Apache-2.0 +@rem -@if "%DEBUG%" == "" @echo off +@if "%DEBUG%"=="" @echo off @rem ########################################################################## @rem @rem Gradle startup script for Windows @@ -31,7 +27,8 @@ if "%OS%"=="Windows_NT" setlocal set DIRNAME=%~dp0 -if "%DIRNAME%" == "" set DIRNAME=. +if "%DIRNAME%"=="" set DIRNAME=. +@rem This is normally unused set APP_BASE_NAME=%~n0 set APP_HOME=%DIRNAME% @@ -46,13 +43,13 @@ if defined JAVA_HOME goto findJavaFromJavaHome set JAVA_EXE=java.exe %JAVA_EXE% -version >NUL 2>&1 -if "%ERRORLEVEL%" == "0" goto execute +if %ERRORLEVEL% equ 0 goto execute -echo. -echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. -echo. -echo Please set the JAVA_HOME variable in your environment to match the -echo location of your Java installation. +echo. 1>&2 +echo ERROR: JAVA_HOME is not set and no 'java' command could be found in your PATH. 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 goto fail @@ -62,11 +59,11 @@ set JAVA_EXE=%JAVA_HOME%/bin/java.exe if exist "%JAVA_EXE%" goto execute -echo. -echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% -echo. -echo Please set the JAVA_HOME variable in your environment to match the -echo location of your Java installation. +echo. 1>&2 +echo ERROR: JAVA_HOME is set to an invalid directory: %JAVA_HOME% 1>&2 +echo. 1>&2 +echo Please set the JAVA_HOME variable in your environment to match the 1>&2 +echo location of your Java installation. 1>&2 goto fail @@ -81,13 +78,15 @@ set CLASSPATH=%APP_HOME%\gradle\wrapper\gradle-wrapper.jar :end @rem End local scope for the variables with windows NT shell -if "%ERRORLEVEL%"=="0" goto mainEnd +if %ERRORLEVEL% equ 0 goto mainEnd :fail rem Set variable GRADLE_EXIT_CONSOLE if you need the _script_ return code instead of rem the _cmd.exe /c_ return code! -if not "" == "%GRADLE_EXIT_CONSOLE%" exit 1 -exit /b 1 +set EXIT_CODE=%ERRORLEVEL% +if %EXIT_CODE% equ 0 set EXIT_CODE=1 +if not ""=="%GRADLE_EXIT_CONSOLE%" exit %EXIT_CODE% +exit /b %EXIT_CODE% :mainEnd if "%OS%"=="Windows_NT" endlocal diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj b/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj deleted file mode 100644 index ce11b077270..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj +++ /dev/null @@ -1,1026 +0,0 @@ -// !$*UTF8*$! -{ - archiveVersion = 1; - classes = { - }; - objectVersion = 56; - objects = { - -/* Begin PBXBuildFile section */ - 0324D68B2BAACB6900DEF36F /* App.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6802BAACB6900DEF36F /* App.swift */; }; - 0324D68C2BAACB6900DEF36F /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6812BAACB6900DEF36F /* ContentView.swift */; }; - 0324D68D2BAACB6900DEF36F /* LogManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6822BAACB6900DEF36F /* LogManager.swift */; }; - 0324D68E2BAACB6900DEF36F /* LogView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6832BAACB6900DEF36F /* LogView.swift */; }; - 0324D68F2BAACB6900DEF36F /* Message.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6842BAACB6900DEF36F /* Message.swift */; }; - 0324D6902BAACB6900DEF36F /* MessageListView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6852BAACB6900DEF36F /* MessageListView.swift */; }; - 0324D6912BAACB6900DEF36F /* MessageView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6862BAACB6900DEF36F /* MessageView.swift */; }; - 0324D6922BAACB6900DEF36F /* ResourceManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6872BAACB6900DEF36F /* ResourceManager.swift */; }; - 0324D6932BAACB6900DEF36F /* ResourceMonitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0324D6882BAACB6900DEF36F /* ResourceMonitor.swift */; }; - 0324D6962BAACB7000DEF36F /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 0324D6942BAACB7000DEF36F /* Assets.xcassets */; }; - 03729EDC2BB1F8DE00152F2E /* LLaMARunner.framework in Embed Frameworks */ = {isa = PBXBuildFile; fileRef = 03729ED52BB1F8DE00152F2E /* LLaMARunner.framework */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; }; - 03729EE12BB1F93800152F2E /* LLaMARunner.mm in Sources */ = {isa = PBXBuildFile; fileRef = 0324D69A2BAACB7C00DEF36F /* LLaMARunner.mm */; }; - 03729EE22BB1F93E00152F2E /* LLaMARunner.h in Headers */ = {isa = PBXBuildFile; fileRef = 0324D6992BAACB7C00DEF36F /* LLaMARunner.h */; settings = {ATTRIBUTES = (Public, ); }; }; - 03729F0A2BB203B300152F2E /* runner.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03729F072BB203B300152F2E /* runner.cpp */; }; - 03729F132BB2042B00152F2E /* sampler.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03729F112BB2042B00152F2E /* sampler.cpp */; }; - 0372C3152C89418E00CD942A /* llava_runner.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 0372C3132C89418E00CD942A /* llava_runner.cpp */; }; - 03CF43962CEC5CEC00C7113B /* backend_coreml in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43952CEC5CEC00C7113B /* backend_coreml */; }; - 03CF43982CEC5CEC00C7113B /* backend_coreml_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43972CEC5CEC00C7113B /* backend_coreml_debug */; }; - 03CF439A2CEC5CEC00C7113B /* backend_mps in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43992CEC5CEC00C7113B /* backend_mps */; }; - 03CF439C2CEC5CEC00C7113B /* backend_mps_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF439B2CEC5CEC00C7113B /* backend_mps_debug */; }; - 03CF439E2CEC5CEC00C7113B /* backend_xnnpack in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF439D2CEC5CEC00C7113B /* backend_xnnpack */; }; - 03CF43A02CEC5CEC00C7113B /* backend_xnnpack_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF439F2CEC5CEC00C7113B /* backend_xnnpack_debug */; }; - 03CF43A22CEC5CEC00C7113B /* executorch_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43A12CEC5CEC00C7113B /* executorch_debug */; }; - 03CF43A42CEC5CEC00C7113B /* kernels_llm in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43A32CEC5CEC00C7113B /* kernels_llm */; }; - 03CF43A62CEC5CEC00C7113B /* kernels_llm_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43A52CEC5CEC00C7113B /* kernels_llm_debug */; }; - 03CF43A82CEC5CEC00C7113B /* kernels_optimized in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43A72CEC5CEC00C7113B /* kernels_optimized */; }; - 03CF43AA2CEC5CEC00C7113B /* kernels_optimized_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43A92CEC5CEC00C7113B /* kernels_optimized_debug */; }; - 03CF43B02CEC5CEC00C7113B /* kernels_quantized in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43AF2CEC5CEC00C7113B /* kernels_quantized */; }; - 03CF43B22CEC5CEC00C7113B /* kernels_quantized_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03CF43B12CEC5CEC00C7113B /* kernels_quantized_debug */; }; - 03D151B82E0E0908007A38BE /* LLaVARunner.mm in Sources */ = {isa = PBXBuildFile; fileRef = 03D151B72E0E0908007A38BE /* LLaVARunner.mm */; }; - 03D151B92E0E0908007A38BE /* LLaVARunner.h in Headers */ = {isa = PBXBuildFile; fileRef = 03D151B62E0E0908007A38BE /* LLaVARunner.h */; settings = {ATTRIBUTES = (Public, ); }; }; - 03D151CA2E0E98C4007A38BE /* sentencepiece.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03D151C92E0E98C4007A38BE /* sentencepiece.cpp */; }; - 03D151CB2E0E98C4007A38BE /* regex_lookahead.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03D151C82E0E98C4007A38BE /* regex_lookahead.cpp */; }; - 03D151CF2E0E9ACB007A38BE /* text_prefiller.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03D151CE2E0E9ACB007A38BE /* text_prefiller.cpp */; }; - 03D151D02E0E9ACB007A38BE /* text_llm_runner.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03D151CD2E0E9ACB007A38BE /* text_llm_runner.cpp */; }; - 03D151D12E0E9ACB007A38BE /* text_decoder_runner.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 03D151CC2E0E9ACB007A38BE /* text_decoder_runner.cpp */; }; - 03D151D92E0E9E43007A38BE /* ExecuTorchTextLLMRunner.mm in Sources */ = {isa = PBXBuildFile; fileRef = 03D151D42E0E9E43007A38BE /* ExecuTorchTextLLMRunner.mm */; }; - 03D151DB2E0E9E43007A38BE /* ExecuTorchTextLLMRunner.h in Headers */ = {isa = PBXBuildFile; fileRef = 03D151D32E0E9E43007A38BE /* ExecuTorchTextLLMRunner.h */; }; - 03D151DC2E0E9E43007A38BE /* ExecuTorchLLM.h in Headers */ = {isa = PBXBuildFile; fileRef = 03D151D22E0E9E43007A38BE /* ExecuTorchLLM.h */; }; - 26A6A4282C8A3769005A761E /* ImagePicker.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26A6A4272C8A3769005A761E /* ImagePicker.swift */; }; - 306A713D2DC1DC0F00936B1F /* token_decoder.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A713C2DC1DC0F00936B1F /* token_decoder.h */; }; - 306A713E2DC1DC0F00936B1F /* regex.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A71392DC1DC0F00936B1F /* regex.h */; }; - 306A713F2DC1DC0F00936B1F /* string_integer_map.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A713B2DC1DC0F00936B1F /* string_integer_map.h */; }; - 306A71402DC1DC0F00936B1F /* re2_regex.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A71382DC1DC0F00936B1F /* re2_regex.h */; }; - 306A71412DC1DC0F00936B1F /* hf_tokenizer.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A71352DC1DC0F00936B1F /* hf_tokenizer.h */; }; - 306A71422DC1DC0F00936B1F /* pre_tokenizer.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A71372DC1DC0F00936B1F /* pre_tokenizer.h */; }; - 306A71432DC1DC0F00936B1F /* pcre2_regex.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A71362DC1DC0F00936B1F /* pcre2_regex.h */; }; - 306A71442DC1DC0F00936B1F /* std_regex.h in Headers */ = {isa = PBXBuildFile; fileRef = 306A713A2DC1DC0F00936B1F /* std_regex.h */; }; - 306A714C2DC1DC3D00936B1F /* std_regex.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A714A2DC1DC3D00936B1F /* std_regex.cpp */; }; - 306A714D2DC1DC3D00936B1F /* re2_regex.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A71482DC1DC3D00936B1F /* re2_regex.cpp */; }; - 306A714E2DC1DC3D00936B1F /* pcre2_regex.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A71462DC1DC3D00936B1F /* pcre2_regex.cpp */; }; - 306A714F2DC1DC3D00936B1F /* hf_tokenizer.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A71452DC1DC3D00936B1F /* hf_tokenizer.cpp */; }; - 306A71502DC1DC3D00936B1F /* regex.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A71492DC1DC3D00936B1F /* regex.cpp */; }; - 306A71512DC1DC3D00936B1F /* pre_tokenizer.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A71472DC1DC3D00936B1F /* pre_tokenizer.cpp */; }; - 306A71522DC1DC3D00936B1F /* token_decoder.cpp in Sources */ = {isa = PBXBuildFile; fileRef = 306A714B2DC1DC3D00936B1F /* token_decoder.cpp */; }; - 3072D5232DC3EA280083FC83 /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = 3072D5222DC3EA280083FC83 /* Constants.swift */; }; - F24909E82E207004001E5B69 /* normalizer.cpp in Sources */ = {isa = PBXBuildFile; fileRef = F24909E72E207004001E5B69 /* normalizer.cpp */; }; - F292B0752D88B0C200BE6839 /* tiktoken.cpp in Sources */ = {isa = PBXBuildFile; fileRef = F292B06F2D88B0C200BE6839 /* tiktoken.cpp */; }; - F292B0762D88B0C200BE6839 /* llama2c_tokenizer.cpp in Sources */ = {isa = PBXBuildFile; fileRef = F292B06C2D88B0C200BE6839 /* llama2c_tokenizer.cpp */; }; - F292B0772D88B0C200BE6839 /* bpe_tokenizer_base.cpp in Sources */ = {isa = PBXBuildFile; fileRef = F292B06A2D88B0C200BE6839 /* bpe_tokenizer_base.cpp */; }; - F292B0882D88B0D200BE6839 /* llama2c_tokenizer.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B07C2D88B0D200BE6839 /* llama2c_tokenizer.h */; }; - F292B0892D88B0D200BE6839 /* tokenizer.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B0832D88B0D200BE6839 /* tokenizer.h */; }; - F292B08B2D88B0D200BE6839 /* result.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B07F2D88B0D200BE6839 /* result.h */; }; - F292B08D2D88B0D200BE6839 /* error.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B07A2D88B0D200BE6839 /* error.h */; }; - F292B08E2D88B0D200BE6839 /* bpe_tokenizer_base.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B0792D88B0D200BE6839 /* bpe_tokenizer_base.h */; }; - F292B08F2D88B0D200BE6839 /* log.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B07D2D88B0D200BE6839 /* log.h */; }; - F292B0912D88B0D200BE6839 /* tiktoken.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B0812D88B0D200BE6839 /* tiktoken.h */; }; - F292B0922D88B0D200BE6839 /* base64.h in Headers */ = {isa = PBXBuildFile; fileRef = F292B0782D88B0D200BE6839 /* base64.h */; }; - F292B1022D88B20C00BE6839 /* llama_tiktoken.cpp in Sources */ = {isa = PBXBuildFile; fileRef = F292B1002D88B20C00BE6839 /* llama_tiktoken.cpp */; }; -/* End PBXBuildFile section */ - -/* Begin PBXContainerItemProxy section */ - 03729ED92BB1F8DE00152F2E /* PBXContainerItemProxy */ = { - isa = PBXContainerItemProxy; - containerPortal = 032C01672AC228E5002955E1 /* Project object */; - proxyType = 1; - remoteGlobalIDString = 03729ED42BB1F8DE00152F2E; - remoteInfo = LLaMARunner; - }; -/* End PBXContainerItemProxy section */ - -/* Begin PBXCopyFilesBuildPhase section */ - 03729EE02BB1F8DE00152F2E /* Embed Frameworks */ = { - isa = PBXCopyFilesBuildPhase; - buildActionMask = 2147483647; - dstPath = ""; - dstSubfolderSpec = 10; - files = ( - 03729EDC2BB1F8DE00152F2E /* LLaMARunner.framework in Embed Frameworks */, - ); - name = "Embed Frameworks"; - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXCopyFilesBuildPhase section */ - -/* Begin PBXFileReference section */ - 0320439D2BB4AC6600050211 /* LLaMA-Info.plist */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.plist.xml; path = "LLaMA-Info.plist"; sourceTree = ""; }; - 0324D6802BAACB6900DEF36F /* App.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = App.swift; sourceTree = ""; }; - 0324D6812BAACB6900DEF36F /* ContentView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = ""; }; - 0324D6822BAACB6900DEF36F /* LogManager.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LogManager.swift; sourceTree = ""; }; - 0324D6832BAACB6900DEF36F /* LogView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = LogView.swift; sourceTree = ""; }; - 0324D6842BAACB6900DEF36F /* Message.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Message.swift; sourceTree = ""; }; - 0324D6852BAACB6900DEF36F /* MessageListView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MessageListView.swift; sourceTree = ""; }; - 0324D6862BAACB6900DEF36F /* MessageView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MessageView.swift; sourceTree = ""; }; - 0324D6872BAACB6900DEF36F /* ResourceManager.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ResourceManager.swift; sourceTree = ""; }; - 0324D6882BAACB6900DEF36F /* ResourceMonitor.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ResourceMonitor.swift; sourceTree = ""; }; - 0324D6942BAACB7000DEF36F /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; - 0324D6992BAACB7C00DEF36F /* LLaMARunner.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = LLaMARunner.h; sourceTree = ""; }; - 0324D69A2BAACB7C00DEF36F /* LLaMARunner.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = LLaMARunner.mm; sourceTree = ""; }; - 035A5E942BB4B523001E0553 /* LLaMA.entitlements */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.plist.entitlements; path = LLaMA.entitlements; sourceTree = ""; }; - 036CAF9D2BB1444500D6C2D5 /* LLaMA.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = LLaMA.app; sourceTree = BUILT_PRODUCTS_DIR; }; - 03729ED52BB1F8DE00152F2E /* LLaMARunner.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = LLaMARunner.framework; sourceTree = BUILT_PRODUCTS_DIR; }; - 03729F072BB203B300152F2E /* runner.cpp */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; name = runner.cpp; path = ../../../examples/models/llama/runner/runner.cpp; sourceTree = ""; }; - 03729F112BB2042B00152F2E /* sampler.cpp */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; path = sampler.cpp; sourceTree = ""; }; - 0372C3132C89418E00CD942A /* llava_runner.cpp */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; name = llava_runner.cpp; path = ../../../examples/models/llava/runner/llava_runner.cpp; sourceTree = ""; }; - 03C5F51C2CE7D35C00D6CE3F /* Release.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = Release.xcconfig; sourceTree = ""; }; - 03C5F51D2CE7D37100D6CE3F /* Debug.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = Debug.xcconfig; sourceTree = ""; }; - 03D151B62E0E0908007A38BE /* LLaVARunner.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = LLaVARunner.h; sourceTree = ""; }; - 03D151B72E0E0908007A38BE /* LLaVARunner.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = LLaVARunner.mm; sourceTree = ""; }; - 03D151C82E0E98C4007A38BE /* regex_lookahead.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = regex_lookahead.cpp; path = src/regex_lookahead.cpp; sourceTree = ""; }; - 03D151C92E0E98C4007A38BE /* sentencepiece.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = sentencepiece.cpp; path = src/sentencepiece.cpp; sourceTree = ""; }; - 03D151CC2E0E9ACB007A38BE /* text_decoder_runner.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = text_decoder_runner.cpp; sourceTree = ""; }; - 03D151CD2E0E9ACB007A38BE /* text_llm_runner.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = text_llm_runner.cpp; sourceTree = ""; }; - 03D151CE2E0E9ACB007A38BE /* text_prefiller.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = text_prefiller.cpp; sourceTree = ""; }; - 03D151D22E0E9E43007A38BE /* ExecuTorchLLM.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = ExecuTorchLLM.h; sourceTree = ""; }; - 03D151D32E0E9E43007A38BE /* ExecuTorchTextLLMRunner.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = ExecuTorchTextLLMRunner.h; sourceTree = ""; }; - 03D151D42E0E9E43007A38BE /* ExecuTorchTextLLMRunner.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = ExecuTorchTextLLMRunner.mm; sourceTree = ""; }; - 26A6A4272C8A3769005A761E /* ImagePicker.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImagePicker.swift; sourceTree = ""; }; - 306A71352DC1DC0F00936B1F /* hf_tokenizer.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = hf_tokenizer.h; sourceTree = ""; }; - 306A71362DC1DC0F00936B1F /* pcre2_regex.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = pcre2_regex.h; sourceTree = ""; }; - 306A71372DC1DC0F00936B1F /* pre_tokenizer.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = pre_tokenizer.h; sourceTree = ""; }; - 306A71382DC1DC0F00936B1F /* re2_regex.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = re2_regex.h; sourceTree = ""; }; - 306A71392DC1DC0F00936B1F /* regex.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = regex.h; sourceTree = ""; }; - 306A713A2DC1DC0F00936B1F /* std_regex.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = std_regex.h; sourceTree = ""; }; - 306A713B2DC1DC0F00936B1F /* string_integer_map.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = string_integer_map.h; sourceTree = ""; }; - 306A713C2DC1DC0F00936B1F /* token_decoder.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = token_decoder.h; sourceTree = ""; }; - 306A71452DC1DC3D00936B1F /* hf_tokenizer.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = hf_tokenizer.cpp; path = src/hf_tokenizer.cpp; sourceTree = ""; }; - 306A71462DC1DC3D00936B1F /* pcre2_regex.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = pcre2_regex.cpp; path = src/pcre2_regex.cpp; sourceTree = ""; }; - 306A71472DC1DC3D00936B1F /* pre_tokenizer.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = pre_tokenizer.cpp; path = src/pre_tokenizer.cpp; sourceTree = ""; }; - 306A71482DC1DC3D00936B1F /* re2_regex.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = re2_regex.cpp; path = src/re2_regex.cpp; sourceTree = ""; }; - 306A71492DC1DC3D00936B1F /* regex.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = regex.cpp; path = src/regex.cpp; sourceTree = ""; }; - 306A714A2DC1DC3D00936B1F /* std_regex.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = std_regex.cpp; path = src/std_regex.cpp; sourceTree = ""; }; - 306A714B2DC1DC3D00936B1F /* token_decoder.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = token_decoder.cpp; path = src/token_decoder.cpp; sourceTree = ""; }; - 3072D5222DC3EA280083FC83 /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = ""; }; - F24909E72E207004001E5B69 /* normalizer.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = normalizer.cpp; path = src/normalizer.cpp; sourceTree = ""; }; - F292B06A2D88B0C200BE6839 /* bpe_tokenizer_base.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = bpe_tokenizer_base.cpp; path = src/bpe_tokenizer_base.cpp; sourceTree = ""; }; - F292B06C2D88B0C200BE6839 /* llama2c_tokenizer.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = llama2c_tokenizer.cpp; path = src/llama2c_tokenizer.cpp; sourceTree = ""; }; - F292B06F2D88B0C200BE6839 /* tiktoken.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; name = tiktoken.cpp; path = src/tiktoken.cpp; sourceTree = ""; }; - F292B0782D88B0D200BE6839 /* base64.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = base64.h; sourceTree = ""; }; - F292B0792D88B0D200BE6839 /* bpe_tokenizer_base.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = bpe_tokenizer_base.h; sourceTree = ""; }; - F292B07A2D88B0D200BE6839 /* error.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = error.h; sourceTree = ""; }; - F292B07C2D88B0D200BE6839 /* llama2c_tokenizer.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = llama2c_tokenizer.h; sourceTree = ""; }; - F292B07D2D88B0D200BE6839 /* log.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = log.h; sourceTree = ""; }; - F292B07F2D88B0D200BE6839 /* result.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = result.h; sourceTree = ""; }; - F292B0812D88B0D200BE6839 /* tiktoken.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = tiktoken.h; sourceTree = ""; }; - F292B0832D88B0D200BE6839 /* tokenizer.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = tokenizer.h; sourceTree = ""; }; - F292B1002D88B20C00BE6839 /* llama_tiktoken.cpp */ = {isa = PBXFileReference; explicitFileType = sourcecode.cpp.cpp; name = llama_tiktoken.cpp; path = ../../../examples/models/llama/tokenizer/llama_tiktoken.cpp; sourceTree = ""; }; -/* End PBXFileReference section */ - -/* Begin PBXFrameworksBuildPhase section */ - 032C016C2AC228E6002955E1 /* Frameworks */ = { - isa = PBXFrameworksBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - runOnlyForDeploymentPostprocessing = 0; - }; - 03729ED22BB1F8DE00152F2E /* Frameworks */ = { - isa = PBXFrameworksBuildPhase; - buildActionMask = 2147483647; - files = ( - 03CF43962CEC5CEC00C7113B /* backend_coreml in Frameworks */, - 03CF43982CEC5CEC00C7113B /* backend_coreml_debug in Frameworks */, - 03CF439A2CEC5CEC00C7113B /* backend_mps in Frameworks */, - 03CF439C2CEC5CEC00C7113B /* backend_mps_debug in Frameworks */, - 03CF439E2CEC5CEC00C7113B /* backend_xnnpack in Frameworks */, - 03CF43A02CEC5CEC00C7113B /* backend_xnnpack_debug in Frameworks */, - 03CF43A22CEC5CEC00C7113B /* executorch_debug in Frameworks */, - 03CF43A42CEC5CEC00C7113B /* kernels_llm in Frameworks */, - 03CF43A62CEC5CEC00C7113B /* kernels_llm_debug in Frameworks */, - 03CF43A82CEC5CEC00C7113B /* kernels_optimized in Frameworks */, - 03CF43AA2CEC5CEC00C7113B /* kernels_optimized_debug in Frameworks */, - 03CF43B02CEC5CEC00C7113B /* kernels_quantized in Frameworks */, - 03CF43B22CEC5CEC00C7113B /* kernels_quantized_debug in Frameworks */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXFrameworksBuildPhase section */ - -/* Begin PBXGroup section */ - 0320439E2BB4AC6600050211 /* SupportingFiles */ = { - isa = PBXGroup; - children = ( - 0320439D2BB4AC6600050211 /* LLaMA-Info.plist */, - 03C5F51D2CE7D37100D6CE3F /* Debug.xcconfig */, - 03C5F51C2CE7D35C00D6CE3F /* Release.xcconfig */, - ); - path = SupportingFiles; - sourceTree = ""; - }; - 0324D6892BAACB6900DEF36F /* Application */ = { - isa = PBXGroup; - children = ( - 3072D5222DC3EA280083FC83 /* Constants.swift */, - 0324D6802BAACB6900DEF36F /* App.swift */, - 0324D6812BAACB6900DEF36F /* ContentView.swift */, - 0324D6822BAACB6900DEF36F /* LogManager.swift */, - 0324D6832BAACB6900DEF36F /* LogView.swift */, - 0324D6842BAACB6900DEF36F /* Message.swift */, - 0324D6852BAACB6900DEF36F /* MessageListView.swift */, - 0324D6862BAACB6900DEF36F /* MessageView.swift */, - 0324D6872BAACB6900DEF36F /* ResourceManager.swift */, - 0324D6882BAACB6900DEF36F /* ResourceMonitor.swift */, - 26A6A4272C8A3769005A761E /* ImagePicker.swift */, - ); - path = Application; - sourceTree = ""; - }; - 0324D68A2BAACB6900DEF36F /* LLaMA */ = { - isa = PBXGroup; - children = ( - 0324D6892BAACB6900DEF36F /* Application */, - 0320439E2BB4AC6600050211 /* SupportingFiles */, - ); - path = LLaMA; - sourceTree = ""; - }; - 0324D6952BAACB7000DEF36F /* LLaMAAssets */ = { - isa = PBXGroup; - children = ( - 0324D6942BAACB7000DEF36F /* Assets.xcassets */, - ); - path = LLaMAAssets; - sourceTree = ""; - }; - 0324D69B2BAACB7C00DEF36F /* Exported */ = { - isa = PBXGroup; - children = ( - 0324D6992BAACB7C00DEF36F /* LLaMARunner.h */, - 0324D69A2BAACB7C00DEF36F /* LLaMARunner.mm */, - 03D151B62E0E0908007A38BE /* LLaVARunner.h */, - 03D151B72E0E0908007A38BE /* LLaVARunner.mm */, - ); - path = Exported; - sourceTree = ""; - }; - 0324D69C2BAACB7C00DEF36F /* LLaMARunner */ = { - isa = PBXGroup; - children = ( - 0324D69B2BAACB7C00DEF36F /* Exported */, - 03D151D82E0E9E43007A38BE /* apple */, - 03729F062BB2035900152F2E /* runner */, - 03729F0F2BB203E100152F2E /* sampler */, - 03729F0E2BB203D700152F2E /* tokenizers */, - ); - path = LLaMARunner; - sourceTree = ""; - }; - 0324D69F2BAACB7C00DEF36F /* LLaMARunner */ = { - isa = PBXGroup; - children = ( - 0324D69C2BAACB7C00DEF36F /* LLaMARunner */, - ); - path = LLaMARunner; - sourceTree = ""; - }; - 032C01662AC228E5002955E1 = { - isa = PBXGroup; - children = ( - 0324D68A2BAACB6900DEF36F /* LLaMA */, - 0324D6952BAACB7000DEF36F /* LLaMAAssets */, - 035A5E952BB4B523001E0553 /* LLaMAEntitlements */, - 0324D69F2BAACB7C00DEF36F /* LLaMARunner */, - 036CAF9D2BB1444500D6C2D5 /* LLaMA.app */, - 03729ED52BB1F8DE00152F2E /* LLaMARunner.framework */, - ); - sourceTree = ""; - }; - 035A5E952BB4B523001E0553 /* LLaMAEntitlements */ = { - isa = PBXGroup; - children = ( - 035A5E942BB4B523001E0553 /* LLaMA.entitlements */, - ); - path = LLaMAEntitlements; - sourceTree = ""; - }; - 03729F062BB2035900152F2E /* runner */ = { - isa = PBXGroup; - children = ( - 0372C3132C89418E00CD942A /* llava_runner.cpp */, - 03729F072BB203B300152F2E /* runner.cpp */, - 03D151CC2E0E9ACB007A38BE /* text_decoder_runner.cpp */, - 03D151CD2E0E9ACB007A38BE /* text_llm_runner.cpp */, - 03D151CE2E0E9ACB007A38BE /* text_prefiller.cpp */, - ); - name = runner; - path = ../../../../../../extension/llm/runner; - sourceTree = ""; - }; - 03729F0E2BB203D700152F2E /* tokenizers */ = { - isa = PBXGroup; - children = ( - F24909E72E207004001E5B69 /* normalizer.cpp */, - F292B06A2D88B0C200BE6839 /* bpe_tokenizer_base.cpp */, - 306A71452DC1DC3D00936B1F /* hf_tokenizer.cpp */, - F292B1002D88B20C00BE6839 /* llama_tiktoken.cpp */, - F292B06C2D88B0C200BE6839 /* llama2c_tokenizer.cpp */, - 306A71462DC1DC3D00936B1F /* pcre2_regex.cpp */, - 306A71472DC1DC3D00936B1F /* pre_tokenizer.cpp */, - 306A71482DC1DC3D00936B1F /* re2_regex.cpp */, - 03D151C82E0E98C4007A38BE /* regex_lookahead.cpp */, - 306A71492DC1DC3D00936B1F /* regex.cpp */, - 03D151C92E0E98C4007A38BE /* sentencepiece.cpp */, - 306A714A2DC1DC3D00936B1F /* std_regex.cpp */, - F292B06F2D88B0C200BE6839 /* tiktoken.cpp */, - 306A714B2DC1DC3D00936B1F /* token_decoder.cpp */, - F292B0862D88B0D200BE6839 /* include */, - ); - name = tokenizers; - path = ../../../../../../extension/llm/tokenizers; - sourceTree = ""; - }; - 03729F0F2BB203E100152F2E /* sampler */ = { - isa = PBXGroup; - children = ( - 03729F112BB2042B00152F2E /* sampler.cpp */, - ); - name = sampler; - path = ../../../../../../extension/llm/sampler; - sourceTree = ""; - }; - 03D151D52E0E9E43007A38BE /* Exported */ = { - isa = PBXGroup; - children = ( - 03D151D22E0E9E43007A38BE /* ExecuTorchLLM.h */, - 03D151D32E0E9E43007A38BE /* ExecuTorchTextLLMRunner.h */, - 03D151D42E0E9E43007A38BE /* ExecuTorchTextLLMRunner.mm */, - ); - path = Exported; - sourceTree = ""; - }; - 03D151D62E0E9E43007A38BE /* ExecuTorchLLM */ = { - isa = PBXGroup; - children = ( - 03D151D52E0E9E43007A38BE /* Exported */, - ); - path = ExecuTorchLLM; - sourceTree = ""; - }; - 03D151D82E0E9E43007A38BE /* apple */ = { - isa = PBXGroup; - children = ( - 03D151D62E0E9E43007A38BE /* ExecuTorchLLM */, - ); - name = apple; - path = ../../../../../../extension/llm/apple; - sourceTree = ""; - }; - F292B0842D88B0D200BE6839 /* tokenizers */ = { - isa = PBXGroup; - children = ( - F292B0782D88B0D200BE6839 /* base64.h */, - F292B0792D88B0D200BE6839 /* bpe_tokenizer_base.h */, - F292B07A2D88B0D200BE6839 /* error.h */, - 306A71352DC1DC0F00936B1F /* hf_tokenizer.h */, - F292B07C2D88B0D200BE6839 /* llama2c_tokenizer.h */, - F292B07D2D88B0D200BE6839 /* log.h */, - 306A71362DC1DC0F00936B1F /* pcre2_regex.h */, - 306A71372DC1DC0F00936B1F /* pre_tokenizer.h */, - 306A71382DC1DC0F00936B1F /* re2_regex.h */, - 306A71392DC1DC0F00936B1F /* regex.h */, - F292B07F2D88B0D200BE6839 /* result.h */, - 306A713A2DC1DC0F00936B1F /* std_regex.h */, - 306A713B2DC1DC0F00936B1F /* string_integer_map.h */, - F292B0812D88B0D200BE6839 /* tiktoken.h */, - 306A713C2DC1DC0F00936B1F /* token_decoder.h */, - F292B0832D88B0D200BE6839 /* tokenizer.h */, - ); - path = tokenizers; - sourceTree = ""; - }; - F292B0852D88B0D200BE6839 /* pytorch */ = { - isa = PBXGroup; - children = ( - F292B0842D88B0D200BE6839 /* tokenizers */, - ); - path = pytorch; - sourceTree = ""; - }; - F292B0862D88B0D200BE6839 /* include */ = { - isa = PBXGroup; - children = ( - F292B0852D88B0D200BE6839 /* pytorch */, - ); - path = include; - sourceTree = ""; - }; -/* End PBXGroup section */ - -/* Begin PBXHeadersBuildPhase section */ - 03729ED02BB1F8DE00152F2E /* Headers */ = { - isa = PBXHeadersBuildPhase; - buildActionMask = 2147483647; - files = ( - 03729EE22BB1F93E00152F2E /* LLaMARunner.h in Headers */, - F292B0882D88B0D200BE6839 /* llama2c_tokenizer.h in Headers */, - F292B0892D88B0D200BE6839 /* tokenizer.h in Headers */, - F292B08B2D88B0D200BE6839 /* result.h in Headers */, - F292B08D2D88B0D200BE6839 /* error.h in Headers */, - 03D151DB2E0E9E43007A38BE /* ExecuTorchTextLLMRunner.h in Headers */, - 03D151DC2E0E9E43007A38BE /* ExecuTorchLLM.h in Headers */, - F292B08E2D88B0D200BE6839 /* bpe_tokenizer_base.h in Headers */, - F292B08F2D88B0D200BE6839 /* log.h in Headers */, - F292B0912D88B0D200BE6839 /* tiktoken.h in Headers */, - F292B0922D88B0D200BE6839 /* base64.h in Headers */, - 306A713D2DC1DC0F00936B1F /* token_decoder.h in Headers */, - 306A713E2DC1DC0F00936B1F /* regex.h in Headers */, - 306A713F2DC1DC0F00936B1F /* string_integer_map.h in Headers */, - 306A71402DC1DC0F00936B1F /* re2_regex.h in Headers */, - 306A71412DC1DC0F00936B1F /* hf_tokenizer.h in Headers */, - 306A71422DC1DC0F00936B1F /* pre_tokenizer.h in Headers */, - 306A71432DC1DC0F00936B1F /* pcre2_regex.h in Headers */, - 306A71442DC1DC0F00936B1F /* std_regex.h in Headers */, - 03D151B92E0E0908007A38BE /* LLaVARunner.h in Headers */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXHeadersBuildPhase section */ - -/* Begin PBXNativeTarget section */ - 032C016E2AC228E6002955E1 /* LLaMA */ = { - isa = PBXNativeTarget; - buildConfigurationList = 032C017D2AC228E7002955E1 /* Build configuration list for PBXNativeTarget "LLaMA" */; - buildPhases = ( - 032C016B2AC228E6002955E1 /* Sources */, - 032C016C2AC228E6002955E1 /* Frameworks */, - 032C016D2AC228E6002955E1 /* Resources */, - 03729EE02BB1F8DE00152F2E /* Embed Frameworks */, - ); - buildRules = ( - ); - dependencies = ( - 03729EDA2BB1F8DE00152F2E /* PBXTargetDependency */, - ); - name = LLaMA; - packageProductDependencies = ( - ); - productName = LLaMA; - productReference = 036CAF9D2BB1444500D6C2D5 /* LLaMA.app */; - productType = "com.apple.product-type.application"; - }; - 03729ED42BB1F8DE00152F2E /* LLaMARunner */ = { - isa = PBXNativeTarget; - buildConfigurationList = 03729EDD2BB1F8DE00152F2E /* Build configuration list for PBXNativeTarget "LLaMARunner" */; - buildPhases = ( - 03641D7C2C409960004AF8DE /* Build Cmake Dependencies */, - 03729ED02BB1F8DE00152F2E /* Headers */, - 03729ED12BB1F8DE00152F2E /* Sources */, - 03729ED22BB1F8DE00152F2E /* Frameworks */, - 03729ED32BB1F8DE00152F2E /* Resources */, - ); - buildRules = ( - ); - dependencies = ( - ); - name = LLaMARunner; - packageProductDependencies = ( - 03CF43952CEC5CEC00C7113B /* backend_coreml */, - 03CF43972CEC5CEC00C7113B /* backend_coreml_debug */, - 03CF43992CEC5CEC00C7113B /* backend_mps */, - 03CF439B2CEC5CEC00C7113B /* backend_mps_debug */, - 03CF439D2CEC5CEC00C7113B /* backend_xnnpack */, - 03CF439F2CEC5CEC00C7113B /* backend_xnnpack_debug */, - 03CF43A12CEC5CEC00C7113B /* executorch_debug */, - 03CF43A32CEC5CEC00C7113B /* kernels_llm */, - 03CF43A52CEC5CEC00C7113B /* kernels_llm_debug */, - 03CF43A72CEC5CEC00C7113B /* kernels_optimized */, - 03CF43A92CEC5CEC00C7113B /* kernels_optimized_debug */, - 03CF43AF2CEC5CEC00C7113B /* kernels_quantized */, - 03CF43B12CEC5CEC00C7113B /* kernels_quantized_debug */, - ); - productName = LLaMARunner; - productReference = 03729ED52BB1F8DE00152F2E /* LLaMARunner.framework */; - productType = "com.apple.product-type.framework"; - }; -/* End PBXNativeTarget section */ - -/* Begin PBXProject section */ - 032C01672AC228E5002955E1 /* Project object */ = { - isa = PBXProject; - attributes = { - BuildIndependentTargetsInParallel = 1; - LastSwiftUpdateCheck = 1540; - LastUpgradeCheck = 1530; - TargetAttributes = { - 032C016E2AC228E6002955E1 = { - CreatedOnToolsVersion = 15.0; - }; - 03729ED42BB1F8DE00152F2E = { - CreatedOnToolsVersion = 15.3; - }; - }; - }; - buildConfigurationList = 032C016A2AC228E5002955E1 /* Build configuration list for PBXProject "LLaMA" */; - compatibilityVersion = "Xcode 14.0"; - developmentRegion = en; - hasScannedForEncodings = 0; - knownRegions = ( - en, - Base, - ); - mainGroup = 032C01662AC228E5002955E1; - packageReferences = ( - 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */, - ); - productRefGroup = 032C01662AC228E5002955E1; - projectDirPath = ""; - projectRoot = ""; - targets = ( - 032C016E2AC228E6002955E1 /* LLaMA */, - 03729ED42BB1F8DE00152F2E /* LLaMARunner */, - ); - }; -/* End PBXProject section */ - -/* Begin PBXResourcesBuildPhase section */ - 032C016D2AC228E6002955E1 /* Resources */ = { - isa = PBXResourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 0324D6962BAACB7000DEF36F /* Assets.xcassets in Resources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; - 03729ED32BB1F8DE00152F2E /* Resources */ = { - isa = PBXResourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXResourcesBuildPhase section */ - -/* Begin PBXShellScriptBuildPhase section */ - 03641D7C2C409960004AF8DE /* Build Cmake Dependencies */ = { - isa = PBXShellScriptBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - inputFileListPaths = ( - ); - inputPaths = ( - ); - name = "Build Cmake Dependencies"; - outputFileListPaths = ( - ); - outputPaths = ( - "$(DERIVED_FILE_DIR)/newOutputFile", - ); - runOnlyForDeploymentPostprocessing = 0; - shellPath = /bin/sh; - shellScript = "set -e\n\nif ! command -v cmake &> /dev/null\nthen\n echo \"Cmake not found, please install Cmake. \\n1. Download Cmake.app from https://cmake.org/download with version > 3.19. \\n2. Install it to Applications/ folder and run `sudo /Applications/CMake.app/Contents/bin/cmake-gui --install` to install CMake commandline tools.\"\n exit 1\nfi\n\nCMAKE_DIR=\"$TEMP_DIR/cmake\"\nrm -rf \"$CMAKE_DIR\"\n\nPLATFORM=\"SIMULATORARM64\"\nDEPLOYMENT_TARGET=\"17.0\"\n\nif [[ \"$PLATFORM_NAME\" == *\"iphoneos\"* ]]; then\n PLATFORM=\"OS64\"\nelif [[ \"$PLATFORM_NAME\" == *\"macos\"* ]]; then\n PLATFORM=\"MAC_ARM64\"\n DEPLOYMENT_TARGET=\"12.0\"\nfi\n\ncmake_build() {\n local src_dir target do_install=0\n local extra_args=()\n local build_dir\n # Parse arguments\n src_dir=\"$1\"\n shift\n target=\"$1\"\n if [[ \"$target\" == \"install\" ]]; then\n # Usage: cmake_build install [extra_args...]\n do_install=1\n shift\n else\n # Usage: cmake_build [install] [extra_args...]\n shift\n if [[ \"$1\" == \"install\" ]]; then\n do_install=1\n shift\n fi\n fi\n # Collect any remaining arguments as extra_args\n extra_args=(\"$@\")\n build_dir=\"$CMAKE_DIR/build/$(basename \"$src_dir\")\"\n mkdir -p \"$build_dir\" || { echo \"Failed to create build dir\"; return 1; }\n pushd \"$build_dir\" > /dev/null || { echo \"Failed to enter build dir\"; return 1; }\n # Platform-specific CMake args\n if [[ \"$PLATFORM\" == \"MAC_ARM64\" ]]; then\n extra_args+=(-DCMAKE_INSTALL_BUNDLEDIR=\"${CMAKE_DIR}/bin\")\n extra_args+=(-DCMAKE_MACOSX_BUNDLE=OFF)\n fi\n # Configure\n cmake -G Xcode \\\n -DCMAKE_BUILD_TYPE=\"Release\" \\\n -DCMAKE_CXX_STANDARD=17 \\\n -DCMAKE_TOOLCHAIN_FILE=\"$SRCROOT/../../../../third-party/ios-cmake/ios.toolchain.cmake\" \\\n -DCMAKE_XCODE_ATTRIBUTE_CLANG_CXX_LANGUAGE_STANDARD=\"c++17\" \\\n -DCMAKE_XCODE_ATTRIBUTE_CLANG_CXX_LIBRARY=\"libc++\" \\\n -DPLATFORM=\"$PLATFORM\" \\\n -DDEPLOYMENT_TARGET=\"$DEPLOYMENT_TARGET\" \\\n -DCMAKE_INSTALL_PREFIX=\"$CMAKE_DIR\" \\\n \"${extra_args[@]}\" \\\n \"$src_dir\" || { echo \"CMake configure failed\"; popd > /dev/null; return 1; }\n # Build\n cmake --build . --config \"Release\" --target $target\n # Install if requested\n if [[ $do_install -eq 1 ]]; then\n cmake --install . --prefix \"$CMAKE_DIR\" || echo \"Ignoring install failures\"\n fi\n}\n\ncmake_build \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/abseil-cpp\" \"install\" \\\n -DABSL_PROPAGATE_CXX_STD=ON\n\ncmake_build \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/re2\" \"install\"\n\ncmake_build \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/pcre2\" \"install\" \\\n -DPCRE2_BUILD_PCRE2_8=ON \\\n -DPCRE2_BUILD_PCRE2_16=OFF \\\n -DPCRE2_BUILD_PCRE2_32=OFF \\\n -DPCRE2_BUILD_TESTS=OFF \\\n -DPCRE2_BUILD_PCRE2GREP=OFF \\\n -DPCRE2_BUILD_PCRE2TEST=OFF \\\n -DPCRE2_BUILD_PCRE2GPERF=OFF \\\n -DPCRE2_BUILD_DOCS=OFF \\\n -DPCRE2_BUILD_LIBPCRE2_PDB=OFF \\\n -DSUPPORT_REGEX_LOOKAHEAD=ON\n\ncmake_build \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/sentencepiece\" \"sentencepiece-static sentencepiece_train-static\" \"install\" \\\n -DSPM_ENABLE_SHARED=OFF \\\n -DSPM_BUILD_TEST=OFF \\\n -DCMAKE_SYSTEM_NAME=\"iOS\"\n \ncmake_build \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/llama.cpp-unicode\" \"install\"\n \n# Include the single header for json.\nmkdir -p \"$CMAKE_DIR/include/nlohmann\"\ncp \"$SRCROOT/../../../../extension/llm/tokenizers/third-party/json/single_include/nlohmann/json.hpp\" \"$CMAKE_DIR/include/nlohmann/json.hpp\"\n\necho \"$(find $CMAKE_DIR/lib -name \"*.a\" | sed -E 's|^.*/lib([^/]+)\\.a|-l\\1|g' | tr '\\n' ' ')\" > \"$CMAKE_DIR/linker_flags\"\n"; - }; -/* End PBXShellScriptBuildPhase section */ - -/* Begin PBXSourcesBuildPhase section */ - 032C016B2AC228E6002955E1 /* Sources */ = { - isa = PBXSourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 0324D6932BAACB6900DEF36F /* ResourceMonitor.swift in Sources */, - 3072D5232DC3EA280083FC83 /* Constants.swift in Sources */, - 0324D68D2BAACB6900DEF36F /* LogManager.swift in Sources */, - 0324D68E2BAACB6900DEF36F /* LogView.swift in Sources */, - 0324D68F2BAACB6900DEF36F /* Message.swift in Sources */, - 0324D6922BAACB6900DEF36F /* ResourceManager.swift in Sources */, - 0324D68C2BAACB6900DEF36F /* ContentView.swift in Sources */, - 0324D6902BAACB6900DEF36F /* MessageListView.swift in Sources */, - 26A6A4282C8A3769005A761E /* ImagePicker.swift in Sources */, - 0324D6912BAACB6900DEF36F /* MessageView.swift in Sources */, - 0324D68B2BAACB6900DEF36F /* App.swift in Sources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; - 03729ED12BB1F8DE00152F2E /* Sources */ = { - isa = PBXSourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 03D151B82E0E0908007A38BE /* LLaVARunner.mm in Sources */, - 03729EE12BB1F93800152F2E /* LLaMARunner.mm in Sources */, - F24909E82E207004001E5B69 /* normalizer.cpp in Sources */, - 0372C3152C89418E00CD942A /* llava_runner.cpp in Sources */, - 03D151CA2E0E98C4007A38BE /* sentencepiece.cpp in Sources */, - 03D151D92E0E9E43007A38BE /* ExecuTorchTextLLMRunner.mm in Sources */, - 03D151CB2E0E98C4007A38BE /* regex_lookahead.cpp in Sources */, - 03D151CF2E0E9ACB007A38BE /* text_prefiller.cpp in Sources */, - 03D151D02E0E9ACB007A38BE /* text_llm_runner.cpp in Sources */, - 03D151D12E0E9ACB007A38BE /* text_decoder_runner.cpp in Sources */, - F292B1022D88B20C00BE6839 /* llama_tiktoken.cpp in Sources */, - F292B0752D88B0C200BE6839 /* tiktoken.cpp in Sources */, - F292B0762D88B0C200BE6839 /* llama2c_tokenizer.cpp in Sources */, - F292B0772D88B0C200BE6839 /* bpe_tokenizer_base.cpp in Sources */, - 03729F0A2BB203B300152F2E /* runner.cpp in Sources */, - 03729F132BB2042B00152F2E /* sampler.cpp in Sources */, - 306A714C2DC1DC3D00936B1F /* std_regex.cpp in Sources */, - 306A714D2DC1DC3D00936B1F /* re2_regex.cpp in Sources */, - 306A714E2DC1DC3D00936B1F /* pcre2_regex.cpp in Sources */, - 306A714F2DC1DC3D00936B1F /* hf_tokenizer.cpp in Sources */, - 306A71502DC1DC3D00936B1F /* regex.cpp in Sources */, - 306A71512DC1DC3D00936B1F /* pre_tokenizer.cpp in Sources */, - 306A71522DC1DC3D00936B1F /* token_decoder.cpp in Sources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXSourcesBuildPhase section */ - -/* Begin PBXTargetDependency section */ - 03729EDA2BB1F8DE00152F2E /* PBXTargetDependency */ = { - isa = PBXTargetDependency; - target = 03729ED42BB1F8DE00152F2E /* LLaMARunner */; - targetProxy = 03729ED92BB1F8DE00152F2E /* PBXContainerItemProxy */; - }; -/* End PBXTargetDependency section */ - -/* Begin XCBuildConfiguration section */ - 032C017B2AC228E7002955E1 /* Debug */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_CXX_LANGUAGE_STANDARD = "c++17"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_ENABLE_OBJC_WEAK = YES; - CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_COMMA = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; - CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; - CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; - CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; - CLANG_WARN_STRICT_PROTOTYPES = YES; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = dwarf; - ENABLE_STRICT_OBJC_MSGSEND = YES; - ENABLE_TESTABILITY = YES; - ENABLE_USER_SCRIPT_SANDBOXING = YES; - "EXCLUDED_ARCHS[sdk=iphonesimulator*]" = x86_64; - GCC_C_LANGUAGE_STANDARD = c17; - GCC_DYNAMIC_NO_PIC = NO; - GCC_NO_COMMON_BLOCKS = YES; - GCC_OPTIMIZATION_LEVEL = 0; - GCC_PREPROCESSOR_DEFINITIONS = ( - "DEBUG=1", - "$(inherited)", - ); - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - IPHONEOS_DEPLOYMENT_TARGET = 17.0; - LD_RUNPATH_SEARCH_PATHS = ( - "$(inherited)", - "@executable_path/Frameworks", - "@loader_path/Frameworks", - ); - LOCALIZATION_PREFERS_STRING_CATALOGS = YES; - MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; - MTL_FAST_MATH = YES; - ONLY_ACTIVE_ARCH = YES; - SDKROOT = iphoneos; - SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)"; - SWIFT_OPTIMIZATION_LEVEL = "-Onone"; - SWIFT_VERSION = 5.0; - }; - name = Debug; - }; - 032C017C2AC228E7002955E1 /* Release */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_CXX_LANGUAGE_STANDARD = "c++17"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_ENABLE_OBJC_WEAK = YES; - CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_COMMA = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; - CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; - CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; - CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; - CLANG_WARN_STRICT_PROTOTYPES = YES; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; - ENABLE_NS_ASSERTIONS = NO; - ENABLE_STRICT_OBJC_MSGSEND = YES; - ENABLE_USER_SCRIPT_SANDBOXING = YES; - "EXCLUDED_ARCHS[sdk=iphonesimulator*]" = x86_64; - GCC_C_LANGUAGE_STANDARD = c17; - GCC_NO_COMMON_BLOCKS = YES; - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - IPHONEOS_DEPLOYMENT_TARGET = 17.0; - LD_RUNPATH_SEARCH_PATHS = ( - "$(inherited)", - "@executable_path/Frameworks", - "@loader_path/Frameworks", - ); - LOCALIZATION_PREFERS_STRING_CATALOGS = YES; - MTL_ENABLE_DEBUG_INFO = NO; - MTL_FAST_MATH = YES; - SDKROOT = iphoneos; - SWIFT_COMPILATION_MODE = wholemodule; - SWIFT_VERSION = 5.0; - VALIDATE_PRODUCT = YES; - }; - name = Release; - }; - 032C017E2AC228E7002955E1 /* Debug */ = { - isa = XCBuildConfiguration; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CODE_SIGN_ENTITLEMENTS = LLaMAEntitlements/LLaMA.entitlements; - CODE_SIGN_IDENTITY = "Apple Development"; - CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; - DEVELOPMENT_TEAM = ""; - ENABLE_PREVIEWS = YES; - GENERATE_INFOPLIST_FILE = YES; - INFOPLIST_FILE = "LLaMA/SupportingFiles/LLaMA-Info.plist"; - INFOPLIST_KEY_CFBundleDisplayName = iLLaMA; - INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES; - INFOPLIST_KEY_NSCameraUsageDescription = ""; - INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; - INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; - INFOPLIST_KEY_UILaunchScreen_Generation = YES; - INFOPLIST_KEY_UIRequiresFullScreen = YES; - INFOPLIST_KEY_UISupportedInterfaceOrientations = UIInterfaceOrientationPortrait; - MARKETING_VERSION = 1.0; - PRODUCT_BUNDLE_IDENTIFIER = org.pytorch.executorch.illama; - PRODUCT_NAME = "$(PROJECT_NAME)"; - PROVISIONING_PROFILE_SPECIFIER = ""; - SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; - SUPPORTS_MACCATALYST = NO; - SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO; - SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO; - SWIFT_EMIT_LOC_STRINGS = YES; - TARGETED_DEVICE_FAMILY = "1,2"; - }; - name = Debug; - }; - 032C017F2AC228E7002955E1 /* Release */ = { - isa = XCBuildConfiguration; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CODE_SIGN_ENTITLEMENTS = LLaMAEntitlements/LLaMA.entitlements; - CODE_SIGN_IDENTITY = "Apple Development"; - CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; - DEVELOPMENT_TEAM = ""; - ENABLE_PREVIEWS = YES; - GENERATE_INFOPLIST_FILE = YES; - INFOPLIST_FILE = "LLaMA/SupportingFiles/LLaMA-Info.plist"; - INFOPLIST_KEY_CFBundleDisplayName = iLLaMA; - INFOPLIST_KEY_LSSupportsOpeningDocumentsInPlace = YES; - INFOPLIST_KEY_NSCameraUsageDescription = ""; - INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; - INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; - INFOPLIST_KEY_UILaunchScreen_Generation = YES; - INFOPLIST_KEY_UIRequiresFullScreen = YES; - INFOPLIST_KEY_UISupportedInterfaceOrientations = UIInterfaceOrientationPortrait; - MARKETING_VERSION = 1.0; - PRODUCT_BUNDLE_IDENTIFIER = org.pytorch.executorch.illama; - PRODUCT_NAME = "$(PROJECT_NAME)"; - PROVISIONING_PROFILE_SPECIFIER = ""; - SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; - SUPPORTS_MACCATALYST = NO; - SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO; - SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO; - SWIFT_EMIT_LOC_STRINGS = YES; - TARGETED_DEVICE_FAMILY = "1,2"; - }; - name = Release; - }; - 03729EDE2BB1F8DE00152F2E /* Debug */ = { - isa = XCBuildConfiguration; - baseConfigurationReference = 03C5F51D2CE7D37100D6CE3F /* Debug.xcconfig */; - buildSettings = { - CODE_SIGN_IDENTITY = ""; - CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; - DEFINES_MODULE = YES; - DEVELOPMENT_TEAM = ""; - DYLIB_COMPATIBILITY_VERSION = 1; - DYLIB_CURRENT_VERSION = 1; - DYLIB_INSTALL_NAME_BASE = "@rpath"; - ENABLE_MODULE_VERIFIER = YES; - GCC_C_LANGUAGE_STANDARD = gnu17; - GCC_PREPROCESSOR_DEFINITIONS = ( - "DEBUG=1", - "ET_USE_TIKTOKEN=1", - "SUPPORT_REGEX_LOOKAHEAD=ON", - "BUILD_WITH_XCODE=1", - ); - GENERATE_INFOPLIST_FILE = YES; - INFOPLIST_KEY_NSHumanReadableCopyright = ""; - INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks"; - IPHONEOS_DEPLOYMENT_TARGET = 17.0; - MARKETING_VERSION = 1.0; - MODULE_VERIFIER_SUPPORTED_LANGUAGES = "objective-c objective-c++"; - MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20"; - PRODUCT_BUNDLE_IDENTIFIER = org.pytorch.executorch.illama.LLaMARunner; - PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)"; - SKIP_INSTALL = YES; - SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; - SUPPORTS_MACCATALYST = NO; - SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO; - SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO; - SWIFT_EMIT_LOC_STRINGS = YES; - TARGETED_DEVICE_FAMILY = "1,2"; - VERSIONING_SYSTEM = "apple-generic"; - VERSION_INFO_PREFIX = ""; - }; - name = Debug; - }; - 03729EDF2BB1F8DE00152F2E /* Release */ = { - isa = XCBuildConfiguration; - baseConfigurationReference = 03C5F51C2CE7D35C00D6CE3F /* Release.xcconfig */; - buildSettings = { - CODE_SIGN_IDENTITY = ""; - CODE_SIGN_STYLE = Automatic; - CURRENT_PROJECT_VERSION = 1; - DEFINES_MODULE = YES; - DEVELOPMENT_TEAM = ""; - DYLIB_COMPATIBILITY_VERSION = 1; - DYLIB_CURRENT_VERSION = 1; - DYLIB_INSTALL_NAME_BASE = "@rpath"; - ENABLE_MODULE_VERIFIER = YES; - GCC_C_LANGUAGE_STANDARD = gnu17; - GCC_PREPROCESSOR_DEFINITIONS = ( - "SUPPORT_REGEX_LOOKAHEAD=ON", - "BUILD_WITH_XCODE=1", - ); - GENERATE_INFOPLIST_FILE = YES; - INFOPLIST_KEY_NSHumanReadableCopyright = ""; - INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks"; - IPHONEOS_DEPLOYMENT_TARGET = 17.0; - MARKETING_VERSION = 1.0; - MODULE_VERIFIER_SUPPORTED_LANGUAGES = "objective-c objective-c++"; - MODULE_VERIFIER_SUPPORTED_LANGUAGE_STANDARDS = "gnu17 gnu++20"; - PRODUCT_BUNDLE_IDENTIFIER = org.pytorch.executorch.illama.LLaMARunner; - PRODUCT_NAME = "$(TARGET_NAME:c99extidentifier)"; - SKIP_INSTALL = YES; - SUPPORTED_PLATFORMS = "iphoneos iphonesimulator"; - SUPPORTS_MACCATALYST = NO; - SUPPORTS_MAC_DESIGNED_FOR_IPHONE_IPAD = NO; - SUPPORTS_XR_DESIGNED_FOR_IPHONE_IPAD = NO; - SWIFT_EMIT_LOC_STRINGS = YES; - TARGETED_DEVICE_FAMILY = "1,2"; - VERSIONING_SYSTEM = "apple-generic"; - VERSION_INFO_PREFIX = ""; - }; - name = Release; - }; -/* End XCBuildConfiguration section */ - -/* Begin XCConfigurationList section */ - 032C016A2AC228E5002955E1 /* Build configuration list for PBXProject "LLaMA" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 032C017B2AC228E7002955E1 /* Debug */, - 032C017C2AC228E7002955E1 /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; - 032C017D2AC228E7002955E1 /* Build configuration list for PBXNativeTarget "LLaMA" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 032C017E2AC228E7002955E1 /* Debug */, - 032C017F2AC228E7002955E1 /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; - 03729EDD2BB1F8DE00152F2E /* Build configuration list for PBXNativeTarget "LLaMARunner" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 03729EDE2BB1F8DE00152F2E /* Debug */, - 03729EDF2BB1F8DE00152F2E /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; -/* End XCConfigurationList section */ - -/* Begin XCRemoteSwiftPackageReference section */ - 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */ = { - isa = XCRemoteSwiftPackageReference; - repositoryURL = "https://github.com/pytorch/executorch"; - requirement = { - branch = "swiftpm-0.8.0.20250724"; - kind = branch; - }; - }; -/* End XCRemoteSwiftPackageReference section */ - -/* Begin XCSwiftPackageProductDependency section */ - 03CF43952CEC5CEC00C7113B /* backend_coreml */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_coreml; - }; - 03CF43972CEC5CEC00C7113B /* backend_coreml_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_coreml_debug; - }; - 03CF43992CEC5CEC00C7113B /* backend_mps */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_mps; - }; - 03CF439B2CEC5CEC00C7113B /* backend_mps_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_mps_debug; - }; - 03CF439D2CEC5CEC00C7113B /* backend_xnnpack */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_xnnpack; - }; - 03CF439F2CEC5CEC00C7113B /* backend_xnnpack_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = backend_xnnpack_debug; - }; - 03CF43A12CEC5CEC00C7113B /* executorch_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = executorch_debug; - }; - 03CF43A32CEC5CEC00C7113B /* kernels_llm */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_llm; - }; - 03CF43A52CEC5CEC00C7113B /* kernels_llm_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_llm_debug; - }; - 03CF43A72CEC5CEC00C7113B /* kernels_optimized */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_optimized; - }; - 03CF43A92CEC5CEC00C7113B /* kernels_optimized_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_optimized_debug; - }; - 03CF43AF2CEC5CEC00C7113B /* kernels_quantized */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_quantized; - }; - 03CF43B12CEC5CEC00C7113B /* kernels_quantized_debug */ = { - isa = XCSwiftPackageProductDependency; - package = 03CF43942CEC5CEC00C7113B /* XCRemoteSwiftPackageReference "executorch" */; - productName = kernels_quantized_debug; - }; -/* End XCSwiftPackageProductDependency section */ - }; - rootObject = 032C01672AC228E5002955E1 /* Project object */; -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMA.xcscheme b/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMA.xcscheme deleted file mode 100644 index 84fa8d52802..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMA.xcscheme +++ /dev/null @@ -1,76 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMARunner.xcscheme b/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMARunner.xcscheme deleted file mode 100644 index d820e0a5f8a..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMARunner.xcscheme +++ /dev/null @@ -1,67 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Constants.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Constants.swift deleted file mode 100644 index 1c2a9d12b97..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Constants.swift +++ /dev/null @@ -1,30 +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. - */ - -import Foundation - -public enum Constants { - public static let qwen3PromptTemplate = """ -<|im_start|>system -You are a helpful assistant. -<|im_end|> -<|im_start|>user -%@<|im_end|> -<|im_start|>assistant - - - - - - -""" - - public static let llama3PromptTemplate = "<|begin_of_text|><|start_header_id|>user<|end_header_id|>%@<|eot_id|><|start_header_id|>assistant<|end_header_id|>" - -public static let phi4PromptTemplate = "<|user|>%@<|end|><|assistant|>" -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift deleted file mode 100644 index c6b8b71dfc1..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift +++ /dev/null @@ -1,605 +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. - */ - -import SwiftUI -import UniformTypeIdentifiers - -import LLaMARunner - -class RunnerHolder: ObservableObject { - var llamaRunner: LLaMARunner? - var llavaRunner: LLaVARunner? -} - -extension UIImage { - func resized(to newSize: CGSize) -> UIImage { - let format = UIGraphicsImageRendererFormat.default() - format.scale = 1 - return UIGraphicsImageRenderer(size: newSize, format: format).image { - _ in draw(in: CGRect(origin: .zero, size: newSize)) - } - } - - func toRGBArray() -> [UInt8]? { - guard let cgImage = self.cgImage else { return nil } - - let width = Int(cgImage.width), height = Int(cgImage.height) - let totalPixels = width * height, bytesPerPixel = 4, bytesPerRow = bytesPerPixel * width - var rgbValues = [UInt8](repeating: 0, count: totalPixels * 3) - var pixelData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) - - guard let context = CGContext( - data: &pixelData, width: width, height: height, bitsPerComponent: 8, - bytesPerRow: bytesPerRow, space: CGColorSpaceCreateDeviceRGB(), - bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue | CGBitmapInfo.byteOrder32Big.rawValue - ) else { return nil } - - context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height)) - - for y in 0.. ModelType { - let filename = (path as NSString).lastPathComponent.lowercased() - if filename.hasPrefix("llama") { - return .llama - } else if filename.hasPrefix("llava") { - return .llava - } else if filename.hasPrefix("qwen3") { - return .qwen3 - } else if filename.hasPrefix("phi4") { - return .phi4 - } - print("Unknown model type in path: \(path). Model filename should start with one of: llama, llava, qwen3, or phi4") - exit(1) - } - } - - private var placeholder: String { - resourceManager.isModelValid ? resourceManager.isTokenizerValid ? "Prompt..." : "Select Tokenizer..." : "Select Model..." - } - - private var title: String { - resourceManager.isModelValid ? resourceManager.isTokenizerValid ? resourceManager.modelName : "Select Tokenizer..." : "Select Model..." - } - - private var modelTitle: String { - resourceManager.isModelValid ? resourceManager.modelName : "Select Model..." - } - - private var tokenizerTitle: String { - resourceManager.isTokenizerValid ? resourceManager.tokenizerName : "Select Tokenizer..." - } - - private var isInputEnabled: Bool { resourceManager.isModelValid && resourceManager.isTokenizerValid } - - var body: some View { - NavigationView { - ZStack { - VStack { - if showingSettings { - VStack(spacing: 20) { - HStack { - VStack(spacing: 10) { - Button(action: { pickerType = .model }) { - Label(modelTitle, systemImage: "doc") - .lineLimit(1) - .truncationMode(.middle) - .frame(maxWidth: 300, alignment: .leading) - } - Button(action: { pickerType = .tokenizer }) { - Label(tokenizerTitle, systemImage: "doc") - .lineLimit(1) - .truncationMode(.middle) - .frame(maxWidth: 300, alignment: .leading) - } - } - .padding() - .background(Color.gray.opacity(0.1)) - .cornerRadius(10) - .fixedSize(horizontal: true, vertical: false) - Spacer() - } - .padding() - } - } - - MessageListView(messages: $messages) - .simultaneousGesture( - DragGesture().onChanged { value in - if value.translation.height > 10 { - hideKeyboard() - } - showingSettings = false - textFieldFocused = false - } - ) - .onTapGesture { - showingSettings = false - textFieldFocused = false - } - - HStack { - Button(action: { - imagePickerSourceType = .photoLibrary - isImagePickerPresented = true - }) { - Image(systemName: "photo.on.rectangle") - .resizable() - .scaledToFit() - .frame(width: 24, height: 24) - } - .background(Color.clear) - .cornerRadius(8) - - Button(action: { - if UIImagePickerController.isSourceTypeAvailable(.camera) { - imagePickerSourceType = .camera - isImagePickerPresented = true - } else { - print("Camera not available") - } - }) { - Image(systemName: "camera") - .resizable() - .scaledToFit() - .frame(width: 24, height: 24) - } - .background(Color.clear) - .cornerRadius(8) - - if resourceManager.isModelValid && ModelType.fromPath(resourceManager.modelPath) == .qwen3 { - Button(action: { - thinkingMode.toggle() - showThinkingModeNotification = true - DispatchQueue.main.asyncAfter(deadline: .now() + 3) { - showThinkingModeNotification = false - } - }) { - Image(systemName: "brain") - .resizable() - .scaledToFit() - .frame(width: 24, height: 24) - .foregroundColor(thinkingMode ? .blue : .gray) - } - .background(Color.clear) - .cornerRadius(8) - } - - TextField(placeholder, text: $prompt, axis: .vertical) - .padding(8) - .background(Color.gray.opacity(0.1)) - .cornerRadius(20) - .lineLimit(1...10) - .overlay( - RoundedRectangle(cornerRadius: 20) - .stroke(isInputEnabled ? Color.blue : Color.gray, lineWidth: 1) - ) - .disabled(!isInputEnabled) - .focused($textFieldFocused) - .onAppear { textFieldFocused = false } - .onTapGesture { - showingSettings = false - } - - Button(action: isGenerating ? stop : generate) { - Image(systemName: isGenerating ? "stop.circle" : "arrowshape.up.circle.fill") - .resizable() - .aspectRatio(contentMode: .fit) - .frame(height: 28) - } - .disabled(isGenerating ? shouldStopGenerating : (!isInputEnabled || prompt.isEmpty)) - } - .padding([.leading, .trailing, .bottom], 10) - } - .sheet(isPresented: $isImagePickerPresented, onDismiss: addSelectedImageMessage) { - ImagePicker(selectedImage: $selectedImage, sourceType: imagePickerSourceType) - .id(imagePickerSourceType.rawValue) - } - - if showThinkingModeNotification { - Text(thinkingMode ? "Thinking mode enabled" : "Thinking mode disabled") - .padding(8) - .background(Color(UIColor.secondarySystemBackground)) - .cornerRadius(8) - .transition(.opacity) - .animation(.easeInOut(duration: 0.2), value: showThinkingModeNotification) - } - } - .navigationBarTitle(title, displayMode: .inline) - .navigationBarItems( - leading: - Button(action: { - showingSettings.toggle() - }) { - Image(systemName: "folder") - .imageScale(.large) - }, - trailing: - HStack { - Menu { - Section(header: Text("Memory")) { - Text("Used: \(resourceMonitor.usedMemory) Mb") - Text("Available: \(resourceMonitor.usedMemory) Mb") - } - } label: { - Text("\(resourceMonitor.usedMemory) Mb") - } - .onAppear { - resourceMonitor.start() - } - .onDisappear { - resourceMonitor.stop() - } - Button(action: { showingLogs = true }) { - Image(systemName: "list.bullet.rectangle") - } - } - ) - .sheet(isPresented: $showingLogs) { - NavigationView { - LogView(logManager: logManager) - } - } - .fileImporter( - isPresented: Binding( - get: { pickerType != nil }, - set: { if !$0 { pickerType = nil } } - ), - allowedContentTypes: allowedContentTypes(), - allowsMultipleSelection: false - ) { [pickerType] result in - handleFileImportResult(pickerType, result) - } - .onAppear { - do { - try resourceManager.createDirectoriesIfNeeded() - } catch { - withAnimation { - messages.append(Message(type: .info, text: "Error creating content directories: \(error.localizedDescription)")) - } - } - } - } - .navigationViewStyle(StackNavigationViewStyle()) - } - - private func addSelectedImageMessage() { - if let selectedImage { - messages.append(Message(image: selectedImage)) - } - } - - private func generate() { - guard !prompt.isEmpty else { return } - isGenerating = true - shouldStopGenerating = false - shouldStopShowingToken = false - let text = prompt.trimmingCharacters(in: .whitespacesAndNewlines) - let seq_len = 768 // text: 256, vision: 768 - let modelPath = resourceManager.modelPath - let tokenizerPath = resourceManager.tokenizerPath - let modelType = ModelType.fromPath(modelPath) - - prompt = "" - hideKeyboard() - showingSettings = false - - messages.append(Message(text: text)) - messages.append(Message(type: modelType == .llama ? .llamagenerated : .llavagenerated)) - - runnerQueue.async { - defer { - DispatchQueue.main.async { - isGenerating = false - selectedImage = nil - } - } - - switch modelType { - case .llama, .qwen3, .phi4: - runnerHolder.llamaRunner = runnerHolder.llamaRunner ?? LLaMARunner(modelPath: modelPath, tokenizerPath: tokenizerPath) - case .llava: - runnerHolder.llavaRunner = runnerHolder.llavaRunner ?? LLaVARunner(modelPath: modelPath, tokenizerPath: tokenizerPath) - } - - guard !shouldStopGenerating else { return } - switch modelType { - case .llama, .qwen3, .phi4: - if let runner = runnerHolder.llamaRunner, !runner.isLoaded() { - var error: Error? - let startLoadTime = Date() - do { - try runner.load() - } catch let loadError { - error = loadError - } - - let loadTime = Date().timeIntervalSince(startLoadTime) - DispatchQueue.main.async { - withAnimation { - var message = messages.removeLast() - message.type = .info - if let error { - message.text = "Model loading failed: error \((error as NSError).code)" - } else { - message.text = "Model loaded in \(String(format: "%.2f", loadTime)) s" - } - messages.append(message) - if error == nil { - messages.append(Message(type: .llamagenerated)) - } - } - } - if error != nil { - return - } - } - case .llava: - if let runner = runnerHolder.llavaRunner, !runner.isLoaded() { - var error: Error? - let startLoadTime = Date() - do { - try runner.load() - } catch let loadError { - error = loadError - } - - let loadTime = Date().timeIntervalSince(startLoadTime) - DispatchQueue.main.async { - withAnimation { - var message = messages.removeLast() - message.type = .info - if let error { - message.text = "Model loading failed: error \((error as NSError).code)" - } else { - message.text = "Model loaded in \(String(format: "%.2f", loadTime)) s" - } - messages.append(message) - if error == nil { - messages.append(Message(type: .llavagenerated)) - } - } - } - if error != nil { - return - } - } - } - - guard !shouldStopGenerating else { - DispatchQueue.main.async { - withAnimation { - _ = messages.removeLast() - } - } - return - } - do { - var tokens: [String] = [] - var rgbArray: [UInt8]? - let MAX_WIDTH = 336.0 - var newHeight = 0.0 - var imageBuffer: UnsafeMutableRawPointer? - - if let img = selectedImage { - let llava_prompt = "\(text) ASSISTANT" - - newHeight = MAX_WIDTH * img.size.height / img.size.width - let resizedImage = img.resized(to: CGSize(width: MAX_WIDTH, height: newHeight)) - rgbArray = resizedImage.toRGBArray() - imageBuffer = UnsafeMutableRawPointer(mutating: rgbArray) - - try runnerHolder.llavaRunner?.generate(imageBuffer!, width: MAX_WIDTH, height: newHeight, prompt: llava_prompt, sequenceLength: seq_len) { token in - - if token != llava_prompt { - if token == "" { - shouldStopGenerating = true - runnerHolder.llavaRunner?.stop() - } else { - tokens.append(token) - if tokens.count > 2 { - let text = tokens.joined() - let count = tokens.count - tokens = [] - DispatchQueue.main.async { - var message = messages.removeLast() - message.text += text - message.tokenCount += count - message.dateUpdated = Date() - messages.append(message) - } - } - if shouldStopGenerating { - runnerHolder.llavaRunner?.stop() - } - } - } - } - } else { - let prompt: String - switch modelType { - case .qwen3: - let basePrompt = String(format: Constants.qwen3PromptTemplate, text) - // If thinking mode is enabled for Qwen, don't skip the special tokens - // and have them be generated. - prompt = thinkingMode ? basePrompt.replacingOccurrences(of: "\n\n\n\n\n", with: "") : basePrompt - case .llama: - prompt = String(format: Constants.llama3PromptTemplate, text) - case .llava: - prompt = String(format: Constants.llama3PromptTemplate, text) - case .phi4: - prompt = String(format: Constants.phi4PromptTemplate, text) - } - - try runnerHolder.llamaRunner?.generate(prompt, sequenceLength: seq_len) { token in - - if token != prompt { - if token == "<|eot_id|>" { - // hack to fix the issue that extension/llm/runner/text_token_generator.h - // keeps generating after <|eot_id|> - shouldStopShowingToken = true - } else if token == "<|im_end|>" { - // Qwen3 specific token. - // Skip. - } else if token == "" { - // Qwen3 specific token. - let textToFlush = tokens.joined() - let flushedTokenCount = tokens.count - tokens = [] - DispatchQueue.main.async { - var message = messages.removeLast() - message.text += textToFlush - message.text += message.text.isEmpty ? "Thinking...\n\n" : "\n\nThinking...\n\n" - message.tokenCount += flushedTokenCount + 1 // + 1 for the start thinking token. - message.dateUpdated = Date() - messages.append(message) - } - } else if token == "" { - // Qwen3 specific token. - let textToFlush = tokens.joined() - let flushedTokenCount = tokens.count - tokens = [] - DispatchQueue.main.async { - var message = messages.removeLast() - message.text += textToFlush - message.text += "\n\nFinished thinking.\n\n" - message.tokenCount += flushedTokenCount + 1 // + 1 for the end thinking token. - message.dateUpdated = Date() - messages.append(message) - } - } else { - tokens.append(token.trimmingCharacters(in: .newlines)) - // Flush tokens in groups of 3 so that it's closer to whole words being generated - // rather than parts of words (tokens). - if tokens.count > 2 { - let text = tokens.joined() - let count = tokens.count - tokens = [] - DispatchQueue.main.async { - var message = messages.removeLast() - message.text += text - message.tokenCount += count - message.dateUpdated = Date() - messages.append(message) - } - } - if shouldStopGenerating { - runnerHolder.llamaRunner?.stop() - } - } - } - } - } - } catch { - DispatchQueue.main.async { - withAnimation { - var message = messages.removeLast() - message.type = .info - message.text = "Text generation failed: error \((error as NSError).code)" - messages.append(message) - } - } - } - } - } - - private func stop() { - shouldStopGenerating = true - } - - private func allowedContentTypes() -> [UTType] { - guard let pickerType else { return [] } - switch pickerType { - case .model: - return [UTType(filenameExtension: "pte")].compactMap { $0 } - case .tokenizer: - return [UTType(filenameExtension: "bin"), UTType(filenameExtension: "model"), UTType(filenameExtension: "json"), ].compactMap { $0 } - } - } - - private func handleFileImportResult(_ pickerType: PickerType?, _ result: Result<[URL], Error>) { - switch result { - case .success(let urls): - guard let url = urls.first, let pickerType else { - withAnimation { - messages.append(Message(type: .info, text: "Failed to select a file")) - } - return - } - runnerQueue.async { - runnerHolder.llamaRunner = nil - runnerHolder.llavaRunner = nil - } - switch pickerType { - case .model: - resourceManager.modelPath = url.path - case .tokenizer: - resourceManager.tokenizerPath = url.path - } - if resourceManager.isModelValid && resourceManager.isTokenizerValid { - showingSettings = false - textFieldFocused = true - } - case .failure(let error): - withAnimation { - messages.append(Message(type: .info, text: "Failed to select a file: \(error.localizedDescription)")) - } - } - } -} - -extension View { - func hideKeyboard() { - UIApplication.shared.sendAction(#selector(UIResponder.resignFirstResponder), to: nil, from: nil, for: nil) - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ImagePicker.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ImagePicker.swift deleted file mode 100644 index 57d71f6686d..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ImagePicker.swift +++ /dev/null @@ -1,50 +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. - */ - -import SwiftUI -import UIKit - -struct ImagePicker: UIViewControllerRepresentable { - class Coordinator: NSObject, UINavigationControllerDelegate, UIImagePickerControllerDelegate { - let parent: ImagePicker - - init(parent: ImagePicker) { - self.parent = parent - } - - func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey: Any]) { - if let image = info[.originalImage] as? UIImage { - parent.selectedImage = image - } - - parent.presentationMode.wrappedValue.dismiss() - } - - func imagePickerControllerDidCancel(_ picker: UIImagePickerController) { - parent.selectedImage = nil - parent.presentationMode.wrappedValue.dismiss() - } - } - - @Environment(\.presentationMode) var presentationMode - @Binding var selectedImage: UIImage? - var sourceType: UIImagePickerController.SourceType = .photoLibrary - - func makeCoordinator() -> Coordinator { - Coordinator(parent: self) - } - - func makeUIViewController(context: Context) -> UIImagePickerController { - let picker = UIImagePickerController() - picker.delegate = context.coordinator - picker.sourceType = sourceType - return picker - } - - func updateUIViewController(_ uiViewController: UIImagePickerController, context: Context) {} -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogManager.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogManager.swift deleted file mode 100644 index 038e2807fed..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogManager.swift +++ /dev/null @@ -1,51 +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. - */ - -import SwiftUI - -import ExecuTorch - -struct LogEntry: Identifiable, Codable { - let id: UUID - let level: Int - let timestamp: TimeInterval - let filename: String - let line: UInt - let message: String -} - -class LogManager: ObservableObject, LogSink { - @AppStorage("logs") private var data = Data() - - @Published var logs: [LogEntry] = [] { - didSet { - data = (try? JSONEncoder().encode(logs)) ?? Data() - } - } - - init() { - logs = (try? JSONDecoder().decode([LogEntry].self, from: data)) ?? [] - Log.shared.add(sink: self) - } - - deinit { - Log.shared.remove(sink: self) - } - - func log(level: LogLevel, timestamp: TimeInterval, filename: String, line: UInt, message: String) { - let log = LogEntry(id: UUID(), level: level.rawValue, timestamp: timestamp, filename: filename, line: line, message: message) - - DispatchQueue.main.sync { - self.logs.append(log) - } - } - - func clear() { - logs.removeAll() - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogView.swift deleted file mode 100644 index 1dc6ebac2b3..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/LogView.swift +++ /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. - */ - -import SwiftUI - -import ExecuTorch - -struct LogView: View { - @ObservedObject var logManager: LogManager - - var body: some View { - ScrollView { - VStack(alignment: .leading) { - ForEach(logManager.logs) { log in - Text("\(format(timestamp: log.timestamp)) \(log.filename):\(log.line)") - .padding(.top) - .foregroundColor(.secondary) - .textSelection(.enabled) - Text(log.message) - .padding(.bottom) - .foregroundColor(color(for: log.level)) - .textSelection(.enabled) - } - } - } - .padding() - .defaultScrollAnchor(.bottom) - .navigationBarTitle("Logs", displayMode: .inline) - .navigationBarItems(trailing: - Button(action: { logManager.clear() }) { - Image(systemName: "trash") - } - ) - } - - private func format(timestamp: TimeInterval) -> String { - let totalSeconds = Int(timestamp) - let hours = (totalSeconds / 3600) % 24 - let minutes = (totalSeconds / 60) % 60 - let seconds = totalSeconds % 60 - let microseconds = Int((timestamp - Double(totalSeconds)) * 1000000) - return String(format: "%02d:%02d:%02d.%06d", hours, minutes, seconds, microseconds) - } - - private func color(for level: Int) -> Color { - switch LogLevel(rawValue: level) { - case .debug: - return .blue - case .info: - return .primary - case .error: - return .red - case .fatal: - return .purple - default: - return .secondary - } - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Message.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Message.swift deleted file mode 100644 index 34ed0d7e933..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/Message.swift +++ /dev/null @@ -1,26 +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. - */ - -import UIKit - -enum MessageType { - case prompted - case llamagenerated // TODO: change this to to something more general, like "textgenerated". - case llavagenerated - case info -} - -struct Message: Identifiable, Equatable { - let id = UUID() - let dateCreated = Date() - var dateUpdated = Date() - var type: MessageType = .prompted - var text = "" - var tokenCount = 0 - var image: UIImage? -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageListView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageListView.swift deleted file mode 100644 index bc0e459123f..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageListView.swift +++ /dev/null @@ -1,88 +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. - */ - -import SwiftUI - -struct MessageListView: View { - @Binding var messages: [Message] - @State private var showScrollToBottomButton = false - @State private var userHasScrolled = false - @State private var keyboardHeight: CGFloat = 0 - - var body: some View { - ScrollViewReader { value in - ScrollView { - VStack { - ForEach(messages) { message in - MessageView(message: message) - .padding([.leading, .trailing], 20) - } - GeometryReader { geometry -> Color in - DispatchQueue.main.async { - let maxY = geometry.frame(in: .global).maxY - let screenHeight = UIScreen.main.bounds.height - keyboardHeight - let isBeyondBounds = maxY > screenHeight - 50 - if showScrollToBottomButton != isBeyondBounds { - showScrollToBottomButton = isBeyondBounds - userHasScrolled = isBeyondBounds - } - } - return Color.clear - } - .frame(height: 0) - } - } - .onChange(of: messages) { - if !userHasScrolled, let lastMessageId = messages.last?.id { - withAnimation { - value.scrollTo(lastMessageId, anchor: .bottom) - } - } - } - .overlay( - Group { - if showScrollToBottomButton { - Button(action: { - withAnimation { - if let lastMessageId = messages.last?.id { - value.scrollTo(lastMessageId, anchor: .bottom) - } - userHasScrolled = false - } - }) { - ZStack { - Circle() - .fill(Color(UIColor.secondarySystemBackground).opacity(0.9)) - .frame(height: 28) - Image(systemName: "arrow.down.circle") - .resizable() - .aspectRatio(contentMode: .fit) - .frame(height: 28) - } - } - .transition(AnyTransition.opacity.animation(.easeInOut(duration: 0.2))) - } - }, - alignment: .bottom - ) - } - .onAppear { - NotificationCenter.default.addObserver(forName: UIResponder.keyboardWillShowNotification, object: nil, queue: .main) { notification in - let keyboardFrame = notification.userInfo?[UIResponder.keyboardFrameEndUserInfoKey] as? CGRect ?? .zero - keyboardHeight = keyboardFrame.height - 40 - } - NotificationCenter.default.addObserver(forName: UIResponder.keyboardWillHideNotification, object: nil, queue: .main) { _ in - keyboardHeight = 0 - } - } - .onDisappear { - NotificationCenter.default.removeObserver(self, name: UIResponder.keyboardWillShowNotification, object: nil) - NotificationCenter.default.removeObserver(self, name: UIResponder.keyboardWillHideNotification, object: nil) - } - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageView.swift deleted file mode 100644 index 542a88377b7..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/MessageView.swift +++ /dev/null @@ -1,73 +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. - */ - -import SwiftUI - -struct MessageView: View { - let message: Message - - var body: some View { - VStack(alignment: .center) { - if message.type == .info { - Text(message.text) - .font(.caption) - .foregroundColor(.secondary) - .padding([.leading, .trailing], 10) - } else { - VStack(alignment: message.type == .llamagenerated || message.type == .llavagenerated ? .leading : .trailing) { - if message.type == .llamagenerated || message.type == .llavagenerated || message.type == .prompted { - Text(message.type == .llamagenerated ? "Llama" : (message.type == .llavagenerated ? "Llava" : "Prompt")) - .font(.caption) - .foregroundColor(.secondary) - .padding(message.type == .llamagenerated || message.type == .llavagenerated ? .trailing : .leading, 20) - } - HStack { - if message.type != .llamagenerated && message.type != .llavagenerated { Spacer() } - if message.text.isEmpty { - if let img = message.image { - Image(uiImage: img) - .resizable() - .scaledToFit() - .frame(maxWidth: 200, maxHeight: 200) - .padding() - .background(Color.gray.opacity(0.2)) - .cornerRadius(8) - .padding(.vertical, 2) - } else { - ProgressView() - .progressViewStyle(CircularProgressViewStyle()) - } - } else { - Text(message.text) - .padding(10) - .foregroundColor(message.type == .llamagenerated || message.type == .llavagenerated ? .primary : .white) - .background(message.type == .llamagenerated || message.type == .llavagenerated ? Color(UIColor.secondarySystemBackground) : Color.blue) - .cornerRadius(20) - .contextMenu { - Button(action: { - UIPasteboard.general.string = message.text - }) { - Text("Copy") - Image(systemName: "doc.on.doc") - } - } - } - if message.type == .llamagenerated || message.type == .llavagenerated { Spacer() } - } - let elapsedTime = message.dateUpdated.timeIntervalSince(message.dateCreated) - if elapsedTime > 0 && message.type != .info { - Text(String(format: "%.1f t/s", Double(message.tokenCount) / elapsedTime)) - .font(.caption) - .foregroundColor(.secondary) - .padding(message.type == .llamagenerated || message.type == .llavagenerated ? .trailing : .leading, 20) - } - }.padding([.leading, .trailing], message.type == .info ? 0 : 10) - } - } - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceManager.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceManager.swift deleted file mode 100644 index 7d3be7975de..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceManager.swift +++ /dev/null @@ -1,37 +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. - */ - -import SwiftUI - -final class ResourceManager: ObservableObject { - @AppStorage("modelPath") var modelPath = "" - @AppStorage("tokenizerPath") var tokenizerPath = "" - private let fileManager = FileManager.default - - var isModelValid: Bool { - fileManager.fileExists(atPath: modelPath) - } - - var isTokenizerValid: Bool { - fileManager.fileExists(atPath: tokenizerPath) - } - - var modelName: String { - URL(fileURLWithPath: modelPath).deletingPathExtension().lastPathComponent - } - - var tokenizerName: String { - URL(fileURLWithPath: tokenizerPath).deletingPathExtension().lastPathComponent - } - - func createDirectoriesIfNeeded() throws { - guard let documentsDirectory = fileManager.urls(for: .documentDirectory, in: .userDomainMask).first else { return } - try fileManager.createDirectory(at: documentsDirectory.appendingPathComponent("models"), withIntermediateDirectories: true, attributes: nil) - try fileManager.createDirectory(at: documentsDirectory.appendingPathComponent("tokenizers"), withIntermediateDirectories: true, attributes: nil) - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceMonitor.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceMonitor.swift deleted file mode 100644 index 3ec16463e8a..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ResourceMonitor.swift +++ /dev/null @@ -1,51 +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. - */ - -import Foundation - -final class ResourceMonitor: ObservableObject { - @Published var usedMemory = 0 - @Published var availableMemory = 0 - private var memoryUpdateTimer: Timer? - - deinit { - stop() - } - - public func start() { - memoryUpdateTimer = Timer.scheduledTimer(withTimeInterval: 1, repeats: true) { [weak self] _ in - self?.updateMemoryUsage() - } - } - - public func stop() { - memoryUpdateTimer?.invalidate() - } - - private func updateMemoryUsage() { - usedMemory = usedMemoryInMB() - availableMemory = availableMemoryInMB() - } - - private func usedMemoryInMB() -> Int { - var info = task_vm_info_data_t() - var count = mach_msg_type_number_t(MemoryLayout.size) / 4 - - let kerr: kern_return_t = withUnsafeMutablePointer(to: &info) { - $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) { - task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count) - } - } - guard kerr == KERN_SUCCESS else { return 0 } - return Int(info.phys_footprint / 0x100000) - } - - private func availableMemoryInMB() -> Int { - return Int(os_proc_available_memory() / 0x100000) - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Debug.xcconfig b/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Debug.xcconfig deleted file mode 100644 index 0c749c04c1a..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Debug.xcconfig +++ /dev/null @@ -1,25 +0,0 @@ -ET_PLATFORM[sdk=iphonesimulator*] = simulator -ET_PLATFORM[sdk=iphoneos*] = ios -ET_PLATFORM[sdk=macos*] = macos - -OTHER_LDFLAGS = $(inherited) \ - -force_load $(BUILT_PRODUCTS_DIR)/libexecutorch_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_coreml_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_mps_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_xnnpack_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_llm_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_optimized_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_quantized_debug_$(ET_PLATFORM).a \ - @$(TEMP_DIR)/cmake/linker_flags - -// LLaMARunner requires additional dependencies built with CMake in a custom run script phase. -// Include headers and libraries from $(TEMP_DIR)/cmake for it. -HEADER_SEARCH_PATHS = $(inherited) \ - $(SRCROOT)/../../../../.. \ - $(TEMP_DIR)/cmake/include \ - $(SRCROOT)/../../../../extension/llm/tokenizers/include \ - $(SRCROOT)/../../../../extension/llm/tokenizers/third-party/sentencepiece \ - $(SRCROOT)/../../../../extension/llm/tokenizers/third-party/sentencepiece/src - -LIBRARY_SEARCH_PATHS = $(inherited) \ - $(TEMP_DIR)/cmake/lib diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/LLaMA-Info.plist b/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/LLaMA-Info.plist deleted file mode 100644 index ff579a6caff..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/LLaMA-Info.plist +++ /dev/null @@ -1,8 +0,0 @@ - - - - - UIFileSharingEnabled - - - diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Release.xcconfig b/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Release.xcconfig deleted file mode 100644 index 899d133127f..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Release.xcconfig +++ /dev/null @@ -1,27 +0,0 @@ -ET_PLATFORM[sdk=iphonesimulator*] = simulator -ET_PLATFORM[sdk=iphoneos*] = ios -ET_PLATFORM[sdk=macos*] = macos - -// Link the Debug version of ExecuTorch runtime to keep the logs. -// Switch to Release for better performance if logs are not needed. -OTHER_LDFLAGS = $(inherited) \ - -force_load $(BUILT_PRODUCTS_DIR)/libexecutorch_debug_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_coreml_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_mps_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libbackend_xnnpack_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_llm_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_optimized_$(ET_PLATFORM).a \ - -force_load $(BUILT_PRODUCTS_DIR)/libkernels_quantized_$(ET_PLATFORM).a \ - @$(TEMP_DIR)/cmake/linker_flags - -// LLaMARunner requires additional dependencies built with CMake in a custom run script phase. -// Include headers and libraries from $(TEMP_DIR)/cmake for it. -HEADER_SEARCH_PATHS = $(inherited) \ - $(SRCROOT)/../../../../.. \ - $(TEMP_DIR)/cmake/include \ - $(SRCROOT)/../../../../extension/llm/tokenizers/include \ - $(SRCROOT)/../../../../extension/llm/tokenizers/third-party/sentencepiece \ - $(SRCROOT)/../../../../extension/llm/tokenizers/third-party/sentencepiece/src - -LIBRARY_SEARCH_PATHS = $(inherited) \ - $(TEMP_DIR)/cmake/lib diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/Contents.json b/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/Contents.json deleted file mode 100644 index f4344003c80..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/Contents.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "images" : [ - { - "filename" : "logo.png", - "idiom" : "universal", - "platform" : "ios", - "size" : "1024x1024" - } - ], - "info" : { - "author" : "xcode", - "version" : 1 - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/logo.png b/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/logo.png deleted file mode 100644 index 60e3e5174e9..00000000000 Binary files a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/AppIcon.appiconset/logo.png and /dev/null differ diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/Contents.json b/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/Contents.json deleted file mode 100644 index 73c00596a7f..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMAAssets/Assets.xcassets/Contents.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "info" : { - "author" : "xcode", - "version" : 1 - } -} diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAEntitlements/LLaMA.entitlements b/examples/demo-apps/apple_ios/LLaMA/LLaMAEntitlements/LLaMA.entitlements deleted file mode 100644 index 99f471672d6..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMAEntitlements/LLaMA.entitlements +++ /dev/null @@ -1,8 +0,0 @@ - - - - - com.apple.developer.kernel.increased-memory-limit - - - diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h deleted file mode 100644 index 1b5ada5138e..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h +++ /dev/null @@ -1,30 +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. - */ - -#import - -NS_ASSUME_NONNULL_BEGIN - -@interface LLaMARunner : NSObject - -- (instancetype)initWithModelPath:(NSString *)modelPath - tokenizerPath:(NSString *)tokenizerPath; -- (BOOL)isLoaded; -- (BOOL)loadWithError:(NSError **)error; -- (BOOL)generate:(NSString *)prompt - sequenceLength:(NSInteger)seq_len - withTokenCallback:(nullable void (^)(NSString *))callback - error:(NSError **)error; -- (void)stop; - -+ (instancetype)new NS_UNAVAILABLE; -- (instancetype)init NS_UNAVAILABLE; - -@end - -NS_ASSUME_NONNULL_END diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm deleted file mode 100644 index a4725246ca0..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm +++ /dev/null @@ -1,94 +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. - */ - -#import "LLaMARunner.h" - -#import -#if BUILD_WITH_XCODE -#import "ExecuTorchTextLLMRunner.h" -#else -#import -#endif -#import - -@interface LLaMARunner () -@end - -@implementation LLaMARunner { - ExecuTorchTextLLMRunner *_runner; -} - -- (instancetype)initWithModelPath:(NSString *)modelPath - tokenizerPath:(NSString *)tokenizerPath { - self = [super init]; - if (self) { - [ExecuTorchLog.sharedLog addSink:self]; - auto tokens = example::get_special_tokens(example::Version::Default); - NSMutableArray *specialTokens = [[NSMutableArray alloc] initWithCapacity:tokens->size()]; - for (const auto &token : *tokens) { - [specialTokens addObject:(NSString *)@(token.c_str())]; - } - _runner = [[ExecuTorchTextLLMRunner alloc] initWithModelPath:modelPath - tokenizerPath:tokenizerPath - specialTokens:specialTokens]; - } - return self; -} - -- (void)dealloc { - [ExecuTorchLog.sharedLog removeSink:self]; -} - -- (BOOL)isLoaded { - return [_runner isLoaded]; -} - -- (BOOL)loadWithError:(NSError**)error { - return [_runner loadWithError:error]; -} - -- (BOOL)generate:(NSString *)prompt - sequenceLength:(NSInteger)seq_len - withTokenCallback:(nullable void (^)(NSString *))callback - error:(NSError **)error { - return [_runner generate:prompt - sequenceLength:seq_len - withTokenCallback:callback - error:error]; -} - -- (void)stop { - [_runner stop]; -} - -#pragma mark - ExecuTorchLogSink - -- (void)logWithLevel:(ExecuTorchLogLevel)level - timestamp:(NSTimeInterval)timestamp - filename:(NSString*)filename - line:(NSUInteger)line - message:(NSString*)message { - NSUInteger totalSeconds = (NSUInteger)timestamp; - NSUInteger hours = (totalSeconds / 3600) % 24; - NSUInteger minutes = (totalSeconds / 60) % 60; - NSUInteger seconds = totalSeconds % 60; - NSUInteger microseconds = (timestamp - totalSeconds) * 1000000; - NSLog( - @"%c %02lu:%02lu:%02lu.%06lu executorch:%s:%zu] %s", - (char)level, - hours, - minutes, - seconds, - microseconds, - filename.UTF8String, - line, - message.UTF8String - ); -} - -@end diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.h b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.h deleted file mode 100644 index 4214a9ad390..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.h +++ /dev/null @@ -1,35 +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. - */ - -#import - -NS_ASSUME_NONNULL_BEGIN - -FOUNDATION_EXPORT NSErrorDomain const LLaVARunnerErrorDomain; - -@interface LLaVARunner : NSObject - -- (instancetype)initWithModelPath:(NSString*)filePath - tokenizerPath:(NSString*)tokenizerPath; -- (BOOL)isLoaded; -- (BOOL)loadWithError:(NSError**)error; -- (BOOL)generate:(void*)imageBuffer - width:(CGFloat)width - height:(CGFloat)height - prompt:(NSString*)prompt - sequenceLength:(NSInteger)seq_len - withTokenCallback:(nullable void (^)(NSString*))callback - error:(NSError**)error; -- (void)stop; - -+ (instancetype)new NS_UNAVAILABLE; -- (instancetype)init NS_UNAVAILABLE; - -@end - -NS_ASSUME_NONNULL_END diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.mm b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.mm deleted file mode 100644 index 8864e45f52b..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaVARunner.mm +++ /dev/null @@ -1,115 +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. - */ - -#import "LLaVARunner.h" - -#import -#import - -using executorch::extension::llm::Image; -using executorch::runtime::Error; - -NSErrorDomain const LLaVARunnerErrorDomain = @"LLaVARunnerErrorDomain"; - -@interface LLaVARunner () -@end - -@implementation LLaVARunner { - std::unique_ptr _runner; -} - -- (instancetype)initWithModelPath:(NSString*)modelPath - tokenizerPath:(NSString*)tokenizerPath { - self = [super init]; - if (self) { - [ExecuTorchLog.sharedLog addSink:self]; - _runner = std::make_unique( - modelPath.UTF8String, tokenizerPath.UTF8String); - } - return self; -} - -- (void)dealloc { - [ExecuTorchLog.sharedLog removeSink:self]; -} - -- (BOOL)isLoaded { - return _runner->is_loaded(); -} - -- (BOOL)loadWithError:(NSError**)error { - const auto status = _runner->load(); - if (status != Error::Ok) { - if (error) { - *error = [NSError errorWithDomain:LLaVARunnerErrorDomain - code:(NSInteger)status - userInfo:nil]; - } - return NO; - } - return YES; -} - -- (BOOL)generate:(void*)imageBuffer - width:(CGFloat)width - height:(CGFloat)height - prompt:(NSString*)prompt - sequenceLength:(NSInteger)seq_len - withTokenCallback:(nullable void (^)(NSString*))callback - error:(NSError**)error { - const auto* data = static_cast(imageBuffer); - const auto status = _runner->generate( - {Image{ - std::vector( - data, data + (int32_t)width * (int32_t)height * 3), - (int32_t)width, - (int32_t)height, - 3}}, - prompt.UTF8String, - seq_len, - [callback](const std::string& token) { callback(@(token.c_str())); }); - if (status != Error::Ok) { - if (error) { - *error = [NSError errorWithDomain:LLaVARunnerErrorDomain - code:(NSInteger)status - userInfo:nil]; - return NO; - } - } - return YES; -} - -- (void)stop { - _runner->stop(); -} - -#pragma mark - ExecuTorchLogSink - -- (void)logWithLevel:(ExecuTorchLogLevel)level - timestamp:(NSTimeInterval)timestamp - filename:(NSString*)filename - line:(NSUInteger)line - message:(NSString*)message { - NSUInteger totalSeconds = (NSUInteger)timestamp; - NSUInteger hours = (totalSeconds / 3600) % 24; - NSUInteger minutes = (totalSeconds / 60) % 60; - NSUInteger seconds = totalSeconds % 60; - NSUInteger microseconds = (timestamp - totalSeconds) * 1000000; - NSLog( - @"%c %02lu:%02lu:%02lu.%06lu executorch:%s:%zu] %s", - (char)level, - hours, - minutes, - seconds, - microseconds, - filename.UTF8String, - line, - message.UTF8String); -} - -@end diff --git a/examples/demo-apps/apple_ios/LLaMA/README.md b/examples/demo-apps/apple_ios/LLaMA/README.md deleted file mode 100644 index 70775d63a68..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/README.md +++ /dev/null @@ -1,47 +0,0 @@ -# ExecuTorch Llama iOS Demo App - -Get hands-on with running LLaMA and LLaVA models — exported via ExecuTorch — natively on your iOS device! - -*Click the image below to see it in action!* - -

- - iOS app running a LlaMA model - -

- -## Requirements -- [Xcode](https://apps.apple.com/us/app/xcode/id497799835?mt=12/) 15.0 or later -- [Cmake](https://cmake.org/download/) 3.19 or later - - Download and open the macOS `.dmg` installer and move the Cmake app to `/Applications` folder. - - Install Cmake command line tools: `sudo /Applications/CMake.app/Contents/bin/cmake-gui --install` -- A development provisioning profile with the [`increased-memory-limit`](https://developer.apple.com/documentation/bundleresources/entitlements/com_apple_developer_kernel_increased-memory-limit) entitlement. - -## Models - -Download already exported LLaMA/LLaVA models along with tokenizers from [HuggingFace](https://huggingface.co/executorch-community) or export your own empowered by [XNNPACK](docs/delegates/xnnpack_README.md) or [MPS](docs/delegates/mps_README.md) backends. - -## Build and Run - -1. Make sure git submodules are up-to-date: - ```bash - git submodule update --init --recursive - ``` - -2. Open the Xcode project: - ```bash - open examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj - ``` - -3. Click the Play button to launch the app in the Simulator. - -4. To run on a device, ensure you have it set up for development and a provisioning profile with the `increased-memory-limit` entitlement. Update the app's bundle identifier to match your provisioning profile with the required capability. - -5. After successfully launching the app, copy the exported ExecuTorch model (`.pte`) and tokenizer (`.model`) files to the iLLaMA folder. Four models are currently supported at the moment - Llama, Qwen3, Phi4-mini, and Llava multimodal. Please ensure that your model `.pte` file starts with `llama`, `qwen3`, `phi4` or `llava` so that the app selects the correct model type. - - - **For the Simulator:** Drag and drop both files onto the Simulator window and save them in the `On My iPhone > iLLaMA` folder. - - **For a Device:** Open a separate Finder window, navigate to the Files tab, drag and drop both files into the iLLaMA folder, and wait for the copying to finish. - -6. Follow the app's UI guidelines to select the model and tokenizer files from the local filesystem and issue a prompt. - -For more details check out the [Using ExecuTorch on iOS](../../../../docs/source/using-executorch-ios.md) page. diff --git a/examples/demo-apps/apple_ios/LLaMA/docs/delegates/mps_README.md b/examples/demo-apps/apple_ios/LLaMA/docs/delegates/mps_README.md deleted file mode 100644 index d6bccc0ef47..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/docs/delegates/mps_README.md +++ /dev/null @@ -1,106 +0,0 @@ -# Building Llama iOS Demo for MPS Backend - -This tutorial covers the end to end workflow for building an iOS demo app using MPS backend on device. -More specifically, it covers: -1. Export and quantization of Llama models against the MPS backend. -2. Building and linking libraries that are required to inference on-device for iOS platform using MPS. -3. Building the iOS demo app itself. - -## Prerequisites -* [Xcode 15](https://developer.apple.com/xcode) -* [iOS 18 SDK](https://developer.apple.com/ios) -* Set up your ExecuTorch repo and environment if you haven’t done so by following the [Setting up ExecuTorch](https://pytorch.org/executorch/main/using-executorch-building-from-source) to set up the repo and dev environment: - -## Setup ExecuTorch -In this section, we will need to set up the ExecuTorch repo first with Conda environment management. Make sure you have Conda available in your system (or follow the instructions to install it [here](https://conda.io/projects/conda/en/latest/user-guide/install/index.html)). The commands below are running on Linux (CentOS). - -Checkout ExecuTorch repo and sync submodules - -``` -git clone -b viable/strict https://github.com/pytorch/executorch.git && cd executorch -``` - -Create either a Python virtual environment: - -``` -python3 -m venv .venv && source .venv/bin/activate && pip install --upgrade pip -``` - -Or a Conda environment - -``` -conda create -n et_mps python=3.10.0 && conda activate et_mps -``` - -Install dependencies - -``` -./install_executorch.sh -``` - -## Prepare Models -In this demo app, we support text-only inference with Llama 3.1, Llama 3, and Llama 2 models. - -Install the required packages to export the model - -``` -./examples/models/llama/install_requirements.sh -``` - -Export the model -``` -python -m extension.llm.export.export_llm base.checkpoint="${MODEL_DIR}/consolidated.00.pth" base.params="${MODEL_DIR}/params.json" model.use_kv_cache=True model.use_sdpa_with_kv_cache=True backend.mps.enabled=True model.dtype_override="fp32" model.enable_dynamic_shape=False quantization.qmode="8da4w" quantization.group_size=32 -``` - -## Pushing Model and Tokenizer - -### Copy the model to Simulator -* Drag&drop the model and tokenizer files onto the Simulator window and save them somewhere inside the iLLaMA folder. -* Pick the files in the app dialog, type a prompt and click the arrow-up button. - -### Copy the model to Device -* Wire-connect the device and open the contents in Finder. -* Navigate to the Files tab and drag & drop the model and tokenizer files onto the iLLaMA folder. -* Wait until the files are copied. - -## Configure the XCode Project - -### Install CMake -Download and open the macOS .dmg installer at https://cmake.org/download and move the Cmake app to /Applications folder. -Install Cmake command line tools: - -``` -sudo /Applications/CMake.app/Contents/bin/cmake-gui --install -``` - - -### Swift Package Manager -The prebuilt ExecuTorch runtime, backend, and kernels are available as a Swift PM package. - -### Xcode -Open the project in Xcode.In Xcode, go to `File > Add Package Dependencies`. Paste the URL of the ExecuTorch repo into the search bar and select it. Make sure to change the branch name to the desired ExecuTorch version, e.g., “swiftpm-0.6.0”, or a branch name in format "swiftpm-." (e.g. "swiftpm-0.7.0-20250401") for a nightly build on a specific date. - -Link your binary with the ExecuTorch runtime and any backends or kernels used by the exported ML model. It is recommended to link the core runtime to the components that use ExecuTorch directly, and link kernels and backends against the main app target. - -Note: To access logs, link against the Debug build of the ExecuTorch runtime, i.e., the executorch_debug framework. For optimal performance, always link against the Release version of the deliverables (those without the _debug suffix), which have all logging overhead removed. - -For more details integrating and Running ExecuTorch on Apple Platforms, checkout this [link](https://pytorch.org/executorch/main/using-executorch-ios). - -

-iOS LLaMA App Swift PM -

- -Then select which ExecuTorch framework should link against which target. - -

-iOS LLaMA App Choosing package -

- -Click “Run” to build the app and run in on your iPhone. If the app successfully run on your device, you should see something like below: - -

-iOS LLaMA App mps -

- -## Reporting Issues -If you encountered any bugs or issues following this tutorial please file a bug/issue here on [Github](https://github.com/pytorch/executorch/issues/new). diff --git a/examples/demo-apps/apple_ios/LLaMA/docs/delegates/xnnpack_README.md b/examples/demo-apps/apple_ios/LLaMA/docs/delegates/xnnpack_README.md deleted file mode 100644 index 7a56d217b82..00000000000 --- a/examples/demo-apps/apple_ios/LLaMA/docs/delegates/xnnpack_README.md +++ /dev/null @@ -1,212 +0,0 @@ -# Building Llama iOS Demo for XNNPACK Backend - -This tutorial covers the end to end workflow for building an iOS demo app using XNNPACK backend on device. -More specifically, it covers: -1. Export and quantization of Llama models against the XNNPACK backend. -2. Building and linking libraries that are required to inference on-device for iOS platform using XNNPACK. -3. Building the iOS demo app itself. - -## Prerequisites -* [Xcode 15](https://developer.apple.com/xcode) -* [iOS 17 SDK](https://developer.apple.com/ios) - -## Setup ExecuTorch -In this section, we will need to set up the ExecuTorch repo first with Conda environment management. Make sure you have Conda available in your system (or follow the instructions to install it [here](https://conda.io/projects/conda/en/latest/user-guide/install/index.html)). The commands below are running on Linux (CentOS). - -Checkout ExecuTorch repo and sync submodules - -``` -git clone -b viable/strict https://github.com/pytorch/executorch.git && cd executorch -``` - -Create either a Python virtual environment: - -``` -python3 -m venv .venv && source .venv/bin/activate && pip install --upgrade pip -``` - -Or a Conda environment: - -``` -conda create -n et_xnnpack python=3.10.0 && conda activate et_xnnpack -``` - -Install dependencies - -``` -./install_executorch.sh -``` - -## Prepare Models -In this demo app, we support text-only inference with up-to-date Llama models and image reasoning inference with LLaVA 1.5. -* You can request and download model weights for Llama through Meta official [website](https://llama.meta.com/). -* For chat use-cases, download the instruct models instead of pretrained. -* Install the required packages to export the model: - -``` -./examples/models/llama/install_requirements.sh -``` - -### For Llama 3.2 1B and 3B SpinQuant models -Meta has released prequantized INT4 SpinQuant Llama 3.2 models that ExecuTorch supports on the XNNPACK backend. -* Export Llama model and generate .pte file as below: -``` -python -m extension.llm.export.export_llm base.model_class="llama3_2" base.checkpoint= base.params= model.use_kv_cache=True model.use_sdpa_with_kv_cache=True backend.xnnpack.enabled=True model.dtype_override="fp32" backend.xnnpack.extended_ops=True base.preq_mode="preq_8da4w_out_8da8w" base.preq_group_size=32 export.max_seq_length=2048 export.max_context_length=2048 base.preq_embedding_quantize=\'8,0\' quantization.use_spin_quant="native" base.metadata='"{\"get_bos_id\":128000, \"get_eos_ids\":[128009, 128001]}"' export.output_name="llama3_2_spinquant.pte" -``` -For convenience, an [exported ExecuTorch SpinQuant model](https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/blob/main/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8.pte) is available on Hugging Face. The export was created using [this detailed recipe notebook](https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-SpinQuant_INT4_EO8-ET/blob/main/Export_Recipe_Llama_3_2_1B_Instruct_SpinQuant_INT4_EO8.ipynb). - -### For Llama 3.2 1B and 3B QAT+LoRA models -Meta has released prequantized INT4 QAT+LoRA Llama 3.2 models that ExecuTorch supports on the XNNPACK backend. -* Export Llama model and generate .pte file as below: -``` -python -m extension.llm.export.export_llm base.model_class="llama3_2" base.checkpoint= base.params= quantization.use_qat=True base.use_lora=16 model.use_kv_cache=True model.use_sdpa_with_kv_cache=True backend.xnnpack.enabled=True model.dtype_override="fp32" backend.xnnpack.extended_ops=True base.preq_mode="preq_8da4w_out_8da8w" base.preq_group_size=32 export.max_seq_length=2048 export.max_context_length=2048 base.preq_embedding_quantize=\'8,0\' base.metadata='"{\"get_bos_id\":128000, \"get_eos_ids\":[128009, 128001]}"' export.output_name="llama3_2_qat_lora.pte" -``` -For convenience, an [exported ExecuTorch QAT+LoRA model](https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-QLORA_INT4_EO8-ET/blob/main/Llama-3.2-1B-Instruct-QLORA_INT4_EO8.pte) is available on Hugging Face. The export was created using [this detailed recipe notebook](https://huggingface.co/executorch-community/Llama-3.2-1B-Instruct-QLORA_INT4_EO8-ET/blob/main/Export_Recipe_Llama_3_2_1B_Instruct_QLORA_INT4_EO8.ipynb). - -### For Llama 3.2 1B and 3B BF16 models -We have supported BF16 as a data type on the XNNPACK backend for Llama 3.2 1B/3B models. -* The 1B model in BF16 format can run on mobile devices with 8GB RAM (iPhone 15 Pro and later). The 3B model will require 12GB+ RAM and hence will not fit on 8GB RAM phones. -* Export Llama model and generate .pte file as below: - -``` -python -m extension.llm.export.export_llm base.model_class="llama3_2" base.checkpoint= base.params= model.use_kv_cache=True model.use_sdpa_with_kv_cache=True backend.xnnpack.enabled=True model.dtype_override="bf16" base.metadata='"{\"get_bos_id\":128000, \"get_eos_ids\":[128009, 128001]}"' export.output_name="llama3_2_bf16.pte" -``` -For convenience, an [exported ExecuTorch bf16 model](https://huggingface.co/executorch-community/Llama-3.2-1B-ET/blob/main/llama3_2-1B.pte) is available on Hugging Face. The export was created using [this detailed recipe notebook](https://huggingface.co/executorch-community/Llama-3.2-1B-ET/blob/main/ExportRecipe_1B.ipynb). - -For more detail using Llama 3.2 lightweight models including prompt template, please go to our official [website](https://www.llama.com/docs/model-cards-and-prompt-formats/llama3_2/#-llama-3.2-lightweight-models-(1b/3b)-). - -### For Llama 3.1 and Llama 2 models - -Export the model -``` -python -m extension.llm.export.export_llm base.checkpoint= base.params= model.use_kv_cache=True model.use_sdpa_with_kv_cache=True backend.xnnpack.enabled=True quantization.qmode="8da4w" quantization.group_size=128 model.dtype_override="fp32" base.metadata='"{\"get_bos_id\":128000, \"get_eos_ids\":[128009, 128001]}"' quantization.embedding_quantize=\'4,32\' export.output_name="llama3_kv_sdpa_xnn_qe_4_32.pte" -``` - -### For LLaVA model -* For the Llava 1.5 model, you can get it from Huggingface [here](https://huggingface.co/llava-hf/llava-1.5-7b-hf). -* Run `examples/models/llava/install_requirements.sh` to install dependencies. -* Run the following command to generate llava.pte, tokenizer.bin and download an image basketball.jpg. - -``` -python -m executorch.examples.models.llava.export_llava --pte-name llava.pte --with-artifacts -``` -* You can find more information [here](https://github.com/pytorch/executorch/tree/main/examples/models/llava). - - -## Configure the XCode Project - -### 1. Install CMake -Download and open the macOS .dmg installer at https://cmake.org/download and move the Cmake app to /Applications folder. -Install Cmake command line tools: - -``` -sudo /Applications/CMake.app/Contents/bin/cmake-gui --install -``` - -### 2. Add ExecuTorch Runtime Package - -There are two options to add ExecuTorch runtime package into your XCode project: - -- [Recommended] Prebuilt package (via Swift Package Manager) -- Manually build the package locally and link them - - -### 2.1 [Recommended] Prebuilt package (via Swift Package Manager) - -The current XCode project is pre-configured to automatically download and link the latest prebuilt package via Swift Package Manager. - -#### (Optional) Changing the prebuilt package version -While we recommended using the latest prebuilt package pre-configured with the XCode project, you can also change the package version manually to your desired version. - -Go to Project Navigator, click on LLaMA. `Project --> LLaMA --> Package Dependencies`, and update the package dependencies to any of the available options below: - -- Branch --> swiftpm-0.7.0.20250401 (amend to match the latest nightly build) -- Branch --> swiftpm-0.6.0 - -### 2.2 Manually build the package locally and link them - -Note: You should only use this step if the prebuilt package doesn't work for your usecase (For example, you require the latest PRs from main, where there are no pre-built package yet) - -If you need to manually build the package, run the following command in your terminal -``` -# Install a compatible version of Buck2 -BUCK2_RELEASE_DATE="2024-12-16" -BUCK2_ARCHIVE="buck2-aarch64-apple-darwin.zst" -BUCK2=".venv/bin/buck2" - -curl -LO "https://github.com/facebook/buck2/releases/download/${BUCK2_RELEASE_DATE}/${BUCK2_ARCHIVE}" -zstd -cdq "$BUCK2_ARCHIVE" > "$BUCK2" && chmod +x "$BUCK2" -rm "$BUCK2_ARCHIVE" - -./scripts/build_apple_frameworks.sh -``` - - After the build finishes successfully, the resulting frameworks can be found in the `cmake-out` directory. Copy them to your project and link them against your targets. - -The following packages should be linked in your app target `LLaMA` (left side, LLaMA --> General --> select LLaMA under "TARGETS" --> scroll down to "Frameworks, Libraries, and Embedded Content") -- backend_coreml -- backend_mps -- backend_xnnpack -- kernels_llm -- kernels_optimized -- kernels_portable -- kernels_quantized - -The following package should be linked in your target app `LLaMARunner` (left side, LLaMA --> General --> select LLaMARunner under "TARGETS" --> scroll down to "Frameworks and Libraries") -- executorch - -

-iOS LLaMA App Choosing package -

- -If you cannot add the package into your app target (it's greyed out), it might have been linked before. You can verify it by checking it from your target app `(LLaMA / LLaMARunner) --> Build Phases --> Link Binary With Libraries`. - - - - More details on integrating and Running ExecuTorch on Apple Platforms, check out the detailed guide [here](https://pytorch.org/executorch/main/using-executorch-ios#local-build). - -### 3. Configure Build Schemes - -The project has two build configurations: -- Debug -- [Recommended] Release - -Navigate to `Product --> Scheme --> Edit Scheme --> Info --> Build Configuration` and update the configuration to "Release". - -We recommend that you only use the Debug build scheme during development, where you might need to access additional logs. Debug build has logging overhead and will impact inferencing performance, while release build has compiler optimizations enabled and all logging overhead removed. - -For more details integrating and Running ExecuTorch on Apple Platforms or building the package locally, checkout this [link](https://pytorch.org/executorch/main/using-executorch-ios). - -### 4. Build and Run the project - -Click the "play" button on top right of your XCode app, or navigate to `Product --> Run` to build and run the app on your phone. - -### 5. Pushing Model and Tokenizer - -There are two options to copy the model (.pte) and tokenizer files (.model) to your app, depending on whether you are running it on a simulator or device. - -#### 5.1 Copy the model and tokenizer to Simulator -* Drag&drop the model and tokenizer files onto the Simulator window and save them somewhere inside the iLLaMA folder. -* Pick the files in the app dialog, type a prompt and click the arrow-up button. - -#### 5.2 Copy the model and tokenizer to Device -* Plug the device into your Mac and open the contents in Finder. -* Navigate to the Files tab and drag & drop the model and tokenizer files onto the iLLaMA folder. -* Wait until the files are copied. - -### 6. Try out the app -Open the iLLaMA app, click the settings button at the top left of the app to select the model and tokenizer files. When the app successfully runs on your device, you should see something like below: - -

-iOS LLaMA App -

- - -For Llava 1.5 models, you can select and image (via image/camera selector button) before typing prompt and send button. - -

-iOS LLaMA App -

- -## Reporting Issues -If you encountered any bugs or issues following this tutorial please file a bug/issue here on [Github](https://github.com/pytorch/executorch/issues/new). diff --git a/examples/demo-apps/react-native/rnllama/README.md b/examples/demo-apps/react-native/rnllama/README.md index 7729f7a153a..46e1c66dc13 100644 --- a/examples/demo-apps/react-native/rnllama/README.md +++ b/examples/demo-apps/react-native/rnllama/README.md @@ -10,7 +10,7 @@ A React Native mobile application for running LLaMA language models using ExecuT - Run LLaMA models directly on device, build the UI using React Native - Tested using Llama 3.2 SpinQuant 1B on iPhone 12 Pro -- The setup is heavily inspired by the [LLaMA iOS app example](https://github.com/pytorch/executorch/tree/main/examples/demo-apps/apple_ios/LLaMA) +- The setup is heavily inspired by the [etLLM app](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) ## Prerequisites diff --git a/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.h b/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.h index 5aaf4bc5724..993b6501a9d 100644 --- a/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.h +++ b/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.h @@ -1,14 +1,14 @@ #ifndef LLaMABridge_h #define LLaMABridge_h -#import +#import #import #import NS_ASSUME_NONNULL_BEGIN @interface LLaMABridge : RCTEventEmitter -@property (nonatomic, strong) LLaMARunner *runner; +@property (nonatomic, strong) ExecuTorchLLMTextRunner *runner; @end NS_ASSUME_NONNULL_END diff --git a/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.mm b/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.mm index 65a3e003f7c..dc5c5c1325c 100644 --- a/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.mm +++ b/examples/demo-apps/react-native/rnllama/ios/LlamaBridge.mm @@ -17,8 +17,9 @@ + (BOOL)requiresMainQueueSetup { resolver:(RCTPromiseResolveBlock)resolve rejecter:(RCTPromiseRejectBlock)reject) { dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ - self.runner = [[LLaMARunner alloc] initWithModelPath:modelPath tokenizerPath:tokenizerPath]; - + self.runner = [[ExecuTorchLLMTextRunner alloc] initWithModelPath:modelPath + tokenizerPath:tokenizerPath]; + NSError *error = nil; if (![self.runner loadWithError:&error]) { reject(@"load_error", error.localizedDescription, error); @@ -36,8 +37,8 @@ + (BOOL)requiresMainQueueSetup { dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{ NSError *error = nil; BOOL success = [self.runner generate:prompt - sequenceLength:[seqLen integerValue] - withTokenCallback:^(NSString *token) { + sequenceLength:[seqLen integerValue] + withTokenCallback:^(NSString *token) { [self sendEventWithName:@"onToken" body:token]; } error:&error]; @@ -56,4 +57,4 @@ + (BOOL)requiresMainQueueSetup { }); } -@end \ No newline at end of file +@end diff --git a/examples/demo-apps/react-native/rnllama/ios/rnllama.xcodeproj/project.pbxproj b/examples/demo-apps/react-native/rnllama/ios/rnllama.xcodeproj/project.pbxproj index 68d8ed3e955..490269d8a4d 100644 --- a/examples/demo-apps/react-native/rnllama/ios/rnllama.xcodeproj/project.pbxproj +++ b/examples/demo-apps/react-native/rnllama/ios/rnllama.xcodeproj/project.pbxproj @@ -7,8 +7,15 @@ objects = { /* Begin PBXBuildFile section */ - 036509DE2E1F7CA700C1BC1B /* LLaMARunner.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 036509D32E1F7C0800C1BC1B /* LLaMARunner.framework */; }; - 036509DF2E1F7CB100C1BC1B /* LLaMARunner.framework in Embed Frameworks */ = {isa = PBXBuildFile; fileRef = 036509D32E1F7C0800C1BC1B /* LLaMARunner.framework */; settings = {ATTRIBUTES = (CodeSignOnCopy, RemoveHeadersOnCopy, ); }; }; + 03F814712E729261002D91CC /* backend_coreml in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814702E729261002D91CC /* backend_coreml */; }; + 03F814732E729261002D91CC /* backend_mps in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814722E729261002D91CC /* backend_mps */; }; + 03F814752E729261002D91CC /* backend_xnnpack in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814742E729261002D91CC /* backend_xnnpack */; }; + 03F814772E729261002D91CC /* executorch_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814762E729261002D91CC /* executorch_debug */; }; + 03F814792E729261002D91CC /* executorch_llm_debug in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814782E729261002D91CC /* executorch_llm_debug */; }; + 03F8147B2E729261002D91CC /* kernels_llm in Frameworks */ = {isa = PBXBuildFile; productRef = 03F8147A2E729261002D91CC /* kernels_llm */; }; + 03F8147D2E729261002D91CC /* kernels_optimized in Frameworks */ = {isa = PBXBuildFile; productRef = 03F8147C2E729261002D91CC /* kernels_optimized */; }; + 03F8147F2E729261002D91CC /* kernels_quantized in Frameworks */ = {isa = PBXBuildFile; productRef = 03F8147E2E729261002D91CC /* kernels_quantized */; }; + 03F814812E729261002D91CC /* kernels_torchao in Frameworks */ = {isa = PBXBuildFile; productRef = 03F814802E729261002D91CC /* kernels_torchao */; }; 13B07FBC1A68108700A75B9A /* AppDelegate.mm in Sources */ = {isa = PBXBuildFile; fileRef = 13B07FB01A68108700A75B9A /* AppDelegate.mm */; }; 13B07FBF1A68108700A75B9A /* Images.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 13B07FB51A68108700A75B9A /* Images.xcassets */; }; 13B07FC11A68108700A75B9A /* main.m in Sources */ = {isa = PBXBuildFile; fileRef = 13B07FB71A68108700A75B9A /* main.m */; }; @@ -21,30 +28,6 @@ E931C67F2CFAF17500DA599B /* LlamaBridge.mm in Sources */ = {isa = PBXBuildFile; fileRef = E931C67E2CFAF17500DA599B /* LlamaBridge.mm */; }; /* End PBXBuildFile section */ -/* Begin PBXContainerItemProxy section */ - 036509D22E1F7C0800C1BC1B /* PBXContainerItemProxy */ = { - isa = PBXContainerItemProxy; - containerPortal = 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */; - proxyType = 2; - remoteGlobalIDString = 03729ED52BB1F8DE00152F2E; - remoteInfo = LLaMARunner; - }; - 036509DC2E1F7C9B00C1BC1B /* PBXContainerItemProxy */ = { - isa = PBXContainerItemProxy; - containerPortal = 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */; - proxyType = 1; - remoteGlobalIDString = 03729ED42BB1F8DE00152F2E; - remoteInfo = LLaMARunner; - }; - 036509E32E1F983A00C1BC1B /* PBXContainerItemProxy */ = { - isa = PBXContainerItemProxy; - containerPortal = 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */; - proxyType = 2; - remoteGlobalIDString = 036CAF9D2BB1444500D6C2D5; - remoteInfo = LLaMA; - }; -/* End PBXContainerItemProxy section */ - /* Begin PBXCopyFilesBuildPhase section */ E931C64C2CFAF07E00DA599B /* Embed Frameworks */ = { isa = PBXCopyFilesBuildPhase; @@ -52,7 +35,6 @@ dstPath = ""; dstSubfolderSpec = 10; files = ( - 036509DF2E1F7CB100C1BC1B /* LLaMARunner.framework in Embed Frameworks */, ); name = "Embed Frameworks"; runOnlyForDeploymentPostprocessing = 0; @@ -60,7 +42,6 @@ /* End PBXCopyFilesBuildPhase section */ /* Begin PBXFileReference section */ - 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */ = {isa = PBXFileReference; lastKnownFileType = "wrapper.pb-project"; name = LLaMA.xcodeproj; path = "/Users/shoumikhin/executorch/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj"; sourceTree = ""; }; 13B07F961A680F5B00A75B9A /* rnllama.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = rnllama.app; sourceTree = BUILT_PRODUCTS_DIR; }; 13B07FAF1A68108700A75B9A /* AppDelegate.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; name = AppDelegate.h; path = rnllama/AppDelegate.h; sourceTree = ""; }; 13B07FB01A68108700A75B9A /* AppDelegate.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; name = AppDelegate.mm; path = rnllama/AppDelegate.mm; sourceTree = ""; }; @@ -86,23 +67,22 @@ isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( - 036509DE2E1F7CA700C1BC1B /* LLaMARunner.framework in Frameworks */, + 03F814792E729261002D91CC /* executorch_llm_debug in Frameworks */, + 03F8147F2E729261002D91CC /* kernels_quantized in Frameworks */, + 03F8147D2E729261002D91CC /* kernels_optimized in Frameworks */, + 03F814772E729261002D91CC /* executorch_debug in Frameworks */, 96905EF65AED1B983A6B3ABC /* libPods-rnllama.a in Frameworks */, + 03F814752E729261002D91CC /* backend_xnnpack in Frameworks */, + 03F814812E729261002D91CC /* kernels_torchao in Frameworks */, + 03F814732E729261002D91CC /* backend_mps in Frameworks */, + 03F8147B2E729261002D91CC /* kernels_llm in Frameworks */, + 03F814712E729261002D91CC /* backend_coreml in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ - 036509CC2E1F7C0800C1BC1B /* Products */ = { - isa = PBXGroup; - children = ( - 036509E42E1F983A00C1BC1B /* LLaMA.app */, - 036509D32E1F7C0800C1BC1B /* LLaMARunner.framework */, - ); - name = Products; - sourceTree = ""; - }; 13B07FAE1A68108700A75B9A /* rnllama */ = { isa = PBXGroup; children = ( @@ -147,7 +127,6 @@ 2D16E6871FA4F8E400B85C8A /* Frameworks */, D65327D7A22EEC0BE12398D9 /* Pods */, D7E4C46ADA2E9064B798F356 /* ExpoModulesProviders */, - 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */, ); indentWidth = 2; sourceTree = ""; @@ -216,7 +195,6 @@ buildRules = ( ); dependencies = ( - 036509DD2E1F7C9B00C1BC1B /* PBXTargetDependency */, ); name = rnllama; productName = rnllama; @@ -246,14 +224,11 @@ Base, ); mainGroup = 83CBB9F61A601CBA00E9B192; + packageReferences = ( + 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */, + ); productRefGroup = 83CBBA001A601CBA00E9B192 /* Products */; projectDirPath = ""; - projectReferences = ( - { - ProductGroup = 036509CC2E1F7C0800C1BC1B /* Products */; - ProjectRef = 036509C92E1F7C0800C1BC1B /* LLaMA.xcodeproj */; - }, - ); projectRoot = ""; targets = ( 13B07F861A680F5B00A75B9A /* rnllama */, @@ -261,23 +236,6 @@ }; /* End PBXProject section */ -/* Begin PBXReferenceProxy section */ - 036509D32E1F7C0800C1BC1B /* LLaMARunner.framework */ = { - isa = PBXReferenceProxy; - fileType = wrapper.framework; - path = LLaMARunner.framework; - remoteRef = 036509D22E1F7C0800C1BC1B /* PBXContainerItemProxy */; - sourceTree = BUILT_PRODUCTS_DIR; - }; - 036509E42E1F983A00C1BC1B /* LLaMA.app */ = { - isa = PBXReferenceProxy; - fileType = wrapper.application; - path = LLaMA.app; - remoteRef = 036509E32E1F983A00C1BC1B /* PBXContainerItemProxy */; - sourceTree = BUILT_PRODUCTS_DIR; - }; -/* End PBXReferenceProxy section */ - /* Begin PBXResourcesBuildPhase section */ 13B07F8E1A680F5B00A75B9A /* Resources */ = { isa = PBXResourcesBuildPhase; @@ -418,14 +376,6 @@ }; /* End PBXSourcesBuildPhase section */ -/* Begin PBXTargetDependency section */ - 036509DD2E1F7C9B00C1BC1B /* PBXTargetDependency */ = { - isa = PBXTargetDependency; - name = LLaMARunner; - targetProxy = 036509DC2E1F7C9B00C1BC1B /* PBXContainerItemProxy */; - }; -/* End PBXTargetDependency section */ - /* Begin XCBuildConfiguration section */ 13B07F941A680F5B00A75B9A /* Debug */ = { isa = XCBuildConfiguration; @@ -640,6 +590,65 @@ defaultConfigurationName = Release; }; /* End XCConfigurationList section */ + +/* Begin XCRemoteSwiftPackageReference section */ + 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */ = { + isa = XCRemoteSwiftPackageReference; + repositoryURL = "https://github.com/pytorch/executorch"; + requirement = { + branch = "swiftpm-0.8.0.20250909"; + kind = branch; + }; + }; +/* End XCRemoteSwiftPackageReference section */ + +/* Begin XCSwiftPackageProductDependency section */ + 03F814702E729261002D91CC /* backend_coreml */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = backend_coreml; + }; + 03F814722E729261002D91CC /* backend_mps */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = backend_mps; + }; + 03F814742E729261002D91CC /* backend_xnnpack */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = backend_xnnpack; + }; + 03F814762E729261002D91CC /* executorch_debug */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = executorch_debug; + }; + 03F814782E729261002D91CC /* executorch_llm_debug */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = executorch_llm_debug; + }; + 03F8147A2E729261002D91CC /* kernels_llm */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = kernels_llm; + }; + 03F8147C2E729261002D91CC /* kernels_optimized */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = kernels_optimized; + }; + 03F8147E2E729261002D91CC /* kernels_quantized */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = kernels_quantized; + }; + 03F814802E729261002D91CC /* kernels_torchao */ = { + isa = XCSwiftPackageProductDependency; + package = 03F8146F2E729261002D91CC /* XCRemoteSwiftPackageReference "executorch" */; + productName = kernels_torchao; + }; +/* End XCSwiftPackageProductDependency section */ }; rootObject = 83CBB9F71A601CBA00E9B192 /* Project object */; } diff --git a/examples/devtools/CMakeLists.txt b/examples/devtools/CMakeLists.txt index e9aa683f1fe..355ff375361 100644 --- a/examples/devtools/CMakeLists.txt +++ b/examples/devtools/CMakeLists.txt @@ -65,6 +65,10 @@ target_link_libraries( portable_kernels ) +if(EXECUTORCH_BUILD_VULKAN) + target_link_libraries(example_runner vulkan_backend) +endif() + if(EXECUTORCH_BUILD_COREML) find_library(ACCELERATE_FRAMEWORK Accelerate) find_library(COREML_FRAMEWORK CoreML) @@ -81,12 +85,10 @@ if(EXECUTORCH_BUILD_COREML) NO_DEFAULT_PATH ) - target_link_libraries( - example_runner "-Wl,-force_load" coremldelegate - ) + target_link_libraries(example_runner "-Wl,-force_load" coremldelegate) target_link_libraries( - example_runner ${PROTOBUF_LITE} ${ACCELERATE_FRAMEWORK} - ${COREML_FRAMEWORK} ${FOUNDATION_FRAMEWORK} ${SQLITE_LIBRARY} + example_runner ${PROTOBUF_LITE} ${ACCELERATE_FRAMEWORK} ${COREML_FRAMEWORK} + ${FOUNDATION_FRAMEWORK} ${SQLITE_LIBRARY} ) endif() diff --git a/examples/llm_manual/export_nanogpt.py b/examples/llm_manual/export_nanogpt.py index 8c948479f2a..9beb041ce27 100644 --- a/examples/llm_manual/export_nanogpt.py +++ b/examples/llm_manual/export_nanogpt.py @@ -15,7 +15,7 @@ from executorch.exir import to_edge from model import GPT -from torch.export import export, export_for_training +from torch.export import export from torch.nn.attention import sdpa_kernel, SDPBackend model = GPT.from_pretrained("gpt2") # use gpt2 weight as pretrained weight @@ -27,7 +27,7 @@ # Trace the model, converting it to a portable intermediate representation. # The torch.no_grad() call tells PyTorch to exclude training-specific logic. with sdpa_kernel([SDPBackend.MATH]), torch.no_grad(): - m = export_for_training( + m = export( model, example_inputs, dynamic_shapes=dynamic_shape, strict=True ).module() traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape, strict=True) diff --git a/examples/llm_pte_finetuning/TARGETS b/examples/llm_pte_finetuning/TARGETS index 5ba24c11203..1acb7adee49 100644 --- a/examples/llm_pte_finetuning/TARGETS +++ b/examples/llm_pte_finetuning/TARGETS @@ -1,9 +1,9 @@ load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("papaya_oncall") -python_library( +runtime.python_library( name = "model_loading_lib", srcs = [ "model_loading_lib.py", @@ -20,7 +20,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "training_lib", srcs = [ "training_lib.py", diff --git a/examples/mediatek/CMakeLists.txt b/examples/mediatek/CMakeLists.txt index 2e79130e5c6..f93954bff9c 100644 --- a/examples/mediatek/CMakeLists.txt +++ b/examples/mediatek/CMakeLists.txt @@ -29,19 +29,16 @@ endif() set(_common_compile_options -Wno-deprecated-declarations -fPIC) # Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/.. - ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 - ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include - ${EXECUTORCH_ROOT}/extension/llm/tokenizers/third-party/json/single_include) +set(_common_include_directories + ${EXECUTORCH_ROOT}/.. ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10 + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/third-party/json/single_include +) # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../executorch_srcs.cmake" -) -extract_sources(${EXECUTORCH_SRCS_FILE}) -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() # Find prebuilt libraries. executorch package should contain portable_ops_lib, # etdump, bundled_program. @@ -73,7 +70,10 @@ if(${ANDROID}) ) target_link_libraries( - mtk_executor_runner ${_executor_runner_libs} executorch neuron_backend + mtk_executor_runner + ${_executor_runner_libs} + executorch + neuron_backend executorch_core extension_evalue_util extension_runner_util @@ -83,42 +83,30 @@ if(${ANDROID}) add_compile_definitions(C10_USING_CUSTOM_GENERATED_MACROS) set(_mtk_oss_executor_runner__srcs ${_executor_runner__srcs}) - list( - TRANSFORM - _mtk_oss_executor_runner__srcs - PREPEND - "${EXECUTORCH_SOURCE_DIR}/" + list(TRANSFORM _mtk_oss_executor_runner__srcs + PREPEND "${EXECUTORCH_SOURCE_DIR}/" ) - list( - FILTER - _mtk_oss_executor_runner__srcs - EXCLUDE REGEX - ".*executor_runner.cpp$" + list(FILTER _mtk_oss_executor_runner__srcs EXCLUDE REGEX + ".*executor_runner.cpp$" ) - list( - PREPEND - _mtk_oss_executor_runner__srcs - ${CMAKE_CURRENT_LIST_DIR}/executor_runner/mtk_oss_executor_runner.cpp + list(PREPEND _mtk_oss_executor_runner__srcs + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/mtk_oss_executor_runner.cpp ) add_executable(mtk_oss_executor_runner ${_mtk_oss_executor_runner__srcs}) - target_include_directories(mtk_oss_executor_runner - PUBLIC - ${_common_include_directories} - ${EXECUTORCH_ROOT}/cmake-android-out/third-party/gflags/include + target_include_directories( + mtk_oss_executor_runner + PUBLIC ${_common_include_directories} + ${EXECUTORCH_ROOT}/cmake-android-out/third-party/gflags/include ) - target_link_libraries(mtk_oss_executor_runner - ${_executor_runner_libs} - extension_module - executorch - neuron_backend - gflags + target_link_libraries( + mtk_oss_executor_runner ${_executor_runner_libs} extension_module + executorch neuron_backend gflags ) - target_compile_options(mtk_oss_executor_runner - PUBLIC - ${_common_compile_options} + target_compile_options( + mtk_oss_executor_runner PUBLIC ${_common_compile_options} ) set(_mtk_llama_executor_runner__srcs ${_mtk_executor_runner__srcs}) @@ -130,30 +118,59 @@ if(${ANDROID}) ) # Build ABSL and RE2 set(EXTENSIONS_LLM_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../extension/llm) - set(THIRD_PARTY_ABSL_DIR ${EXTENSIONS_LLM_DIR}/tokenizers/third-party/abseil-cpp) + set(THIRD_PARTY_ABSL_DIR + ${EXTENSIONS_LLM_DIR}/tokenizers/third-party/abseil-cpp + ) set(THIRD_PARTY_RE2_DIR ${EXTENSIONS_LLM_DIR}/tokenizers/third-party/re2) + set(THIRD_PARTY_JSON_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../third-party/json) + set(THIRD_PARTY_UNICODE_DIR + ${EXTENSIONS_LLM_DIR}/tokenizers/third-party/llama.cpp-unicode + ) + set(THIRD_PARTY_PCRE2_DIR ${EXTENSIONS_LLM_DIR}/tokenizers/third-party/pcre2) set(ABSL_ENABLE_INSTALL ON) set(ABSL_PROPAGATE_CXX_STD ON) set(_pic_flag ${CMAKE_POSITION_INDEPENDENT_CODE}) set(CMAKE_POSITION_INDEPENDENT_CODE ON) add_subdirectory( - ${THIRD_PARTY_ABSL_DIR} ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/abseil + ${THIRD_PARTY_ABSL_DIR} + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/abseil + ) + add_subdirectory( + ${THIRD_PARTY_RE2_DIR} + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/re2 + ) + add_subdirectory( + ${THIRD_PARTY_JSON_DIR} + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/json + ) + add_subdirectory( + ${THIRD_PARTY_UNICODE_DIR} + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/llama.cpp-unicode ) add_subdirectory( - ${THIRD_PARTY_RE2_DIR} ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/re2 + ${THIRD_PARTY_PCRE2_DIR} + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/pcre2 ) set(CMAKE_POSITION_INDEPENDENT_CODE ${_pic_flag}) # Build tokenizers + set(SUPPORT_REGEX_LOOKAHEAD ON) set(LLAMA2_TOKENIZER_DIR ${EXTENSIONS_LLM_DIR}/tokenizers) add_library(tokenizer STATIC) target_include_directories( - tokenizer PUBLIC ${_common_include_directories} ${THIRD_PARTY_ABSL_DIR} - ${THIRD_PARTY_RE2_DIR} ${LLAMA2_TOKENIZER_DIR}/include - ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/pcre2 - ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include + tokenizer + PUBLIC ${_common_include_directories} + ${THIRD_PARTY_ABSL_DIR} + ${THIRD_PARTY_RE2_DIR} + ${LLAMA2_TOKENIZER_DIR}/include + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/third-party/pcre2 + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include + ${THIRD_PARTY_JSON_DIR} + ${THIRD_PARTY_UNICODE_DIR}/include + ${THIRD_PARTY_PCRE2_DIR} ) target_link_libraries(tokenizer PRIVATE re2::re2) + target_sources( tokenizer PRIVATE @@ -162,9 +179,55 @@ if(${ANDROID}) ${LLAMA2_TOKENIZER_DIR}/src/regex.cpp ${LLAMA2_TOKENIZER_DIR}/src/bpe_tokenizer_base.cpp ${LLAMA2_TOKENIZER_DIR}/src/re2_regex.cpp + ${LLAMA2_TOKENIZER_DIR}/src/hf_tokenizer.cpp + ${LLAMA2_TOKENIZER_DIR}/src/pre_tokenizer.cpp + ${LLAMA2_TOKENIZER_DIR}/src/token_decoder.cpp + ${LLAMA2_TOKENIZER_DIR}/src/normalizer.cpp + ${LLAMA2_TOKENIZER_DIR}/third-party/llama.cpp-unicode/src/unicode.cpp + ${LLAMA2_TOKENIZER_DIR}/third-party/llama.cpp-unicode/src/unicode-data.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../models/llama/tokenizer/llama_tiktoken.cpp ) + # Add support for regex_lookahead + set(PCRE2_STATIC_PIC ON) + set(PCRE2_BUILD_PCRE2_8 ON) + set(PCRE2_BUILD_PCRE2_16 OFF) + set(PCRE2_BUILD_PCRE2_32 OFF) + set(PCRE2_BUILD_TESTS OFF) + set(PCRE2_BUILD_PCRE2GREP OFF) + set(PCRE2_BUILD_PCRE2TEST OFF) + set(PCRE2_BUILD_PCRE2GPERF OFF) + set(PCRE2_BUILD_DOCS OFF) + set(PCRE2_BUILD_LIBPCRE2_PDB OFF) + + # Set the INTERFACE_INCLUDE_DIRECTORIES property for pcre2-8-static + set_target_properties( + pcre2-8-static + PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES + $ + ) + add_library( + regex_lookahead STATIC + ${LLAMA2_TOKENIZER_DIR}/src/pcre2_regex.cpp + ${LLAMA2_TOKENIZER_DIR}/src/regex_lookahead.cpp + ${LLAMA2_TOKENIZER_DIR}/src/std_regex.cpp + ) + add_library(tokenizer::regex_lookahead ALIAS regex_lookahead) + target_link_libraries(regex_lookahead PUBLIC pcre2-8-static) + target_include_directories( + regex_lookahead + PUBLIC $ + ) + target_link_libraries(tokenizer PUBLIC regex_lookahead) + install( + TARGETS regex_lookahead pcre2-8-static + EXPORT tokenizers-targets + ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} + ) + # Include directory for neuron headers include_directories( BEFORE ${_common_include_directories} @@ -178,12 +241,12 @@ if(${ANDROID}) add_executable(mtk_llama_executor_runner ${_mtk_llama_executor_runner__srcs}) target_link_libraries( - mtk_llama_executor_runner - ${_executor_runner_libs} - neuron_backend - gflags + mtk_llama_executor_runner ${_executor_runner_libs} neuron_backend gflags mtk_llama_executor_lib - tokenizer + ) + target_link_libraries( + mtk_llama_executor_runner tokenizer + $ ) target_compile_options( mtk_llama_executor_runner PUBLIC ${_common_compile_options} diff --git a/examples/mediatek/README.md b/examples/mediatek/README.md old mode 100644 new mode 100755 index 876d40adf7e..ffc91e10e60 --- a/examples/mediatek/README.md +++ b/examples/mediatek/README.md @@ -26,44 +26,61 @@ examples/mediatek # Examples Build Instructions ## Environment Setup -- Follow the instructions of **Prerequisites** and **Setup** in `backends/mediatek/scripts/README.md`. +- Follow the instructions in `backends/mediatek/README.md` to build the backend library `libneuron_backend.so`. -- Build required libraries by `backends/mediatek/scripts/mtk_build.sh` before building examples. - -## Build MediaTek Examples -1. Build the backend and the examples by exedcuting the script: +## Build MediaTek Runners +1. Build the mediatek model runner by executing the script: ```bash ./mtk_build_examples.sh ``` +This will generate the required runners in `executorch/cmake-android-out/examples/mediatek/` -## LLaMa Example Instructions +## Model Export Instructions ##### Note: Verify that localhost connection is available before running AoT Flow -1. Exporting Models to `.pte` -- In the `examples/mediatek directory`, run: +1. Download Required Files +- Download the model files from the official Hugging Face website, and move the files to the respective folder in `examples/mediatek/models/llm_models/weights/` **EXCEPT** the `config.json` file. + - The `config.json` file is already included in the model folders, which may include some modifications required for the model exportation. +- Include the calibration data (if any) under `aot_utils/llm_utils/prompts/` + +2. Exporting Models to `.pte` +- In the `examples/mediatek/ directory`, run: ```bash -source shell_scripts/export_llama.sh +source shell_scripts/export_.sh ``` - Defaults: - - `model_name` = llama3 + - `model_name` = Depends on model family. Check respective `shell_scripts/export_.sh` for info. - `num_chunks` = 4 - `prompt_num_tokens` = 128 - - `cache_size` = 1024 - - `calibration_set_name` = None + - `cache_size` = 512 + - `calibration_data_file` = None + - `precision` = A16W4 + - `platform` = DX4 + - Argument Explanations/Options: - - `model_name`: llama2/llama3 - **Note: Currently Only Tested on Llama2 7B Chat and Llama3 8B Instruct.** - - `num_chunks`: Number of chunks to split the model into. Each chunk contains the same number of decoder layers. Will result in `num_chunks` number of `.pte` files being generated. Typical values are 1, 2 and 4. + - `model_name`: View list 'Available model names' below. + - `num_chunks`: Number of chunks to split the model into. Each chunk contains the same number of decoder layers. Typical values are 1, 2 and 4. - `prompt_num_tokens`: Number of tokens (> 1) consumed each forward pass for the prompt processing stage. - `cache_size`: Cache Size. - - `calibration_set_name`: Name of calibration dataset with extension that is found inside the `aot_utils/llm_utils/prompts` directory. Example: `alpaca.txt`. If `"None"`, will use dummy data to calibrate. + - `calibration_data_file`: Name of calibration dataset with extension that is found inside the `aot_utils/llm_utils/prompts/` directory. Example: `alpaca.txt`. If `"None"`, will use dummy data to calibrate. + - `precision`: Quantization precision for the model. Available options are `["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"]` + - `platform`: The platform of the device. `DX4` for Mediatek Dimensity 9400 and `DX3` for Mediatek Dimensity 9300. **Note: Export script example only tested on `.txt` file.** -2. `.pte` files will be generated in `examples/mediatek/pte` - - Users should expect `num_chunks*2` number of pte files (half of them for prompt and half of them for generation). - - Generation `.pte` files have "`1t`" in their names. - - Additionally, an embedding bin file will be generated in the weights folder where the `config.json` can be found in. [`examples/mediatek/models/llm_models/weights//embedding__fp32.bin`] +- Available model names: + - Llama: + - llama3.2-3b, llama3.2-1b, llama3, llama2 + - Qwen: + - Qwen3-4B, Qwen3-1.7B, Qwen2-7B-Instruct, Qwen2.5-3B, Qwen2.5-0.5B-Instruct, Qwen2-1.5B-Instruct + - Gemma: + - gemma2, gemma3 + - Phi: + - phi3.5, phi4 + +3. `.pte` files will be generated in `examples/mediatek/pte/` + - Users should expect `num_chunks` number of pte files. + - An embedding bin file will be generated in the weights folder where the `config.json` can be found in. [`examples/mediatek/models/llm_models/weights//embedding__fp32.bin`] - eg. For `llama3-8B-instruct`, embedding bin generated in `examples/mediatek/models/llm_models/weights/llama3-8B-instruct/` - - AoT flow will take roughly 2.5 hours (114GB RAM for `num_chunks=4`) to complete (Results will vary by device/hardware configurations) + - AoT flow will take around 30 minutes to 2.5 hours to complete (Results will vary depending on device/hardware configurations and model sizes) ### oss 1. Exporting Model to `.pte` @@ -71,29 +88,34 @@ source shell_scripts/export_llama.sh ``` - Argument Options: - - `model_name`: deeplabv3/edsr/inceptionv3/inceptionv4/mobilenetv2/mobilenetv3/resnet18/resnet50 + - `model_name`: deeplabv3/edsr/inceptionv3/inceptionv4/mobilenetv2/mobilenetv3/resnet18/resnet50/dcgan/wav2letter/vit_b_16/mobilebert/emformer_rnnt/bert/distilbert # Runtime -## Environment Setup - -To set up the build environment for the `mtk_executor_runner`: - -1. Navigate to the `backends/mediatek/scripts` directory within the repository. -2. Follow the detailed build steps provided in that location. -3. Upon successful completion of the build steps, the `mtk_executor_runner` binary will be generated. - ## Deploying and Running on the Device ### Pushing Files to the Device -Transfer the `.pte` model files and the `mtk_executor_runner` binary to your Android device using the following commands: +Transfer the directory containing the `.pte` model files, the `run__sample.sh` script, the `embedding__fp32.bin`, the tokenizer file, the `mtk_llama_executor_runner` binary and the 3 `.so` files to your Android device using the following commands: ```bash -adb push mtk_executor_runner -adb push .pte +adb push mtk_llama_executor_runner +adb push examples/mediatek/executor_runner/run__sample.sh +adb push embedding__fp32.bin +adb push tokenizer.model +adb push ``` -Make sure to replace `` with the actual name of your model file. And, replace the `` with the desired detination on the device. +Make sure to replace `` with the actual name of your directory containing pte files. And, replace the `` with the desired detination on the device. + +At this point your phone directory should have the following files: +- libneuron_backend.so +- libneuronusdk_adapter.mtk.so +- libneuron_buffer_allocator.so +- mtk_llama_executor_runner +- +- tokenizer.json / tokenizer.model(for llama3) / tokenizer.bin(for phi3 and gemma2) +- embedding__fp32.bin +- run__sample.sh ##### Note: For oss models, please push additional files to your Android device ```bash @@ -107,12 +129,13 @@ for i in input*bin; do adb push "$i" ; done; Execute the model on your Android device by running: ```bash -adb shell "/data/local/tmp/mtk_executor_runner --model_path /data/local/tmp/.pte --iteration " +adb shell +cd +sh run__sample.sh ``` +#### Note: The `mtk_llama_executor_runner` is applicable to the models listed in `examples/mediatek/models/llm_models/weights/`. -In the command above, replace `` with the name of your model file and `` with the desired number of iterations to run the model. - -##### Note: For llama models, please use `mtk_llama_executor_runner`. Refer to `examples/mediatek/executor_runner/run_llama3_sample.sh` for reference. +##### Note: For non-LLM models, please run `adb shell "/data/local/tmp/mtk_executor_runner --model_path /data/local/tmp/.pte --iteration "`. ##### Note: For oss models, please use `mtk_oss_executor_runner`. ```bash adb shell "/data/local/tmp/mtk_oss_executor_runner --model_path /data/local/tmp/.pte --input_list /data/local/tmp/input_list.txt --output_folder /data/local/tmp/output_" diff --git a/examples/mediatek/aot_utils/llm_utils/loggingprinter.py b/examples/mediatek/aot_utils/llm_utils/loggingprinter.py index 63e976f15ae..0c77437726e 100644 --- a/examples/mediatek/aot_utils/llm_utils/loggingprinter.py +++ b/examples/mediatek/aot_utils/llm_utils/loggingprinter.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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 sys diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter.py b/examples/mediatek/aot_utils/llm_utils/preformatter.py index 1ed3da5fea0..6e74b5eb377 100644 --- a/examples/mediatek/aot_utils/llm_utils/preformatter.py +++ b/examples/mediatek/aot_utils/llm_utils/preformatter.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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 json import os from typing import Union diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter_templates/gemma.json b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/gemma.json new file mode 100755 index 00000000000..4de16b0ce84 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/gemma.json @@ -0,0 +1,6 @@ +{ + "description": "Template used by gemma.", + "prompt_input": "user\n{instruction}\nmodel\n", + "prompt_no_input": "user\n{instruction}\nmodel\n", + "response_split": "model\n" +} \ No newline at end of file diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi3.json b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi3.json new file mode 100755 index 00000000000..17912e8fc6c --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi3.json @@ -0,0 +1,6 @@ +{ + "description": "Template used by Phi-3.", + "prompt_input": "<|system|>\nYou are a helpful AI assistant. Please provide safe, ethical and accurate information to the user.\n<|user|>\n {instruction} \n <|assistant|>", + "prompt_no_input": "<|system|>\nYou are a helpful AI assistant. Please provide safe, ethical and accurate information to the user.\n<|user|>\n {instruction} \n <|assistant|>", + "response_split": "<|assistant|>" +} diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi4_reasoning.json b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi4_reasoning.json new file mode 100644 index 00000000000..d802657816f --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/phi4_reasoning.json @@ -0,0 +1,6 @@ +{ + "description": "Template used by Phi-4.", + "prompt_input": "<|system|>Your name is Phi, an AI expert developed by Microsoft.<|end|><|user|>{instruction}<|end|><|assistant|>", + "prompt_no_input": "<|system|>Your name is Phi, an AI expert developed by Microsoft.<|end|><|user|>{instruction}<|end|><|assistant|>", + "response_split": "<|assistant|>" +} diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen.json b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen.json new file mode 100755 index 00000000000..30ce72da429 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen.json @@ -0,0 +1,6 @@ +{ + "description": "Template used by Qwen.", + "prompt_input": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n", + "prompt_no_input": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n", + "response_split": "<|im_start|>assistant\n" +} diff --git a/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen3.json b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen3.json new file mode 100644 index 00000000000..11394818169 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/preformatter_templates/qwen3.json @@ -0,0 +1,6 @@ +{ + "description": "Template used by Qwen3.", + "prompt_input": "<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n", + "prompt_no_input": "<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n", + "response_split": "assistant\n" +} diff --git a/examples/mediatek/aot_utils/llm_utils/sanity_checks.py b/examples/mediatek/aot_utils/llm_utils/sanity_checks.py index 0c39d4fd9ac..84b2e3a7534 100644 --- a/examples/mediatek/aot_utils/llm_utils/sanity_checks.py +++ b/examples/mediatek/aot_utils/llm_utils/sanity_checks.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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 os import sys @@ -127,7 +133,14 @@ def check_supported_model(config): "qwen", "qwen1.5", "qwen2", + "qwen3", "milm", + "phi3", + "phi4", + "gemma1", + "gemma2", + "gemma3", + "whisper", ] if not isinstance(config, BaseConfig): raise RuntimeError( diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/english_normalizer.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/english_normalizer.py new file mode 100644 index 00000000000..c4552e87b59 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/english_normalizer.py @@ -0,0 +1,649 @@ +# Copyright 2022 The OpenAI team and The HuggingFace Team. All rights reserved. # noqa: CPY001 +# Most of the code is copy pasted from the original whisper repository +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Module to contain all english normalizer classes.""" + +# flake8: noqa: C901 + + +import re +import unicodedata +from fractions import Fraction +from typing import Iterator, List, Match, Optional, Union + +import regex + +# non-ASCII letters that are not separated by "NFKD" normalization +ADDITIONAL_DIACRITICS = { + "œ": "oe", + "Œ": "OE", + "ø": "o", + "Ø": "O", + "æ": "ae", + "Æ": "AE", + "ß": "ss", + "ẞ": "SS", + "đ": "d", + "Đ": "D", + "ð": "d", + "Ð": "D", + "þ": "th", + "Þ": "th", + "ł": "l", + "Ł": "L", +} + + +def remove_symbols_and_diacritics(s: str, keep=""): + """Remove symbols and diacritics. + + Replace any other markers, symbols, and punctuations with a space, and drop any diacritics (category 'Mn' and some + manual mappings). + """ + + def replace_character(char): + if char in keep: + return char + if char in ADDITIONAL_DIACRITICS: + return ADDITIONAL_DIACRITICS[char] + + if unicodedata.category(char) == "Mn": + return "" + + if unicodedata.category(char)[0] in "MSP": + return " " + + return char + + return "".join(replace_character(c) for c in unicodedata.normalize("NFKD", s)) + + +def remove_symbols(s: str): + """Replace any other markers, symbols, punctuations with a space, keeping diacritics.""" + return "".join( + " " if unicodedata.category(c)[0] in "MSP" else c + for c in unicodedata.normalize("NFKC", s) + ) + + +class BasicTextNormalizer: + """BasicTextNormalizer class.""" + + def __init__(self, remove_diacritics: bool = False, split_letters: bool = False): + """Initialize the BasicTextNormalizer class.""" + self.clean = ( + remove_symbols_and_diacritics if remove_diacritics else remove_symbols + ) + self.split_letters = split_letters + + def __call__(self, s: str): + """Run the BasicTextNormalizer.""" + s = s.lower() + s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets + s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis + s = self.clean(s).lower() + + if self.split_letters: + s = " ".join(regex.findall(r"\X", s, regex.U)) + + return re.sub( + r"\s+", " ", s + ) # replace any successive whitespace characters with a space + + +class EnglishNumberNormalizer: + """Convert any spelled-out numbers into arabic numbers. + + While handling: + + - remove any commas + - keep the suffixes such as: `1960s`, `274th`, `32nd`, etc. + - spell out currency symbols after the number. e.g. `$20 million` -> `20000000 dollars` + - spell out `one` and `ones` + - interpret successive single-digit numbers as nominal: `one oh one` -> `101` + """ + + def __init__(self): + """Initialize the EnglishNumberNormalizer class.""" + super().__init__() + + self.zeros = {"o", "oh", "zero"} + # fmt: off + self.ones = { + name: i + for i, name in enumerate( + [ + "one", + "two", + "three", + "four", + "five", + "six", + "seven", + "eight", + "nine", + "ten", + "eleven", + "twelve", + "thirteen", + "fourteen", + "fifteen", + "sixteen", + "seventeen", + "eighteen", + "nineteen" + ], + start=1, + ) + } + # fmt: on + self.ones_plural = { + "sixes" if name == "six" else name + "s": (value, "s") + for name, value in self.ones.items() + } + self.ones_ordinal = { + "zeroth": (0, "th"), + "first": (1, "st"), + "second": (2, "nd"), + "third": (3, "rd"), + "fifth": (5, "th"), + "twelfth": (12, "th"), + **{ + name + ("h" if name.endswith("t") else "th"): (value, "th") + for name, value in self.ones.items() + if value > 3 and value != 5 and value != 12 + }, + } + self.ones_suffixed = {**self.ones_plural, **self.ones_ordinal} + + self.tens = { + "twenty": 20, + "thirty": 30, + "forty": 40, + "fifty": 50, + "sixty": 60, + "seventy": 70, + "eighty": 80, + "ninety": 90, + } + self.tens_plural = { + name.replace("y", "ies"): (value, "s") for name, value in self.tens.items() + } + self.tens_ordinal = { + name.replace("y", "ieth"): (value, "th") + for name, value in self.tens.items() + } + self.tens_suffixed = {**self.tens_plural, **self.tens_ordinal} + + self.multipliers = { + "hundred": 100, + "thousand": 1_000, + "million": 1_000_000, + "billion": 1_000_000_000, + "trillion": 1_000_000_000_000, + "quadrillion": 1_000_000_000_000_000, + "quintillion": 1_000_000_000_000_000_000, + "sextillion": 1_000_000_000_000_000_000_000, + "septillion": 1_000_000_000_000_000_000_000_000, + "octillion": 1_000_000_000_000_000_000_000_000_000, + "nonillion": 1_000_000_000_000_000_000_000_000_000_000, + "decillion": 1_000_000_000_000_000_000_000_000_000_000_000, + } + self.multipliers_plural = { + name + "s": (value, "s") for name, value in self.multipliers.items() + } + self.multipliers_ordinal = { + name + "th": (value, "th") for name, value in self.multipliers.items() + } + self.multipliers_suffixed = { + **self.multipliers_plural, + **self.multipliers_ordinal, + } + self.decimals = {*self.ones, *self.tens, *self.zeros} + + self.preceding_prefixers = { + "minus": "-", + "negative": "-", + "plus": "+", + "positive": "+", + } + self.following_prefixers = { + "pound": "£", + "pounds": "£", + "euro": "€", + "euros": "€", + "dollar": "$", + "dollars": "$", + "cent": "¢", + "cents": "¢", + } + self.prefixes = set( + list(self.preceding_prefixers.values()) + + list(self.following_prefixers.values()) + ) + self.suffixers = { + "per": {"cent": "%"}, + "percent": "%", + } + self.specials = {"and", "double", "triple", "point"} + + self.words = { + key + for mapping in [ + self.zeros, + self.ones, + self.ones_suffixed, + self.tens, + self.tens_suffixed, + self.multipliers, + self.multipliers_suffixed, + self.preceding_prefixers, + self.following_prefixers, + self.suffixers, + self.specials, + ] + for key in mapping + } + self.literal_words = {"one", "ones"} + + def process_words(self, words: List[str]) -> Iterator[str]: + """Process words.""" + prefix: Optional[str] = None + value: Optional[Union[str, int]] = None + skip = False + + def to_fraction(s: str): + try: + return Fraction(s) + except ValueError: + return None + + def output(result: Union[str, int]): + nonlocal prefix, value + result = str(result) + if prefix is not None: + result = prefix + result + value = None + prefix = None + return result + + if len(words) == 0: + return + + for i, current in enumerate(words): + prev = words[i - 1] if i != 0 else None + next_ = words[i + 1] if i != len(words) - 1 else None + if skip: + skip = False + continue + + next_is_numeric = next_ is not None and re.match(r"^\d+(\.\d+)?$", next_) + has_prefix = current[0] in self.prefixes + current_without_prefix = current[1:] if has_prefix else current + if re.match(r"^\d+(\.\d+)?$", current_without_prefix): + # arabic numbers (potentially with signs and fractions) + f = to_fraction(current_without_prefix) + if f is None: + raise ValueError("Converting the fraction failed") + + if value is not None: + if isinstance(value, str) and value.endswith("."): + # concatenate decimals / ip address components + value = str(value) + str(current) + continue + else: + yield output(value) + + prefix = current[0] if has_prefix else prefix + value = f.numerator if f.denominator == 1 else current_without_prefix + elif current not in self.words: + # non-numeric words + if value is not None: + yield output(value) + yield output(current) + elif current in self.zeros: + value = str(value or "") + "0" + elif current in self.ones: + ones = self.ones[current] + + if value is None: + value = ones + elif isinstance(value, str) or prev in self.ones: + value = ( + value[:-1] + str(ones) + if prev in self.tens and ones < 10 + else str(value) + str(ones) + ) + elif ones < 10: + if value % 10 == 0: + value += ones + else: + value = str(value) + str(ones) + else: # eleven to nineteen + if value % 100 == 0: + value += ones + else: + value = str(value) + str(ones) + elif current in self.ones_suffixed: + # ordinal or cardinal; yield the number right away + ones, suffix = self.ones_suffixed[current] + if value is None: + yield output(str(ones) + suffix) + elif isinstance(value, str) or prev in self.ones: + if prev in self.tens and ones < 10: + yield output(value[:-1] + str(ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + elif ones < 10: + if value % 10 == 0: + yield output(str(value + ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + else: # eleven to nineteen + if value % 100 == 0: + yield output(str(value + ones) + suffix) + else: + yield output(str(value) + str(ones) + suffix) + value = None + elif current in self.tens: + tens = self.tens[current] + if value is None: + value = tens + elif isinstance(value, str): + value = str(value) + str(tens) + else: + if value % 100 == 0: + value += tens + else: + value = str(value) + str(tens) + elif current in self.tens_suffixed: + # ordinal or cardinal; yield the number right away + tens, suffix = self.tens_suffixed[current] + if value is None: + yield output(str(tens) + suffix) + elif isinstance(value, str): + yield output(str(value) + str(tens) + suffix) + else: + if value % 100 == 0: + yield output(str(value + tens) + suffix) + else: + yield output(str(value) + str(tens) + suffix) + elif current in self.multipliers: + multiplier = self.multipliers[current] + if value is None: + value = multiplier + elif isinstance(value, str) or value == 0: + f = to_fraction(value) + p = f * multiplier if f is not None else None + if f is not None and p.denominator == 1: + value = p.numerator + else: + yield output(value) + value = multiplier + else: + before = value // 1000 * 1000 + residual = value % 1000 + value = before + residual * multiplier + elif current in self.multipliers_suffixed: + multiplier, suffix = self.multipliers_suffixed[current] + if value is None: + yield output(str(multiplier) + suffix) + elif isinstance(value, str): + f = to_fraction(value) + p = f * multiplier if f is not None else None + if f is not None and p.denominator == 1: + yield output(str(p.numerator) + suffix) + else: + yield output(value) + yield output(str(multiplier) + suffix) + else: # int + before = value // 1000 * 1000 + residual = value % 1000 + value = before + residual * multiplier + yield output(str(value) + suffix) + value = None + elif current in self.preceding_prefixers: + # apply prefix (positive, minus, etc.) if it precedes a number + if value is not None: + yield output(value) + + if next_ in self.words or next_is_numeric: + prefix = self.preceding_prefixers[current] + else: + yield output(current) + elif current in self.following_prefixers: + # apply prefix (dollars, cents, etc.) only after a number + if value is not None: + prefix = self.following_prefixers[current] + yield output(value) + else: + yield output(current) + elif current in self.suffixers: + # apply suffix symbols (percent -> '%') + if value is not None: + suffix = self.suffixers[current] + if isinstance(suffix, dict): + if next_ in suffix: + yield output(str(value) + suffix[next_]) + skip = True + else: + yield output(value) + yield output(current) + else: + yield output(str(value) + suffix) + else: + yield output(current) + elif current in self.specials: + if next_ not in self.words and not next_is_numeric: + # apply special handling only if the next_ word can be numeric + if value is not None: + yield output(value) + yield output(current) + elif current == "and": + # ignore "and" after hundreds, thousands, etc. + if prev not in self.multipliers: + if value is not None: + yield output(value) + yield output(current) + elif current == "double" or current == "triple": + if next_ in self.ones or next_ in self.zeros: + repeats = 2 if current == "double" else 3 + ones = self.ones.get(next_, 0) + value = str(value or "") + str(ones) * repeats + skip = True + else: + if value is not None: + yield output(value) + yield output(current) + elif current == "point": + if next_ in self.decimals or next_is_numeric: + value = str(value or "") + "." + else: + # should all have been covered at this point + raise ValueError(f"Unexpected token: {current}") + else: + # all should have been covered at this point + raise ValueError(f"Unexpected token: {current}") + + if value is not None: + yield output(value) + + def preprocess(self, s: str): + """Replace ' and a half' with ' point five'.""" + results = [] + + segments = re.split(r"\band\s+a\s+half\b", s) + for i, segment in enumerate(segments): + if len(segment.strip()) == 0: + continue + if i == len(segments) - 1: + results.append(segment) + else: + results.append(segment) + last_word = segment.rsplit(maxsplit=2)[-1] + if last_word in self.decimals or last_word in self.multipliers: + results.append("point five") + else: + results.append("and a half") + + s = " ".join(results) + + # put a space at number/letter boundary + s = re.sub(r"([a-z])([0-9])", r"\1 \2", s) + s = re.sub(r"([0-9])([a-z])", r"\1 \2", s) + + # but remove spaces which could be a suffix + return re.sub(r"([0-9])\s+(st|nd|rd|th|s)\b", r"\1\2", s) + + def postprocess(self, s: str): + """Post-process.""" + + def combine_cents(m: Match): + try: + currency = m.group(1) + integer = m.group(2) + cents = int(m.group(3)) + return f"{currency}{integer}.{cents:02d}" + except ValueError: + return m.string + + def extract_cents(m: Match): + try: + return f"¢{int(m.group(1))}" + except ValueError: + return m.string + + # apply currency postprocessing; "$2 and ¢7" -> "$2.07" + s = re.sub(r"([€£$])([0-9]+) (?:and )?¢([0-9]{1,2})\b", combine_cents, s) + s = re.sub(r"[€£$]0.([0-9]{1,2})\b", extract_cents, s) + + # write "one(s)" instead of "1(s)", just for the readability + return re.sub(r"\b1(s?)\b", r"one\1", s) + + def __call__(self, s: str): + """Run the EnglishNumberNormalizer.""" + s = self.preprocess(s) + s = " ".join(word for word in self.process_words(s.split()) if word is not None) + return self.postprocess(s) + + +class EnglishSpellingNormalizer: + """Applies British-American spelling mappings as listed in [1]. + + [1] https://www.tysto.com/uk-us-spelling-list.html @lint-ignore + """ + + def __init__(self, english_spelling_mapping): + """Initialize the EnglishSpellingNormalizer.""" + self.mapping = english_spelling_mapping + + def __call__(self, s: str): + """Run the EnglishSpellingNormalizer.""" + return " ".join(self.mapping.get(word, word) for word in s.split()) + + +class EnglishTextNormalizer: + """EnglishTextNormalizer class.""" + + def __init__(self, english_spelling_mapping): + """Initialize the EnglishTextNormalizer.""" + self.ignore_patterns = r"\b(hmm|mm|mhm|mmm|uh|um)\b" + self.replacers = { + # common contractions + r"\bwon't\b": "will not", + r"\bcan't\b": "can not", + r"\blet's\b": "let us", + r"\bain't\b": "aint", + r"\by'all\b": "you all", + r"\bwanna\b": "want to", + r"\bgotta\b": "got to", + r"\bgonna\b": "going to", + r"\bi'ma\b": "i am going to", + r"\bimma\b": "i am going to", + r"\bwoulda\b": "would have", + r"\bcoulda\b": "could have", + r"\bshoulda\b": "should have", + r"\bma'am\b": "madam", + # contractions in titles/prefixes + r"\bmr\b": "mister ", + r"\bmrs\b": "missus ", + r"\bst\b": "saint ", + r"\bdr\b": "doctor ", + r"\bprof\b": "professor ", + r"\bcapt\b": "captain ", + r"\bgov\b": "governor ", + r"\bald\b": "alderman ", + r"\bgen\b": "general ", + r"\bsen\b": "senator ", + r"\brep\b": "representative ", + r"\bpres\b": "president ", + r"\brev\b": "reverend ", + r"\bhon\b": "honorable ", + r"\basst\b": "assistant ", + r"\bassoc\b": "associate ", + r"\blt\b": "lieutenant ", + r"\bcol\b": "colonel ", + r"\bjr\b": "junior ", + r"\bsr\b": "senior ", + r"\besq\b": "esquire ", + # prefect tenses, ideally it should be any past participles, but it's harder.. + r"'d been\b": " had been", + r"'s been\b": " has been", + r"'d gone\b": " had gone", + r"'s gone\b": " has gone", + r"'d done\b": " had done", # "'s done" is ambiguous + r"'s got\b": " has got", + # general contractions + r"n't\b": " not", + r"'re\b": " are", + r"'s\b": " is", + r"'d\b": " would", + r"'ll\b": " will", + r"'t\b": " not", + r"'ve\b": " have", + r"'m\b": " am", + } + self.standardize_numbers = EnglishNumberNormalizer() + self.standardize_spellings = EnglishSpellingNormalizer(english_spelling_mapping) + + def __call__(self, s: str): + """Run the EnglishTextNormalizer.""" + s = s.lower() + + s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets + s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis + s = re.sub(self.ignore_patterns, "", s) + s = re.sub( + r"\s+'", "'", s + ) # standardize when there's a space before an apostrophe + + for pattern, replacement in self.replacers.items(): + s = re.sub(pattern, replacement, s) + + s = re.sub(r"(\d),(\d)", r"\1\2", s) # remove commas between digits + s = re.sub(r"\.([^0-9]|$)", r" \1", s) # remove periods not followed by numbers + s = remove_symbols_and_diacritics( + s, keep=".%$¢€£" + ) # keep some symbols for numerics + + s = self.standardize_numbers(s) + s = self.standardize_spellings(s) + + # now remove prefix/suffix symbols that are not preceded/followed by numbers + s = re.sub(r"[.$¢€£]([^0-9])", r" \1", s) + s = re.sub(r"([^0-9])%", r"\1 ", s) + + return re.sub( + r"\s+", " ", s + ) # replace any successive whitespace characters with a space diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma.py new file mode 100644 index 00000000000..cd63b446997 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma.py @@ -0,0 +1,352 @@ +# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved. # noqa: CPY001 +# +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Define Gemma Tokenizer Class.""" + +import os +from shutil import copyfile +from typing import Any, Dict, List, Optional, Tuple + +import sentencepiece as spm +from tokenizers import AddedToken + +from .tokenization_utils import PreTrainedTokenizer + +VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"} + +SPIECE_UNDERLINE = "▁" + + +class GemmaTokenizer(PreTrainedTokenizer): + """Construct a Gemma tokenizer. + + Based on byte-level Byte-Pair-Encoding. The default padding token is unset as there is + no padding token in the original model. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The end of sequence token. + pad_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. + sp_model_kwargs (`Dict[str, Any]`, `Optional`, *optional*): + Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for + SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things, + to set: + + - `enable_sampling`: Enable subword regularization. + - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout. + + - `nbest_size = {0,1}`: No sampling is performed. + - `nbest_size > 1`: samples from the nbest_size results. + - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice) + using forward-filtering-and-backward-sampling algorithm. + + - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for + BPE-dropout. + + add_bos_token (`bool`, *optional*, defaults to `True`): + Whether or not to add an `bos_token` at the start of sequences. + add_eos_token (`bool`, *optional*, defaults to `False`): + Whether or not to add an `eos_token` at the end of sequences. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like + extra spaces. + use_default_system_prompt (`bool`, *optional*, defaults to `False`): + Whether or not the default system prompt for Gemma should be used. + spaces_between_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to add spaces between special tokens. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + unk_token="", + bos_token="", + eos_token="", + pad_token="", + sp_model_kwargs: Optional[Dict[str, Any]] = None, + add_bos_token=True, + add_eos_token=False, + clean_up_tokenization_spaces=False, + use_default_system_prompt=False, + spaces_between_special_tokens=False, + **kwargs, + ): + """Initialize the GemmaTokenizer class.""" + self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs + bos_token = ( + AddedToken(bos_token, normalized=False, special=True) + if isinstance(bos_token, str) + else bos_token + ) + eos_token = ( + AddedToken(eos_token, normalized=False, special=True) + if isinstance(eos_token, str) + else eos_token + ) + unk_token = ( + AddedToken(unk_token, normalized=False, special=True) + if isinstance(unk_token, str) + else unk_token + ) + pad_token = ( + AddedToken(pad_token, normalized=False, special=True) + if isinstance(pad_token, str) + else pad_token + ) + + self.vocab_file = vocab_file + self.add_bos_token = add_bos_token + self.add_eos_token = add_eos_token + self.use_default_system_prompt = use_default_system_prompt + self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs) + self.sp_model.Load(vocab_file) + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + pad_token=pad_token, + add_bos_token=add_bos_token, + add_eos_token=add_eos_token, + sp_model_kwargs=sp_model_kwargs, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + use_default_system_prompt=use_default_system_prompt, + spaces_between_special_tokens=spaces_between_special_tokens, + **kwargs, + ) + + def __getstate__(self): + """Gets the GemmaTokenizer state.""" + state = self.__dict__.copy() + state["sp_model"] = None + state["sp_model_proto"] = self.sp_model.serialized_model_proto() + return state + + def __setstate__(self, d): + """Sets the GemmaTokenizer state.""" + self.__dict__.update(d) + self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs) + self.sp_model.LoadFromSerializedProto(self.sp_model_proto) + + @property + def vocab_size(self): + """Returns vocab size.""" + return self.sp_model.get_piece_size() + + def get_vocab(self): + """Returns vocab as a dict.""" + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} + vocab.update(self.added_tokens_encoder) + return vocab + + def tokenize(self, text, **kwargs) -> List[str]: + """Tokenize function. + + Simply calls PreTrainedTokenizer's method. + + Args: + text: TextInput + kwargs: Additional keyword arguments. + """ + return super().tokenize(text, **kwargs) + + def _tokenize(self, text, **kwargs): + return self.sp_model.encode(text, out_type=str) + + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.sp_model.piece_to_id(token) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.sp_model.IdToPiece(index) + + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + current_sub_tokens = [] + out_string = "" + for token in tokens: + # make sure that special tokens are not decoded using sentencepiece model + if token in self._added_tokens_encoder: + out_string += self.sp_model.decode(current_sub_tokens) + token + current_sub_tokens = [] + else: + current_sub_tokens.append(token) + out_string += self.sp_model.decode(current_sub_tokens) + return out_string + + def save_vocabulary( + self, save_directory, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save the vocabulary and special tokens file to a directory. + + Args: + save_directory (`str`): + The directory in which to save the vocabulary. + filename_prefix (`str`): + The prefix for filename. + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + if not os.path.isdir(save_directory): + raise NotADirectoryError( + f"Vocabulary path ({save_directory}) should be a directory" + ) + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath( + out_vocab_file + ) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + elif not os.path.isfile(self.vocab_file): + with open(out_vocab_file, "wb") as fi: + content_spiece_model = self.sp_model.serialized_model_proto() + fi.write(content_spiece_model) + + return (out_vocab_file,) + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + """Build inputs with special tokens.""" + bos_token_id = [self.bos_token_id] if self.add_bos_token else [] + eos_token_id = [self.eos_token_id] if self.add_eos_token else [] + + output = bos_token_id + token_ids_0 + eos_token_id + + if token_ids_1 is not None: + output = output + bos_token_id + token_ids_1 + eos_token_id + + return output + + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """Retrieve sequence ids from a token list that has no special tokens added. + + This method is called when adding special tokens using the tokenizer `prepare_for_model` method. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + bos_token_id = [1] if self.add_bos_token else [] + eos_token_id = [1] if self.add_eos_token else [] + + if token_ids_1 is None: + return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id + return ( + bos_token_id + + ([0] * len(token_ids_0)) + + eos_token_id + + bos_token_id + + ([0] * len(token_ids_1)) + + eos_token_id + ) + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """Creates a mask from the two sequences passed to be used in a sequence-pair classification task. + + An ALBERT sequence pair mask has the following format: + + ``` + 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 + | first sequence | second sequence | + ``` + + if token_ids_1 is None, only returns the first portion of the mask (0s). + + Args: + token_ids_0 (`List[int]`): + List of ids. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s). @lint-ignore + """ + bos_token_id = [self.bos_token_id] if self.add_bos_token else [] + eos_token_id = [self.eos_token_id] if self.add_eos_token else [] + + output = [0] * len(bos_token_id + token_ids_0 + eos_token_id) + + if token_ids_1 is not None: + output += [1] * len(bos_token_id + token_ids_1 + eos_token_id) + + return output + + def _decode( + self, + token_ids: List[int], + skip_special_tokens: bool = False, + spaces_between_special_tokens: bool = False, + **kwargs, + ) -> str: + sub_texts = [] + current_sub_text = [] + for ids in token_ids: + if skip_special_tokens and ids in self.all_special_ids: + continue + if ids in self._added_tokens_decoder: + if current_sub_text: + sub_texts.append(self.sp_model.decode(current_sub_text)) + sub_texts.append(self._added_tokens_decoder[ids].content) + current_sub_text = [] + else: + current_sub_text.append(ids) + if current_sub_text: + sub_texts.append(self.sp_model.decode(current_sub_text)) + + sub_texts = ( + " ".join(sub_texts) if spaces_between_special_tokens else "".join(sub_texts) + ) + + return sub_texts.replace(SPIECE_UNDERLINE, " ") + + +__all__ = ["GemmaTokenizer"] diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma_fast.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma_fast.py new file mode 100644 index 00000000000..007e9684e82 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gemma_fast.py @@ -0,0 +1,207 @@ +# Copyright 2024 The HuggingFace Inc. team. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Define Fast Gemma Tokenizer Class.""" + +import os +from shutil import copyfile +from typing import Optional, Tuple + +import tokenizers +from tokenizers import processors + +from .tokenization_gemma import GemmaTokenizer +from .tokenization_utils_fast import PreTrainedTokenizerFast + +_, major, minor = tokenizers.__version__.split(".")[:3] +if (int(major) == 13 and int(minor) < 3) or int(major) < 13: + raise RuntimeError("Requires tokenizers>=0.13.3") + + +VOCAB_FILES_NAMES = { + "vocab_file": "tokenizer.model", + "tokenizer_file": "tokenizer.json", +} + + +class GemmaTokenizerFast(PreTrainedTokenizerFast): + """Construct a Gemma tokenizer fast. Based on byte-level Byte-Pair-Encoding. + + This uses notably ByteFallback and no prefix space. Normalization is applied to replace `" "` with `"▁"` + + ```python + >>> from transformers import GemmaTokenizerFast + + >>> tokenizer = GemmaTokenizerFast.from_pretrained("hf-internal-testing/dummy-gemma") + >>> tokenizer.encode("Hello this is a test") + [2, 4521, 736, 603, 476, 2121] + ``` + + If you want to change the `bos_token` or the `eos_token`, make sure to specify them when initializing the model, or + call `tokenizer.update_post_processor()` to make sure that the post-processing is correctly done (otherwise the + values of the first token and final token of an encoded sequence will not be correct). For more details, checkout + [post-processors] (https://huggingface.co/docs/tokenizers/api/post-processors) documentation. + + + This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .model extension) that + contains the vocabulary necessary to instantiate a tokenizer. + tokenizer_file (`str`, *optional*): + [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that + contains everything needed to load the tokenizer. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like + extra spaces. + unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token. + eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `""`): + The end of sequence token. + pad_token (`str`, *optional*, defaults to `""`): + The padding token + add_bos_token (`bool`, *optional*, defaults to `True`): + Whether or not to add an `bos_token` at the start of sequences. + add_eos_token (`bool`, *optional*, defaults to `False`): + Whether or not to add an `eos_token` at the end of sequences. + """ + + vocab_files_names = VOCAB_FILES_NAMES + slow_tokenizer_class = GemmaTokenizer + padding_side = "left" + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file=None, + tokenizer_file=None, + clean_up_tokenization_spaces=False, + unk_token="", + bos_token="", + eos_token="", + pad_token="", + add_bos_token=True, + add_eos_token=False, + **kwargs, + ): + """Initialize the GemmaTokenizerFast class.""" + super().__init__( + vocab_file=vocab_file, + tokenizer_file=tokenizer_file, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + add_bos_token=add_bos_token, + add_eos_token=add_eos_token, + **kwargs, + ) + self._add_bos_token = add_bos_token + self._add_eos_token = add_eos_token + self.update_post_processor() + self.vocab_file = vocab_file + + @property + def can_save_slow_tokenizer(self) -> bool: + """Check if able to save slow tokenizer.""" + return os.path.isfile(self.vocab_file) if self.vocab_file else False + + # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.update_post_processor + def update_post_processor(self): + """Updates the underlying post processor with the current `bos_token` and `eos_token`.""" + bos = self.bos_token + bos_token_id = self.bos_token_id + if bos is None and self.add_bos_token: + raise ValueError("add_bos_token = True but bos_token = None") + + eos = self.eos_token + eos_token_id = self.eos_token_id + if eos is None and self.add_eos_token: + raise ValueError("add_eos_token = True but eos_token = None") + + single = f'{(bos + ":0 ") if self.add_bos_token else ""}$A:0{(" " + eos + ":0") if self.add_eos_token else ""}' + pair = f'{single}{(" " + bos + ":1") if self.add_bos_token else ""} $B:1{(" " + eos + ":1") if self.add_eos_token else ""}' # noqa: E501 + + special_tokens = [] + if self.add_bos_token: + special_tokens.append((bos, bos_token_id)) + if self.add_eos_token: + special_tokens.append((eos, eos_token_id)) + self._tokenizer.post_processor = processors.TemplateProcessing( + single=single, pair=pair, special_tokens=special_tokens + ) + + @property + def add_eos_token(self): + """Tokenizer's add_eos_token property.""" + return self._add_eos_token + + @property + def add_bos_token(self): + """Tokenizer's add_bos_token property.""" + return self._add_bos_token + + @add_eos_token.setter + def add_eos_token(self, value): + self._add_eos_token = value + self.update_post_processor() + + @add_bos_token.setter + def add_bos_token(self, value): + self._add_bos_token = value + self.update_post_processor() + + # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.save_vocabulary + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save vocabulary.""" + if not self.can_save_slow_tokenizer: + raise ValueError( + "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow " + "tokenizer." + ) + + if not os.path.isdir(save_directory): + raise NotADirectoryError( + f"Vocabulary path ({save_directory}) should be a directory" + ) + out_vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file): + copyfile(self.vocab_file, out_vocab_file) + + return (out_vocab_file,) + + # Copied from transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast.build_inputs_with_special_tokens + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + """Build inputs with special tokens.""" + bos_token_id = [self.bos_token_id] if self.add_bos_token else [] + eos_token_id = [self.eos_token_id] if self.add_eos_token else [] + + output = bos_token_id + token_ids_0 + eos_token_id + + if token_ids_1 is not None: + output = output + bos_token_id + token_ids_1 + eos_token_id + + return output diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gpt2.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gpt2.py new file mode 100644 index 00000000000..2da3939a54b --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_gpt2.py @@ -0,0 +1,369 @@ +# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for OpenAI GPT.""" + +import json +import os +from functools import lru_cache +from typing import List, Optional, Tuple + +import regex as re +from tokenizers import AddedToken + +from .tokenization_utils import PreTrainedTokenizer + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", +} + + +@lru_cache() +def bytes_to_unicode(): + """Returns list of utf-8 byte and a mapping to unicode strings. + + We specifically avoids mapping to whitespace/control characters the bpe code barfs on. + + The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab + if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for + decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup + tables between utf-8 bytes and unicode strings. + """ + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + """Return set of symbol pairs in a word. + + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +class GPT2Tokenizer(PreTrainedTokenizer): + """Construct a GPT-2 tokenizer. Based on byte-level Byte-Pair-Encoding. + + This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will + be encoded differently whether it is at the beginning of the sentence (without space) or not: + + ```python + >>> from transformers import GPT2Tokenizer + + >>> tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2") + >>> tokenizer("Hello world")["input_ids"] + [15496, 995] + + >>> tokenizer(" Hello world")["input_ids"] + [18435, 995] + ``` + + You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you + call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance. + + + + When used with `is_split_into_words=True`, this tokenizer will add a space before each word (even the first one). + + + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + merges_file (`str`): + Path to the merges file. + errors (`str`, *optional*, defaults to `"replace"`): + Paradigm to follow when decoding bytes to UTF-8. See + [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The beginning of sequence token. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + pad_token (`str`, *optional*): + The token used for padding, for example when batching sequences of different lengths. + add_prefix_space (`bool`, *optional*, defaults to `False`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. (GPT2 tokenizer detect beginning of words by the preceding space). + add_bos_token (`bool`, *optional*, defaults to `False`): + Whether or not to add an initial beginning of sentence token to the input. This allows to treat the leading + word just as any other word. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + merges_file, + errors="replace", + unk_token="<|endoftext|>", + bos_token="<|endoftext|>", + eos_token="<|endoftext|>", + pad_token=None, + add_prefix_space=False, + add_bos_token=False, + **kwargs, + ): + """Initialize the GPT2Tokenizer class.""" + bos_token = ( + AddedToken(bos_token, lstrip=False, rstrip=False) + if isinstance(bos_token, str) + else bos_token + ) + eos_token = ( + AddedToken(eos_token, lstrip=False, rstrip=False) + if isinstance(eos_token, str) + else eos_token + ) + unk_token = ( + AddedToken(unk_token, lstrip=False, rstrip=False) + if isinstance(unk_token, str) + else unk_token + ) + pad_token = ( + AddedToken(pad_token, lstrip=False, rstrip=False) + if isinstance(pad_token, str) + else pad_token + ) + + self.add_bos_token = add_bos_token + + with open(vocab_file, encoding="utf-8") as vocab_handle: + self.encoder = json.load(vocab_handle) + self.decoder = {v: k for k, v in self.encoder.items()} + self.errors = errors # how to handle errors in decoding + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + with open(merges_file, encoding="utf-8") as merges_handle: + bpe_merges = merges_handle.read().split("\n")[1:-1] + bpe_merges = [tuple(merge.split()) for merge in bpe_merges] + self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) + self.cache = {} + self.add_prefix_space = add_prefix_space + + # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions + self.pat = re.compile( + r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" + ) + + super().__init__( + errors=errors, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + add_prefix_space=add_prefix_space, + add_bos_token=add_bos_token, + **kwargs, + ) + + @property + def vocab_size(self): + """Gets the vocab size.""" + return len(self.encoder) + + def get_vocab(self): + """Gets the vocab.""" + return dict(self.encoder, **self.added_tokens_encoder) + + def bpe(self, token): + """Byte pair encoding.""" + if token in self.cache: + return self.cache[token] + word = tuple(token) + pairs = get_pairs(word) + + if not pairs: + return token + + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + else: + new_word.extend(word[i:j]) + i = j + + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): + """Build inputs with special tokens.""" + bos_token_ids = [self.bos_token_id] if self.add_bos_token else [] + + output = bos_token_ids + token_ids_0 + + if token_ids_1 is None: + return output + + return output + bos_token_ids + token_ids_1 + + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """Retrieves sequence ids from a token list that has no special tokens added. + + This method is called when adding special tokens using the tokenizer `prepare_for_model` or `encode_plus` + methods. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + if not self.add_bos_token: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=False, + ) + + if token_ids_1 is None: + return [1] + ([0] * len(token_ids_0)) + return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + + def _tokenize(self, text): + """Tokenize a string.""" + bpe_tokens = [] + for token in re.findall(self.pat, text): + token = "".join( + self.byte_encoder[b] for b in token.encode("utf-8") + ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) + return bpe_tokens + + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.encoder.get(token, self.encoder.get(self.unk_token)) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.decoder.get(index) + + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + return bytearray([self.byte_decoder[c] for c in text]).decode( + "utf-8", errors=self.errors + ) + + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save vocab.""" + if not os.path.isdir(save_directory): + raise NotADirectoryError( + f"Vocabulary path ({save_directory}) should be a directory" + ) + vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + merge_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["merges_file"], + ) + + with open(vocab_file, "w", encoding="utf-8") as f: + f.write( + json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + + "\n" + ) + + index = 0 + with open(merge_file, "w", encoding="utf-8") as writer: + writer.write("#version: 0.2\n") + for bpe_tokens, token_index in sorted( + self.bpe_ranks.items(), key=lambda kv: kv[1] + ): + if index != token_index: + # logger.warning( + # f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive." + # " Please check that the tokenizer is not corrupted!" + # ) + index = token_index + writer.write(" ".join(bpe_tokens) + "\n") + index += 1 + + return vocab_file, merge_file + + def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs): + """Prepare for tokenization.""" + add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space) + if is_split_into_words or add_prefix_space: + text = " " + text + return (text, kwargs) diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2.py new file mode 100644 index 00000000000..5224cedb5fa --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2.py @@ -0,0 +1,363 @@ +# Copyright 2024 The Qwen team, Alibaba Group and The HuggingFace Inc. team. All rights reserved. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Define QWen2 Tokenizer Class.""" + +import json +import os +import unicodedata +from functools import lru_cache +from typing import Optional, Tuple + +import regex as re +from tokenizers import AddedToken + +from .tokenization_utils import PreTrainedTokenizer + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", +} + + +MAX_MODEL_INPUT_SIZES = {"qwen/qwen-tokenizer": 32768} + +PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" # noqa: E501 + + +@lru_cache() +# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode +def bytes_to_unicode(): + """Returns list of utf-8 byte and a mapping to unicode strings. + + We specifically avoids mapping to whitespace/control characters the bpe code barfs on. + + The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab + if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for + decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup + tables between utf-8 bytes and unicode strings. + """ + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs +def get_pairs(word): + """Return set of symbol pairs in a word. + + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +class Qwen2Tokenizer(PreTrainedTokenizer): + """Construct a Qwen2 tokenizer. Based on byte-level Byte-Pair-Encoding. + + Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will + be encoded differently whether it is at the beginning of the sentence (without space) or not: + + ```python + >>> from transformers import Qwen2Tokenizer + + >>> tokenizer = Qwen2Tokenizer.from_pretrained("Qwen/Qwen-tokenizer") + >>> tokenizer("Hello world")["input_ids"] + [9707, 1879] + + >>> tokenizer(" Hello world")["input_ids"] + [21927, 1879] + ``` + This is expected. + + You should not use GPT2Tokenizer instead, because of the different pretokenization rules. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to + this superclass for more information regarding those methods. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + merges_file (`str`): + Path to the merges file. + errors (`str`, *optional*, defaults to `"replace"`): + Paradigm to follow when decoding bytes to UTF-8. See + [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str`, *optional*): + The beginning of sequence token. Not applicable for this tokenizer. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The token used for padding, for example when batching sequences of different lengths. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`): + Whether or not the model should cleanup the spaces that were added when splitting the input text during the + tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces. + split_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the special tokens should be split during the tokenization process. The default behavior is + to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then + `tokenizer.tokenize("<|endoftext|>") = + ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` + will be give `['<', '|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` + tokenizers for the moment. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + merges_file, + errors="replace", + unk_token="<|endoftext|>", + bos_token=None, + eos_token="<|endoftext|>", + pad_token="<|endoftext|>", + clean_up_tokenization_spaces=False, + split_special_tokens=False, + **kwargs, + ): + """Initializes the Qwen2Tokenizer.""" + # Qwen vocab does not contain control tokens; added tokens need to be special + bos_token = ( + AddedToken( + bos_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(bos_token, str) + else bos_token + ) + eos_token = ( + AddedToken( + eos_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(eos_token, str) + else eos_token + ) + unk_token = ( + AddedToken( + unk_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(unk_token, str) + else unk_token + ) + pad_token = ( + AddedToken( + pad_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(pad_token, str) + else pad_token + ) + + with open(vocab_file, encoding="utf-8") as vocab_handle: + self.encoder = json.load(vocab_handle) + self.decoder = {v: k for k, v in self.encoder.items()} + self.errors = errors # how to handle errors in decoding + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + bpe_merges = [] + with open(merges_file, encoding="utf-8") as merges_handle: + for i, line in enumerate(merges_handle): + line = line.strip() + if (i == 0 and line.startswith("#version:")) or not line: + continue + bpe_merges.append(tuple(line.split())) + self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) + # NOTE: the cache can grow without bound and will get really large for long running processes + # (esp. for texts of language that do not use space between word, e.g. Chinese); technically + # not a memory leak but appears as one. + # GPT2Tokenizer has the same problem, so let's be consistent. + self.cache = {} + + self.pat = re.compile(PRETOKENIZE_REGEX) + + super().__init__( + errors=errors, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + unk_token=unk_token, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + @property + def vocab_size(self) -> int: + """Gets vocab size.""" + return len(self.encoder) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab + def get_vocab(self): + """Get vocab.""" + return dict(self.encoder, **self.added_tokens_encoder) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe + def bpe(self, token): + """Byte pair encoding.""" + if token in self.cache: + return self.cache[token] + word = tuple(token) + pairs = get_pairs(word) + + if not pairs: + return token + + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + else: + new_word.extend(word[i:j]) + i = j + + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize + def _tokenize(self, text): + """Tokenize a string.""" + bpe_tokens = [] + for token in re.findall(self.pat, text): + token = "".join( + self.byte_encoder[b] for b in token.encode("utf-8") + ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) + return bpe_tokens + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.encoder.get(token, self.encoder.get(self.unk_token)) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab.""" + return self.decoder.get(index) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + return bytearray([self.byte_decoder[c] for c in text]).decode( + "utf-8", errors=self.errors + ) + + def decode( + self, + token_ids, + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = False, + spaces_between_special_tokens: bool = False, + **kwargs, + ) -> str: + """Decode.""" + # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers + # and cannot be configured elsewhere, but it should default to False for Qwen2Tokenizer + return super().decode( + token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + spaces_between_special_tokens=spaces_between_special_tokens, + **kwargs, + ) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save only the vocabulary of the tokenizer (vocabulary). + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + if not os.path.isdir(save_directory): + raise NotADirectoryError( + f"Vocabulary path ({save_directory}) should be a directory" + ) + vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + merge_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["merges_file"], + ) + + with open(vocab_file, "w", encoding="utf-8") as f: + f.write( + json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + + "\n" + ) + + index = 0 + with open(merge_file, "w", encoding="utf-8") as writer: + writer.write("#version: 0.2\n") + for bpe_tokens, token_index in sorted( + self.bpe_ranks.items(), key=lambda kv: kv[1] + ): + if index != token_index: + # logger.warning( + # f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive." + # " Please check that the tokenizer is not corrupted!" + # ) + index = token_index + writer.write(" ".join(bpe_tokens) + "\n") + index += 1 + + return vocab_file, merge_file + + def prepare_for_tokenization(self, text, **kwargs): + """Prepare for tokenization.""" + text = unicodedata.normalize("NFC", text) + return (text, kwargs) diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2_fast.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2_fast.py new file mode 100644 index 00000000000..41f9cf63a03 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_qwen2_fast.py @@ -0,0 +1,146 @@ +# Copyright 2024 The Qwen team, Alibaba Group and The HuggingFace Inc. team. All rights reserved. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Define Fast QWen2 Tokenizer Class.""" + +from typing import Optional, Tuple + +from tokenizers import AddedToken + +from .tokenization_qwen2 import Qwen2Tokenizer +from .tokenization_utils_fast import PreTrainedTokenizerFast + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", + "tokenizer_file": "tokenizer.json", +} + + +MAX_MODEL_INPUT_SIZES = {"qwen/qwen-tokenizer": 32768} + + +class Qwen2TokenizerFast(PreTrainedTokenizerFast): + """Construct a "fast" Qwen2 tokenizer (backed by HuggingFace's *tokenizers* library). + + Based on byte-level Byte-Pair-Encoding. + + Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will + be encoded differently whether it is at the beginning of the sentence (without space) or not: + + ```python + >>> from transformers import Qwen2TokenizerFast + + >>> tokenizer = Qwen2TokenizerFast.from_pretrained("Qwen/Qwen-tokenizer") + >>> tokenizer("Hello world")["input_ids"] + [9707, 1879] + + >>> tokenizer(" Hello world")["input_ids"] + [21927, 1879] + ``` + This is expected. + + This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should + refer to this superclass for more information regarding those methods. + + Args: + vocab_file (`str`, *optional*): + Path to the vocabulary file. + merges_file (`str`, *optional*): + Path to the merges file. + tokenizer_file (`str`, *optional*): + Path to [tokenizers](https://github.com/huggingface/tokenizers) file (generally has a .json extension) that + contains everything needed to load the tokenizer. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. Not applicable to this tokenizer. + bos_token (`str`, *optional*): + The beginning of sequence token. Not applicable for this tokenizer. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The token used for padding, for example when batching sequences of different lengths. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + slow_tokenizer_class = Qwen2Tokenizer + + def __init__( + self, + vocab_file=None, + merges_file=None, + tokenizer_file=None, + unk_token="<|endoftext|>", + bos_token=None, + eos_token="<|endoftext|>", + pad_token="<|endoftext|>", + **kwargs, + ): + """Initializes the Qwen2TokenizerFast.""" + # We need to at least pass vocab_file and merges_file to base class + # in case a slow tokenizer needs to be initialized; other can be + # configured through files. + # following GPT2TokenizerFast, also adding unk_token, bos_token, and eos_token + + bos_token = ( + AddedToken( + bos_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(bos_token, str) + else bos_token + ) + eos_token = ( + AddedToken( + eos_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(eos_token, str) + else eos_token + ) + unk_token = ( + AddedToken( + unk_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(unk_token, str) + else unk_token + ) + pad_token = ( + AddedToken( + pad_token, lstrip=False, rstrip=False, special=True, normalized=False + ) + if isinstance(pad_token, str) + else pad_token + ) + + super().__init__( + vocab_file=vocab_file, + merges_file=merges_file, + tokenizer_file=tokenizer_file, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + **kwargs, + ) + + # Copied from transformers.models.gpt2.tokenization_gpt2_fast.GPT2TokenizerFast.save_vocabulary + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save only the vocabulary of the tokenizer (vocabulary). + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + files = self._tokenizer.model.save(save_directory, name=filename_prefix) + return tuple(files) diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils.py new file mode 100644 index 00000000000..15fbf37bb9a --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils.py @@ -0,0 +1,1072 @@ +# Copyright 2020 The HuggingFace Inc. team. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for python tokenizers.""" + +# flake8: noqa: C901 + +import itertools +import re +from collections import OrderedDict +from typing import Any, Dict, List, Optional, overload, Tuple, Union + +from .tokenization_utils_base import ( + AddedToken, + BatchEncoding, + ENCODE_KWARGS_DOCSTRING, + ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING, + EncodedInput, + EncodedInputPair, + INIT_TOKENIZER_DOCSTRING, + PreTokenizedInput, + PreTokenizedInputPair, + PreTrainedTokenizerBase, + TextInput, + TextInputPair, + TruncationStrategy, +) +from .utils import add_end_docstrings, PaddingStrategy, TensorType + +# Slow tokenizers are saved in a vocabulary plus three separated files +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +ADDED_TOKENS_FILE = "added_tokens.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" + + +class Trie: + """Trie in Python. + + Creates a Trie out of a list of words. The trie is used to split on `added_tokens` in one pass + Loose reference https://en.wikipedia.org/wiki/Trie + """ + + def __init__(self, *args): + """Initialize Trie.""" + self.data = {} + self._tokens = set() + self._termination_char = "" + self.update(*args) + + def update(self, *args): + """Updates the Trie with new tokens provided as arguments. + + Args: + *args: Variable number of words to be added to the Trie. + """ + for token in tuple(*args): + self.add(token) + + def add(self, word: str): + """Passes over every char (utf-8) on word and recursively adds it to the internal `data` trie representation. + + The special key `""` in `self._termination_char` is used to represent termination. + + This function is idempotent, adding twice the same word will leave the trie unchanged + + Example: + ```python + >>> trie = Trie() + >>> trie.add("Hello 友達") + >>> trie.data + {"H": {"e": {"l": {"l": {"o": {" ": {"友": {"達": {"": 1}}}}}}}}} + + >>> trie.add("Hello") + >>> trie.data + {"H": {"e": {"l": {"l": {"o": {"": 1, " ": {"友": {"達": {"": 1}}}}}}}}} + ``` + """ + if not word: + # Prevent empty string + return + + self._tokens.add(word) + ref = self.data + for char in word: + ref[char] = ref.setdefault(char, {}) + ref = ref[char] + ref[self._termination_char] = 1 + + def split(self, text: str) -> List[str]: + """Will look for the words added to the trie within `text`. + + Output is the original string splitted along the boundaries of the words found. + + This trie will match the longest possible word first ! + + Example: + ```python + >>> trie = Trie() + >>> trie.split("[CLS] This is a extra_id_100") + ["[CLS] This is a extra_id_100"] + + >>> trie.add("[CLS]") + >>> trie.add("extra_id_1") + >>> trie.add("extra_id_100") + >>> trie.split("[CLS] This is a extra_id_100") + ["[CLS]", " This is a ", "extra_id_100"] + ``` + """ + # indexes are counted left of the chars index. + # "hello", index 0, is left of h, index 1 is between h and e. + # index 5 is right of the "o". + + # States are going to capture every possible start (indexes as above) + # as keys, and have as values, a pointer to the position in the trie + # where we're at. This is a partial match for now. + # This enables to keep track of multiple matches while we're iterating + # the string + # If the trie contains, "blowing", and "lower" and we encounter the + # string "blower", we need to split into ["b", "lower"]. + # This is where we need to keep track of multiple possible starts. + states = OrderedDict() + + # This will contain every indices where we need + # to cut. + # We force to cut at offset 0 and len(text) (added later) + offsets = [0] + + # This is used by the lookahead which needs to skip over + # some text where the full match exceeded the place in the initial + # for loop + skip = 0 + # Main loop, Giving this algorithm O(n) complexity + for current, current_char in enumerate(text): + if skip and current < skip: + # Prevents the lookahead for matching twice + # like extra_id_100 and id_100 + continue + + # This will track every state + # that stop matching, we need to stop tracking them. + # If we look at "lowball", we're going to match "l" (add it to states), "o", "w", then + # fail on "b", we need to remove 0 from the valid states. + to_remove = set() + # Whenever we found a match, we need to drop everything + # this is a greedy algorithm, it will match on the first found token + reset = False + + # In this case, we already have partial matches (But unfinished) + for start, trie_pointer in states.items(): + if "" in trie_pointer: + # This is a final match, we need to reset and + # store the results in `offsets`. + + # Lookahead to match longest first + # Important in case of extra_id_1 vs extra_id_100 + # Here we are also actively looking for other earlier partial + # matches + # "[CLS]", "L", we need to match CLS even if L is special + for lookstart, looktrie_pointer in states.items(): + if lookstart > start: + # This partial match is later, we can stop looking + break + if lookstart < start: + # This partial match is earlier, the trie pointer + # was already updated, so index is + 1 + lookahead_index = current + 1 + end = current + 1 + else: + # Here lookstart == start and + # looktrie_pointer == trie_pointer + # It wasn't updated yet so indices are current ones + lookahead_index = current + end = current + next_char = ( + text[lookahead_index] + if lookahead_index < len(text) + else None + ) + if "" in looktrie_pointer: + start = lookstart + end = lookahead_index + skip = lookahead_index + + while next_char in looktrie_pointer: + looktrie_pointer = looktrie_pointer[next_char] + lookahead_index += 1 + if "" in looktrie_pointer: + start = lookstart + end = lookahead_index + skip = lookahead_index + + if lookahead_index == len(text): + # End of string + break + next_char = text[lookahead_index] + # End lookahead + + # Storing and resetting + offsets.append(start) + offsets.append(end) + reset = True + break + if current_char in trie_pointer: + # The current character being looked at has a match within the trie + # update the pointer (it will be stored back into states later). + trie_pointer = trie_pointer[current_char] + + # Storing back the new pointer into the states. + # Partial matches got longer by one. + states[start] = trie_pointer + else: + # The new character has not match in the trie, we need + # to stop keeping track of this partial match. + # We can't do it directly within the loop because of how + # python iteration works + to_remove.add(start) + + # Either clearing the full start (we found a real match) + # Or clearing only the partial matches that didn't work. + if reset: + states = {} + else: + for start in to_remove: + del states[start] + + # If this character is a starting character within the trie + # start keeping track of this partial match. + if current >= skip and current_char in self.data: + states[current] = self.data[current_char] + + # We have a cut at the end with states. + for start, trie_pointer in states.items(): + if "" in trie_pointer: + # This is a final match, we need to reset and + # store the results in `offsets`. + end = len(text) + offsets.append(start) + offsets.append(end) + # Longest cut is always the one with lower start so the first + # item so we need to break. + break + + return self.cut_text(text, offsets) + + def cut_text(self, text, offsets): + """Cut text.""" + # We have all the offsets now, we just need to do the actual splitting. + # We need to eventually add the first part of the string and the eventual + # last part. + offsets.append(len(text)) + tokens = [] + start = 0 + for end in offsets: + if start > end: + # logger.error( + # "There was a bug in Trie algorithm in tokenization. Attempting to recover. Please report it" + # " anyway." + # ) + continue + if start == end: + # This might happen if there's a match at index 0 + # we're also preventing zero-width cuts in case of two + # consecutive matches + continue + tokens.append(text[start:end]) + start = end + + return tokens + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizer(PreTrainedTokenizerBase): + """Base class for all slow tokenizers. + + Inherits from [`~tokenization_utils_base.PreTrainedTokenizerBase`]. + + Handle all the shared methods for tokenization and special tokens as well as methods downloading/caching/loading + pretrained tokenizers as well as adding tokens to the vocabulary. + + This class also contain the added tokens in a unified way on top of all tokenizers so we don't have to handle the + specific vocabulary augmentation methods of the various underlying dictionary structures (BPE, sentencepiece...). + """ + + def __init__(self, **kwargs): + """Initialize PreTrainedTokenizer.""" + # 1. Init the parent class + + self.tokens_trie = Trie() + + # 2. init `_added_tokens_decoder` if child class did not + if not hasattr(self, "_added_tokens_decoder"): + self._added_tokens_decoder: Dict[int, AddedToken] = {} + + # 3. if a `added_tokens_decoder` is passed, we are loading from a saved tokenizer, we overwrite + self._added_tokens_decoder.update(kwargs.pop("added_tokens_decoder", {})) + self._added_tokens_encoder: Dict[str, int] = { + k.content: v for v, k in self._added_tokens_decoder.items() + } + + # 4 init the parent class + super().__init__(**kwargs) + + # 4. If some of the special tokens are not part of the vocab, we add them, at the end. + # the order of addition is the same as self.SPECIAL_TOKENS_ATTRIBUTES following `tokenizers` + self._add_tokens( + [ + token + for token in self.all_special_tokens_extended + if token not in self._added_tokens_encoder + ], + special_tokens=True, + ) + + self._decode_use_source_tokenizer = False + + @property + def is_fast(self) -> bool: + """Is fast tokenizer property.""" + return False + + @property + def vocab_size(self) -> int: + """`int`: Size of the base vocabulary (without the added tokens).""" + raise NotImplementedError + + @property + def added_tokens_encoder(self) -> Dict[str, int]: + """Returns the sorted mapping from string to index. + + The added tokens encoder is cached for performance optimisation in `self._added_tokens_encoder` + for the slow tokenizers. + """ + return { + k.content: v + for v, k in sorted( + self._added_tokens_decoder.items(), key=lambda item: item[0] + ) + } + + @property + def added_tokens_decoder(self) -> Dict[int, AddedToken]: + """Returns the added tokens in the vocabulary as a dictionary of index to AddedToken. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return dict( + sorted(self._added_tokens_decoder.items(), key=lambda item: item[0]) + ) + + @added_tokens_decoder.setter + def added_tokens_decoder( + self, value: Dict[int, Union[AddedToken, str]] + ) -> Dict[int, AddedToken]: + # Always raise an error if string because users should define the behavior + for index, token in value.items(): + if not isinstance(token, (str, AddedToken)) or not isinstance(index, int): + raise TypeError( + f"The provided `added_tokens_decoder` has an element of type {index.__class__, token.__class__}, " + f"should be a dict of {int, Union[AddedToken, str]}" + ) + + self._added_tokens_decoder[index] = ( + AddedToken(token) if isinstance(token, str) else token + ) + self._added_tokens_encoder[str(token)] = index + self._update_total_vocab_size() + + def get_added_vocab(self) -> Dict[str, int]: + """Returns the added tokens in the vocabulary as a dictionary of token to index. + + Results might be different from the fast call because for now we always add the tokens + even if they are already in the vocabulary. This is something we should change. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return self._added_tokens_encoder + + def __len__(self): + """Size of the full vocabulary with the added tokens.""" + return self.total_vocab_size + + def _update_total_vocab_size(self): + """Update the size of the full vocabulary with the added tokens. + + Counts the `keys` and not the `values` because + otherwise if there is a hole in the vocab, we will add tokenizers at a wrong index. This operation is slow and + is only updated when adding tokens. + """ + self.total_vocab_size = len(self.get_vocab()) + + def _add_tokens( + self, + new_tokens: Union[List[str], List[AddedToken]], + special_tokens: bool = False, + ) -> int: + """Add a list of new tokens to the tokenizer class. + + If the new tokens are not in the vocabulary, they are added to + it with indices starting from length of the current vocabulary. Special tokens are sometimes already in the + vocab which is why they have to be handled specifically. + + Args: + new_tokens (`List[str]`or `List[tokenizers.AddedToken]`): + Token(s) to add in vocabulary. A token is counted as added if it's not already in the vocabulary + (tested by checking if the tokenizer assign the index of the `unk_token` to them). If a token is part + of the vocabulary then we simply mark this token as an `AddedToken` which allows to control the + stripping and normalization of this token. This is NOT possible in `tokenizers`. + special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the tokens should be added as special tokens. + + Returns: + `int`: The number of tokens actually added to the vocabulary. + + Examples: + ```python + # Let's see how to increase the vocabulary of Bert model and tokenizer + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + model = BertModel.from_pretrained("google-bert/bert-base-uncased") + + num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) + print("We have added", num_added_toks, "tokens") + # Note: resize_token_embeddings expects to receive the full size of the new vocabulary, i.e. the length of the + tokenizer.model.resize_token_embeddings(len(tokenizer)) + ``` + """ + added_tokens = 0 + if new_tokens is None: + return added_tokens + # TODO this is fairly slow to improve! + current_vocab = self.get_vocab().copy() + new_idx = len( + current_vocab + ) # only call this once, len gives the last index + 1 + for token in new_tokens: + if not isinstance(token, (str, AddedToken)): + raise TypeError(f"Token {token} is not a string but a {type(token)}.") + if str(token) == "": + continue + if isinstance(token, str): + if token in self._added_tokens_encoder: + continue + # very important for fast and slow equivalence! + is_special = token in self.all_special_tokens or special_tokens + token = AddedToken( + token, + rstrip=False, + lstrip=False, + normalized=not is_special, + special=is_special, + ) + elif special_tokens: + # doing token.special=True changes the normalization! will fix in rust + # this is important and the only reason why the AddedTokens in each class are normalized by default + token.__setstate__({"special": True, "normalized": token.normalized}) + if token in self._added_tokens_decoder: + continue + if ( + not token.special + and token.normalized + and getattr(self, "do_lower_case", False) + ): + # Normalize if requested + token.content = token.content.lower() + if token.content not in current_vocab: + token_index = new_idx + added_tokens + current_vocab[token.content] = token_index + added_tokens += 1 + else: + token_index = current_vocab[token.content] + + if token.special and str(token) not in self.all_special_tokens: + self._special_tokens_map["additional_special_tokens"].append(token) + # the setter automatically updates the reverse map + self._added_tokens_decoder[token_index] = token + self._added_tokens_encoder[token.content] = token_index + + self._update_trie() + self._update_total_vocab_size() + return added_tokens + + def _update_trie(self, unique_no_split_tokens: Optional[str] = []): + for token in self._added_tokens_decoder.values(): + if token not in self.tokens_trie._tokens: # noqa: SLF001 + self.tokens_trie.add(token.content) + for token in unique_no_split_tokens: + if token not in self.tokens_trie._tokens: # noqa: SLF001 + self.tokens_trie.add(token) + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + """Returns the number of added tokens when encoding a sequence with special tokens. + + + + This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put + this inside your training loop. + + + + Args: + pair (`bool`, *optional*, defaults to `False`): + Whether the number of added tokens should be computed in the case of a sequence pair or a single + sequence. + + Returns: + `int`: Number of special tokens added to sequences. + """ + token_ids_0 = [] + token_ids_1 = [] + return len( + self.build_inputs_with_special_tokens( + token_ids_0, token_ids_1 if pair else None + ) + ) + + def tokenize(self, text: TextInput, **kwargs) -> List[str]: + """Converts a string into a sequence of tokens, using the tokenizer. + + Split in words for word-based vocabulary or sub-words for sub-word-based vocabularies + (BPE/SentencePieces/WordPieces). Takes care of added tokens. + + Args: + text (`str`): + The sequence to be encoded. + **kwargs (additional keyword arguments): + Passed along to the model-specific `prepare_for_tokenization` preprocessing method. + + Returns: + `List[str]`: The list of tokens. + """ + split_special_tokens = kwargs.pop( + "split_special_tokens", self.split_special_tokens + ) + + text, kwargs = self.prepare_for_tokenization(text, **kwargs) + + if hasattr(self, "do_lower_case") and self.do_lower_case: + # convert non-special tokens to lowercase. Might be super slow as well? + escaped_special_toks = [ + re.escape(s_tok) for s_tok in (self.all_special_tokens) + ] + escaped_special_toks += [ + re.escape(s_tok.content) + for s_tok in (self._added_tokens_decoder.values()) + if not s_tok.special and s_tok.normalized + ] + pattern = r"(" + r"|".join(escaped_special_toks) + r")|" + r"(.+?)" + text = re.sub( + pattern, lambda m: m.groups()[0] or m.groups()[1].lower(), text + ) + + if split_special_tokens: + no_split_token = [] + tokens = [text] + else: + no_split_token = ( + self._added_tokens_encoder.keys() + ) # don't split on any of the added tokens + # "This is something else" + tokens = self.tokens_trie.split(text) + + # ["This is something", "", " else"] + for i, token in enumerate(tokens): + if token in no_split_token: + tok_extended = self._added_tokens_decoder.get( + self._added_tokens_encoder[token], None + ) + left = tokens[i - 1] if i > 0 else None + right = tokens[i + 1] if i < len(tokens) - 1 else None + if isinstance(tok_extended, AddedToken): + if tok_extended.rstrip and right: + # A bit counter-intuitive but we strip the left of the string + # since tok_extended.rstrip means the special token is eating all white spaces on its right + tokens[i + 1] = right.lstrip() # noqa: B909 + # Strip white spaces on the left + if tok_extended.lstrip and left: + tokens[i - 1] = left.rstrip() # Opposite here # noqa: B909 + if tok_extended.single_word and left and left[-1] != " ": + tokens[i - 1] += token + tokens[i] = "" + elif tok_extended.single_word and right and right[0] != " ": + tokens[i + 1] = token + tokens[i + 1] # noqa: B909 + tokens[i] = "" + else: + raise ValueError( + f"{tok_extended} cannot be tokenized because it was not properly added" + f" to the tokenizer. This means that it is not an `AddedToken` but a {type(tok_extended)}" + ) + # ["This is something", "", "else"] + tokenized_text = [] + for token in tokens: + # Need to skip eventual empty (fully stripped) tokens + if not token: + continue + if token in no_split_token: + tokenized_text.append(token) + else: + tokenized_text.extend(self._tokenize(token)) + # ["This", " is", " something", "", "else"] + return tokenized_text + + def _tokenize(self, text, **kwargs): + """Converts a string into a sequence of tokens (string), using the tokenizer. + + Split in words for word-based vocabulary or sub-words for sub-word-based vocabularies + (BPE/SentencePieces/WordPieces). + + Do NOT take care of added tokens. + """ + raise NotImplementedError + + def convert_tokens_to_ids( + self, tokens: Union[str, List[str]] + ) -> Union[int, List[int]]: + """Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids), using the vocabulary. + + Args: + tokens (`str` or `List[str]`): One or several token(s) to convert to token id(s). + + Returns: + `int` or `List[int]`: The token id or list of token ids. + """ # noqa: E501 + if tokens is None: + return None + + if isinstance(tokens, str): + return self._convert_token_to_id_with_added_voc(tokens) + + ids = [] + for token in tokens: + ids.append(self._convert_token_to_id_with_added_voc(token)) + return ids + + def _convert_token_to_id_with_added_voc(self, token): + if token is None: + return None + + if token in self._added_tokens_encoder: + return self._added_tokens_encoder[token] + return self._convert_token_to_id(token) + + def _convert_token_to_id(self, token): + raise NotImplementedError + + def _encode_plus( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + def get_input_ids(text): + if isinstance(text, str): + tokens = self.tokenize(text, **kwargs) + return self.convert_tokens_to_ids(tokens) + if ( + isinstance(text, (list, tuple)) + and len(text) > 0 + and isinstance(text[0], str) + ): + if is_split_into_words: + tokens = list( + itertools.chain( + *( + self.tokenize(t, is_split_into_words=True, **kwargs) + for t in text + ) + ) + ) + return self.convert_tokens_to_ids(tokens) + return self.convert_tokens_to_ids(text) + if ( + isinstance(text, (list, tuple)) + and len(text) > 0 + and isinstance(text[0], int) + ): + return text + if is_split_into_words: + raise ValueError( + f"Input {text} is not valid. Should be a string or a list/tuple of strings when" + " `is_split_into_words=True`." + ) + raise ValueError( + f"Input {text} is not valid. Should be a string, a list/tuple of strings or a list/tuple of integers." + ) + + if return_offsets_mapping: + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast. " + "More information on available tokenizers at " + "https://github.com/huggingface/transformers/pull/2674" + ) + + first_ids = get_input_ids(text) + second_ids = get_input_ids(text_pair) if text_pair is not None else None + + return self.prepare_for_model( + first_ids, + pair_ids=second_ids, + add_special_tokens=add_special_tokens, + padding=padding_strategy.value, + truncation=truncation_strategy.value, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + prepend_batch_axis=True, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + List[TextInput], + List[TextInputPair], + List[PreTokenizedInput], + List[PreTokenizedInputPair], + List[EncodedInput], + List[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + def get_input_ids(text): + if isinstance(text, str): + tokens = self.tokenize(text, **kwargs) + return self.convert_tokens_to_ids(tokens) + if ( + isinstance(text, (list, tuple)) + and len(text) > 0 + and isinstance(text[0], str) + ): + if is_split_into_words: + tokens = list( + itertools.chain( + *( + self.tokenize(t, is_split_into_words=True, **kwargs) + for t in text + ) + ) + ) + return self.convert_tokens_to_ids(tokens) + return self.convert_tokens_to_ids(text) + if ( + isinstance(text, (list, tuple)) + and len(text) > 0 + and isinstance(text[0], int) + ): + return text + raise ValueError( + "Input is not valid. Should be a string, a list/tuple of strings or a list/tuple of integers." + ) + + if return_offsets_mapping: + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast." + ) + + input_ids = [] + for ids_or_pair_ids in batch_text_or_text_pairs: + if not isinstance(ids_or_pair_ids, (list, tuple)) or ( + is_split_into_words + and not isinstance(ids_or_pair_ids[0], (list, tuple)) + ): + ids, pair_ids = ids_or_pair_ids, None + else: + ids, pair_ids = ids_or_pair_ids + + first_ids = get_input_ids(ids) + second_ids = get_input_ids(pair_ids) if pair_ids is not None else None + input_ids.append((first_ids, second_ids)) + + batch_outputs = self._batch_prepare_for_model( + input_ids, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=return_tensors, + verbose=verbose, + split_special_tokens=split_special_tokens, + ) + + return BatchEncoding(batch_outputs) + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING + ) + def _batch_prepare_for_model( + self, + batch_ids_pairs: List[Union[PreTokenizedInputPair, Tuple[List[int], None]]], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[str] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + ) -> BatchEncoding: + """Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. + + It adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens. + """ + batch_outputs = {} + for first_ids, second_ids in batch_ids_pairs: + outputs = self.prepare_for_model( + first_ids, + second_ids, + add_special_tokens=add_special_tokens, + padding=PaddingStrategy.DO_NOT_PAD.value, # we pad in batch afterward + truncation=truncation_strategy.value, + max_length=max_length, + stride=stride, + pad_to_multiple_of=None, # we pad in batch afterward + padding_side=None, # we pad in batch afterward + return_attention_mask=False, # we pad in batch afterward + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=None, # We convert the whole batch to tensors at the end + prepend_batch_axis=False, + verbose=verbose, + split_special_tokens=split_special_tokens, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + batch_outputs = self.pad( + batch_outputs, + padding=padding_strategy.value, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def prepare_for_tokenization( + self, text: str, is_split_into_words: bool = False, **kwargs + ) -> Tuple[str, Dict[str, Any]]: + """Performs any necessary transformations before tokenization. + + This method should pop the arguments from kwargs and return the remaining `kwargs` as well. We test the + `kwargs` at the end of the encoding process to be sure all the arguments have been used. + + Args: + text (`str`): + The text to prepare. + is_split_into_words (`bool`, *optional*, defaults to `False`): + Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the + tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace) + which it will tokenize. This is useful for NER or token classification. + kwargs (`Dict[str, Any]`, *optional*): + Keyword arguments to use for the tokenization. + + Returns: + `Tuple[str, Dict[str, Any]]`: The prepared text and the unused kwargs. + """ + return (text, kwargs) + + def get_special_tokens_mask( + self, + token_ids_0: List, + token_ids_1: Optional[List] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """Retrieves sequence ids from a token list that has no special tokens added. + + This method is called when adding + special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods. + + Args: + token_ids_0 (`List[int]`): + List of ids of the first sequence. + token_ids_1 (`List[int]`, *optional*): + List of ids of the second sequence. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + if token_ids_1 is not None: + raise ValueError( + "You should not supply a second sequence if the provided sequence of " + "ids is already formatted with special tokens for the model." + ) + + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + return [0] * ((len(token_ids_1) if token_ids_1 else 0) + len(token_ids_0)) + + @overload + def convert_ids_to_tokens( + self, ids: int, skip_special_tokens: bool = False + ) -> str: ... + + @overload + def convert_ids_to_tokens( + self, ids: List[int], skip_special_tokens: bool = False + ) -> List[str]: ... + + def convert_ids_to_tokens( + self, ids: Union[int, List[int]], skip_special_tokens: bool = False + ) -> Union[str, List[str]]: + """Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and added tokens. + + Args: + ids (`int` or `List[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `List[str]`: The decoded token(s). + """ # noqa: E501 + if isinstance(ids, int): + if ids in self._added_tokens_decoder: + return self._added_tokens_decoder[ids].content + return self._convert_id_to_token(ids) + tokens = [] + for index in ids: + index = int(index) + if skip_special_tokens and index in self.all_special_ids: + continue + if index in self._added_tokens_decoder: + tokens.append(self._added_tokens_decoder[index].content) + else: + tokens.append(self._convert_id_to_token(index)) + return tokens + + def _convert_id_to_token(self, index: int) -> str: + raise NotImplementedError + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Conver tokens to string.""" + return " ".join(tokens) + + def _decode( + self, + token_ids: Union[int, List[int]], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + spaces_between_special_tokens: bool = True, + **kwargs, + ) -> str: + self._decode_use_source_tokenizer = kwargs.pop("use_source_tokenizer", False) + + filtered_tokens = self.convert_ids_to_tokens( + token_ids, skip_special_tokens=skip_special_tokens + ) + # If given is a single id, prevents splitting the string in upcoming loop + if isinstance(filtered_tokens, str): + filtered_tokens = [filtered_tokens] + + legacy_added_tokens = set(self._added_tokens_encoder.keys()) - set( + self.all_special_tokens + ) | { + token + for token in self.additional_special_tokens + if self.convert_tokens_to_ids(token) >= self.vocab_size + } + # To avoid mixing byte-level and unicode for byte-level BPT + # we need to build string separately for added tokens and byte-level tokens + # cf. https://github.com/huggingface/transformers/issues/1133 + sub_texts = [] + current_sub_text = [] + for token in filtered_tokens: + if skip_special_tokens and token in self.all_special_tokens: + continue + if token in legacy_added_tokens: + if current_sub_text: + string = self.convert_tokens_to_string(current_sub_text) + if len(string) > 0: + sub_texts.append(string) + current_sub_text = [] + sub_texts.append(token) + else: + current_sub_text.append(token) + if current_sub_text: + sub_texts.append(self.convert_tokens_to_string(current_sub_text)) + + text = ( + " ".join(sub_texts) if spaces_between_special_tokens else "".join(sub_texts) + ) + + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + return self.clean_up_tokenization(text) + return text diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_base.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_base.py new file mode 100644 index 00000000000..e620c6f99cf --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_base.py @@ -0,0 +1,4415 @@ +# Copyright 2020 The HuggingFace Inc. team. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Base classes common to both the slow and the fast tokenization classes. + +PreTrainedTokenizerBase (host all the user fronting encoding methods), +Special token mixing (host the special tokens logic), and +BatchEncoding (wrap the dictionary of output with special method for the Fast tokenizers) +""" +# flake8: noqa: C901 + +import copy +import json +import os +import re +import warnings +from collections import UserDict +from collections.abc import Mapping, Sized +from contextlib import contextmanager +from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union + +import numpy as np +import torch +from tokenizers import AddedToken, Encoding as EncodingFast + +from .utils import ( + add_end_docstrings, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + ExplicitEnum, + is_flax_available, + is_jax_tensor, + is_mlx_available, + is_numpy_array, + is_offline_mode, + is_protobuf_available, + is_remote_url, + is_tf_available, + is_tf_tensor, + is_torch_available, + is_torch_device, + is_torch_tensor, + PaddingStrategy, + TensorType, + to_py_obj, +) + + +def import_protobuf_decode_error(*_): + """Import the DecodeError from protobuf if available. + + Args: + _ (Any): Placeholder argument. + + Returns: + DecodeError: The DecodeError class from protobuf. + + Raises: + ImportError: If protobuf is not available. + """ + if is_protobuf_available: + from google.protobuf.message import DecodeError + + return DecodeError + raise ImportError() + + +VERY_LARGE_INTEGER = int( + 1e30 +) # This is used to set the max input length for a model with infinite size input +LARGE_INTEGER = int( + 1e20 +) # This is used when we need something big but slightly smaller than VERY_LARGE_INTEGER + +# Define type aliases and NamedTuples +TextInput = str +PreTokenizedInput = List[str] +EncodedInput = List[int] +TextInputPair = Tuple[str, str] +PreTokenizedInputPair = Tuple[List[str], List[str]] +EncodedInputPair = Tuple[List[int], List[int]] + +# Define type aliases for text-related non-text modalities +AudioInput = Union[ + "np.ndarray", "torch.Tensor", List["np.ndarray"], List["torch.Tensor"] +] + +# Slow tokenizers used to be saved in three separated files +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +ADDED_TOKENS_FILE = "added_tokens.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" +CHAT_TEMPLATE_FILE = "chat_template.jinja" + +# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file +FULL_TOKENIZER_FILE = "tokenizer.json" +_re_tokenizer_file = re.compile(r"tokenizer\.(.*)\.json") + + +class TruncationStrategy(ExplicitEnum): + """Possible values for the `truncation` argument in [`PreTrainedTokenizerBase.__call__`]. + + Useful for tab-completion in + an IDE. + """ + + ONLY_FIRST = "only_first" + ONLY_SECOND = "only_second" + LONGEST_FIRST = "longest_first" + DO_NOT_TRUNCATE = "do_not_truncate" + + +class CharSpan(NamedTuple): + """Character span in the original string. + + Args: + start (`int`): Index of the first character in the original string. + end (`int`): Index of the character following the last character in the original string. + """ + + start: int + end: int + + +class TokenSpan(NamedTuple): + """Token span in an encoded string (list of tokens). + + Args: + start (`int`): Index of the first token in the span. + end (`int`): Index of the token following the last token in the span. + """ + + start: int + end: int + + +class BatchEncoding(UserDict): + """Batch Encoding Class. + + Holds the output of the [`~tokenization_utils_base.PreTrainedTokenizerBase.__call__`], + [`~tokenization_utils_base.PreTrainedTokenizerBase.encode_plus`] and + [`~tokenization_utils_base.PreTrainedTokenizerBase.batch_encode_plus`] methods (tokens, attention_masks, etc). + + This class is derived from a python dictionary and can be used as a dictionary. In addition, this class exposes + utility methods to map from word/character space to token space. + + Args: + data (`dict`, *optional*): + Dictionary of lists/arrays/tensors returned by the `__call__`/`encode_plus`/`batch_encode_plus` methods + ('input_ids', 'attention_mask', etc.). + encoding (`tokenizers.Encoding` or `Sequence[tokenizers.Encoding]`, *optional*): + If the tokenizer is a fast tokenizer which outputs additional information like mapping from word/character + space to token space the `tokenizers.Encoding` instance or list of instance (for batches) hold this + information. + tensor_type (`Union[None, str, TensorType]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + prepend_batch_axis (`bool`, *optional*, defaults to `False`): + Whether or not to add a batch axis when converting to tensors (see `tensor_type` above). Note that this + parameter has an effect if the parameter `tensor_type` is set, *otherwise has no effect*. + n_sequences (`Optional[int]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + """ + + def __init__( + self, + data: Optional[Dict[str, Any]] = None, + encoding: Optional[Union[EncodingFast, Sequence[EncodingFast]]] = None, + tensor_type: Union[str, TensorType, None] = None, + prepend_batch_axis: bool = False, + n_sequences: Optional[int] = None, + ): + """Initialize BatchEncoding.""" + super().__init__(data) + + if isinstance(encoding, EncodingFast): + encoding = [encoding] + + self._encodings = encoding + + if n_sequences is None and encoding is not None and len(encoding): + n_sequences = encoding[0].n_sequences + + self._n_sequences = n_sequences + + self.convert_to_tensors( + tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis + ) + + @property + def n_sequences(self) -> Optional[int]: + """Number of sequences. + + `Optional[int]`: The number of sequences used to generate each sample from the batch encoded in this + [`BatchEncoding`]. Currently can be one of `None` (unknown), `1` (a single sentence) or `2` (a pair of + sentences) + """ + return self._n_sequences + + @property + def is_fast(self) -> bool: + """Boolean if tokenizer is fast or not. + + `bool`: Indicate whether this [`BatchEncoding`] was generated from the result of a [`PreTrainedTokenizerFast`] + or not. + """ + return self._encodings is not None + + def __getitem__(self, item: Union[int, str]) -> Union[Any, EncodingFast]: + """Implement getter. + + If the key is a string, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', + etc.). + + If the key is an integer, get the `tokenizers.Encoding` for batch item with index `key`. + + If the key is a slice, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', etc.) + with the constraint of slice. + """ + if isinstance(item, str): + return self.data[item] + if self._encodings is not None: + return self._encodings[item] + if isinstance(item, slice): + return {key: self.data[key][item] for key in self.data} + raise KeyError( + "Invalid key. Only three types of key are available: " + "(1) string, (2) integers for backend Encoding, and (3) slices for data subsetting." + ) + + def __getattr__(self, item: str): + """Attribute getter.""" + try: + return self.data[item] + except KeyError as e: + raise AttributeError from e + + def __getstate__(self): + """State getter.""" + return {"data": self.data, "encodings": self._encodings} + + def __setstate__(self, state): + """State setter.""" + if "data" in state: + self.data = state["data"] + + if "encodings" in state: + self._encodings = state["encodings"] + + def keys(self): + """Get dict keys.""" + return self.data.keys() + + def values(self): + """Get dict values.""" + return self.data.values() + + def items(self): + """Get dict itmes.""" + return self.data.items() + + # After this point: + # Extended properties and methods only available for fast (Rust-based) tokenizers + # provided by HuggingFace tokenizers library. + + @property + def encodings(self) -> Optional[List[EncodingFast]]: + """Encodings property. + + `Optional[List[tokenizers.Encoding]]`: The list all encodings from the tokenization process. Returns `None` if + the input was tokenized through Python (i.e., not a fast) tokenizer. + """ + return self._encodings + + def tokens(self, batch_index: int = 0) -> List[str]: + """List of tokens. + + Return the list of tokens (sub-parts of the input strings after word/subword splitting and before conversion to + integer indices) at a given batch index (only works for the output of a fast tokenizer). + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[str]`: The list of tokens at that index. + """ + if not self._encodings: + raise ValueError( + "tokens() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].tokens + + def sequence_ids(self, batch_index: int = 0) -> List[Optional[int]]: + """Return a list mapping the tokens to the id of their original sentences. + + - `None` for special tokens added around or between sequences, + - `0` for tokens corresponding to words in the first sequence, + - `1` for tokens corresponding to words in the second sequence when a pair of sequences was jointly + encoded. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the sequence id corresponding to each token. Special tokens added + by the tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding + sequence. + """ + if not self._encodings: + raise ValueError( + "sequence_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].sequence_ids + + def words(self, batch_index: int = 0) -> List[Optional[int]]: + """Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the + tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word + (several tokens will be mapped to the same word index if they are parts of that word). + """ + if not self._encodings: + raise ValueError( + "words() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast` class)." + ) + warnings.warn( + "`BatchEncoding.words()` property is deprecated and should be replaced with the identical, " + "but more self-explanatory `BatchEncoding.word_ids()` property.", + FutureWarning, + stacklevel=2, + ) + return self.word_ids(batch_index) + + def word_ids(self, batch_index: int = 0) -> List[Optional[int]]: + """Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the + tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word + (several tokens will be mapped to the same word index if they are parts of that word). + """ + if not self._encodings: + raise ValueError( + "word_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].word_ids + + def token_to_sequence( + self, batch_or_token_index: int, token_index: Optional[int] = None + ) -> int: + """Get the index of the sequence represented by the given token. + + In the general use case, this method returns `0` + for a single sequence or the first sequence of a pair, and `1` for the second sequence of a pair + + Can be called as: + + - `self.token_to_sequence(token_index)` if batch size is 1 + - `self.token_to_sequence(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + if not self._encodings: + raise ValueError( + "token_to_sequence() is not available when using Python based tokenizers" + ) + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_sequence(token_index) + + def token_to_word( + self, batch_or_token_index: int, token_index: Optional[int] = None + ) -> int: + """Get the index of the word corresponding (i.e. comprising) to an encoded token in a sequence of the batch. + + Can be called as: + + - `self.token_to_word(token_index)` if batch size is 1 + - `self.token_to_word(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + if not self._encodings: + raise ValueError( + "token_to_word() is not available when using Python based tokenizers" + ) + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_word(token_index) + + def word_to_tokens( + self, + batch_or_word_index: int, + word_index: Optional[int] = None, + sequence_index: int = 0, + ) -> Optional[TokenSpan]: + """Get the encoded token span corresponding to a word in a sequence of the batch. + + Token spans are returned as a [`~tokenization_utils_base.TokenSpan`] with: + + - **start** -- Index of the first token. + - **end** -- Index of the token following the last token. + + Can be called as: + + - `self.word_to_tokens(word_index, sequence_index: int = 0)` if batch size is 1 + - `self.word_to_tokens(batch_index, word_index, sequence_index: int = 0)` if batch size is greater or equal to + 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the word in the sequence. + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + ([`~tokenization_utils_base.TokenSpan`], *optional*): Span of tokens in the encoded sequence. Returns + `None` if no tokens correspond to the word. This can happen especially when the token is a special token + that has been used to format the tokenization. For example when we add a class token at the very beginning + of the tokenization. + """ + if not self._encodings: + raise ValueError( + "word_to_tokens() is not available when using Python based tokenizers" + ) + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if word_index < 0: + word_index = self._seq_len + word_index + span = self._encodings[batch_index].word_to_tokens(word_index, sequence_index) + return TokenSpan(*span) if span is not None else None + + def token_to_chars( + self, batch_or_token_index: int, token_index: Optional[int] = None + ) -> CharSpan: + """Get the character span corresponding to an encoded token in a sequence of the batch. + + Character spans are returned as a [`~tokenization_utils_base.CharSpan`] with: + + - **start** -- Index of the first character in the original string associated to the token. + - **end** -- Index of the character following the last character in the original string associated to the + token. + + Can be called as: + + - `self.token_to_chars(token_index)` if batch size is 1 + - `self.token_to_chars(batch_index, token_index)` if batch size is greater or equal to 1 + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token or tokens in + the sequence. + + Returns: + [`~tokenization_utils_base.CharSpan`]: Span of characters in the original string, or None, if the token + (e.g. , ) doesn't correspond to any chars in the origin string. + """ + if not self._encodings: + raise ValueError( + "token_to_chars() is not available when using Python based tokenizers" + ) + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + span_indices = self._encodings[batch_index].token_to_chars(token_index) + + return CharSpan(*span_indices) if span_indices is not None else None + + def char_to_token( + self, + batch_or_char_index: int, + char_index: Optional[int] = None, + sequence_index: int = 0, + ) -> int: + """Get the index of the token in the output comprising a character in the original string. + + Can be called as: + + - `self.char_to_token(char_index)` if batch size is 1 + - `self.char_to_token(batch_index, char_index)` if batch size is greater or equal to 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int`: Index of the token, or None if the char index refers to a whitespace only token and whitespace is + trimmed with `trim_offsets=True`. + """ + if not self._encodings: + raise ValueError( + "char_to_token() is not available when using Python based tokenizers" + ) + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_token(char_index, sequence_index) + + def word_to_chars( + self, + batch_or_word_index: int, + word_index: Optional[int] = None, + sequence_index: int = 0, + ) -> CharSpan: + """Get the character span in the original string corresponding to given word in a sequence of the batch. + + Character spans are returned as a CharSpan NamedTuple with: + + - start: index of the first character in the original string + - end: index of the character following the last character in the original string + + Can be called as: + + - `self.word_to_chars(word_index)` if batch size is 1 + - `self.word_to_chars(batch_index, word_index)` if batch size is greater or equal to 1 + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + `CharSpan` or `List[CharSpan]`: Span(s) of the associated character or characters in the string. CharSpan + are NamedTuple with: + + - start: index of the first character associated to the token in the original string + - end: index of the character following the last character associated to the token in the original + string + """ + if not self._encodings: + raise ValueError( + "word_to_chars() is not available when using Python based tokenizers" + ) + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + return CharSpan( + *(self._encodings[batch_index].word_to_chars(word_index, sequence_index)) + ) + + def char_to_word( + self, + batch_or_char_index: int, + char_index: Optional[int] = None, + sequence_index: int = 0, + ) -> int: + """Get the word in the original string corresponding to a character in the original string. + + Can be called as: + + - `self.char_to_word(char_index)` if batch size is 1 + - `self.char_to_word(batch_index, char_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the character in the original string. + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the character in the + original string. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int` or `List[int]`: Index or indices of the associated encoded token(s). + """ + if not self._encodings: + raise ValueError( + "char_to_word() is not available when using Python based tokenizers" + ) + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_word(char_index, sequence_index) + + def convert_to_tensors( + self, + tensor_type: Optional[Union[str, TensorType]] = None, + prepend_batch_axis: bool = False, + ): + """Convert the inner content to tensors. + + Args: + tensor_type (`str` or [`~utils.TensorType`], *optional*): + The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If + `None`, no modification is done. + prepend_batch_axis (`int`, *optional*, defaults to `False`): + Whether or not to add the batch dimension during the conversion. + """ + if tensor_type is None: + return self + + # Convert to TensorType + if not isinstance(tensor_type, TensorType): + tensor_type = TensorType(tensor_type) + + # Get a function reference for the correct framework + if tensor_type == TensorType.TENSORFLOW: + if not is_tf_available(): + raise ImportError( + "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed." + ) + import tensorflow as tf + + as_tensor = tf.constant + is_tensor = tf.is_tensor + elif tensor_type == TensorType.PYTORCH: + if not is_torch_available: + raise ImportError( + "Unable to convert output to PyTorch tensors format, PyTorch is not installed." + ) + import torch + + is_tensor = torch.is_tensor + + def as_tensor(value, dtype=None): + if isinstance(value, list) and isinstance(value[0], np.ndarray): + return torch.from_numpy(np.array(value)) + return torch.tensor(value) + + elif tensor_type == TensorType.JAX: + if not is_flax_available(): + raise ImportError( + "Unable to convert output to JAX tensors format, JAX is not installed." + ) + import jax.numpy as jnp + + as_tensor = jnp.array + is_tensor = is_jax_tensor + + elif tensor_type == TensorType.MLX: + if not is_mlx_available(): + raise ImportError( + "Unable to convert output to MLX tensors format, MLX is not installed." + ) + import mlx.core as mx + + as_tensor = mx.array + + def is_tensor(obj): + return isinstance(obj, mx.array) + + else: + + def as_tensor(value, dtype=None): + if isinstance(value, (list, tuple)) and isinstance( + value[0], (list, tuple, np.ndarray) + ): + value_lens = [len(val) for val in value] + if len(set(value_lens)) > 1 and dtype is None: + # we have a ragged list so handle explicitly + value = as_tensor( + [np.asarray(val) for val in value], dtype=object + ) + return np.asarray(value, dtype=dtype) + + is_tensor = is_numpy_array + + # Do the tensor conversion in batch + for key, value in self.items(): + try: + if prepend_batch_axis: + value = [value] + + if not is_tensor(value): + tensor = as_tensor(value) + + # Removing this for now in favor of controlling the shape with `prepend_batch_axis` + # # at-least2d + # if tensor.ndim > 2: + # tensor = tensor.squeeze(0) + # elif tensor.ndim < 2: + # tensor = tensor[None, :] + + self[key] = tensor + except Exception as e: + if key == "overflowing_tokens": + raise ValueError( + "Unable to create tensor returning overflowing tokens of different lengths. " + "Please see if a fast version of this tokenizer is available to have this feature available." + ) from e + raise ValueError( + "Unable to create tensor, you should probably activate truncation and/or padding with" + " 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your" + f" features (`{key}` in this case) have excessive nesting (inputs type `list` where type `int` is" + " expected)." + ) from e + + return self + + def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding": + """Send all values to device by calling `v.to(device)` (PyTorch only). + + Args: + device (`str` or `torch.device`): The device to put the tensors on. + + Returns: + [`BatchEncoding`]: The same instance after modification. + """ + import torch + + # This check catches things like APEX blindly calling "to" on all inputs to a module + # Otherwise it passes the casts down and casts the LongTensor containing the token idxs + # into a HalfTensor + if ( + isinstance(device, str) + or is_torch_device(device) + or isinstance(device, int) + ): + self.data = { + k: v.to(device=device) if isinstance(v, torch.Tensor) else v + for k, v in self.data.items() + } + else: + print( + f"Attempting to cast a BatchEncoding to type {device!s}. This is not supported." + ) + return self + + +class SpecialTokensMixin: + """Special Tokens Mixin. + + A mixin derived by [`PreTrainedTokenizer`] and [`PreTrainedTokenizerFast`] to handle specific behaviors related to + special tokens. In particular, this class hold the attributes which can be used to directly access these special + tokens in a model-independent manner and allow to set and update the special tokens. + + Args: + bos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the beginning of a sentence. + eos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the end of a sentence. + unk_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing an out-of-vocabulary token. + sep_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token separating two different sentences in the same input (used by BERT for instance). + pad_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. + cls_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the class of the input (used by BERT for instance). + mask_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing a masked token (used by masked-language modeling pretraining objectives, like + BERT). + additional_special_tokens (tuple or list of `str` or `tokenizers.AddedToken`, *optional*): + A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be + skipped when decoding if `skip_special_tokens` is set to `True`. + """ + + SPECIAL_TOKENS_ATTRIBUTES = [ + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", + "additional_special_tokens", + ] + + def __init__(self, verbose=False, **kwargs): + """Initialize SpecialTokensMixin class.""" + self._pad_token_type_id = 0 + self.verbose = verbose + self.special_tokens_map_ = dict.fromkeys(self.SPECIAL_TOKENS_ATTRIBUTES) + self.special_tokens_map_["additional_special_tokens"] = ( + [] + ) # for BC where it defaults to empty list + + # We directly set the hidden value to allow initialization with special tokens + # which are not yet in the vocabulary. Necessary for serialization/de-serialization + # TODO clean this up at some point (probably by switching to fast tokenizers) + + for key, value in kwargs.items(): + if value is None: + continue + if key in self.SPECIAL_TOKENS_ATTRIBUTES: + if key == "additional_special_tokens": + assert isinstance( + value, (list, tuple) + ), f"Value {value} is not a list or tuple" + assert all( + isinstance(t, (str, AddedToken)) for t in value + ), "One of the tokens is not a string or an AddedToken" + setattr(self, key, value) + elif isinstance(value, (str, AddedToken)): + setattr(self, key, value) + else: + raise TypeError( + f"Special token {key} has to be either str or AddedToken but got: {type(value)}" + ) + + def sanitize_special_tokens(self) -> int: + """The `sanitize_special_tokens` is now deprecated.""" + print("The `sanitize_special_tokens` will be removed in transformers v5.") + return self.add_tokens(self.all_special_tokens_extended, special_tokens=True) + + def add_special_tokens( + self, + special_tokens_dict: Dict[str, Union[str, AddedToken]], + replace_additional_special_tokens=True, + ) -> int: + """Adds special tokens. + + Add a dictionary of special tokens (eos, pad, cls, etc.) to the encoder and link them to class attributes. If + special tokens are NOT in the vocabulary, they are added to it (indexed starting from the last index of the + current vocabulary). + + When adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix of the + model so that its embedding matrix matches the tokenizer. + + In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method. + + Using `add_special_tokens` will ensure your special tokens can be used in several ways: + + - Special tokens can be skipped when decoding using `skip_special_tokens = True`. + - Special tokens are carefully handled by the tokenizer (they are never split), similar to `AddedTokens`. + - You can easily refer to special tokens using tokenizer class attributes like `tokenizer.cls_token`. This + makes it easy to develop model-agnostic training and fine-tuning scripts. + + When possible, special tokens are already registered for provided pretrained models (for instance + [`BertTokenizer`] `cls_token` is already registered to be :obj*'[CLS]'* and XLM's one is also registered to be + `''`). + + Args: + special_tokens_dict (dictionary *str* to *str* or `tokenizers.AddedToken`): + Keys should be in the list of predefined special attributes: [`bos_token`, `eos_token`, `unk_token`, + `sep_token`, `pad_token`, `cls_token`, `mask_token`, `additional_special_tokens`]. + + Tokens are only added if they are not already in the vocabulary (tested by checking if the tokenizer + assign the index of the `unk_token` to them). + replace_additional_special_tokens (`bool`, *optional*,, defaults to `True`): + If `True`, the existing list of additional special tokens will be replaced by the list provided in + `special_tokens_dict`. Otherwise, `self.special_tokens_map_["additional_special_tokens"]` is + just extended. In the former case, the tokens will NOT be removed from the + tokenizer's full vocabulary - they are only being flagged + as non-special tokens. Remember, this only affects which tokens are skipped during decoding, not the + `added_tokens_encoder` and `added_tokens_decoder`. This means that the previous + `additional_special_tokens` are still added tokens, and will not be split by the model. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + ```python + # Let's see how to add a new classification token to GPT-2 + tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2") + model = GPT2Model.from_pretrained("openai-community/gpt2") + + special_tokens_dict = {"cls_token": ""} + + num_added_toks = tokenizer.add_special_tokens(special_tokens_dict) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., + the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + + assert tokenizer.cls_token == "" + ``` + """ + if not special_tokens_dict: + return 0 + + added_tokens = [] + for key, value in special_tokens_dict.items(): + assert ( + key in self.SPECIAL_TOKENS_ATTRIBUTES + ), f"Key {key} is not a special token" + + if self.verbose: + print(f"Assigning {value} to the {key} key of the tokenizer") + + if key == "additional_special_tokens": + assert isinstance(value, (list, tuple)) and all( + isinstance(t, (str, AddedToken)) for t in value + ), f"Tokens {value} for key {key} should all be str or AddedToken instances" + + to_add = [] + for token in value: + if isinstance(token, str): + # for legacy purpose we default to stripping. `test_add_tokens_tokenizer` depends on this + token = AddedToken( + token, + rstrip=False, + lstrip=False, + normalized=False, + special=True, + ) + if ( + not replace_additional_special_tokens + and str(token) in self.additional_special_tokens + ): + continue + to_add.append(token) + if replace_additional_special_tokens and len(to_add) > 0: + setattr(self, key, list(to_add)) + else: + self.special_tokens_map_["additional_special_tokens"].extend(to_add) + added_tokens += to_add + + else: + if not isinstance(value, (str, AddedToken)): + raise ValueError( + f"Token {value} for key {key} should be a str or an AddedToken instance" + ) + if isinstance(value, (str)): + # for legacy purpose we default to stripping. `False` depends on this + value = AddedToken( + value, + rstrip=False, + lstrip=False, + normalized=False, + special=True, + ) + if isinstance(value, AddedToken): + setattr(self, key, value) + if value not in added_tokens: + added_tokens.append(value) + + # if we are adding tokens that were not part of the vocab, we ought to add them + return self.add_tokens(added_tokens, special_tokens=True) + + def add_tokens( + self, + new_tokens: Union[str, AddedToken, List[Union[str, AddedToken]]], + special_tokens: bool = False, + ) -> int: + """Add tokens. + + Add a list of new tokens to the tokenizer class. If the new tokens are not in the vocabulary, they are added to + it with indices starting from length of the current vocabulary and will be isolated before the tokenization + algorithm is applied. Added tokens and tokens from the vocabulary of the tokenization algorithm are therefore + not treated in the same way. + + Note, when adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix + of the model so that its embedding matrix matches the tokenizer. + + In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method. + + Args: + new_tokens (`str`, `tokenizers.AddedToken` or a list of *str* or `tokenizers.AddedToken`): + Tokens are only added if they are not already in the vocabulary. `tokenizers.AddedToken` wraps a string + token to let you personalize its behavior: whether this token should only match against a single word, + whether this token should strip all potential whitespaces on the left side, whether this token should + strip all potential whitespaces on the right side, etc. + special_tokens (`bool`, *optional*, defaults to `False`): + Can be used to specify if the token is a special token. This mostly change the normalization behavior + (special tokens like CLS or [MASK] are usually not lower-cased for instance). + + See details for `tokenizers.AddedToken` in HuggingFace tokenizers library. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + ```python + # Let's see how to increase the vocabulary of Bert model and tokenizer + tokenizer = BertTokenizerFast.from_pretrained("google-bert/bert-base-uncased") + model = BertModel.from_pretrained("google-bert/bert-base-uncased") + + num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., + the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + ``` + """ + if not new_tokens: + return 0 + + if not isinstance(new_tokens, (list, tuple)): + new_tokens = [new_tokens] + + return self._add_tokens(new_tokens, special_tokens=special_tokens) + + def _add_tokens( + self, + new_tokens: Union[List[str], List[AddedToken]], + special_tokens: bool = False, + ) -> int: + raise NotImplementedError + + @property + def pad_token_type_id(self) -> int: + """`int`: Id of the padding token type in the vocabulary.""" + return self._pad_token_type_id + + def __setattr__(self, key, value): + """Attribute setter.""" + key_without_id = key + key_is_special_id = key.endswith(("_id", "_ids")) + if key_is_special_id: + key_without_id = key[:-3] if not key.endswith("_ids") else key[:-4] + + if self.__dict__.get("special_tokens_map_", None) is not None and any( + name in self.__dict__["special_tokens_map_"] + for name in [key, key_without_id] + ): + if key_is_special_id: + if value is not None: + value = ( + self.convert_ids_to_tokens(value) + if key != "additional_special_tokens" + else [self.convert_ids_to_tokens(val) for val in value] + ) + key = key_without_id + + if ( + key != "additional_special_tokens" + and not isinstance(value, (str, AddedToken)) + and value is not None + ): + raise ValueError(f"Cannot set a non-string value as the {key}") + self.special_tokens_map_[key] = value + else: + super().__setattr__(key, value) + + def __getattr__(self, key): + """Attribute getter.""" + key_without_id = key + key_is_special_id = key.endswith(("_id", "_ids")) + if key_is_special_id: + key_without_id = key[:-3] if not key.endswith("_ids") else key[:-4] + + if self.__dict__.get("special_tokens_map_", None) is not None and any( + name in self.__dict__["special_tokens_map_"] + for name in [key, key_without_id] + ): + special_tokens_map_ = self.__dict__["special_tokens_map_"] + if not key_is_special_id: + if special_tokens_map_[key] is None: + return None + value = special_tokens_map_[key] + return ( + str(value) + if key != "additional_special_tokens" + else [str(tok) for tok in value] + ) + attr_as_tokens = getattr(self, key_without_id) + return ( + self.convert_tokens_to_ids(attr_as_tokens) + if attr_as_tokens is not None + else None + ) + + if key not in self.__dict__: + raise AttributeError(f"{self.__class__.__name__} has no attribute {key}") + return super().__getattr__(key) + + @property + def special_tokens_map(self) -> Dict[str, Union[str, List[str]]]: + """Special tokens map. + + `Dict[str, Union[str, List[str]]]`: A dictionary mapping special token class attributes (`cls_token`, + `unk_token`, etc.) to their values (`''`, `''`, etc.). + + Convert potential tokens of `tokenizers.AddedToken` type to string. + """ + set_attr = {} + for attr in self.SPECIAL_TOKENS_ATTRIBUTES: + attr_value = getattr(self, attr) + if attr_value: + set_attr[attr] = attr_value + return set_attr + + @property + def special_tokens_map_extended( + self, + ) -> Dict[str, Union[str, AddedToken, List[Union[str, AddedToken]]]]: + """Extended special tokens map. + + `Dict[str, Union[str, tokenizers.AddedToken, List[Union[str, tokenizers.AddedToken]]]]`: A dictionary mapping + special token class attributes (`cls_token`, `unk_token`, etc.) to their values (`''`, `''`, etc.). + + Don't convert tokens of `tokenizers.AddedToken` type to string so they can be used to control more finely how + special tokens are tokenized. + """ + set_attr = {} + for attr in self.SPECIAL_TOKENS_ATTRIBUTES: + attr_value = self.special_tokens_map_[attr] + if attr_value: + set_attr[attr] = attr_value + return set_attr + + @property + def all_special_tokens_extended(self) -> List[Union[str, AddedToken]]: + """All extended special tokens map. + + `List[Union[str, tokenizers.AddedToken]]`: All the special tokens (`''`, `''`, etc.), the order has + nothing to do with the index of each tokens. If you want to know the correct indices, check + `self.added_tokens_encoder`. We can't create an order anymore as the keys are `AddedTokens` and not `Strings`. + + Don't convert tokens of `tokenizers.AddedToken` type to string so they can be used to control more finely how + special tokens are tokenized. + """ + all_tokens = [] + seen = set() + for value in self.special_tokens_map_extended.values(): + if isinstance(value, (list, tuple)): + tokens_to_add = [token for token in value if str(token) not in seen] + else: + tokens_to_add = [value] if str(value) not in seen else [] + seen.update(map(str, tokens_to_add)) + all_tokens.extend(tokens_to_add) + return all_tokens + + @property + def all_special_tokens(self) -> List[str]: + """All special tokens. + + `List[str]`: A list of the unique special tokens (`''`, `''`, ..., etc.). + + Convert tokens of `tokenizers.AddedToken` type to string. + """ + return [str(s) for s in self.all_special_tokens_extended] + + @property + def all_special_ids(self) -> List[int]: + """All special ids. + + `List[int]`: List the ids of the special tokens(`''`, `''`, etc.) mapped to class attributes. + """ + all_toks = self.all_special_tokens + return self.convert_tokens_to_ids(all_toks) + + def _set_model_specific_special_tokens(self, special_tokens: List[str]): + """Sets model specific tokens. + + Adds new special tokens to the "SPECIAL_TOKENS_ATTRIBUTES" list which will be part + of "self.special_tokens" and saved as a special token in tokenizer's config. + This allows us to dynamically add new model-type specific tokens after initilizing the tokenizer. + For example: if the model tokenizers is multimodal, we can support special image or audio tokens. + """ + self.SPECIAL_TOKENS_ATTRIBUTES = self.SPECIAL_TOKENS_ATTRIBUTES + list( + special_tokens.keys() + ) + for key, value in special_tokens.items(): + if isinstance(value, (str, AddedToken)): + self.special_tokens_map_[key] = value + else: + raise TypeError( + f"Special token {key} has to be either str or AddedToken but got: {type(value)}" + ) + + +ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to add special tokens when encoding the sequences. This will use the underlying + `PretrainedTokenizerBase.build_inputs_with_special_tokens` function, which defines which tokens are + automatically added to the input ids. This is usefull if you want to add `bos` or `eos` tokens + automatically. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + stride (`int`, *optional*, defaults to 0): + If set to a number along with `max_length`, the overflowing tokens returned when + `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence + returned to provide some overlap between truncated and overflowing sequences. The value of this + argument defines the number of overlapping tokens. + is_split_into_words (`bool`, *optional*, defaults to `False`): + Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the + tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace) + which it will tokenize. This is useful for NER or token classification. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. Requires `padding` to be activated. + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. +""" # noqa: E501 + +ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) @lint-ignore + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) @lint-ignore + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) @lint-ignore + + - **token_type_ids** -- List of token type ids to be fed to a model (when `return_token_type_ids=True` or + if *"token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) @lint-ignore + + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) @lint-ignore + + - **overflowing_tokens** -- List of overflowing tokens sequences (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **num_truncated_tokens** -- Number of tokens truncated (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **special_tokens_mask** -- List of 0s and 1s, with 1 specifying added special tokens and 0 specifying + regular sequence tokens (when `add_special_tokens=True` and `return_special_tokens_mask=True`). + - **length** -- The length of the inputs (when `return_length=True`) +""" + + +INIT_TOKENIZER_DOCSTRING = r""" + Class attributes (overridden by derived classes) + + - **vocab_files_names** (`Dict[str, str]`) -- A dictionary with, as keys, the `__init__` keyword name of each + vocabulary file required by the model, and as associated values, the filename for saving the associated file + (string). + - **pretrained_vocab_files_map** (`Dict[str, Dict[str, str]]`) -- A dictionary of dictionaries, with the + high-level keys being the `__init__` keyword name of each vocabulary file required by the model, the + low-level being the `short-cut-names` of the pretrained models with, as associated values, the `url` to the + associated pretrained vocabulary file. + - **model_input_names** (`List[str]`) -- A list of inputs expected in the forward pass of the model. + - **padding_side** (`str`) -- The default value for the side on which the model should have padding applied. + Should be `'right'` or `'left'`. + - **truncation_side** (`str`) -- The default value for the side on which the model should have truncation + applied. Should be `'right'` or `'left'`. + + Args: + model_max_length (`int`, *optional*): + The maximum length (in number of tokens) for the inputs to the transformer model. When the tokenizer is + loaded with [`~tokenization_utils_base.PreTrainedTokenizerBase.from_pretrained`], this will be set to the + value stored for the associated model in `max_model_input_sizes` (see above). If no value is provided, will + default to VERY_LARGE_INTEGER (`int(1e30)`). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + truncation_side (`str`, *optional*): + The side on which the model should have truncation applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + chat_template (`str`, *optional*): + A Jinja template string that will be used to format lists of chat messages. See + https://huggingface.co/docs/transformers/chat_templating for a full description. + model_input_names (`List[string]`, *optional*): + The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or + `"attention_mask"`). Default value is picked from the class attribute of the same name. + bos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the beginning of a sentence. Will be associated to `self.bos_token` and + `self.bos_token_id`. + eos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the end of a sentence. Will be associated to `self.eos_token` and + `self.eos_token_id`. + unk_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing an out-of-vocabulary token. Will be associated to `self.unk_token` and + `self.unk_token_id`. + sep_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token separating two different sentences in the same input (used by BERT for instance). Will be + associated to `self.sep_token` and `self.sep_token_id`. + pad_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. Will be associated to `self.pad_token` and `self.pad_token_id`. + cls_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the class of the input (used by BERT for instance). Will be associated to + `self.cls_token` and `self.cls_token_id`. + mask_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing a masked token (used by masked-language modeling pretraining objectives, like + BERT). Will be associated to `self.mask_token` and `self.mask_token_id`. + additional_special_tokens (tuple or list of `str` or `tokenizers.AddedToken`, *optional*): + A tuple or a list of additional special tokens. Add them here to ensure they are skipped when decoding with + `skip_special_tokens` is set to True. If they are not part of the vocabulary, they will be added at the end + of the vocabulary. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`): + Whether or not the model should cleanup the spaces that were added when splitting the input text during the + tokenization process. + split_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the special tokens should be split during the tokenization process. Passing will affect the + internal state of the tokenizer. The default behavior is to not split special tokens. This means that if + `` is the `bos_token`, then `tokenizer.tokenize("") = ['`]. Otherwise, if + `split_special_tokens=True`, then `tokenizer.tokenize("")` will be give `['<','s', '>']`. +""" + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizerBase(SpecialTokensMixin): + """Base class for [`PreTrainedTokenizer`] and [`PreTrainedTokenizerFast`]. + + Handles shared (mostly boiler plate) methods for those two classes. + """ + + vocab_files_names: Dict[str, str] = {} + pretrained_vocab_files_map: Dict[str, Dict[str, str]] = {} + _auto_class: Optional[str] = None + + # first name has to correspond to main model input name + # to make sure `tokenizer.pad(...)` works correctly + model_input_names: List[str] = ["input_ids", "token_type_ids", "attention_mask"] + padding_side: str = "right" + truncation_side: str = "right" + slow_tokenizer_class = None + + def __init__(self, **kwargs): + """Initialize PreTrainedTokenizerBase.""" + # inputs and kwargs for saving and re-loading (see ``from_pretrained`` and ``save_pretrained``) + self.init_inputs = () + for key in kwargs: + if hasattr(self, key) and callable(getattr(self, key)): + raise AttributeError( + f"{key} conflicts with the method {key} in {self.__class__.__name__}" + ) + + self.init_kwargs = copy.deepcopy(kwargs) + self.name_or_path = kwargs.pop("name_or_path", "") + self._processor_class = kwargs.pop("processor_class", None) + + # For backward compatibility we fallback to set model_max_length from max_len if provided + model_max_length = kwargs.pop("model_max_length", kwargs.pop("max_len", None)) + self.model_max_length = ( + model_max_length if model_max_length is not None else VERY_LARGE_INTEGER + ) + + # Padding and truncation side are right by default and overridden in subclasses. If specified in the kwargs, it + # is changed. + self.padding_side = kwargs.pop("padding_side", self.padding_side) + if self.padding_side not in ["right", "left"]: + raise ValueError( + f"Padding side should be selected between 'right' and 'left', current value: {self.padding_side}" + ) + + self.truncation_side = kwargs.pop("truncation_side", self.truncation_side) + if self.truncation_side not in ["right", "left"]: + raise ValueError( + f"Truncation side should be selected between 'right' and 'left', current value: {self.truncation_side}" + ) + + self.model_input_names = kwargs.pop("model_input_names", self.model_input_names) + + # By default, cleaning tokenization spaces for both fast and slow tokenizers + self.clean_up_tokenization_spaces = kwargs.pop( + "clean_up_tokenization_spaces", False + ) + + # By default, do not split special tokens for both fast and slow tokenizers + self.split_special_tokens = kwargs.pop("split_special_tokens", False) + + self.deprecation_warnings = ( + {} + ) # Use to store when we have already noticed a deprecation warning. + self._in_target_context_manager = False + + # Stores a Jinja template that formats chat histories into tokenizable strings + self.chat_template = kwargs.pop("chat_template", None) + if isinstance(self.chat_template, (list, tuple)): + # Chat templates are stored as lists of dicts with fixed key names, + # we reconstruct that into a single dict while loading them. + self.chat_template = { + template["name"]: template["template"] + for template in self.chat_template + } + + super().__init__(**kwargs) + + self.extra_special_tokens = kwargs.pop("extra_special_tokens", {}) + self._set_model_specific_special_tokens( + special_tokens=self.extra_special_tokens + ) + + @property + def max_len_single_sentence(self) -> int: + """`int`: The maximum length of a sentence that can be fed to the model.""" + return self.model_max_length - self.num_special_tokens_to_add(pair=False) + + @property + def max_len_sentences_pair(self) -> int: + """`int`: The maximum combined length of a pair of sentences that can be fed to the model.""" + return self.model_max_length - self.num_special_tokens_to_add(pair=True) + + @max_len_single_sentence.setter + def max_len_single_sentence(self, value) -> int: + # For backward compatibility, allow to try to setup 'max_len_single_sentence'. + if ( + value == self.model_max_length - self.num_special_tokens_to_add(pair=False) + and self.verbose + ): + if not self.deprecation_warnings.get("max_len_single_sentence", False): + print( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_single_sentence"] = True + else: + raise ValueError( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + + @max_len_sentences_pair.setter + def max_len_sentences_pair(self, value) -> int: + # For backward compatibility, allow to try to setup 'max_len_sentences_pair'. + if ( + value == self.model_max_length - self.num_special_tokens_to_add(pair=True) + and self.verbose + ): + if not self.deprecation_warnings.get("max_len_sentences_pair", False): + print( + "Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_sentences_pair"] = True + else: + raise ValueError( + "Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up." + ) + + def _set_processor_class(self, processor_class: str): + """Sets processor class as an attribute.""" + self._processor_class = processor_class + + @property + def added_tokens_decoder(self) -> Dict[int, AddedToken]: + """Added tokens decoder property.""" + raise NotImplementedError() + + def __repr__(self) -> str: + """Alternate representation.""" + added_tokens_decoder_rep = "\n\t".join( + [f"{k}: {v.__repr__()}," for k, v in self.added_tokens_decoder.items()] + ) + return ( + f"{self.__class__.__name__}(name_or_path='{self.name_or_path}'," + f" vocab_size={self.vocab_size}, model_max_length={self.model_max_length}, is_fast={self.is_fast}," + f" padding_side='{self.padding_side}', truncation_side='{self.truncation_side}'," + f" special_tokens={self.special_tokens_map}, " + f"clean_up_tokenization_spaces={self.clean_up_tokenization_spaces}," + " added_tokens_decoder={\n\t" + added_tokens_decoder_rep + "\n}\n)" + ) + + def __len__(self) -> int: + """Length property.""" + raise NotImplementedError() + + def get_vocab(self) -> Dict[str, int]: + """Returns the vocabulary as a dictionary of token to index. + + `tokenizer.get_vocab()[token]` is equivalent to `tokenizer.convert_tokens_to_ids(token)` when `token` is in the + vocab. + + Returns: + `Dict[str, int]`: The vocabulary. + """ + raise NotImplementedError() + + def apply_chat_template( + self, + conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]], + tools: Optional[List[Dict]] = None, + documents: Optional[List[Dict[str, str]]] = None, + chat_template: Optional[str] = None, + add_generation_prompt: bool = False, + continue_final_message: bool = False, + tokenize: bool = True, + padding: bool = False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_dict: bool = False, + return_assistant_tokens_mask: bool = False, + tokenizer_kwargs: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]: + """Converts a list of dictionaries with `"role"` and `"content"` keys to a list of token ids. + + This method is intended for use with chat models, and will read the tokenizer's chat_template attribute to + determine the format and control tokens to use when converting. + + Args: + conversation (Union[List[Dict[str, str]], List[List[Dict[str, str]]]]): A list of dicts + with "role" and "content" keys, representing the chat history so far. + tools (`List[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON + Schema, giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + documents (`List[Dict[str, str]]`, *optional*): + A list of dicts representing documents that will be accessible to the model if it is performing RAG + (retrieval-augmented generation). If the template does not support RAG, this argument will have no + effect. We recommend that each document should be a dict containing "title" and "text" keys. Please + see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG) + for examples of passing documents with chat templates. + chat_template (`str`, *optional*): + A Jinja template to use for this conversion. It is usually not necessary to pass anything to this + argument, as the model's template will be used by default. + add_generation_prompt (bool, *optional*): + If this is set, a prompt with the token(s) that indicate + the start of an assistant message will be appended to the formatted output. This is useful when you want + to generate a response from the model. + Note that this argument will be passed to the chat template, and so it must be supported in the + template for this argument to have any effect. + continue_final_message (bool, *optional*): + If this is set, the chat will be formatted so that the final + message in the chat is open-ended, without any EOS tokens. The model will continue this message + rather than starting a new one. This allows you to "prefill" part of + the model's response for it. Cannot be used at the same time as `add_generation_prompt`. + tokenize (`bool`, defaults to `True`): + Whether to tokenize the output. If `False`, the output will be a string. + padding (`bool`, defaults to `False`): + Whether to pad sequences to the maximum length. Has no effect if tokenize is `False`. + truncation (`bool`, defaults to `False`): + Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`. + max_length (`int`, *optional*): + Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If + not specified, the tokenizer's `max_length` attribute will be used as a default. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable + values are: + - `'tf'`: Return TensorFlow `tf.Tensor` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + return_dict (`bool`, defaults to `False`): + Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`. + tokenizer_kwargs (`Dict[str: Any]`, *optional*): Additional kwargs to pass to the tokenizer. + return_assistant_tokens_mask (`bool`, defaults to `False`): + Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant, + the mask will contain 1. For user and system tokens, the mask will contain 0. + This functionality is only available for chat templates that support it via the `{% generation %}` + keyword. + **kwargs: Additional kwargs to pass to the template renderer. Will be accessible by the chat template. + + Returns: + `Union[List[int], Dict]`: A list of token ids representing the tokenized chat so far, including control + tokens. This + output is ready to pass to the model, either directly or via methods like `generate()`. If `return_dict` is + set, will return a dict of tokenizer outputs instead. + """ + if return_dict and not tokenize: + raise ValueError( + "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict " + "of tokenizer outputs to return." + ) + + if return_assistant_tokens_mask and not return_dict: + raise ValueError( + "`return_assistant_tokens_mask=True` is incompatible with `return_dict=False`" + ) + + if tokenizer_kwargs is None: + tokenizer_kwargs = {} + + chat_template = self.get_chat_template(chat_template, tools) + + if return_assistant_tokens_mask and not re.search( + r"\{\%-?\s*generation\s*-?\%\}", chat_template + ): + print( + "return_assistant_tokens_mask==True but chat template does not contain `{% generation %}` keyword." + ) + + # Compilation function uses a cache to avoid recompiling the same template + # compiled_template = _compile_jinja_template(chat_template) + compiled_template = chat_template + + if isinstance(conversation, (list, tuple)) and ( + isinstance(conversation[0], (list, tuple)) + or hasattr(conversation[0], "messages") + ): + conversations = conversation + is_batched = True + else: + conversations = [conversation] + is_batched = False + + if continue_final_message: + if add_generation_prompt: + raise ValueError( + "continue_final_message and add_generation_prompt are not compatible. Use continue_final_message " + "when you want the model to continue the final message, and add_generation_prompt when you want to " + "add a header that will prompt it to start a new assistant message instead." + ) + if return_assistant_tokens_mask: + raise ValueError( + "continue_final_message is not compatible with return_assistant_tokens_mask." + ) + + # We accept either JSON schemas or functions for tools. If we get functions, we convert them to schemas + if tools is not None: + tool_schemas = [] + for tool in tools: + if isinstance(tool, dict): + tool_schemas.append(tool) + else: + raise ValueError( + "Tools should either be a JSON schema, or a callable function with type hints " + "and a docstring suitable for auto-conversion to a schema." + ) + else: + tool_schemas = None + + if documents is not None: + for document in documents: + if not isinstance(document, dict): + raise TypeError( + "Documents should be a list of dicts with 'title' and 'text' keys!" + ) + + rendered = [] + all_generation_indices = [] + template_kwargs = { + **self.special_tokens_map, + **kwargs, + } # kwargs overwrite special tokens if both are present + for chat in conversations: + if hasattr(chat, "messages"): + # Indicates it's a Conversation object + chat = chat.messages + if return_assistant_tokens_mask: + raise NotImplementedError() + rendered_chat = compiled_template.render( + messages=chat, + tools=tool_schemas, + documents=documents, + add_generation_prompt=add_generation_prompt, + **template_kwargs, + ) + if continue_final_message: + final_message = chat[-1]["content"] + if isinstance(final_message, (list, tuple)): + final_message = final_message[-1]["text"] + try: + rendered_chat = rendered_chat[ + : rendered_chat.rindex(final_message) + len(final_message) + ] + except: # noqa: E722 + # Some chat templates like Llama-3.1 trim messages before rendering, so we must do the same here. + final_message = final_message.strip() + rendered_chat = rendered_chat[ + : rendered_chat.rindex(final_message) + len(final_message) + ] + rendered.append(rendered_chat) + + if not is_batched: + rendered = rendered[0] + + if tokenize: + out = self( + rendered, + padding=padding, + truncation=truncation, + max_length=max_length, + add_special_tokens=False, + return_tensors=return_tensors, + **tokenizer_kwargs, + ) + if return_dict: + if return_assistant_tokens_mask: + assistant_masks = [] + input_ids = ( + out["input_ids"] + if is_batched or return_tensors + else [out["input_ids"]] + ) + for i in range(len(input_ids)): + current_mask = [0] * len(input_ids[i]) + for ( + assistant_start_char, + assistant_end_char, + ) in all_generation_indices[i]: + start_token = out.char_to_token(i, assistant_start_char) + end_token = out.char_to_token(i, assistant_end_char - 1) + if start_token is None: + # start_token is out of bounds maybe due to truncation. + break + for token_id in range( + start_token, + end_token + 1 if end_token else len(input_ids[i]), + ): + current_mask[token_id] = 1 + assistant_masks.append(current_mask) + out["assistant_masks"] = ( + assistant_masks if is_batched else assistant_masks[0] + ) + return out + return out["input_ids"] + return rendered + + def get_chat_template( + self, chat_template: Optional[str] = None, tools: Optional[List[Dict]] = None + ) -> str: + """Retrieve the chat template string used for tokenizing chat messages. + + This template is used + internally by the `apply_chat_template` method and can also be used externally to retrieve the model's chat + template for better generation tracking. + + Args: + chat_template (`str`, *optional*): + A Jinja template or the name of a template to use for this conversion. + It is usually not necessary to pass anything to this argument, + as the model's template will be used by default. + tools (`List[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON + Schema, giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + + Returns: + `str`: The chat template string. + """ + # First, handle the cases when the model has a dict of multiple templates + if isinstance(self.chat_template, dict): + template_dict = self.chat_template + if chat_template is not None and chat_template in template_dict: + # The user can pass the name of a template to the chat template argument instead of an entire template + chat_template = template_dict[chat_template] + elif chat_template is None: + if tools is not None and "tool_use" in template_dict: + chat_template = template_dict["tool_use"] + elif "default" in template_dict: + chat_template = template_dict["default"] + else: + raise ValueError( + "This model has multiple chat templates with no default specified! Please either pass a chat " + "template or the name of the template you wish to use to the `chat_template` argument." + f" Available template names are {sorted(template_dict.keys())}." + ) + + elif chat_template is None: + # These are the cases when the model has a single template + # priority: `chat_template` argument > `tokenizer.chat_template` + if self.chat_template is not None: + chat_template = self.chat_template + else: + raise ValueError( + "Cannot use chat template functions because tokenizer.chat_template is not set and no template " + "argument was passed! For information about writing templates and setting the " + "tokenizer.chat_template attribute, please see the documentation at " + "https://huggingface.co/docs/transformers/main/en/chat_templating" + ) + + return chat_template + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + *init_inputs, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + trust_remote_code=False, + **kwargs, + ): + """Instantiate a [`~tokenization_utils_base.PreTrainedTokenizerBase`] from a predefined tokenizer. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + + - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co. + - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved + using the [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`] method, e.g., + `./my_model_directory/`. + - (**Deprecated**, not applicable to all derived classes) A path or url to a single saved vocabulary + file (if and only if the tokenizer only requires a single vocabulary file like Bert or XLNet), e.g., + `./my_model_directory/vocab.txt`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded predefined tokenizer vocabulary files should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download the vocabulary files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request. @lint-ignore + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + local_files_only (`bool`, *optional*, defaults to `False`): + Whether or not to only rely on local files and not to attempt to download any files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + subfolder (`str`, *optional*): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for + facebook/rag-token-base), specify it here. + inputs (additional positional arguments, *optional*): + Will be passed along to the Tokenizer `__init__` method. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether or not to allow for custom models defined on the Hub in their own modeling files. This option + should only be set to `True` for repositories you trust and in which you have read the code, as it will + execute code present on the Hub on your local machine. + init_inputs (additional arguments): . + kwargs (additional keyword arguments, *optional*): + Will be passed to the Tokenizer `__init__` method. Can be used to set special tokens like `bos_token`, + `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`, + `additional_special_tokens`. See parameters in the `__init__` for more details. + + + + Passing `token=True` is required when you want to use a private model. + + + + Examples: + ```python + # We can't instantiate directly the base class *PreTrainedTokenizerBase* so let's show our examples on a derived + # class: BertTokenizer + # Download vocabulary from huggingface.co and cache. + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + + # Download vocabulary from huggingface.co (user-uploaded) and cache. + tokenizer = BertTokenizer.from_pretrained("dbmdz/bert-base-german-cased") + + # If vocabulary files are in a directory (e.g. tokenizer was saved using + # *save_pretrained('./test/saved_model/')*) + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/") + + # If the tokenizer uses a single vocabulary file, you can point directly to this file + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/my_vocab.txt") + + # You can link tokens to special vocabulary when instantiating + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased", unk_token="") + # You should be sure '' is in the vocabulary when doing that. + # Otherwise use tokenizer.add_special_tokens({'unk_token': ''}) instead) + assert tokenizer.unk_token == "" + ``` + """ + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + use_auth_token = kwargs.pop("use_auth_token", None) + subfolder = kwargs.pop("subfolder", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + commit_hash = kwargs.pop("_commit_hash", None) + gguf_file = kwargs.get("gguf_file") + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. " + "Please use `token` instead.", + FutureWarning, + stacklevel=2, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + user_agent = { + "file_type": "tokenizer", + "from_auto_class": from_auto_class, + "is_fast": "Fast" in cls.__name__, + } + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode and not local_files_only: + # print("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + vocab_files = {} + init_configuration = {} + + is_local = os.path.isdir(pretrained_model_name_or_path) + single_file_id = None + if os.path.isfile(pretrained_model_name_or_path) or is_remote_url( + pretrained_model_name_or_path + ): + if len(cls.vocab_files_names) > 1 and not gguf_file: + raise ValueError( + f"Calling {cls.__name__}.from_pretrained() with the path to a single file or url is not " + "supported for this tokenizer. Use a model identifier or the path to a directory instead." + ) + warnings.warn( + f"Calling {cls.__name__}.from_pretrained() with the path to a single file or url is deprecated and " + "won't be possible anymore in v5. Use a model identifier or the path to a directory instead.", + FutureWarning, + stacklevel=2, + ) + file_id = next(iter(cls.vocab_files_names.keys())) + + vocab_files[file_id] = pretrained_model_name_or_path + single_file_id = file_id + else: + if gguf_file: + vocab_files["vocab_file"] = gguf_file + else: + # At this point pretrained_model_name_or_path is either a directory or a model identifier name + additional_files_names = { + "added_tokens_file": ADDED_TOKENS_FILE, # kept only for legacy + "special_tokens_map_file": SPECIAL_TOKENS_MAP_FILE, # kept only for legacy + "tokenizer_config_file": TOKENIZER_CONFIG_FILE, + # tokenizer_file used to initialize a slow from a fast. + # Properly copy the `addedTokens` instead of adding in random orders + "tokenizer_file": FULL_TOKENIZER_FILE, + "chat_template_file": CHAT_TEMPLATE_FILE, + } + vocab_files = {**cls.vocab_files_names, **additional_files_names} + if "tokenizer_file" in vocab_files: + # Try to get the tokenizer config to see if there are versioned tokenizer files. + fast_tokenizer_file = FULL_TOKENIZER_FILE + resolved_config_file = cached_file( + pretrained_model_name_or_path, + TOKENIZER_CONFIG_FILE, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + token=token, + revision=revision, + local_files_only=local_files_only, + subfolder=subfolder, + user_agent=user_agent, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + _commit_hash=commit_hash, + ) + # commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + if resolved_config_file is not None: + with open(resolved_config_file, encoding="utf-8") as reader: + tokenizer_config = json.load(reader) + if "fast_tokenizer_files" in tokenizer_config: + fast_tokenizer_file = get_fast_tokenizer_file( + tokenizer_config["fast_tokenizer_files"] + ) + vocab_files["tokenizer_file"] = fast_tokenizer_file + + # Get files from url, cache, or disk depending on the case + resolved_vocab_files = {} + unresolved_files = [] + for file_id, file_path in vocab_files.items(): + if file_path is None: + resolved_vocab_files[file_id] = None + elif single_file_id == file_id: + if os.path.isfile(file_path): + resolved_vocab_files[file_id] = file_path + else: + resolved_vocab_files[file_id] = cached_file( + pretrained_model_name_or_path, + file_path, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + _commit_hash=commit_hash, + ) + # commit_hash = extract_commit_hash(resolved_vocab_files[file_id], commit_hash) + + if len(unresolved_files) > 0: + print( + f"Can't load following files from cache: {unresolved_files} and cannot check if these " + "files are necessary for the tokenizer to operate." + ) + + # If one passes a GGUF file path to `gguf_file` there is no need for this check as the tokenizer will be + # loaded directly from the GGUF file. + if ( + all( + full_file_name is None + for full_file_name in resolved_vocab_files.values() + ) + and not gguf_file + ): + raise OSError( + f"Can't load tokenizer for '{pretrained_model_name_or_path}'. If you were trying to load it from " + "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " + f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " + f"containing all relevant files for a {cls.__name__} tokenizer." + ) + + for file_id, _ in vocab_files.items(): + if file_id not in resolved_vocab_files: + continue + + # if is_local: + # print(f"loading file {file_path}") + # else: + # print(f"loading file {file_path} from cache at {resolved_vocab_files[file_id]}") + + return cls._from_pretrained( + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + _commit_hash=commit_hash, + _is_local=is_local, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def _from_pretrained( + cls, + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=None, + cache_dir=None, + local_files_only=False, + _commit_hash=None, + _is_local=False, + trust_remote_code=False, + **kwargs, + ): + # We instantiate fast tokenizers based on a slow tokenizer if we don't have access to the tokenizer.json + # file or if `from_slow` is set to True. + from_slow = kwargs.get("from_slow", False) + gguf_file = kwargs.get("gguf_file") + has_tokenizer_file = ( + resolved_vocab_files.get("tokenizer_file", None) is not None + ) + + # If one passes a GGUF file path to `gguf_file` there is no need for this check as the tokenizer will be + # loaded directly from the GGUF file. + if ( + (from_slow or not has_tokenizer_file) + and cls.slow_tokenizer_class is not None + and not gguf_file + ): + slow_tokenizer = ( + cls.slow_tokenizer_class + )._from_pretrained( # noqa: SLF001 + copy.deepcopy(resolved_vocab_files), + pretrained_model_name_or_path, + copy.deepcopy(init_configuration), + *init_inputs, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + _commit_hash=_commit_hash, + **(copy.deepcopy(kwargs)), + ) + else: + slow_tokenizer = None + + # Prepare tokenizer initialization kwargs + # Did we saved some inputs and kwargs to reload ? + tokenizer_config_file = resolved_vocab_files.pop("tokenizer_config_file", None) + if tokenizer_config_file is not None: + with open( + tokenizer_config_file, encoding="utf-8" + ) as tokenizer_config_handle: + init_kwargs = json.load(tokenizer_config_handle) + # First attempt. We get tokenizer_class from tokenizer_config to check mismatch between tokenizers. + config_tokenizer_class = init_kwargs.get("tokenizer_class") + init_kwargs.pop("tokenizer_class", None) + if not has_tokenizer_file: + init_kwargs.pop("tokenizer_file", None) + saved_init_inputs = init_kwargs.pop("init_inputs", ()) + if not init_inputs: + init_inputs = saved_init_inputs + else: + config_tokenizer_class = None + init_kwargs = init_configuration + + # If an independent chat template file exists, it takes priority over template entries in the tokenizer config + chat_template_file = resolved_vocab_files.pop("chat_template_file", None) + if chat_template_file is not None: + with open(chat_template_file) as chat_template_handle: + init_kwargs["chat_template"] = ( + chat_template_handle.read() + ) # Clobbers any template in the config + + if not _is_local: + if "auto_map" in init_kwargs: + # For backward compatibility with odl format. + if isinstance(init_kwargs["auto_map"], (tuple, list)): + init_kwargs["auto_map"] = {"AutoTokenizer": init_kwargs["auto_map"]} + init_kwargs["auto_map"] = add_model_info_to_auto_map( + init_kwargs["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in init_kwargs: + init_kwargs["custom_pipelines"] = add_model_info_to_custom_pipelines( + init_kwargs["custom_pipelines"], pretrained_model_name_or_path + ) + + if config_tokenizer_class is None: + # Matt: This entire block is only used to decide if the tokenizer class matches the class in the repo. + # If not, it raises a warning, but otherwise continues. Since we mostly load tokenizers with + # AutoTokenizer these days, it seems like a lot of work (and a source of bugs) for little gain. + # Maybe we can just remove this entirely? + raise NotImplementedError("Unable to auto-detect tokenizer") + from .models.auto.configuration_auto import AutoConfig # tests_ignore + + # Second attempt. If we have not yet found tokenizer_class, let's try to use the config. + try: + config = AutoConfig.from_pretrained( + pretrained_model_name_or_path, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + trust_remote_code=trust_remote_code, + _commit_hash=_commit_hash, + ) + config_tokenizer_class = config.tokenizer_class + except (OSError, ValueError, KeyError): + # skip if an error occurred. + config = None + if config_tokenizer_class is None: + # Third attempt. If we have not yet found the original type of the tokenizer, + # we are loading we see if we can infer it from the type of the configuration file + from .models.auto.tokenization_auto import ( # tests_ignore + TOKENIZER_MAPPING_NAMES, + ) + + if hasattr(config, "model_type"): + model_type = config.model_type + else: + # Fallback: use pattern matching on the string. + model_type = None + for pattern in TOKENIZER_MAPPING_NAMES: + if pattern in str(pretrained_model_name_or_path): + model_type = pattern + break + + if model_type is not None: + config_tokenizer_class, config_tokenizer_class_fast = ( + TOKENIZER_MAPPING_NAMES.get(model_type, (None, None)) + ) + if config_tokenizer_class is None: + config_tokenizer_class = config_tokenizer_class_fast + + if config_tokenizer_class is not None and cls.__name__.replace( + "Fast", "" + ) != config_tokenizer_class.replace("Fast", ""): + print( + "The tokenizer class you load from this checkpoint is not the same type as the class this" + " function is called from. It may result in unexpected tokenization. \nThe tokenizer class you" + f" load from this checkpoint is '{config_tokenizer_class}'. \nThe class this function is called" + f" from is '{cls.__name__}'." + ) + + # Update with newly provided kwargs + init_kwargs.update(kwargs) + + # Merge resolved_vocab_files arguments in init_kwargs. + added_tokens_file = resolved_vocab_files.pop("added_tokens_file", None) + special_tokens_map_file = resolved_vocab_files.pop( + "special_tokens_map_file", None + ) + for args_name, file_path in resolved_vocab_files.items(): + if args_name not in init_kwargs: + init_kwargs[args_name] = file_path + tokenizer_file = resolved_vocab_files.pop("tokenizer_file", None) + + if slow_tokenizer is not None: + init_kwargs["__slow_tokenizer"] = slow_tokenizer + init_kwargs["name_or_path"] = pretrained_model_name_or_path + + # Handle tokenizer serialization of added and special tokens + added_tokens_decoder: Dict[int, AddedToken] = {} + added_tokens_map: Dict[str, AddedToken] = {} + # if we have info on the slow added tokens + if "added_tokens_decoder" in init_kwargs: + for idx, token in init_kwargs["added_tokens_decoder"].items(): + if isinstance(token, dict): + token = AddedToken(**token) + if isinstance(token, AddedToken): + added_tokens_decoder[int(idx)] = token + added_tokens_map[str(token)] = token + else: + raise ValueError( + f"Found a {token.__class__} in the saved `added_tokens_decoder`, should be a dictionary or " + "an AddedToken instance" + ) + else: + # begin legacy: read the added_tokens_file and update kwargs with special_tokens_map if modified + if special_tokens_map_file is not None: + with open( + special_tokens_map_file, encoding="utf-8" + ) as special_tokens_map_handle: + special_tokens_map = json.load(special_tokens_map_handle) + for key, value in special_tokens_map.items(): + if kwargs.get(key): + # This value has already been redefined by the kwargs + # We keep this new value and ignore the one stored in the special_tokens_map_file + continue + if isinstance(value, dict): + value["special"] = True + value = AddedToken(**value) + elif key == "additional_special_tokens" and isinstance( + value, list + ): + additional_special_tokens = ( + init_kwargs.pop("additional_special_tokens", []) or [] + ) + for token in value: + if isinstance(token, dict): + token["special"] = True + token = AddedToken(**token) + if token not in additional_special_tokens: + additional_special_tokens.append(token) + value = additional_special_tokens + init_kwargs[key] = value + + # slow -> slow|fast, legacy: convert the `"added_tokens.json"` file to `added_tokens_decoder`. + # this is for legacy purpose. We don't add the tokens after init for efficiency. + if added_tokens_file is not None: + special_tokens = [] + for key in cls.SPECIAL_TOKENS_ATTRIBUTES & init_kwargs.keys(): + if init_kwargs[key] is not None: + if key == "additional_special_tokens": + special_tokens += [str(token) for token in init_kwargs[key]] + else: + special_tokens.append(str(init_kwargs[key])) + + with open(added_tokens_file, encoding="utf-8") as added_tokens_handle: + added_tok_encoder = json.load(added_tokens_handle) + for str_token, index in added_tok_encoder.items(): + # if index not in added_tokens_decoder and str_token not in added_tokens_map: + special = str_token in special_tokens + added_tokens_decoder[index] = AddedToken( + str_token, + rstrip=False, + lstrip=False, + normalized=not special, + special=special, + ) + added_tokens_map[str(token)] = added_tokens_decoder[index] + + # allows converting a fast -> slow: add the `tokenizer.json`'s `"added_tokens"` to the slow tokenizer + # if `tokenizer_config.json` is `None` + if tokenizer_file is not None: + # This is for slow so can be done before + with open(tokenizer_file, encoding="utf-8") as tokenizer_file_handle: + tokenizer_file_handle = json.load(tokenizer_file_handle) + added_tokens = tokenizer_file_handle.pop("added_tokens") + for serialized_tokens in added_tokens: + idx = serialized_tokens.pop("id") + added_tokens_decoder[idx] = AddedToken(**serialized_tokens) + added_tokens_map[str(added_tokens_decoder[idx])] = ( + added_tokens_decoder[idx] + ) + # end legacy + + # Passing AddedTokens and not strings to the class to prevent it from casting the string to a different AddedToken # noqa: E501 + # convert {'__type': 'AddedToken', 'content': '', 'lstrip': False, 'normalized': True, ...} to AddedTokens + init_kwargs["added_tokens_decoder"] = added_tokens_decoder + init_kwargs = cls.convert_added_tokens(init_kwargs, save=False) + for key in cls.SPECIAL_TOKENS_ATTRIBUTES & init_kwargs.keys(): + if ( + added_tokens_map != {} + and init_kwargs[key] is not None + and key != "additional_special_tokens" + ): + init_kwargs[key] = added_tokens_map.get( + str(init_kwargs[key]), init_kwargs[key] + ) + + # Instantiate the tokenizer. + try: + tokenizer = cls(*init_inputs, **init_kwargs) + except import_protobuf_decode_error(): + print( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(Google protobuf error: Tried to load SPM model with non-SPM vocab file).", + ) + return False + except RuntimeError as e: + if "sentencepiece_processor.cc" in str(e): + print( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(SentencePiece RuntimeError: Tried to load SPM model with non-SPM vocab file).", + ) + return False + except OSError as e: + raise OSError( + "Unable to load vocabulary from file. " + "Please check that the provided vocabulary is accessible and not corrupted." + ) from e + + if ( + added_tokens_decoder != {} + and max(list(added_tokens_decoder.keys())[-1], 0) > tokenizer.vocab_size + ): + print( + "Special tokens have been added in the vocabulary, make sure the associated word embeddings are" + " fine-tuned or trained." + ) + return tokenizer + + @staticmethod + def _eventually_correct_t5_max_length( + pretrained_model_name_or_path, max_model_length, init_max_model_length + ): + # This method should be deleted in Transformers v5 + # Its only purpose is to potentially throw a warning + # that incorrectly defined max lengths of T5's tokenizer are used + # which we will correct in Transformers v5. + return max_model_length + + @classmethod + def convert_added_tokens( + cls, obj: Union[AddedToken, Any], save=False, add_type_field=True + ): + """Convert added tokens.""" + if isinstance(obj, dict) and "__type" in obj and obj["__type"] == "AddedToken": + obj.pop("__type") + return AddedToken(**obj) + if isinstance(obj, AddedToken) and save: + obj = obj.__getstate__() + if add_type_field: + obj["__type"] = "AddedToken" + else: + # Don't save "special" for previous tokenizers + obj.pop("special") + return obj + if isinstance(obj, (list, tuple)): + return [ + cls.convert_added_tokens(o, save=save, add_type_field=add_type_field) + for o in obj + ] + if isinstance(obj, dict): + return { + k: cls.convert_added_tokens(v, save=save, add_type_field=add_type_field) + for k, v in obj.items() + } + return obj + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + legacy_format: Optional[bool] = None, + filename_prefix: Optional[str] = None, + push_to_hub: bool = False, + **kwargs, + ) -> Tuple[str]: + """Save the full tokenizer state. + + This method make sure the full tokenizer can then be re-loaded using the + [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] class method.. + + Warning,None This won't save modifications you may have applied to the tokenizer after the instantiation (for + instance, modifying `tokenizer.do_lower_case` after creation). + + Args: + save_directory (`str` or `os.PathLike`): The path to a directory where the tokenizer will be saved. + legacy_format (`bool`, *optional*): + Only applicable for a fast tokenizer. If unset (default), will save the tokenizer in the unified JSON + format as well as in legacy format if it exists, i.e. with tokenizer specific vocabulary and a separate + added_tokens files. + + If `False`, will only save the tokenizer in the unified JSON format. This format is incompatible with + "slow" tokenizers (not powered by the *tokenizers* library), so the tokenizer will not be able to be + loaded in the corresponding "slow" tokenizer. + + If `True`, will save the tokenizer in legacy format. If the "slow" tokenizer doesn't exits, a value + error is raised. + filename_prefix (`str`, *optional*): + A prefix to add to the names of the files saved by the tokenizer. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + + Returns: + A tuple of `str`: The files saved. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use " + "`token` instead.", + FutureWarning, + stacklevel=2, + ) + if kwargs.get("token") is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + if os.path.isfile(save_directory): + raise NotADirectoryError( + f"Vocabulary path ({save_directory}) should be a directory" + ) + + os.makedirs(save_directory, exist_ok=True) + + special_tokens_map_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + SPECIAL_TOKENS_MAP_FILE, + ) + tokenizer_config_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + TOKENIZER_CONFIG_FILE, + ) + chat_template_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + CHAT_TEMPLATE_FILE, + ) + + tokenizer_config = copy.deepcopy(self.init_kwargs) + + # Let's save the init kwargs + target_keys = set(self.init_kwargs.keys()) + # Let's save the special tokens map (only the strings) + target_keys.update(["model_max_length", "clean_up_tokenization_spaces"]) + + for k in target_keys: + if hasattr(self, k): + tokenizer_config[k] = getattr(self, k) + + # Let's make sure we properly save the special tokens + tokenizer_config.update(self.special_tokens_map) + if "extra_special_tokens" not in tokenizer_config: + tokenizer_config["extra_special_tokens"] = self.extra_special_tokens + tokenizer_config.update(self.extra_special_tokens) + + if self.chat_template is not None: + if isinstance(self.chat_template, dict): + # Chat template dicts are saved to the config as lists of dicts with fixed key names. + # They will be reconstructed as a single dict during loading. + # We're trying to discourage chat template dicts, and they are always + # saved in the config, never as single files. + tokenizer_config["chat_template"] = [ + {"name": k, "template": v} for k, v in self.chat_template.items() + ] + elif kwargs.get("save_raw_chat_template"): + with open(chat_template_file, "w", encoding="utf-8") as f: + f.write(self.chat_template) + print(f"chat template saved in {chat_template_file}") + if "chat_template" in tokenizer_config: + tokenizer_config.pop( + "chat_template" + ) # To ensure it doesn't somehow end up in the config too + else: + tokenizer_config["chat_template"] = self.chat_template + + if len(self.init_inputs) > 0: + tokenizer_config["init_inputs"] = copy.deepcopy(self.init_inputs) + for file_id in self.vocab_files_names: + tokenizer_config.pop(file_id, None) + + # no typefields, this way old fast and slow can load it + tokenizer_config = self.convert_added_tokens( + tokenizer_config, add_type_field=True, save=True + ) + + # Process added tokens seperatly: allows previous versions to ignore it! + added_tokens = {} + for key, value in self.added_tokens_decoder.items(): + added_tokens[key] = value.__getstate__() + tokenizer_config["added_tokens_decoder"] = added_tokens + + # Add tokenizer class to the tokenizer config to be able to reload it with from_pretrained + tokenizer_class = self.__class__.__name__ + # Remove the Fast at the end unless we have a special `PreTrainedTokenizerFast` + if ( + tokenizer_class.endswith("Fast") + and tokenizer_class != "PreTrainedTokenizerFast" + ): + tokenizer_class = tokenizer_class[:-4] + tokenizer_config["tokenizer_class"] = tokenizer_class + if getattr(self, "_auto_map", None) is not None: + tokenizer_config["auto_map"] = self._auto_map + if getattr(self, "_processor_class", None) is not None: + tokenizer_config["processor_class"] = self._processor_class + + # remove private information + if "name_or_path" in tokenizer_config: + tokenizer_config.pop("name_or_path") + tokenizer_config.pop("special_tokens_map_file", None) + tokenizer_config.pop("tokenizer_file", None) + if "device_map" in tokenizer_config: + tokenizer_config.pop("device_map") + + with open(tokenizer_config_file, "w", encoding="utf-8") as f: + out_str = ( + json.dumps( + tokenizer_config, indent=2, sort_keys=True, ensure_ascii=False + ) + + "\n" + ) + f.write(out_str) + print(f"tokenizer config file saved in {tokenizer_config_file}") + + # Sanitize AddedTokens in special_tokens_map + + # kept for forward compatibility, will be removed in transoformers 5. Typefields are not saved for FC, special + # should not be save either + write_dict = self.convert_added_tokens( + self.special_tokens_map_extended, save=True, add_type_field=False + ) + with open(special_tokens_map_file, "w", encoding="utf-8") as f: + out_str = ( + json.dumps(write_dict, indent=2, sort_keys=True, ensure_ascii=False) + + "\n" + ) + f.write(out_str) + print(f"Special tokens file saved in {special_tokens_map_file}") + + file_names = (tokenizer_config_file, special_tokens_map_file) + + return self._save_pretrained( + save_directory=save_directory, + file_names=file_names, + legacy_format=legacy_format, + filename_prefix=filename_prefix, + ) + + def _save_pretrained( + self, + save_directory: Union[str, os.PathLike], + file_names: Tuple[str], + legacy_format: Optional[bool] = None, + filename_prefix: Optional[str] = None, + ) -> Tuple[str]: + """Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens. + + Fast tokenizers can also be saved in a unique JSON file containing {config + vocab + added-tokens} using the + specific [`~tokenization_utils_fast.PreTrainedTokenizerFast._save_pretrained`] + """ + if legacy_format is False: + raise ValueError( + "Only fast tokenizers (instances of PreTrainedTokenizerFast) can be saved in non legacy format." + ) + + save_directory = str(save_directory) + + added_tokens_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + ADDED_TOKENS_FILE, + ) + # the new get_added_vocab() also returns special tokens and tokens that have an index < vocab_size + added_vocab = { + tok: index + for tok, index in self.added_tokens_encoder.items() + if index >= self.vocab_size + } + if added_vocab: + with open(added_tokens_file, "w", encoding="utf-8") as f: + out_str = ( + json.dumps( + added_vocab, indent=2, sort_keys=True, ensure_ascii=False + ) + + "\n" + ) + f.write(out_str) + print(f"added tokens file saved in {added_tokens_file}") + + vocab_files = self.save_vocabulary( + save_directory, filename_prefix=filename_prefix + ) + + return file_names + vocab_files + (added_tokens_file,) + + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save only the vocabulary of the tokenizer (vocabulary + added tokens). + + This method won't save the configuration and special token mappings of the tokenizer. Use + [`~PreTrainedTokenizerFast._save_pretrained`] to save the whole state of the tokenizer. + + Args: + save_directory (`str`): + The directory in which to save the vocabulary. + filename_prefix (`str`, *optional*): + An optional prefix to add to the named of the saved files. + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + raise NotImplementedError + + def tokenize( + self, + text: str, + pair: Optional[str] = None, + add_special_tokens: bool = False, + **kwargs, + ) -> List[str]: + """Converts a string into a sequence of tokens, replacing unknown tokens with the `unk_token`. + + Args: + text (`str`): + The sequence to be encoded. + pair (`str`, *optional*): + A second sequence to be encoded with the first. + add_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to add the special tokens associated with the corresponding model. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific encode method. See details in + [`~PreTrainedTokenizerBase.__call__`] + + Returns: + `List[str]`: The list of tokens. + """ + raise NotImplementedError + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, + """ + **kwargs: Passed along to the `.tokenize()` method. + """, + """ + Returns: + `List[int]`, `torch.Tensor`, `tf.Tensor` or `np.ndarray`: The tokenized ids of the text. + """, + ) + def encode( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> List[int]: + """Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. + + Same as doing `self.convert_tokens_to_ids(self.tokenize(text))`. + """ + encoded_inputs = self.encode_plus( + text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + padding_side=padding_side, + return_tensors=return_tensors, + **kwargs, + ) + + return encoded_inputs["input_ids"] + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + """Number of special tokens to add.""" + raise NotImplementedError + + def _get_padding_truncation_strategies( + self, + padding=False, + truncation=None, + max_length=None, + pad_to_multiple_of=None, + verbose=True, + **kwargs, + ): + """Find the correct padding/truncation strategy.""" + old_truncation_strategy = kwargs.pop("truncation_strategy", "do_not_truncate") + old_pad_to_max_length = kwargs.pop("pad_to_max_length", False) + + # Backward compatibility for previous behavior, maybe we should deprecate it: + # If you only set max_length, it activates truncation for max_length + if max_length is not None and padding is False and truncation is None: + if verbose: + if not self.deprecation_warnings.get( + "Truncation-not-explicitly-activated", False + ): + print( + "Truncation was not explicitly activated but `max_length` is provided a specific value, please" + " use `truncation=True` to explicitly truncate examples to max length. Defaulting to" + " 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the" + " tokenizer you can select this strategy more precisely by providing a specific strategy to" + " `truncation`." + ) + self.deprecation_warnings["Truncation-not-explicitly-activated"] = True + truncation = "longest_first" + + # Get padding strategy + if padding is False and old_pad_to_max_length: + if verbose: + warnings.warn( + "The `pad_to_max_length` argument is deprecated and will be removed in a future version, " + "use `padding=True` or `padding='longest'` to pad to the longest sequence in the batch, or " + "use `padding='max_length'` to pad to a max length. In this case, you can give a specific " + "length with `max_length` (e.g. `max_length=45`) or leave max_length to None to pad to the " + "maximal input size of the model (e.g. 512 for Bert).", + FutureWarning, + stacklevel=2, + ) + padding_strategy = ( + PaddingStrategy.LONGEST + if max_length is None + else PaddingStrategy.MAX_LENGTH + ) + elif padding is not False: + if padding is True: + if verbose: + if max_length is not None and ( + truncation is None + or truncation is False + or truncation == "do_not_truncate" + ): + warnings.warn( + "`max_length` is ignored when `padding`=`True` and there is no truncation strategy. " + "To pad to max length, use `padding='max_length'`.", + stacklevel=2, + ) + if old_pad_to_max_length is not False: + warnings.warn( + "Though `pad_to_max_length` = `True`, it is ignored because `padding`=`True`.", + stacklevel=2, + ) + padding_strategy = ( + PaddingStrategy.LONGEST + ) # Default to pad to the longest sequence in the batch + elif not isinstance(padding, PaddingStrategy): + padding_strategy = PaddingStrategy(padding) + elif isinstance(padding, PaddingStrategy): + padding_strategy = padding + else: + padding_strategy = PaddingStrategy.DO_NOT_PAD + + # Get truncation strategy + if truncation is None and old_truncation_strategy != "do_not_truncate": + if verbose: + warnings.warn( + "The `truncation_strategy` argument is deprecated and will be removed in a future version, use" + " `truncation=True` to truncate examples to a max length. You can give a specific length with" + " `max_length` (e.g. `max_length=45`) or leave max_length to None to truncate to the maximal input" + " size of the model (e.g. 512 for Bert). If you have pairs of inputs, you can give a specific" + " truncation strategy selected among `truncation='only_first'` (will only truncate the first" + " sentence in the pairs) `truncation='only_second'` (will only truncate the second sentence in the" + " pairs) or `truncation='longest_first'` (will iteratively remove tokens from the longest sentence" + " in the pairs).", + FutureWarning, + stacklevel=2, + ) + truncation_strategy = TruncationStrategy(old_truncation_strategy) + elif truncation is not False and truncation is not None: + if truncation is True: + truncation_strategy = ( + TruncationStrategy.LONGEST_FIRST + ) # Default to truncate the longest sequences in pairs of inputs + elif not isinstance(truncation, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation) + elif isinstance(truncation, TruncationStrategy): + truncation_strategy = truncation + else: + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + + # Set max length if needed + if max_length is None: + if padding_strategy == PaddingStrategy.MAX_LENGTH: + if self.model_max_length > LARGE_INTEGER: + if verbose: + if not self.deprecation_warnings.get( + "Asking-to-pad-to-max_length", False + ): + print( + "Asking to pad to max_length but no maximum length is provided and the model has no" + " predefined maximum length. Default to no padding." + ) + self.deprecation_warnings["Asking-to-pad-to-max_length"] = True + padding_strategy = PaddingStrategy.DO_NOT_PAD + else: + max_length = self.model_max_length + + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE: + if self.model_max_length > LARGE_INTEGER: + if verbose: + if not self.deprecation_warnings.get( + "Asking-to-truncate-to-max_length", False + ): + print( + "Asking to truncate to max_length but no maximum length is provided and the model has" + " no predefined maximum length. Default to no truncation." + ) + self.deprecation_warnings[ + "Asking-to-truncate-to-max_length" + ] = True + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + else: + max_length = self.model_max_length + + # Test if we have a padding token + if padding_strategy != PaddingStrategy.DO_NOT_PAD and ( + self.pad_token is None or self.pad_token_id < 0 + ): + raise ValueError( + "Asking to pad but the tokenizer does not have a padding token. " + "Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` " + "or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`." + ) + + # Check that we will truncate to a multiple of pad_to_multiple_of if both are provided + if ( + truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE + and padding_strategy != PaddingStrategy.DO_NOT_PAD + and pad_to_multiple_of is not None + and max_length is not None + and (max_length % pad_to_multiple_of != 0) + ): + raise ValueError( + "Truncation and padding are both activated but " + f"truncation length ({max_length}) is not a multiple of pad_to_multiple_of ({pad_to_multiple_of})." + ) + + return padding_strategy, truncation_strategy, max_length, kwargs + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING + ) + def __call__( # noqa: D417 + self, + text: Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ] = None, + text_pair: Optional[ + Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ] + ] = None, + text_target: Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ] = None, + text_pair_target: Optional[ + Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ] + ] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """Main method to tokenize and prepare for the model one or several sequence(s). + + Args: + text (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_pair (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_target (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. Each sequence can be a string or a + list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), + you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_pair_target (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. Each sequence can be a string or a + list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), + you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + """ + # To avoid duplicating + all_kwargs = { + "add_special_tokens": add_special_tokens, + "padding": padding, + "truncation": truncation, + "max_length": max_length, + "stride": stride, + "is_split_into_words": is_split_into_words, + "pad_to_multiple_of": pad_to_multiple_of, + "padding_side": padding_side, + "return_tensors": return_tensors, + "return_token_type_ids": return_token_type_ids, + "return_attention_mask": return_attention_mask, + "return_overflowing_tokens": return_overflowing_tokens, + "return_special_tokens_mask": return_special_tokens_mask, + "return_offsets_mapping": return_offsets_mapping, + "return_length": return_length, + "split_special_tokens": kwargs.pop( + "split_special_tokens", self.split_special_tokens + ), + "verbose": verbose, + } + all_kwargs.update(kwargs) + if text is None and text_target is None: + raise ValueError("You need to specify either `text` or `text_target`.") + if text is not None: + # The context manager will send the inputs as normal texts and not text_target, but we shouldn't change the + # input mode in this case. + if not self._in_target_context_manager: + self._switch_to_input_mode() + encodings = self._call_one(text=text, text_pair=text_pair, **all_kwargs) + if text_target is not None: + self._switch_to_target_mode() + target_encodings = self._call_one( + text=text_target, text_pair=text_pair_target, **all_kwargs + ) + # Leave back tokenizer in input mode + self._switch_to_input_mode() + + if text_target is None: + return encodings + if text is None: + return target_encodings + + encodings["labels"] = target_encodings["input_ids"] + return encodings + + def _call_one( + self, + text: Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ], + text_pair: Optional[ + Union[ + TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] + ] + ] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + # Input type checking for clearer error + def _is_valid_text_input(t): + if isinstance(t, str): + # Strings are fine + return True + if isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... empty + return True + if isinstance(t[0], str): + # ... list of strings + return True + if isinstance(t[0], (list, tuple)): + # ... list with an empty list or with a list of strings + return len(t[0]) == 0 or isinstance(t[0][0], str) + return False + return False + + if not _is_valid_text_input(text): + raise ValueError( + "text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if text_pair is not None and not _is_valid_text_input(text_pair): + raise ValueError( + "text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if is_split_into_words: + is_batched = ( + isinstance(text, (list, tuple)) + and text + and isinstance(text[0], (list, tuple)) + ) + else: + is_batched = isinstance(text, (list, tuple)) + + if is_batched: + if isinstance(text_pair, str): + raise TypeError( + "when tokenizing batches of text, `text_pair` must be a list or tuple with the same length as" + " `text`." + ) + if text_pair is not None and len(text) != len(text_pair): + raise ValueError( + f"batch length of `text`: {len(text)} does not match batch length of `text_pair`: {len(text_pair)}." + ) + batch_text_or_text_pairs = ( + list(zip(text, text_pair)) if text_pair is not None else text + ) + return self.batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + return self.encode_plus( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING + ) + def encode_plus( # noqa: D417 + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """Tokenize and prepare for the model a sequence or a pair of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + text (`str`, `List[str]` or (for non-fast tokenizers) `List[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + text_pair (`str`, `List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + """ + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = ( + self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + ) + + return self._encode_plus( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=kwargs.pop( + "split_special_tokens", self.split_special_tokens + ), + **kwargs, + ) + + def _encode_plus( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + raise NotImplementedError + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING + ) + def batch_encode_plus( # noqa: D417 + self, + batch_text_or_text_pairs: Union[ + List[TextInput], + List[TextInputPair], + List[PreTokenizedInput], + List[PreTokenizedInputPair], + List[EncodedInput], + List[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + """Tokenize and prepare for the model a list of sequences or a list of pairs of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + batch_text_or_text_pairs (`List[str]`, `List[Tuple[str, str]]`, `List[List[str]]`, + `List[Tuple[List[str], List[str]]]`, and for not-fast tokenizers, also `List[List[int]]`, + `List[Tuple[List[int], List[int]]]`): + Batch of sequences or pair of sequences to be encoded. This can be a list of + string/string-sequences/int-sequences or a list of pair of string/string-sequences/int-sequence (see + details in `encode_plus`). + """ + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = ( + self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + ) + + return self._batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + List[TextInput], + List[TextInputPair], + List[PreTokenizedInput], + List[PreTokenizedInputPair], + List[EncodedInput], + List[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + raise NotImplementedError + + def pad( + self, + encoded_inputs: Union[ + BatchEncoding, + List[BatchEncoding], + Dict[str, EncodedInput], + Dict[str, List[EncodedInput]], + List[Dict[str, EncodedInput]], + ], + padding: Union[bool, str, PaddingStrategy] = True, + max_length: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + verbose: bool = True, + ) -> BatchEncoding: + """Pad a single encoded input or a batch of encoded inputs up to predefined length. + + Padding side (left/right) padding token ids are defined at the tokenizer level (with `self.padding_side`, + `self.pad_token_id` and `self.pad_token_type_id`). + + Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the + text followed by a call to the `pad` method to get a padded encoding. + + + + If the `encoded_inputs` passed are dictionary of numpy arrays, PyTorch tensors or TensorFlow tensors, the + result will use the same type unless you provide a different tensor type with `return_tensors`. In the case of + PyTorch tensors, you will lose the specific device of your tensors however. + + + + Args: + encoded_inputs: ([`BatchEncoding`], list of [`BatchEncoding`], `Dict[str, List[int]]`, `Dict[str, + List[List[int]]` or `List[Dict[str, List[int]]]`): + Tokenized inputs. Can represent one input ([`BatchEncoding`] or `Dict[str, List[int]]`) or a batch of + tokenized inputs (list of [`BatchEncoding`], *Dict[str, List[List[int]]]* or *List[Dict[str, + List[int]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader + collate function. + + Instead of `List[int]` you can have tensors (numpy arrays, PyTorch tensors or TensorFlow tensors), see + the note above for the return type. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) @lint-ignore + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + """ + if self.__class__.__name__.endswith( + "Fast" + ) and not self.deprecation_warnings.get( + "Asking-to-pad-a-fast-tokenizer", False + ): + print( + f"You're using a {self.__class__.__name__} tokenizer. Please note that with a fast tokenizer," + " using the `__call__` method is faster than using a method to encode the text followed by a call" + " to the `pad` method to get a padded encoding." + ) + self.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True + + # If we have a list of dicts, let's convert it in a dict of lists + # We do this to allow using this method as a collate_fn function in PyTorch Dataloader + if isinstance(encoded_inputs, (list, tuple)) and isinstance( + encoded_inputs[0], Mapping + ): + encoded_inputs = { + key: [example[key] for example in encoded_inputs] + for key in encoded_inputs[0] + } + + # The model's main input name, usually `input_ids`, has been passed for padding + if self.model_input_names[0] not in encoded_inputs: + raise ValueError( + "You should supply an encoding or a list of encodings to this method " + f"that includes {self.model_input_names[0]}, but you provided {list(encoded_inputs.keys())}" + ) + + required_input = encoded_inputs[self.model_input_names[0]] + + if required_input is None or ( + isinstance(required_input, Sized) and len(required_input) == 0 + ): + if return_attention_mask: + encoded_inputs["attention_mask"] = [] + return encoded_inputs + + # If we have PyTorch/TF/NumPy tensors/arrays as inputs, we cast them as python objects + # and rebuild them afterwards if no return_tensors is specified + # Note that we lose the specific device the tensor may be on for PyTorch + + first_element = required_input[0] + if isinstance(first_element, (list, tuple)): + # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element. + for item in required_input: + if len(item) != 0: + first_element = item[0] + break + # At this state, if `first_element` is still a list/tuple, it's an empty one so there is nothing to do. + if not isinstance(first_element, (int, list, tuple)): + if is_tf_tensor(first_element): + return_tensors = "tf" if return_tensors is None else return_tensors + elif is_torch_tensor(first_element): + return_tensors = "pt" if return_tensors is None else return_tensors + elif isinstance(first_element, np.ndarray): + return_tensors = "np" if return_tensors is None else return_tensors + else: + raise ValueError( + f"type of {first_element} unknown: {type(first_element)}. " + "Should be one of a python, numpy, pytorch or tensorflow object." + ) + + for key, value in encoded_inputs.items(): + encoded_inputs[key] = to_py_obj(value) + + # Convert padding_strategy in PaddingStrategy + padding_strategy, _, max_length, _ = self._get_padding_truncation_strategies( + padding=padding, max_length=max_length, verbose=verbose + ) + + required_input = encoded_inputs[self.model_input_names[0]] + if required_input and not isinstance(required_input[0], (list, tuple)): + encoded_inputs = self._pad( + encoded_inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + return BatchEncoding(encoded_inputs, tensor_type=return_tensors) + + batch_size = len(required_input) + assert all( + len(v) == batch_size for v in encoded_inputs.values() + ), "Some items in the output dictionary have a different batch size than others." + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = max(len(inputs) for inputs in required_input) + padding_strategy = PaddingStrategy.MAX_LENGTH + + batch_outputs = {} + for i in range(batch_size): + inputs = {k: v[i] for k, v in encoded_inputs.items()} + outputs = self._pad( + inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """Create the token type IDs corresponding to the sequences passed. + + [What are token type IDs?](../glossary#token-type-ids) @lint-ignore + + Should be overridden in a subclass if the model has a special way of building those. + + Args: + token_ids_0 (`List[int]`): The first tokenized sequence. + token_ids_1 (`List[int]`, *optional*): The second tokenized sequence. + + Returns: + `List[int]`: The token type ids. + """ + if token_ids_1 is None: + return len(token_ids_0) * [0] + return [0] * len(token_ids_0) + [1] * len(token_ids_1) + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """Build model inputs from a sequence for classification tasks by concatenating and adding special tokens. + + This implementation does not add special tokens and this method should be overridden in a subclass. + + Args: + token_ids_0 (`List[int]`): The first tokenized sequence. + token_ids_1 (`List[int]`, *optional*): The second tokenized sequence. + + Returns: + `List[int]`: The model input with special tokens. + """ + if token_ids_1 is None: + return token_ids_0 + return token_ids_0 + token_ids_1 + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING + ) + def prepare_for_model( # noqa: D417 + self, + ids: List[int], + pair_ids: Optional[List[int]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy] = None, + max_length: Optional[int] = None, + stride: int = 0, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + prepend_batch_axis: bool = False, + **kwargs, + ) -> BatchEncoding: + """Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. + + It adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens. Please Note, for *pair_ids* + different than `None` and *truncation_strategy = longest_first* or `True`, it is not possible to return + overflowing tokens. Such a combination of arguments will raise an error. + + Args: + ids (`List[int]`): + Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and + `convert_tokens_to_ids` methods. + pair_ids (`List[int]`, *optional*): + Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize` + and `convert_tokens_to_ids` methods. + """ + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = ( + self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + ) + + pair = bool(pair_ids is not None) + len_ids = len(ids) + len_pair_ids = len(pair_ids) if pair else 0 + + if return_token_type_ids and not add_special_tokens: + raise ValueError( + "Asking to return token_type_ids while setting add_special_tokens to False " + "results in an undefined behavior. Please set add_special_tokens to True or " + "set return_token_type_ids to None." + ) + + if ( + return_overflowing_tokens + and truncation_strategy == TruncationStrategy.LONGEST_FIRST + and pair_ids is not None + ): + raise ValueError( + "Not possible to return overflowing tokens for pair of sequences with the " + "`longest_first`. Please select another truncation strategy than `longest_first`, " + "for instance `only_second` or `only_first`." + ) + + # Load from model defaults + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + encoded_inputs = {} + + # Compute the total size of the returned encodings + total_len = ( + len_ids + + len_pair_ids + + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0) + ) + + # Truncation: Handle max sequence length + overflowing_tokens = [] + if ( + truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE + and max_length + and total_len > max_length + ): + ids, pair_ids, overflowing_tokens = self.truncate_sequences( + ids, + pair_ids=pair_ids, + num_tokens_to_remove=total_len - max_length, + truncation_strategy=truncation_strategy, + stride=stride, + ) + + if return_overflowing_tokens: + encoded_inputs["overflowing_tokens"] = overflowing_tokens + encoded_inputs["num_truncated_tokens"] = total_len - max_length + + # Add special tokens + if add_special_tokens: + sequence = self.build_inputs_with_special_tokens(ids, pair_ids) + token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids) + else: + sequence = ids + pair_ids if pair else ids + token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else []) + + # Build output dictionary + encoded_inputs["input_ids"] = sequence + if return_token_type_ids: + encoded_inputs["token_type_ids"] = token_type_ids + if return_special_tokens_mask: + if add_special_tokens: + encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask( + ids, pair_ids + ) + else: + encoded_inputs["special_tokens_mask"] = [0] * len(sequence) + + # Check lengths + self._eventual_warn_about_too_long_sequence( + encoded_inputs["input_ids"], max_length, verbose + ) + + # Padding + if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask: + encoded_inputs = self.pad( + encoded_inputs, + max_length=max_length, + padding=padding_strategy.value, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + if return_length: + encoded_inputs["length"] = len(encoded_inputs["input_ids"]) + + return BatchEncoding( + encoded_inputs, + tensor_type=return_tensors, + prepend_batch_axis=prepend_batch_axis, + ) + + def truncate_sequences( + self, + ids: List[int], + pair_ids: Optional[List[int]] = None, + num_tokens_to_remove: int = 0, + truncation_strategy: Union[str, TruncationStrategy] = "longest_first", + stride: int = 0, + ) -> Tuple[List[int], List[int], List[int]]: + """Truncates a sequence pair in-place following the strategy. + + Args: + ids (`List[int]`): + Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and + `convert_tokens_to_ids` methods. + pair_ids (`List[int]`, *optional*): + Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize` + and `convert_tokens_to_ids` methods. + num_tokens_to_remove (`int`, *optional*, defaults to 0): + Number of tokens to remove using the truncation strategy. + truncation_strategy: (`str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to + `'longest_first'`): + The strategy to follow for truncation. Can be: + + - `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will truncate + token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a + batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths greater + than the model maximum admissible input size). + stride (`int`, *optional*, defaults to 0): + If set to a positive number, the overflowing tokens returned will contain some tokens from the main + sequence returned. The value of this argument defines the number of additional tokens. + + Returns: + `Tuple[List[int], List[int], List[int]]`: The truncated `ids`, the truncated `pair_ids` and the list of + overflowing tokens. Note: The *longest_first* strategy returns empty list of overflowing tokens if a pair + of sequences (or a batch of pairs) is provided. + """ + if num_tokens_to_remove <= 0: + return ids, pair_ids, [] + + if not isinstance(truncation_strategy, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation_strategy) + + overflowing_tokens = [] + if truncation_strategy == TruncationStrategy.ONLY_FIRST or ( + truncation_strategy == TruncationStrategy.LONGEST_FIRST and pair_ids is None + ): + if len(ids) > num_tokens_to_remove: + window_len = min(len(ids), stride + num_tokens_to_remove) + if self.truncation_side == "left": + overflowing_tokens = ids[:window_len] + ids = ids[num_tokens_to_remove:] + elif self.truncation_side == "right": + overflowing_tokens = ids[-window_len:] + ids = ids[:-num_tokens_to_remove] + else: + raise ValueError( + f"invalid truncation strategy: {self.truncation_side}, use 'left' or 'right'." + ) + + else: + error_msg = ( + f"We need to remove {num_tokens_to_remove} to truncate the input " + f"but the first sequence has a length {len(ids)}. " + ) + if truncation_strategy == TruncationStrategy.ONLY_FIRST: + error_msg = ( + error_msg + "Please select another truncation strategy than " + f"{truncation_strategy}, for instance 'longest_first' or 'only_second'." + ) + raise RuntimeError(error_msg) + elif truncation_strategy == TruncationStrategy.LONGEST_FIRST: + print( + "Be aware, overflowing tokens are not returned for the setting you have chosen," + f" i.e. sequence pairs with the '{TruncationStrategy.LONGEST_FIRST.value}' " + "truncation strategy. So the returned list will always be empty even if some " + "tokens have been removed." + ) + len_pair_ids = len(pair_ids) if pair_ids is not None else 0 + len_ids = len(ids) + first_remove = min(abs(len_pair_ids - len_ids), num_tokens_to_remove) + second_remove = num_tokens_to_remove - first_remove + if len_ids > len_pair_ids: + ids_to_move = first_remove + second_remove // 2 + pair_ids_to_move = second_remove - second_remove // 2 + else: + ids_to_move = second_remove // 2 + pair_ids_to_move = first_remove + second_remove - (second_remove // 2) + + if self.truncation_side == "right": + ids = ids[:-ids_to_move] if ids_to_move > 0 else ids + pair_ids = ( + pair_ids[:-pair_ids_to_move] + if pair_ids is not None and pair_ids_to_move > 0 + else pair_ids + ) + elif self.truncation_side == "left": + ids = ids[ids_to_move:] + pair_ids = pair_ids[pair_ids_to_move:] if pair_ids is not None else None + else: + raise ValueError(f"invalid truncation strategy:{self.truncation_side}") + + elif ( + truncation_strategy == TruncationStrategy.ONLY_SECOND + and pair_ids is not None + ): + if len(pair_ids) > num_tokens_to_remove: + window_len = min(len(pair_ids), stride + num_tokens_to_remove) + if self.truncation_side == "right": + overflowing_tokens = pair_ids[-window_len:] + pair_ids = pair_ids[:-num_tokens_to_remove] + elif self.truncation_side == "left": + overflowing_tokens = pair_ids[:window_len] + pair_ids = pair_ids[num_tokens_to_remove:] + else: + raise ValueError( + f"invalid truncation strategy:{self.truncation_side}" + ) + else: + raise RuntimeError( + f"We need to remove {num_tokens_to_remove} to truncate the input " + f"but the second sequence has a length {len(pair_ids)}. " + f"Please select another truncation strategy than {truncation_strategy}, " + "for instance 'longest_first' or 'only_first'." + ) + + return (ids, pair_ids, overflowing_tokens) + + def _pad( + self, + encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], + max_length: Optional[int] = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + ) -> dict: + """Pad encoded inputs (on left/right and up to predefined length or max length in the batch). + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + padding_strategy: PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in `padding_side` argument: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + required_input = encoded_inputs[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if ( + max_length is not None + and pad_to_multiple_of is not None + and (max_length % pad_to_multiple_of != 0) + ): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = ( + padding_strategy != PaddingStrategy.DO_NOT_PAD + and len(required_input) != max_length + ) + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(required_input) + + if needs_to_be_padded: + difference = max_length - len(required_input) + padding_side = ( + padding_side if padding_side is not None else self.padding_side + ) + + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = ( + encoded_inputs["attention_mask"] + [0] * difference + ) + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = ( + encoded_inputs["token_type_ids"] + + [self.pad_token_type_id] * difference + ) + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = ( + encoded_inputs["special_tokens_mask"] + [1] * difference + ) + encoded_inputs[self.model_input_names[0]] = ( + required_input + [self.pad_token_id] * difference + ) + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [ + 0 + ] * difference + encoded_inputs["attention_mask"] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [ + self.pad_token_type_id + ] * difference + encoded_inputs["token_type_ids"] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [ + 1 + ] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs[self.model_input_names[0]] = [ + self.pad_token_id + ] * difference + required_input + else: + raise ValueError(f"Invalid padding strategy:{padding_side}") + + return encoded_inputs + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens in a single string. + + The most simple way to do it is `" ".join(tokens)` but we + often want to remove sub-word tokenization artifacts at the same time. + + Args: + tokens (`List[str]`): The token to join in a string. + + Returns: + `str`: The joined tokens. + """ + raise NotImplementedError + + def batch_decode( + self, + sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> List[str]: + """Convert a list of lists of token ids into a list of strings by calling decode. + + Args: + sequences (`Union[List[int], List[List[int]], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `List[str]`: The list of decoded sentences. + """ + return [ + self.decode( + seq, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + for seq in sequences + ] + + def decode( + self, + token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> str: + """Converts a sequence of ids in a string. + + Using the tokenizer and vocabulary with options to remove special + tokens and clean up tokenization spaces. + + Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `str`: The decoded sentence. + """ + # Convert inputs to python lists + token_ids = to_py_obj(token_ids) + + return self._decode( + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + def _decode( + self, + token_ids: Union[int, List[int]], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> str: + raise NotImplementedError + + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """Retrieves sequence ids from a token list that has no special tokens added. + + This method is called when adding + special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods. + + Args: + token_ids_0 (`List[int]`): + List of ids of the first sequence. + token_ids_1 (`List[int]`, *optional*): + List of ids of the second sequence. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + assert already_has_special_tokens and token_ids_1 is None, ( + "You cannot use ``already_has_special_tokens=False`` with this tokenizer. " + "Please use a slow (full python) tokenizer to activate this argument. " + "Or set `return_special_tokens_mask=True` when calling the encoding method " + "to get the special tokens mask in any tokenizer. " + ) + + all_special_ids = self.all_special_ids # cache the property + + return [1 if token in all_special_ids else 0 for token in token_ids_0] + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + """Clean up simple English tokenization artifacts like spaces before punctuations and abbreviated forms. + + Args: + out_string (`str`): The text to clean up. + + Returns: + `str`: The cleaned-up string. + """ + return ( + out_string.replace(" .", ".") + .replace(" ?", "?") + .replace(" !", "!") + .replace(" ,", ",") + .replace(" ' ", "'") + .replace(" n't", "n't") + .replace(" 'm", "'m") + .replace(" 's", "'s") + .replace(" 've", "'ve") + .replace(" 're", "'re") + ) + + def _eventual_warn_about_too_long_sequence( + self, ids: List[int], max_length: Optional[int], verbose: bool + ): + """Eventually warn about too long sequences. + + Depending on the input and internal state we might trigger a warning about a sequence that is too long for its + corresponding model + + Args: + ids (`List[str]`): The ids produced by the tokenization + max_length (`int`, *optional*): The max_length desired (does not trigger a warning if it is set) + verbose (`bool`): Whether or not to print more information and warnings. + + """ + if max_length is None and len(ids) > self.model_max_length and verbose: + if not self.deprecation_warnings.get( + "sequence-length-is-longer-than-the-specified-maximum", False + ): + print( + "Token indices sequence length is longer than the specified maximum sequence length " + f"for this model ({len(ids)} > {self.model_max_length}). Running this sequence through the model " + "will result in indexing errors" + ) + self.deprecation_warnings[ + "sequence-length-is-longer-than-the-specified-maximum" + ] = True + + def _switch_to_input_mode(self): + """Private method to put the tokenizer in input mode (when it has different modes for input/outputs).""" + + def _switch_to_target_mode(self): + """Private method to put the tokenizer in target mode (when it has different modes for input/outputs).""" + + @contextmanager + def as_target_tokenizer(self): + """Temporarily sets the tokenizer for encoding the targets. + + Useful for tokenizer associated to + sequence-to-sequence models that need a slightly different processing for the labels. + """ + warnings.warn( + "`as_target_tokenizer` is deprecated and will be removed in v5 of Transformers. You can tokenize your " + "labels by using the argument `text_target` of the regular `__call__` method (either in the same call as " + "your input texts if you use the same keyword arguments, or in a separate call.", + stacklevel=2, + ) + self._switch_to_target_mode() + self._in_target_context_manager = True + yield + self._in_target_context_manager = False + self._switch_to_input_mode() + + @classmethod + def register_for_auto_class(cls, auto_class="AutoTokenizer"): + """Register this class with a given auto class. + + This should only be used for custom tokenizers as the ones in the + library are already mapped with `AutoTokenizer`. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoTokenizer"`): + The auto class to register this new tokenizer with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def prepare_seq2seq_batch( + self, + src_texts: List[str], + tgt_texts: Optional[List[str]] = None, + max_length: Optional[int] = None, + max_target_length: Optional[int] = None, + padding: str = "longest", + return_tensors: Optional[str] = None, + truncation: bool = True, + **kwargs, + ) -> BatchEncoding: + """Prepare model inputs for translation. For best performance, translate one sentence at a time. + + Arguments: + src_texts (`List[str]`): + List of documents to summarize or source language texts. + tgt_texts (`list`, *optional*): + List of summaries or target language texts. + max_length (`int`, *optional*): + Controls the maximum length for encoder inputs (documents to summarize or source language texts) If + left unset or set to `None`, this will use the predefined model maximum length if a maximum length is + required by one of the truncation/padding parameters. If the model has no specific maximum input length + (like XLNet) truncation/padding to a maximum length will be deactivated. + max_target_length (`int`, *optional*): + Controls the maximum length of decoder inputs (target language texts or summaries) If left unset or set + to `None`, this will use the max_length value. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + truncation (`bool`, `str` or [`TruncationStrategy`], *optional*, defaults to `True`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + **kwargs: + Additional keyword arguments passed along to `self.__call__`. + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to the encoder. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model. + - **labels** -- List of token ids for tgt_texts. + + The full set of keys `[input_ids, attention_mask, labels]`, will only be returned if tgt_texts is passed. + Otherwise, input_ids, attention_mask will be the only keys. + """ + # docstyle-ignore + formatted_warning = """ +`prepare_seq2seq_batch` is deprecated and will be removed in version 5 of HuggingFace Transformers. Use the regular +`__call__` method to prepare your inputs and targets. + +Here is a short example: + +model_inputs = tokenizer(src_texts, text_target=tgt_texts, ...) + +If you either need to use different keyword arguments for the source and target texts, you should do two calls like +this: + +model_inputs = tokenizer(src_texts, ...) +labels = tokenizer(text_target=tgt_texts, ...) +model_inputs["labels"] = labels["input_ids"] + +See the documentation of your specific tokenizer for more details on the specific arguments to the tokenizer of choice. +For a more complete example, see the implementation of `prepare_seq2seq_batch`. +""" + warnings.warn(formatted_warning, FutureWarning, stacklevel=2) + # mBART-specific kwargs that should be ignored by other models. + kwargs.pop("src_lang", None) + kwargs.pop("tgt_lang", None) + if max_length is None: + max_length = self.model_max_length + model_inputs = self( + src_texts, + add_special_tokens=True, + return_tensors=return_tensors, + max_length=max_length, + padding=padding, + truncation=truncation, + **kwargs, + ) + if tgt_texts is None: + return model_inputs + # Process tgt_texts + if max_target_length is None: + max_target_length = max_length + with self.as_target_tokenizer(): + labels = self( + tgt_texts, + add_special_tokens=True, + return_tensors=return_tensors, + padding=padding, + max_length=max_target_length, + truncation=truncation, + **kwargs, + ) + model_inputs["labels"] = labels["input_ids"] + return model_inputs + + +def get_fast_tokenizer_file(tokenization_files: List[str]) -> str: + """Get the tokenization file to use for this version of transformers. + + Args: + tokenization_files (`List[str]`): The list of available configuration files. + + Returns: + `str`: The tokenization file to use. + """ + tokenizer_files_map = {} + for file_name in tokenization_files: + search = _re_tokenizer_file.search(file_name) + if search is not None: + v = search.groups()[0] + tokenizer_files_map[v] = file_name + + return FULL_TOKENIZER_FILE diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_fast.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_fast.py index b61466606a0..65c9f39bb1f 100644 --- a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_fast.py +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_utils_fast.py @@ -629,7 +629,6 @@ def _encode_plus( return_offsets_mapping=return_offsets_mapping, return_length=return_length, verbose=verbose, - **kwargs, ) # Return tensor is None, then we can remove the leading batch axis diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_whisper.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_whisper.py new file mode 100644 index 00000000000..1da13d2a782 --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/tokenization_whisper.py @@ -0,0 +1,1507 @@ +# Copyright 2022 The HuggingFace Inc. team. All rights reserved. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenization classes for Whisper.""" + +# flake8: noqa: C901 + +import json +import os +import warnings +from functools import lru_cache +from typing import List, Optional, Tuple, Union + +import numpy as np +import regex as re + +from .english_normalizer import BasicTextNormalizer, EnglishTextNormalizer +from .tokenization_utils import AddedToken, PreTrainedTokenizer + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "tokenizer_file": "tokenizer.json", + "merges_file": "merges.txt", + "normalizer_file": "normalizer.json", +} + + +MAX_MODEL_INPUT_SIZES = { + "openai/whisper-base": 448, +} + + +# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode +def bytes_to_unicode(): + """Returns list of utf-8 byte and a mapping to unicode strings. + + We specifically avoids mapping to whitespace/control characters the bpe code barfs on. + + The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab + if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for + decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup + tables between utf-8 bytes and unicode strings. + """ + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs +def get_pairs(word): + """Return set of symbol pairs in a word. + + Word is represented as tuple of symbols (symbols being variable-length strings). + """ + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +LANGUAGES = { + "en": "english", + "zh": "chinese", + "de": "german", + "es": "spanish", + "ru": "russian", + "ko": "korean", + "fr": "french", + "ja": "japanese", + "pt": "portuguese", + "tr": "turkish", + "pl": "polish", + "ca": "catalan", + "nl": "dutch", + "ar": "arabic", + "sv": "swedish", + "it": "italian", + "id": "indonesian", + "hi": "hindi", + "fi": "finnish", + "vi": "vietnamese", + "he": "hebrew", + "uk": "ukrainian", + "el": "greek", + "ms": "malay", + "cs": "czech", + "ro": "romanian", + "da": "danish", + "hu": "hungarian", + "ta": "tamil", + "no": "norwegian", + "th": "thai", + "ur": "urdu", + "hr": "croatian", + "bg": "bulgarian", + "lt": "lithuanian", + "la": "latin", + "mi": "maori", + "ml": "malayalam", + "cy": "welsh", + "sk": "slovak", + "te": "telugu", + "fa": "persian", + "lv": "latvian", + "bn": "bengali", + "sr": "serbian", + "az": "azerbaijani", + "sl": "slovenian", + "kn": "kannada", + "et": "estonian", + "mk": "macedonian", + "br": "breton", + "eu": "basque", + "is": "icelandic", + "hy": "armenian", + "ne": "nepali", + "mn": "mongolian", + "bs": "bosnian", + "kk": "kazakh", + "sq": "albanian", + "sw": "swahili", + "gl": "galician", + "mr": "marathi", + "pa": "punjabi", + "si": "sinhala", + "km": "khmer", + "sn": "shona", + "yo": "yoruba", + "so": "somali", + "af": "afrikaans", + "oc": "occitan", + "ka": "georgian", + "be": "belarusian", + "tg": "tajik", + "sd": "sindhi", + "gu": "gujarati", + "am": "amharic", + "yi": "yiddish", + "lo": "lao", + "uz": "uzbek", + "fo": "faroese", + "ht": "haitian creole", + "ps": "pashto", + "tk": "turkmen", + "nn": "nynorsk", + "mt": "maltese", + "sa": "sanskrit", + "lb": "luxembourgish", + "my": "myanmar", + "bo": "tibetan", + "tl": "tagalog", + "mg": "malagasy", + "as": "assamese", + "tt": "tatar", + "haw": "hawaiian", + "ln": "lingala", + "ha": "hausa", + "ba": "bashkir", + "jw": "javanese", + "su": "sundanese", + "yue": "cantonese", +} + +# language code lookup by name, with a few language aliases +TO_LANGUAGE_CODE = { + **{language: code for code, language in LANGUAGES.items()}, + "burmese": "my", + "valencian": "ca", + "flemish": "nl", + "haitian": "ht", + "letzeburgesch": "lb", + "pushto": "ps", + "panjabi": "pa", + "moldavian": "ro", + "moldovan": "ro", + "sinhalese": "si", + "castilian": "es", + "mandarin": "zh", +} + +TASK_IDS = ["translate", "transcribe"] + + +class WhisperTokenizer(PreTrainedTokenizer): + """Construct a Whisper tokenizer. + + This tokenizer inherits from [`PreTrainedTokenizer`] which contains some of the main methods. Users should refer to + the superclass for more information regarding such methods. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + merges_file (`str`): + Path to the merges file. + normalizer_file (`str`, *optional*): + Path to the normalizer_file file. + errors (`str`, *optional*, defaults to `"replace"`): + Paradigm to follow when decoding bytes to UTF-8. See + [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information. + unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this + token instead. + bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The beginning of sequence token. The `decoder_start_token_id` is used to set the first token as + `"<|startoftranscript|>"` when generating. + eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`): + The end of sequence token. + pad_token (`str`, *optional*): + The token used for padding, for example when batching sequences of different lengths. + add_prefix_space (`bool`, *optional*, defaults to `False`): + Whether or not to add an initial space to the input. This allows to treat the leading word just as any + other word. + language (`str`, *optional*): + The language of the transcription text. The corresponding language id token is appended to the start of the + sequence for multilingual speech recognition and speech translation tasks, e.g. for Spanish the token + `"<|es|>"` is appended to the start of sequence. This should be used for multilingual fine-tuning only. + task (`str`, *optional*): + Task identifier to append at the start of sequence (if any). This should be used for mulitlingual + fine-tuning, with `"transcribe"` for speech recognition and `"translate"` for speech translation. + predict_timestamps (`bool`, *optional*, defaults to `False`): + Whether to omit the `<|notimestamps|>` token at the start of the sequence. + """ + + vocab_files_names = VOCAB_FILES_NAMES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + merges_file, + normalizer_file=None, + errors="replace", + unk_token="<|endoftext|>", + bos_token="<|endoftext|>", + eos_token="<|endoftext|>", + pad_token=None, + add_prefix_space=False, + language=None, + task=None, + predict_timestamps=False, + **kwargs, + ): + """Initialize WhisperTokenizer.""" + bos_token = ( + AddedToken( + bos_token, lstrip=False, rstrip=False, normalized=False, special=True + ) + if isinstance(bos_token, str) + else bos_token + ) + eos_token = ( + AddedToken( + eos_token, lstrip=False, rstrip=False, normalized=False, special=True + ) + if isinstance(eos_token, str) + else eos_token + ) + unk_token = ( + AddedToken( + unk_token, lstrip=False, rstrip=False, normalized=False, special=True + ) + if isinstance(unk_token, str) + else unk_token + ) + pad_token = ( + AddedToken( + pad_token, lstrip=False, rstrip=False, normalized=False, special=True + ) + if isinstance(pad_token, str) + else pad_token + ) + + with open(vocab_file, encoding="utf-8") as vocab_handle: + self.encoder = json.load(vocab_handle) + self.decoder = {v: k for k, v in self.encoder.items()} + self.errors = errors # how to handle errors in decoding + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + with open(merges_file, encoding="utf-8") as merges_handle: + bpe_merges = merges_handle.read().split("\n")[1:-1] + bpe_merges = [tuple(merge.split()) for merge in bpe_merges] + self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) + self.cache = {} + self.add_prefix_space = add_prefix_space + + if normalizer_file is not None: + with open(normalizer_file, encoding="utf-8") as vocab_handle: + self.english_spelling_normalizer = json.load(vocab_handle) + else: + self.english_spelling_normalizer = None + + # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions + self.pat = re.compile( + r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" + ) + self.timestamp_pat = re.compile(r"<\|(\d+\.\d+)\|>") + + self.language = language + super().__init__( + errors=errors, + unk_token=unk_token, + bos_token=bos_token, + eos_token=eos_token, + pad_token=pad_token, + add_prefix_space=add_prefix_space, + **kwargs, + ) + + self.task = task + self.predict_timestamps = predict_timestamps + + @property + def vocab_size(self) -> int: + """Get length of vocab.""" + return len(self.encoder) + + def get_vocab(self): + """Gets the vocab.""" + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} + vocab.update(self.added_tokens_encoder) + return vocab + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe with GPT2 -> Whisper + def bpe(self, token): + """Byte pair encoding.""" + if token in self.cache: + return self.cache[token] + word = tuple(token) + pairs = get_pairs(word) + + if not pairs: + return token + + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + else: + new_word.extend(word[i:j]) + i = j + + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + new_word = tuple(new_word) + word = new_word + if len(word) == 1: + break + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + def set_prefix_tokens( + self, + language: Optional[str] = None, + task: Optional[str] = None, + predict_timestamps: Optional[bool] = None, + ): + """Override the prefix tokens appended to the start of the label sequence. + + This method can be used standalone to update the prefix tokens as required when fine-tuning. Example: + + ```python + >>> # instantiate the tokenizer and set the prefix token to Spanish + >>> tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-tiny", language="spanish") + >>> # now switch the prefix token from Spanish to French + >>> tokenizer.set_prefix_tokens(language="french") + ``` + + Args: + language (`str`, *optional*, defaults to `None`): + The language of the transcription text. + task (`str`, *optional*, defaults to `None`): + Task identifier to append at the start of sequence (if any). + predict_timestamps (`bool`, *optional*, defaults to `None`): + Whether to omit the `<|notimestamps|>` token at the start of the sequence. + """ + self.language = language if language is not None else self.language + self.task = task if task is not None else self.task + self.predict_timestamps = ( + predict_timestamps + if predict_timestamps is not None + else self.predict_timestamps + ) + + @property + def prefix_tokens(self) -> List[int]: + """Gets prefix tokens.""" + bos_token_id = self.convert_tokens_to_ids("<|startoftranscript|>") + translate_token_id = self.convert_tokens_to_ids("<|translate|>") + transcribe_token_id = self.convert_tokens_to_ids("<|transcribe|>") + notimestamps_token_id = self.convert_tokens_to_ids("<|notimestamps|>") + langs = tuple(LANGUAGES.keys()) + + if self.language is not None: + self.language = self.language.lower() + if self.language in TO_LANGUAGE_CODE: + language_id = TO_LANGUAGE_CODE[self.language] + elif self.language in TO_LANGUAGE_CODE.values(): + language_id = self.language + else: + is_language_code = len(self.language) == 2 + raise ValueError( + f"Unsupported language: {self.language}. Language should be one of:" + f" {list(TO_LANGUAGE_CODE.values()) if is_language_code else list(TO_LANGUAGE_CODE.keys())}." + ) + + if self.task is not None and self.task not in TASK_IDS: + raise ValueError( + f"Unsupported task: {self.task}. Task should be in: {TASK_IDS}" + ) + + bos_sequence = [bos_token_id] + if self.language is not None: + bos_sequence.append(bos_token_id + 1 + langs.index(language_id)) + if self.task is not None: + bos_sequence.append( + transcribe_token_id if self.task == "transcribe" else translate_token_id + ) + if not self.predict_timestamps: + bos_sequence.append(notimestamps_token_id) + return bos_sequence + + def build_inputs_with_special_tokens( + self, token_ids_0, token_ids_1=None + ) -> List[int]: + """Build model inputs from a sequence by appending eos_token_id.""" + if token_ids_1 is None: + # return self.prefix_tokens + token_ids_0 + [self.eos_token_id] + return self.prefix_tokens + token_ids_0 + # We don't expect to process pairs, but leave the pair logic for API consistency + # return self.prefix_tokens + token_ids_0 + token_ids_1 + [self.eos_token_id] + return self.prefix_tokens + token_ids_0 + token_ids_1 + + def get_special_tokens_mask( + self, + token_ids_0: List[int], + token_ids_1: Optional[List[int]] = None, + already_has_special_tokens: bool = False, + ) -> List[int]: + """Retrieve sequence ids from a token list that has no special tokens added. + + This method is called when adding special tokens using the tokenizer `prepare_for_model` method. + + Args: + token_ids_0 (`List[int]`): + List of IDs. + token_ids_1 (`List[int]`, *optional*): + Optional second list of IDs for sequence pairs. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ) + + prefix_ones = [1] * len(self.prefix_tokens) + suffix_ones = [1] + if token_ids_1 is None: + return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones + return ( + prefix_ones + + ([0] * len(token_ids_0)) + + ([0] * len(token_ids_1)) + + suffix_ones + ) + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize with GPT2 -> Whisper + def _tokenize(self, text): + """Tokenize a string.""" + bpe_tokens = [] + for token in re.findall(self.pat, text): + token = "".join( + self.byte_encoder[b] for b in token.encode("utf-8") + ) # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) + return bpe_tokens + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id with GPT2 -> Whisper + def _convert_token_to_id(self, token): + """Converts a token (str) in an id using the vocab.""" + return self.encoder.get(token, self.encoder.get(self.unk_token)) + + def _convert_id_to_token(self, index): + """Converts an index (integer) in a token (str) using the vocab. + + Whisper's base tokenizer always decodes OOV tokens as "", thus we do not use the `unk_token` here. + """ + return self.decoder.get(index, "") + + def _normalize(self, text): + warnings.warn( + "The private method `_normalize` is deprecated and will be removed in v5 of Transformers." + "You can normalize an input string using the Whisper English normalizer using the `normalize` method.", + stacklevel=2, + ) + return self.normalize(text) + + def _basic_normalize(self, text, remove_diacritics=False): + warnings.warn( + "The private method `_basic_normalize` is deprecated and will be removed in v5 of Transformers." + "You can normalize an input string using the Whisper basic normalizer using the `basic_normalize` method.", + stacklevel=2, + ) + return self.basic_normalize(text, remove_diacritics=remove_diacritics) + + def normalize(self, text): + """Normalize a given string using the `EnglishTextNormalizer` class. + + Which preforms commons transformation on english text. + """ + normalizer = EnglishTextNormalizer(self.english_spelling_normalizer) + return normalizer(text) + + @staticmethod + def basic_normalize(text, remove_diacritics=False): + """Normalize a given string using the `BasicTextNormalizer` class. + + Which preforms commons transformation on multilingual text. + """ + normalizer = BasicTextNormalizer(remove_diacritics=remove_diacritics) + return normalizer(text) + + def _decode_with_timestamps( + self, + token_ids, + skip_special_tokens=False, + time_precision=0.02, + segment_size=1500, + ) -> str: + """Timestamp tokens are above the special tokens' id range and are ignored by `decode()`. + + This method decodes given tokens with timestamps tokens annotated, e.g. "<|1.08|>". + """ + timestamp_begin = self.all_special_ids[-1] + 1 + outputs = [[]] + + cur_max_timestamp = 0.0 + prev_segments_len = 0.0 + penultimate_timestamp = 0.0 + + for i, token in enumerate(token_ids): + if token >= timestamp_begin: + timestamp = float((token - timestamp_begin) * time_precision) + + if timestamp < cur_max_timestamp: + # next segment has started + last_was_single_ending = i >= 2 and not ( + token_ids[i - 1] >= timestamp_begin + and token_ids[i - 2] >= timestamp_begin + ) + if last_was_single_ending: + prev_segments_len += time_precision * segment_size + else: + cur_max_timestamp = penultimate_timestamp + prev_segments_len += penultimate_timestamp + outputs = outputs[:-2] + + penultimate_timestamp = cur_max_timestamp + cur_max_timestamp = timestamp + + outputs.append(f"<|{(timestamp + prev_segments_len):.2f}|>") + outputs.append([]) + else: + outputs[-1].append(token) + outputs = [ + ( + s + if isinstance(s, str) + else self.decode( + s, + skip_special_tokens=skip_special_tokens, + decode_with_timestamps=False, + ) + ) + for s in outputs + ] + return "".join(outputs) + + def _compute_offsets(self, token_ids, time_precision=0.02, segment_size=1500): + """Compute offsets for a given tokenized input. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + time_precision (`float`, *optional*, defaults to 0.02): + The time ratio to convert from token to time. + segment_size (`int`, *optional*, defaults to 1500): + The number of features in the input mel spectrogram. + """ + offsets = [] + # ensure torch tensor of token ids is placed on cpu + if "torch" in str(type(token_ids)) and ( + hasattr(token_ids, "cpu") and callable(token_ids.cpu) + ): + token_ids = token_ids.cpu() + token_ids = np.array(token_ids) + if token_ids.shape[0] > 1 and len(token_ids.shape) > 1: + raise ValueError("Can only process a single input at a time") + timestamp_begin = self.all_special_ids[-1] + 1 + timestamp_tokens = token_ids >= timestamp_begin + + consecutive = np.where(timestamp_tokens[:-1] & timestamp_tokens[1:])[0] + 1 + if consecutive.shape[0] == 0 and timestamp_tokens.sum() <= 1: + # either there are no timestamps or there are no consecutive ones + return [] + if np.where(timestamp_tokens)[0][-1] + 1 not in consecutive: + # we add the final timestamp if it is not already in the list + consecutive = np.append(consecutive, np.where(timestamp_tokens)[0][-1] + 1) + + last_slice = np.where(timestamp_tokens)[0][0] + cur_max_timestamp = 0 + prev_segments_len = 0 + for current_slice in consecutive: + sliced_tokens = token_ids[last_slice:current_slice] + if len(sliced_tokens) > 1: + start_timestamp_position = sliced_tokens[0].item() - timestamp_begin + end_timestamp_position = sliced_tokens[-1].item() - timestamp_begin + + if start_timestamp_position < cur_max_timestamp: + # next segment has started + is_single_ending = last_slice >= 2 and not ( + token_ids[last_slice - 2] >= timestamp_begin + and token_ids[last_slice - 1] >= timestamp_begin + ) + if is_single_ending: + prev_segments_len += segment_size + else: + prev_segments_len += cur_max_timestamp + + cur_max_timestamp = end_timestamp_position + + # strip timestamp tokens from the text output + sliced_tokens = self._preprocess_token_ids(sliced_tokens) + text = self._decode(sliced_tokens) + text = self._filter_timestamp_ids(text) + offsets.append( + { + "text": text, + "timestamp": ( + start_timestamp_position * time_precision + + prev_segments_len * time_precision, + end_timestamp_position * time_precision + + prev_segments_len * time_precision, + ), + } + ) + last_slice = current_slice + + return offsets + + @lru_cache # noqa: B019 + def timestamp_ids(self, time_precision=0.02): + """Compute the timestamp token ids for a given precision and save to least-recently used (LRU) cache. + + Args: + time_precision (`float`, *optional*, defaults to 0.02): + The time ratio to convert from token to time. + """ + return self.convert_tokens_to_ids( + [("<|%.2f|>" % (i * time_precision)) for i in range(1500 + 1)] + ) + + def _preprocess_token_ids(self, token_ids, skip_special_tokens: bool = False): + """Pre-process the token ids for decoding by removing the prompt tokens ids and timestamp token ids. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Typically, obtained using the `__call__` method of the tokenizer. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens from the token ids. If `True`, the prompt token ids will be + removed. + """ + if skip_special_tokens: + prompt_token_id = self.convert_tokens_to_ids("<|startofprev|>") + decoder_start_token_id = self.convert_tokens_to_ids("<|startoftranscript|>") + token_ids = self._strip_prompt( + token_ids, prompt_token_id, decoder_start_token_id + ) + + return token_ids + + def _filter_timestamp_ids(self, token_ids): + return re.sub(self.timestamp_pat, "", token_ids) + + def decode( + self, + token_ids, + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + output_offsets: bool = False, + time_precision: float = 0.02, + decode_with_timestamps: bool = True, + normalize: bool = False, + basic_normalize: bool = False, + remove_diacritics: bool = False, + **kwargs, + ) -> str: + """Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special tokens and clean up tokenization spaces. + + Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. Will remove the previous tokens (pre-prompt) + if present. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces` (available in the `tokenizer_config`). + output_offsets (`bool`, *optional*, defaults to `False`): + Whether or not to output the offsets of the tokens. This should only be set if the model predicted + timestamps. If there are previous tokens (pre-prompt) to decode, they will only appear in the decoded + text if they contain timestamp tokens. + time_precision (`float`, *optional*, defaults to 0.02): + The time ratio to convert from token to time. + decode_with_timestamps (`bool`, *optional*, defaults to `False`): + Whether or not to decode with timestamps included in the raw text. + normalize (`bool`, *optional*, defaults to `False`): + Whether or not to apply the English text normalizer to the decoded text. Only applicable when the + target text is in English. Otherwise, the basic text normalizer should be applied. + basic_normalize (`bool`, *optional*, defaults to `False`): + Whether or not to apply the Basic text normalizer to the decoded text. Applicable to multilingual + target text. + remove_diacritics (`bool`, *optional*, defaults to `False`): + Whether or not to remove diacritics when applying the Basic text normalizer. Removing diacritics may + destroy information in the decoded text, hence it should be used with caution. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `str`: The decoded sentence. + """ # noqa: E501 + filtered_ids = self._preprocess_token_ids( + token_ids, + skip_special_tokens=skip_special_tokens, + ) + + text = super().decode( + filtered_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + normalize=normalize, + basic_normalize=basic_normalize, + remove_diacritics=remove_diacritics, + **kwargs, + ) + if decode_with_timestamps: + # legacy method to decode timestamps when not included in the tokenizer vocabulary + text = self._decode_with_timestamps( + filtered_ids, + time_precision=time_precision, + skip_special_tokens=skip_special_tokens, + ) + else: + text = self._filter_timestamp_ids(text) + + # retrieve offsets + if output_offsets: + offsets = self._compute_offsets(token_ids, time_precision=time_precision) + return {"text": text, "offsets": offsets} + return text + + def _decode( + self, + token_ids: Union[int, List[int]], + skip_special_tokens: bool = False, + normalize: bool = False, + basic_normalize: bool = False, + remove_diacritics: bool = False, + **kwargs, + ) -> str: + self._decode_use_source_tokenizer = kwargs.pop("use_source_tokenizer", False) + filtered_tokens = self.convert_ids_to_tokens( + token_ids, skip_special_tokens=skip_special_tokens + ) + + # To avoid mixing byte-level and unicode for byte-level BPT + # we need to build string separately for added tokens and byte-level tokens + # cf. https://github.com/huggingface/transformers/issues/1133 + sub_texts = [] + current_sub_text = [] + for token in filtered_tokens: + if skip_special_tokens and token in self.all_special_ids: + continue + if token in self.added_tokens_encoder: + if current_sub_text: + sub_texts.append(self.convert_tokens_to_string(current_sub_text)) + current_sub_text = [] + sub_texts.append(token) + else: + current_sub_text.append(token) + if current_sub_text: + sub_texts.append(self.convert_tokens_to_string(current_sub_text)) + + text = "".join(sub_texts) + + if normalize: + return self.normalize(text) + if basic_normalize: + return self.basic_normalize(text, remove_diacritics=remove_diacritics) + return text + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string with GPT2 -> Whisper + def convert_tokens_to_string(self, tokens): + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + return bytearray([self.byte_decoder[c] for c in text]).decode( + "utf-8", errors=self.errors + ) + + def save_vocabulary( + self, save_directory: str, filename_prefix: Optional[str] = None + ) -> Tuple[str]: + """Save vocab.""" + if not os.path.isdir(save_directory): + return None + vocab_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["vocab_file"], + ) + merge_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["merges_file"], + ) + normalizer_file = os.path.join( + save_directory, + (filename_prefix + "-" if filename_prefix else "") + + VOCAB_FILES_NAMES["normalizer_file"], + ) + + with open(vocab_file, "w", encoding="utf-8") as f: + f.write( + json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + + "\n" + ) + + index = 0 + with open(merge_file, "w", encoding="utf-8") as writer: + writer.write("#version: 0.2\n") + for bpe_tokens, token_index in sorted( + self.bpe_ranks.items(), key=lambda kv: kv[1] + ): + if index != token_index: + index = token_index + writer.write(" ".join(bpe_tokens) + "\n") + index += 1 + + if self.english_spelling_normalizer is not None: + with open(normalizer_file, "w", encoding="utf-8") as f: + f.write( + json.dumps( + self.english_spelling_normalizer, + indent=2, + sort_keys=True, + ensure_ascii=False, + ) + + "\n" + ) + + return vocab_file, merge_file, normalizer_file + + # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.prepare_for_tokenization with GPT2 -> Whisper + def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs): + """Prepares text for tokenization.""" + add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space) + if is_split_into_words or add_prefix_space: + text = " " + text + return (text, kwargs) + + def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True): + """Get decoder prompt ids.""" + self.set_prefix_tokens( + task=task, language=language, predict_timestamps=not no_timestamps + ) + # prefix tokens are of the form: <|startoftranscript|> <|lang_id|> <|task|> <|notimestamps|> + # we don't want to force the bos token at position 1, as this is the starting token + # when we generate, so we slice the prefix tokens to: <|lang_id|> <|task|> <|notimestamps|> + # to get the forced tokens + forced_tokens = self.prefix_tokens[1:] + return [(rank + 1, token) for rank, token in enumerate(forced_tokens)] + + def _decode_asr( + self, model_outputs, *, return_timestamps, return_language, time_precision + ): + return _decode_asr( + self, + model_outputs, + return_timestamps=return_timestamps, + return_language=return_language, + time_precision=time_precision, + ) + + def get_prompt_ids(self, text: str, return_tensors="np"): + """Converts prompt text to IDs that can be passed to [`~WhisperForConditionalGeneration.generate`].""" + batch_encoding = self( + "<|startofprev|>", " " + text.strip(), add_special_tokens=False + ) + + # Check for special tokens + prompt_text_ids = batch_encoding["input_ids"][1:] + special_token_id = next( + (x for x in prompt_text_ids if x >= self.all_special_ids[0]), None + ) + if special_token_id is not None: + token = self.convert_ids_to_tokens(special_token_id) + raise ValueError( + f"Encountered text in the prompt corresponding to disallowed special token: {token}." + ) + + batch_encoding.convert_to_tensors(tensor_type=return_tensors) + return batch_encoding["input_ids"] + + def _strip_prompt( + self, token_ids: List[int], prompt_token_id: int, decoder_start_token_id: int + ): + if not isinstance(token_ids, list): + token_ids = self._convert_to_list(token_ids) + + # handle case of empty token_ids for decoding with timestamps. + # at this point token_ids is a list, so it is safe to use if not check. + if not token_ids: + return token_ids + + has_prompt = token_ids[0] == prompt_token_id + if has_prompt: + if decoder_start_token_id in token_ids: + return token_ids[token_ids.index(decoder_start_token_id) :] + return [] + + return token_ids + + @staticmethod + def _convert_to_list(token_ids): + # convert type to ndarray if necessary + if hasattr(token_ids, "numpy"): + if "torch" in str(type(token_ids)): + token_ids = token_ids.cpu().numpy() + elif "tensorflow" in str(type(token_ids)): + token_ids = token_ids.numpy() + elif "jaxlib" in str(type(token_ids)): + token_ids = token_ids.tolist() + # now the token ids are either a numpy array, or a list of lists + if isinstance(token_ids, np.ndarray): + token_ids = token_ids.tolist() + return token_ids + + +def _decode_asr( + tokenizer, model_outputs, *, return_timestamps, return_language, time_precision +): + """Internal method meant to only be used by asr pipeline. + + Handles all the little quirks specific to whisper to handle the various options not allowed in other seq2seq models. + """ + # =========== Overview ============ + # - iterate over all outputs + # - all tokens within output + # - Each token can be + # - language token + # - special token + # - timestamp token + # - text token + # - We accumulate the text tokens. + # - We split on end timestamps + # - Lots of complexity comes from stride and timestamps + + last_language = None + + def new_chunk(): + return {"language": last_language, "timestamp": [None, None], "text": ""} + + # Welcome to the state machine ! + chunks = [] + chunk = new_chunk() + time_offset = 0.0 + timestamp_begin = tokenizer.convert_tokens_to_ids("<|notimestamps|>") + 1 + previous_tokens = [] + previous_token_timestamps = [] + skip = False + right_stride_start = None + + all_special_ids = set(tokenizer.all_special_ids) + prompt_token_id = tokenizer.convert_tokens_to_ids("<|startofprev|>") + decoder_start_token_id = tokenizer.convert_tokens_to_ids("<|startoftranscript|>") + # - iterate over all outputs + for _chunk_id, output in enumerate(model_outputs): + # We can drop everything to Python list, it's going to make + # our lives easier + token_ids = output["tokens"][0].tolist() + # (possibly) remove the prompt from the token ids + token_ids = tokenizer._strip_prompt( + token_ids, prompt_token_id, decoder_start_token_id + ) # noqa: SLF001 + if return_timestamps == "word": + token_timestamps = output["token_timestamps"][0].tolist() + + # Those keep track of timestamps within strides + # Which need to be skipped and resolve all tokens in a single + # chunk. + last_timestamp = None + first_timestamp = timestamp_begin + + if "stride" in output: + chunk_len, stride_left, stride_right = output["stride"] + # Offset the timings to account for the other `model_outputs`. + time_offset -= stride_left + right_stride_start = chunk_len - stride_right + + # Keeping track of timestamps within strides + # We're going to NOT split on those, and delay until we're + # out of BOTH stride. Otherwise lots of issues occur and + # corner cases + if stride_left: + first_timestamp = stride_left / time_precision + timestamp_begin + if stride_right: + for token in reversed(token_ids): + if token >= timestamp_begin: + # There can be several token in the right stride + # But the last one is ALWAYS going to be skipped + if ( + last_timestamp is not None + and (token - timestamp_begin) * time_precision + < right_stride_start + ): + break + last_timestamp = token + + current_tokens = [] + current_token_timestamps = [] + + # - all tokens within output + for i, token in enumerate(token_ids): + # 4 possible states for each token + # - 1/ Language code + # - 2/ all other special tokens (which we ignore) + # - 3/ Timestamp + # - 4/ Regular text + if token in all_special_ids: + # Either language code or other + text = tokenizer.decode([token]) + # Removing outer shell <|XX|> + text = text[2:-2] + language = LANGUAGES.get(text) + if language is not None: + # 1/ Indeed some language + # TODO Handle when language is different from the previous + # one, and we cannot use timestamped tokens to create chunks + if ( + last_language + and language != last_language + and not return_timestamps + ): + previous_tokens.append(current_tokens) + resolved_tokens = _find_longest_common_sequence(previous_tokens) + resolved_text = tokenizer.decode(resolved_tokens) + chunk["text"] = resolved_text + chunks.append(chunk) + + # Flush all our temporary context + previous_tokens = [] + current_tokens = [] + chunk = new_chunk() + chunk["language"] = language + last_language = language + else: + # 2/ This is a regular special token, ignoring it + pass + elif token >= timestamp_begin: + # 3/ Timestamp token + time = (token - timestamp_begin) * time_precision + time_offset + time = round(time, 2) + if last_timestamp and token >= last_timestamp: + # Whisper outputted a timestamp token, but it falls within + # our stride, so we're going to skip it for the time being + # and resolve this later + # Skip is necessary because timestamp tokens always come + # by pair, so we need to skip the next one too (which would mark the start of another chunk). + skip = True + elif skip or (previous_tokens and token < first_timestamp): + skip = False + elif chunk["timestamp"][0] is None: + chunk["timestamp"][0] = time + else: + # This is the end of the timestamp chunk + if time == chunk["timestamp"][0]: + # This is a bug in timestamp token output + # where we're taking the duplicate token + # as a stop where it should be a start. + # This is an issue in the underlying model output + # Let's just skip it so it becomes de-factor + # a start agin + pass + else: + chunk["timestamp"][1] = time + # Handling merges. + previous_tokens.append(current_tokens) + if return_timestamps == "word": + previous_token_timestamps.append(current_token_timestamps) + resolved_tokens, resolved_token_timestamps = ( + _find_longest_common_sequence( + previous_tokens, previous_token_timestamps + ) + ) + resolved_text = tokenizer.decode(resolved_tokens) + chunk["text"] = resolved_text + if return_timestamps == "word": + chunk["words"] = _collate_word_timestamps( + tokenizer, + resolved_tokens, + resolved_token_timestamps, + last_language, + return_language, + ) + chunks.append(chunk) + + # Flush all our temporary context + previous_tokens = [] + current_tokens = [] + previous_token_timestamps = [] + current_token_timestamps = [] + chunk = new_chunk() + else: + # 4/ Regular token + # We just append to the list of all tokens so we can handle + # merges later and decode into text. + current_tokens.append(token) + if return_timestamps == "word": + start_time = round(token_timestamps[i] + time_offset, 2) + if i + 1 < len(token_timestamps): + end_time = round(token_timestamps[i + 1] + time_offset, 2) + else: + end_time = None # should never happen + current_token_timestamps.append((start_time, end_time)) + + if "stride" in output: + time_offset += chunk_len - stride_right + + # Leftover tokens + if current_tokens: + previous_tokens.append(current_tokens) + if return_timestamps == "word": + previous_token_timestamps.append(current_token_timestamps) + elif not (any(p for p in previous_tokens)): + chunk = new_chunk() + previous_tokens = [] + current_tokens = [] + previous_token_timestamps = [] + current_token_timestamps = [] + + if previous_tokens: + # Happens when we don't use timestamps + resolved_tokens, resolved_token_timestamps = _find_longest_common_sequence( + previous_tokens, previous_token_timestamps + ) + resolved_text = tokenizer.decode(resolved_tokens) + chunk["text"] = resolved_text + if return_timestamps == "word": + chunk["words"] = _collate_word_timestamps( + tokenizer, + resolved_tokens, + resolved_token_timestamps, + last_language, + return_language, + ) + chunks.append(chunk) + + # Preparing and cleaning up the pipeline output + full_text = "".join(chunk["text"] for chunk in chunks) + if return_timestamps or return_language: + for chunk in chunks: + if not return_timestamps: + chunk.pop("timestamp") + else: + chunk["timestamp"] = tuple(chunk["timestamp"]) + if not return_language: + chunk.pop("language") + + if return_timestamps == "word": + new_chunks = [] + for chunk in chunks: + new_chunks.extend(chunk["words"]) + optional = {"chunks": new_chunks} + else: + optional = {"chunks": chunks} + else: + optional = {} + return full_text, optional + + +def _find_longest_common_sequence(sequences, token_timestamp_sequences=None): + # It would be much harder to do O(n) because of fault tolerance. + # We actually have a really good property which is that the total sequence + # MUST be those subsequences in order. + # If token_timestamp_sequences is provided, will split those sequences in + # exactly the same way. + + left_sequence = sequences[0] + left_length = len(left_sequence) + total_sequence = [] + + if token_timestamp_sequences: + left_token_timestamp_sequence = token_timestamp_sequences[0] + total_token_timestamp_sequence = [] + + for seq_idx, right_sequence in enumerate(sequences[1:]): + # index = 0 + max_ = 0.0 + max_indices = (left_length, left_length, 0, 0) + # Here we're sliding matches + # [a, b, c, d] + # [c, d, f] + # = [c] == [d] + # + # [a, b, c, d] + # [c, d, f] + # = [c, d] == [c, d] + # + # + # [a, b, c, d] + # [c, d, f] + # + # = [b, c, d] == [c, d, f] + # + # [a, b, c, d] + # [c, d, f] + # + # [a, b, c] == [c, d, f] + # + # [a, b, c, d] + # [d, f] + # + # [a, b] == [d, f] + # + # [a, b, c, d] + # [f] + # + # [a] == [f] + right_length = len(right_sequence) + for i in range(1, left_length + right_length): + # epsilon to favor long perfect matches + eps = i / 10000.0 + + # Slightly convoluted because we don't want out of bound indices + # This will be necessary for a small conflict resolution optimization + # later + left_start = max(0, left_length - i) + left_stop = min(left_length, left_length + right_length - i) + left = np.array(left_sequence[left_start:left_stop]) + + right_start = max(0, i - left_length) + right_stop = min(right_length, i) + right = np.array(right_sequence[right_start:right_stop]) + + # We can only match subsequences of the same size. + if len(left) != len(right): + raise RuntimeError( + "There is a bug within whisper `decode_asr` function, " + "please report it. Dropping to prevent bad inference." + ) + + if token_timestamp_sequences: + # Get length of longest subsequence of tokens that match + # and have timestamps that are in order + matches = sum( + 1 + for idx, elem in enumerate(left) + if ( + elem == right[idx] + and left_token_timestamp_sequence[left_start + idx] + <= token_timestamp_sequences[seq_idx + 1][right_start + idx] + ) + ) + + else: + matches = np.sum(left == right) + + matching = matches / i + eps + if matches > 1 and matching > max_: + max_ = matching + max_indices = (left_start, left_stop, right_start, right_stop) + + (left_start, left_stop, right_start, right_stop) = max_indices + + # This is a small conflict optimization since those sequences overlap + # in audio. + # We're going to give more confidence to the left sequence + # for the left of the overlap, + # and to the right of the sequence, for the right of the overlap + left_mid = (left_stop + left_start) // 2 + right_mid = (right_stop + right_start) // 2 + total_sequence.extend(left_sequence[:left_mid]) + left_sequence = right_sequence[right_mid:] + left_length = len(left_sequence) + + if token_timestamp_sequences: + total_token_timestamp_sequence.extend( + left_token_timestamp_sequence[:left_mid] + ) + left_token_timestamp_sequence = token_timestamp_sequences[seq_idx + 1][ + right_mid: + ] + + total_sequence.extend(left_sequence) + + if token_timestamp_sequences is None: + return total_sequence + + if len(token_timestamp_sequences) > 0: + total_token_timestamp_sequence.extend(left_token_timestamp_sequence) + return total_sequence, total_token_timestamp_sequence + return total_sequence, [] + + +def _collate_word_timestamps( + tokenizer, tokens, token_timestamps, language, return_language +): + words, _, token_indices = _combine_tokens_into_words(tokenizer, tokens, language) + + optional_language_field = {"language": language} if return_language else {} + + return [ + { + "text": word, + "timestamp": ( + token_timestamps[indices[0]][0], + token_timestamps[indices[-1]][1], + ), + **optional_language_field, + } + for word, indices in zip(words, token_indices) + ] + + +def _combine_tokens_into_words( + tokenizer, + tokens: List[int], + language: Optional[str] = None, + prepend_punctuations: str = "\"'“¡¿([{-", + append_punctuations: str = "\"'.。,，!！?？:：”)]}、", # noqa: RUF001 +): + """Groups tokens by word. + + Returns a tuple containing a list of strings with the words, + and a list of `token_id` sequences with the tokens making up each word. + """ + if language is None: + language = tokenizer.language + if language is None: + language = "english" + + if language in {"chinese", "japanese", "thai", "lao", "myanmar", "cantonese"}: + # These languages don't typically use spaces. + words, word_tokens, token_indices = _split_tokens_on_unicode(tokenizer, tokens) + else: + words, word_tokens, token_indices = _split_tokens_on_spaces(tokenizer, tokens) + + _merge_punctuations( + words, word_tokens, token_indices, prepend_punctuations, append_punctuations + ) + return words, word_tokens, token_indices + + +def _split_tokens_on_unicode(tokenizer, tokens: List[int]): + """Combine tokens into words by splitting at any position where the tokens are decoded as valid unicode points.""" + decoded_full = tokenizer.decode(tokens, decode_with_timestamps=True) + replacement_char = "\ufffd" + + words = [] + word_tokens = [] + token_indices = [] + current_tokens = [] + current_indices = [] + unicode_offset = 0 + + for token_idx, token in enumerate(tokens): + current_tokens.append(token) + current_indices.append(token_idx) + decoded = tokenizer.decode(current_tokens, decode_with_timestamps=True) + + if ( + replacement_char not in decoded + or decoded_full[unicode_offset + decoded.index(replacement_char)] + == replacement_char + ): + words.append(decoded) + word_tokens.append(current_tokens) + token_indices.append(current_indices) + current_tokens = [] + current_indices = [] + unicode_offset += len(decoded) + + return words, word_tokens, token_indices + + +def _split_tokens_on_spaces(tokenizer, tokens: List[int]): + """Combine tokens into words by splitting at whitespace and punctuation tokens.""" + subwords, subword_tokens_list, subword_indices_list = _split_tokens_on_unicode( + tokenizer, tokens + ) + words = [] + word_tokens = [] + token_indices = [] + + for subword, subword_tokens, subword_indices in zip( + subwords, subword_tokens_list, subword_indices_list + ): + special = subword_tokens[0] >= tokenizer.eos_token_id + with_space = subword.startswith(" ") + punctuation = subword.strip() in "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~" + + if special or with_space or punctuation or len(words) == 0: + words.append(subword) + word_tokens.append(subword_tokens) + token_indices.append(subword_indices) + else: + words[-1] = words[-1] + subword + word_tokens[-1].extend(subword_tokens) + token_indices[-1].extend(subword_indices) + + return words, word_tokens, token_indices + + +def _merge_punctuations(words, tokens, indices, prepended, appended): + """Merges punctuation tokens with neighboring words.""" + # prepend punctuations + i = len(words) - 2 + j = len(words) - 1 + while i >= 0: + if words[i].startswith(" ") and words[i].strip() in prepended: + words[j] = words[i] + words[j] + tokens[j] = tokens[i] + tokens[j] + indices[j] = indices[i] + indices[j] + words[i] = "" + tokens[i] = [] + indices[i] = [] + else: + j = i + i -= 1 + + # append punctuations + i = 0 + j = 1 + while j < len(words): + if not words[i].endswith(" ") and words[j] in appended: + words[i] += words[j] + tokens[i] += tokens[j] + indices[i] += indices[j] + words[j] = "" + tokens[j] = [] + indices[j] = [] + else: + i = j + j += 1 + + # remove elements that are now empty + words[:] = [word for word in words if word] + tokens[:] = [token for token in tokens if token] + indices[:] = [idx for idx in indices if idx] + + +__all__ = ["WhisperTokenizer"] diff --git a/examples/mediatek/aot_utils/llm_utils/tokenizers_/utils.py b/examples/mediatek/aot_utils/llm_utils/tokenizers_/utils.py new file mode 100644 index 00000000000..a137e2c982e --- /dev/null +++ b/examples/mediatek/aot_utils/llm_utils/tokenizers_/utils.py @@ -0,0 +1,514 @@ +# Copyright (C) 2024 MediaTek Inc. All rights reserved. +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenizer utilities.""" + +# flake8: noqa: C901 + +import functools +import os +import sys +import types +from collections import UserDict +from enum import Enum +from pathlib import Path +from typing import Dict, Optional, Union +from uuid import uuid4 + +import huggingface_hub +import numpy as np +from huggingface_hub import _CACHED_NO_EXIST, hf_hub_download, try_to_load_from_cache + +get_torch_version = False +is_flax_available = False +is_mlx_available = False +is_tf_available = False +is_torch_available = True +is_torch_fx_proxy = False +is_offline_mode = True +is_protobuf_available = False + + +def is_remote_url(_): + """Check if the given input is a remote URL. + + Args: + _ (Any): The input to check. + + Returns: + bool: Always returns False. + """ + return False + + +def add_model_info_to_auto_map(auto_map, repo_id): + """Adds the information of the repo_id to a given auto map.""" + for key, value in auto_map.items(): + if isinstance(value, (tuple, list)): + auto_map[key] = [ + f"{repo_id}--{v}" if (v is not None and "--" not in v) else v + for v in value + ] + elif value is not None and "--" not in value: + auto_map[key] = f"{repo_id}--{value}" + + return auto_map + + +def add_model_info_to_custom_pipelines(custom_pipeline, repo_id): + """Adds the information of the repo_id to a given custom pipeline.""" + # {custom_pipelines : {task: {"impl": "path.to.task"},...} } + for task in custom_pipeline: + if "impl" in custom_pipeline[task]: + module = custom_pipeline[task]["impl"] + if "--" not in module: + custom_pipeline[task]["impl"] = f"{repo_id}--{module}" + return custom_pipeline + + +def infer_framework_from_repr(x): + """Guess the framework from repr. + + Tries to guess the framework of an object `x` from its repr (brittle but will help in `is_tensor` to try the + frameworks in a smart order, without the need to import the frameworks). + """ + representation = str(type(x)) + if representation.startswith(" str: + """Formats a user-agent string with basic info about a request.""" + ua = f"python/{sys.version.split()[0]}; session_id/{uuid4().hex}" + if huggingface_hub.constants.HF_HUB_DISABLE_TELEMETRY: + return ua + "; telemetry/off" + # CI will set this value to True + if isinstance(user_agent, dict): + ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items()) + elif isinstance(user_agent, str): + ua += "; " + user_agent + return ua + + +def cached_file( + path_or_repo_id: Union[str, os.PathLike], + filename: str, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[Dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + local_files_only: bool = False, + subfolder: str = "", + repo_type: Optional[str] = None, + user_agent: Optional[Union[str, Dict[str, str]]] = None, + _raise_exceptions_for_gated_repo: bool = True, + _raise_exceptions_for_missing_entries: bool = True, + _raise_exceptions_for_connection_errors: bool = True, + _commit_hash: Optional[str] = None, + **deprecated_kwargs, +) -> Optional[str]: + """Tries to locate a file in a local folder and repo, downloads and cache it if necessary. + + Args: + path_or_repo_id (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a model repo on huggingface.co. + - a path to a *directory* potentially containing the file. + filename (`str`): + The name of the file to locate in `path_or_repo`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded pretrained model configuration should be cached if the standard + cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128', + 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request. @lint-ignore + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + subfolder (`str`, *optional*, defaults to `""`): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can + specify the folder name here. + repo_type (`str`, *optional*): + Specify the repo type (useful when downloading from a space for instance). + user_agent (`str`): Use agent. + deprecated_kwargs: Deprecated kwargs. + + + + Passing `token=True` is required when you want to use a private model. + + + + Returns: + `Optional[str]`: Returns the resolved file (to the cache folder if downloaded from a repo). + + Examples: + ```python + # Download a model weight from the Hub and cache it. + model_weights_file = cached_file("google-bert/bert-base-uncased", "pytorch_model.bin") + ``` + """ + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if is_offline_mode and not local_files_only: + # print("Offline mode: forcing local_files_only=True") + local_files_only = True + if subfolder is None: + subfolder = "" + + path_or_repo_id = str(path_or_repo_id) + full_filename = os.path.join(subfolder, filename) + if os.path.isdir(path_or_repo_id): + resolved_file = os.path.join(os.path.join(path_or_repo_id, subfolder), filename) + if not os.path.isfile(resolved_file): + if _raise_exceptions_for_missing_entries and filename not in [ + "config.json", + f"{subfolder}/config.json", + ]: + raise OSError( + f"{path_or_repo_id} does not appear to have a file named {full_filename}. Checkout " + f"'https://huggingface.co/{path_or_repo_id}/tree/{revision}' for available files." + ) + return None + return resolved_file + + if cache_dir is None: + cache_dir = os.getenv( + "TRANSFORMERS_CACHE", + os.getenv( + "PYTORCH_TRANSFORMERS_CACHE", huggingface_hub.constants.HF_HUB_CACHE + ), + ) + if isinstance(cache_dir, Path): + cache_dir = str(cache_dir) + + if _commit_hash is not None and not force_download: + # If the file is cached under that commit hash, we return it directly. + resolved_file = try_to_load_from_cache( + path_or_repo_id, + full_filename, + cache_dir=cache_dir, + revision=_commit_hash, + repo_type=repo_type, + ) + if resolved_file is not None: + if resolved_file is not _CACHED_NO_EXIST: + return resolved_file + if not _raise_exceptions_for_missing_entries: + return None + raise OSError(f"Could not locate {full_filename} inside {path_or_repo_id}.") + + user_agent = http_user_agent(user_agent) + + return hf_hub_download( + path_or_repo_id, + filename, + subfolder=None if len(subfolder) == 0 else subfolder, + repo_type=repo_type, + revision=revision, + cache_dir=cache_dir, + user_agent=user_agent, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + + +def copy_func(f): + """Returns a copy of a function f.""" + # Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard) + g = types.FunctionType( + f.__code__, + f.__globals__, + name=f.__name__, + argdefs=f.__defaults__, + closure=f.__closure__, + ) + g = functools.update_wrapper(g, f) + g.__kwdefaults__ = f.__kwdefaults__ + return g diff --git a/examples/mediatek/aot_utils/llm_utils/utils.py b/examples/mediatek/aot_utils/llm_utils/utils.py index a7c242559a1..6c00bbb8446 100644 --- a/examples/mediatek/aot_utils/llm_utils/utils.py +++ b/examples/mediatek/aot_utils/llm_utils/utils.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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 os import sys @@ -292,54 +298,120 @@ def generate_mask( valid_input, batch_size=1, mask_value=-100.0, + medusa_mask=None, + sliding_window=False, + sliding_window_size=None, pytorch=True, + bita_prefix_length=0, + bita_draft_length=0, ): + assert ( valid_cache <= cache_size ), "valid_cache must be less than or equal to cache_size" - assert ( - valid_input <= input_length - ), "valid_input must be less than or equal to input_length" - # Cache mask portion - valid = np.zeros((1, 1, 1, valid_cache + input_length), dtype=np.float32) - cache_mask = np.full( - (1, 1, 1, cache_size - valid_cache), mask_value, dtype=np.float32 - ) - cache_mask = np.concatenate((cache_mask, valid), axis=-1) - cache_mask_final_shape = np.broadcast_to( - cache_mask, (batch_size, 1, input_length, cache_size + input_length) - ) - # Attention mask portion - mask_cond = np.arange(valid_input) - triangle = mask_cond >= (mask_cond + 1).reshape(valid_input, 1) - small_attention_mask = triangle.astype(np.float32) * mask_value - attention_mask = np.pad( - small_attention_mask, - (0, input_length - valid_input), - "constant", - constant_values=mask_value, - ) - attention_mask_with_cache = np.concatenate( - [np.zeros((input_length, cache_size), dtype=np.float32), attention_mask], - axis=-1, - ) - attention_mask_final_shape = np.broadcast_to( - attention_mask_with_cache[None, None, :, :], - (batch_size, 1, input_length, cache_size + input_length), - ) + if sliding_window: # SWA + assert sliding_window_size is not None, "sliding_window_size is None" + # Full mask: (input_length, cache_size + input_length) + full_mask_length = cache_size + input_length + # 1. generate invalid input portion + invalid_input_part = np.full( + (batch_size, 1, input_length - valid_input, full_mask_length), + mask_value, + dtype=np.float32, + ) + + # 2. valid input + invalid cache part + invalid_cache_part = np.full( + (batch_size, 1, valid_input, cache_size - valid_cache), + mask_value, + dtype=np.float32, + ) + + # 3. valid input + invalid input part + valid_input_vs_invalid_input_part = np.full( + (batch_size, 1, valid_input, input_length - valid_input), + mask_value, + dtype=np.float32, + ) + + # 4. valid input + valid_cache, valid_input part + valid_part = np.full( + (batch_size, 1, valid_input, valid_cache + valid_input), + mask_value, + dtype=np.float32, + ) + + # swap to zeros when within sliding window size + for i in range(valid_part.shape[2]): + end_position = valid_cache + i + start_position = max(0, end_position - sliding_window_size + 1) + valid_part[:, :, i, start_position : end_position + 1] = 0 + + # combine everything together + final_mask = np.concatenate( + [invalid_cache_part, valid_part, valid_input_vs_invalid_input_part], axis=-1 + ) + final_mask = np.concatenate([final_mask, invalid_input_part], axis=-2) + combined_mask = final_mask + + else: # Normal attention + assert ( + valid_input <= input_length + ), "valid_input must be less than or equal to input_length" + # Cache mask portion + valid = np.zeros((1, 1, 1, valid_cache + input_length), dtype=np.float32) + cache_mask = np.full( + (1, 1, 1, cache_size - valid_cache), mask_value, dtype=np.float32 + ) + cache_mask = np.concatenate((cache_mask, valid), axis=-1) + cache_mask_final_shape = np.broadcast_to( + cache_mask, (batch_size, 1, input_length, cache_size + input_length) + ) + + # Attention mask portion + mask_cond = np.arange(valid_input) + triangle = mask_cond >= (mask_cond + 1).reshape(valid_input, 1) + small_attention_mask = triangle.astype(np.float32) * mask_value + attention_mask = np.pad( + small_attention_mask, + (0, input_length - valid_input), + "constant", + constant_values=mask_value, + ) + attention_mask_with_cache = np.concatenate( + [np.zeros((input_length, cache_size), dtype=np.float32), attention_mask], + axis=-1, + ) + attention_mask_final_shape = np.broadcast_to( + attention_mask_with_cache[None, None, :, :], + (batch_size, 1, input_length, cache_size + input_length), + ) - combined_mask = attention_mask_final_shape + cache_mask_final_shape + combined_mask = attention_mask_final_shape + cache_mask_final_shape + combined_mask[ + :, + :, + -bita_draft_length:, + -valid_cache + - input_length + - bita_prefix_length : -valid_cache + - input_length, + ] = 0 + + if medusa_mask is not None: + medusa_len = medusa_mask.size(-1) + combined_mask[:, :, -medusa_len:, -medusa_len:][medusa_mask == 0] = mask_value if pytorch: return torch.from_numpy(combined_mask.copy()) return combined_mask.copy() -def get_dest_path(output_folder, exp_name, shape, chunk_idx): - dest_folder_root = output_folder + f"_{shape}" +def get_dest_path(output_folder, exp_name, shape=None, chunk_idx=0): + dest_folder_root = output_folder + f"{f'_{shape}' if shape is not None else ''}" os.makedirs(dest_folder_root, exist_ok=True) - fname = f"{exp_name}_{shape}_{chunk_idx}.pte" + fname = f"{exp_name}{f'_{shape}' if shape is not None else ''}_{chunk_idx}.pte" dest_path = os.path.join(dest_folder_root, fname) return dest_path @@ -387,36 +459,319 @@ def get_export_shapes(shapes): return export_shapes, max_num_token, max_cache_size -def get_master_rot_emb(config, dtype): - rot_dim = int(config.hidden_size / config.num_attention_heads) - length = config.max_position_embeddings +def get_master_pos_emb(config, weight_dir, dtype, **kwargs): + """Generates the master positional embeddings. + + Args: + config (object): The configuration object containing model parameters. + dtype (type): The data type of the embeddings. + kwargs: Other kwargs. + + Returns: + np.ndarray or torch.Tensor: The master positional embeddings. + """ + + state_dict = kwargs.pop("state_dict", None) + + if config.model_type == "whisper_decoder": + weight_dir = weight_dir + if weight_dir is None and state_dict is None: + raise ValueError("Expect either weight_dir or state_dict, but got neither") + + expected_path = os.path.join(weight_dir, "embedding_pos_fp16.bin") + if os.path.exists(expected_path): + embedding_weight = np.fromfile(expected_path, dtype=np.float16).reshape( + 1, config.max_position_embeddings, config.hidden_size + ) + else: + embedding_weight = ( + state_dict["model.decoder.embed_positions.weight"].unsqueeze(0).numpy() + ) + embedding_weight.astype(np.float16).tofile(expected_path) + + if isinstance(dtype, torch.dtype): + return torch.tensor(embedding_weight).to(dtype) + return embedding_weight.astype(dtype) + + +def get_master_rot_emb(config, dtype, **kwargs): + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1) + if partial_rotary_factor is None: + partial_rotary_factor = 1 + + rot_dim = ( + int(config.head_dim * partial_rotary_factor) + if config.model_type not in ["gemma2", "gemma3", "qwen2", "qwen3"] + else config.head_dim + ) + length = int(config.max_position_embeddings * config.ntk_scaling_factor) if config.ntk_scaling_factor != 1.0: base = (10000 * config.ntk_scaling_factor) ** (rot_dim / (rot_dim - 2)) else: base = 10000 - inv_freq = 1.0 / ( - base ** (np.arange(0, rot_dim, 2, dtype=np.float32) / rot_dim) - ) # (rot_dim/2) - t = np.arange(length, dtype=np.float32) # (len) - freqs = np.einsum("i,j->ij", t, inv_freq) # (len, rot_dim/2) - emb = np.concatenate((freqs, freqs), axis=-1) # (len, rot_dim) - master_cos = np.cos(emb)[None, None, :, :] # (1,1,len,rot_dim) - master_sin = np.sin(emb)[None, None, :, :] # (1,1,len,rot_dim) + if getattr(config, "rope_scaling", None) is not None: + if config.rope_scaling["type"] == "longrope": + short_factor = config.rope_scaling["short_factor"] + long_factor = config.rope_scaling["long_factor"] + original_max_position_embeddings = config.original_max_position_embeddings + length = original_max_position_embeddings + assert original_max_position_embeddings is not None, ( + "Using longrope for rope scaling but" + "original_max_position_embeddings is None." + ) + + ext_factors_long = torch.tensor(long_factor, dtype=torch.float32) + ext_factors = torch.tensor(short_factor, dtype=torch.float32) - rot_emb = np.concatenate((master_cos, master_sin), axis=1) + inv_freq_shape = ( + torch.arange(0, rot_dim, 2, dtype=torch.int64).float() / rot_dim + ) + inv_freq = 1.0 / (ext_factors * base**inv_freq_shape) + inv_freq = inv_freq.unsqueeze(1) + t = torch.arange(length, dtype=torch.float32) + t = t.unsqueeze(0) + + inv_freq_long = 1.0 / (ext_factors_long * base**inv_freq_shape) + inv_freq_long = inv_freq_long.unsqueeze(1) + t_long = torch.arange( + config.max_position_embeddings - length, dtype=torch.float32 + ) + t_long = t_long.unsqueeze(0) + + # Short factor + freqs = (inv_freq.float() @ t.float()).transpose(0, 1) + emb = torch.cat((freqs, freqs), dim=-1) + + scale = config.max_position_embeddings / original_max_position_embeddings + if scale <= 1.0: + scaling_factor = 1.0 + else: + scaling_factor = math.sqrt( + 1 + math.log(scale) / math.log(original_max_position_embeddings) + ) + master_cos = emb.cos() * scaling_factor + master_sin = emb.sin() * scaling_factor + master_cos = master_cos[None, None, :, :] + master_sin = master_sin[None, None, :, :] + + # Long factor + freqs_long = (inv_freq_long.float() @ t_long.float()).transpose(0, 1) + emb_long = torch.cat((freqs_long, freqs_long), dim=-1) + + scale = config.max_position_embeddings / original_max_position_embeddings + if scale <= 1.0: + scaling_factor = 1.0 + else: + scaling_factor = math.sqrt( + 1 + math.log(scale) / math.log(original_max_position_embeddings) + ) + master_cos_long = emb_long.cos() * scaling_factor + master_sin_long = emb_long.sin() * scaling_factor + master_cos_long = master_cos_long[None, None, :, :] + master_sin_long = master_sin_long[None, None, :, :] + + master_cos = torch.cat((master_cos, master_cos_long), dim=2) + master_sin = torch.cat((master_sin, master_sin_long), dim=2) + rot_emb = torch.cat((master_cos, master_sin), dim=1) + + elif config.rope_scaling["type"] == "mrope": + # position_ids = config.e.mrope_position_ids + # rope_delta = config.e.mrope_delta + position_ids = kwargs.get("qwen2_vl_position_ids") + rope_delta = kwargs.get("qwen2_vl_mrope_delta") + assert ( + position_ids is not None + ), "Must pass position_ids when using Qwen2-VL mrope." + assert ( + rope_delta is not None + ), "Must pass rope_delta when using Qwen2-VL mrope." + + mrope_section = config.rope_scaling["mrope_section"] * 2 + cache_position = torch.arange( + position_ids.shape[2], + config.max_position_embeddings, + dtype=torch.float32, + ) + delta = cache_position + rope_delta + delta = delta.view(1, -1).expand(1, -1) + delta = delta.unsqueeze(0).expand(3, -1, -1) + position_ids = torch.cat([position_ids, delta], dim=-1) + + inv_freq = 1.0 / ( + base ** (torch.arange(0, rot_dim, 2, dtype=torch.float32) / rot_dim) + ) # (rot_dim/2) + inv_freq_expanded = ( + inv_freq[None, None, :, None] + .float() + .expand(3, position_ids.shape[1], -1, 1) + ) + position_ids_expanded = position_ids[ + :, :, None, : + ].float() # (3, bs, 1, positions) + freqs = ( + inv_freq_expanded.float() @ position_ids_expanded.float() + ).transpose(2, 3) + emb = torch.cat((freqs, freqs), dim=-1) + master_cos = emb.cos() + master_sin = emb.sin() + + cos = torch.cat( + [ + m[i % 3] + for i, m in enumerate(master_cos.split(mrope_section, dim=-1)) + ], + dim=-1, + ).unsqueeze(1) + sin = torch.cat( + [ + m[i % 3] + for i, m in enumerate(master_sin.split(mrope_section, dim=-1)) + ], + dim=-1, + ).unsqueeze(1) + rot_emb = torch.cat([cos, sin], dim=1) + + elif config.rope_scaling["type"] == "yarn": + scaling_factor = config.rope_scaling["factor"] + extrapolation_factor = config.rope_scaling.get("extrapolation_factor", 1) + attn_factor = config.rope_scaling.get("attn_factor", 1) + beta_fast = config.rope_scaling.get("beta_fast", 32) + beta_slow = config.rope_scaling.get("beta_slow", 1) + original_max_position_embeddings = config.original_max_position_embeddings + assert ( + original_max_position_embeddings is not None + ), "Using Yarn for rope scaling but original max_position_embeddings is None." + + def _yarn_find_correction_dim( + num_rotations, dim, base, max_position_embeddings + ): + return ( + dim + * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) + ) / (2 * math.log(base)) + + def _yarn_find_correction_range( + low_rot, high_rot, dim, base, max_position_embeddings + ): + low = math.floor( + _yarn_find_correction_dim( + low_rot, dim, base, max_position_embeddings + ) + ) + high = math.ceil( + _yarn_find_correction_dim( + high_rot, dim, base, max_position_embeddings + ) + ) + return max(low, 0), min(high, dim - 1) # Clamp values just in case + + def _yarn_linear_ramp_mask(min_val, max_val, dim): + if min_val == max_val: + max_val += 0.001 # Prevent singularity + + linear_func = (np.arange(dim, dtype=np.float32) - min_val) / ( + max_val - min_val + ) + return np.clip(linear_func, 0, 1) + + def _yarn_get_softmax_scale(config, scaling_factor): + if scaling_factor <= 1 or config.model_type != "llama": + return 1.0 + return 0.1 * math.log(scaling_factor) + 1.0 + + # yarn has 3 parts: interpolation, extrapolation and linear ramp + inv_freq_extrapolation = 1.0 / ( + base ** (np.arange(0, rot_dim, 2, dtype=np.float32) / rot_dim) + ) # (rot_dim/2) + inv_freq_interpolation = inv_freq_extrapolation / scaling_factor + + low, high = _yarn_find_correction_range( + beta_fast, beta_slow, rot_dim, base, original_max_position_embeddings + ) + inv_freq_mask = ( + 1 - _yarn_linear_ramp_mask(low, high, rot_dim // 2).astype(np.float32) + ) * extrapolation_factor + inv_freq = ( + inv_freq_interpolation * (1 - inv_freq_mask) + + inv_freq_extrapolation * inv_freq_mask + ) + + t = np.arange(length, dtype=np.float32) # (len) + freqs = np.einsum("i,j->ij", t, inv_freq) # (len, rot_dim/2) + emb = np.concatenate((freqs, freqs), axis=-1) # (len, rot_dim) + additional_softmax_scale = ( + _yarn_get_softmax_scale(config, scaling_factor) * attn_factor + ) + master_cos = ( + np.cos(emb)[None, None, :, :] * additional_softmax_scale + ) # (1,1,len,rot_dim) + master_sin = ( + np.sin(emb)[None, None, :, :] * additional_softmax_scale + ) # (1,1,len,rot_dim) + + rot_emb = np.concatenate((master_cos, master_sin), axis=1) + + elif config.rope_scaling["type"] == "linear": + inv_freq = 1.0 / ( + base ** (np.arange(0, rot_dim, 2, dtype=np.float32) / rot_dim) + ) # (rot_dim/2) + inv_freq /= config.rope_scaling["factor"] + t = np.arange(length, dtype=np.float32) # (len) + freqs = np.einsum("i,j->ij", t, inv_freq) # (len, rot_dim/2) + emb = np.concatenate((freqs, freqs), axis=-1) # (len, rot_dim) + master_cos = np.cos(emb)[None, None, :, :] # (1,1,len,rot_dim) + master_sin = np.sin(emb)[None, None, :, :] # (1,1,len,rot_dim) + rot_emb = np.concatenate((master_cos, master_sin), axis=1) + + else: + assert False, ( + f"Rope scaling only supports longrope, mrope, yarn and linear," + f'but got {config.rope_scaling["type"]}' + ) + + else: + inv_freq = 1.0 / ( + base ** (np.arange(0, rot_dim, 2, dtype=np.float32) / rot_dim) + ) # (rot_dim/2) + t = np.arange(length, dtype=np.float32) # (len) + freqs = np.einsum("i,j->ij", t, inv_freq) # (len, rot_dim/2) + emb = np.concatenate((freqs, freqs), axis=-1) # (len, rot_dim) + master_cos = np.cos(emb)[None, None, :, :] # (1,1,len,rot_dim) + master_sin = np.sin(emb)[None, None, :, :] # (1,1,len,rot_dim) + + rot_emb = np.concatenate((master_cos, master_sin), axis=1) if isinstance(dtype, torch.dtype): + if isinstance(rot_emb, torch.Tensor): + return rot_emb.to(dtype) return torch.from_numpy(rot_emb).to(dtype) - else: - return rot_emb.astype(dtype) + if isinstance(rot_emb, torch.Tensor): + return rot_emb.detach().cpu().numpy().astype(dtype) + return rot_emb.astype(dtype) def get_normalized_config(config_filepath): config_file = json.load(open(config_filepath, "r")) - if config_file["model_type"] == "llama": - from models.llm_models.configuration_llama import LlamaConfig as config_class + if config_file.get("llm", None) is None: + if config_file["model_type"] == "llama": + from models.llm_models.configuration_llama import ( + LlamaConfig as config_class, + ) + elif config_file["model_type"] in ["qwen3", "qwen2"]: + from models.llm_models.configuration_qwen import QwenConfig as config_class + elif config_file["model_type"] in ["phi3", "phi4"]: + from models.llm_models.configuration_phi import PhiConfig as config_class + elif config_file["model_type"] in ["gemma1", "gemma2", "gemma3"]: + from models.llm_models.configuration_gemma import ( + GemmaConfig as config_class, + ) + else: + if config_file["llm"]["model_type"] == "whisper_decoder": + from models.llm_models.configuration_whisper import ( + WhisperConfig as config_class, + ) config = config_class(**config_file, verbose=False) return config @@ -470,25 +825,105 @@ def resolve_model_classes( ): config_file = json.load(open(config_filepath, "r")) weight_dir = get_dirname(config_filepath) - if config_file["model_type"] == "llama": - from models.llm_models.configuration_llama import LlamaConfig as config_class - from models.llm_models.modeling_llama import LlamaModelChunk as chunk_class + + if config_file.get("llm", None) is None: + if config_file["model_type"] == "llama": + from models.llm_models.configuration_llama import ( + LlamaConfig as config_class, + ) + from models.llm_models.modeling_llama import LlamaModelChunk as chunk_class + elif config_file["model_type"] == "qwen2": + from models.llm_models.configuration_qwen import QwenConfig as config_class + from models.llm_models.modeling_qwen import Qwen2ModelChunk as chunk_class + elif config_file["model_type"] == "qwen3": + from models.llm_models.configuration_qwen import QwenConfig as config_class + from models.llm_models.modeling_qwen import Qwen3ModelChunk as chunk_class + elif config_file["model_type"] == "phi3": + from models.llm_models.configuration_phi import PhiConfig as config_class + from models.llm_models.modeling_phi import Phi3ModelChunk as chunk_class + elif config_file["model_type"] == "phi4": + from models.llm_models.configuration_phi import PhiConfig as config_class + from models.llm_models.modeling_phi import Phi4ModelChunk as chunk_class + elif config_file["model_type"] == "gemma2": + from models.llm_models.configuration_gemma import ( + GemmaConfig as config_class, + ) + from models.llm_models.modeling_gemma import Gemma2ModelChunk as chunk_class + elif config_file["model_type"] == "gemma3": + from models.llm_models.configuration_gemma import ( + GemmaConfig as config_class, + ) + from models.llm_models.modeling_gemma import Gemma3ModelChunk as chunk_class + else: + if config_file["llm"]["model_type"] == "whisper_decoder": + from models.llm_models.configuration_whisper import ( + WhisperConfig as config_class, + ) + from models.llm_models.modeling_whisper import ( + WhisperDecoderModelChunk as decoder_class, + WhisperEncoderModel as encoder_class, + ) + + chunk_class = [encoder_class, decoder_class] config = config_class(**config_file, response_handler=response_handler) if bypass_tokenizer: return config, weight_dir, chunk_class else: if config.tokenizer == "default": - if config_file["model_type"] == "llama": - from aot_utils.llm_utils.tokenizers_.tokenization_llama import ( - LlamaTokenizer as tokenizer_class, - ) + if config_file.get("llm", None) is None: + if config_file["model_type"] in ["llama", "phi3"]: + from aot_utils.llm_utils.tokenizers_.tokenization_llama import ( + LlamaTokenizer as tokenizer_class, + ) + elif config_file["model_type"] in ["qwen3", "qwen2"]: + from aot_utils.llm_utils.tokenizers_.tokenization_qwen2_fast import ( + Qwen2TokenizerFast as tokenizer_class, + ) + elif config_file["model_type"] in ["gemma3", "gemma2", "gemma1"]: + from aot_utils.llm_utils.tokenizers_.tokenization_gemma_fast import ( + GemmaTokenizerFast as tokenizer_class, + ) + elif config_file["model_type"] == "phi4": + from aot_utils.llm_utils.tokenizers_.tokenization_gpt2 import ( + GPT2Tokenizer as tokenizer_class, + ) + else: + if config_file["llm"]["model_type"] == "whisper_decoder": + from aot_utils.llm_utils.tokenizers_.tokenization_whisper import ( + WhisperTokenizer as tokenizer_class, + ) else: if config.tokenizer == "llama": from aot_utils.llm_utils.tokenizers_.tokenization_llama import ( LlamaTokenizer as tokenizer_class, ) + elif config.tokenizer == "qwen2": + from aot_utils.llm_utils.tokenizers_.tokenization_qwen2 import ( + Qwen2Tokenizer as tokenizer_class, + ) + elif config.tokenizer == "qwen2_fast": + from aot_utils.llm_utils.tokenizers_.tokenization_qwen2_fast import ( + Qwen2TokenizerFast as tokenizer_class, + ) + elif config.tokenizer == "gemma": + from aot_utils.llm_utils.tokenizers_.tokenization_gemma import ( + GemmaTokenizer as tokenizer_class, + ) + elif config.tokenizer == "gemma_fast": + from aot_utils.llm_utils.tokenizers_.tokenization_gemma_fast import ( + GemmaTokenizerFast as tokenizer_class, + ) + elif config.tokenizer == "gpt2": + from aot_utils.llm_utils.tokenizers_.tokenization_gpt2 import ( + GPT2Tokenizer as tokenizer_class, + ) elif config.tokenizer == "pretrained_fast": from aot_utils.llm_utils.tokenizers_.tokenization_utils_fast import ( PreTrainedTokenizerFast as tokenizer_class, ) + elif config.tokenizer == "whisper_decoder": + from aot_utils.llm_utils.tokenizers_.tokenization_whisper import ( + WhisperTokenizer as tokenizer_class, + ) + return config, weight_dir, tokenizer_class, chunk_class diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/audio_paths.txt b/examples/mediatek/aot_utils/mllm_utils/audio/audio_paths.txt new file mode 100644 index 00000000000..98b430a7f92 --- /dev/null +++ b/examples/mediatek/aot_utils/mllm_utils/audio/audio_paths.txt @@ -0,0 +1,3 @@ +/ssd1/joe/executorch/examples/mediatek/aot_utils/mllm_utils/audio/nggyu.mp3 +/ssd1/joe/executorch/examples/mediatek/aot_utils/mllm_utils/audio/bruce.mp3 +/ssd1/joe/executorch/examples/mediatek/aot_utils/mllm_utils/audio/Jimmy.mp3 diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/bruce.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/bruce.mp3 new file mode 100644 index 00000000000..b9607ad6d5c Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/bruce.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/french.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/french.mp3 new file mode 100644 index 00000000000..44d69ded534 Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/french.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/nggyu.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/nggyu.mp3 new file mode 100644 index 00000000000..6ccbe6284f3 Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/nggyu.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/sample-4.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/sample-4.mp3 new file mode 100644 index 00000000000..f6dadd3b92c Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/sample-4.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/seven_miles_cut.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/seven_miles_cut.mp3 new file mode 100644 index 00000000000..f88e084303d Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/seven_miles_cut.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/audio/those_years_cut.mp3 b/examples/mediatek/aot_utils/mllm_utils/audio/those_years_cut.mp3 new file mode 100644 index 00000000000..5445cd2425c Binary files /dev/null and b/examples/mediatek/aot_utils/mllm_utils/audio/those_years_cut.mp3 differ diff --git a/examples/mediatek/aot_utils/mllm_utils/preprocessor_whisper.py b/examples/mediatek/aot_utils/mllm_utils/preprocessor_whisper.py new file mode 100644 index 00000000000..b9e88a9e8e2 --- /dev/null +++ b/examples/mediatek/aot_utils/mllm_utils/preprocessor_whisper.py @@ -0,0 +1,324 @@ +# Copyright 2022 The HuggingFace Inc. team. # noqa: CPY001 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Preprocessor classes.""" + +# flake8: noqa: C901 + +from typing import List, Optional, Union + +import numpy as np +from transformers.audio_utils import mel_filter_bank, spectrogram, window_function +from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor +from transformers.feature_extraction_utils import BatchFeature + + +class WhisperAudioProcessor(SequenceFeatureExtractor): + """Define Whisper Audio preprocessor.""" + + model_input_names = ["input_features"] + + def __init__( + self, + feature_size=80, + sampling_rate=16000, + hop_length=160, + chunk_length=30, + n_fft=400, + padding_value=0.0, + return_attention_mask=False, + **kwargs, + ) -> None: + """Initializes the WhisperAudioProcessor. + + Args: + feature_size (`int`, *optional*, defaults to 80): + The feature dimension of the extracted features. + sampling_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + hop_length (`int`, *optional*, defaults to 160): + Length of the overlaping windows for the STFT used to obtain the Mel Frequency coefficients. + chunk_length (`int`, *optional*, defaults to 30): + The maximum number of chuncks of `sampling_rate` samples used to trim and pad longer or shorter audio + sequences. + n_fft (`int`, *optional*, defaults to 400): + Size of the Fourier transform. + padding_value (`float`, *optional*, defaults to 0.0): + Padding value used to pad the audio. Should correspond to silences. + return_attention_mask (`bool`, *optional*): + Pad inputs to max length with silence token (zero) and no attention mask + kwargs: Additional keyword arguments. + """ + super().__init__( + feature_size=feature_size, + sampling_rate=sampling_rate, + padding_value=padding_value, + return_attention_mask=return_attention_mask, + **kwargs, + ) + self.n_fft = n_fft + self.hop_length = hop_length + self.chunk_length = chunk_length + self.n_samples = chunk_length * sampling_rate + self.nb_max_frames = self.n_samples // hop_length + self.sampling_rate = sampling_rate + self.mel_filters = mel_filter_bank( + num_frequency_bins=1 + n_fft // 2, + num_mel_filters=feature_size, + min_frequency=0.0, + max_frequency=8000.0, + sampling_rate=sampling_rate, + norm="slaney", + mel_scale="slaney", + ) + + def _np_extract_fbank_features( + self, waveform_batch: np.array, device: str + ) -> np.ndarray: + """Extract fbank features. + + Compute the log-mel spectrogram of the provided audio, gives similar results to Whisper's original torch + implementation with 1e-5 tolerance. + """ + if device != "cpu": + raise ValueError( + f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator " + "devices requires torch, which is not installed. Either set `device='cpu'`, or " + "install torch according to the official instructions: https://pytorch.org/get-started/locally/", + ) + log_spec_batch = [] + for waveform in waveform_batch: + log_spec = spectrogram( + waveform, + window_function(self.n_fft, "hann"), + frame_length=self.n_fft, + hop_length=self.hop_length, + power=2.0, + mel_filters=self.mel_filters, + log_mel="log10", + ) + log_spec = log_spec[:, :-1] + log_spec = np.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + log_spec_batch.append(log_spec) + return np.array(log_spec_batch) + + @staticmethod + def zero_mean_unit_var_norm( + input_values: List[np.ndarray], + attention_mask: List[np.ndarray], + padding_value: float = 0.0, + ) -> List[np.ndarray]: + """Every array in the list is normalized to have zero mean and unit variance.""" + if attention_mask is not None: + attention_mask = np.array(attention_mask, np.int32) + normed_input_values = [] + + for vector, length in zip(input_values, attention_mask.sum(-1)): + normed_slice = (vector - vector[:length].mean()) / np.sqrt( + vector[:length].var() + 1e-7 + ) + if length < normed_slice.shape[0]: + normed_slice[length:] = padding_value + + normed_input_values.append(normed_slice) + else: + normed_input_values = [ + (x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values + ] + + return normed_input_values + + def preprocess( + self, + raw_speech: Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]], + truncation: bool = True, + pad_to_multiple_of: Optional[int] = None, + return_tensors: Optional[Union[str]] = None, + return_attention_mask: Optional[bool] = None, + padding: Optional[str] = "max_length", + max_length: Optional[int] = None, + sampling_rate: Optional[int] = None, + do_normalize: Optional[bool] = None, + device: Optional[str] = "cpu", + return_token_timestamps: Optional[bool] = None, + **kwargs, + ) -> BatchFeature: + """Preprocess function. + + Main method to featurize and prepare for the model one or several sequence(s). Implementation uses PyTorch for + the STFT computation if available, otherwise a slower NumPy based one. + + Args: + raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`): + The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not + stereo, i.e. single float per timestep. + truncation (`bool`, *optional*, default to `True`): + Activates truncation to cut input sequences longer than *max_length* to *max_length*. + pad_to_multiple_of (`int`, *optional*, defaults to None): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific feature_extractor's default. + + [What are attention masks?](../glossary#attention-mask) @lint-ignore + + + + For Whisper models, `attention_mask` should always be passed for batched inference, to avoid subtle + bugs. + + + + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + padding (`str`, *optional*): + Padding strategy. + max_length (`int`, *optional*): + Maximum length. + sampling_rate (`int`, *optional*): + The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass + `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition + pipeline. + padding_value (`float`, *optional*, defaults to 0.0): + The value that is used to fill the padding values / vectors. + do_normalize (`bool`, *optional*, defaults to `False`): + Whether or not to zero-mean unit-variance normalize the input. Normalizing can help to significantly + improve the performance of the model. + device (`str`, *optional*, defaults to `'cpu'`): + Specifies the device for computation of the log-mel spectrogram of audio signals in the + `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda") + return_token_timestamps (`bool`, *optional*, defaults to `None`): + Whether or not to return the number of frames of the input raw_speech. + These num_frames can be used by the model to compute word level timestamps. + kwargs: Additional keyword arguments. + """ + if sampling_rate is not None: + if sampling_rate != self.sampling_rate: + raise ValueError( + f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a" + f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input" + f" was sampled with {self.sampling_rate} and not {sampling_rate}." + ) + else: + print( + "It is strongly recommended to pass the `sampling_rate` argument to this function. " + "Failing to do so can result in silent errors that might be hard to debug." + ) + + is_batched_numpy = ( + isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1 + ) + if is_batched_numpy and len(raw_speech.shape) > 2: + raise ValueError( + f"Only mono-channel audio is supported for input to {self}" + ) + is_batched = is_batched_numpy or ( + isinstance(raw_speech, (list, tuple)) + and (isinstance(raw_speech[0], (np.ndarray, tuple, list))) + ) + + if is_batched: + raw_speech = [ + np.asarray([speech], dtype=np.float32).T for speech in raw_speech + ] + elif not is_batched and not isinstance(raw_speech, np.ndarray): + raw_speech = np.asarray(raw_speech, dtype=np.float32) + elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype( + np.float64 + ): + raw_speech = raw_speech.astype(np.float32) + + # always return batch + if not is_batched: + raw_speech = [np.asarray([raw_speech]).T] + + batched_speech = BatchFeature({"input_features": raw_speech}) + + # convert into correct format for padding + + padded_inputs = self.pad( + batched_speech, + padding=padding, + max_length=max_length if max_length else self.n_samples, + truncation=truncation, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask or do_normalize, + ) + + # zero-mean and unit-variance normalization + if do_normalize: + padded_inputs["input_features"] = self.zero_mean_unit_var_norm( + padded_inputs["input_features"], + attention_mask=padded_inputs["attention_mask"], + padding_value=self.padding_value, + ) + padded_inputs["input_features"] = np.stack( + padded_inputs["input_features"], axis=0 + ) + + # make sure list is in array format + input_features = padded_inputs.get("input_features").transpose(2, 0, 1) + + extract_fbank_features = self._np_extract_fbank_features + input_features = extract_fbank_features(input_features[0], device) + + if isinstance(input_features[0], List): + padded_inputs["input_features"] = [ + np.asarray(feature, dtype=np.float32) for feature in input_features + ] + + else: + padded_inputs["input_features"] = input_features + + if return_attention_mask: + # rescale from sample (48000) to feature (3000) + padded_inputs["attention_mask_audio"] = padded_inputs["attention_mask"][ + :, :: self.hop_length + ] + padded_inputs.pop( + "attention_mask" + ) # change name to prevent conflict with vision + input_lengths = padded_inputs["attention_mask_audio"].sum(-1)[0] + input_lengths = (input_lengths - 1) // 2 + 1 + if "audio_output_lengths" in kwargs: + kwargs["audio_output_lengths"].append((input_lengths - 2) // 2 + 1) + else: + kwargs["audio_output_lengths"] = [(input_lengths - 2) // 2 + 1] + batch_size, _, max_mel_seq_len = input_features.shape + max_seq_len = (max_mel_seq_len - 2) // 2 + 1 + # Create a sequence tensor of shape (batch_size, max_seq_len) + seq_range = np.arange(0, max_seq_len) + lengths_expand = np.repeat(input_lengths, max_seq_len) + padding_mask = seq_range >= lengths_expand + padded_inputs["attention_mask_audio"] = np.reshape( + padding_mask * -200.0, (batch_size, 1, max_seq_len) + ) + + if return_token_timestamps is not None: + padded_inputs["num_frames"] = [ + len(raw_speech_i) // self.hop_length for raw_speech_i in raw_speech + ] + + if return_tensors is not None: + padded_inputs = padded_inputs.convert_to_tensors(return_tensors) + + return padded_inputs, kwargs diff --git a/examples/mediatek/aot_utils/oss_utils/utils.py b/examples/mediatek/aot_utils/oss_utils/utils.py index e365309f10c..8aef4bc53e5 100755 --- a/examples/mediatek/aot_utils/oss_utils/utils.py +++ b/examples/mediatek/aot_utils/oss_utils/utils.py @@ -5,9 +5,11 @@ # LICENSE file in the root directory of this source tree. import os +import random from typing import Optional import torch +import transformers from executorch import exir from executorch.backends.mediatek import ( NeuropilotPartitioner, @@ -33,9 +35,7 @@ def build_executorch_binary( if quant_dtype not in Precision: raise AssertionError(f"No support for Precision {quant_dtype}.") - captured_model = torch.export.export_for_training( - model, inputs, strict=True - ).module() + captured_model = torch.export.export(model, inputs, strict=True).module() annotated_model = prepare_pt2e(captured_model, quantizer) print("Quantizing the model...") # calibration @@ -44,6 +44,7 @@ def build_executorch_binary( quantized_model = convert_pt2e(annotated_model, fold_quantize=False) aten_dialect = torch.export.export(quantized_model, inputs, strict=True) else: + print("Using float model...") aten_dialect = torch.export.export(model, inputs, strict=True) from executorch.exir.program._program import to_edge_transform_and_lower @@ -73,3 +74,58 @@ def make_output_dir(path: str): os.remove(os.path.join(path, f)) os.removedirs(path) os.makedirs(path) + + +def get_masked_language_model_dataset(dataset_path, tokenizer, data_size, shuffle=True): + + def get_data_loader(): + class MaskedSentencesDataset(torch.utils.data.Dataset): + def __init__(self, dataset_path, tokenizer, data_size) -> None: + self.data_size = data_size + self.dataset = self._get_val_dataset(dataset_path, data_size, tokenizer) + + def _get_val_dataset(self, dataset_path, data_size, tokenizer): + data_collator = transformers.DataCollatorForLanguageModeling( + tokenizer=tokenizer + ) + with open(dataset_path, "r") as f: + texts = f.read().split("\n") + texts = [ + text for text in random.choices(texts, k=2000) if len(text) > 1 + ] + dataset = data_collator([tokenizer(text) for text in texts]) + return dataset + + def __getitem__(self, idx): + return ( + self.dataset["input_ids"][idx].to(torch.int32), + self.dataset["attention_mask"][idx].to(torch.float32), + self.dataset["labels"][idx], + ) + + def __len__(self): + return self.data_size + + dataset = MaskedSentencesDataset(dataset_path, tokenizer, data_size) + return torch.utils.data.DataLoader( + dataset, + shuffle=shuffle, + ) + + # prepare input data + inputs, targets = [], [] + data_loader = get_data_loader() + for data in data_loader: + if len(inputs) >= data_size: + break + input_ids = data[0] + attention_mask = data[1] + target = data[2][0] + indice = [i for i, x in enumerate(target) if x != -100] + # continue if no mask annotated + if len(indice) == 0: + continue + inputs.append((input_ids, attention_mask)) + targets.append(target) + + return inputs, targets diff --git a/examples/mediatek/executor_runner/llama_runner/LlamaConfig.h b/examples/mediatek/executor_runner/llama_runner/LlamaConfig.h index f512d59b5c5..2688ade477e 100644 --- a/examples/mediatek/executor_runner/llama_runner/LlamaConfig.h +++ b/examples/mediatek/executor_runner/llama_runner/LlamaConfig.h @@ -24,8 +24,11 @@ struct LlamaModelOptions { size_t hidden_size = 4096; size_t num_head = 32; size_t num_layer = 32; + size_t head_dim = 0; + size_t window_size = 0; size_t max_token_length = 2048; - double rot_emb_base = 10000.0f; + double partial_rotary_factor = 1.0; + double rot_emb_base = 10000.0; // Types LLMType model_input_type = LLMType::INT16; @@ -40,6 +43,7 @@ struct LlamaModelPaths { std::string token_embedding_path; std::vector prompt_model_paths; std::vector gen_model_paths; + std::vector model_package_paths; }; } // namespace example diff --git a/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp b/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp index bf4b3eefdde..b154759e245 100644 --- a/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp +++ b/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp @@ -21,6 +21,7 @@ #include "LlamaConfig.h" #include "LlamaModelChunk.h" +#include "Utils.h" #include "llm_helper/include/llm_types.h" #include "llm_helper/include/mask_builder.h" @@ -42,39 +43,60 @@ inline std::vector getIndexRange( LlamaModelChunk::LlamaModelChunk( const ModelPathMap& modelPathMap, const LlamaModelOptions& modelOptions, + const bool useSharedWeights, const size_t initBatchSize, const size_t numCache, const size_t numRotEmbInputs, + const bool enableSWA, const RotaryEmbeddingMasterLut* rotEmbMasterLut) : ModelChunk(modelPathMap, initBatchSize), + kIsSharedWeightsUsed(useSharedWeights), kMaxTokenLength(modelOptions.max_token_length), kCacheLength(modelOptions.cache_size), kMaskType(modelOptions.mask_type), kRotEmbMasterLut(rotEmbMasterLut), kCacheType(modelOptions.cache_type), - kCacheTypeSize(llm_helper::getLLMTypeSize(kCacheType)), - kMaskInputIndex(1), - kRotEmbInputIndexes(getIndexRange(2, numRotEmbInputs)), - kCacheInputIndexes( - getIndexRange(kRotEmbInputIndexes.back() + 1, numCache)), - kCacheOutputIndexes(getIndexRange(1, numCache)) {} + kWindowSize(modelOptions.window_size), -LlamaModelChunk::~LlamaModelChunk() {} + kRotEmbInputCount(numRotEmbInputs), + kCacheCount(numCache), + enableSWA(enableSWA), + kCacheTypeSize(llm_helper::getLLMTypeSize(kCacheType)) {} -size_t LlamaModelChunk::GetExpectedInputCount() const { - const size_t rotEmbInputCount = kRotEmbInputIndexes.size(); - const size_t cacheInputCount = kCacheInputIndexes.size(); - return 2 + rotEmbInputCount + cacheInputCount; -} +LlamaModelChunk::~LlamaModelChunk() {} -size_t LlamaModelChunk::GetExpectedOutputCount() const { - const size_t cacheOutputCount = kCacheOutputIndexes.size(); - return 1 + cacheOutputCount; +std::string LlamaModelChunk::SelectMethod( + const std::vector& methodNames) const { + const size_t curTokenSize = GetModelId(); + for (const auto& methodName : methodNames) { + const auto matches = utils::extract_substr(methodName, "([0-9]+)t[0-9]+c"); + if (matches.empty()) { + continue; + } + ET_CHECK_MSG( + matches.size() == 2, "Invalid method name: %s", methodName.c_str()); + // Extract the first match group as token size + const size_t methodTokenSize = + static_cast(std::atol(matches[1].c_str())); + if (curTokenSize == methodTokenSize) { + ET_LOG( + Debug, + "Selected method \"%s\" for token size %zu", + methodName.c_str(), + curTokenSize); + return methodName; + } + } + ET_LOG( + Error, + "Unable to find suitable method, fallback to use the first method."); + return {}; } void LlamaModelChunk::Initialize() { LoadModels(); GetModelIoInfo(); + defineIOs(); CheckIoCount(); PrepareCacheIOs(); AllocateIoBuffers(); @@ -86,19 +108,111 @@ void LlamaModelChunk::Initialize() { mIsInitialized = true; } +void LlamaModelChunk::defineIOs() { + // Inputs + defineInput(IOKind::Embedding); + defineInput(IOKind::Mask); + defineInput(IOKind::SWAMask, enableSWA); + defineInput(IOKind::RotEmb, kRotEmbInputCount); + defineInput(IOKind::KVCache, kCacheCount); + // Outputs + defineOutput(IOKind::Logits); + defineOutput(IOKind::KVCache, kCacheCount); +} + +void LlamaModelChunk::defineInput(const IOKind kind, const size_t count) { + ET_CHECK_MSG( + !hasInput(kind), + "Input kind has already been defined: %d", + static_cast(kind)); + const auto startIdx = mExpectedNumInputs; + std::vector indexes(count); + std::iota(indexes.begin(), indexes.end(), startIdx); + mInputIndexes[kind] = std::move(indexes); + mExpectedNumInputs += count; +} + +void LlamaModelChunk::defineOutput(const IOKind kind, const size_t count) { + ET_CHECK_MSG( + !hasOutput(kind), + "Output kind has already been defined: %d", + static_cast(kind)); + const auto startIdx = mExpectedNumOutputs; + std::vector indexes(count); + std::iota(indexes.begin(), indexes.end(), startIdx); + mOutputIndexes[kind] = std::move(indexes); + mExpectedNumOutputs += count; +} + +bool LlamaModelChunk::hasInput(const IOKind kind) const { + return (mInputIndexes.find(kind) != mInputIndexes.end()) && + !mInputIndexes.at(kind).empty(); +} + +bool LlamaModelChunk::hasOutput(const IOKind kind) const { + return (mOutputIndexes.find(kind) != mOutputIndexes.end()) && + !mOutputIndexes.at(kind).empty(); +} + +const std::vector& LlamaModelChunk::getInputIndexes( + const IOKind kind) const { + ET_CHECK_MSG( + hasInput(kind), "Check failed for input kind %d", static_cast(kind)); + return mInputIndexes.at(kind); +} + +const std::vector& LlamaModelChunk::getOutputIndexes( + const IOKind kind) const { + ET_CHECK_MSG( + hasOutput(kind), + "Check failed for output kind %d", + static_cast(kind)); + return mOutputIndexes.at(kind); +} + +size_t LlamaModelChunk::getInputIndex(const IOKind kind, const size_t pos) + const { + const auto& inputIndexes = getInputIndexes(kind); + ET_CHECK_MSG( + pos < inputIndexes.size(), "getInputIndex(): Index out of range"); + return inputIndexes[pos]; +} + +size_t LlamaModelChunk::getOutputIndex(const IOKind kind, const size_t pos) + const { + const auto& outputIndexes = getOutputIndexes(kind); + ET_CHECK_MSG( + pos < outputIndexes.size(), "getOutputIndex(): Index out of range"); + return outputIndexes[pos]; +} + +size_t LlamaModelChunk::getNumInputsFor(const IOKind kind) const { + if (!hasInput(kind)) { + return 0; + } + return mInputIndexes.at(kind).size(); +} + +size_t LlamaModelChunk::getNumOutputsFor(const IOKind kind) const { + if (!hasOutput(kind)) { + return 0; + } + return mOutputIndexes.at(kind).size(); +} + void LlamaModelChunk::CheckIoCount() { const auto& method = GetModelMethod(); const size_t modelInputCount = method.inputs_size(); const size_t modelOutputCount = method.outputs_size(); ET_CHECK_MSG( - modelInputCount == GetExpectedInputCount(), + modelInputCount == mExpectedNumInputs, "Number of inputs does not match (expected %zu but got %zu).", - GetExpectedInputCount(), + mExpectedNumInputs, modelInputCount); ET_CHECK_MSG( - modelOutputCount == GetExpectedOutputCount(), + modelOutputCount == mExpectedNumOutputs, "Number of outputs does not match (expected %zu but got %zu).", - GetExpectedOutputCount(), + mExpectedNumOutputs, modelOutputCount); } @@ -110,7 +224,8 @@ bool LlamaModelChunk::HotSwapModel(const size_t tokenBatchSize) { mMaskBuilder->markMaskDirty(); // Update mask size - const auto newMaskSizeBytes = mInputBufferInfos[kMaskInputIndex].nbytesUsed; + const auto newMaskSizeBytes = + mInputBufferInfos[getInputIndex(IOKind::Mask)].nbytesUsed; mMaskBuilder->updateMaskSize(newMaskSizeBytes); return status; @@ -171,7 +286,7 @@ void LlamaModelChunk::LeftPaddingCachePostprocess() { const size_t zeroCount = mCurrentPadSize * strideSizeBytes; // Fill padded sections with zeros - for (const auto cacheInputIdx : kCacheInputIndexes) { + for (const auto cacheInputIdx : getInputIndexes(IOKind::KVCache)) { auto cacheBuffer = reinterpret_cast(mInputBufferInfos[cacheInputIdx].data); for (size_t rowIdx = 0; rowIdx < numRows; rowIdx++) { @@ -212,7 +327,7 @@ void LlamaModelChunk::RollbackCache( const size_t numRows = GetCacheNumRows(); // Shift right and truncate rollbackTokCount, then fill left with zeros - for (const auto cacheInputIdx : kCacheInputIndexes) { + for (const auto cacheInputIdx : getInputIndexes(IOKind::KVCache)) { auto cacheBuffer = reinterpret_cast(mInputBufferInfos[cacheInputIdx].data); @@ -247,7 +362,26 @@ void LlamaModelChunk::UpdatePosEmbAndMask(const size_t numInputToken) { if (mCurrentTokenIndex > 0 && GetLeftPadding() > 0) { ET_LOG(Fatal, "Left-padding is only allowed in the first prompt pass."); } + auto isMaskUpdatable = mMaskBuilder->getMaskUpdateStatus(); + if (enableSWA) { + const auto& swaMaskBufferInfo = + mInputBufferInfos[getInputIndex(IOKind::SWAMask)]; + const auto swaMaskBuffer = swaMaskBufferInfo.data; + const auto swaMaskSizeBytes = swaMaskBufferInfo.nbytesUsed; + mMaskBuilder->setMaskBuffer(swaMaskBuffer, swaMaskSizeBytes); + mMaskBuilder->enableSlidingWindow(kWindowSize); + mMaskBuilder->buildMask(mTokenBatchSize, mCurrentTokenIndex); + } + // Pass same isMaskUpdatable to both mask + mMaskBuilder->setIsMaskUpdatable(isMaskUpdatable); + const auto& maskBufferInfo = mInputBufferInfos[getInputIndex(IOKind::Mask)]; + const auto maskBuffer = maskBufferInfo.data; + const auto maskSizeBytes = maskBufferInfo.nbytesUsed; + mMaskBuilder->setMaskBuffer(maskBuffer, maskSizeBytes); + mMaskBuilder->disableSlidingWindow(); mMaskBuilder->updateMask(mTokenBatchSize, mCurrentTokenIndex, numInputToken); + + mMaskBuilder->resetPadLength(); SetPosEmbed(mCurrentTokenIndex); } @@ -282,8 +416,9 @@ void LlamaModelChunk::SetPosEmbed(const size_t tokenIndex) { auto getRotEmbInputs = [&]() { std::vector rotEmbInputs; - rotEmbInputs.reserve(kRotEmbInputIndexes.size()); - for (const auto inputIdx : kRotEmbInputIndexes) + const auto& RotEmbInputIndexes = getInputIndexes(IOKind::RotEmb); + rotEmbInputs.reserve(RotEmbInputIndexes.size()); + for (const auto inputIdx : RotEmbInputIndexes) rotEmbInputs.push_back(mInputBufferInfos[inputIdx].data); return rotEmbInputs; }; @@ -298,13 +433,14 @@ void LlamaModelChunk::SetPosEmbed(const size_t tokenIndex) { void LlamaModelChunk::PrepareCacheIOs() { // Get cache shape const auto method_meta = GetModelMethod().method_meta(); - const auto firstInCacheIdx = kCacheInputIndexes.front(); + const auto firstInCacheIdx = getInputIndex(IOKind::KVCache); mCacheShape = method_meta.input_tensor_meta(firstInCacheIdx)->sizes(); // Link cache IOs - const size_t numCaches = kCacheInputIndexes.size(); + const size_t numCaches = getNumInputsFor(IOKind::KVCache); for (size_t i = 0; i < numCaches; i++) { - this->LinkModelIO(kCacheInputIndexes[i], kCacheOutputIndexes[i]); + this->LinkModelIO( + getInputIndex(IOKind::KVCache, i), getOutputIndex(IOKind::KVCache, i)); } } @@ -325,17 +461,31 @@ size_t LlamaModelChunk::GetCacheStrideSize() const { } void LlamaModelChunk::InitMaskBuilder() { - const auto& maskBufferInfo = mInputBufferInfos[kMaskInputIndex]; + mMaskBuilder = std::make_unique(kMaskType, kCacheLength); + // SWA Mask + if (enableSWA) { + const auto& swaMaskBufferInfo = + mInputBufferInfos[getInputIndex(IOKind::SWAMask)]; + const auto swaMaskBuffer = swaMaskBufferInfo.data; + const auto swaMaskSizeBytes = swaMaskBufferInfo.nbytesUsed; + mMaskBuilder->setMaskBuffer(swaMaskBuffer, swaMaskSizeBytes); + mMaskBuilder->enableSlidingWindow(kWindowSize); + mMaskBuilder->buildMask(mTokenBatchSize, mCurrentTokenIndex); + } + // Global Mask + const auto& maskBufferInfo = mInputBufferInfos[getInputIndex(IOKind::Mask)]; const auto maskBuffer = maskBufferInfo.data; const auto maskSizeBytes = maskBufferInfo.nbytesUsed; - mMaskBuilder = std::make_unique( - maskBuffer, maskSizeBytes, kMaskType, kCacheLength); + mMaskBuilder->setMaskBuffer(maskBuffer, maskSizeBytes); + mMaskBuilder->disableSlidingWindow(); mMaskBuilder->buildMask(mTokenBatchSize, mCurrentTokenIndex); + + mMaskBuilder->resetPadLength(); } void LlamaModelChunk::InitCache() { // Zero initialization - for (const auto cacheIdx : kCacheInputIndexes) { + for (const auto cacheIdx : getInputIndexes(IOKind::KVCache)) { const auto& inputCacheInfo = mInputBufferInfos[cacheIdx]; char* cacheBuffer = reinterpret_cast(inputCacheInfo.data); const size_t cacheSizeBytes = inputCacheInfo.nbytes; diff --git a/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.h b/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.h index 0a5002199db..d8cf08e45af 100644 --- a/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.h +++ b/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.h @@ -37,6 +37,19 @@ using ModelIndexMap = std::unordered_map; // Llama decoder chunk class LlamaModelChunk : public ModelChunk { + public: + enum class IOKind { + Embedding, + Mask, + SWAMask, + RotEmb, + KVCache, + Logits, + + // Cache Eviction + Attention + }; + private: static constexpr size_t kCacheLengthDim = 2; @@ -44,9 +57,11 @@ class LlamaModelChunk : public ModelChunk { explicit LlamaModelChunk( const ModelPathMap& modelPathMap, const LlamaModelOptions& modelOptions, + const bool useSharedWeights, const size_t initBatchSize, const size_t numCache, const size_t numRotEmbInputs, + const bool enableSWA, const RotaryEmbeddingMasterLut* rotEmbMasterLut); ~LlamaModelChunk(); @@ -97,18 +112,36 @@ class LlamaModelChunk : public ModelChunk { const size_t numSeenToken); private: - void CheckIoCount(); + // IO Setup and Query + virtual void defineIOs(); + + void defineInput(const IOKind kind, const size_t count = 1); + void defineOutput(const IOKind kind, const size_t count = 1); + + bool hasInput(const IOKind kind) const; + bool hasOutput(const IOKind kind) const; + + const std::vector& getInputIndexes(const IOKind kind) const; + const std::vector& getOutputIndexes(const IOKind kind) const; - size_t GetExpectedInputCount() const; + size_t getInputIndex(const IOKind kind, const size_t pos = 0) const; + size_t getOutputIndex(const IOKind kind, const size_t pos = 0) const; - size_t GetExpectedOutputCount() const; + size_t getNumInputsFor(const IOKind kind) const; + size_t getNumOutputsFor(const IOKind kind) const; + + void CheckIoCount(); private: - // Input/Output Indexes - const size_t kMaskInputIndex; - const std::vector kRotEmbInputIndexes; - const std::vector kCacheInputIndexes; - const std::vector kCacheOutputIndexes; + bool AllowModelsCoexist() const override { + return kIsSharedWeightsUsed; + } + + std::string SelectMethod( + const std::vector& methodNames) const override; + + // Whether shared weights is used + bool kIsSharedWeightsUsed = false; // Cache TensorShape mCacheShape; @@ -130,9 +163,22 @@ class LlamaModelChunk : public ModelChunk { // Mask builder std::unique_ptr mMaskBuilder; + const size_t kWindowSize; + const bool enableSWA; // Keep track of token index. Its value can also be viewed as numSeenToken. size_t mCurrentTokenIndex = 0; + + protected: + const size_t kRotEmbInputCount; + const size_t kCacheCount; + const size_t kSWAMaskCount = 0; + + // IO Indexes + size_t mExpectedNumInputs = 0; + size_t mExpectedNumOutputs = 0; + std::unordered_map> mInputIndexes; + std::unordered_map> mOutputIndexes; }; } // namespace example diff --git a/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp b/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp index 2254241a001..e6e07f27ffd 100644 --- a/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp +++ b/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp @@ -24,12 +24,12 @@ void LlamaRuntime::Initialize( const LlamaModelOptions& modelOptions, const LlamaModelPaths& modelPaths) { mModelOptions = modelOptions; - const size_t numChunk = modelPaths.gen_model_paths.size(); - const size_t numCache = 2 * modelOptions.num_layer / numChunk; - ET_CHECK_MSG(numChunk > 0, "No model to initialize"); // Initialize rotary embedding master lookup table - const size_t rotEmbDim = modelOptions.hidden_size / modelOptions.num_head; + const size_t headDim = modelOptions.head_dim + ? modelOptions.head_dim + : (modelOptions.hidden_size / modelOptions.num_head); + const size_t rotEmbDim = headDim * modelOptions.partial_rotary_factor; mRotEmbMasterLut = std::make_unique( modelOptions.rot_emb_type, modelOptions.max_token_length, @@ -37,12 +37,39 @@ void LlamaRuntime::Initialize( modelOptions.rot_emb_base); mRotEmbMasterLut->generate(); + const bool useSharedWeights = !modelPaths.model_package_paths.empty(); + + ET_CHECK_MSG( + !useSharedWeights || + modelPaths.prompt_model_paths.empty() && + modelPaths.gen_model_paths.empty(), + "The paths for both prompt and gen model paths should be empty when shared weights is used."); + + const size_t numChunk = useSharedWeights + ? modelPaths.model_package_paths.size() + : modelPaths.gen_model_paths.size(); + ET_CHECK_MSG(numChunk > 0, "No model to initialize"); + const size_t numCache = 2 * modelOptions.num_layer / numChunk; + constexpr size_t numRotEmbInputs = 1; - const bool usePromptModel = !modelPaths.prompt_model_paths.empty(); + const bool usePromptModel = !modelPaths.prompt_model_paths.empty() || + !modelPaths.model_package_paths.empty(); const size_t initBatchSize = usePromptModel ? modelOptions.prompt_token_batch_size : 1; mTokenBatchSize = initBatchSize; + // Enable SWA if window size is not 0 + const bool enableSWA = (modelOptions.window_size != 0); + + // Get effective prompt and gen model paths + const auto& [prompt_model_paths, gen_model_paths] = [&] { + if (useSharedWeights) { + return std::pair{ + modelPaths.model_package_paths, modelPaths.model_package_paths}; + } + return std::pair{modelPaths.prompt_model_paths, modelPaths.gen_model_paths}; + }(); + for (size_t chunkIdx = 0; chunkIdx < numChunk; chunkIdx++) { ModelPathMap modelPathMap; auto addModelPath = [&](const auto& modelPaths, const size_t batchSize) { @@ -50,15 +77,16 @@ void LlamaRuntime::Initialize( return; modelPathMap[batchSize] = modelPaths[chunkIdx]; }; - addModelPath( - modelPaths.prompt_model_paths, modelOptions.prompt_token_batch_size); - addModelPath(modelPaths.gen_model_paths, 1); + addModelPath(prompt_model_paths, modelOptions.prompt_token_batch_size); + addModelPath(gen_model_paths, 1); auto llamaChunk = std::make_unique( modelPathMap, modelOptions, + useSharedWeights, initBatchSize, numCache, numRotEmbInputs, + enableSWA, mRotEmbMasterLut.get()); mLlamaModelChunks.push_back(std::move(llamaChunk)); } diff --git a/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp b/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp index 7f67fb4ca79..0b4d38ee851 100644 --- a/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp +++ b/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp @@ -41,20 +41,171 @@ using executorch::runtime::Tag; static constexpr size_t kMethodAllocatorPoolSize = 4 * 1024U * 1024U; // 4MB -// ExecuTorch model instance -// The member ordering affects the order of destruction. -struct ModelInstance { - std::unique_ptr program; +// ExecuTorch model instance with cacheable program. +class ModelInstance { + public: + explicit ModelInstance(const std::string& modelPath) { + if (mCachedPrograms.find(modelPath) != mCachedPrograms.end()) { + auto cachedProgram = mCachedPrograms.at(modelPath).lock(); + if (cachedProgram) { + mProgramInstance = cachedProgram; + ET_LOG( + Debug, "Loaded existing program from cache: %s", modelPath.c_str()); + return; + } else { + mCachedPrograms.erase(modelPath); // Expired + } + } + ET_LOG(Debug, "Loading model from scratch: %s", modelPath.c_str()); + mProgramInstance = std::make_shared(); + + // Create a loader to get the data of the program file. There are other + // 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. + Result loader = FileDataLoader::from(modelPath.c_str()); + ET_CHECK_MSG( + loader.ok(), + "FileDataLoader::from() failed: 0x%" PRIx32, + loader.error()); + // Extract the data loader out to a persistent storage before loading the + // program. + mProgramInstance->dataLoader = + std::make_unique(std::move(loader.get())); + + // Parse the program file. This is immutable, and can also be reused between + // multiple execution invocations across multiple threads. + Result program_loaded = + Program::load(mProgramInstance->dataLoader.get()); + ET_CHECK_MSG( + program_loaded.ok(), + "Failed to parse model file %s", + modelPath.c_str()); + ET_LOG(Debug, "Model file %s is loaded.", modelPath.c_str()); + + // Extract program out to a persistent storage before calling any of its + // methods. + mProgramInstance->program = + std::make_unique(std::move(program_loaded.get())); + mCachedPrograms.emplace(modelPath, mProgramInstance); + } + + Method& GetMethod() { + ET_CHECK_MSG(mMethod != nullptr, "Method is not loaded."); + return *mMethod; + } + + const Method& GetMethod() const { + ET_CHECK_MSG(mMethod != nullptr, "Method is not loaded."); + return *mMethod; + } + + Program& GetProgram() { + return *(mProgramInstance->program); + } + + const Program& GetProgram() const { + return *(mProgramInstance->program); + } - std::vector> planned_buffers; - std::vector> planned_spans; + std::vector GetMethodNames() const { + std::vector methodNames; + for (size_t i = 0; i < GetProgram().num_methods(); i++) { + const auto method_name_result = GetProgram().get_method_name(i); + ET_CHECK_MSG(method_name_result.ok(), "Program has no method %zu", i); + methodNames.emplace_back(*method_name_result); + } + return methodNames; + } - std::vector method_allocator_pool; - std::unique_ptr method_allocator; - std::unique_ptr planned_memory; - std::unique_ptr memory_manager; + void LoadFirstMethod() { + // Use the first method in the program. + const char* method_name = nullptr; + const auto method_name_result = GetProgram().get_method_name(0); + ET_CHECK_MSG(method_name_result.ok(), "Program has no methods"); + method_name = *method_name_result; + ET_LOG(Debug, "Loading the first method."); + LoadMethod(method_name); + } - std::unique_ptr method; + void LoadMethod(const std::string& method_name) { + ET_CHECK_MSG(!mMethod, "Method is already loaded."); + const auto method_name_cstr = method_name.c_str(); + + // MethodMeta describes the memory requirements of the method. + Result method_meta = GetProgram().method_meta(method_name_cstr); + ET_CHECK_MSG( + method_meta.ok(), + "Failed to get method_meta for %s: 0x%x", + method_name_cstr, + (unsigned int)method_meta.error()); + + mMethodAllocatorPool.resize(kMethodAllocatorPoolSize); + mMethodAllocator = std::make_unique( + kMethodAllocatorPoolSize, mMethodAllocatorPool.data()); + mMethodAllocator->enable_profiling("method allocator"); + + size_t num_memory_planned_buffers = + method_meta->num_memory_planned_buffers(); + for (size_t id = 0; id < num_memory_planned_buffers; ++id) { + // .get() will always succeed because id < num_memory_planned_buffers. + size_t buffer_size = static_cast( + method_meta->memory_planned_buffer_size(id).get()); + ET_LOG( + Debug, "Setting up planned buffer %zu, size %zu.", id, buffer_size); + mPlannedBuffers.push_back(std::make_unique(buffer_size)); + mPlannedSpans.push_back({mPlannedBuffers.back().get(), buffer_size}); + } + mPlannedMemory = std::make_unique( + Span>{mPlannedSpans.data(), mPlannedSpans.size()}); + + // Assemble all of the allocators into the MemoryManager that the Executor + // will use. + auto& neuron_allocator = GET_NEURON_ALLOCATOR; + mMemoryManager = std::make_unique( + mMethodAllocator.get(), + mPlannedMemory.get(), + dynamic_cast(&neuron_allocator)); + + ET_LOG(Debug, "Loading method %s", method_name_cstr); + Result method = + GetProgram().load_method(method_name_cstr, mMemoryManager.get()); + ET_CHECK_MSG( + method.ok(), + "Loading of method %s failed with status 0x%" PRIx32, + method_name_cstr, + method.error()); + + mMethod = std::make_unique(std::move(method.get())); + } + + private: + ModelInstance(const ModelInstance&) = delete; + ModelInstance& operator=(ModelInstance&&) = delete; + ModelInstance& operator=(const ModelInstance&) = delete; + + private: + // The member ordering below affects the order of destruction. + + struct ProgramInstance { + std::unique_ptr dataLoader; + std::unique_ptr program; + }; + std::shared_ptr mProgramInstance; + + std::vector> mPlannedBuffers; + std::vector> mPlannedSpans; + + std::vector mMethodAllocatorPool; + std::unique_ptr mMethodAllocator; + std::unique_ptr mPlannedMemory; + std::unique_ptr mMemoryManager; + + std::unique_ptr mMethod; + + // Maps .pte file paths to the cached program instances. + inline static std::unordered_map> + mCachedPrograms; }; void ModelChunk::Initialize() { @@ -487,94 +638,18 @@ void ModelChunk::ReleaseIoBuffers() { Method& ModelChunk::GetModelMethod() { auto modelInstance = reinterpret_cast(GetModelInstance()); - return *(modelInstance->method); + return modelInstance->GetMethod(); } // Override the virtual functions void* ModelChunk::CreateModelInstance(const std::string& modelPath) { - auto modelInstance = new ModelInstance; - - // Create a loader to get the data of the program file. There are other - // 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. - Result loader = FileDataLoader::from(modelPath.c_str()); - ET_CHECK_MSG( - loader.ok(), "FileDataLoader::from() failed: 0x%" PRIx32, loader.error()); - - // Parse the program file. This is immutable, and can also be reused between - // multiple execution invocations across multiple threads. - Result program_loaded = Program::load(&loader.get()); - if (!program_loaded.ok()) { - ET_LOG(Error, "Failed to parse model file %s", modelPath.c_str()); - return nullptr; - } - ET_LOG(Debug, "Model file %s is loaded.", modelPath.c_str()); - - // Extract program out to a persistent storage before calling any of its - // methods. - modelInstance->program = - std::make_unique(std::move(program_loaded.get())); - auto& program = modelInstance->program; - - // Use the first method in the program. - const char* method_name = nullptr; - { - const auto method_name_result = program->get_method_name(0); - ET_CHECK_MSG(method_name_result.ok(), "Program has no methods"); - method_name = *method_name_result; - } - ET_LOG(Debug, "Using method %s", method_name); - - // MethodMeta describes the memory requirements of the method. - Result method_meta = program->method_meta(method_name); - ET_CHECK_MSG( - method_meta.ok(), - "Failed to get method_meta for %s: 0x%x", - method_name, - (unsigned int)method_meta.error()); - - modelInstance->method_allocator_pool.resize(kMethodAllocatorPoolSize); - modelInstance->method_allocator = std::make_unique( - kMethodAllocatorPoolSize, modelInstance->method_allocator_pool.data()); - auto& method_allocator = modelInstance->method_allocator; - method_allocator->enable_profiling("method allocator"); - - auto& planned_buffers = modelInstance->planned_buffers; // Owns the memory - auto& planned_spans = modelInstance->planned_spans; // Passed to the allocator - - size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers(); - for (size_t id = 0; id < num_memory_planned_buffers; ++id) { - // .get() will always succeed because id < num_memory_planned_buffers. - size_t buffer_size = - static_cast(method_meta->memory_planned_buffer_size(id).get()); - ET_LOG(Debug, "Setting up planned buffer %zu, size %zu.", id, buffer_size); - planned_buffers.push_back(std::make_unique(buffer_size)); - planned_spans.push_back({planned_buffers.back().get(), buffer_size}); + auto modelInstance = new ModelInstance(modelPath); + const auto selectedMethod = SelectMethod(modelInstance->GetMethodNames()); + if (!selectedMethod.empty()) { + modelInstance->LoadMethod(selectedMethod); + } else { + modelInstance->LoadFirstMethod(); // Load the first available method } - modelInstance->planned_memory = std::make_unique( - Span>{planned_spans.data(), planned_spans.size()}); - auto& planned_memory = modelInstance->planned_memory; - - // Assemble all of the allocators into the MemoryManager that the Executor - // will use. - auto& neuron_allocator = GET_NEURON_ALLOCATOR; - modelInstance->memory_manager = std::make_unique( - method_allocator.get(), - planned_memory.get(), - dynamic_cast(&neuron_allocator)); - auto& memory_manager = modelInstance->memory_manager; - - ET_LOG(Debug, "Begin loading method %s", method_name); - Result method = - program->load_method(method_name, memory_manager.get()); - ET_CHECK_MSG( - method.ok(), - "Loading of method %s failed with status 0x%" PRIx32, - method_name, - method.error()); - ET_LOG(Debug, "Method loaded."); - - modelInstance->method = std::make_unique(std::move(method.get())); return modelInstance; } diff --git a/examples/mediatek/executor_runner/llama_runner/ModelChunk.h b/examples/mediatek/executor_runner/llama_runner/ModelChunk.h index 67d9e30b5f1..66a1f4586bb 100644 --- a/examples/mediatek/executor_runner/llama_runner/ModelChunk.h +++ b/examples/mediatek/executor_runner/llama_runner/ModelChunk.h @@ -94,6 +94,11 @@ class ModelChunk : protected MultiTokenSizeModelLoader { executorch::runtime::Method& GetModelMethod(); private: + virtual std::string SelectMethod( + const std::vector& methodNames) const { + return {}; // Default choose the first available one. + } + // Override the virtual functions void* CreateModelInstance(const std::string& modelPath) override; diff --git a/examples/mediatek/executor_runner/llama_runner/Utils.h b/examples/mediatek/executor_runner/llama_runner/Utils.h index 24e8a4d6e50..9afd036a21e 100644 --- a/examples/mediatek/executor_runner/llama_runner/Utils.h +++ b/examples/mediatek/executor_runner/llama_runner/Utils.h @@ -59,6 +59,22 @@ static std::vector split(const std::string& str, const char sep) { return tokens; } +static std::vector extract_substr( + const std::string& str, + const std::string& pattern) { + std::vector tokens; + const std::regex token_pattern(pattern); + std::smatch matches; + auto cur = str.cbegin(); + while (std::regex_search(cur, str.cend(), matches, token_pattern)) { + for (const auto& match : matches) { + tokens.push_back(match.str()); + } + cur = matches.suffix().first; + } + return tokens; +} + static std::string read_file(const std::string& filepath) { std::ifstream file(filepath); std::stringstream buffer; diff --git a/examples/mediatek/executor_runner/llama_runner/llm_helper/include/mask_builder.h b/examples/mediatek/executor_runner/llama_runner/llm_helper/include/mask_builder.h index 5ab6741c11c..0839443aa56 100644 --- a/examples/mediatek/executor_runner/llama_runner/llm_helper/include/mask_builder.h +++ b/examples/mediatek/executor_runner/llama_runner/llm_helper/include/mask_builder.h @@ -17,6 +17,10 @@ namespace llm_helper { class MaskBuilder { public: + // Construct mask builder without specifying the mask buffer. + explicit MaskBuilder(const LLMType maskType, const size_t cacheLength); + + // Construct mask builder and initialize mask buffer. explicit MaskBuilder( void* maskBuffer, const size_t maskSizeBytes, @@ -25,6 +29,13 @@ class MaskBuilder { ~MaskBuilder(); + // Specify the mask buffer to build/update mask. + MaskBuilder& setMaskBuffer(void* maskBuffer, const size_t maskSizeBytes); + + MaskBuilder& enableSlidingWindow(const size_t windowSize); + + MaskBuilder& disableSlidingWindow(); + // Build mask from scratch. void buildMask(const size_t tokenBatchSize, const size_t numSeenToken); @@ -46,6 +57,12 @@ class MaskBuilder { // alignment. void updateMaskSize(const size_t sizeBytes); + void resetPadLength(); + + bool getMaskUpdateStatus(); + + void setIsMaskUpdatable(const bool status); + private: template void buildMask(const size_t tokenBatchSize, const size_t numSeenToken); @@ -73,6 +90,10 @@ class MaskBuilder { size_t mRightPadLength = 0; bool mIsMaskUpdatable = false; + + // Sliding Window Attention (SWA) size. + // Set to non-zero to enable SWA, or set to zero to disable SWA. + size_t mSlidingWindowSize = 0; }; } // namespace llm_helper diff --git a/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp b/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp index e83e8b37082..905e00a3c16 100644 --- a/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp +++ b/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp @@ -31,6 +31,13 @@ __DECL_MASK__(__fp16, 0, -100) __DECL_MASK__(float, 0, -100) #undef __DECL_MASK__ +MaskBuilder::MaskBuilder(const LLMType maskType, const size_t cacheLength) + : mMaskBuffer(nullptr), + mMaskSizeBytes(0), + kMaskType(maskType), + kMaskTypeSize(getLLMTypeSize(maskType)), + kCacheLength(cacheLength) {} + MaskBuilder::MaskBuilder( void* maskBuffer, const size_t maskSizeBytes, @@ -44,6 +51,24 @@ MaskBuilder::MaskBuilder( MaskBuilder::~MaskBuilder() {} +MaskBuilder& MaskBuilder::setMaskBuffer( + void* maskBuffer, + const size_t maskSizeBytes) { + mMaskBuffer = maskBuffer; + mMaskSizeBytes = maskSizeBytes; + return *this; +} + +MaskBuilder& MaskBuilder::enableSlidingWindow(const size_t windowSize) { + mSlidingWindowSize = windowSize; + return *this; +} + +MaskBuilder& MaskBuilder::disableSlidingWindow() { + mSlidingWindowSize = 0; + return *this; +} + void MaskBuilder::updateMaskSize(const size_t sizeBytes) { mMaskSizeBytes = sizeBytes; } @@ -59,10 +84,11 @@ void MaskBuilder::buildMask( constexpr auto maskTrue = MaskVal::kTrue; constexpr auto maskFalse = MaskVal::kFalse; const size_t maskLength = kCacheLength + tokenBatchSize; + const size_t curWindowSize = + mSlidingWindowSize ? mSlidingWindowSize : maskLength; // The mask is a combination (concat) of input cache mask and attention mask - const size_t startTrueIdx = - kCacheLength - std::min(kCacheLength, numSeenToken); + const size_t numVisibleCacheTokens = std::min(kCacheLength, numSeenToken); const size_t rowSize = mMaskSizeBytes / tokenBatchSize / kMaskTypeSize; @@ -82,22 +108,19 @@ void MaskBuilder::buildMask( // There are tokenBatchSize number of rows for (size_t inTokIdx = 0; inTokIdx < tokenBatchSize; inTokIdx++) { const auto& rowIdx = inTokIdx; // For clarity + const size_t attnTrueCount = inTokIdx + 1; auto curMaskBuffer = reinterpret_cast(mMaskBuffer) + rowIdx * rowSize; - size_t i = 0; // Buffer write index - // Set the (rectangle) input cache mask - while (i < startTrueIdx) + const size_t rowTrueCount = + std::min(curWindowSize, numVisibleCacheTokens + attnTrueCount); + const size_t lastTrueIdx = kCacheLength + attnTrueCount - 1; + const size_t firstTrueIdx = lastTrueIdx - rowTrueCount + 1; + size_t i = 0; + while (i < firstTrueIdx) curMaskBuffer[i++] = maskFalse; - while (i < kCacheLength) + while (i <= lastTrueIdx) curMaskBuffer[i++] = maskTrue; - - // Set the (triangle) attention mask - const size_t attnTrueCount = inTokIdx + 1; - for (size_t counter = 0; counter < attnTrueCount; counter++) { - curMaskBuffer[i++] = maskTrue; - } - // Fill the remaining with False while (i < maskLength) curMaskBuffer[i++] = maskFalse; } @@ -246,7 +269,6 @@ bool MaskBuilder::adjustMaskForPadding(const size_t tokenBatchSize) { const size_t maskPadCount = std::min(mLeftPadLength, inTokIdx + 1); std::fill(curMaskBuffer, curMaskBuffer + maskPadCount, maskFalse); } - mLeftPadLength = 0; // Reset pad length } else if (mRightPadLength > 0) { // Mask the padded rows const auto startIdx = tokenBatchSize - mRightPadLength; @@ -254,10 +276,25 @@ bool MaskBuilder::adjustMaskForPadding(const size_t tokenBatchSize) { auto curMaskBuffer = maskBuffer + inTokIdx * rowSize; std::fill(curMaskBuffer, curMaskBuffer + maskLength, maskFalse); } - mRightPadLength = 0; // Reset pad length } return true; // Mask is modified for padding } +void MaskBuilder::resetPadLength() { + if (mLeftPadLength > 0) { + mLeftPadLength = 0; + } else if (mRightPadLength > 0) { + mRightPadLength = 0; + } +} + +bool MaskBuilder::getMaskUpdateStatus() { + return mIsMaskUpdatable; +} + +void MaskBuilder::setIsMaskUpdatable(const bool status) { + mIsMaskUpdatable = status; +} + } // namespace llm_helper } // namespace example diff --git a/examples/mediatek/executor_runner/mtk_llama_executor_runner.cpp b/examples/mediatek/executor_runner/mtk_llama_executor_runner.cpp index 733cc8c3465..dd189165f32 100644 --- a/examples/mediatek/executor_runner/mtk_llama_executor_runner.cpp +++ b/examples/mediatek/executor_runner/mtk_llama_executor_runner.cpp @@ -30,6 +30,7 @@ #include "llama_runner/llm_helper/include/llm_types.h" #include +#include #include #include @@ -42,10 +43,13 @@ DEFINE_uint64(cache_size, 1024, "Model cache size."); DEFINE_uint64(hidden_size, 4096, "Model hidden size."); DEFINE_uint64(num_head, 32, "Number of attention heads in each layer."); DEFINE_uint64(num_layer, 32, "Number of layers in the model."); +DEFINE_uint64(head_dim, 0, "Head dimension of the model."); +DEFINE_uint64(window_size, 0, "Window size of Sliding Window Attention."); DEFINE_uint64( max_token_length, 2048, "Maximum token length that the model supports."); +DEFINE_double(partial_rotary_factor, 1, "Partial rotary factor of the model."); DEFINE_double( rot_emb_base, 10000, @@ -66,14 +70,12 @@ DEFINE_string( token_embedding_path, "embedding.bin", "Input token embedding lookup table path."); +DEFINE_string(prompt_model_paths, "", "Comma-separated prompt model paths."); +DEFINE_string(gen_model_paths, "", "Comma-separated generative model paths."); DEFINE_string( - prompt_model_paths, - "model_128t.pte", - "Comma-separated prompt model paths."); -DEFINE_string( - gen_model_paths, - "model_1t.pte", - "Comma-separated generative model paths."); + model_package_paths, + "", + "Comma-separated weight-shared model package paths."); // Tokenizer DEFINE_string(tokenizer_path, "tokenizer.model", "tokenizer.model vocab path."); @@ -104,6 +106,7 @@ using example::utils::Timer; using example::utils::to_string; using executorch::runtime::Error; using executorch::runtime::Result; +using tokenizers::HFTokenizer; using tokenizers::Llama2cTokenizer; using tokenizers::Tokenizer; @@ -115,7 +118,10 @@ LlamaModelOptions get_model_options() { .hidden_size = FLAGS_hidden_size, .num_head = FLAGS_num_head, .num_layer = FLAGS_num_layer, + .head_dim = FLAGS_head_dim, + .window_size = FLAGS_window_size, .max_token_length = FLAGS_max_token_length, + .partial_rotary_factor = FLAGS_partial_rotary_factor, .rot_emb_base = FLAGS_rot_emb_base, // Types @@ -132,7 +138,9 @@ LlamaModelPaths get_model_paths() { .tokenizer_path = FLAGS_tokenizer_path, .token_embedding_path = FLAGS_token_embedding_path, .prompt_model_paths = split(FLAGS_prompt_model_paths, ','), - .gen_model_paths = split(FLAGS_gen_model_paths, ',')}; + .gen_model_paths = split(FLAGS_gen_model_paths, ','), + .model_package_paths = split(FLAGS_model_package_paths, ','), + }; return model_paths; } @@ -271,6 +279,8 @@ Error inference( const auto input_tokens = std::move(encode_res.get()); std::cout << "\n[Input Prompt]\n" << prompt << std::endl; + std::cout << "\n[Input Prompt Tokens]\n" + << to_string(input_tokens) << std::endl; // Run prompt mode (pre-fill) auto prefill_res = digest_prompt(llama_runtime, tokenizer, input_tokens); @@ -288,6 +298,8 @@ std::unique_ptr load_tokenizer() { tokenizer = std::make_unique(); } else if (FLAGS_tokenizer_type == "tiktoken") { tokenizer = example::get_tiktoken_for_llama(); + } else if (FLAGS_tokenizer_type == "hf") { + tokenizer = std::make_unique(); } ET_CHECK_MSG( tokenizer, "Invalid tokenizer type: %s", FLAGS_tokenizer_type.c_str()); @@ -311,7 +323,8 @@ int main(int argc, char** argv) { LlamaModelOptions model_options = get_model_options(); LlamaModelPaths model_paths = get_model_paths(); - if (model_paths.prompt_model_paths.empty()) { + if (model_paths.prompt_model_paths.empty() && + model_paths.model_package_paths.empty()) { model_options.prompt_token_batch_size = 1; ET_LOG( Info, diff --git a/examples/mediatek/executor_runner/run_gemma2_sample.sh b/examples/mediatek/executor_runner/run_gemma2_sample.sh new file mode 100644 index 00000000000..d6c565f7b88 --- /dev/null +++ b/examples/mediatek/executor_runner/run_gemma2_sample.sh @@ -0,0 +1,71 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=2304 +NUM_HEAD=8 +HEAD_DIM=256 +NUM_LAYER=26 +MAX_TOKEN_LENGTH=8192 +ROT_EMB_BASE=10000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=256000 +BOS_TOKEN=2 +EOS_TOKEN=1 +TOKENIZER_TYPE=bpe # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3, "hf" for huggingface tokenizer + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer_gemma2.bin" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_gemma2_2b_it_fp32.bin" + +# Comma-Separated Paths +WEIGHT_SHARED_MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/gemma2_2b_it_A16W4_1_chunks/gemma2_2b_it_A16W4_1_chunks_0.pte," + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt_gemma.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --head_dim=$HEAD_DIM \ + --num_head=$NUM_HEAD \ + --num_layer=$NUM_LAYER \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$WEIGHT_SHARED_MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_gemma3_sample.sh b/examples/mediatek/executor_runner/run_gemma3_sample.sh new file mode 100644 index 00000000000..761746ff7ee --- /dev/null +++ b/examples/mediatek/executor_runner/run_gemma3_sample.sh @@ -0,0 +1,74 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=1152 +NUM_HEAD=4 +NUM_LAYER=26 +HEAD_DIM=256 +WINDOW_SIZE=512 +MAX_TOKEN_LENGTH=32768 +ROT_EMB_BASE=1000000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=262144 +BOS_TOKEN=2 +EOS_TOKEN=1 +TOKENIZER_TYPE=bpe # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3, "hf" for huggingface tokenizer + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer_gemma3.bin" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_gemma3_1b_it_fp32.bin" + +# Comma-Separated Paths +MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/gemma3_1b_it_A16W8_1_chunks/gemma3_1b_it_A16W8_1_chunks_0.pte," + + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt_gemma.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --num_head=$NUM_HEAD \ + --num_layer=$NUM_LAYER \ + --head_dim=$HEAD_DIM \ + --window_size=$WINDOW_SIZE \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_llama3_sample.sh b/examples/mediatek/executor_runner/run_llama3_sample.sh index d9d6ea43e3c..eeb1089d792 100644 --- a/examples/mediatek/executor_runner/run_llama3_sample.sh +++ b/examples/mediatek/executor_runner/run_llama3_sample.sh @@ -36,18 +36,12 @@ TOKENIZER_PATH="/data/local/tmp/llama3/tokenizer.model" TOKEN_EMBEDDING_PATH="/data/local/tmp/llama3/embedding_llama3_8b_instruct_fp32.bin" # Comma-Separated Paths -PROMPT_MODEL_PATHS="\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_128t512c_0.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_128t512c_1.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_128t512c_2.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_128t512c_3.pte," - -# Comma-Separated Paths -GEN_MODEL_PATHS="\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_1t512c_0.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_1t512c_1.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_1t512c_2.pte,\ -/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_1t512c_3.pte," +WEIGHT_SHARED_MODEL_PACKAGE_PATHS="\ +/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_0.pte,\ +/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_1.pte,\ +/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_2.pte,\ +/data/local/tmp/llama3/llama3_8b_SC_sym4W_sym16A_4_chunks_Overall_3.pte,\ +" PROMPT_FILE=/data/local/tmp/llama3/sample_prompt.txt @@ -75,6 +69,5 @@ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD --tokenizer_type=$TOKENIZER_TYPE \ --tokenizer_path=$TOKENIZER_PATH \ --token_embedding_path=$TOKEN_EMBEDDING_PATH \ - --prompt_model_paths=$PROMPT_MODEL_PATHS \ - --gen_model_paths=$GEN_MODEL_PATHS \ + --model_package_paths=$WEIGHT_SHARED_MODEL_PACKAGE_PATHS \ --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_phi3_sample.sh b/examples/mediatek/executor_runner/run_phi3_sample.sh new file mode 100644 index 00000000000..73edb99614a --- /dev/null +++ b/examples/mediatek/executor_runner/run_phi3_sample.sh @@ -0,0 +1,72 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=3072 +NUM_HEAD=32 +NUM_LAYER=32 +MAX_TOKEN_LENGTH=4096 +ROT_EMB_BASE=10000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=32000 +BOS_TOKEN=1 +EOS_TOKEN=32000 +TOKENIZER_TYPE=bpe # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3 + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer.bin" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_phi3.5-mini-instruct_fp32.bin" + +# Comma-Separated Paths +WEIGHT_SHARED_MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/phi3.5-mini-instruct_A16W4_4_chunks/phi3.5-mini-instruct_A16W4_4_chunks_0.pte,\ +/data/local/tmp/et_mtk/phi3.5-mini-instruct_A16W4_4_chunks/phi3.5-mini-instruct_A16W4_4_chunks_1.pte,\ +/data/local/tmp/et_mtk/phi3.5-mini-instruct_A16W4_4_chunks/phi3.5-mini-instruct_A16W4_4_chunks_2.pte,\ +/data/local/tmp/et_mtk/phi3.5-mini-instruct_A16W4_4_chunks/phi3.5-mini-instruct_A16W4_4_chunks_3.pte," + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt_phi3.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --num_head=$NUM_HEAD \ + --num_layer=$NUM_LAYER \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$WEIGHT_SHARED_MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_phi4_sample.sh b/examples/mediatek/executor_runner/run_phi4_sample.sh new file mode 100644 index 00000000000..a6d9824e178 --- /dev/null +++ b/examples/mediatek/executor_runner/run_phi4_sample.sh @@ -0,0 +1,74 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=3072 +NUM_HEAD=24 +NUM_LAYER=32 +PARTIAL_ROTARY_FACTOR=0.75 +MAX_TOKEN_LENGTH=4096 +ROT_EMB_BASE=10000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=200064 +BOS_TOKEN=199999 +EOS_TOKEN=199999 +TOKENIZER_TYPE=hf # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3, "hf" for huggingface tokenizer + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer_phi4.json" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_phi-4-mini-reasoning_fp32.bin" + +# Comma-Separated Paths +MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/phi-4-mini-reasoning_A16W8_4_chunks/phi-4-mini-reasoning_A16W8_4_chunks_0.pte,\ +/data/local/tmp/et_mtk/phi-4-mini-reasoning_A16W8_4_chunks/phi-4-mini-reasoning_A16W8_4_chunks_1.pte,\ +/data/local/tmp/et_mtk/phi-4-mini-reasoning_A16W8_4_chunks/phi-4-mini-reasoning_A16W8_4_chunks_2.pte,\ +/data/local/tmp/et_mtk/phi-4-mini-reasoning_A16W8_4_chunks/phi-4-mini-reasoning_A16W8_4_chunks_3.pte," + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt_phi4.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner_longrope \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --num_head=$NUM_HEAD \ + --partial_rotary_factor=$PARTIAL_ROTARY_FACTOR \ + --num_layer=$NUM_LAYER \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_qwen2_sample.sh b/examples/mediatek/executor_runner/run_qwen2_sample.sh new file mode 100644 index 00000000000..9d7c2daa311 --- /dev/null +++ b/examples/mediatek/executor_runner/run_qwen2_sample.sh @@ -0,0 +1,72 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=3584 +NUM_HEAD=28 +NUM_LAYER=28 +MAX_TOKEN_LENGTH=32768 +ROT_EMB_BASE=1000000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=152064 +BOS_TOKEN=151643 +EOS_TOKEN=151645 +TOKENIZER_TYPE=hf # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3, "hf" for huggingface tokenizer + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer_qwen3.json" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_Qwen2-7B-Instruct_fp32.bin" + +# Comma-Separated Paths +WEIGHT_SHARED_MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/Qwen2-7B-Instruct_A16W4_4_chunks/Qwen2-7B-Instruct_A16W4_4_chunks_0.pte,\ +/data/local/tmp/et_mtk/Qwen2-7B-Instruct_A16W4_4_chunks/Qwen2-7B-Instruct_A16W4_4_chunks_1.pte,\ +/data/local/tmp/et_mtk/Qwen2-7B-Instruct_A16W4_4_chunks/Qwen2-7B-Instruct_A16W4_4_chunks_2.pte,\ +/data/local/tmp/et_mtk/Qwen2-7B-Instruct_A16W4_4_chunks/Qwen2-7B-Instruct_A16W4_4_chunks_3.pte," + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --num_head=$NUM_HEAD \ + --num_layer=$NUM_LAYER \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$WEIGHT_SHARED_MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/executor_runner/run_qwen3_sample.sh b/examples/mediatek/executor_runner/run_qwen3_sample.sh new file mode 100644 index 00000000000..4ce4b2b37f7 --- /dev/null +++ b/examples/mediatek/executor_runner/run_qwen3_sample.sh @@ -0,0 +1,74 @@ + # Copyright (c) 2024 MediaTek Inc. + # + # Licensed under the BSD License (the "License"); you may not use this file + # except in compliance with the License. See the license file in the root + # directory of this source tree for more details. + +# Runtime +MAX_RESPONSE=200 + +# Model External +PROMPT_TOKEN_BATCH_SIZE=128 +CACHE_SIZE=512 + +# Model Internals +HIDDEN_SIZE=2560 +NUM_HEAD=32 +NUM_LAYER=36 +HEAD_DIM=128 +MAX_TOKEN_LENGTH=40960 +ROT_EMB_BASE=1000000 + +# Model IO Types +INPUT_TYPE=fp32 +OUTPUT_TYPE=fp32 +CACHE_TYPE=fp32 +MASK_TYPE=fp32 +ROT_EMB_TYPE=fp32 + +# Tokenizer +VOCAB_SIZE=151936 +BOS_TOKEN=151643 +EOS_TOKEN=151645 +TOKENIZER_TYPE=hf # Use "bpe" for LLAMA2, "tiktoken" for LLAMA3 + +# Paths +TOKENIZER_PATH="/data/local/tmp/et_mtk/tokenizer_qwen3.json" +TOKEN_EMBEDDING_PATH="/data/local/tmp/et_mtk/embedding_Qwen3-4B_fp32.bin" + +# Comma-Separated Paths +WEIGHT_SHARED_MODEL_PACKAGE_PATHS="\ +/data/local/tmp/et_mtk/Qwen3-4B_A16W4_4_chunks/Qwen3-4B_A16W4_4_chunks_0.pte,\ +/data/local/tmp/et_mtk/Qwen3-4B_A16W4_4_chunks/Qwen3-4B_A16W4_4_chunks_1.pte,\ +/data/local/tmp/et_mtk/Qwen3-4B_A16W4_4_chunks/Qwen3-4B_A16W4_4_chunks_2.pte,\ +/data/local/tmp/et_mtk/Qwen3-4B_A16W4_4_chunks/Qwen3-4B_A16W4_4_chunks_3.pte," + +PROMPT_FILE=/data/local/tmp/et_mtk/prompt_qwen3.txt + +chmod +x mtk_llama_executor_runner + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$PWD + +./mtk_llama_executor_runner \ + --max_response=$MAX_RESPONSE \ + --prompt_token_batch_size=$PROMPT_TOKEN_BATCH_SIZE \ + --cache_size=$CACHE_SIZE \ + --hidden_size=$HIDDEN_SIZE \ + --num_head=$NUM_HEAD \ + --num_layer=$NUM_LAYER \ + --head_dim=$HEAD_DIM \ + --max_token_length=$MAX_TOKEN_LENGTH \ + --rot_emb_base=$ROT_EMB_BASE \ + --input_type=$INPUT_TYPE \ + --output_type=$OUTPUT_TYPE \ + --cache_type=$CACHE_TYPE \ + --mask_type=$MASK_TYPE \ + --rot_emb_type=$ROT_EMB_TYPE \ + --vocab_size=$VOCAB_SIZE \ + --bos_token=$BOS_TOKEN \ + --eos_token=$EOS_TOKEN \ + --tokenizer_type=$TOKENIZER_TYPE \ + --tokenizer_path=$TOKENIZER_PATH \ + --token_embedding_path=$TOKEN_EMBEDDING_PATH \ + --model_package_paths=$WEIGHT_SHARED_MODEL_PACKAGE_PATHS \ + --prompt_file=$PROMPT_FILE \ No newline at end of file diff --git a/examples/mediatek/model_export_scripts/bert.py b/examples/mediatek/model_export_scripts/bert.py new file mode 100755 index 00000000000..420f4776833 --- /dev/null +++ b/examples/mediatek/model_export_scripts/bert.py @@ -0,0 +1,82 @@ +# Copyright (c) MediaTek 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 argparse +import os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) + +from aot_utils.oss_utils.utils import ( + build_executorch_binary, + get_masked_language_model_dataset, +) +from transformers import AutoModelForMaskedLM, AutoTokenizer + + +def main(args): + # ensure the working directory exist. + os.makedirs(args.artifact, exist_ok=True) + data_size = 100 + + tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased") + inputs, targets = get_masked_language_model_dataset( + args.dataset, tokenizer, data_size + ) + + # build pte + module = AutoModelForMaskedLM.from_pretrained( + "google-bert/bert-base-uncased" + ).eval() + pte_filename = "bert_mtk" + + build_executorch_binary( + module, + inputs[0], + f"{args.artifact}/{pte_filename}", + inputs, + skip_op_name={"aten_embedding_default", "aten_where_self"}, + ) + + # save data to inference on device + input_list_file = f"{args.artifact}/input_list.txt" + with open(input_list_file, "w") as f: + for i in range(len(inputs)): + f.write(f"input_{i}_0.bin input_{i}_1.bin\n") + for idx, data in enumerate(inputs): + for i, d in enumerate(data): + file_name = f"{args.artifact}/input_{idx}_{i}.bin" + d.detach().numpy().tofile(file_name) + for idx, data in enumerate(targets): + file_name = f"{args.artifact}/golden_{idx}_0.bin" + data.detach().numpy().tofile(file_name) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts and output by this example. Default ./bert", + default="./bert", + type=str, + ) + parser.add_argument( + "-d", + "--dataset", + help=( + "path to the validation text. " + "e.g. --dataset wikisent2.txt " + "for https://www.kaggle.com/datasets/mikeortman/wikipedia-sentences" + ), + default="wikisent2.txt", + type=str, + required=False, + ) + + args = parser.parse_args() + main(args) diff --git a/examples/mediatek/model_export_scripts/distilbert.py b/examples/mediatek/model_export_scripts/distilbert.py new file mode 100755 index 00000000000..81553d2376d --- /dev/null +++ b/examples/mediatek/model_export_scripts/distilbert.py @@ -0,0 +1,82 @@ +# Copyright (c) MediaTek 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 argparse +import os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) + +from aot_utils.oss_utils.utils import ( + build_executorch_binary, + get_masked_language_model_dataset, +) +from transformers import AutoModelForMaskedLM, AutoTokenizer + + +def main(args): + # ensure the working directory exist. + os.makedirs(args.artifact, exist_ok=True) + data_size = 100 + + tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased") + inputs, targets = get_masked_language_model_dataset( + args.dataset, tokenizer, data_size + ) + + # build pte + module = AutoModelForMaskedLM.from_pretrained( + "distilbert/distilbert-base-uncased" + ).eval() + pte_filename = "distilbert_mtk" + + build_executorch_binary( + module, + inputs[0], + f"{args.artifact}/{pte_filename}", + inputs, + skip_op_name={"aten_embedding_default", "aten_where_self"}, + ) + + # save data to inference on device + input_list_file = f"{args.artifact}/input_list.txt" + with open(input_list_file, "w") as f: + for i in range(len(inputs)): + f.write(f"input_{i}_0.bin input_{i}_1.bin\n") + for idx, data in enumerate(inputs): + for i, d in enumerate(data): + file_name = f"{args.artifact}/input_{idx}_{i}.bin" + d.detach().numpy().tofile(file_name) + for idx, data in enumerate(targets): + file_name = f"{args.artifact}/golden_{idx}_0.bin" + data.detach().numpy().tofile(file_name) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts and output by this example. Default ./distilbert", + default="./distilbert", + type=str, + ) + parser.add_argument( + "-d", + "--dataset", + help=( + "path to the validation text. " + "e.g. --dataset wikisent2.txt " + "for https://www.kaggle.com/datasets/mikeortman/wikipedia-sentences" + ), + default="wikisent2.txt", + type=str, + required=False, + ) + + args = parser.parse_args() + main(args) diff --git a/examples/mediatek/model_export_scripts/gemma.py b/examples/mediatek/model_export_scripts/gemma.py new file mode 100644 index 00000000000..f08420d4823 --- /dev/null +++ b/examples/mediatek/model_export_scripts/gemma.py @@ -0,0 +1,581 @@ +# Copyright (c) MediaTek 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 os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) +import argparse +import struct +import warnings + +import torch + +from aot_utils.llm_utils.preformatter import Preformatter +from aot_utils.llm_utils.sanity_checks import ( + check_all_chunks_same_num_layer, + check_between_inclusive, + check_exist, + check_ext, + check_old_arg, + check_shapes, + check_supported_model, + check_supported_tokenizer, + check_tokenizer_exist, + check_weights_exist, +) +from aot_utils.llm_utils.utils import ( + dump_embedding_lut_for_cmdline, + generate_mask, + get_dest_path, + get_dirname, + get_embedding_layer, + get_exp_name, + get_export_shapes, + get_master_rot_emb, + get_normalized_config, + load_checkpoints, + resolve_model_classes, +) +from datasets import load_dataset +from executorch import exir +from executorch.backends.mediatek import ( + NeuropilotPartitioner, + NeuropilotQuantizer, + Precision, +) +from executorch.exir.backend.backend_api import ( + MethodProgramsPartitionerSpec, + to_backend, +) +from executorch.exir.backend.backend_details import CompileSpec +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from tqdm import tqdm + +warnings.filterwarnings("ignore") + + +def get_argument_parser(): + parser = argparse.ArgumentParser( + description="Run Export to ET for suppoorted LLM models.", allow_abbrev=False + ) + parser.add_argument( + "config", + type=str, + help="[Required] Model config json file. " + "Model config must be in same directory as all model weight bins and tokenizer files.", + ) + parser.add_argument( + "-p", + "--precision", + type=str, + default="A16W8", + choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], + help="Precision to quantize entire model to.", + ) + parser.add_argument( + "--platform", + type=str, + default="DX4", + choices=["DX3", "DX4"], + help="Chip model of the inference device. " + "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", + ) + parser.add_argument( + "-d", + "--dataset", + type=str, + default=None, + help="Calibration dataset name or path to dataset. Defaults to None to use random inputs", + ) + parser.add_argument( + "-n", + "--num_chunks", + type=int, + default=4, + help="Number of chunks to cut the model into. Defaults to 4.", + ) + parser.add_argument( + "-r", + "--response_cap", + type=int, + default=9, + help="Max Number of Response Tokens to save during calibration. Defaults to 9.", + ) + parser.add_argument( + "--preformatter", + type=str, + default=None, + help="Preformatter Template to use to wrap input with. Defaults to None.", + ) + parser.add_argument( + "-shapes", + nargs="+", + help="[Required] Expected input shapes to reconfigure TFLites to. Space separated list of " + "shapes in the format: xtyc (e.g. 32t512c)", + ) + + return parser + + +# flake8: noqa: F405 +def args_sanity_checks(args): + check_old_arg(args.config) + check_exist(args.config, "Config file") + check_ext(args.config, ".json", "Config file") + config = get_normalized_config(args.config) + + weight_dir = get_dirname(args.config) + check_tokenizer_exist(weight_dir) + check_weights_exist(weight_dir) + + check_supported_model(config) + check_supported_tokenizer(config) + + if args.preformatter is not None: + check_exist(args.preformatter, "Preformatter json file") + check_ext(args.preformatter, ".json", "preformatter") + + if args.dataset is not None: + check_exist(args.dataset) + + check_between_inclusive(args.num_chunks, 1, config.num_hidden_layers, "num_chunks") + + check_shapes(args.shapes) + + +def print_args(args, exp_name): + print("Please check if all arguments are correct:") + print(f"Config file: {args.config}") + print(f"Output pte folder: pte/{exp_name}") + print(f"Quantization precision: {args.precision}") + print(f"Preformatter: {args.preformatter}") + print(f"Calibration Dataset: {args.dataset}") + print(f"Max Response Tokens: {args.response_cap}") + print(f"Number of chunks: {args.num_chunks}") + print(f"Export shape(s): {args.shapes}") + print(f"Platform: {args.platform}") + print() + + +def apply_preformatter(inp, preformatter=None): + formatted_text = preformatter.generate_prompt(inp["text"]) + inp["text"] = formatted_text + print(f"Formatted Prompt:\n{formatted_text}") + return inp + + +def tokenize_dataset(inp, tokenizer): + text = inp["text"] + inp_encoded = tokenizer(text, return_tensors="pt") # dict + inp_encoded.pop("attention_mask") + inp_encoded = inp_encoded["input_ids"] + inp_encoded = inp_encoded.to(torch.int32) + inp["input_ids"] = inp_encoded + inp.pop("text") + return inp + + +def reset_cache( + num_chunks, num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size +): + cache = [] + for i in range(num_chunks): + curr_chunk_cache = torch.zeros( + ( + 2 * num_blocks_per_chunk[i], + num_key_value_heads, + max_cache_size, # generate fixed cache as torch dynamic shape cannot handle 2 dynamic dim + head_dim, + ), + dtype=torch.float32, + ) + cache.append(curr_chunk_cache) + return cache + + +def forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + batch_name, +): + for chunk_idx in range(len(models)): + cache_in = cache[chunk_idx] + + try: + model_input_dict[str(chunk_idx)] = { + **model_input_dict[str(chunk_idx)], + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + }, + } + except: + model_input_dict[str(chunk_idx)] = { + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + } + } + + if isinstance(mask, dict): + global_mask = mask["GLOBAL"] + local_mask = mask["SLIDING_LOCAL"] + with torch.no_grad(): + model_out = models[chunk_idx]( + hidden_state, + global_mask, + local_mask, + pos_emb, + *torch.split(cache_in, 1, dim=0), + ) + else: + with torch.no_grad(): + model_out = models[chunk_idx]( + hidden_state, mask, pos_emb, *torch.split(cache_in, 1, dim=0) + ) + + hidden_state = model_out[0] + cache[chunk_idx] = torch.cat( + model_out[1 : 1 + 2 * num_blocks_per_chunk[chunk_idx]], dim=0 + ).clone() + return hidden_state, cache + + +def prepare_model_inputs( + inp, + models, + embedding_layer, + master_rot_emb, + num_blocks_per_chunk, + num_key_value_heads, + head_dim, + max_cache_size, + eos_token_id_tensor, + response_cap, + window_size, +): + model_input_dict = {str(i): None for i in range(len(models))} + input_ids = inp.pop("input_ids") + hidden_state = embedding_layer(torch.tensor(input_ids)) + input_length = hidden_state.shape[1] + # Assume fixed cache size + global_mask = generate_mask(max_cache_size, 0, input_length, input_length) + if window_size is not None: + local_mask = generate_mask( + max_cache_size, + 0, + input_length, + input_length, + sliding_window=True, + sliding_window_size=window_size, + ) + mask = {"SLIDING_LOCAL": local_mask, "GLOBAL": global_mask} + else: + mask = global_mask + pos_emb = master_rot_emb[:, :, :input_length, :] + # cache shape: num chunks of 2*num_block, num kv heads, c, head dim + cache = reset_cache( + len(models), num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size + ) # empty kv + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + "prompt", + ) + next_token_logits = logits[:, -1, :] # last layer logits + next_token = torch.argmax(next_token_logits, dim=-1) + response_count = 0 + seq_length = input_length + while True: + curr_input_id = next_token[:, None].to(torch.int32) + input_length = curr_input_id.shape[1] + hidden_state = embedding_layer(curr_input_id) + global_mask = generate_mask( + max_cache_size, seq_length, input_length, input_length + ) + if window_size is not None: + local_mask = generate_mask( + max_cache_size, + seq_length, + input_length, + input_length, + sliding_window=True, + sliding_window_size=window_size, + ) + mask = {"SLIDING_LOCAL": local_mask, "GLOBAL": global_mask} + else: + mask = global_mask + pos_emb = master_rot_emb[:, :, seq_length : seq_length + input_length, :] + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + f"response{response_count}", + ) + next_token_logits = logits[:, -1, :] + next_token = torch.argmax(next_token_logits, dim=-1) + if next_token == eos_token_id_tensor: + print(f"Found EOS on batch: {response_count}") + break + + response_count += 1 + seq_length += input_length + if response_count == response_cap: + break + + return model_input_dict + + +def calibrate_model(model, cal_dataset, chunk_idx: str): + with torch.no_grad(): + for inp in tqdm(cal_dataset, desc="Calibrating Model: "): + # pass prompt and response + for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "): + if inp[chunk_idx][batch] is not None: + inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"]) + pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"]) + cache = torch.tensor(inp[chunk_idx][batch]["cache"]) + mask = inp[chunk_idx][batch]["mask"] + if isinstance(mask, dict): + global_mask = torch.tensor(mask["GLOBAL"]) + local_mask = torch.tensor(mask["SLIDING_LOCAL"]) + model( + inputs_embeds, + global_mask, + local_mask, + pos_emb, + *torch.split(cache, 1, dim=0), + ) + else: + mask = torch.tensor(mask) + model( + inputs_embeds, mask, pos_emb, *torch.split(cache, 1, dim=0) + ) + + +def export_to_et_ir( + output_folder, + exp_name, + model, + precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset=None, +): + print(f"Exporting Chunk {chunk_idx} to PTE") + example_inputs, dynamic_shapes = model.get_example_inputs( + max_num_token, max_cache_size, True + ) + print("Getting pre autograd ATen Dialect Graph") + pre_autograd_aten_dialect = torch.export.export_for_training( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() # NOTE: Will be replaced with export + quantizer = NeuropilotQuantizer() + quantizer.setup_precision(getattr(Precision, precision)) + prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) + # at this point quant min max are inf + if cal_dataset is not None: + calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) + else: + prepared_graph(*example_inputs) # dummy calibration + converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) + + method_to_edge_program = {} + method_to_partitioner = {} + edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) + + model_shared_key_name = f"{exp_name}_{chunk_idx}" + + # Fixed Shape Export Here + for shape, ntok_and_cache in export_shapes.items(): + model_fname = f"{exp_name}_{shape}_{chunk_idx}" + example_inputs = model.get_example_inputs(*ntok_and_cache) + print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") + aten_dialect: exir.ExportedProgram = torch.export.export( + converted_graph, example_inputs, strict=True + ) + + method_to_edge_program[f"{model_fname}"] = exir.to_edge( + aten_dialect + ).exported_program() + del aten_dialect + + compile_spec = [ + CompileSpec("gno", b"LTS"), + CompileSpec("gno-exp", b""), + CompileSpec("gno-non-4d-tiling", b""), + CompileSpec("ImportForever", struct.pack("?", True)), + CompileSpec("platform-config", platform_b), + CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), + ] + method_to_partitioner[f"{model_fname}"] = NeuropilotPartitioner(compile_spec) + + print("Delegating Edge Program to Neuropilot Backend") + delegated_program = to_backend( + MethodProgramsPartitionerSpec(method_to_edge_program, method_to_partitioner) + ) + + edge_manager = exir.EdgeProgramManager( + delegated_program, compile_config=edge_compile_config + ) + del delegated_program + + print("Transforming delegated program to executorch backend") + executorch_program = edge_manager.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, + ) + ) + del edge_manager + print(f"\n Model Size: {len(executorch_program.buffer)}") + + dest_path = get_dest_path(output_folder, exp_name, None, chunk_idx) + print(f"{exp_name} ET Model chunk {chunk_idx} Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) + + +def main(): + parser = get_argument_parser() + args = parser.parse_args() + args_sanity_checks(args) + if args.dataset is None: + exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" + else: + exp_name = ( + f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" + ) + if args.platform == "DX4": + platform_b = b"mt6991" + elif args.platform == "DX3": + platform_b = b"mt6989" + else: + raise ValueError( + f"Platform should be either DX3 or DX4, but got {args.platform}" + ) + print_args(args, exp_name) + + config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( + args.config + ) + tokenizer = tokenizer_class.from_pretrained(weight_dir) + if args.preformatter is not None: + preformatter = Preformatter(args.preformatter) + + head_dim = int(config.head_dim) + + # Evenly distribute the layers across chunks. + num_blocks_per_chunk = [ + (config.num_hidden_layers // args.num_chunks) + + (i < (config.num_hidden_layers % args.num_chunks)) + for i in range(args.num_chunks) + ] + check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 + + output_folder = os.path.join("pte", exp_name) + + # Load all collected checkpoint files into one giant state_dict + state_dict = load_checkpoints(weight_dir) + + dump_embedding_lut_for_cmdline(weight_dir, state_dict, config) + + export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) + print(f"export shapes: {export_shapes}") + print(f"Max Num Token: {max_num_token}") + print(f"Max Cache Size: {max_cache_size}") + + if args.dataset is not None: + embedding_layer = get_embedding_layer(config, weight_dir, state_dict) + + # Instantiate model chunks + print("Instantiating submodels") + models = [] + for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): + chunk = chunk_class( + config, + num_blocks, + chunk_idx=chunk_idx, + dtype=torch.float32, + include_tail=(chunk_idx == args.num_chunks - 1), + jit_trace=True, + ) + chunk = chunk.load_weights(state_dict, sum(num_blocks_per_chunk[:chunk_idx])) + models.append(chunk) + + cal_dataset = None + if args.dataset is not None: + cal_dataset = load_dataset("text", data_files=args.dataset, split="train") + master_rot_emb = get_master_rot_emb(config, dtype=torch.float32) + if args.preformatter is not None: + cal_dataset = cal_dataset.map( + apply_preformatter, fn_kwargs={"preformatter": preformatter} + ) + cal_dataset = cal_dataset.map( + tokenize_dataset, fn_kwargs={"tokenizer": tokenizer} + ) + window_size = getattr(config, "sliding_window_attention_size", None) + print("Preparing Model Calibration Inputs...") + cal_dataset = cal_dataset.map( + prepare_model_inputs, + fn_kwargs={ + "models": models, + "embedding_layer": embedding_layer, + "master_rot_emb": master_rot_emb, + "num_blocks_per_chunk": num_blocks_per_chunk, + "num_key_value_heads": config.num_key_value_heads, + "head_dim": head_dim, + "max_cache_size": max_cache_size, + "eos_token_id_tensor": torch.tensor(tokenizer.eos_token_id), + "response_cap": args.response_cap, + "window_size": window_size, + }, + ) + + for chunk_idx, chunk in enumerate(models): + export_to_et_ir( + output_folder, + exp_name, + chunk, + args.precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/mediatek/model_export_scripts/llama.py b/examples/mediatek/model_export_scripts/llama.py index 6a098e2a9b1..a58c87fa0d6 100644 --- a/examples/mediatek/model_export_scripts/llama.py +++ b/examples/mediatek/model_export_scripts/llama.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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 os import sys @@ -42,6 +48,10 @@ NeuropilotQuantizer, Precision, ) +from executorch.exir.backend.backend_api import ( + MethodProgramsPartitionerSpec, + to_backend, +) from executorch.exir.backend.backend_details import CompileSpec from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from tqdm import tqdm @@ -67,6 +77,14 @@ def get_argument_parser(): choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], help="Precision to quantize entire model to.", ) + parser.add_argument( + "--platform", + type=str, + default="DX4", + choices=["DX3", "DX4"], + help="Chip model of the inference device. " + "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", + ) parser.add_argument( "-d", "--dataset", @@ -140,6 +158,7 @@ def print_args(args, exp_name): print(f"Max Response Tokens: {args.response_cap}") print(f"Number of chunks: {args.num_chunks}") print(f"Export shape(s): {args.shapes}") + print(f"Platform: {args.platform}") print() @@ -311,6 +330,7 @@ def export_to_et_ir( max_cache_size, chunk_idx, export_shapes, + platform_b, cal_dataset=None, ): print(f"Exporting Chunk {chunk_idx} to PTE") @@ -318,7 +338,7 @@ def export_to_et_ir( max_num_token, max_cache_size, True ) print("Getting pre autograd ATen Dialect Graph") - pre_autograd_aten_dialect = torch.export.export_for_training( + pre_autograd_aten_dialect = torch.export.export( model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True ).module() # NOTE: Will be replaced with export quantizer = NeuropilotQuantizer() @@ -331,52 +351,64 @@ def export_to_et_ir( prepared_graph(*example_inputs) # dummy calibration converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) - print("Getting ATen Dialect Graph") + method_to_edge_program = {} + method_to_partitioner = {} + edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) + + model_shared_key_name = f"{exp_name}_{chunk_idx}" + # Fixed Shape Export Here for shape, ntok_and_cache in export_shapes.items(): - dest_path = get_dest_path(output_folder, exp_name, shape, chunk_idx) - print(f"Exporting Shape {shape} to:\n{dest_path}") + model_fname = f"{exp_name}_{shape}_{chunk_idx}" example_inputs = model.get_example_inputs(*ntok_and_cache) + print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") aten_dialect: exir.ExportedProgram = torch.export.export( converted_graph, example_inputs, strict=True ) - print("Lowering to Edge Dialect Graph") - edge_program: exir.EdgeProgramManager = exir.to_edge( - aten_dialect, - compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), - ) + method_to_edge_program[f"{model_fname}"] = exir.to_edge( + aten_dialect + ).exported_program() del aten_dialect - print("Delegating Edge Program to Neuropilot Backend") compile_spec = [ CompileSpec("gno", b"LTS"), CompileSpec("gno-exp", b""), CompileSpec("gno-non-4d-tiling", b""), CompileSpec("ImportForever", struct.pack("?", True)), - CompileSpec("platform-config", b"mt6989"), + CompileSpec("platform-config", platform_b), + CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), ] - partitioner = NeuropilotPartitioner(compile_spec) - delegated_program = edge_program.to_backend(partitioner) - print("Exported Delegated Program:") - print(delegated_program.exported_program()) - del edge_program - - print("Transforming delegated program to executorch backend") - executorch_program = delegated_program.to_executorch( - config=exir.ExecutorchBackendConfig( - memory_planning_pass=exir.passes.MemoryPlanningPass( - alloc_graph_input=False, - alloc_graph_output=False, - ), - extract_delegate_segments=True, - ) + method_to_partitioner[f"{model_fname}"] = NeuropilotPartitioner(compile_spec) + + print("Delegating Edge Program to Neuropilot Backend") + delegated_program = to_backend( + MethodProgramsPartitionerSpec(method_to_edge_program, method_to_partitioner) + ) + + edge_manager = exir.EdgeProgramManager( + delegated_program, compile_config=edge_compile_config + ) + del delegated_program + + print("Transforming delegated program to executorch backend") + executorch_program = edge_manager.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, ) + ) + del edge_manager + print(f"\n Model Size: {len(executorch_program.buffer)}") - print(f"ET Model Dest: {dest_path}\n") - os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) - with open(dest_path, "wb") as file: - file.write(executorch_program.buffer) + dest_path = get_dest_path(output_folder, exp_name, None, chunk_idx) + print(f"{exp_name} ET Model chunk {chunk_idx} Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) def main(): @@ -389,6 +421,14 @@ def main(): exp_name = ( f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" ) + if args.platform == "DX4": + platform_b = b"mt6991" + elif args.platform == "DX3": + platform_b = b"mt6989" + else: + raise ValueError( + f"Platform should be either DX3 or DX4, but got {args.platform}" + ) print_args(args, exp_name) config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( @@ -475,6 +515,7 @@ def main(): max_cache_size, chunk_idx, export_shapes, + platform_b, cal_dataset, ) diff --git a/examples/mediatek/model_export_scripts/phi.py b/examples/mediatek/model_export_scripts/phi.py new file mode 100644 index 00000000000..062a4c1c571 --- /dev/null +++ b/examples/mediatek/model_export_scripts/phi.py @@ -0,0 +1,524 @@ +# Copyright (c) MediaTek 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 os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) +import argparse +import struct +import warnings + +import torch + +from aot_utils.llm_utils.preformatter import Preformatter +from aot_utils.llm_utils.sanity_checks import ( + check_all_chunks_same_num_layer, + check_between_inclusive, + check_exist, + check_ext, + check_old_arg, + check_shapes, + check_supported_model, + check_supported_tokenizer, + check_tokenizer_exist, + check_weights_exist, +) +from aot_utils.llm_utils.utils import ( + dump_embedding_lut_for_cmdline, + generate_mask, + get_dest_path, + get_dirname, + get_embedding_layer, + get_exp_name, + get_export_shapes, + get_master_rot_emb, + get_normalized_config, + load_checkpoints, + resolve_model_classes, +) +from datasets import load_dataset +from executorch import exir +from executorch.backends.mediatek import ( + NeuropilotPartitioner, + NeuropilotQuantizer, + Precision, +) +from executorch.exir.backend.backend_api import ( + MethodProgramsPartitionerSpec, + to_backend, +) +from executorch.exir.backend.backend_details import CompileSpec +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from tqdm import tqdm + +warnings.filterwarnings("ignore") + + +def get_argument_parser(): + parser = argparse.ArgumentParser( + description="Run Export to ET for suppoorted LLM models.", allow_abbrev=False + ) + parser.add_argument( + "config", + type=str, + help="[Required] Model config json file. " + "Model config must be in same directory as all model weight bins and tokenizer files.", + ) + parser.add_argument( + "-p", + "--precision", + type=str, + default="A16W8", + choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], + help="Precision to quantize entire model to.", + ) + parser.add_argument( + "--platform", + type=str, + default="DX4", + choices=["DX3", "DX4"], + help="Chip model of the inference device. " + "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", + ) + parser.add_argument( + "-d", + "--dataset", + type=str, + default=None, + help="Calibration dataset name or path to dataset. Defaults to None to use random inputs", + ) + parser.add_argument( + "-n", + "--num_chunks", + type=int, + default=4, + help="Number of chunks to cut the model into. Defaults to 4.", + ) + parser.add_argument( + "-r", + "--response_cap", + type=int, + default=9, + help="Max Number of Response Tokens to save during calibration. Defaults to 9.", + ) + parser.add_argument( + "--preformatter", + type=str, + default=None, + help="Preformatter Template to use to wrap input with. Defaults to None.", + ) + parser.add_argument( + "-shapes", + nargs="+", + help="[Required] Expected input shapes to reconfigure TFLites to. Space separated list of " + "shapes in the format: xtyc (e.g. 32t512c)", + ) + + return parser + + +# flake8: noqa: F405 +def args_sanity_checks(args): + check_old_arg(args.config) + check_exist(args.config, "Config file") + check_ext(args.config, ".json", "Config file") + config = get_normalized_config(args.config) + + weight_dir = get_dirname(args.config) + check_tokenizer_exist(weight_dir) + check_weights_exist(weight_dir) + + check_supported_model(config) + check_supported_tokenizer(config) + + if args.preformatter is not None: + check_exist(args.preformatter, "Preformatter json file") + check_ext(args.preformatter, ".json", "preformatter") + + if args.dataset is not None: + check_exist(args.dataset) + + check_between_inclusive(args.num_chunks, 1, config.num_hidden_layers, "num_chunks") + + check_shapes(args.shapes) + + +def print_args(args, exp_name): + print("Please check if all arguments are correct:") + print(f"Config file: {args.config}") + print(f"Output pte folder: pte/{exp_name}") + print(f"Quantization precision: {args.precision}") + print(f"Preformatter: {args.preformatter}") + print(f"Calibration Dataset: {args.dataset}") + print(f"Max Response Tokens: {args.response_cap}") + print(f"Number of chunks: {args.num_chunks}") + print(f"Export shape(s): {args.shapes}") + print(f"Platform: {args.platform}") + print() + + +def apply_preformatter(inp, preformatter=None): + formatted_text = preformatter.generate_prompt(inp["text"]) + inp["text"] = formatted_text + print(f"Formatted Prompt:\n{formatted_text}") + return inp + + +def tokenize_dataset(inp, tokenizer): + text = inp["text"] + inp_encoded = tokenizer(text, return_tensors="pt") # dict + inp_encoded.pop("attention_mask") + inp_encoded = inp_encoded["input_ids"] + inp_encoded = inp_encoded.to(torch.int32) + inp["input_ids"] = inp_encoded + inp.pop("text") + return inp + + +def reset_cache( + num_chunks, num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size +): + cache = [] + for i in range(num_chunks): + curr_chunk_cache = torch.zeros( + ( + 2 * num_blocks_per_chunk[i], + num_key_value_heads, + max_cache_size, # generate fixed cache as torch dynamic shape cannot handle 2 dynamic dim + head_dim, + ), + dtype=torch.float32, + ) + cache.append(curr_chunk_cache) + return cache + + +def forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + batch_name, +): + for chunk_idx in range(len(models)): + cache_in = cache[chunk_idx] + + try: + model_input_dict[str(chunk_idx)] = { + **model_input_dict[str(chunk_idx)], + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + }, + } + except: + model_input_dict[str(chunk_idx)] = { + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + } + } + with torch.no_grad(): + model_out = models[chunk_idx]( + hidden_state, mask, pos_emb, *torch.split(cache_in, 1, dim=0) + ) + hidden_state = model_out[0] + cache[chunk_idx] = torch.cat( + model_out[1 : 1 + 2 * num_blocks_per_chunk[chunk_idx]], dim=0 + ).clone() + return hidden_state, cache + + +def prepare_model_inputs( + inp, + models, + embedding_layer, + master_rot_emb, + num_blocks_per_chunk, + num_key_value_heads, + head_dim, + max_cache_size, + eos_token_id_tensor, + response_cap, +): + model_input_dict = {str(i): None for i in range(len(models))} + input_ids = inp.pop("input_ids") + hidden_state = embedding_layer(torch.tensor(input_ids)) + input_length = hidden_state.shape[1] + # Assume fixed cache size + mask = generate_mask(max_cache_size, 0, input_length, input_length) + pos_emb = master_rot_emb[:, :, :input_length, :] + # cache shape: num chunks of 2*num_block, num kv heads, c, head dim + cache = reset_cache( + len(models), num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size + ) # empty kv + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + "prompt", + ) + next_token_logits = logits[:, -1, :] # last layer logits + next_token = torch.argmax(next_token_logits, dim=-1) + response_count = 0 + seq_length = input_length + while True: + curr_input_id = next_token[:, None].to(torch.int32) + input_length = curr_input_id.shape[1] + hidden_state = embedding_layer(curr_input_id) + mask = generate_mask(max_cache_size, seq_length, input_length, input_length) + pos_emb = master_rot_emb[:, :, seq_length : seq_length + input_length, :] + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + f"response{response_count}", + ) + next_token_logits = logits[:, -1, :] + next_token = torch.argmax(next_token_logits, dim=-1) + if next_token == eos_token_id_tensor: + print(f"Found EOS on batch: {response_count}") + break + + response_count += 1 + seq_length += input_length + if response_count == response_cap: + break + + return model_input_dict + + +def calibrate_model(model, cal_dataset, chunk_idx: str): + with torch.no_grad(): + for inp in tqdm(cal_dataset, desc="Calibrating Model: "): + # pass prompt and response + for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "): + if inp[chunk_idx][batch] is not None: + inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"]) + mask = torch.tensor(inp[chunk_idx][batch]["mask"]) + pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"]) + cache = torch.tensor(inp[chunk_idx][batch]["cache"]) + model(inputs_embeds, mask, pos_emb, *torch.split(cache, 1, dim=0)) + + +def export_to_et_ir( + output_folder, + exp_name, + model, + precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset=None, +): + print(f"Exporting Chunk {chunk_idx} to PTE") + example_inputs, dynamic_shapes = model.get_example_inputs( + max_num_token, max_cache_size, True + ) + print("Getting pre autograd ATen Dialect Graph") + pre_autograd_aten_dialect = torch.export.export_for_training( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() # NOTE: Will be replaced with export + quantizer = NeuropilotQuantizer() + quantizer.setup_precision(getattr(Precision, precision)) + prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) + # at this point quant min max are inf + if cal_dataset is not None: + calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) + else: + prepared_graph(*example_inputs) # dummy calibration + converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) + + method_to_edge_program = {} + method_to_partitioner = {} + edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) + + model_shared_key_name = f"{exp_name}_{chunk_idx}" + + # Fixed Shape Export Here + for shape, ntok_and_cache in export_shapes.items(): + model_fname = f"{exp_name}_{shape}_{chunk_idx}" + example_inputs = model.get_example_inputs(*ntok_and_cache) + print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") + aten_dialect: exir.ExportedProgram = torch.export.export( + converted_graph, example_inputs, strict=True + ) + + method_to_edge_program[f"{model_fname}"] = exir.to_edge( + aten_dialect + ).exported_program() + del aten_dialect + + compile_spec = [ + CompileSpec("gno", b"LTS"), + CompileSpec("gno-exp", b""), + CompileSpec("gno-non-4d-tiling", b""), + CompileSpec("ImportForever", struct.pack("?", True)), + CompileSpec("platform-config", platform_b), + CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), + ] + method_to_partitioner[f"{model_fname}"] = NeuropilotPartitioner(compile_spec) + + print("Delegating Edge Program to Neuropilot Backend") + delegated_program = to_backend( + MethodProgramsPartitionerSpec(method_to_edge_program, method_to_partitioner) + ) + + edge_manager = exir.EdgeProgramManager( + delegated_program, compile_config=edge_compile_config + ) + del delegated_program + + print("Transforming delegated program to executorch backend") + executorch_program = edge_manager.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, + ) + ) + del edge_manager + print(f"\n Model Size: {len(executorch_program.buffer)}") + + dest_path = get_dest_path(output_folder, exp_name, None, chunk_idx) + print(f"{exp_name} ET Model chunk {chunk_idx} Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) + + +def main(): + parser = get_argument_parser() + args = parser.parse_args() + args_sanity_checks(args) + if args.dataset is None: + exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" + else: + exp_name = ( + f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" + ) + if args.platform == "DX4": + platform_b = b"mt6991" + elif args.platform == "DX3": + platform_b = b"mt6989" + else: + raise ValueError( + f"Platform should be either DX3 or DX4, but got {args.platform}" + ) + print_args(args, exp_name) + + config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( + args.config + ) + tokenizer = tokenizer_class.from_pretrained(weight_dir) + if args.preformatter is not None: + preformatter = Preformatter(args.preformatter) + + head_dim = int(config.head_dim) + + # Evenly distribute the layers across chunks. + num_blocks_per_chunk = [ + (config.num_hidden_layers // args.num_chunks) + + (i < (config.num_hidden_layers % args.num_chunks)) + for i in range(args.num_chunks) + ] + check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 + + output_folder = os.path.join("pte", exp_name) + + # Load all collected checkpoint files into one giant state_dict + state_dict = load_checkpoints(weight_dir) + + dump_embedding_lut_for_cmdline(weight_dir, state_dict, config) + + export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) + print(f"export shapes: {export_shapes}") + print(f"Max Num Token: {max_num_token}") + print(f"Max Cache Size: {max_cache_size}") + + if args.dataset is not None: + embedding_layer = get_embedding_layer(config, weight_dir, state_dict) + + # Instantiate model chunks + print("Instantiating submodels") + models = [] + for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): + chunk = chunk_class( + config, + num_blocks, + chunk_idx=chunk_idx, + dtype=torch.float32, + include_tail=(chunk_idx == args.num_chunks - 1), + jit_trace=True, + ) + chunk = chunk.load_weights(state_dict, sum(num_blocks_per_chunk[:chunk_idx])) + models.append(chunk) + + cal_dataset = None + if args.dataset is not None: + cal_dataset = load_dataset("text", data_files=args.dataset, split="train") + master_rot_emb = get_master_rot_emb(config, dtype=torch.float32) + if args.preformatter is not None: + cal_dataset = cal_dataset.map( + apply_preformatter, fn_kwargs={"preformatter": preformatter} + ) + cal_dataset = cal_dataset.map( + tokenize_dataset, fn_kwargs={"tokenizer": tokenizer} + ) + print("Preparing Model Calibration Inputs...") + cal_dataset = cal_dataset.map( + prepare_model_inputs, + fn_kwargs={ + "models": models, + "embedding_layer": embedding_layer, + "master_rot_emb": master_rot_emb, + "num_blocks_per_chunk": num_blocks_per_chunk, + "num_key_value_heads": config.num_key_value_heads, + "head_dim": head_dim, + "max_cache_size": max_cache_size, + "eos_token_id_tensor": torch.tensor(tokenizer.eos_token_id), + "response_cap": args.response_cap, + }, + ) + + for chunk_idx, chunk in enumerate(models): + export_to_et_ir( + output_folder, + exp_name, + chunk, + args.precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/mediatek/model_export_scripts/qwen.py b/examples/mediatek/model_export_scripts/qwen.py new file mode 100644 index 00000000000..21a97dba6a2 --- /dev/null +++ b/examples/mediatek/model_export_scripts/qwen.py @@ -0,0 +1,525 @@ +# Copyright (c) MediaTek 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 os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) +import argparse +import struct +import warnings + +import torch + +from aot_utils.llm_utils.preformatter import Preformatter +from aot_utils.llm_utils.sanity_checks import ( + check_all_chunks_same_num_layer, + check_between_inclusive, + check_exist, + check_ext, + check_old_arg, + check_shapes, + check_supported_model, + check_supported_tokenizer, + check_tokenizer_exist, + check_weights_exist, +) +from aot_utils.llm_utils.utils import ( + dump_embedding_lut_for_cmdline, + generate_mask, + get_dest_path, + get_dirname, + get_embedding_layer, + get_exp_name, + get_export_shapes, + get_master_rot_emb, + get_normalized_config, + load_checkpoints, + resolve_model_classes, +) +from datasets import load_dataset +from executorch import exir +from executorch.backends.mediatek import ( + NeuropilotPartitioner, + NeuropilotQuantizer, + Precision, +) +from executorch.exir.backend.backend_api import ( + MethodProgramsPartitionerSpec, + to_backend, +) +from executorch.exir.backend.backend_details import CompileSpec +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from tqdm import tqdm + +warnings.filterwarnings("ignore") + + +def get_argument_parser(): + parser = argparse.ArgumentParser( + description="Run Export to ET for suppoorted LLM models.", allow_abbrev=False + ) + parser.add_argument( + "config", + type=str, + help="[Required] Model config json file. " + "Model config must be in same directory as all model weight bins and tokenizer files.", + ) + parser.add_argument( + "-p", + "--precision", + type=str, + default="A16W8", + choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], + help="Precision to quantize entire model to.", + ) + parser.add_argument( + "--platform", + type=str, + default="DX4", + choices=["DX3", "DX4"], + help="Chip model of the inference device. " + "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", + ) + parser.add_argument( + "-d", + "--dataset", + type=str, + default=None, + help="Calibration dataset name or path to dataset. Defaults to None to use random inputs", + ) + parser.add_argument( + "-n", + "--num_chunks", + type=int, + default=4, + help="Number of chunks to cut the model into. Defaults to 4.", + ) + parser.add_argument( + "-r", + "--response_cap", + type=int, + default=9, + help="Max Number of Response Tokens to save during calibration. Defaults to 9.", + ) + parser.add_argument( + "--preformatter", + type=str, + default=None, + help="Preformatter Template to use to wrap input with. Defaults to None.", + ) + parser.add_argument( + "-shapes", + nargs="+", + help="[Required] Expected input shapes to reconfigure TFLites to. Space separated list of " + "shapes in the format: xtyc (e.g. 32t512c)", + ) + + return parser + + +# flake8: noqa: F405 +def args_sanity_checks(args): + check_old_arg(args.config) + check_exist(args.config, "Config file") + check_ext(args.config, ".json", "Config file") + config = get_normalized_config(args.config) + + weight_dir = get_dirname(args.config) + check_tokenizer_exist(weight_dir) + check_weights_exist(weight_dir) + + check_supported_model(config) + check_supported_tokenizer(config) + + if args.preformatter is not None: + check_exist(args.preformatter, "Preformatter json file") + check_ext(args.preformatter, ".json", "preformatter") + + if args.dataset is not None: + check_exist(args.dataset) + + check_between_inclusive(args.num_chunks, 1, config.num_hidden_layers, "num_chunks") + + check_shapes(args.shapes) + + +def print_args(args, exp_name): + print("Please check if all arguments are correct:") + print(f"Config file: {args.config}") + print(f"Output pte folder: pte/{exp_name}") + print(f"Quantization precision: {args.precision}") + print(f"Preformatter: {args.preformatter}") + print(f"Calibration Dataset: {args.dataset}") + print(f"Max Response Tokens: {args.response_cap}") + print(f"Number of chunks: {args.num_chunks}") + print(f"Export shape(s): {args.shapes}") + print(f"Platform: {args.platform}") + print() + + +def apply_preformatter(inp, preformatter=None): + formatted_text = preformatter.generate_prompt(inp["text"]) + inp["text"] = formatted_text + print(f"Formatted Prompt:\n{formatted_text}") + return inp + + +def tokenize_dataset(inp, tokenizer): + text = inp["text"] + inp_encoded = tokenizer(text, return_tensors="pt") # dict + inp_encoded.pop("attention_mask") + inp_encoded = inp_encoded["input_ids"] + inp_encoded = inp_encoded.to(torch.int32) + inp["input_ids"] = inp_encoded + inp.pop("text") + return inp + + +def reset_cache( + num_chunks, num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size +): + cache = [] + for i in range(num_chunks): + curr_chunk_cache = torch.zeros( + ( + 2 * num_blocks_per_chunk[i], + num_key_value_heads, + max_cache_size, # generate fixed cache as torch dynamic shape cannot handle 2 dynamic dim + head_dim, + ), + dtype=torch.float32, + ) + cache.append(curr_chunk_cache) + return cache + + +def forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + batch_name, +): + for chunk_idx in range(len(models)): + cache_in = cache[chunk_idx] + + try: + model_input_dict[str(chunk_idx)] = { + **model_input_dict[str(chunk_idx)], + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + }, + } + except: + model_input_dict[str(chunk_idx)] = { + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + } + } + with torch.no_grad(): + model_out = models[chunk_idx]( + hidden_state, mask, pos_emb, *torch.split(cache_in, 1, dim=0) + ) + hidden_state = model_out[0] + cache[chunk_idx] = torch.cat( + model_out[1 : 1 + 2 * num_blocks_per_chunk[chunk_idx]], dim=0 + ).clone() + return hidden_state, cache + + +def prepare_model_inputs( + inp, + models, + embedding_layer, + master_rot_emb, + num_blocks_per_chunk, + num_key_value_heads, + head_dim, + max_cache_size, + eos_token_id_tensor, + response_cap, +): + model_input_dict = {str(i): None for i in range(len(models))} + input_ids = inp.pop("input_ids") + hidden_state = embedding_layer(torch.tensor(input_ids)) + input_length = hidden_state.shape[1] + # Assume fixed cache size + mask = generate_mask(max_cache_size, 0, input_length, input_length) + pos_emb = master_rot_emb[:, :, :input_length, :] + # cache shape: num chunks of 2*num_block, num kv heads, c, head dim + cache = reset_cache( + len(models), num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size + ) # empty kv + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + "prompt", + ) + next_token_logits = logits[:, -1, :] # last layer logits + next_token = torch.argmax(next_token_logits, dim=-1) + response_count = 0 + seq_length = input_length + while True: + curr_input_id = next_token[:, None].to(torch.int32) + input_length = curr_input_id.shape[1] + hidden_state = embedding_layer(curr_input_id) + mask = generate_mask(max_cache_size, seq_length, input_length, input_length) + pos_emb = master_rot_emb[:, :, seq_length : seq_length + input_length, :] + logits, cache = forward_and_save( + models, + hidden_state, + cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + f"response{response_count}", + ) + next_token_logits = logits[:, -1, :] + next_token = torch.argmax(next_token_logits, dim=-1) + + if next_token == eos_token_id_tensor: + print(f"Found EOS on batch: {response_count}") + break + + response_count += 1 + seq_length += input_length + if response_count == response_cap: + break + + return model_input_dict + + +def calibrate_model(model, cal_dataset, chunk_idx: str): + with torch.no_grad(): + for inp in tqdm(cal_dataset, desc="Calibrating Model: "): + # pass prompt and response + for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "): + if inp[chunk_idx][batch] is not None: + inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"]) + mask = torch.tensor(inp[chunk_idx][batch]["mask"]) + pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"]) + cache = torch.tensor(inp[chunk_idx][batch]["cache"]) + model(inputs_embeds, mask, pos_emb, *torch.split(cache, 1, dim=0)) + + +def export_to_et_ir( + output_folder, + exp_name, + model, + precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset=None, +): + print(f"Exporting Chunk {chunk_idx} to PTE") + example_inputs, dynamic_shapes = model.get_example_inputs( + max_num_token, max_cache_size, True + ) + print("Getting pre autograd ATen Dialect Graph") + pre_autograd_aten_dialect = torch.export.export_for_training( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() # NOTE: Will be replaced with export + quantizer = NeuropilotQuantizer() + quantizer.setup_precision(getattr(Precision, precision)) + prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) + # at this point quant min max are inf + if cal_dataset is not None: + calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) + else: + prepared_graph(*example_inputs) # dummy calibration + converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) + + method_to_edge_program = {} + method_to_partitioner = {} + edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False) + + model_shared_key_name = f"{exp_name}_{chunk_idx}" + + # Fixed Shape Export Here + for shape, ntok_and_cache in export_shapes.items(): + model_fname = f"{exp_name}_{shape}_{chunk_idx}" + example_inputs = model.get_example_inputs(*ntok_and_cache) + print(f"Getting ATen Dialect Graph for {exp_name} {shape} chunk {chunk_idx}") + aten_dialect: exir.ExportedProgram = torch.export.export( + converted_graph, example_inputs, strict=True + ) + + method_to_edge_program[f"{model_fname}"] = exir.to_edge( + aten_dialect + ).exported_program() + del aten_dialect + + compile_spec = [ + CompileSpec("gno", b"LTS"), + CompileSpec("gno-exp", b""), + CompileSpec("gno-non-4d-tiling", b""), + CompileSpec("ImportForever", struct.pack("?", True)), + CompileSpec("platform-config", platform_b), + CompileSpec("ExtractSharedBlobKey", model_shared_key_name.encode()), + ] + method_to_partitioner[f"{model_fname}"] = NeuropilotPartitioner(compile_spec) + + print("Delegating Edge Program to Neuropilot Backend") + delegated_program = to_backend( + MethodProgramsPartitionerSpec(method_to_edge_program, method_to_partitioner) + ) + + edge_manager = exir.EdgeProgramManager( + delegated_program, compile_config=edge_compile_config + ) + del delegated_program + + print("Transforming delegated program to executorch backend") + executorch_program = edge_manager.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, + ) + ) + del edge_manager + print(f"\n Model Size: {len(executorch_program.buffer)}") + + dest_path = get_dest_path(output_folder, exp_name, None, chunk_idx) + print(f"{exp_name} ET Model chunk {chunk_idx} Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) + + +def main(): + parser = get_argument_parser() + args = parser.parse_args() + args_sanity_checks(args) + if args.dataset is None: + exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" + else: + exp_name = ( + f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" + ) + if args.platform == "DX4": + platform_b = b"mt6991" + elif args.platform == "DX3": + platform_b = b"mt6989" + else: + raise ValueError( + f"Platform should be either DX3 or DX4, but got {args.platform}" + ) + print_args(args, exp_name) + + config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( + args.config + ) + tokenizer = tokenizer_class.from_pretrained(weight_dir) + if args.preformatter is not None: + preformatter = Preformatter(args.preformatter) + + head_dim = int(config.head_dim) + + # Evenly distribute the layers across chunks. + num_blocks_per_chunk = [ + (config.num_hidden_layers // args.num_chunks) + + (i < (config.num_hidden_layers % args.num_chunks)) + for i in range(args.num_chunks) + ] + check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 + + output_folder = os.path.join("pte", exp_name) + + # Load all collected checkpoint files into one giant state_dict + state_dict = load_checkpoints(weight_dir) + + dump_embedding_lut_for_cmdline(weight_dir, state_dict, config) + + export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) + print(f"export shapes: {export_shapes}") + print(f"Max Num Token: {max_num_token}") + print(f"Max Cache Size: {max_cache_size}") + + if args.dataset is not None: + embedding_layer = get_embedding_layer(config, weight_dir, state_dict) + + # Instantiate model chunks + print("Instantiating submodels") + models = [] + for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): + chunk = chunk_class( + config, + num_blocks, + chunk_idx=chunk_idx, + dtype=torch.float32, + include_tail=(chunk_idx == args.num_chunks - 1), + jit_trace=True, + ) + chunk = chunk.load_weights(state_dict, sum(num_blocks_per_chunk[:chunk_idx])) + models.append(chunk) + + cal_dataset = None + if args.dataset is not None: + cal_dataset = load_dataset("text", data_files=args.dataset, split="train") + master_rot_emb = get_master_rot_emb(config, dtype=torch.float32) + if args.preformatter is not None: + cal_dataset = cal_dataset.map( + apply_preformatter, fn_kwargs={"preformatter": preformatter} + ) + cal_dataset = cal_dataset.map( + tokenize_dataset, fn_kwargs={"tokenizer": tokenizer} + ) + print("Preparing Model Calibration Inputs...") + cal_dataset = cal_dataset.map( + prepare_model_inputs, + fn_kwargs={ + "models": models, + "embedding_layer": embedding_layer, + "master_rot_emb": master_rot_emb, + "num_blocks_per_chunk": num_blocks_per_chunk, + "num_key_value_heads": config.num_key_value_heads, + "head_dim": head_dim, + "max_cache_size": max_cache_size, + "eos_token_id_tensor": torch.tensor(tokenizer.eos_token_id), + "response_cap": args.response_cap, + }, + ) + + for chunk_idx, chunk in enumerate(models): + export_to_et_ir( + output_folder, + exp_name, + chunk, + args.precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/mediatek/model_export_scripts/whisper.py b/examples/mediatek/model_export_scripts/whisper.py new file mode 100644 index 00000000000..5c4da2e6079 --- /dev/null +++ b/examples/mediatek/model_export_scripts/whisper.py @@ -0,0 +1,689 @@ +# Copyright (c) MediaTek 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 os +import sys + +if os.getcwd() not in sys.path: + sys.path.append(os.getcwd()) +import argparse +import struct +import warnings + +import torch + +from aot_utils.llm_utils.sanity_checks import ( + check_all_chunks_same_num_layer, + check_between_inclusive, + check_exist, + check_ext, + check_old_arg, + check_shapes, + check_supported_model, + check_supported_tokenizer, + check_tokenizer_exist, + check_weights_exist, +) +from aot_utils.llm_utils.utils import ( + dump_embedding_lut_for_cmdline, + generate_mask, + get_dest_path, + get_dirname, + get_embedding_layer, + get_exp_name, + get_export_shapes, + get_master_pos_emb, + get_normalized_config, + load_checkpoints, + resolve_model_classes, +) +from aot_utils.mllm_utils.preprocessor_whisper import WhisperAudioProcessor +from datasets import Audio, Dataset, load_dataset +from executorch import exir +from executorch.backends.mediatek import ( + NeuropilotPartitioner, + NeuropilotQuantizer, + Precision, +) +from executorch.exir.backend.backend_details import CompileSpec +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from tqdm import tqdm + + +warnings.filterwarnings("ignore") + + +def get_argument_parser(): + parser = argparse.ArgumentParser( + description="Run Export to ET for suppoorted LLM models.", allow_abbrev=False + ) + parser.add_argument( + "config", + type=str, + help="[Required] Model config json file. " + "Model config must be in same directory as all model weight bins and tokenizer files.", + ) + parser.add_argument( + "-p", + "--precision", + type=str, + default="A16W8", + choices=["A16W4", "A16W8", "A16W16", "A8W4", "A8W8"], + help="Precision to quantize entire model to.", + ) + parser.add_argument( + "--platform", + type=str, + default="DX4", + choices=["DX3", "DX4"], + help="Chip model of the inference device. " + "DX3 for Dimensity 9300, DX4 for Dimensity 9400.", + ) + parser.add_argument( + "-d", + "--dataset", + type=str, + default=None, + help="Calibration dataset name or path to dataset. Defaults to None to use random inputs", + ) + parser.add_argument( + "-n", + "--num_chunks", + type=int, + default=4, + help="Number of chunks to cut the model into. Defaults to 4.", + ) + parser.add_argument( + "-r", + "--response_cap", + type=int, + default=9, + help="Max Number of Response Tokens to save during calibration. Defaults to 9.", + ) + parser.add_argument( + "--preformatter", + type=str, + default=None, + help="Preformatter Template to use to wrap input with. Defaults to None.", + ) + parser.add_argument( + "-shapes", + nargs="+", + help="[Required] Expected input shapes to reconfigure TFLites to. Space separated list of " + "shapes in the format: xtyc (e.g. 32t512c)", + ) + + return parser + + +# flake8: noqa: F405 +def args_sanity_checks(args): + check_old_arg(args.config) + check_exist(args.config, "Config file") + check_ext(args.config, ".json", "Config file") + config = get_normalized_config(args.config) + + weight_dir = get_dirname(args.config) + check_tokenizer_exist(weight_dir) + check_weights_exist(weight_dir) + + check_supported_model(config) + check_supported_tokenizer(config) + + if args.preformatter is not None: + check_exist(args.preformatter, "Preformatter json file") + check_ext(args.preformatter, ".json", "preformatter") + + if args.dataset is not None: + check_exist(args.dataset) + + check_between_inclusive( + args.num_chunks, 1, config.llm.num_hidden_layers, "num_chunks" + ) + + check_shapes(args.shapes) + + +def print_args(args, exp_name): + print("Please check if all arguments are correct:") + print(f"Config file: {args.config}") + print(f"Output pte folder: pte/{exp_name}") + print(f"Quantization precision: {args.precision}") + print(f"Preformatter: {args.preformatter}") + print(f"Calibration Dataset: {args.dataset}") + print(f"Max Response Tokens: {args.response_cap}") + print(f"Number of chunks: {args.num_chunks}") + print(f"Export shape(s): {args.shapes}") + print(f"Platform: {args.platform}") + print() + + +def apply_preformatter(inp, preformatter=None): + formatted_text = preformatter.generate_prompt(inp["text"]) + inp["text"] = formatted_text + print(f"Formatted Prompt:\n{formatted_text}") + return inp + + +def tokenize_dataset(inp, tokenizer): + text = inp["text"] + inp_encoded = tokenizer(text, return_tensors="pt") # dict + inp_encoded.pop("attention_mask") + inp_encoded = inp_encoded["input_ids"] + inp_encoded = inp_encoded.to(torch.int32) + inp["input_ids"] = inp_encoded + inp.pop("text") + return inp + + +def process_audio(inp, preprocessor, encoder): + audio_path = inp["text"] + if audio_path.endswith("mp3"): + print(f"Loading audio from path: {audio_path}") + mm_input = Dataset.from_dict({"audio": [str(audio_path)]}).cast_column( + "audio", Audio(sampling_rate=16000) + )[0]["audio"]["array"] + else: + print(f"Expected mm input filepath to end with 'mp3', but got {audio_path}") + + processed, _ = preprocessor.preprocess( + mm_input, sampling_rate=16000, return_tensors="pt" + ) + inputs = processed["input_features"].to(torch.float32) + inp["encoder_inputs"] = inputs + _, cross_cache = encoder(inputs) + + inp["encoded_audio"] = processed + inp["cross_cache"] = cross_cache + + return inp + + +def prepare_encoder_inputs(inp): + inp["hidden_states"] = inp["encoder_inputs"] + inp.pop("encoder_inputs") + return inp + + +def add_instruction(inp, tokenizer): + prompt = "" + inp_encoded = tokenizer(prompt, return_tensors="pt") + inp_encoded.pop("attention_mask") + inp_encoded = inp_encoded["input_ids"] + inp_encoded = inp_encoded.to(torch.int32) + inp["input_ids"] = inp_encoded + return inp + + +def reset_cache( + num_chunks, num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size +): + cache = [] + for i in range(num_chunks): + curr_chunk_cache = torch.zeros( + ( + 2 * num_blocks_per_chunk[i], + num_key_value_heads, + max_cache_size, # generate fixed cache as torch dynamic shape cannot handle 2 dynamic dim + head_dim, + ), + dtype=torch.float32, + ) + cache.append(curr_chunk_cache) + return cache + + +def forward_and_save( + models, + hidden_state, + cache, + cross_cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + batch_name, +): + for chunk_idx in range(len(models)): + cache_in = cache[chunk_idx] + num_layers = int(cross_cache.shape[0] / 2) + num_blocks = num_blocks_per_chunk[chunk_idx] + cross_key = cross_cache[ + chunk_idx * num_blocks : (chunk_idx + 1) * num_blocks, ... + ] + cross_value = cross_cache[ + chunk_idx * num_blocks + + num_layers : (chunk_idx + 1) * num_blocks + + num_layers, + ..., + ] + cross_cache_in = torch.cat([cross_key, cross_value], dim=0) + + try: + model_input_dict[str(chunk_idx)] = { + **model_input_dict[str(chunk_idx)], + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + "cross_cache": cross_cache_in, + }, + } + except: + model_input_dict[str(chunk_idx)] = { + batch_name: { + "hidden_state": hidden_state, + "mask": mask, + "pos_emb": pos_emb, + "cache": cache_in, + "cross_cache": cross_cache_in, + } + } + with torch.no_grad(): + model_out = models[chunk_idx]( + hidden_state, + mask, + pos_emb, + cross_cache_in, + *torch.split(cache_in, 1, dim=0), + ) + hidden_state = model_out[0] + cache[chunk_idx] = torch.cat( + model_out[1 : 1 + 2 * num_blocks_per_chunk[chunk_idx]], dim=0 + ).clone() + return hidden_state, cache + + +def prepare_model_inputs( + inp, + models, + embedding_layer, + master_pos_emb, + num_blocks_per_chunk, + num_key_value_heads, + head_dim, + max_cache_size, + eos_token_id_tensor, + response_cap, +): + model_input_dict = {str(i): None for i in range(len(models))} + input_ids = inp.pop("input_ids") + hidden_state = embedding_layer(torch.tensor(input_ids)) + input_length = hidden_state.shape[1] + # Assume fixed cache size + mask = generate_mask(max_cache_size, 0, input_length, input_length) + pos_emb = master_pos_emb[:, :input_length, :] + cross_cache = inp.pop("cross_cache") + cross_cache = torch.tensor(cross_cache) + # cache shape: num chunks of 2*num_block, num kv heads, c, head dim + cache = reset_cache( + len(models), num_key_value_heads, num_blocks_per_chunk, head_dim, max_cache_size + ) + + logits, cache = forward_and_save( + models, + hidden_state, + cache, + cross_cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + "prompt", + ) + next_token_logits = logits[:, -1, :] # last layer logits + next_token = torch.argmax(next_token_logits, dim=-1) + + response_count = 0 + seq_length = input_length + while True: + curr_input_id = next_token[:, None].to(torch.int32) + input_length = curr_input_id.shape[1] + hidden_state = embedding_layer(curr_input_id) + mask = generate_mask(max_cache_size, seq_length, input_length, input_length) + pos_emb = master_pos_emb[:, seq_length : seq_length + input_length, :] + + logits, cache = forward_and_save( + models, + hidden_state, + cache, + cross_cache, + mask, + pos_emb, + model_input_dict, + num_blocks_per_chunk, + f"response{response_count}", + ) + next_token_logits = logits[:, -1, :] + next_token = torch.argmax(next_token_logits, dim=-1) + + if next_token == eos_token_id_tensor: + print(f"Found EOS on batch: {response_count}") + break + + response_count += 1 + seq_length += input_length + if response_count == response_cap: + break + + return model_input_dict + + +def calibrate_model(model, cal_dataset, chunk_idx: str): + with torch.no_grad(): + for inp in tqdm(cal_dataset, desc="Calibrating Model: "): + # pass prompt and response + for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "): + if inp[chunk_idx][batch] is not None: + inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"]) + mask = torch.tensor(inp[chunk_idx][batch]["mask"]) + pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"]) + cache = torch.tensor(inp[chunk_idx][batch]["cache"]) + cross_cache = torch.tensor(inp[chunk_idx][batch]["cross_cache"]) + model( + inputs_embeds, + mask, + pos_emb, + cross_cache, + *torch.split(cache, 1, dim=0), + ) + + +def export_to_et_ir( + output_folder, + exp_name, + model, + precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset=None, +): + print(f"Exporting Chunk {chunk_idx} to PTE") + example_inputs, dynamic_shapes = model.get_example_inputs( + max_num_token, max_cache_size, True + ) + print("Getting pre autograd ATen Dialect Graph") + pre_autograd_aten_dialect = torch.export.export_for_training( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True + ).module() # NOTE: Will be replaced with export + quantizer = NeuropilotQuantizer() + quantizer.setup_precision(getattr(Precision, precision)) + prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) + # at this point quant min max are inf + if cal_dataset is not None: + calibrate_model(prepared_graph, cal_dataset, str(chunk_idx)) + else: + prepared_graph(*example_inputs) # dummy calibration + converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) + + print("Getting ATen Dialect Graph") + # Fixed Shape Export Here + for shape, ntok_and_cache in export_shapes.items(): + dest_path = get_dest_path(output_folder, exp_name, shape, chunk_idx) + print(f"Exporting Shape {shape} to:\n{dest_path}") + example_inputs = model.get_example_inputs(*ntok_and_cache) + aten_dialect: exir.ExportedProgram = torch.export.export( + converted_graph, example_inputs, strict=True + ) + + print("Lowering to Edge Dialect Graph") + edge_program: exir.EdgeProgramManager = exir.to_edge( + aten_dialect, + compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), + ) + del aten_dialect + + print("Delegating Edge Program to Neuropilot Backend") + compile_spec = [ + CompileSpec("gno", b"LTS"), + CompileSpec("gno-exp", b""), + CompileSpec("gno-non-4d-tiling", b""), + CompileSpec("ImportForever", struct.pack("?", True)), + CompileSpec("platform-config", platform_b), + ] + partitioner = NeuropilotPartitioner(compile_spec) + delegated_program = edge_program.to_backend(partitioner) + print("Exported Delegated Program:") + print(delegated_program.exported_program()) + del edge_program + + print("Transforming delegated program to executorch backend") + executorch_program = delegated_program.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, + ) + ) + + print(f"ET Model Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) + + +def export_encoder_to_et_ir( + enc_output_folder, + enc_exp_name, + model, + precision, + num_mel_bins, + platform_b, + cal_dataset=None, +): + print(f"Exporting Encoder to PTE") + example_inputs = model.get_example_inputs(num_mel_bins) + print("Getting pre autograd ATen Dialect Graph") + pre_autograd_aten_dialect = torch.export.export_for_training( + model, example_inputs, strict=True + ).module() # NOTE: Will be replaced with export + quantizer = NeuropilotQuantizer() + quantizer.setup_precision(getattr(Precision, precision)) + prepared_graph = prepare_pt2e(pre_autograd_aten_dialect, quantizer) + # at this point quant min max are inf + if cal_dataset is not None: + with torch.no_grad(): + for inp in tqdm(cal_dataset, desc="Calibrating Model: "): + prepared_graph(torch.tensor(inp["hidden_states"])) + else: + prepared_graph(*example_inputs) # dummy calibration + converted_graph = convert_pt2e(prepared_graph, fold_quantize=False) + + print("Getting ATen Dialect Graph") + # Fixed Shape Export Here + + file_name = f"{enc_exp_name}.pte" + dest_path = os.path.join(enc_output_folder, file_name) + print(f"Exporting encoder to:\n{dest_path}") + example_inputs = model.get_example_inputs(num_mel_bins) + aten_dialect: exir.ExportedProgram = torch.export.export( + converted_graph, example_inputs, strict=True + ) + + print("Lowering to Edge Dialect Graph") + edge_program: exir.EdgeProgramManager = exir.to_edge( + aten_dialect, + compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), + ) + del aten_dialect + + print("Delegating Edge Program to Neuropilot Backend") + compile_spec = [ + CompileSpec("gno", b"LTS"), + CompileSpec("gno-exp", b""), + CompileSpec("gno-non-4d-tiling", b""), + CompileSpec("ImportForever", struct.pack("?", True)), + CompileSpec("platform-config", platform_b), + ] + partitioner = NeuropilotPartitioner(compile_spec) + delegated_program = edge_program.to_backend(partitioner) + print("Exported Delegated Program:") + print(delegated_program.exported_program()) + del edge_program + + print("Transforming delegated program to executorch backend") + executorch_program = delegated_program.to_executorch( + config=exir.ExecutorchBackendConfig( + memory_planning_pass=exir.passes.MemoryPlanningPass( + alloc_graph_input=False, + alloc_graph_output=False, + ), + extract_delegate_segments=True, + ) + ) + + print(f"ET Model Dest: {dest_path}\n") + os.makedirs(dest_path.rsplit("/", 1)[0], exist_ok=True) + with open(dest_path, "wb") as file: + file.write(executorch_program.buffer) + + +def main(): + parser = get_argument_parser() + args = parser.parse_args() + args_sanity_checks(args) + if args.dataset is None: + exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" + else: + exp_name = ( + f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" + ) + if args.platform == "DX4": + platform_b = b"mt6991" + elif args.platform == "DX3": + platform_b = b"mt6989" + else: + raise ValueError( + f"Platform should be either DX3 or DX4, but got {args.platform}" + ) + print_args(args, exp_name) + enc_exp_name = f"{get_exp_name(args.config)}_{args.precision}_encoder" + + config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( + args.config + ) + encoder_class = chunk_class[0] + decoder_class = chunk_class[1] + tokenizer = tokenizer_class.from_pretrained( + weight_dir, language="english", task="transcribe" + ) + + preprocessor_attr = config.p.__dict__ + preprocessor = WhisperAudioProcessor(**preprocessor_attr) + + head_dim = config.llm.head_dim + + # Evenly distribute the layers across chunks. + num_blocks_per_chunk = [ + (config.llm.num_hidden_layers // args.num_chunks) + + (i < (config.llm.num_hidden_layers % args.num_chunks)) + for i in range(args.num_chunks) + ] + check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 + + output_folder = os.path.join("pte", exp_name) + enc_output_folder = os.path.join("pte", enc_exp_name) + + # Load all collected checkpoint files into one giant state_dict + state_dict = load_checkpoints(weight_dir) + + dump_embedding_lut_for_cmdline(weight_dir, state_dict, config.llm) + + export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) + print(f"export shapes: {export_shapes}") + print(f"Max Num Token: {max_num_token}") + print(f"Max Cache Size: {max_cache_size}") + + embedding_layer = get_embedding_layer(config.llm, weight_dir, state_dict) + + # Instantiate model chunks + print("Instantiating submodels") + + encoder = encoder_class(config.encoder) + encoder = encoder.load_weights(state_dict) + + models = [] + for chunk_idx, num_blocks in enumerate(num_blocks_per_chunk): + decoder_chunk = decoder_class( + config.llm, + num_blocks, + chunk_idx=chunk_idx, + dtype=torch.float32, + include_tail=(chunk_idx == args.num_chunks - 1), + jit_trace=True, + ) + decoder_chunk = decoder_chunk.load_weights( + state_dict, sum(num_blocks_per_chunk[:chunk_idx]) + ) + models.append(decoder_chunk) + + cal_dataset = None + if args.dataset is not None: + cal_dataset = load_dataset("text", data_files=args.dataset, split="train") + master_pos_emb = get_master_pos_emb(config.llm, weight_dir, dtype=torch.float32) + cal_dataset = cal_dataset.map( + process_audio, + fn_kwargs={ + "preprocessor": preprocessor, + "encoder": encoder, + }, + ) + + encoder_cal_dataset = cal_dataset.map(prepare_encoder_inputs) + + cal_dataset = cal_dataset.map( # keys: encoded_audio, encoder_inputs, cross_cache, input_ids + add_instruction, fn_kwargs={"tokenizer": tokenizer} + ) + + print("Preparing Model Calibration Inputs...") + cal_dataset = cal_dataset.map( + prepare_model_inputs, + fn_kwargs={ + "models": models, + "embedding_layer": embedding_layer, + "master_pos_emb": master_pos_emb, + "num_blocks_per_chunk": num_blocks_per_chunk, + "num_key_value_heads": config.llm.num_key_value_heads, + "head_dim": head_dim, + "max_cache_size": max_cache_size, + "eos_token_id_tensor": torch.tensor(tokenizer.eos_token_id), + "response_cap": args.response_cap, + }, + ) + + export_encoder_to_et_ir( + enc_output_folder, + enc_exp_name, + encoder, + args.precision, + config.encoder.num_mel_bins, + platform_b, + encoder_cal_dataset, + ) + + for chunk_idx, chunk in enumerate(models): + export_to_et_ir( + output_folder, + exp_name, + chunk, + args.precision, + max_num_token, + max_cache_size, + chunk_idx, + export_shapes, + platform_b, + cal_dataset, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/mediatek/models/llm_models/configuration_base.py b/examples/mediatek/models/llm_models/configuration_base.py index 354f733cc17..81fd6166c71 100644 --- a/examples/mediatek/models/llm_models/configuration_base.py +++ b/examples/mediatek/models/llm_models/configuration_base.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + from abc import ABC, abstractmethod @@ -24,6 +30,7 @@ def __init__(self): self.use_stable_embedding = False self.tie_word_embeddings = False self.combine_qkv = False + self.use_qk_norm = False self.tokenizer = "default" diff --git a/examples/mediatek/models/llm_models/configuration_gemma.py b/examples/mediatek/models/llm_models/configuration_gemma.py new file mode 100644 index 00000000000..b8a24995933 --- /dev/null +++ b/examples/mediatek/models/llm_models/configuration_gemma.py @@ -0,0 +1,151 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +from models.llm_models.configuration_base import BaseConfig + + +# flake8: noqa: C901 + + +class GemmaConfig(BaseConfig): + def __init__( + self, + vocab_size=None, + hidden_size=None, + intermediate_size=None, + num_hidden_layers=None, + num_attention_heads=None, + max_position_embeddings=None, + norm="RMSNorm", + position_embedding="rope", + norm_eps=1e-6, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + unk_token_id=0, + use_stable_embedding=False, + tie_word_embeddings=True, + combine_qkv=False, + response_handler=None, + model_type=None, + **kwargs, + ): + super().__init__() + + self.model_type = model_type + self.vocab_size = vocab_size + if self.vocab_size is None: + raise KeyError("vocab_size is required but missing from config.json") + self.hidden_size = hidden_size + if self.hidden_size is None: + raise KeyError("hidden_size is required but missing from config.json") + self.intermediate_size = intermediate_size + if self.intermediate_size is None: + raise KeyError("intermediate_size is required but missing from config.json") + self.num_hidden_layers = num_hidden_layers + if self.num_hidden_layers is None: + raise KeyError("num_hidden_layers is required but missing from config.json") + self.num_attention_heads = num_attention_heads + if self.num_attention_heads is None: + raise KeyError( + "num_attention_heads is required but missing from config.json" + ) + self.num_key_value_heads = kwargs.pop( + "num_key_value_heads", self.num_attention_heads + ) + if self.num_attention_heads % self.num_key_value_heads != 0: + raise RuntimeError( + f"num_attention_heads ({self.num_attention_heads}) must be exactly " + f"divisible by num_key_value_heads ({self.num_key_value_heads})" + ) + if norm not in ["RMSNorm", "LayerNorm"]: + raise ValueError("norm must be one of: RMSNorm (default) or LayerNorm") + self.norm = norm + self.norm_eps = kwargs.pop("rms_norm_eps", norm_eps) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.unk_token_id = unk_token_id + self.head_dim = kwargs.pop( + "head_dim", self.hidden_size // self.num_attention_heads + ) + + if position_embedding not in ["rope", "alibi"]: + raise ValueError("Positional embedding must be one of: rope, alibi") + self.position_embedding = position_embedding + self.ntk_scaling_factor = kwargs.pop("ntk_scaling_factor", 1.0) + if self.ntk_scaling_factor != 1.0 and self.position_embedding != "rope": + raise KeyError("ntk_scaling_factor is strictly for position_embedding=rope") + self.max_position_embeddings = max_position_embeddings + if self.max_position_embeddings is None and self.position_embedding == "rope": + raise KeyError( + "max_position_embeddings is required for position_embedding=rope but missing from config.json" + ) + + self.use_stable_embedding = use_stable_embedding + self.tie_word_embeddings = tie_word_embeddings + self.combine_qkv = combine_qkv + + self.tokenizer = kwargs.pop("tokenizer", self.tokenizer) + + if self.model_type == "gemma2": + self.attn_logit_softcapping = kwargs.pop("attn_logit_softcapping", 50.0) + self.final_logit_softcapping = kwargs.pop("final_logit_softcapping", 30.0) + self.query_pre_attn_scalar = kwargs.pop("query_pre_attn_scalar", 256) + elif self.model_type == "gemma3": + self.sliding_window_attention_size = kwargs.pop( + "sliding_window_attention_size", 1024 + ) + self.global_local_attention_pattern = kwargs.pop( + "global_local_attention_pattern", None + ) + self.sliding_window_pattern = kwargs.pop("sliding_window_pattern", None) + self.rope_local_base_freq = kwargs.pop("rope_local_base_freq", 10000) + self.query_pre_attn_scalar = kwargs.pop("query_pre_attn_scalar", 256) + self.use_vision = kwargs.pop("use_vision", False) + self.use_qk_norm = kwargs.pop("use_qk_norm", True) + self.use_res_clamp = kwargs.pop("use_res_clamp", True) + self.rope_theta = kwargs.pop("rope_theta", 1000000) + self.rope_scaling = kwargs.pop("rope_scaling", None) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} config:") + print(f"Hidden size: {self.hidden_size}") + print(f"Intermediate size: {self.intermediate_size}") + print(f"Num layers: {self.num_hidden_layers}") + print(f"Num attention heads: {self.num_attention_heads}") + print(f"Num KV heads: {self.num_key_value_heads}") + print(f"Head Dim: {self.head_dim}") + if self.model_type == "gemma3": + print(f"Size sliding window: {self.sliding_window_attention_size}") + print(f"Sliding window pattern: {self.sliding_window_pattern}") + print(f"Positional embedding: {self.position_embedding}") + if self.position_embedding == "rope": + print(f"Max pos emb: {self.max_position_embeddings}") + if self.ntk_scaling_factor != 1.0: + print(f"NTK scaling factor: {self.ntk_scaling_factor}") + print(f"Norm type: {self.norm}") + print(f"Norm epsilon: {self.norm_eps}") + print(f"BOS token id: {self.bos_token_id}") + print(f"EOS token id: {self.eos_token_id}") + print(f"PAD token id: {self.pad_token_id}") + print(f"UNK token id: {self.unk_token_id}") + print(f"Vocab size: {self.vocab_size}") + print(f"Use stable embedding: {self.use_stable_embedding}") + print(f"Use qk norm: {self.use_qk_norm}") + print(f"Tie word embeddings: {self.tie_word_embeddings}") + print(f"Combine QKV: {self.combine_qkv}") + if self.tokenizer != "default": + print(f"Tokenizer: {self.tokenizer}") + print() diff --git a/examples/mediatek/models/llm_models/configuration_llama.py b/examples/mediatek/models/llm_models/configuration_llama.py index 159caac6624..e9194fc5a18 100644 --- a/examples/mediatek/models/llm_models/configuration_llama.py +++ b/examples/mediatek/models/llm_models/configuration_llama.py @@ -1,22 +1,9 @@ -# coding=utf-8 -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# Copyright (c) MediaTek Inc. +# All rights reserved # -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +# 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 from models.llm_models.configuration_base import BaseConfig @@ -83,6 +70,9 @@ def __init__( self.eos_token_id = eos_token_id self.pad_token_id = pad_token_id self.unk_token_id = unk_token_id + self.head_dim = kwargs.pop( + "head_dim", self.hidden_size // self.num_attention_heads + ) if position_embedding not in ["rope", "alibi"]: raise ValueError("Positional embedding must be one of: rope, alibi") @@ -115,6 +105,7 @@ def print_config(self, response_handler): print(f"Num layers: {self.num_hidden_layers}") print(f"Num attention heads: {self.num_attention_heads}") print(f"Num KV heads: {self.num_key_value_heads}") + print(f"Head Dim: {self.head_dim}") print(f"Positional embedding: {self.position_embedding}") if self.position_embedding == "rope": print(f"Max pos emb: {self.max_position_embeddings}") diff --git a/examples/mediatek/models/llm_models/configuration_phi.py b/examples/mediatek/models/llm_models/configuration_phi.py new file mode 100644 index 00000000000..1dea7b81a21 --- /dev/null +++ b/examples/mediatek/models/llm_models/configuration_phi.py @@ -0,0 +1,133 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +from models.llm_models.configuration_base import BaseConfig + + +# flake8: noqa: C901 + + +class PhiConfig(BaseConfig): + def __init__( + self, + vocab_size=None, + hidden_size=None, + intermediate_size=None, + num_hidden_layers=None, + num_attention_heads=None, + max_position_embeddings=None, + norm="RMSNorm", + position_embedding="rope", + norm_eps=1e-6, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + unk_token_id=0, + use_stable_embedding=False, + tie_word_embeddings=False, + combine_qkv=True, + response_handler=None, + **kwargs, + ): + super().__init__() + + self.model_type = kwargs.pop("model_type", "phi3") + self.vocab_size = vocab_size + if self.vocab_size is None: + raise KeyError("vocab_size is required but missing from config.json") + self.hidden_size = hidden_size + if self.hidden_size is None: + raise KeyError("hidden_size is required but missing from config.json") + self.intermediate_size = intermediate_size + if self.intermediate_size is None: + raise KeyError("intermediate_size is required but missing from config.json") + self.num_hidden_layers = num_hidden_layers + if self.num_hidden_layers is None: + raise KeyError("num_hidden_layers is required but missing from config.json") + self.num_attention_heads = num_attention_heads + if self.num_attention_heads is None: + raise KeyError( + "num_attention_heads is required but missing from config.json" + ) + self.num_key_value_heads = kwargs.pop( + "num_key_value_heads", self.num_attention_heads + ) + if self.num_attention_heads % self.num_key_value_heads != 0: + raise RuntimeError( + f"num_attention_heads ({self.num_attention_heads}) must be exactly " + f"divisible by num_key_value_heads ({self.num_key_value_heads})" + ) + if norm not in ["RMSNorm", "LayerNorm"]: + raise ValueError("norm must be one of: RMSNorm (default) or LayerNorm") + self.norm = norm + self.norm_eps = kwargs.pop("rms_norm_eps", norm_eps) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.unk_token_id = unk_token_id + self.head_dim = kwargs.pop( + "head_dim", self.hidden_size // self.num_attention_heads + ) + + # Rotary embedding + self.partial_rotary_factor = kwargs.pop("partial_rotary_factor", None) + self.rope_scaling = kwargs.pop("rope_scaling", None) + self.original_max_position_embeddings = kwargs.pop( + "original_max_position_embeddings", None + ) + + if position_embedding not in ["rope", "alibi"]: + raise ValueError("Positional embedding must be one of: rope, alibi") + self.position_embedding = position_embedding + self.ntk_scaling_factor = kwargs.pop("ntk_scaling_factor", 1.0) + if self.ntk_scaling_factor != 1.0 and self.position_embedding != "rope": + raise KeyError("ntk_scaling_factor is strictly for position_embedding=rope") + self.max_position_embeddings = max_position_embeddings + if self.max_position_embeddings is None and self.position_embedding == "rope": + raise KeyError( + "max_position_embeddings is required for position_embedding=rope but missing from config.json" + ) + + self.use_stable_embedding = use_stable_embedding + self.tie_word_embeddings = tie_word_embeddings + self.combine_qkv = combine_qkv + + self.tokenizer = kwargs.pop("tokenizer", self.tokenizer) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} config:") + print(f"Hidden size: {self.hidden_size}") + print(f"Intermediate size: {self.intermediate_size}") + print(f"Num layers: {self.num_hidden_layers}") + print(f"Num attention heads: {self.num_attention_heads}") + print(f"Num KV heads: {self.num_key_value_heads}") + print(f"Head Dim: {self.head_dim}") + print(f"Positional embedding: {self.position_embedding}") + if self.position_embedding == "rope": + print(f"Max pos emb: {self.max_position_embeddings}") + if self.ntk_scaling_factor != 1.0: + print(f"NTK scaling factor: {self.ntk_scaling_factor}") + print(f"Norm type: {self.norm}") + print(f"Norm epsilon: {self.norm_eps}") + print(f"BOS token id: {self.bos_token_id}") + print(f"EOS token id: {self.eos_token_id}") + print(f"PAD token id: {self.pad_token_id}") + print(f"UNK token id: {self.unk_token_id}") + print(f"Vocab size: {self.vocab_size}") + print(f"Use stable embedding: {self.use_stable_embedding}") + print(f"Tie word embeddings: {self.tie_word_embeddings}") + print(f"Combine QKV: {self.combine_qkv}") + if self.tokenizer != "default": + print(f"Tokenizer: {self.tokenizer}") + print() diff --git a/examples/mediatek/models/llm_models/configuration_qwen.py b/examples/mediatek/models/llm_models/configuration_qwen.py new file mode 100644 index 00000000000..0e96999b435 --- /dev/null +++ b/examples/mediatek/models/llm_models/configuration_qwen.py @@ -0,0 +1,131 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +from models.llm_models.configuration_base import BaseConfig + + +# flake8: noqa: C901 + + +class QwenConfig(BaseConfig): + def __init__( + self, + vocab_size=None, + hidden_size=None, + intermediate_size=None, + num_hidden_layers=None, + num_attention_heads=None, + max_position_embeddings=None, + norm="RMSNorm", + position_embedding="rope", + norm_eps=1e-6, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + unk_token_id=0, + use_stable_embedding=False, + tie_word_embeddings=False, + combine_qkv=False, + response_handler=None, + model_type=None, + **kwargs, + ): + super().__init__() + + self.model_type = model_type + self.vocab_size = vocab_size + if self.vocab_size is None: + raise KeyError("vocab_size is required but missing from config.json") + self.hidden_size = hidden_size + if self.hidden_size is None: + raise KeyError("hidden_size is required but missing from config.json") + self.intermediate_size = intermediate_size + if self.intermediate_size is None: + raise KeyError("intermediate_size is required but missing from config.json") + self.num_hidden_layers = num_hidden_layers + if self.num_hidden_layers is None: + raise KeyError("num_hidden_layers is required but missing from config.json") + self.num_attention_heads = num_attention_heads + if self.num_attention_heads is None: + raise KeyError( + "num_attention_heads is required but missing from config.json" + ) + self.num_key_value_heads = kwargs.pop( + "num_key_value_heads", self.num_attention_heads + ) + if self.num_attention_heads % self.num_key_value_heads != 0: + raise RuntimeError( + f"num_attention_heads ({self.num_attention_heads}) must be exactly " + f"divisible by num_key_value_heads ({self.num_key_value_heads})" + ) + if norm not in ["RMSNorm", "LayerNorm"]: + raise ValueError("norm must be one of: RMSNorm (default) or LayerNorm") + self.norm = norm + self.norm_eps = kwargs.pop("rms_norm_eps", norm_eps) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.unk_token_id = unk_token_id + self.head_dim = kwargs.pop( + "head_dim", self.hidden_size // self.num_attention_heads + ) + + if position_embedding not in ["rope", "alibi"]: + raise ValueError("Positional embedding must be one of: rope, alibi") + self.position_embedding = position_embedding + self.ntk_scaling_factor = kwargs.pop("ntk_scaling_factor", 1.0) + if self.ntk_scaling_factor != 1.0 and self.position_embedding != "rope": + raise KeyError("ntk_scaling_factor is strictly for position_embedding=rope") + self.max_position_embeddings = max_position_embeddings + if self.max_position_embeddings is None and self.position_embedding == "rope": + raise KeyError( + "max_position_embeddings is required for position_embedding=rope but missing from config.json" + ) + + self.use_stable_embedding = use_stable_embedding + self.tie_word_embeddings = tie_word_embeddings + self.combine_qkv = combine_qkv + + self.tokenizer = kwargs.pop("tokenizer", self.tokenizer) + + if self.model_type == "qwen3": + self.use_qk_norm = kwargs.pop("use_qk_norm", True) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} config:") + print(f"Hidden size: {self.hidden_size}") + print(f"Intermediate size: {self.intermediate_size}") + print(f"Num layers: {self.num_hidden_layers}") + print(f"Num attention heads: {self.num_attention_heads}") + print(f"Num KV heads: {self.num_key_value_heads}") + print(f"Head Dim: {self.head_dim}") + print(f"Positional embedding: {self.position_embedding}") + if self.position_embedding == "rope": + print(f"Max pos emb: {self.max_position_embeddings}") + if self.ntk_scaling_factor != 1.0: + print(f"NTK scaling factor: {self.ntk_scaling_factor}") + print(f"Norm type: {self.norm}") + print(f"Norm epsilon: {self.norm_eps}") + print(f"BOS token id: {self.bos_token_id}") + print(f"EOS token id: {self.eos_token_id}") + print(f"PAD token id: {self.pad_token_id}") + print(f"UNK token id: {self.unk_token_id}") + print(f"Vocab size: {self.vocab_size}") + print(f"Use stable embedding: {self.use_stable_embedding}") + print(f"Use qk norm: {self.use_qk_norm}") + print(f"Tie word embeddings: {self.tie_word_embeddings}") + print(f"Combine QKV: {self.combine_qkv}") + if self.tokenizer != "default": + print(f"Tokenizer: {self.tokenizer}") + print() diff --git a/examples/mediatek/models/llm_models/configuration_whisper.py b/examples/mediatek/models/llm_models/configuration_whisper.py new file mode 100644 index 00000000000..71f6606097e --- /dev/null +++ b/examples/mediatek/models/llm_models/configuration_whisper.py @@ -0,0 +1,270 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +from models.llm_models.configuration_base import BaseConfig + + +# flake8: noqa: C901 + + +class WhisperEncoderConfig(BaseConfig): + def __init__( + self, + model_type="whisper", + encoder_ffn_dim=None, + encoder_layers=None, + encoder_attention_heads=None, + d_model=None, + decoder_layers=None, + num_mel_bins=None, + max_source_positions=None, + max_position_embeddings=None, + head_dim=None, + response_handler=None, + **kwargs, + ): + super().__init__() + self.model_type = model_type + self.intermediate_size = encoder_ffn_dim + if self.intermediate_size is None: + raise KeyError("encoder_ffn_dim is required but missing from config.json") + + self.num_hidden_layers = encoder_layers + if self.num_hidden_layers is None: + raise KeyError("encoder_layers is required but missing from config.json") + + self.num_attention_heads = encoder_attention_heads + if self.num_attention_heads is None: + raise KeyError( + "encoder_attention_heads is required but missing from config.json" + ) + + self.hidden_size = d_model + if self.hidden_size is None: + raise KeyError("d_model is required but missing from config.json") + + self.decoder_num_layers = decoder_layers + if self.decoder_num_layers is None: + raise KeyError("decoder_layers is required but missing from config.json") + + self.num_mel_bins = num_mel_bins + if self.num_mel_bins is None: + raise KeyError("num_mel_bins is required but missing from config.json") + + self.max_source_positions = max_source_positions + if self.max_source_positions is None: + raise KeyError( + "max_source_positions is required but missing from config.json" + ) + + if head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + else: + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.norm = kwargs.pop("norm", "LayerNorm") + self.num_key_value_heads = kwargs.pop( + "num_key_value_heads", self.num_attention_heads + ) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} encoder config:") + print(f"Hidden size: {self.hidden_size}") + print(f"Intermediate size: {self.intermediate_size}") + print(f"Num layers: {self.num_hidden_layers}") + print(f"Num attention heads: {self.num_attention_heads}") + print(f"Head dim: {self.head_dim}") + print(f"Decoder num layers: {self.decoder_num_layers}") + print(f"Max source pos: {self.max_source_positions}") + if self.max_position_embeddings is not None: + print(f"Max pos emb: {self.max_position_embeddings}") + print(f"Num mel bins: {self.num_mel_bins}") + print() + + +class WhisperDecoderConfig(BaseConfig): + def __init__( + self, + vocab_size=None, + hidden_size=None, + intermediate_size=None, + num_hidden_layers=None, + num_attention_heads=None, + max_position_embeddings=None, + norm="LayerNorm", + position_embedding="rope", + norm_eps=1e-6, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + unk_token_id=0, + use_stable_embedding=False, + tie_word_embeddings=True, + combine_qkv=False, + response_handler=None, + model_type=None, + **kwargs, + ): + super().__init__() + + self.model_type = model_type + self.vocab_size = vocab_size + if self.vocab_size is None: + raise KeyError("vocab_size is required but missing from config.json") + self.hidden_size = hidden_size + if self.hidden_size is None: + raise KeyError("hidden_size is required but missing from config.json") + self.intermediate_size = intermediate_size + if self.intermediate_size is None: + raise KeyError("intermediate_size is required but missing from config.json") + self.num_hidden_layers = num_hidden_layers + if self.num_hidden_layers is None: + raise KeyError("num_hidden_layers is required but missing from config.json") + self.num_attention_heads = num_attention_heads + if self.num_attention_heads is None: + raise KeyError( + "num_attention_heads is required but missing from config.json" + ) + self.num_key_value_heads = kwargs.pop( + "num_key_value_heads", self.num_attention_heads + ) + if self.num_attention_heads % self.num_key_value_heads != 0: + raise RuntimeError( + f"num_attention_heads ({self.num_attention_heads}) must be exactly " + f"divisible by num_key_value_heads ({self.num_key_value_heads})" + ) + if norm not in ["RMSNorm", "LayerNorm"]: + raise ValueError("norm must be one of: RMSNorm (default) or LayerNorm") + self.norm = norm + self.norm_eps = kwargs.pop("rms_norm_eps", norm_eps) + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + self.pad_token_id = pad_token_id + self.unk_token_id = unk_token_id + self.head_dim = kwargs.pop( + "head_dim", self.hidden_size // self.num_attention_heads + ) + + if position_embedding not in ["rope", "alibi"]: + raise ValueError("Positional embedding must be one of: rope, alibi") + self.position_embedding = position_embedding + self.ntk_scaling_factor = kwargs.pop("ntk_scaling_factor", 1.0) + if self.ntk_scaling_factor != 1.0 and self.position_embedding != "rope": + raise KeyError("ntk_scaling_factor is strictly for position_embedding=rope") + self.max_position_embeddings = max_position_embeddings + if self.max_position_embeddings is None and self.position_embedding == "rope": + raise KeyError( + "max_position_embeddings is required for position_embedding=rope but missing from config.json" + ) + + self.use_stable_embedding = use_stable_embedding + self.tie_word_embeddings = tie_word_embeddings + self.combine_qkv = combine_qkv + + self.tokenizer = kwargs.pop("tokenizer", self.tokenizer) + self.max_source_positions = kwargs.pop("max_source_positions", None) + if self.max_source_positions is None: + raise KeyError( + "max_source_positions is required but missing from config.json" + ) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} llm config:") + print(f"Hidden size: {self.hidden_size}") + print(f"Intermediate size: {self.intermediate_size}") + print(f"Num layers: {self.num_hidden_layers}") + print(f"Num attention heads: {self.num_attention_heads}") + print(f"Num KV heads: {self.num_key_value_heads}") + print(f"Head Dim: {self.head_dim}") + print(f"Positional embedding: {self.position_embedding}") + if self.position_embedding == "rope": + print(f"Max pos emb: {self.max_position_embeddings}") + if self.ntk_scaling_factor != 1.0: + print(f"NTK scaling factor: {self.ntk_scaling_factor}") + print(f"Norm type: {self.norm}") + print(f"Norm epsilon: {self.norm_eps}") + print(f"BOS token id: {self.bos_token_id}") + print(f"EOS token id: {self.eos_token_id}") + print(f"PAD token id: {self.pad_token_id}") + print(f"UNK token id: {self.unk_token_id}") + print(f"Vocab size: {self.vocab_size}") + print(f"Use stable embedding: {self.use_stable_embedding}") + print(f"Use qk norm: {self.use_qk_norm}") + print(f"Tie word embeddings: {self.tie_word_embeddings}") + print(f"Combine QKV: {self.combine_qkv}") + if self.tokenizer != "default": + print(f"Tokenizer: {self.tokenizer}") + print() + + +class WhisperPreprocessorConfig: + def __init__(self, response_handler=None, **kwargs): + self.model_type = kwargs.pop("model_type", "whisper") + assert self.model_type == ( + "whisper" + ), f"Expected model_type to be whisper but got {self.model_type} instead" + + self.feature_size = kwargs.pop("feature_size", 80) + self.sampling_rate = kwargs.pop("sampling_rate", 16000) + self.return_attention_mask = kwargs.pop("return_attention_mask", False) + + if response_handler is None: + response_handler = nullcontext() + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + with response_handler: + print(f"{self.model_type} preprocessor config:") + print(f"Feature Size: {self.feature_size}") + print(f"Sampling Rate: {self.sampling_rate}") + print(f"Audio Attn Mask: {self.return_attention_mask}") + print() + + +class WhisperConfig(BaseConfig): + def __init__(self, **kwargs): + encoder_config = kwargs.pop("encoder", None) + if encoder_config is None: + raise KeyError("Encoder config not found in config.json") + llm_config = kwargs.pop("llm", None) + if llm_config is None: + raise KeyError("llm config not found in config.json") + p_config = kwargs.pop("preprocessor", None) + if p_config is None: + raise KeyError("preprocessor config not found in config.json") + + self.encoder = WhisperEncoderConfig(**encoder_config, verbose=False) + self.llm = WhisperDecoderConfig(**llm_config, verbose=False) + self.p = WhisperPreprocessorConfig(**p_config, verbose=False) + self.model_type = "whisper" + self.tokenizer = "default" + + response_handler = kwargs.pop("response_handler", None) + if response_handler is None: + response_handler = nullcontext() + + if kwargs.pop("verbose", True): + self.print_config(response_handler) + + def print_config(self, response_handler): + self.encoder.print_config(response_handler) + self.llm.print_config(response_handler) + self.p.print_config(response_handler) diff --git a/examples/mediatek/models/llm_models/modeling_base.py b/examples/mediatek/models/llm_models/modeling_base.py index ee6b419b931..9855ecb1c31 100644 --- a/examples/mediatek/models/llm_models/modeling_base.py +++ b/examples/mediatek/models/llm_models/modeling_base.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + from abc import ABC, abstractmethod import torch diff --git a/examples/mediatek/models/llm_models/modeling_common.py b/examples/mediatek/models/llm_models/modeling_common.py index bf1f50f58ef..bfbe925257e 100644 --- a/examples/mediatek/models/llm_models/modeling_common.py +++ b/examples/mediatek/models/llm_models/modeling_common.py @@ -1,3 +1,9 @@ +# Copyright (c) MediaTek 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. + """Common backbone across multiple models""" import math @@ -36,6 +42,23 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x * torch.sigmoid(1.702 * x) +class FastGelu(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.gelu(x, approximate="tanh") + + +class TorchGelu(nn.Module): + def __init__(self): + super().__init__() + self.gelu = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.gelu(x) + + class MLP(nn.Module): def __init__(self, config: BaseConfig): super().__init__() @@ -56,15 +79,18 @@ def forward(self, x): class Attention(nn.Module): - def __init__(self, config: BaseConfig): + def __init__(self, config: BaseConfig, jit_trace=False): super().__init__() self.config = config self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads - self.head_dim = self.hidden_size // self.num_heads + self.head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) self.attn_scale = math.sqrt(self.head_dim) + self.jit_trace = jit_trace if config.combine_qkv: self.qkv_proj = nn.Linear( @@ -81,6 +107,10 @@ def __init__(self, config: BaseConfig): ) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size) + if config.use_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.norm_eps) + def apply_rotary_pos_emb_mtk(self, q, k, cos, sin): q1 = q[..., : q.shape[-1] // 2] q2 = q[..., q.shape[-1] // 2 :] @@ -105,6 +135,20 @@ def repeat_kv(self, hidden_states, batch, q_len, n_rep): hidden_states = hidden_states.repeat(1, 1, n_rep, 1) return hidden_states.view(batch, self.num_heads, q_len, self.head_dim) + def apply_k_norm(self, k): + if self.jit_trace: + return self.k_norm(k) + else: + dtype = k.dtype + return self.k_norm(k.to(torch.float32)).to(dtype) + + def apply_q_norm(self, q): + if self.jit_trace: + return self.q_norm(q) + else: + dtype = q.dtype + return self.q_norm(q.to(torch.float32)).to(dtype) + def forward( self, hidden_states, # (b, t, 4096) @@ -160,6 +204,10 @@ def forward( .transpose(1, 2) ) + if self.config.use_qk_norm: + query_states = self.apply_q_norm(query_states) + key_states = self.apply_k_norm(key_states) + if self.config.position_embedding == "rope": cos, sin = torch.split(pos_emb, 1, dim=1) query_states, key_states = self.apply_rotary_pos_emb_mtk( @@ -185,7 +233,7 @@ def forward( attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2) - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads) attn_output = self.o_proj(attn_output) key_states_out = key_states_out[:, :, q_len:, :] @@ -207,7 +255,7 @@ def __init__( self.hidden_size = config.hidden_size self.return_attn = return_attn self.jit_trace = jit_trace - self.self_attn = attn_class(config) + self.self_attn = attn_class(config, jit_trace=jit_trace) self.mlp = mlp_class(config) if config.norm == "RMSNorm": self.input_norm = RMSNorm(config.hidden_size, eps=config.norm_eps).float() @@ -281,7 +329,9 @@ def __init__( super().__init__( config, num_blocks, chunk_idx, dtype, include_tail, return_attn, jit_trace ) - self.head_dim = config.hidden_size // config.num_attention_heads + self.head_dim = getattr( + self.config, "head_dim", config.hidden_size // config.num_attention_heads + ) self.layers = nn.ModuleList( [ decoder_class(config, return_attn=return_attn, jit_trace=jit_trace) @@ -364,6 +414,7 @@ def load_weights(self, state_dict, state_dict_start_idx): temp_key = None input_norm_subkey = None post_attention_norm_subkey = None + q_norm_subkey = None for key in state_dict_keys: if expected_subkey in key: temp_key = key @@ -379,6 +430,8 @@ def load_weights(self, state_dict, state_dict_start_idx): and "post_attention" in key ): post_attention_norm_subkey = key.split(".")[-2] + if f"layers.{state_dict_start_idx}" in key and "q_norm" in key: + q_norm_subkey = key if temp_key is None: raise KeyError( f"Cannot find layer {state_dict_start_idx}'s o_proj weight inside state_dict. " @@ -393,9 +446,18 @@ def load_weights(self, state_dict, state_dict_start_idx): if post_attention_norm_subkey is None: raise KeyError( f"Cannot find layer {state_dict_start_idx}'s post attention norm weight inside state_dict." - f" Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " + f"Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " "post_attention inside the key string." ) + if q_norm_subkey is not None and self.config.use_qk_norm is False: + model_name = "".join( + [char for char in self.config.model_type if not char.isdigit()] + ) + raise KeyError( + f"Layer {state_dict_start_idx}'s q_norm weight detected inside state_dict but not used." + f"Please set use_qk_norm to True in configuration_{model_name}.py" + ) + prefix = temp_key.split(expected_subkey)[0] fake_weights = False @@ -461,7 +523,8 @@ def load_weights(self, state_dict, state_dict_start_idx): dtype=self.dtype, ), f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( - self.config.hidden_size, dtype=self.dtype + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, ), f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( self.config.num_key_value_heads * self.head_dim, @@ -517,6 +580,19 @@ def load_weights(self, state_dict, state_dict_start_idx): }, } + if self.config.use_qk_norm: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_norm.weight": torch.rand( + self.config.head_dim, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.self_attn.k_norm.weight": torch.rand( + self.head_dim, dtype=self.dtype + ), + }, + } + if self.config.norm == "LayerNorm": temp_state_dict = { **temp_state_dict, @@ -562,7 +638,11 @@ def load_weights(self, state_dict, state_dict_start_idx): ), f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", - torch.zeros(self.config.hidden_size, dtype=self.dtype), + torch.zeros( + self.config.head_dim + * self.config.num_attention_heads, + dtype=self.dtype, + ), ), f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", @@ -627,6 +707,19 @@ def load_weights(self, state_dict, state_dict_start_idx): }, } + if self.config.use_qk_norm: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_norm.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_norm.weight" + ), + }, + } + if self.config.norm == "LayerNorm": temp_state_dict = { **temp_state_dict, @@ -725,7 +818,7 @@ def load_weights(self, state_dict, state_dict_start_idx): def get_example_inputs( self, num_token: int = 128, cache_size: int = 512, get_dym_shape=False ): - head_dim = int(self.config.hidden_size / self.config.num_attention_heads) + head_dim = int(self.head_dim) example_inputs = ( torch.randn( 1, num_token, self.config.hidden_size, device="cpu", dtype=torch.float32 diff --git a/examples/mediatek/models/llm_models/modeling_gemma.py b/examples/mediatek/models/llm_models/modeling_gemma.py new file mode 100644 index 00000000000..e1ee4d70584 --- /dev/null +++ b/examples/mediatek/models/llm_models/modeling_gemma.py @@ -0,0 +1,1570 @@ +# Copyright (c) MediaTek 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 math + +import numpy as np +import torch +from models.llm_models.configuration_gemma import GemmaConfig + +from models.llm_models.modeling_common import ( + Attention, + DecoderLayer, + FastGelu, + MLP, + ModelChunk, +) +from torch import nn +from torch.export import Dim + +np.random.seed(42) + +# flake8: noqa: C901 + + +def scaled_tanh(hidden_states, scale): + output = torch.tanh(hidden_states / scale) + return scale * output + + +class GemmaRMSNorm(nn.Module): + def __init__( + self, + hidden_size, + eps=1e-6, + with_scale=True, + decompose=False, + ): + super().__init__() + self.with_scale = with_scale + if self.with_scale: + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.decompose = decompose + + def forward(self, hidden_states): + variance = hidden_states.pow(2).mean(-1, keepdim=True) + if self.decompose: + hidden_states = hidden_states / torch.sqrt(variance + self.variance_epsilon) + else: + hidden_states = hidden_states * torch.rsqrt( + variance + self.variance_epsilon + ) + + if self.with_scale: + return (1.0 + self.weight) * hidden_states + return hidden_states + + +class GemmaMLP(MLP): + def __init__(self, config: GemmaConfig): + super().__init__(config) + self.gelu = FastGelu() + + def forward(self, x): + gate = self.gate_proj(x) + up = self.up_proj(x) + pre_down = self.gelu(gate) * up + down = self.down_proj(pre_down) + + return down + + +class GemmaAttention(Attention): + def __init__( + self, + config: GemmaConfig, + jit_trace=False, + ): + super().__init__( + config, + jit_trace, + ) + + +class Gemma2Attention(Attention): + def __init__( + self, + config: GemmaConfig, + jit_trace=False, + ): + super().__init__(config, jit_trace) + self.query_pre_attn_scalar = torch.tensor( + self.config.query_pre_attn_scalar**0.5 + ) + + def forward( + self, + hidden_states, # (b, t, 4096) + mask, # (b, 1, t, c+t) + pos_emb, # (b, 2, t, head dim) + past_key, # (b, num kv heads, c, head dim) + past_value, # (b, num kv heads, c, head dim) + ): + bsz, q_len, _ = hidden_states.size() + c_len = past_key.size()[2] + + if self.config.combine_qkv: + proj = self.qkv_proj(hidden_states) + query_states = ( + proj[:, :, : self.config.hidden_size] + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + proj[ + :, + :, + self.config.hidden_size : self.config.hidden_size + + self.num_key_value_heads * self.head_dim, + ] + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + proj[ + :, + :, + self.config.hidden_size + + self.num_key_value_heads * self.head_dim :, + ] + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + else: + query_states = ( + self.q_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + + if self.config.position_embedding == "rope": + cos, sin = torch.split(pos_emb, 1, dim=1) + query_states, key_states = self.apply_rotary_pos_emb_mtk( + query_states, key_states, cos, sin + ) + + key_states = torch.cat([past_key, key_states], dim=2) + value_states = torch.cat([past_value, value_states], dim=2) + key_states_out = key_states + value_states_out = value_states + if self.num_key_value_groups > 1: + key_states = self.repeat_kv( + key_states, bsz, q_len + c_len, self.num_key_value_groups + ) + value_states = self.repeat_kv( + value_states, bsz, q_len + c_len, self.num_key_value_groups + ) + # NOTE: gemma2 scales with query pre attn scalar instead of attn scale + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) + attn_weights = attn_weights / self.query_pre_attn_scalar + + # NOTE: Unique to gemma 2 + attn_weights = scaled_tanh(attn_weights, self.config.attn_logit_softcapping) + + attn_weights = attn_weights + mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads) + attn_output = self.o_proj(attn_output) + + key_states_out = key_states_out[:, :, q_len:, :] + value_states_out = value_states_out[:, :, q_len:, :] + + return attn_output, key_states_out, value_states_out + + +class Gemma3Attention(Attention): + def __init__(self, config: GemmaConfig, jit_trace=False): + super().__init__(config, jit_trace) + self.query_pre_attn_scalar = torch.tensor( + self.config.query_pre_attn_scalar**0.5 + ) + if config.use_qk_norm: + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.norm_eps) + + def forward( + self, + hidden_states, # (b, t, 4096) + mask, # (b, 1, t, c+t) + pos_emb, # (b, 2, t, head dim) + past_key, # (b, num kv heads, c, head dim) + past_value, # (b, num kv heads, c, head dim) + ): + bsz, q_len, _ = hidden_states.size() + c_len = past_key.size()[2] + + if self.config.combine_qkv: + proj = self.qkv_proj(hidden_states) + query_states = ( + proj[:, :, : self.config.hidden_size] + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + proj[ + :, + :, + self.config.hidden_size : self.config.hidden_size + + self.num_key_value_heads * self.head_dim, + ] + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + proj[ + :, + :, + self.config.hidden_size + + self.num_key_value_heads * self.head_dim :, + ] + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + else: + query_states = ( + self.q_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + + # NOTE: gemma3 applies qk_norm before rot emb + if self.config.use_qk_norm: + query_states = self.apply_q_norm(query_states) + key_states = self.apply_k_norm(key_states) + + if self.config.position_embedding == "rope": + cos, sin = torch.split(pos_emb, 1, dim=1) + query_states, key_states = self.apply_rotary_pos_emb_mtk( + query_states, key_states, cos, sin + ) + + key_states = torch.cat([past_key, key_states], dim=2) + value_states = torch.cat([past_value, value_states], dim=2) + key_states_out = key_states + value_states_out = value_states + if self.num_key_value_groups > 1: + key_states = self.repeat_kv( + key_states, bsz, q_len + c_len, self.num_key_value_groups + ) + value_states = self.repeat_kv( + value_states, bsz, q_len + c_len, self.num_key_value_groups + ) + # NOTE: gemma3 scales with query pre attn scalar instead of attn scale + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) + attn_weights = attn_weights / self.query_pre_attn_scalar + + attn_weights = attn_weights + mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads) + attn_output = self.o_proj(attn_output) + + key_states_out = key_states_out[:, :, q_len:, :] + value_states_out = value_states_out[:, :, q_len:, :] + + return attn_output, key_states_out, value_states_out + + +class GemmaDecoderLayer(DecoderLayer): + def __init__( + self, + config: GemmaConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=GemmaAttention, + mlp_class=GemmaMLP, + ) + + +class Gemma2DecoderLayer(DecoderLayer): + def __init__( + self, + config: GemmaConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=Gemma2Attention, + mlp_class=GemmaMLP, + ) + self.input_norm = GemmaRMSNorm(config.hidden_size, eps=config.norm_eps).float() + self.post_attention_norm = GemmaRMSNorm( + config.hidden_size, eps=config.norm_eps + ).float() + self.pre_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.norm_eps + ).float() + self.post_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.norm_eps + ).float() + + def forward( + self, + hidden_states, # (b, t, hidden_dim) + mask, # (b, 1, t, c+t) + pos_emb, # (b, 2, t, head_dim) + past_key, # (b, num_kv_head, c, head_dim) + past_value, # (b, num_kv_head, c, head_dim) + ): + residual = hidden_states + if self.jit_trace: + hidden_states = self.input_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.input_norm(hidden_states.to(torch.float32)).to(dtype) + + layer_device = hidden_states.device + + # Self Attention + attn_output, present_key, present_value = self.self_attn( + hidden_states=hidden_states.to(layer_device), + mask=mask.to(layer_device), + pos_emb=pos_emb.to(layer_device), + past_key=past_key.to(layer_device), + past_value=past_value.to(layer_device), + ) + + if self.jit_trace: + hidden_states = self.post_attention_norm(attn_output) + else: + dtype = hidden_states.dtype + hidden_states = self.post_attention_norm(attn_output.to(torch.float32)).to( + dtype + ) + + hidden_states = residual.to(layer_device) + hidden_states + + # Fully Connected + residual = hidden_states + if self.jit_trace: + hidden_states = self.pre_feedforward_layernorm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.pre_feedforward_layernorm( + hidden_states.to(torch.float32) + ).to(dtype) + + hidden_states = self.mlp(hidden_states) + + if self.jit_trace: + hidden_states = self.post_feedforward_layernorm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.post_feedforward_layernorm( + hidden_states.to(torch.float32) + ).to(dtype) + + hidden_states = residual + hidden_states + + if self.return_attn: + return hidden_states, present_key, present_value, attn_output + return hidden_states, present_key, present_value + + +class Gemma3DecoderLayer(DecoderLayer): + def __init__( + self, + config: GemmaConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=Gemma3Attention, + mlp_class=GemmaMLP, + ) + self.input_norm = GemmaRMSNorm(config.hidden_size, eps=config.norm_eps) + self.post_attention_norm = GemmaRMSNorm(config.hidden_size, eps=config.norm_eps) + self.pre_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.norm_eps + ) + self.post_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.norm_eps + ) + self.use_res_clamp = config.use_res_clamp + + def forward( + self, + hidden_states, # (b, t, hidden_dim) + mask, # (b, 1, t, c+t) + pos_emb, # (b, 2, t, head_dim) + past_key, # (b, num_kv_head, c, head_dim) + past_value, # (b, num_kv_head, c, head_dim) + ): + residual = hidden_states + if self.jit_trace: + hidden_states = self.input_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.input_norm(hidden_states.to(torch.float32)).to(dtype) + + layer_device = hidden_states.device + + # Self Attention + attn_output, present_key, present_value = self.self_attn( + hidden_states=hidden_states.to(layer_device), + mask=mask.to(layer_device), + pos_emb=pos_emb.to(layer_device), + past_key=past_key.to(layer_device), + past_value=past_value.to(layer_device), + ) + + if self.jit_trace: + hidden_states = self.post_attention_norm(attn_output) + else: + dtype = hidden_states.dtype + hidden_states = self.post_attention_norm(attn_output.to(torch.float32)).to( + dtype + ) + + hidden_states = residual.to(layer_device) + hidden_states + + if self.use_res_clamp: + hidden_states = torch.clamp( + hidden_states, + torch.finfo(torch.float16).min, + torch.finfo(torch.float16).max, + ) + + # Fully Connected + residual = hidden_states + if self.jit_trace: + hidden_states = self.pre_feedforward_layernorm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.pre_feedforward_layernorm( + hidden_states.to(torch.float32) + ).to(dtype) + + hidden_states = self.mlp(hidden_states) + + if self.jit_trace: + hidden_states = self.post_feedforward_layernorm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.post_feedforward_layernorm( + hidden_states.to(torch.float32) + ).to(dtype) + + hidden_states = residual + hidden_states + + if self.use_res_clamp: + hidden_states = torch.clamp( + hidden_states, + torch.finfo(torch.float16).min, + torch.finfo(torch.float16).max, + ) + + if self.return_attn: + return hidden_states, present_key, present_value, attn_output + return hidden_states, present_key, present_value + + +class GemmaModelChunk(ModelChunk): + def __init__( + self, + config: GemmaConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=GemmaDecoderLayer, + ) + + +class Gemma2ModelChunk(ModelChunk): + def __init__( + self, + config: GemmaConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=Gemma2DecoderLayer, + ) + if self.include_tail: + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.norm_eps).float() + + def forward(self, inputs_embeds, mask, pos_emb, *cache): + if not self.jit_trace: + assert ( + len(cache) == 2 * self.num_blocks + ), f"split cache wrong number of input caches: {len(cache)} != 2*{self.num_blocks}" + assert ( + cache[0].shape[0] == inputs_embeds.size()[0] + ), f"split cache batch size mismatch: {cache[0].shape[0]} != {inputs_embeds.size()[0]}" + + inputs_embeds = inputs_embeds.to(self.device_list[0]) + + if self.config.use_stable_embedding and self.chunk_idx == 0: + if self.jit_trace: + inputs_embeds = self.embed_layer_norm(inputs_embeds) + else: + inputs_embeds = self.embed_layer_norm( + inputs_embeds.to(torch.float32) + ).to(self.dtype) + + if self.chunk_idx == 0: + inputs_embeds = inputs_embeds * math.sqrt(self.config.hidden_size) + + hidden_states = inputs_embeds + + next_key_cache = [] + next_value_cache = [] + if self.return_attn: + attn_outputs = [] + + # decoder layers + for idx, decoder_layer in enumerate(self.layers): + decoder_outputs = decoder_layer( + hidden_states.to(self.device_list[idx]), + mask=mask.to(self.device_list[idx]), + pos_emb=pos_emb.to(self.device_list[idx]), + past_key=cache[idx].to(self.device_list[idx]), + past_value=cache[self.num_blocks + idx].to(self.device_list[idx]), + ) + hidden_states = decoder_outputs[0] + next_key_cache.append(decoder_outputs[1].to(inputs_embeds.device)) + next_value_cache.append(decoder_outputs[2].to(inputs_embeds.device)) + if self.return_attn: + attn_outputs.append(decoder_outputs[3].to(inputs_embeds.device)) + + if self.include_tail: + if self.jit_trace: + hidden_states = self.norm(hidden_states) + else: + hidden_states = self.norm(hidden_states.to(torch.float32)).to( + self.dtype + ) + hidden_states = self.lm_head(hidden_states) + hidden_states = scaled_tanh( + hidden_states, self.config.final_logit_softcapping + ) + + if self.return_attn: + return hidden_states, *next_key_cache, *next_value_cache, *attn_outputs + return hidden_states, *next_key_cache, *next_value_cache + + def load_weights(self, state_dict, state_dict_start_idx): + if state_dict is None: + fake_weights = True + else: + expected_subkey = f"layers.{state_dict_start_idx}.self_attn.o_proj.weight" + state_dict_keys = list(state_dict.keys()) + temp_key = None + input_norm_subkey = None + post_attention_norm_subkey = None + for key in state_dict_keys: + if expected_subkey in key: + temp_key = key + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "input" in key + ): + input_norm_subkey = key.split(".")[-2] + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "post_attention" in key + ): + post_attention_norm_subkey = key.split(".")[-2] + if temp_key is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s o_proj weight inside state_dict. " + f"Please ensure o_proj weight key contains: {expected_subkey}" + ) + if input_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s input norm weight inside state_dict. " + f"Please ensure input norm weight key contains: layers.{state_dict_start_idx}, norm, and input inside" + " the key string." + ) + if post_attention_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s post attention norm weight inside state_dict." + f" Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " + "post_attention inside the key string." + ) + prefix = temp_key.split(expected_subkey)[0] + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + if self.config.use_stable_embedding and self.chunk_idx == 0: + if fake_weights: + temp_state_dict = { + "embed_layer_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "embed_layer_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + } + else: + temp_state_dict = { + "embed_layer_norm.weight": state_dict.pop( + f"{prefix}embed_layer_norm.weight" + ).to(torch.float32), + "embed_layer_norm.bias": state_dict.pop( + f"{prefix}embed_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + } + else: + temp_state_dict = {} + + for inner_layer_idx in range(self.num_blocks): + if fake_weights: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = torch.rand( + 3 * self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.head_dim * self.config.num_attention_heads, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.pre_feedforward_layernorm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_feedforward_layernorm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + # Not Fake Weights + else: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.weight" + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.bias", + torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ), + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros( + self.config.head_dim + * self.config.num_attention_heads, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.up_proj.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.pre_feedforward_layernorm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.pre_feedforward_layernorm.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.post_feedforward_layernorm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.post_feedforward_layernorm.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_proj.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.up_proj.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + if self.include_tail: + if fake_weights: + temp_state_dict = { + **temp_state_dict, + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "lm_head.weight": torch.rand( + self.config.vocab_size, + self.config.hidden_size, + dtype=self.dtype, + ), + "lm_head.bias": torch.zeros( + self.config.vocab_size, dtype=self.dtype + ), + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + if self.config.tie_word_embeddings: + lm_head_weight_key = f"{prefix}embed_tokens.weight" + lm_head_bias_key = f"{prefix}embed_tokens.bias" + else: + lm_head_weight_key = "lm_head.weight" + lm_head_bias_key = "lm_head.bias" + temp_state_dict = { + **temp_state_dict, + **{ + "lm_head.weight": state_dict.pop(lm_head_weight_key), + "norm.weight": state_dict.pop(f"{prefix}norm.weight").to( + torch.float32 + ), + "lm_head.bias": state_dict.pop( + lm_head_bias_key, + torch.zeros(self.config.vocab_size, dtype=self.dtype), + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for chunk {self.chunk_idx}") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_blocks): + self.layers[i].to(self.device_list[i]) + if self.config.use_stable_embedding and self.chunk_idx == 0: + self.embed_layer_norm.to(self.device_list[0]) + if self.include_tail: + self.norm.to(self.device_list[-1]) + self.lm_head.to(self.device_list[-1]) + self.eval() + + return self + + +class Gemma3ModelChunk(ModelChunk): + def __init__( + self, + config: GemmaConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=Gemma3DecoderLayer, + ) + if self.include_tail: + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.norm_eps).float() + self.sliding_window_pattern = self.config.sliding_window_pattern + self.global_local_attention_pattern = self.config.global_local_attention_pattern + + def forward(self, inputs_embeds, global_mask, local_mask, pos_emb, *cache): + if not self.jit_trace: + assert ( + len(cache) == 2 * self.num_blocks + ), f"split cache wrong number of input caches: {len(cache)} != 2*{self.num_blocks}" + assert ( + cache[0].shape[0] == inputs_embeds.size()[0] + ), f"split cache batch size mismatch: {cache[0].shape[0]} != {inputs_embeds.size()[0]}" + + inputs_embeds = inputs_embeds.to(self.device_list[0]) + + if self.config.use_stable_embedding and self.chunk_idx == 0: + if self.jit_trace: + inputs_embeds = self.embed_layer_norm(inputs_embeds) + else: + inputs_embeds = self.embed_layer_norm( + inputs_embeds.to(torch.float32) + ).to(self.dtype) + + if self.chunk_idx == 0: + inputs_embeds = inputs_embeds * math.sqrt(self.config.hidden_size) + + hidden_states = inputs_embeds + + next_key_cache = [] + next_value_cache = [] + if self.return_attn: + attn_outputs = [] + + # decoder layers + for idx, decoder_layer in enumerate(self.layers): + outer_layer_idx = self.chunk_idx * self.num_blocks + idx + mask_pos = outer_layer_idx % self.sliding_window_pattern + mask_type = self.global_local_attention_pattern[mask_pos] + if mask_type == "SLIDING_LOCAL": + mask = local_mask + elif mask_type == "GLOBAL": + mask = global_mask + else: + raise TypeError(f"Unrecognized mask type: {mask_type}") + + decoder_outputs = decoder_layer( + hidden_states.to(self.device_list[idx]), + mask=mask.to(self.device_list[idx]), + pos_emb=pos_emb.to(self.device_list[idx]), + past_key=cache[idx].to(self.device_list[idx]), + past_value=cache[self.num_blocks + idx].to(self.device_list[idx]), + ) + hidden_states = decoder_outputs[0] + next_key_cache.append(decoder_outputs[1].to(inputs_embeds.device)) + next_value_cache.append(decoder_outputs[2].to(inputs_embeds.device)) + if self.return_attn: + attn_outputs.append(decoder_outputs[3].to(inputs_embeds.device)) + + if self.include_tail: + if self.jit_trace: + hidden_states = self.norm(hidden_states) + else: + hidden_states = self.norm(hidden_states.to(torch.float32)).to( + self.dtype + ) + hidden_states = self.lm_head(hidden_states) + + if self.return_attn: + return hidden_states, *next_key_cache, *next_value_cache, *attn_outputs + return hidden_states, *next_key_cache, *next_value_cache + + def load_weights(self, state_dict, state_dict_start_idx): + if state_dict is None: + fake_weights = True + else: + expected_subkey = f"layers.{state_dict_start_idx}.self_attn.o_proj.weight" + state_dict_keys = list(state_dict.keys()) + temp_key = None + input_norm_subkey = None + post_attention_norm_subkey = None + q_norm_subkey = None + for key in state_dict_keys: + if expected_subkey in key: + temp_key = key + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "input" in key + ): + input_norm_subkey = key.split(".")[-2] + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "post_attention" in key + ): + post_attention_norm_subkey = key.split(".")[-2] + if f"layers.{state_dict_start_idx}" in key and "q_norm" in key: + q_norm_subkey = key + if temp_key is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s o_proj weight inside state_dict. " + f"Please ensure o_proj weight key contains: {expected_subkey}" + ) + if input_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s input norm weight inside state_dict. " + f"Please ensure input norm weight key contains: layers.{state_dict_start_idx}, norm, and input inside" + " the key string." + ) + if post_attention_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s post attention norm weight inside state_dict." + f"Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " + "post_attention inside the key string." + ) + if q_norm_subkey is not None and self.config.use_qk_norm is False: + model_name = "".join( + [char for char in self.config.model_type if not char.isdigit()] + ) + raise KeyError( + f"Layer {state_dict_start_idx}'s q_norm weight detected inside state_dict but not used." + f"Please set use_qk_norm to True in configuration_{model_name}.py" + ) + + prefix = temp_key.split(expected_subkey)[0] + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + if self.config.use_stable_embedding and self.chunk_idx == 0: + if fake_weights: + temp_state_dict = { + "embed_layer_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "embed_layer_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + } + else: + temp_state_dict = { + "embed_layer_norm.weight": state_dict.pop( + f"{prefix}embed_layer_norm.weight" + ).to(torch.float32), + "embed_layer_norm.bias": state_dict.pop( + f"{prefix}embed_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + } + else: + temp_state_dict = {} + + for inner_layer_idx in range(self.num_blocks): + if fake_weights: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = torch.rand( + 3 * self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.pre_feedforward_layernorm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_feedforward_layernorm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + }, + } + + if self.config.use_qk_norm: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_norm.weight": torch.rand( + self.config.head_dim, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.self_attn.k_norm.weight": torch.rand( + self.head_dim, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + else: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.weight" + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.bias", + torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ), + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros( + self.config.head_dim + * self.config.num_attention_heads, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.up_proj.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.pre_feedforward_layernorm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.pre_feedforward_layernorm.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.post_feedforward_layernorm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.post_feedforward_layernorm.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.gate_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_proj.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.up_proj.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + }, + } + + if self.config.use_qk_norm: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_norm.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_norm.weight" + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + if self.include_tail: + if fake_weights: + temp_state_dict = { + **temp_state_dict, + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "lm_head.weight": torch.rand( + self.config.vocab_size, + self.config.hidden_size, + dtype=self.dtype, + ), + "lm_head.bias": torch.zeros( + self.config.vocab_size, dtype=self.dtype + ), + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + if self.config.tie_word_embeddings: + lm_head_weight_key = f"{prefix}embed_tokens.weight" + lm_head_bias_key = f"{prefix}embed_tokens.bias" + else: + lm_head_weight_key = "lm_head.weight" + lm_head_bias_key = "lm_head.bias" + temp_state_dict = { + **temp_state_dict, + **{ + "lm_head.weight": state_dict.pop(lm_head_weight_key), + "norm.weight": state_dict.pop(f"{prefix}norm.weight").to( + torch.float32 + ), + "lm_head.bias": state_dict.pop( + lm_head_bias_key, + torch.zeros(self.config.vocab_size, dtype=self.dtype), + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for chunk {self.chunk_idx}") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_blocks): + self.layers[i].to(self.device_list[i]) + if self.config.use_stable_embedding and self.chunk_idx == 0: + self.embed_layer_norm.to(self.device_list[0]) + if self.include_tail: + self.norm.to(self.device_list[-1]) + self.lm_head.to(self.device_list[-1]) + self.eval() + + return self + + def get_example_inputs( + self, num_token: int = 128, cache_size: int = 512, get_dym_shape=False + ): + head_dim = int(self.head_dim) + example_inputs = ( + torch.randn( + 1, num_token, self.config.hidden_size, device="cpu", dtype=torch.float32 + ), + torch.randn( # Global mask + 1, + 1, + num_token, + cache_size + num_token, + device="cpu", + dtype=torch.float32, + ), + torch.randn( # Local mask + 1, + 1, + num_token, + cache_size + num_token, + device="cpu", + dtype=torch.float32, + ), + torch.randn(1, 2, num_token, head_dim, device="cpu", dtype=torch.float32), + *[ + torch.randn( + 1, + self.config.num_key_value_heads, + cache_size, + head_dim, + device="cpu", + dtype=torch.float32, + ) + for _ in range(2 * self.num_blocks) + ], + ) + # Specify dims that would be dynamic during calibration + # Note: Assume cache size fixed shape as torch dynamic shape cannot handle dim 3 being + # combination of 2 dynamic dims + if get_dym_shape: + nt = Dim("num_token", max=num_token) + cache_dims = tuple(({} for _ in range(2 * self.num_blocks))) + dynamic_shapes = ( + {0: Dim.STATIC, 1: nt, 2: Dim.STATIC}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: nt + cache_size}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: nt + cache_size}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: Dim.STATIC}, + cache_dims, + ) + return example_inputs, dynamic_shapes + + return example_inputs diff --git a/examples/mediatek/models/llm_models/modeling_llama.py b/examples/mediatek/models/llm_models/modeling_llama.py index ac4f4f4e749..aef81c1b7e8 100644 --- a/examples/mediatek/models/llm_models/modeling_llama.py +++ b/examples/mediatek/models/llm_models/modeling_llama.py @@ -1,23 +1,8 @@ -# coding=utf-8 -# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# Copyright (c) MediaTek Inc. +# All rights reserved # -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" PyTorch LLaMA model.""" +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. import numpy as np import torch @@ -34,8 +19,15 @@ def __init__(self, config: LlamaConfig): class LlamaAttention(Attention): - def __init__(self, config: LlamaConfig): - super().__init__(config) + def __init__( + self, + config: LlamaConfig, + jit_trace=False, + ): + super().__init__( + config, + jit_trace, + ) class LlamaDecoderLayer(DecoderLayer): diff --git a/examples/mediatek/models/llm_models/modeling_phi.py b/examples/mediatek/models/llm_models/modeling_phi.py new file mode 100644 index 00000000000..c05b9f4b17e --- /dev/null +++ b/examples/mediatek/models/llm_models/modeling_phi.py @@ -0,0 +1,901 @@ +# Copyright (c) MediaTek 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 numpy as np +import torch +from models.llm_models.configuration_phi import PhiConfig + +from models.llm_models.modeling_common import Attention, DecoderLayer, ModelChunk +from torch import nn +from torch.export import Dim + +np.random.seed(42) + +# flake8: noqa: C901 + + +class PhiMLP(nn.Module): + def __init__(self, config: PhiConfig): + super().__init__() + hidden_size = config.hidden_size + intermediate_size = config.intermediate_size + + self.gate_up_proj = nn.Linear(hidden_size, 2 * intermediate_size) + self.down_proj = nn.Linear(intermediate_size, hidden_size) + + def forward(self, x): + up = self.gate_up_proj(x) + gate, up = up.chunk(2, dim=-1) + pre_down = gate * torch.sigmoid(gate) * up + down = self.down_proj(pre_down) + + return down + + +class Phi3Attention(Attention): + def __init__(self, config: PhiConfig, jit_trace=None): + super().__init__(config, jit_trace) + + +class Phi4Attention(Attention): + def __init__(self, config: PhiConfig, jit_trace=None): + super().__init__(config, jit_trace) + + def apply_rotary_pos_emb_mtk(self, q, k, cos, sin): + """Apply rotary positional embedding to query and key states. + + Phi4-mini uses partial rotary embedding, which only rotate some of the q and k. + """ + + rotary_dim = cos.shape[-1] + + # Use split + q_rot, q_pass = torch.split(q, rotary_dim, dim=-1) + q1, q2 = torch.split(q_rot, rotary_dim // 2, dim=-1) + q_rotated = torch.cat((-q2, q1), dim=-1) + + k_rot, k_pass = torch.split(k, rotary_dim, dim=-1) + k1, k2 = torch.split(k_rot, rotary_dim // 2, dim=-1) + k_rotated = torch.cat((-k2, k1), dim=-1) + + q_embed = q_rot * cos + q_rotated * sin + k_embed = k_rot * cos + k_rotated * sin + + q_embed = torch.cat((q_embed, q_pass), dim=-1) + k_embed = torch.cat((k_embed, k_pass), dim=-1) + + return q_embed, k_embed + + +class Phi3DecoderLayer(DecoderLayer): + def __init__( + self, + config: PhiConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=Phi3Attention, + mlp_class=PhiMLP, + ) + + +class Phi4DecoderLayer(DecoderLayer): + def __init__( + self, + config: PhiConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=Phi4Attention, + mlp_class=PhiMLP, + ) + + +class Phi3ModelChunk(ModelChunk): + def __init__( + self, + config: PhiConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=Phi3DecoderLayer, + ) + + def load_weights(self, state_dict, state_dict_start_idx): + if state_dict is None: + fake_weights = True + else: + expected_subkey = f"layers.{state_dict_start_idx}.self_attn.o_proj.weight" + state_dict_keys = list(state_dict.keys()) + temp_key = None + input_norm_subkey = None + post_attention_norm_subkey = None + for key in state_dict_keys: + if expected_subkey in key: + temp_key = key + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "input" in key + ): + input_norm_subkey = key.split(".")[-2] + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "post_attention" in key + ): + post_attention_norm_subkey = key.split(".")[-2] + if temp_key is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s o_proj weight inside state_dict. " + f"Please ensure o_proj weight key contains: {expected_subkey}" + ) + if input_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s input norm weight inside state_dict. " + f"Please ensure input norm weight key contains: layers.{state_dict_start_idx}, norm, and input inside" + " the key string." + ) + if post_attention_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s post attention norm weight inside state_dict." + f" Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " + "post_attention inside the key string." + ) + prefix = temp_key.split(expected_subkey)[0] + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + if self.config.use_stable_embedding and self.chunk_idx == 0: + if fake_weights: + temp_state_dict = { + "embed_layer_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "embed_layer_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + } + else: + temp_state_dict = { + "embed_layer_norm.weight": state_dict.pop( + f"{prefix}embed_layer_norm.weight" + ).to(torch.float32), + "embed_layer_norm.bias": state_dict.pop( + f"{prefix}embed_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + } + else: + temp_state_dict = {} + + for inner_layer_idx in range(self.num_blocks): + if fake_weights: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = torch.rand( + 3 * self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.weight": torch.rand( # Check Again + 2 * self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.bias": torch.zeros( + 2 * self.config.intermediate_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + else: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.weight" + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.bias", + torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ), + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_up_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_up_proj.bias", + torch.zeros( + 2 * self.config.intermediate_size, dtype=self.dtype + ), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + if self.include_tail: + if fake_weights: + temp_state_dict = { + **temp_state_dict, + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "lm_head.weight": torch.rand( + self.config.vocab_size, + self.config.hidden_size, + dtype=self.dtype, + ), + "lm_head.bias": torch.zeros( + self.config.vocab_size, dtype=self.dtype + ), + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + if self.config.tie_word_embeddings: + lm_head_weight_key = f"{prefix}embed_tokens.weight" + lm_head_bias_key = f"{prefix}embed_tokens.bias" + else: + lm_head_weight_key = "lm_head.weight" + lm_head_bias_key = "lm_head.bias" + temp_state_dict = { + **temp_state_dict, + **{ + "lm_head.weight": state_dict.pop(lm_head_weight_key), + "norm.weight": state_dict.pop(f"{prefix}norm.weight").to( + torch.float32 + ), + "lm_head.bias": state_dict.pop( + lm_head_bias_key, + torch.zeros(self.config.vocab_size, dtype=self.dtype), + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for chunk {self.chunk_idx}") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_blocks): + self.layers[i].to(self.device_list[i]) + if self.config.use_stable_embedding and self.chunk_idx == 0: + self.embed_layer_norm.to(self.device_list[0]) + if self.include_tail: + self.norm.to(self.device_list[-1]) + self.lm_head.to(self.device_list[-1]) + self.eval() + + return self + + +class Phi4ModelChunk(ModelChunk): + def __init__( + self, + config: PhiConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=Phi4DecoderLayer, + ) + self.partial_rotary_factor = config.partial_rotary_factor + + def load_weights(self, state_dict, state_dict_start_idx): + if state_dict is None: + fake_weights = True + else: + expected_subkey = f"layers.{state_dict_start_idx}.self_attn.o_proj.weight" + state_dict_keys = list(state_dict.keys()) + temp_key = None + input_norm_subkey = None + post_attention_norm_subkey = None + for key in state_dict_keys: + if expected_subkey in key: + temp_key = key + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "input" in key + ): + input_norm_subkey = key.split(".")[-2] + if ( + f"layers.{state_dict_start_idx}" in key + and "norm" in key + and "post_attention" in key + ): + post_attention_norm_subkey = key.split(".")[-2] + if temp_key is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s o_proj weight inside state_dict. " + f"Please ensure o_proj weight key contains: {expected_subkey}" + ) + if input_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s input norm weight inside state_dict. " + f"Please ensure input norm weight key contains: layers.{state_dict_start_idx}, norm, and input inside" + " the key string." + ) + if post_attention_norm_subkey is None: + raise KeyError( + f"Cannot find layer {state_dict_start_idx}'s post attention norm weight inside state_dict." + f" Please ensure post attention norm weight key contains: layers.{state_dict_start_idx}, norm, and " + "post_attention inside the key string." + ) + prefix = temp_key.split(expected_subkey)[0] + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + if self.config.use_stable_embedding and self.chunk_idx == 0: + if fake_weights: + temp_state_dict = { + "embed_layer_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "embed_layer_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + } + else: + temp_state_dict = { + "embed_layer_norm.weight": state_dict.pop( + f"{prefix}embed_layer_norm.weight" + ).to(torch.float32), + "embed_layer_norm.bias": state_dict.pop( + f"{prefix}embed_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + } + else: + temp_state_dict = {} + + for inner_layer_idx in range(self.num_blocks): + if fake_weights: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = torch.rand( + 3 * self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.weight": torch.rand( # Check Again + 2 * self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.bias": torch.zeros( + 2 * self.config.intermediate_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + else: + if self.config.combine_qkv: + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.weight" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.weight" + ) + temp_state_dict[ + f"layers.{inner_layer_idx}.self_attn.qkv_proj.bias" + ] = state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.qkv_proj.bias", + torch.zeros( + (2 * self.config.num_key_value_heads * self.head_dim) + + self.config.hidden_size, + dtype=self.dtype, + ), + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_up_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.o_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.gate_up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.gate_up_proj.bias", + torch.zeros( + 2 * self.config.intermediate_size, dtype=self.dtype + ), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.mlp.down_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + if self.include_tail: + if fake_weights: + temp_state_dict = { + **temp_state_dict, + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "lm_head.weight": torch.rand( + self.config.vocab_size, + self.config.hidden_size, + dtype=self.dtype, + ), + "lm_head.bias": torch.zeros( + self.config.vocab_size, dtype=self.dtype + ), + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + if self.config.tie_word_embeddings: + lm_head_weight_key = f"{prefix}embed_tokens.weight" + lm_head_bias_key = f"{prefix}embed_tokens.bias" + else: + lm_head_weight_key = "lm_head.weight" + lm_head_bias_key = "lm_head.bias" + temp_state_dict = { + **temp_state_dict, + **{ + "lm_head.weight": state_dict.pop(lm_head_weight_key), + "norm.weight": state_dict.pop(f"{prefix}norm.weight").to( + torch.float32 + ), + "lm_head.bias": state_dict.pop( + lm_head_bias_key, + torch.zeros(self.config.vocab_size, dtype=self.dtype), + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for chunk {self.chunk_idx}") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_blocks): + self.layers[i].to(self.device_list[i]) + if self.config.use_stable_embedding and self.chunk_idx == 0: + self.embed_layer_norm.to(self.device_list[0]) + if self.include_tail: + self.norm.to(self.device_list[-1]) + self.lm_head.to(self.device_list[-1]) + self.eval() + + return self + + def get_example_inputs( + self, num_token: int = 128, cache_size: int = 512, get_dym_shape=False + ): + head_dim = int(self.head_dim) + example_inputs = ( + torch.randn( + 1, num_token, self.config.hidden_size, device="cpu", dtype=torch.float32 + ), + torch.randn( + 1, + 1, + num_token, + cache_size + num_token, + device="cpu", + dtype=torch.float32, + ), + torch.randn( + 1, + 2, + num_token, + int(head_dim * self.partial_rotary_factor), + device="cpu", + dtype=torch.float32, + ), + *[ + torch.randn( + 1, + self.config.num_key_value_heads, + cache_size, + head_dim, + device="cpu", + dtype=torch.float32, + ) + for _ in range(2 * self.num_blocks) + ], + ) + # Specify dims that would be dynamic during calibration + # Note: Assume cache size fixed shape as torch dynamic shape cannot handle dim 3 being + # combination of 2 dynamic dims + if get_dym_shape: + nt = Dim("num_token", max=num_token) + cache_dims = tuple(({} for _ in range(2 * self.num_blocks))) + dynamic_shapes = ( + {0: Dim.STATIC, 1: nt, 2: Dim.STATIC}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: nt + cache_size}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: Dim.STATIC}, + cache_dims, + ) + return example_inputs, dynamic_shapes + + return example_inputs diff --git a/examples/mediatek/models/llm_models/modeling_qwen.py b/examples/mediatek/models/llm_models/modeling_qwen.py new file mode 100644 index 00000000000..8ba712464ed --- /dev/null +++ b/examples/mediatek/models/llm_models/modeling_qwen.py @@ -0,0 +1,106 @@ +# Copyright (c) MediaTek 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 numpy as np +import torch +from models.llm_models.configuration_qwen import QwenConfig + +from models.llm_models.modeling_common import Attention, DecoderLayer, MLP, ModelChunk + +np.random.seed(42) + + +class QwenMLP(MLP): + def __init__(self, config: QwenConfig): + super().__init__(config) + + +class QwenAttention(Attention): + def __init__(self, config: QwenConfig, jit_trace=False): + super().__init__(config, jit_trace) + + +class Qwen3Attention(Attention): + def __init__(self, config: QwenConfig, jit_trace=False): + super().__init__(config, jit_trace) + + +class QwenDecoderLayer(DecoderLayer): + def __init__( + self, + config: QwenConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=QwenAttention, + mlp_class=QwenMLP, + ) + + +class Qwen3DecoderLayer(DecoderLayer): + def __init__( + self, + config: QwenConfig, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + return_attn, + jit_trace, + attn_class=Qwen3Attention, + mlp_class=QwenMLP, + ) + + +class Qwen2ModelChunk(ModelChunk): + def __init__( + self, + config: QwenConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=QwenDecoderLayer, + ) + + +class Qwen3ModelChunk(ModelChunk): + def __init__( + self, + config: QwenConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=Qwen3DecoderLayer, + ) diff --git a/examples/mediatek/models/llm_models/modeling_whisper.py b/examples/mediatek/models/llm_models/modeling_whisper.py new file mode 100644 index 00000000000..15562cfc4a9 --- /dev/null +++ b/examples/mediatek/models/llm_models/modeling_whisper.py @@ -0,0 +1,1241 @@ +# Copyright (c) MediaTek 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 numpy as np +import torch +from models.llm_models.configuration_whisper import ( + WhisperConfig, + WhisperDecoderConfig, + WhisperEncoderConfig, +) + +from models.llm_models.modeling_common import Attention, ModelChunk, RMSNorm, TorchGelu +from torch import nn +from torch.export import Dim + +np.random.seed(42) + +# flake8: noqa: C901 + + +class WhisperMLP(nn.Module): + def __init__(self, config: WhisperConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size) + self.gelu = TorchGelu() + + def forward(self, x): + up = self.up_proj(x) + pre_down = self.gelu(up) + down = self.down_proj(pre_down) + + return down + + +class WhisperEncoderAttention(nn.Module): + def __init__(self, config: WhisperEncoderConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.scale = self.head_dim**-0.5 + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim) + self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim) + self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size) + + def forward( + self, + hidden_states, + ): + bsz, q_len, _ = hidden_states.size() + + query_states = ( + (self.q_proj(hidden_states) * self.scale) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .reshape(self.num_heads, q_len, self.head_dim) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .reshape(self.num_heads, q_len, self.head_dim) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .reshape(self.num_heads, q_len, self.head_dim) + ) + + attn_weights = torch.matmul(query_states, key_states.transpose(1, 2)) + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(0, 1) + attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads) + attn_output = self.o_proj(attn_output) + + return attn_output + + +class WhisperDecoderAttention(Attention): + def __init__(self, config: WhisperConfig): + super().__init__(config) + self.scale = self.head_dim**-0.5 + + def forward( + self, + hidden_states, + mask, + past_key, + past_value, + ): + bsz, q_len, _ = hidden_states.size() + c_len = past_key.size()[2] + + query_states = ( + (self.q_proj(hidden_states) * self.scale) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + key_states = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + value_states = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_key_value_heads, self.head_dim) + .transpose(1, 2) + ) + + key_states = torch.cat([past_key, key_states], dim=2) + value_states = torch.cat([past_value, value_states], dim=2) + key_states_out = key_states + value_states_out = value_states + if self.num_key_value_groups > 1: + key_states = self.repeat_kv( + key_states, bsz, q_len + c_len, self.num_key_value_groups + ) + value_states = self.repeat_kv( + value_states, bsz, q_len + c_len, self.num_key_value_groups + ) + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) + + attn_weights = attn_weights + mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.head_dim * self.num_heads) + attn_output = self.o_proj(attn_output) + + key_states_out = key_states_out[:, :, q_len:, :] + value_states_out = value_states_out[:, :, q_len:, :] + + return attn_output, key_states_out, value_states_out + + +class WhisperCrossAttention(nn.Module): + def __init__(self, config: WhisperConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.scale = self.head_dim**-0.5 + self.q_proj = nn.Linear(self.hidden_size, self.hidden_size) + self.o_proj = nn.Linear(self.hidden_size, self.hidden_size) + + def forward(self, hidden_states, cross_key, cross_value): + bsz, q_len, _ = hidden_states.size() + + query_states = ( + (self.q_proj(hidden_states) * self.scale) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + + key_states = cross_key + value_states = cross_value + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) + attn_weights = nn.functional.softmax(attn_weights, dim=-1) + + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + + attn_output = self.o_proj(attn_output) + + return attn_output + + +class WhisperEncoderLayer(nn.Module): + def __init__( + self, + config: WhisperEncoderConfig, + jit_trace=False, + attn_class=WhisperEncoderAttention, + mlp_class=WhisperMLP, + ): + super().__init__() + self.config = config + self.jit_trace = jit_trace + self.hidden_size = config.hidden_size + self.self_attn = attn_class(config) + self.mlp = mlp_class(config) + + norm_class = RMSNorm if config.norm == "RMSNorm" else nn.LayerNorm + self.input_norm = norm_class(self.hidden_size, eps=config.norm_eps).float() + self.post_attention_norm = norm_class( + self.hidden_size, eps=config.norm_eps + ).float() + + def forward( + self, + hidden_states, + ): + residual = hidden_states + + if self.jit_trace: + hidden_states = self.input_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.input_norm(hidden_states.to(torch.float32)).to(dtype) + + layer_device = hidden_states.device + + attn_outputs = self.self_attn(hidden_states.to(layer_device)) + hidden_states = residual.to(layer_device) + attn_outputs + + residual = hidden_states + if self.jit_trace: + hidden_states = self.post_attention_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.post_attention_norm( + hidden_states.to(torch.float32) + ).to(dtype) + + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class WhisperEncoderKVLayer(nn.Module): + def __init__(self, config: WhisperEncoderConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.k_proj = nn.Linear(self.hidden_size, self.hidden_size) + self.v_proj = nn.Linear(self.hidden_size, self.hidden_size) + + def forward(self, hidden_states): + bsz, q_len, _ = hidden_states.size() + + k_out = ( + self.k_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + v_out = ( + self.v_proj(hidden_states) + .view(bsz, q_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + + return k_out, v_out + + +class WhisperDecoderLayer(nn.Module): + def __init__( + self, + config: WhisperConfig, + return_attn=False, + jit_trace=False, + attn_class=WhisperDecoderAttention, + mlp_class=WhisperMLP, + cross_attn_class=WhisperCrossAttention, + ): + super().__init__() + self.hidden_size = config.hidden_size + self.return_attn = return_attn + self.jit_trace = jit_trace + self.self_attn = attn_class(config) + self.cross_attn = cross_attn_class(config) + self.mlp = mlp_class(config) + + norm_class = RMSNorm if config.norm == "RMSNorm" else nn.LayerNorm + self.input_norm = norm_class(self.hidden_size, eps=config.norm_eps).float() + self.post_attention_norm = norm_class( + self.hidden_size, eps=config.norm_eps + ).float() + self.post_cross_attention_norm = norm_class( + self.hidden_size, eps=config.norm_eps + ).float() + + def forward( + self, + hidden_states, + mask, + past_key, + past_value, + cross_key, + cross_value, + ): + residual = hidden_states + if self.jit_trace: + hidden_states = self.input_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.input_norm(hidden_states.to(torch.float32)).to(dtype) + + layer_device = hidden_states.device + + # Self Attention + attn_output, present_key, present_value = self.self_attn( + hidden_states=hidden_states.to(layer_device), + mask=mask.to(layer_device), + past_key=past_key.to(layer_device), + past_value=past_value.to(layer_device), + ) + hidden_states = residual.to(layer_device) + attn_output + residual = hidden_states + + if self.jit_trace: + hidden_states = self.post_attention_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.post_attention_norm( + hidden_states.to(torch.float32) + ).to(dtype) + + # Cross Attention + hidden_states = self.cross_attn( + hidden_states=hidden_states.to(layer_device), + cross_key=cross_key.to(layer_device), + cross_value=cross_value.to(layer_device), + ) + hidden_states = residual.to(layer_device) + hidden_states + residual = hidden_states + if self.jit_trace: + hidden_states = self.post_cross_attention_norm(hidden_states) + else: + dtype = hidden_states.dtype + hidden_states = self.post_cross_attention_norm( + hidden_states.to(torch.float32) + ).to(dtype) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + if self.return_attn: + return hidden_states, present_key, present_value, attn_output + return hidden_states, present_key, present_value + + +class WhisperEncoderModel(nn.Module): + def __init__( + self, + config: WhisperEncoderConfig, + dtype=torch.float32, + jit_trace=False, + encoder_class=WhisperEncoderLayer, + ): + super().__init__() + self.config = config + self.conv1 = nn.Conv1d( + config.num_mel_bins, config.hidden_size, kernel_size=3, padding=1 + ) + self.conv2 = nn.Conv1d( + config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1 + ) + self.gelu1 = TorchGelu() + self.gelu2 = TorchGelu() + self.embed_positions = nn.Embedding( + config.max_source_positions, config.hidden_size + ) + self.num_layers = config.num_hidden_layers + self.dtype = dtype + self.jit_trace = jit_trace + self.head_dim = config.head_dim + + self.layers = nn.ModuleList( + [encoder_class(config, self.jit_trace) for _ in range(self.num_layers)] + ) + + norm_class = RMSNorm if config.norm == "RMSNorm" else nn.LayerNorm + self.norm = norm_class(config.hidden_size, eps=config.norm_eps).float() + self.encoder_tails_kv = nn.ModuleList( + [WhisperEncoderKVLayer(config) for _ in range(config.decoder_num_layers)] + ) + + def forward(self, input_embeds): + input_embeds = self.gelu1(self.conv1(input_embeds)) + input_embeds = self.gelu2(self.conv2(input_embeds)).permute(0, 2, 1) + embed_pos = self.embed_positions.weight + input_embeds = input_embeds + embed_pos.reshape( + (1, embed_pos.shape[0], embed_pos.shape[1]) + ) + + hidden_states = input_embeds + for idx, encoder_layer in enumerate(self.layers): + encoder_outputs = encoder_layer(hidden_states.to(self.device_list[idx])) + hidden_states = encoder_outputs + + cross_key_cache = [] + cross_value_cache = [] + if self.jit_trace: + hidden_states = self.norm(hidden_states) + else: + hidden_states = self.norm(hidden_states.to(torch.float32)).to(self.dtype) + + for idx, encoder_tails_kv_layer in enumerate(self.encoder_tails_kv): + k_out, v_out = encoder_tails_kv_layer( + hidden_states.to(self.device_list[idx]) + ) + cross_key_cache.append(k_out.to(input_embeds.device)) + cross_value_cache.append(v_out.to(input_embeds.device)) + + k_out = torch.cat([k for k in cross_key_cache], dim=0) + v_out = torch.cat([v for v in cross_value_cache], dim=0) + cross_cache = torch.cat([k_out, v_out], dim=0) + + return hidden_states, cross_cache + + def load_weights(self, state_dict): + state_dict_start_idx = 0 + if state_dict is None: + fake_weights = True + else: + input_norm_subkey = "self_attn_layer_norm" + post_attention_norm_subkey = "final_layer_norm" + prefix = "model.encoder." + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + temp_state_dict = {} + + for inner_layer_idx in range(self.num_layers): + if fake_weights: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.head_dim * self.config.num_attention_heads, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + # Not Fake Weights + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.config.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.config.head_dim, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.out_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc2.weight" + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc1.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.out_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc2.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc1.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + + for dec_idx in range(self.config.decoder_num_layers): + if fake_weights: + temp_state_dict = { + **temp_state_dict, + **{ + f"encoder_tails_kv.{dec_idx}.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"encoder_tails_kv.{dec_idx}.k_proj.bias": torch.rand( + self.config.hidden_size, dtype=self.dtype + ), + f"encoder_tails_kv.{dec_idx}.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"encoder_tails_kv.{dec_idx}.v_proj.bias": torch.rand( + self.config.hidden_size, dtype=self.dtype + ), + }, + } + + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"encoder_tails_kv.{dec_idx}.k_proj.weight": state_dict.pop( + f"model.decoder.layers.{dec_idx}.encoder_attn.k_proj.weight" + ).to(torch.float32), + f"encoder_tails_kv.{dec_idx}.k_proj.bias": state_dict.pop( + f"model.decoder.layers.{dec_idx}.encoder_attn.k_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"encoder_tails_kv.{dec_idx}.v_proj.weight": state_dict.pop( + f"model.decoder.layers.{dec_idx}.encoder_attn.v_proj.weight" + ).to(torch.float32), + f"encoder_tails_kv.{dec_idx}.v_proj.bias": state_dict.pop( + f"model.decoder.layers.{dec_idx}.encoder_attn.v_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + }, + } + + if fake_weights: + temp_state_dict = { + **temp_state_dict, + **{ + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "conv1.weight": torch.rand( + self.config.hidden_size, + self.config.num_mel_bins, + 3, + dtype=self.dtype, + ), + "conv2.weight": torch.rand( + self.config.hidden_size, + self.config.num_mel_bins, + 3, + dtype=self.dtype, + ), + "conv1.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + "conv2.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + "embed_positions.weight": torch.rand( + self.config.max_source_positions, + self.config.hidden_size, + dtype=self.dtype, + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + temp_state_dict = { + **temp_state_dict, + **{ + "norm.weight": state_dict.pop(f"{prefix}layer_norm.weight").to( + torch.float32 + ), + "conv1.weight": state_dict.pop(f"{prefix}conv1.weight").to( + torch.float32 + ), + "conv2.weight": state_dict.pop(f"{prefix}conv2.weight").to( + torch.float32 + ), + "conv1.bias": state_dict.pop( + f"{prefix}conv1.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + "conv2.bias": state_dict.pop( + f"{prefix}conv2.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + "embed_positions.weight": state_dict.pop( + f"{prefix}embed_positions.weight" + ).to(torch.float32), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for encoder") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_layers): + self.layers[i].to(self.device_list[i]) + self.norm.to(self.device_list[-1]) + for i in range(self.config.decoder_num_layers): + self.encoder_tails_kv[i].to(self.device_list[i]) + self.conv1.to(self.device_list[0]) + self.conv2.to(self.device_list[0]) + self.embed_positions.to(self.device_list[0]) + self.eval() + + return self + + def get_example_inputs(self, num_mel_bins: int = 128, cache_size: int = 512): + example_inputs = ( + torch.randn(1, num_mel_bins, 3000, device="cpu", dtype=torch.float32), + ) + + return example_inputs + + +class WhisperDecoderModelChunk(ModelChunk): + def __init__( + self, + config: WhisperDecoderConfig, + num_blocks, + chunk_idx, + dtype=torch.float32, + include_tail=False, + return_attn=False, + jit_trace=False, + ): + super().__init__( + config, + num_blocks, + chunk_idx, + dtype, + include_tail, + return_attn, + jit_trace, + decoder_class=WhisperDecoderLayer, + ) + + def forward(self, inputs_embeds, mask, pos_emb, cross_cache, *cache): + if not self.jit_trace: + assert ( + len(cache) == 2 * self.num_blocks + ), f"split cache wrong number of input caches: {len(cache)} != 2*{self.num_blocks}" + assert ( + cache[0].shape[0] == inputs_embeds.size()[0] + ), f"split cache batch size mismatch: {cache[0].shape[0]} != {inputs_embeds.size()[0]}" + + inputs_embeds = inputs_embeds.to(self.device_list[0]) + + if self.config.use_stable_embedding and self.chunk_idx == 0: + if self.jit_trace: + inputs_embeds = self.embed_layer_norm(inputs_embeds) + else: + inputs_embeds = self.embed_layer_norm( + inputs_embeds.to(torch.float32) + ).to(self.dtype) + + if self.chunk_idx == 0: + inputs_embeds = inputs_embeds + pos_emb.to(self.device_list[0]) + hidden_states = inputs_embeds + + cross_cache = [*torch.split(cross_cache, 1, dim=0)] + + next_key_cache = [] + next_value_cache = [] + if self.return_attn: + attn_outputs = [] + + # decoder layers + for idx, decoder_layer in enumerate(self.layers): + decoder_outputs = decoder_layer( + hidden_states.to(self.device_list[idx]), + mask=mask.to(self.device_list[idx]), + past_key=cache[idx].to(self.device_list[idx]), + past_value=cache[self.num_blocks + idx].to(self.device_list[idx]), + cross_key=cross_cache[idx].to(self.device_list[idx]), + cross_value=cross_cache[self.num_blocks + idx].to( + self.device_list[idx] + ), + ) + + hidden_states = decoder_outputs[0] + next_key_cache.append(decoder_outputs[1].to(inputs_embeds.device)) + next_value_cache.append(decoder_outputs[2].to(inputs_embeds.device)) + if self.return_attn: + attn_outputs.append(decoder_outputs[3].to(inputs_embeds.device)) + + if self.include_tail: + if self.jit_trace: + hidden_states = self.norm(hidden_states) + else: + hidden_states = self.norm(hidden_states.to(torch.float32)).to( + self.dtype + ) + hidden_states = self.lm_head(hidden_states) + + if self.return_attn: + return hidden_states, *next_key_cache, *next_value_cache, *attn_outputs + return hidden_states, *next_key_cache, *next_value_cache + + def load_weights(self, state_dict, state_dict_start_idx): + if state_dict is None: + fake_weights = True + else: + input_norm_subkey = "self_attn_layer_norm" + post_attention_norm_subkey = "encoder_attn_layer_norm" + prefix = "model.decoder." + fake_weights = False + + outer_layer_idx = state_dict_start_idx + self.device_list = [] + + if self.config.use_stable_embedding and self.chunk_idx == 0: + if fake_weights: + temp_state_dict = { + "embed_layer_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "embed_layer_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + } + else: + temp_state_dict = { + "embed_layer_norm.weight": state_dict.pop( + f"{prefix}embed_layer_norm.weight" + ).to(torch.float32), + "embed_layer_norm.bias": state_dict.pop( + f"{prefix}embed_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + } + else: + temp_state_dict = {} + + for inner_layer_idx in range(self.num_blocks): + if fake_weights: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": torch.rand( + self.config.head_dim * self.config.num_attention_heads, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": torch.rand( + self.config.num_key_value_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": torch.zeros( + self.config.num_key_value_heads * self.head_dim, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.cross_attn.q_proj.weight": torch.rand( + self.config.num_attention_heads * self.head_dim, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.cross_attn.q_proj.bias": torch.zeros( + self.config.num_attention_heads * self.head_dim, + dtype=self.dtype, + ), + }, + } + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.cross_attn.o_proj.weight": torch.rand( + self.config.hidden_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": torch.rand( + self.config.hidden_size, + self.config.intermediate_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": torch.rand( + self.config.intermediate_size, + self.config.hidden_size, + dtype=self.dtype, + ), + f"layers.{inner_layer_idx}.input_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_cross_attention_norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.cross_attn.o_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": torch.zeros( + self.config.hidden_size, dtype=self.dtype + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + f"layers.{inner_layer_idx}.post_cross_attention_norm.bias": torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ), + }, + } + + # Not Fake Weights + else: + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.weight" + ), + f"layers.{inner_layer_idx}.self_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.q_proj.bias", + torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.k_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.k_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.config.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.self_attn.v_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.v_proj.bias", + torch.zeros( + self.config.num_key_value_heads * self.config.head_dim, + dtype=self.dtype, + ), + ), + f"layers.{inner_layer_idx}.cross_attn.q_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.encoder_attn.q_proj.weight" + ), + f"layers.{inner_layer_idx}.cross_attn.q_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.encoder_attn.q_proj.bias", + torch.zeros( + self.config.head_dim * self.config.num_attention_heads, + dtype=self.dtype, + ), + ), + }, + } + + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.self_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.out_proj.weight" + ), + f"layers.{inner_layer_idx}.cross_attn.o_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.encoder_attn.out_proj.weight" + ), + f"layers.{inner_layer_idx}.mlp.down_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc2.weight" + ), + f"layers.{inner_layer_idx}.mlp.up_proj.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc1.weight" + ), + f"layers.{inner_layer_idx}.input_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.weight" + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.post_cross_attention_norm.weight": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.final_layer_norm.weight" + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.self_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.self_attn.out_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.cross_attn.o_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.encoder_attn.out_proj.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.down_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc2.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ), + f"layers.{inner_layer_idx}.mlp.up_proj.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.fc1.bias", + torch.zeros( + self.config.intermediate_size, dtype=self.dtype + ), + ), + }, + } + + if self.config.norm == "LayerNorm": + temp_state_dict = { + **temp_state_dict, + **{ + f"layers.{inner_layer_idx}.input_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{input_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32), + f"layers.{inner_layer_idx}.post_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.{post_attention_norm_subkey}.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + f"layers.{inner_layer_idx}.post_cross_attention_norm.bias": state_dict.pop( + f"{prefix}layers.{outer_layer_idx}.final_layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to( + torch.float32 + ), + }, + } + + if torch.cuda.device_count() == 0 or self.jit_trace: + self.device_list.append("cpu") + else: + device_id = outer_layer_idx // ( + self.config.num_hidden_layers // torch.cuda.device_count() + + (self.config.num_hidden_layers % torch.cuda.device_count() != 0) + ) + self.device_list.append(f"cuda:{device_id}") + outer_layer_idx += 1 + + if self.include_tail: + if fake_weights: + temp_state_dict = { + **temp_state_dict, + "norm.weight": torch.rand( + self.config.hidden_size, dtype=torch.float32 + ), + "lm_head.weight": torch.rand( + self.config.vocab_size, + self.config.hidden_size, + dtype=self.dtype, + ), + "lm_head.bias": torch.zeros( + self.config.vocab_size, dtype=self.dtype + ), + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = torch.zeros( + self.config.hidden_size, dtype=torch.float32 + ) + else: + if self.config.tie_word_embeddings: + lm_head_weight_key = f"{prefix}embed_tokens.weight" + lm_head_bias_key = f"{prefix}embed_tokens.bias" + else: + lm_head_weight_key = "lm_head.weight" + lm_head_bias_key = "lm_head.bias" + temp_state_dict = { + **temp_state_dict, + **{ + "lm_head.weight": state_dict.pop(lm_head_weight_key), + "norm.weight": state_dict.pop(f"{prefix}layer_norm.weight").to( + torch.float32 + ), + "lm_head.bias": state_dict.pop( + lm_head_bias_key, + torch.zeros(self.config.vocab_size, dtype=self.dtype), + ), + }, + } + if self.config.norm == "LayerNorm": + temp_state_dict["norm.bias"] = state_dict.pop( + f"{prefix}layer_norm.bias", + torch.zeros(self.config.hidden_size, dtype=self.dtype), + ).to(torch.float32) + + print(f"Loading weights for chunk {self.chunk_idx}") + if temp_state_dict.keys() != self.state_dict().keys(): + temp_state_dict_only_keys = [ + x for x in temp_state_dict.keys() if x not in self.state_dict().keys() + ] + model_only_keys = [ + x for x in self.state_dict().keys() if x not in temp_state_dict.keys() + ] + raise RuntimeError( + f"model state dict keys don't match with state_dict to load into model.\nModel only keys:{model_only_keys}\nstate_dict only keys:{temp_state_dict_only_keys}" + ) + self.load_state_dict(temp_state_dict) + for i in range(self.num_blocks): + self.layers[i].to(self.device_list[i]) + if self.config.use_stable_embedding and self.chunk_idx == 0: + self.embed_layer_norm.to(self.device_list[0]) + if self.include_tail: + self.norm.to(self.device_list[-1]) + self.lm_head.to(self.device_list[-1]) + self.eval() + + return self + + def get_example_inputs( + self, num_token: int = 128, cache_size: int = 512, get_dym_shape=False + ): + head_dim = int(self.head_dim) + example_inputs = ( + torch.randn( + 1, num_token, self.config.hidden_size, device="cpu", dtype=torch.float32 + ), + torch.randn( + 1, + 1, + num_token, + cache_size + num_token, + device="cpu", + dtype=torch.float32, + ), + torch.randn( + 1, num_token, self.config.hidden_size, device="cpu", dtype=torch.float32 + ), + torch.randn( + 2 * self.num_blocks, + self.config.num_key_value_heads, + 1500, + head_dim, + device="cpu", + dtype=torch.float32, + ), + *[ + torch.randn( + 1, + self.config.num_key_value_heads, + cache_size, + head_dim, + device="cpu", + dtype=torch.float32, + ) + for _ in range(2 * self.num_blocks) + ], + ) + # Specify dims that would be dynamic during calibration + # Note: Assume cache size fixed shape as torch dynamic shape cannot handle dim 3 being + # combination of 2 dynamic dims + if get_dym_shape: + nt = Dim("num_token", max=num_token) + cache_dims = tuple(({} for _ in range(2 * self.num_blocks))) + dynamic_shapes = ( + {0: Dim.STATIC, 1: nt, 2: Dim.STATIC}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: nt, 3: nt + cache_size}, + {0: Dim.STATIC, 1: nt, 2: Dim.STATIC}, + {0: Dim.STATIC, 1: Dim.STATIC, 2: Dim.STATIC, 3: Dim.STATIC}, + cache_dims, + ) + return example_inputs, dynamic_shapes + + return example_inputs diff --git a/examples/mediatek/models/llm_models/weights/Qwen2-1.5B-Instruct/config.json b/examples/mediatek/models/llm_models/weights/Qwen2-1.5B-Instruct/config.json new file mode 100755 index 00000000000..8b9b7f44417 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen2-1.5B-Instruct/config.json @@ -0,0 +1,28 @@ +{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 1536, + "initializer_range": 0.02, + "intermediate_size": 8960, + "max_position_embeddings": 32768, + "max_window_layers": 28, + "model_type": "qwen2", + "num_attention_heads": 12, + "num_hidden_layers": 28, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.40.1", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936, + "mask_value": -10000 +} diff --git a/examples/mediatek/models/llm_models/weights/Qwen2-7B-Instruct/config.json b/examples/mediatek/models/llm_models/weights/Qwen2-7B-Instruct/config.json new file mode 100644 index 00000000000..ae681edbd94 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen2-7B-Instruct/config.json @@ -0,0 +1,27 @@ +{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "initializer_range": 0.02, + "intermediate_size": 18944, + "max_position_embeddings": 32768, + "max_window_layers": 28, + "model_type": "qwen2", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 131072, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.2", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 152064 +} diff --git a/examples/mediatek/models/llm_models/weights/Qwen2.5-0.5B-Instruct/config.json b/examples/mediatek/models/llm_models/weights/Qwen2.5-0.5B-Instruct/config.json new file mode 100644 index 00000000000..ca193797a82 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen2.5-0.5B-Instruct/config.json @@ -0,0 +1,28 @@ +{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 896, + "initializer_range": 0.02, + "intermediate_size": 4864, + "max_position_embeddings": 32768, + "max_window_layers": 21, + "model_type": "qwen2", + "num_attention_heads": 14, + "num_hidden_layers": 24, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.43.1", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936, + "mask_value": -10000 +} diff --git a/examples/mediatek/models/llm_models/weights/Qwen2.5-3B/config.json b/examples/mediatek/models/llm_models/weights/Qwen2.5-3B/config.json new file mode 100644 index 00000000000..acd0065677f --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen2.5-3B/config.json @@ -0,0 +1,28 @@ +{ + "architectures": [ + "Qwen2ForCausalLM" + ], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151643, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 32768, + "max_window_layers": 36, + "model_type": "qwen2", + "num_attention_heads": 16, + "num_hidden_layers": 36, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.40.1", + "use_cache": true, + "use_mrope": false, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/examples/mediatek/models/llm_models/weights/Qwen3-1.7B/config.json b/examples/mediatek/models/llm_models/weights/Qwen3-1.7B/config.json new file mode 100644 index 00000000000..d0f5de97999 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen3-1.7B/config.json @@ -0,0 +1,31 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 6144, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 16, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936, + "tokenizer": "qwen2" +} \ No newline at end of file diff --git a/examples/mediatek/models/llm_models/weights/Qwen3-4B/config.json b/examples/mediatek/models/llm_models/weights/Qwen3-4B/config.json new file mode 100644 index 00000000000..e49eccdc32f --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/Qwen3-4B/config.json @@ -0,0 +1,30 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "max_position_embeddings": 40960, + "max_window_layers": 36, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} \ No newline at end of file diff --git a/examples/mediatek/models/llm_models/weights/gemma2_2b_it/config.json b/examples/mediatek/models/llm_models/weights/gemma2_2b_it/config.json new file mode 100755 index 00000000000..927a3582851 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/gemma2_2b_it/config.json @@ -0,0 +1,32 @@ +{ + "architectures": [ + "Gemma2ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "attn_logit_softcapping": 50.0, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "head_dim": 256, + "hidden_act": "gelu_pytorch_tanh", + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 2304, + "initializer_range": 0.02, + "intermediate_size": 9216, + "max_position_embeddings": 8192, + "model_type": "gemma2", + "num_attention_heads": 8, + "num_hidden_layers": 26, + "num_key_value_heads": 4, + "pad_token_id": 0, + "query_pre_attn_scalar": 256, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.4", + "use_cache": true, + "vocab_size": 256000 +} diff --git a/examples/mediatek/models/llm_models/weights/gemma3_1b_it/config.json b/examples/mediatek/models/llm_models/weights/gemma3_1b_it/config.json new file mode 100755 index 00000000000..e5883df9b2a --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/gemma3_1b_it/config.json @@ -0,0 +1,42 @@ +{ + "architectures": [ + "Gemma3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "attn_logit_softcapping": null, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": [ + 1, + 106 + ], + "rope_scaling": { + "factor":1.0, + "type": "linear" + }, + "final_logit_softcapping": null, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 1152, + "initializer_range": 0.02, + "intermediate_size": 6912, + "max_position_embeddings": 32768, + "model_type": "gemma3", + "num_attention_heads": 4, + "num_hidden_layers": 26, + "num_key_value_heads": 1, + "pad_token_id": 0, + "query_pre_attn_scalar": 256, + "rms_norm_eps": 1e-06, + "rope_local_base_freq": 10000, + "rope_theta": 1000000, + "sliding_window_attention_size": 512, + "sliding_window_pattern": 6, + "global_local_attention_pattern": ["SLIDING_LOCAL","SLIDING_LOCAL","SLIDING_LOCAL","SLIDING_LOCAL","SLIDING_LOCAL","GLOBAL"], + "torch_dtype": "bfloat16", + "transformers_version": "4.50.0.dev0", + "use_cache": true, + "vocab_size": 262144, + "use_res_clamp": false +} \ No newline at end of file diff --git a/examples/mediatek/models/llm_models/weights/phi-4-mini-reasoning/config.json b/examples/mediatek/models/llm_models/weights/phi-4-mini-reasoning/config.json new file mode 100755 index 00000000000..57bcfc4af8a --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/phi-4-mini-reasoning/config.json @@ -0,0 +1,143 @@ +{ + "architectures": [ + "Phi3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi3.Phi3Config", + "AutoModelForCausalLM": "modeling_phi3.Phi3ForCausalLM", + "AutoTokenizer": "Xenova/gpt-4o" + }, + "bos_token_id": 199999, + "embd_pdrop": 0.0, + "eos_token_id": 199999, + "full_attn_mod": 1, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 8192, + "interpolate_factor": 1, + "lm_head_bias": false, + "max_position_embeddings": 131072, + "mlp_bias": false, + "model_type": "phi4", + "num_attention_heads": 24, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "original_max_position_embeddings": 4096, + "pad_token_id": 199999, + "partial_rotary_factor": 0.75, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "long_factor": [ + 1, + 1.118320672, + 1.250641126, + 1.398617824, + 1.564103225, + 1.74916897, + 1.956131817, + 2.187582649, + 2.446418898, + 2.735880826, + 3.059592084, + 3.421605075, + 3.826451687, + 4.279200023, + 4.785517845, + 5.351743533, + 5.984965424, + 6.693110555, + 7.485043894, + 8.370679318, + 9.36110372, + 10.4687158, + 11.70738129, + 13.09260651, + 14.64173252, + 16.37415215, + 18.31155283, + 20.47818807, + 22.90118105, + 25.61086418, + 28.64115884, + 32.03, + 32.1, + 32.13, + 32.23, + 32.6, + 32.61, + 32.64, + 32.66, + 32.7, + 32.71, + 32.93, + 32.97, + 33.28, + 33.49, + 33.5, + 44.16, + 47.77 + ], + "short_factor": [ + 1, + 1.118320672, + 1.250641126, + 1.398617824, + 1.564103225, + 1.74916897, + 1.956131817, + 2.187582649, + 2.446418898, + 2.735880826, + 3.059592084, + 3.421605075, + 3.826451687, + 4.279200023, + 4.785517845, + 5.351743533, + 5.984965424, + 6.693110555, + 7.485043894, + 8.370679318, + 9.36110372, + 10.4687158, + 11.70738129, + 13.09260651, + 14.64173252, + 16.37415215, + 18.31155283, + 20.47818807, + 22.90118105, + 25.61086418, + 28.64115884, + 32.03, + 32.1, + 32.13, + 32.23, + 32.6, + 32.61, + 32.64, + 32.66, + 32.7, + 32.71, + 32.93, + 32.97, + 33.28, + 33.49, + 33.5, + 44.16, + 47.77 + ], + "type": "longrope" + }, + "rope_theta": 10000.0, + "sliding_window": 262144, + "tie_word_embeddings": true, + "torch_dtype": "bfloat16", + "transformers_version": "4.50.0", + "use_cache": false, + "vocab_size": 200064 +} diff --git a/examples/mediatek/models/llm_models/weights/phi3.5-mini-instruct/config.json b/examples/mediatek/models/llm_models/weights/phi3.5-mini-instruct/config.json new file mode 100644 index 00000000000..bea4d195d75 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/phi3.5-mini-instruct/config.json @@ -0,0 +1,138 @@ +{ + "_name_or_path": "Phi-3.5-mini-instruct", + "architectures": [ + "Phi3ForCausalLM" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi3.Phi3Config", + "AutoModelForCausalLM": "modeling_phi3.Phi3ForCausalLM" + }, + "bos_token_id": 1, + "embd_pdrop": 0.0, + "eos_token_id": 32000, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 4096, + "model_type": "phi3", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "original_max_position_embeddings": 4096, + "pad_token_id": 32000, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "long_factor": [ + 1.0800000429153442, + 1.1100000143051147, + 1.1399999856948853, + 1.340000033378601, + 1.5899999141693115, + 1.600000023841858, + 1.6200000047683716, + 2.620000123977661, + 3.2300000190734863, + 3.2300000190734863, + 4.789999961853027, + 7.400000095367432, + 7.700000286102295, + 9.09000015258789, + 12.199999809265137, + 17.670000076293945, + 24.46000099182129, + 28.57000160217285, + 30.420001983642578, + 30.840002059936523, + 32.590003967285156, + 32.93000411987305, + 42.320003509521484, + 44.96000289916992, + 50.340003967285156, + 50.45000457763672, + 57.55000305175781, + 57.93000411987305, + 58.21000289916992, + 60.1400032043457, + 62.61000442504883, + 62.62000274658203, + 62.71000289916992, + 63.1400032043457, + 63.1400032043457, + 63.77000427246094, + 63.93000411987305, + 63.96000289916992, + 63.970001220703125, + 64.02999877929688, + 64.06999969482422, + 64.08000183105469, + 64.12000274658203, + 64.41000366210938, + 64.4800033569336, + 64.51000213623047, + 64.52999877929688, + 64.83999633789062 + ], + "short_factor": [ + 1.0, + 1.0199999809265137, + 1.0299999713897705, + 1.0299999713897705, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0699999332427979, + 1.0999999046325684, + 1.1099998950958252, + 1.1599998474121094, + 1.1599998474121094, + 1.1699998378753662, + 1.2899998426437378, + 1.339999794960022, + 1.679999828338623, + 1.7899998426437378, + 1.8199998140335083, + 1.8499997854232788, + 1.8799997568130493, + 1.9099997282028198, + 1.9399996995925903, + 1.9899996519088745, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0799996852874756, + 2.0899996757507324, + 2.189999580383301, + 2.2199995517730713, + 2.5899994373321533, + 2.729999542236328, + 2.749999523162842, + 2.8399994373321533 + ], + "type": "longrope" + }, + "rope_theta": 10000.0, + "sliding_window": 262144, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.43.3", + "use_cache": true, + "attention_bias": false, + "vocab_size": 32064 +} diff --git a/examples/mediatek/models/llm_models/weights/whisper-large-v3-turbo/config_refactor.json b/examples/mediatek/models/llm_models/weights/whisper-large-v3-turbo/config_refactor.json new file mode 100644 index 00000000000..57bb3956042 --- /dev/null +++ b/examples/mediatek/models/llm_models/weights/whisper-large-v3-turbo/config_refactor.json @@ -0,0 +1,316 @@ +{ + "preprocessor": { + "model_type": "whisper", + "feature_size": 128, + "sampling_rate": 16000 + }, + "pre_encoder_hook": { + "name": "numpy_to_torch" + }, + "encoder": { + "model_type": "whisper", + "num_mel_bins": 128, + "encoder_ffn_dim": 5120, + "encoder_layers": 32, + "encoder_attention_heads": 20, + "d_model": 1280, + "decoder_layers": 4, + "max_source_positions": 1500, + "max_position_embeddings": 448 + }, + "llm": { + "model_type": "whisper_decoder", + "bos_token_id": 50257, + "eos_token_id": 50257, + "hidden_size": 1280, + "intermediate_size": 5120, + "max_position_embeddings": 448, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "num_hidden_layers": 4, + "pad_token_id": 50256, + "vocab_size": 51866, + "decoder_start_token_id": 50258, + "begin_suppress_tokens": [ + 220, + 50256 + ], + "mask_feature_length": 10, + "mask_feature_min_masks": 0, + "mask_feature_prob": 0.0, + "mask_time_length": 10, + "mask_time_min_masks": 2, + "mask_time_prob": 0.05, + "max_length": 448, + "max_source_positions": 1500, + "max_target_positions": 448, + "median_filter_width": 7, + "num_mel_bins": 128 + }, + "logits_processor":{ + "name": "whisper", + "alignment_heads": [ + [ + 7, + 0 + ], + [ + 10, + 17 + ], + [ + 12, + 18 + ], + [ + 13, + 12 + ], + [ + 16, + 1 + ], + [ + 17, + 14 + ], + [ + 19, + 11 + ], + [ + 21, + 4 + ], + [ + 24, + 1 + ], + [ + 25, + 6 + ] + ], + "begin_suppress_tokens": [ + 220, + 50257 + ], + "bos_token_id": 50257, + "decoder_start_token_id": 50258, + "eos_token_id": 50257, + "forced_decoder_ids": [ + [ + 1, + null + ], + [ + 2, + 50360 + ] + ], + "is_multilingual": true, + "lang_to_id": { + "<|af|>": 50327, + "<|am|>": 50334, + "<|ar|>": 50272, + "<|as|>": 50350, + "<|az|>": 50304, + "<|ba|>": 50355, + "<|be|>": 50330, + "<|bg|>": 50292, + "<|bn|>": 50302, + "<|bo|>": 50347, + "<|br|>": 50309, + "<|bs|>": 50315, + "<|ca|>": 50270, + "<|cs|>": 50283, + "<|cy|>": 50297, + "<|da|>": 50285, + "<|de|>": 50261, + "<|el|>": 50281, + "<|en|>": 50259, + "<|es|>": 50262, + "<|et|>": 50307, + "<|eu|>": 50310, + "<|fa|>": 50300, + "<|fi|>": 50277, + "<|fo|>": 50338, + "<|fr|>": 50265, + "<|gl|>": 50319, + "<|gu|>": 50333, + "<|haw|>": 50352, + "<|ha|>": 50354, + "<|he|>": 50279, + "<|hi|>": 50276, + "<|hr|>": 50291, + "<|ht|>": 50339, + "<|hu|>": 50286, + "<|hy|>": 50312, + "<|id|>": 50275, + "<|is|>": 50311, + "<|it|>": 50274, + "<|ja|>": 50266, + "<|jw|>": 50356, + "<|ka|>": 50329, + "<|kk|>": 50316, + "<|km|>": 50323, + "<|kn|>": 50306, + "<|ko|>": 50264, + "<|la|>": 50294, + "<|lb|>": 50345, + "<|ln|>": 50353, + "<|lo|>": 50336, + "<|lt|>": 50293, + "<|lv|>": 50301, + "<|mg|>": 50349, + "<|mi|>": 50295, + "<|mk|>": 50308, + "<|ml|>": 50296, + "<|mn|>": 50314, + "<|mr|>": 50320, + "<|ms|>": 50282, + "<|mt|>": 50343, + "<|my|>": 50346, + "<|ne|>": 50313, + "<|nl|>": 50271, + "<|nn|>": 50342, + "<|no|>": 50288, + "<|oc|>": 50328, + "<|pa|>": 50321, + "<|pl|>": 50269, + "<|ps|>": 50340, + "<|pt|>": 50267, + "<|ro|>": 50284, + "<|ru|>": 50263, + "<|sa|>": 50344, + "<|sd|>": 50332, + "<|si|>": 50322, + "<|sk|>": 50298, + "<|sl|>": 50305, + "<|sn|>": 50324, + "<|so|>": 50326, + "<|sq|>": 50317, + "<|sr|>": 50303, + "<|su|>": 50357, + "<|sv|>": 50273, + "<|sw|>": 50318, + "<|ta|>": 50287, + "<|te|>": 50299, + "<|tg|>": 50331, + "<|th|>": 50289, + "<|tk|>": 50341, + "<|tl|>": 50348, + "<|tr|>": 50268, + "<|tt|>": 50351, + "<|uk|>": 50280, + "<|ur|>": 50290, + "<|uz|>": 50337, + "<|vi|>": 50278, + "<|yi|>": 50335, + "<|yo|>": 50325, + "<|yue|>": 50358, + "<|zh|>": 50260 + }, + "max_initial_timestamp_index": 50, + "max_length": 448, + "no_timestamps_token_id": 50364, + "pad_token_id": 50257, + "prev_sot_token_id": 50362, + "return_timestamps": false, + "suppress_tokens": [ + 1, + 2, + 7, + 8, + 9, + 10, + 14, + 25, + 26, + 27, + 28, + 29, + 31, + 58, + 59, + 60, + 61, + 62, + 63, + 90, + 91, + 92, + 93, + 359, + 503, + 522, + 542, + 873, + 893, + 902, + 918, + 922, + 931, + 1350, + 1853, + 1982, + 2460, + 2627, + 3246, + 3253, + 3268, + 3536, + 3846, + 3961, + 4183, + 4667, + 6585, + 6647, + 7273, + 9061, + 9383, + 10428, + 10929, + 11938, + 12033, + 12331, + 12562, + 13793, + 14157, + 14635, + 15265, + 15618, + 16553, + 16604, + 18362, + 18956, + 20075, + 21675, + 22520, + 26130, + 26161, + 26435, + 28279, + 29464, + 31650, + 32302, + 32470, + 36865, + 42863, + 47425, + 49870, + 50254, + 50258, + 50359, + 50360, + 50361, + 50362, + 50363 + ], + "task_to_id": { + "transcribe": 50360, + "translate": 50359 + }, + "transformers_version": "4.36.0.dev0" + } +} diff --git a/examples/mediatek/shell_scripts/export_gemma.sh b/examples/mediatek/shell_scripts/export_gemma.sh new file mode 100644 index 00000000000..776a9f187d6 --- /dev/null +++ b/examples/mediatek/shell_scripts/export_gemma.sh @@ -0,0 +1,44 @@ +model=${1:-'gemma2'} +chunks=${2:-4} +tok=${3:-128} +cache=${4:-512} +cal=${5:-None} +pres=${6:-A16W4} +plat=${7:-DX4} + +if [ $model = "gemma2" ] +then + config_path=gemma2_2b_it/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/gemma.json" +elif [ $model = "gemma3" ] +then + config_path=gemma3_1b_it/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/gemma.json" +fi + +if [ $cal = "None" ] +then + data="" +else + data="-d aot_utils/llm_utils/prompts/${cal}" +fi + +echo "Model: $model" +echo "Config Path: $config_path" +echo "Num Chunks: $chunks" +echo "Num Tokens: $tok" +echo "Cache Size: $cache" +echo "Precision: $pres" +echo "Calibration Dataset: $cal" +echo "Preformatter: $pref" +echo "Platform: $plat" + +python3 model_export_scripts/gemma.py \ + models/llm_models/weights/${config_path} \ + -p $pres \ + --num_chunks $chunks \ + ${data} \ + ${pref} \ + -shapes ${tok}t${cache}c 1t${cache}c \ + --platform $plat + \ No newline at end of file diff --git a/examples/mediatek/shell_scripts/export_llama.sh b/examples/mediatek/shell_scripts/export_llama.sh index d92e6198b4a..74e857cb008 100644 --- a/examples/mediatek/shell_scripts/export_llama.sh +++ b/examples/mediatek/shell_scripts/export_llama.sh @@ -4,6 +4,7 @@ tok=${3:-128} cache=${4:-512} cal=${5:-None} pres=${6:-A16W4} +plat=${7:-DX4} if [ $model = "llama3.2-3b" ] then @@ -38,6 +39,7 @@ echo "Cache Size: $cache" echo "Precision: $pres" echo "Calibration Dataset: $cal" echo "Preformatter: $pref" +echo "Platform: $plat" python3 model_export_scripts/llama.py \ models/llm_models/weights/${config_path} \ @@ -45,4 +47,6 @@ python3 model_export_scripts/llama.py \ --num_chunks $chunks \ ${data} \ ${pref} \ - -shapes ${tok}t${cache}c 1t${cache}c + -shapes ${tok}t${cache}c 1t${cache}c \ + --platform $plat + \ No newline at end of file diff --git a/examples/mediatek/shell_scripts/export_oss.sh b/examples/mediatek/shell_scripts/export_oss.sh index 3e00506c05a..5be5736f169 100755 --- a/examples/mediatek/shell_scripts/export_oss.sh +++ b/examples/mediatek/shell_scripts/export_oss.sh @@ -41,4 +41,10 @@ then elif [ $model = "emformer_rnnt" ] then python3 model_export_scripts/emformer_rnnt.py +elif [ $model = "bert" ] +then + python3 model_export_scripts/bert.py +elif [ $model = "distilbert" ] +then + python3 model_export_scripts/distilbert.py fi diff --git a/examples/mediatek/shell_scripts/export_phi.sh b/examples/mediatek/shell_scripts/export_phi.sh new file mode 100644 index 00000000000..e009bcb211b --- /dev/null +++ b/examples/mediatek/shell_scripts/export_phi.sh @@ -0,0 +1,44 @@ +model=${1:-'phi4'} +chunks=${2:-4} +tok=${3:-128} +cache=${4:-512} +cal=${5:-None} +pres=${6:-A16W4} +plat=${7:-DX4} + +if [ $model = "phi3.5" ] +then + config_path=phi3.5-mini-instruct/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/phi3.json" +elif [ $model = "phi4" ] +then + config_path=phi-4-mini-reasoning/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/phi4_reasoning.json" +fi + +if [ $cal = "None" ] +then + data="" +else + data="-d aot_utils/llm_utils/prompts/${cal}" +fi + +echo "Model: $model" +echo "Config Path: $config_path" +echo "Num Chunks: $chunks" +echo "Num Tokens: $tok" +echo "Cache Size: $cache" +echo "Precision: $pres" +echo "Calibration Dataset: $cal" +echo "Preformatter: $pref" +echo "Platform: $plat" + +python3 model_export_scripts/phi.py \ + models/llm_models/weights/${config_path} \ + -p $pres \ + --num_chunks $chunks \ + ${data} \ + ${pref} \ + -shapes ${tok}t${cache}c 1t${cache}c \ + --platform $plat + \ No newline at end of file diff --git a/examples/mediatek/shell_scripts/export_qwen.sh b/examples/mediatek/shell_scripts/export_qwen.sh new file mode 100755 index 00000000000..37db56969c7 --- /dev/null +++ b/examples/mediatek/shell_scripts/export_qwen.sh @@ -0,0 +1,60 @@ +model=${1:-'Qwen3-4B'} +chunks=${2:-4} +tok=${3:-128} +cache=${4:-512} +cal=${5:-None} +pres=${6:-A16W4} +plat=${7:-DX4} + +if [ $model = "Qwen3-4B" ] +then + config_path=Qwen3-4B/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen3.json" +elif [ $model = "Qwen3-1.7B" ] +then + config_path=Qwen3-1.7B/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen3.json" +elif [ $model = "Qwen2-7B-Instruct" ] +then + config_path=Qwen2-7B-Instruct/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen.json" +elif [ $model = "Qwen2.5-3B" ] +then + config_path=Qwen2.5-3B/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen.json" +elif [ $model = "Qwen2.5-0.5B-Instruct" ] +then + config_path=Qwen2.5-0.5B-Instruct/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen.json" +elif [ $model = "Qwen2-1.5B-Instruct" ] +then + config_path=Qwen2-1.5B-Instruct/config.json + pref="--preformatter aot_utils/llm_utils/preformatter_templates/qwen.json" +fi + +if [ $cal = "None" ] +then + data="" +else + data="-d aot_utils/llm_utils/prompts/${cal}" +fi + +echo "Model: $model" +echo "Config Path: $config_path" +echo "Num Chunks: $chunks" +echo "Num Tokens: $tok" +echo "Cache Size: $cache" +echo "Precision: $pres" +echo "Calibration Dataset: $cal" +echo "Preformatter: $pref" +echo "Platform: $plat" + +python3 model_export_scripts/qwen.py \ + models/llm_models/weights/${config_path} \ + -p $pres \ + --num_chunks $chunks \ + ${data} \ + ${pref} \ + -shapes ${tok}t${cache}c 1t${cache}c \ + --platform $plat + \ No newline at end of file diff --git a/examples/mediatek/shell_scripts/export_whisper.sh b/examples/mediatek/shell_scripts/export_whisper.sh new file mode 100644 index 00000000000..1a58e220db8 --- /dev/null +++ b/examples/mediatek/shell_scripts/export_whisper.sh @@ -0,0 +1,42 @@ +model=${1:-'whisper-large-v3-turbo'} +chunks=${2:-1} +tok=${3:-10} +cache=${4:-512} +cal=${5:-None} +pres=${6:-A16W4} +plat=${7:-DX4} + +if [ $model = "whisper-large-v3-turbo" ] +then + config_path=whisper-large-v3-turbo/config_refactor.json + +elif [ $model = "whisper-large-v3" ] +then + config_path=whisper-large-v3/config_refactor.json + +fi + +if [ $cal = "None" ] +then + data="" +else + data="-d aot_utils/mllm_utils/audio/${cal}" +fi + +echo "Model: $model" +echo "Config Path: $config_path" +echo "Num Chunks: $chunks" +echo "Num Tokens: $tok" +echo "Cache Size: $cache" +echo "Precision: $pres" +echo "Calibration Dataset: $cal" +echo "Preformatter: $pref" +echo "Platform: $plat" + +python3 model_export_scripts/whisper.py \ + models/llm_models/weights/${config_path} \ + -p $pres \ + --num_chunks $chunks \ + ${data} \ + -shapes ${tok}t${cache}c 1t${cache}c \ + --platform $plat diff --git a/examples/models/__init__.py b/examples/models/__init__.py index 76469846608..b997ba1bf6e 100644 --- a/examples/models/__init__.py +++ b/examples/models/__init__.py @@ -1,6 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. -# Copyright 2024 Arm Limited and/or its affiliates. +# Copyright 2024-2025 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. @@ -35,14 +35,16 @@ class Model(str, Enum): ResNet50 = "resnet50" Llava = "llava" EfficientSam = "efficient_sam" - Qwen25 = "qwen2_5" + Qwen25 = "qwen2_5_1_5b" Phi4Mini = "phi_4_mini" + SmolLM2 = "smollm2" def __str__(self) -> str: return self.value class Backend(str, Enum): + Xnnpack = "xnnpack" XnnpackQuantizationDelegation = "xnnpack-quantization-delegation" CoreMlExportOnly = "coreml" CoreMlExportAndTest = "coreml-test" # AOT export + test with runner @@ -82,6 +84,7 @@ def __str__(self) -> str: str(Model.EfficientSam): ("efficient_sam", "EfficientSAM"), str(Model.Qwen25): ("qwen2_5", "Qwen2_5Model"), str(Model.Phi4Mini): ("phi_4_mini", "Phi4MiniModel"), + str(Model.SmolLM2): ("smollm2", "SmolLM2Model"), } __all__ = [ diff --git a/examples/models/checkpoint.py b/examples/models/checkpoint.py index 57a5b0ffaca..4583b105732 100644 --- a/examples/models/checkpoint.py +++ b/examples/models/checkpoint.py @@ -1,11 +1,14 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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-unsafe +import json +import os from pathlib import Path from typing import Any, Dict, Optional @@ -16,7 +19,7 @@ def get_default_model_resource_dir(model_file_path: str) -> Path: """ Get the default path to resouce files (which contain files such as the checkpoint and param files), either: - 1. Uses the path from pkg_resources, only works with buck2 + 1. Uses the path from importlib.resources, only works with buck2 2. Uses default path located in examples/models/llama/params Expected to be called from with a `model.py` file located in a @@ -33,22 +36,22 @@ def get_default_model_resource_dir(model_file_path: str) -> Path: """ try: - import pkg_resources + import importlib.resources as _resources - # 1st way: If we can import this path, we are running with buck2 and all resources can be accessed with pkg_resources. + # 1st way: If we can import this path, we are running with buck2 and all resources can be accessed with importlib.resources. # pyre-ignore from executorch.examples.models.llama import params # noqa # Get the model name from the cwd, assuming that this module is called from a path such as # examples/models//model.py. model_name = Path(model_file_path).parent.name - resource_dir = Path( - pkg_resources.resource_filename( - f"executorch.examples.models.{model_name}", "params" - ) - ) - except: - # 2nd way. + model_dir = _resources.files(f"executorch.examples.models.{model_name}") + with _resources.as_file(model_dir) as model_path: + resource_dir = model_path / "params" + assert resource_dir.exists() + + except Exception: + # 2nd way: resource_dir = Path(model_file_path).absolute().parent / "params" return resource_dir @@ -73,3 +76,39 @@ def get_checkpoint_dtype(checkpoint: Dict[str, Any]) -> Optional[torch.dtype]: f"Mixed dtype model. Dtype of {first_key}: {first.dtype}. Mismatches in the checkpoint: {mismatched_dtypes}" ) return dtype + + +def load_checkpoint_from_pytorch_model(input_dir: str) -> Dict: + index_path = os.path.join(input_dir, "pytorch_model.bin.index.json") + if os.path.exists(index_path): + # Sharded checkpoint. + with open(index_path, "r") as f: + index = json.load(f) + weight_map = index["weight_map"] + checkpoint_shards = sorted(set(weight_map.values())) + + # Load all the shards into memory + shard_to_weights = {} + for shard in checkpoint_shards: + shard_to_weights[shard] = torch.load( + os.path.join(input_dir, shard), + weights_only=True, + map_location=torch.device("cpu"), + ) + + # Merge tensors into consolidated state dict. + merged_state_dict = {} + for weight_name, shard in weight_map.items(): + tensor = shard_to_weights[shard][weight_name] + merged_state_dict[weight_name] = tensor + return merged_state_dict + + # Single checkpoint + model_path = os.path.join(input_dir, "pytorch_model.bin") + if os.path.exists(model_path): + state_dict = torch.load( + model_path, weights_only=True, map_location=torch.device("cpu") + ) + return state_dict + + raise FileNotFoundError(f"Could not find pytorch_model checkpoint in {input_dir}") diff --git a/examples/models/gemma3/__init__.py b/examples/models/gemma3/__init__.py new file mode 100644 index 00000000000..ae34db47954 --- /dev/null +++ b/examples/models/gemma3/__init__.py @@ -0,0 +1,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 executorch.examples.models.gemma3.convert_weights import convert_weights +from executorch.examples.models.llama.model import Llama2Model + + +class Gemma3Model(Llama2Model): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + +__all__ = [ + "Gemma3Model", + "convert_weights", +] diff --git a/examples/models/gemma3/config/1b_config.json b/examples/models/gemma3/config/1b_config.json new file mode 100644 index 00000000000..3a9e673716b --- /dev/null +++ b/examples/models/gemma3/config/1b_config.json @@ -0,0 +1,23 @@ +{ + "dim": 1152, + "ffn_dim_multiplier": 1, + "hidden_dim": 6912, + "n_heads": 4, + "head_dim": 256, + "n_kv_heads": 1, + "n_layers": 26, + "act_fn": "gelu_approx", + "norm_type": "gemma3", + "norm_eps": 1e-06, + "post_attention_norm": true, + "post_ffn_norm": true, + "rope_theta": 1000000.0, + "use_scaled_rope": false, + "apply_embedding": true, + "embedding_scale_factor": 33.941125497, + "vocab_size": 262144, + "use_hf_rope": true, + "attention_qkv_bias": false, + "use_qk_norm": true, + "qk_norm_before_rope": true +} diff --git a/examples/models/gemma3/convert_weights.py b/examples/models/gemma3/convert_weights.py new file mode 100644 index 00000000000..ed44b9eb1cc --- /dev/null +++ b/examples/models/gemma3/convert_weights.py @@ -0,0 +1,110 @@ +import argparse + +import json +import os +from typing import Dict + +import torch +from safetensors.torch import load_file + +from torchtune.models.convert_weights import get_mapped_key + + +# Weight mappings from Gemma 3's checkpoint to ExecuTorch's transformer parameters. +_GEMMA3_TO_EXECUTORCH = { + "model.embed_tokens.weight": "tok_embeddings.weight", + "model.norm.weight": "norm.weight", + "model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight", + "model.layers.{}.self_attn.k_norm.weight": "layers.{}.attention.k_norm_fn.weight", + "model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight", + "model.layers.{}.self_attn.q_norm.weight": "layers.{}.attention.q_norm_fn.weight", + "model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight", + "model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight", + "model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight", + "model.layers.{}.post_attention_layernorm.weight": "layers.{}.post_attention_norm.weight", + "model.layers.{}.pre_feedforward_layernorm.weight": "layers.{}.ffn_norm.weight", + "model.layers.{}.post_feedforward_layernorm.weight": "layers.{}.post_ffn_norm.weight", + "model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight", + "model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight", + "model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight", +} + + +def gemma3_to_executorch( + state_dict: Dict[str, torch.Tensor] +) -> Dict[str, torch.Tensor]: + """ + Convert the state dict so that it matches what ExecuTorch's transformer definition expects. + """ + converted_state_dict = {} + for key, value in state_dict.items(): + new_key = get_mapped_key(key, _GEMMA3_TO_EXECUTORCH) + converted_state_dict[new_key] = value + converted_state_dict["output.weight"] = converted_state_dict[ + "tok_embeddings.weight" + ] + return converted_state_dict + + +def load_checkpoint_from_safetensors(input_dir: str) -> Dict: + index_path = os.path.join(input_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + # Sharded checkpoint. + with open(index_path, "r") as f: + index = json.load(f) + weight_map = index["weight_map"] + checkpoint_shards = sorted(set(weight_map.values())) + + # Load all the shards into memory + shard_to_weights = {} + for shard in checkpoint_shards: + shard_to_weights[shard] = load_file(os.path.join(input_dir, shard)) + + # Merge tensors into consolidated state dict. + merged_state_dict = {} + for weight_name, shard in weight_map.items(): + tensor = shard_to_weights[shard][weight_name] + merged_state_dict[weight_name] = tensor + return merged_state_dict + else: + # Single checkpoint. + state_dict = load_file(os.path.join(input_dir, "model.safetensors")) + return state_dict + + +def load_checkpoint(input_dir: str) -> Dict: + pytorch_path = os.path.join(input_dir, "pytorch_model.bin") + if os.path.exists(pytorch_path): + print("Loading checkpoint from PyTorch .bin file") + return torch.load(pytorch_path, map_location="cpu", weights_only=True) + print("Loading checkpoint from safetensors directory") + return load_checkpoint_from_safetensors(input_dir) + + +def convert_weights(input_dir: str, output_file: str) -> None: + print("Loading checkpoint...") + sd = load_checkpoint(input_dir) + print("Converting checkpoint...") + sd = gemma3_to_executorch(sd) + print("Saving checkpoint...") + torch.save(sd, output_file) + print("Done.") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert Gemma3 weights to ExecuTorch transformer format." + ) + parser.add_argument( + "input_dir", + type=str, + help="Path to directory containing safetensor checkpoint files, or PyTorch checkpoint file.", + ) + parser.add_argument("output", type=str, help="Path to the output checkpoint") + + args = parser.parse_args() + convert_weights(args.input_dir, args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/models/lfm2/README.md b/examples/models/lfm2/README.md new file mode 100644 index 00000000000..8f2a84be376 --- /dev/null +++ b/examples/models/lfm2/README.md @@ -0,0 +1,62 @@ +## Summary +[LFM2](https://huggingface.co/collections/LiquidAI/lfm2-686d721927015b2ad73eaa38) is a new generation of hybrid models developed by [Liquid AI](https://www.liquid.ai/) and available in 3 variants - 350M, 700M, 1.2B. + +## Instructions + +LFM2 uses the same example code as optimized Llama model, while the checkpoint, model params, and tokenizer are different. Please see the [Llama README page](../llama/README.md) for details. +LFM2 is a hybrid model, where some attention layers are replaced with short convolutions. + +### Example export +Here is a basic example for exporting LFM2, although please refer to the Llama README's [Step 2: Prepare model](../llama/README.md#step-2-prepare-model) for more advanced usage. + +Export 350m to XNNPack, quantized with 8da4w: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml \ + +base.model_class="lfm2_350m" \ + +base.params="examples/models/lfm2/config/lfm2_350m_config.json" \ + +export.output_name="lfm2_350m_8da4w.pte" +``` + +Export 700m to XNNPack, quantized with 8da4w: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml \ + +base.model_class="lfm2_700m" \ + +base.params="examples/models/lfm2/config/lfm2_700m_config.json" \ + +export.output_name="lfm2_700m_8da4w.pte" +``` + +Export 1_2b to XNNPack, quantized with 8da4w: +``` +python -m extension.llm.export.export_llm \ + --config examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml \ + +base.model_class="lfm2_1_2b" \ + +base.params="examples/models/lfm2/config/lfm2_1_2b_config.json" \ + +export.output_name="lfm2_1_2b_8da4w.pte" +``` +### Example run +With ExecuTorch pybindings: +``` +python -m examples.models.llama.runner.native \ + --model lfm2_700m \ + --pte lfm2_700m_8da4w.pte \ + --tokenizer ~/.cache/huggingface/hub/models--LiquidAI--LFM2-700M/snapshots/ab260293733f05dd4ce22399bea1cae2cf9b272d/tokenizer.json \ + --tokenizer_config ~/.cache/huggingface/hub/models--LiquidAI--LFM2-700M/snapshots/ab260293733f05dd4ce22399bea1cae2cf9b272d/tokenizer_config.json \ + --prompt "<|startoftext|><|im_start|>user\nWho are you?<|im_end|>\n<|im_start|>assistant\n" \ + --params examples/models/lfm2/config/lfm2_700m_config.json \ + --max_len 128 \ + -kv \ + --temperature 0.3 +``` + +With ExecuTorch's sample c++ runner (see the Llama README's [Step 3: Run on your computer to validate](../llama/README.md#step-3-run-on-your-computer-to-validate) to build the runner): +``` +cmake-out/examples/models/llama/llama_main \ + --model_path lfm2_700m_8da4w.pte \ + --tokenizer_path ~/.cache/huggingface/hub/models--LiquidAI--LFM2-700M/snapshots/ab260293733f05dd4ce22399bea1cae2cf9b272d/tokenizer.json \ + --prompt="<|startoftext|><|im_start|>user\nWho are you?<|im_end|>\n<|im_start|>assistant\n" \ + --temperature 0.3 +``` + +To run the model on an example iOS or Android app, see the Llama README's [Step 5: Build Mobile apps](../llama/README.md#step-5-build-mobile-apps) section. diff --git a/examples/models/lfm2/__init__.py b/examples/models/lfm2/__init__.py new file mode 100644 index 00000000000..1efdc55af81 --- /dev/null +++ b/examples/models/lfm2/__init__.py @@ -0,0 +1,10 @@ +# 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.examples.models.lfm2.convert_weights import convert_weights +from executorch.examples.models.lfm2.short_conv import ShortConvBlock + +__all__ = [ + "convert_weights", + "ShortConvBlock", +] diff --git a/examples/models/lfm2/config/lfm2_1_2b_config.json b/examples/models/lfm2/config/lfm2_1_2b_config.json new file mode 100644 index 00000000000..015b6940ed9 --- /dev/null +++ b/examples/models/lfm2/config/lfm2_1_2b_config.json @@ -0,0 +1,34 @@ +{ + "dim": 2048, + "ffn_dim_multiplier": 1, + "hidden_dim": 8192, + "n_heads": 32, + "n_kv_heads": 8, + "n_layers": 16, + "norm_eps": 1e-5, + "rope_theta": 1000000.0, + "use_scaled_rope": false, + "vocab_size": 65536, + "use_hf_rope": true, + "use_qk_norm": true, + "qk_norm_before_rope": true, + "layer_types": [ + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "conv" + ] +} diff --git a/examples/models/lfm2/config/lfm2_350m_config.json b/examples/models/lfm2/config/lfm2_350m_config.json new file mode 100644 index 00000000000..4cd2b2a1830 --- /dev/null +++ b/examples/models/lfm2/config/lfm2_350m_config.json @@ -0,0 +1,34 @@ +{ + "dim": 1024, + "ffn_dim_multiplier": 1, + "hidden_dim": 4608, + "n_heads": 16, + "n_kv_heads": 8, + "n_layers": 16, + "norm_eps": 1e-5, + "rope_theta": 1000000.0, + "use_scaled_rope": false, + "vocab_size": 65536, + "use_hf_rope": true, + "use_qk_norm": true, + "qk_norm_before_rope": true, + "layer_types": [ + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "conv" + ] +} diff --git a/examples/models/lfm2/config/lfm2_700m_config.json b/examples/models/lfm2/config/lfm2_700m_config.json new file mode 100644 index 00000000000..9f3afc0f121 --- /dev/null +++ b/examples/models/lfm2/config/lfm2_700m_config.json @@ -0,0 +1,34 @@ +{ + "dim": 1536, + "ffn_dim_multiplier": 1, + "hidden_dim": 6912, + "n_heads": 24, + "n_kv_heads": 8, + "n_layers": 16, + "norm_eps": 1e-5, + "rope_theta": 1000000.0, + "use_scaled_rope": false, + "vocab_size": 65536, + "use_hf_rope": true, + "use_qk_norm": true, + "qk_norm_before_rope": true, + "layer_types": [ + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "full_attention", + "conv", + "conv" + ] +} diff --git a/examples/models/lfm2/config/lfm2_xnnpack_fp32.yaml b/examples/models/lfm2/config/lfm2_xnnpack_fp32.yaml new file mode 100644 index 00000000000..9dd93821326 --- /dev/null +++ b/examples/models/lfm2/config/lfm2_xnnpack_fp32.yaml @@ -0,0 +1,12 @@ +base: + metadata: '{"get_bos_id": 1, "get_eos_ids":[7]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: fp32 + +backend: + xnnpack: + enabled: True + extended_ops: True diff --git a/examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml b/examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml new file mode 100644 index 00000000000..0925a5989be --- /dev/null +++ b/examples/models/lfm2/config/lfm2_xnnpack_q8da4w.yaml @@ -0,0 +1,15 @@ +base: + metadata: '{"get_bos_id": 1, "get_eos_ids":[7]}' + +model: + use_kv_cache: True + use_sdpa_with_kv_cache: True + dtype_override: fp32 + +quantization: + qmode: 8da4w + +backend: + xnnpack: + enabled: True + extended_ops: True diff --git a/examples/models/lfm2/convert_weights.py b/examples/models/lfm2/convert_weights.py new file mode 100644 index 00000000000..d23f2aa7c49 --- /dev/null +++ b/examples/models/lfm2/convert_weights.py @@ -0,0 +1,74 @@ +import os +from typing import Dict + +import torch +from safetensors.torch import load_file + +from torchtune.models.convert_weights import get_mapped_key + +_LFM_2_TO_META = { + "model.embed_tokens.weight": "tok_embeddings.weight", + "model.embedding_norm.weight": "norm.weight", + "model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight", + "model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight", + "model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight", + "model.layers.{}.self_attn.out_proj.weight": "layers.{}.attention.wo.weight", + "model.layers.{}.self_attn.k_layernorm.weight": "layers.{}.attention.k_norm_fn.weight", + "model.layers.{}.self_attn.q_layernorm.weight": "layers.{}.attention.q_norm_fn.weight", + "model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight", + "model.layers.{}.operator_norm.weight": "layers.{}.attention_norm.weight", +} + + +def lfm_2_to_meta(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert a state dict from LFM2 HF format to Meta's format. This function + doesn't handle any sharding or splitting of state dicts. It follows the + state_dict IN -> state_dict OUT pattern. + + Args: + state_dict (Dict[str, torch.Tensor]): State dict in LFM2 HF format. + + Returns: + Dict[str, torch.Tensor]: State dict in Meta's format. + """ + converted_state_dict = {} + + for key, value in state_dict.items(): + try: + new_key = get_mapped_key(key, _LFM_2_TO_META) + except: + new_key = key.removeprefix("model.") + + # split in_proj + if new_key.endswith(".conv.in_proj.weight"): + for name, split_value in zip( + ["B_proj", "C_proj", "x_proj"], torch.chunk(value, 3, dim=0) + ): + converted_state_dict[new_key.replace("in_proj", name)] = split_value + else: + converted_state_dict[new_key] = value + + # If lm_head.weight is not present in state dict, assume tied embeddings + if "lm_head.weight" not in state_dict: + converted_state_dict["output.weight"] = converted_state_dict[ + "tok_embeddings.weight" + ] + + return converted_state_dict + + +def load_checkpoint(input_dir: str) -> Dict: + print("Loading checkpoint from safetensors directory") + state_dict = load_file(os.path.join(input_dir, "model.safetensors")) + return state_dict + + +def convert_weights(input_dir: str, output_file: str) -> None: + print("Loading checkpoint...") + sd = load_checkpoint(input_dir) + print("Converting checkpoint...") + sd = lfm_2_to_meta(sd) + print("Saving checkpoint...") + torch.save(sd, output_file) + print("Done.") diff --git a/examples/models/lfm2/short_conv.py b/examples/models/lfm2/short_conv.py new file mode 100644 index 00000000000..ae04580d6c6 --- /dev/null +++ b/examples/models/lfm2/short_conv.py @@ -0,0 +1,112 @@ +from typing import Optional + +import torch +from executorch.examples.models.llama.attention import ForwardOptions +from executorch.examples.models.llama.feed_forward import FeedForward + +from executorch.examples.models.llama.norm import RMSNorm +from torch import nn + + +class ShortConv(nn.Module): + def __init__( + self, + dim: int, + L_cache: int = 3, + bias: bool = False, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.dim = dim + self.L_cache = L_cache + self.device = device + self.dtype = dtype + self.bias = bias + + self.conv = nn.Conv1d( + dim, + dim, + kernel_size=L_cache, + padding=0, ## we don't need padding since we handle it manually + groups=dim, + bias=bias, + ) + + conv_state = torch.zeros( + 1, ## batch size is assumed to be 1 for now + dim, + L_cache - 1, + device="cpu", + ) + self.register_buffer("conv_state", conv_state) + + ## better performance in Executorch with separate projections + self.B_proj = nn.Linear(dim, dim, bias=bias) + self.C_proj = nn.Linear(dim, dim, bias=bias) + self.x_proj = nn.Linear(dim, dim, bias=bias) + + self.out_proj = nn.Linear(dim, dim, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch_size, seqlen, dim = x.size() + assert batch_size == 1, "batch_size must be 1" + + B = self.B_proj(x).transpose(-1, -2) # (batch_size, dim, seq_len) + C = self.C_proj(x).transpose(-1, -2) # (batch_size, dim, seq_len) + x = self.x_proj(x).transpose(-1, -2) # (batch_size, dim, seq_len) + + Bx = B * x # (batch_size, dim, seq_len) + + ## This is where we handle padding + ## By default, the conv_state is initialized to 0. + # So, assuming prefill is done on an empty cache, concatenating conv_state to the beginning of the sequence acts similary to + ## using nn.Conv1d(padding=L_cache-1) (for prefill) without no manual padding. + ## However, the manual padding has the added benefit of being correct during decode, when the cache is not initialized to 0. + Bx = torch.cat( + [self.conv_state, Bx], dim=-1 + ) # (batch_size, dim, seq_len + L_cache - 1) + + ## Update the conv_state + new_conv_state = Bx[ + ..., -(self.L_cache - 1) : + ] # (batch_size, dim, L_cache - 1) + with torch.no_grad(): + self.conv_state.copy_(new_conv_state) + + conv_out = self.conv(Bx)[..., : x.size(-1)] # (batch_size, dim, seq_len) + y = C * conv_out # (batch_size, dim, seq_len) + + y = y.transpose(-1, -2) # (batch_size, seq_len, dim) + y = y.contiguous() # (batch_size, seq_len, dim) + y = self.out_proj(y) # (batch_size, seq_len, dim) + return y + + def reset_cache(self): + self.conv_state.zero_() + + +class ShortConvBlock(nn.Module): + def __init__(self, dim: int, hidden_dim: int, norm_eps: float): + super().__init__() + self.L_cache = 3 # hardcode 3 for now + self.conv = ShortConv(dim, self.L_cache, bias=False) + self.feed_forward = FeedForward(dim, hidden_dim) + self.ffn_norm = RMSNorm(dim, norm_eps) + # use attention_norm norm instead of operator_norm to unify with TransformerBlock + self.attention_norm = RMSNorm(dim, norm_eps) + + def forward( + self, + x, + freqs_cos=None, + freqs_sin=None, + _unused_attn_options: Optional[ForwardOptions] = None, + ): # x: 1xN + h = self.conv.forward(self.attention_norm(x)) + h = x + h + out = h + self.feed_forward(self.ffn_norm(h)) + return out, None + + def reset_cache(self): + self.conv.reset_cache() diff --git a/examples/models/llama/CMakeLists.txt b/examples/models/llama/CMakeLists.txt index efa9c8e4009..e7c73c0cffc 100644 --- a/examples/models/llama/CMakeLists.txt +++ b/examples/models/llama/CMakeLists.txt @@ -15,7 +15,7 @@ # ~~~ # It should also be cmake-lint clean. # -cmake_minimum_required(VERSION 3.24) # 3.24 is required for WHOLE_ARCHIVE +cmake_minimum_required(VERSION 3.24) # 3.24 is required for WHOLE_ARCHIVE project(llama_runner) # Duplicating options as root CMakeLists.txt @@ -37,7 +37,7 @@ cmake_dependent_option( "NOT EXECUTORCH_BUILD_ARM_BAREMETAL" OFF ) -option(EXECUTORCH_BUILD_TORCHAO "Build the torchao kernels" OFF) +option(EXECUTORCH_BUILD_KERNELS_TORCHAO_MPS "Build the torchao MPS kernels" OFF) if(NOT PYTHON_EXECUTABLE) set(PYTHON_EXECUTABLE python3) @@ -115,23 +115,21 @@ if(TARGET custom_ops) list(APPEND link_libraries custom_ops) endif() -if(EXECUTORCH_BUILD_TORCHAO) +if(TARGET torchao_ops_executorch) + executorch_target_link_options_shared_lib(torchao_ops_executorch) + list(APPEND link_libraries torchao_ops_executorch) +endif() + +if(EXECUTORCH_BUILD_KERNELS_TORCHAO_MPS) # Currently only enable this on Arm-based Macs - if(CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") - set(TORCHAO_BUILD_ATEN_OPS OFF) - set(TORCHAO_BUILD_EXECUTORCH_OPS ON) - set(TORCHAO_BUILD_CPU_AARCH64 ON) - set(TORCHAO_ENABLE_ARM_NEON_DOT ON) - add_subdirectory( - ${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental - ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental - ) - executorch_target_link_options_shared_lib(torchao_ops_executorch) - list(APPEND link_libraries torchao_ops_executorch) + if(CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR STREQUAL + "arm64" + ) if(EXECUTORCH_BUILD_MPS) add_subdirectory( ${CMAKE_CURRENT_SOURCE_DIR}/../../../third-party/ao/torchao/experimental/ops/mps - ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental/ops/mps) + ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/ao/torchao/experimental/ops/mps + ) executorch_target_link_options_shared_lib(torchao_ops_mps_executorch) list(APPEND link_libraries torchao_ops_mps_executorch) endif() @@ -218,9 +216,12 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") endif() endif() -target_include_directories( - llama_main - PUBLIC ${_common_include_directories} -) +target_include_directories(llama_main PUBLIC ${_common_include_directories}) target_link_libraries(llama_main PUBLIC llama_runner ${link_libraries}) target_compile_options(llama_main PUBLIC ${_common_compile_options}) +if(APPLE) + target_link_options(llama_main PRIVATE -Wl,-rpath,@loader_path) +elseif(UNIX) + set_target_properties(llama_main PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'") +endif() +# Windows doesn't need rpath - DLLs are found via standard Windows search order diff --git a/examples/models/llama/README.md b/examples/models/llama/README.md index 05d20249382..dba0cf8d8a8 100644 --- a/examples/models/llama/README.md +++ b/examples/models/llama/README.md @@ -168,7 +168,7 @@ LLAMA_CHECKPOINT=path/to/consolidated.00.pth LLAMA_PARAMS=path/to/params.json python -m extension.llm.export.export_llm \ - --config examples/models/llamaconfig/llama_bf16.yaml \ + --config examples/models/llama/config/llama_bf16.yaml \ +base.model_class="llama3_2" \ +base.checkpoint="${LLAMA_CHECKPOINT:?}" \ +base.params="${LLAMA_PARAMS:?}" \ @@ -333,18 +333,20 @@ adb shell "cd /data/local/tmp/llama && ./llama_main --model_path --t ### iOS -Please refer to [this tutorial](https://pytorch.org/executorch/main/llm/llama-demo-ios) to for full instructions on building the iOS LLAMA Demo App. Rename `tokenizer.model` file to `tokenizer.bin` because the demo app looks for the tokenizer file with .bin extension. +Please refer to [this tutorial](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) to for full instructions on building the iOS etLLM Demo App. ### Android Please refer to [this tutorial](https://pytorch.org/executorch/main/llm/llama-demo-android) to for full instructions on building the Android LLAMA Demo App. ## Running with low-bit kernels -We now give instructions for quantizating and running your model with low-bit kernels. These are still experimental, and require you do development on an Arm-based Mac, and install executorch from source with the environment variable EXECUTORCH_BUILD_TORCHAO=1 defined: +We now give instructions for quantizating and running your model with low-bit kernels. These are still experimental, and require you do development on an Arm-based Mac, and install executorch from source with the environment variable EXECUTORCH_BUILD_KERNELS_TORCHAO=1 defined: ``` -EXECUTORCH_BUILD_TORCHAO=1 python install_executorch.py +EXECUTORCH_BUILD_KERNELS_TORCHAO=1 python install_executorch.py ``` +(If you'd like lowbit to use KleidiAI when available, you can instead install with `EXECUTORCH_BUILD_KERNELS_TORCHAO=1 TORCHAO_BUILD_KLEIDIAI=1 python install_executorch.py`.) + Also note that low-bit quantization often requires QAT (quantization-aware training) to give good quality results. First export your model for lowbit quantization (step 2 above): @@ -394,9 +396,12 @@ cmake -DPYTHON_EXECUTABLE=python \ -DEXECUTORCH_BUILD_XNNPACK=OFF \ -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_TORCHAO=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out . -cmake --build cmake-out -j16 --target install --config Release +cmake --build cmake-out -j16 --config Release --target install ``` Next install the llama runner with torchao kernels enabled (similar to step 3.2 above): @@ -404,11 +409,6 @@ Next install the llama runner with torchao kernels enabled (similar to step 3.2 ``` cmake -DPYTHON_EXECUTABLE=python \ -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ - -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ - -DEXECUTORCH_BUILD_XNNPACK=OFF \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ - -DEXECUTORCH_BUILD_TORCHAO=ON \ -Bcmake-out/examples/models/llama \ examples/models/llama cmake --build cmake-out/examples/models/llama -j16 --config Release diff --git a/examples/models/llama/TARGETS b/examples/models/llama/TARGETS index 9ea683e4174..fe26b2f08e0 100644 --- a/examples/models/llama/TARGETS +++ b/examples/models/llama/TARGETS @@ -13,6 +13,24 @@ runtime.python_library( name = "llama_transformer", srcs = [ "llama_transformer.py", + ], + _is_external_target = True, + base_module = "executorch.examples.models.llama", + visibility = [ + "//executorch/...", + "@EXECUTORCH_CLIENTS", + ], + deps = [ + ":transformer_modules", + "//caffe2:torch", + "//executorch/examples/models/lfm2:lfm2", + ], +) + +runtime.python_library( + name = "transformer_modules", + srcs = [ + "feed_forward.py", "lora.py", "rope.py", "attention.py", @@ -25,9 +43,6 @@ runtime.python_library( "//executorch/...", "@EXECUTORCH_CLIENTS", ], - deps = [ - "//caffe2:torch", - ], ) runtime.python_library( @@ -84,6 +99,7 @@ runtime.python_binary( ], deps = [ ":export_library", + ":eval_library", "//caffe2:torch", "//executorch/extension/pybindings:aten_lib", "//executorch/extension/llm/export:export_llm_lib", @@ -116,7 +132,6 @@ runtime.python_library( "source_transformation/rope.py", "source_transformation/sdpa.py", "source_transformation/spin_quant.py", - "source_transformation/vulkan_rope.py", "source_transformation/attention_sink.py", ], ) @@ -153,10 +168,10 @@ runtime.python_library( "//caffe2:torch", "//executorch/extension/llm/export/config:llm_config", "//executorch/backends/vulkan/_passes:vulkan_passes", + "//executorch/exir/passes:external_constants_pass", "//executorch/exir/passes:init_mutable_pass", "//executorch/examples/models:model_base", "//executorch/examples/models:models", - "//executorch/exir/passes:init_mutable_pass", "//executorch/extension/llm/custom_ops:custom_ops_aot_py", "//executorch/extension/llm/export:export_lib", # one definition has to be included in the user of the libarary diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index 6f23456eaaa..6e3f7cb9fb2 100644 --- a/examples/models/llama/attention.py +++ b/examples/models/llama/attention.py @@ -331,6 +331,7 @@ def __init__( args: ModelArgs, layer_id: int, rope: Rope, + **_kwargs: Any, ): """ Multi-head attention layer. diff --git a/examples/models/llama/eval_llama_lib.py b/examples/models/llama/eval_llama_lib.py index 991ff72ae43..03f8f5cd759 100644 --- a/examples/models/llama/eval_llama_lib.py +++ b/examples/models/llama/eval_llama_lib.py @@ -223,7 +223,7 @@ def gen_eval_wrapper( ) else: # TODO: use manager.pre_autograd_graph_module for the eval to remove the if-else branch - # for quantizers. Currently export_for_training only works with --kv_cache, but + # for quantizers. Currently export only works with --kv_cache, but # fails without the kv_cache mode model = ( manager.model.eval().to(device="cuda") diff --git a/examples/models/llama/evaluate/eager_eval.py b/examples/models/llama/evaluate/eager_eval.py index ff9cac16c88..da4742cfc96 100644 --- a/examples/models/llama/evaluate/eager_eval.py +++ b/examples/models/llama/evaluate/eager_eval.py @@ -10,6 +10,7 @@ import torch from lm_eval.models.huggingface import HFLM as eval_wrapper +from pytorch_tokenizers.hf_tokenizer import HuggingFaceTokenizer from pytorch_tokenizers.llama2c import Llama2cTokenizer as SentencePieceTokenizer from pytorch_tokenizers.tiktoken import TiktokenTokenizer as Tiktoken @@ -24,7 +25,7 @@ class EagerEvalWrapper(eval_wrapper): def __init__( self, model: nn.Module, - tokenizer: Union[SentencePieceTokenizer, Tiktoken], + tokenizer: Union[SentencePieceTokenizer, Tiktoken, HuggingFaceTokenizer], max_seq_length: Optional[int] = None, use_kv_cache: bool = False, ): diff --git a/examples/models/llama/experimental/generate.py b/examples/models/llama/experimental/generate.py index 01b5d6668c3..f97b4c543b2 100644 --- a/examples/models/llama/experimental/generate.py +++ b/examples/models/llama/experimental/generate.py @@ -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. -# Adapted from gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py +# Adapted from gpt-fast: https://github.com/meta-pytorch/gpt-fast/blob/main/generate.py import argparse from typing import Optional, Tuple diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index a0cb7dab0ea..7192204a141 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -17,25 +17,22 @@ import shlex from enum import Enum from functools import partial + +from importlib import resources as _resources from json import JSONDecodeError from pathlib import Path from typing import Callable, List, Optional, Union -import pkg_resources import torch from executorch.devtools.backend_debug import print_delegation_info - from executorch.devtools.etrecord import generate_etrecord as generate_etrecord_func from executorch.examples.models.llama.hf_download import ( download_and_convert_hf_checkpoint, ) from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass - from executorch.extension.llm.export.builder import DType, LLMEdgeManager - from executorch.extension.llm.export.config.llm_config import LlmConfig - from executorch.extension.llm.export.partitioner_lib import ( get_coreml_partitioner, get_mps_partitioner, @@ -43,7 +40,6 @@ get_vulkan_partitioner, get_xnnpack_partitioner, ) - from executorch.extension.llm.export.quantizer_lib import ( get_coreml_quantizer, get_pt2e_quantization_params, @@ -52,7 +48,6 @@ get_vulkan_quantizer, ) from executorch.util.activation_memory_profiler import generate_memory_trace - from omegaconf import DictConfig from ..model_factory import EagerModelFactory @@ -60,20 +55,16 @@ fuse_layer_norms, get_model_with_r1_r2, ) - from .source_transformation.attention import replace_attention_to_attention_sha from .source_transformation.custom_kv_cache import ( replace_kv_cache_with_custom_kv_cache, replace_kv_cache_with_quantized_kv_cache, replace_kv_cache_with_ring_kv_cache, ) - from .source_transformation.quantize import ( get_quant_embedding_transform, get_quant_weight_transform, ) -from .source_transformation.rms_norm import replace_rms_norm_with_native_rms_norm - from .source_transformation.rope import materialze_broadcast_of_rope_freq_cis from .source_transformation.sdpa import ( replace_causal_mask, @@ -85,7 +76,6 @@ replace_sdpa_with_quantized_sdpa, replace_sdpa_with_simple_sdpa, ) -from .source_transformation.vulkan_rope import replace_with_vulkan_rotary_emb IS_FBCODE = True # os.environ.get("FBCODE_PLATFORM", False) FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -103,21 +93,29 @@ "llama3_1", "llama3_2", "static_llama", - "qwen2_5", + "qwen2_5_0_5b", + "qwen2_5_1_5b", "qwen3_0_6b", "qwen3_1_7b", "qwen3_4b", "phi_4_mini", "smollm2", + "lfm2_350m", # hybrid + "lfm2_700m", # hybrid + "lfm2_1_2b", # hybrid ] TORCHTUNE_DEFINED_MODELS = ["llama3_2_vision"] HUGGING_FACE_REPO_IDS = { - "qwen2_5": "Qwen/Qwen2.5-1.5B", + "qwen2_5_0_5b": "Qwen/Qwen2.5-0.5B", + "qwen2_5_1_5b": "Qwen/Qwen2.5-1.5B", "phi_4_mini": "microsoft/Phi-4-mini-instruct", "smollm2": "HuggingFaceTB/SmolLM-135M", "qwen3_0_6b": "Qwen/Qwen3-0.6B", "qwen3_1_7b": "Qwen/Qwen3-1.7B", "qwen3_4b": "Qwen/Qwen3-4B", + "lfm2_350m": "LiquidAI/LFM2-350M", + "lfm2_700m": "LiquidAI/LFM2-700M", + "lfm2_1_2b": "LiquidAI/LFM2-1.2B", } @@ -132,7 +130,7 @@ def set_pkg_name(name: str) -> None: def get_resource_path(resource_name) -> str: - return pkg_resources.resource_filename(pkg_name, resource_name) + return str(_resources.files(pkg_name).joinpath(resource_name)) def set_verbosity(val): @@ -420,6 +418,7 @@ def build_args_parser() -> argparse.ArgumentParser: help="Delegate more operators beyond DQLinear to the xnnpack backend. Requires -X or --xnnpack to be set.", ) parser.add_argument("-V", "--vulkan", action="store_true") + parser.add_argument("--vulkan-force-fp16", action="store_true") parser.add_argument("--mps", action="store_true") parser.add_argument("--coreml", action="store_true") parser.add_argument( @@ -578,7 +577,7 @@ def canonical_path(path: Union[str, Path], *, dir: bool = False) -> str: print("not FBCODE") return path[4:] else: - return_val = pkg_resources.resource_filename(pkg_name, path[4:]) + return_val = str(_resources.files(pkg_name).joinpath(path[4:])) if verbose_export(): print(f"canonical name is: {return_val}") return return_val @@ -605,22 +604,16 @@ def export_llama( model_name = llm_config.base.model_class.value if not llm_config.base.checkpoint and model_name in HUGGING_FACE_REPO_IDS: repo_id = HUGGING_FACE_REPO_IDS[model_name] - if model_name == "qwen2_5": - from executorch.examples.models.qwen2_5 import ( # pyre-ignore[21] - convert_weights, - ) + if model_name.startswith("qwen2_5"): + from executorch.examples.models.qwen2_5 import convert_weights elif model_name.startswith("qwen3"): - from executorch.examples.models.qwen3 import ( # pyre-ignore[21] - convert_weights, - ) + from executorch.examples.models.qwen3 import convert_weights elif model_name == "phi_4_mini": - from executorch.examples.models.phi_4_mini import ( # pyre-ignore[21] - convert_weights, - ) + from executorch.examples.models.phi_4_mini import convert_weights elif model_name == "smollm2": - from executorch.examples.models.smollm2 import ( # pyre-ignore[21] - convert_weights, - ) + from executorch.examples.models.smollm2 import convert_weights + elif model_name.startswith("lfm2"): + from executorch.examples.models.lfm2 import convert_weights else: raise ValueError( f"Converting weights to meta format for {model_name} is not yet supported" @@ -790,7 +783,7 @@ def get_quantizer_and_quant_params(llm_config): def _qmode_type(value): - choices = ["int8", "8da4w", "8da4w-gptq", "vulkan_4w"] + choices = ["int8", "8da4w", "8da4w-gptq", "4w"] patterns = [r"torchao:8da(\d+)w", r"torchao:fpa(\d+)w"] if value in choices: @@ -864,9 +857,7 @@ def _to_edge_and_lower_llama_xnnpack( # TODO: Enable generating ETRecord with XNNPack and to_edge_transform_and_lower(). if generate_etrecord: - raise NotImplementedError( - "export_llama does not support XNNPack and generating ETRecord at the moment." - ) + builder_exported.generate_etrecord = True builder = builder_exported.pt2e_quantize(quantizers).to_edge_transform_and_lower( partitioners @@ -874,6 +865,8 @@ def _to_edge_and_lower_llama_xnnpack( if verbose: print_delegation_info(builder.edge_manager.exported_program().graph_module) + # we need builder.export_program + return builder.to_executorch(passes=additional_passes) @@ -893,6 +886,7 @@ def _to_edge_and_lower_llama( # noqa: C901 use_kv_cache: bool = False, embedding_quantize: Optional[str] = None, pt2e_quantize: Optional[str] = None, + vulkan_force_fp16: bool = False, coreml_ios: int = 15, coreml_quantize: Optional[str] = None, coreml_compute_units: str = "cpu_only", @@ -913,6 +907,7 @@ def _to_edge_and_lower_llama( # noqa: C901 get_vulkan_partitioner( dtype_override, enable_dynamic_shape, + vulkan_force_fp16, ) ) modelname = f"vulkan_{modelname}" @@ -947,7 +942,6 @@ def _to_edge_and_lower_llama( # noqa: C901 AnnotateStack, ConvertBmmToMatmul, FoldQDQ, - RecomposeRmsNorm, TagQuantIO, ) @@ -989,7 +983,6 @@ def _to_edge_and_lower_llama( # noqa: C901 dep_table = get_passes_dependency_for_capture_program() passes_job[AnnotateStack][QCOM_PASS_ACTIVATE_KEY] = True passes_job[ConvertBmmToMatmul][QCOM_PASS_ACTIVATE_KEY] = True - passes_job[RecomposeRmsNorm][QCOM_PASS_ACTIVATE_KEY] = True passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ "get_quant_io_dtype_fn" @@ -1078,6 +1071,31 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 llm_config.backend.xnnpack.enabled = True if llm_config.backend.xnnpack.enabled: + if llm_config.export.foundation_weights_file is not None: + gen_tag_fn: Callable[[torch.fx.Node], Optional[str]] = lambda x: ( + llm_config.export.foundation_weights_file + if "lora" not in x.name + else None + ) + + from executorch.exir.passes.external_constants_pass import ( + delegate_external_constants_pass_unlifted, + external_constants_pass, + ) + + assert ( + builder_exported.pre_autograd_graph_module is not None + ), "pre_autograd_graph_module shouldn't be None here" + delegate_external_constants_pass_unlifted( + module=builder_exported.pre_autograd_graph_module, + gen_tag_fn=gen_tag_fn, + ) + + # Also add a pass for 'to_executorch' to tag weights that aren't delegated. + additional_passes.append( + partial(external_constants_pass, gen_tag_fn=gen_tag_fn) + ) + builder = _to_edge_and_lower_llama_xnnpack( builder_exported, modelname, @@ -1110,6 +1128,7 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager: # noqa: C901 if llm_config.quantization.pt2e_quantize else None ), + vulkan_force_fp16=llm_config.backend.vulkan.force_fp16, coreml_ios=llm_config.backend.coreml.ios, coreml_quantize=( llm_config.backend.coreml.quantize.value @@ -1422,14 +1441,12 @@ def _get_source_transforms( # noqa transforms.append(get_model_with_r1_r2(optimized_rotation_path)) transforms.append(replace_attention_to_attention_sha) transforms.append(replace_causal_mask) - transforms.append(replace_rms_norm_with_native_rms_norm) # pyre-fixme[16]: Module `backends` has no attribute `qualcomm`. transforms.append(convert_linear_to_conv2d) else: transforms.append(replace_kv_cache_with_simple_kv_cache) transforms.append(replace_sdpa_with_flex_sdpa) transforms.append(replace_causal_mask) - transforms.append(replace_rms_norm_with_native_rms_norm) if optimized_rotation_path: transforms.append(fuse_layer_norms) transforms.append(get_model_with_r1_r2(optimized_rotation_path)) @@ -1450,9 +1467,6 @@ def _get_source_transforms( # noqa transforms.append(replace_sdpa_with_simple_sdpa) transforms.append(replace_kv_cache_with_coreml_kv_cache) - if vulkan: - transforms.append(replace_with_vulkan_rotary_emb) - if local_global_attention: transforms.append( partial( diff --git a/examples/models/llama/feed_forward.py b/examples/models/llama/feed_forward.py new file mode 100644 index 00000000000..21a4e27df04 --- /dev/null +++ b/examples/models/llama/feed_forward.py @@ -0,0 +1,13 @@ +import torch.nn.functional as F +from torch import nn + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py index a53e1716375..3a325d0f4f8 100644 --- a/examples/models/llama/llama_transformer.py +++ b/examples/models/llama/llama_transformer.py @@ -17,26 +17,13 @@ ATTENTION_REGISTRY, ForwardOptions, ) - +from executorch.examples.models.llama.feed_forward import FeedForward from executorch.examples.models.llama.model_args import ModelArgs from executorch.examples.models.llama.norm import RMSNorm from executorch.examples.models.llama.rope import Rope from torch import nn -class FeedForward(nn.Module): - def __init__(self, args: ModelArgs): - super().__init__() - assert args.hidden_dim is not None - hidden_dim: int = args.hidden_dim - self.w1 = nn.Linear(args.dim, hidden_dim, bias=False) - self.w2 = nn.Linear(hidden_dim, args.dim, bias=False) - self.w3 = nn.Linear(args.dim, hidden_dim, bias=False) - - def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) - - class ConditionalFeedForward(nn.Module): def __init__(self, args: ModelArgs): super().__init__() @@ -99,10 +86,15 @@ def __init__(self, args: ModelArgs, attention: Attention): self.dim = args.dim self.head_dim = args.head_dim self.attention = attention + + assert ( + args.hidden_dim is not None + ), "`hidden_dim` must be set in ModelArgs to construct a TransformerBlock." if args.moe: self.block_sparse_moe = MOEFeedForward(args) else: - self.feed_forward = FeedForward(args) + self.feed_forward = FeedForward(dim=args.dim, hidden_dim=args.hidden_dim) + self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps) self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) @@ -121,7 +113,7 @@ def from_type(cls, layer_id, args, rope) -> "TransformerBlock": f"Available: {list(ATTENTION_REGISTRY.keys())}" ) cls = ATTENTION_REGISTRY[args.attention_type] - attention = cls(args, layer_id, rope) + attention = cls(args, layer_id, rope, **args.attention_kwargs) return TransformerBlock(args, attention) def forward(self, x, freqs_cos, freqs_sin, attn_options: ForwardOptions): # x: 1xN @@ -255,8 +247,25 @@ def construct_transformer(model_args: ModelArgs) -> Transformer: layers = torch.nn.ModuleList() cls = ATTENTION_REGISTRY[model_args.attention_type] for layer_id in range(model_args.n_layers): - attention = cls(model_args, layer_id, rope) - transformer_block = TransformerBlock(model_args, attention) - layers.append(transformer_block) + # hybrid models define layer_types + if model_args.layer_types and model_args.layer_types[layer_id] == "conv": + from executorch.examples.models.lfm2.short_conv import ShortConvBlock + + assert ( + model_args.hidden_dim is not None + ), "`hidden_dim` must be set in ModelArgs to construct a TransformerBlock." + layers.append( + ShortConvBlock( + dim=model_args.dim, + hidden_dim=model_args.hidden_dim, + norm_eps=model_args.norm_eps, + ) + ) + else: + attention = cls( + model_args, layer_id, rope, **model_args.attention_kwargs + ) # pyre-ignore[45] + transformer_block = TransformerBlock(model_args, attention) + layers.append(transformer_block) return Transformer(model_args, layers, rope) diff --git a/examples/models/llama/model_args.py b/examples/models/llama/model_args.py index 1335aaf609e..04d29f91ac6 100644 --- a/examples/models/llama/model_args.py +++ b/examples/models/llama/model_args.py @@ -1,5 +1,38 @@ +import dataclasses from dataclasses import dataclass -from typing import Dict, Optional +from enum import Enum +from functools import partial +from typing import Any, Dict, Optional + +import torch.nn.functional as F + + +class ActFn(Enum): + SILU = "silu" + GELU = "gelu" + GELU_APPROX = "gelu_approx" + + @classmethod + def from_string(cls, value: str) -> "ActFn": + """Convert string to ActFn enum.""" + try: + return cls(value) + except ValueError: + valid_values = [e.value for e in cls] + raise ValueError( + f"Invalid activation function: {value}. Valid options: {valid_values}" + ) + + def get_function(self): + """Return the corresponding activation function.""" + if self == ActFn.SILU: + return F.silu + elif self == ActFn.GELU: + return F.gelu + elif self == ActFn.GELU_APPROX: + return partial(F.gelu, approximate="tanh") + else: + raise ValueError(f"Unsupported activation function: {self}") @dataclass @@ -14,6 +47,8 @@ class ModelArgs: multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2 ffn_dim_multiplier: Optional[float] = None norm_eps: float = 1e-5 + post_attention_norm: bool = False + post_ffn_norm: bool = False max_batch_size: int = 1 max_seq_len: int = 2048 max_context_len: int = 2048 @@ -21,6 +56,8 @@ class ModelArgs: num_experts: int = 8 # Number of experts num_activated_experts: int = 2 # Number of experts to activate attention_type: str = "mha" # Attention type, registered in attention.py + norm_type: str = "rmsnorm" # Normalization type, registered in norm.py + act_fn: ActFn = dataclasses.field(default=ActFn.SILU) # Activation function type attention_qkv_bias: bool = False use_kv_cache: bool = False # Use key/value cache use_sdpa_with_kv_cache_op: bool = ( @@ -36,10 +73,14 @@ class ModelArgs: # A dictionary mapping from pruned token-id to original token-id output_prune_map: Optional[Dict[int, int]] = None apply_embedding: bool = True # Use embedding inside the transformer + embedding_scale_factor: float = 1.0 # Multiple by which to scale embeddings. apply_output: bool = True # Use output layer (unembedding) inside the transformer use_qk_norm: bool = False # apply normalization to q and k in the attention qk_norm_before_rope: bool = False # when to apply qk norm use_hf_rope: bool = False # Use HuggingFace's RoPE implementation + no_rope_layer_interval: Optional[int] = ( + None # Interval at which to skip RoPE. From Rope to Nope and Back Again: A New Hybrid Attention Strategy (https://huggingface.co/papers/2501.18795). + ) partial_rotary_factor: float = 1.0 rope_theta: Optional[float] = ( None # The official name to override self.rope_freq_base. @@ -66,6 +107,12 @@ class ModelArgs: target_modules: Optional[list] = None peft_type: Optional[str] = None # PEFT type. base_model_name_or_path: Optional[str] = None # Base model name or path. + kv_io_bit_width: Optional[int] = ( + None # KV cache bit width. This is for QNN backend only for now. + ) + attention_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) + # Hybrid models can have layer types different from attention + layer_types: Optional[list] = None def __post_init__(self): if self.n_kv_heads is None: @@ -96,3 +143,7 @@ def find_multiple(n: int, k: int) -> int: if self.head_dim is None: self.head_dim = self.dim // self.n_heads + + # Convert string act_fn to enum if needed + if isinstance(self.act_fn, str): + self.act_fn = ActFn.from_string(self.act_fn) diff --git a/examples/models/llama/non_cpu_backends.md b/examples/models/llama/non_cpu_backends.md index f414582a3c1..6a4c5e16cc3 100644 --- a/examples/models/llama/non_cpu_backends.md +++ b/examples/models/llama/non_cpu_backends.md @@ -10,7 +10,7 @@ Export: python -m examples.models.llama2.export_llama --checkpoint llama3.pt --params params.json -kv --disable_dynamic_shape --mps --use_sdpa_with_kv_cache -d fp32 -qmode 8da4w -G 32 --embedding-quantize 4,32 ``` -After exporting the MPS model .pte file, the [iOS LLAMA](https://pytorch.org/executorch/main/llm/llama-demo-ios) app can support running the model. ` --embedding-quantize 4,32` is an optional args for quantizing embedding to reduce the model size. +After exporting the MPS model .pte file, the [iOS etLLM](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) app can support running the model. ` --embedding-quantize 4,32` is an optional args for quantizing embedding to reduce the model size. ### CoreML Export: diff --git a/examples/models/llama/rope.py b/examples/models/llama/rope.py index f788b8f5032..8c0d5db6a80 100644 --- a/examples/models/llama/rope.py +++ b/examples/models/llama/rope.py @@ -9,7 +9,7 @@ import math from functools import partial -from typing import Optional, Tuple +from typing import Optional, Tuple, Union import torch from executorch.examples.models.llama.model_args import ModelArgs @@ -47,9 +47,10 @@ def precompute_freqs_cis( use_scaled: bool = False, scale_factor: Optional[int] = None, high_freq_factor: int = 4, + device: Union[str, torch.device] = "cpu", ): freqs = 1.0 / ( - theta ** (torch.arange(0, dim, 2, device="cpu")[: (dim // 2)].float() / dim) + theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim) ) t = torch.arange(end, device=freqs.device) # pyre-ignore if use_scaled: diff --git a/examples/models/llama/runner/CMakeLists.txt b/examples/models/llama/runner/CMakeLists.txt index ebe1fb201f2..7c6c5413ab3 100644 --- a/examples/models/llama/runner/CMakeLists.txt +++ b/examples/models/llama/runner/CMakeLists.txt @@ -23,13 +23,10 @@ endif() include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) -# The buck-based executorch_srcs.cmake setup was crossing package -# boundaries and trying to build stuff from -# executorch/extension/llm/runner and tokenizers. Just set up sources -# manually. -set(llama_runner_srcs - runner.cpp - ../tokenizer/llama_tiktoken.cpp) +# The buck-based executorch_srcs.cmake setup was crossing package boundaries and +# trying to build stuff from executorch/extension/llm/runner and tokenizers. +# Just set up sources manually. +set(llama_runner_srcs runner.cpp ../tokenizer/llama_tiktoken.cpp) if(CMAKE_TOOLCHAIN_IOS OR ANDROID @@ -59,6 +56,6 @@ target_link_libraries(llama_runner PUBLIC tokenizers::tokenizers) target_include_directories( llama_runner PUBLIC ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include - ${EXECUTORCH_ROOT}/.. + ${EXECUTORCH_ROOT}/.. ) target_compile_options(llama_runner PUBLIC ${_preprocessor_flag}) diff --git a/examples/models/llama/runner/runner.h b/examples/models/llama/runner/runner.h index 09a166b0109..f07cd4e8ee8 100644 --- a/examples/models/llama/runner/runner.h +++ b/examples/models/llama/runner/runner.h @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -33,6 +34,7 @@ std::unique_ptr create_llama_runner( float temperature = -1.0f); std::unique_ptr load_llama_tokenizer( - const std::string& tokenizer_path); + const std::string& tokenizer_path, + Version version = Version::Default); } // namespace example diff --git a/examples/models/llama/runner/static_attention_io_manager.h b/examples/models/llama/runner/static_attention_io_manager.h index 41c826773fa..a696d92c40c 100644 --- a/examples/models/llama/runner/static_attention_io_manager.h +++ b/examples/models/llama/runner/static_attention_io_manager.h @@ -328,7 +328,7 @@ class StaticAttentionMask { } void set_causal_mask() { - for (size_t i = 0; i < input_len_ - 1; i++) { + for (size_t i = 0; i < input_len_; i++) { auto* p = data_ + (cache_len_ + input_len_) * i; std::fill(p + cache_len_, p + cache_len_ + 1 + i, zero_val_); std::fill(p + cache_len_ + 1 + i, p + cache_len_ + input_len_, mask_val_); @@ -576,6 +576,10 @@ class StaticAttentionIOManager { } } + size_t input_pos() const { + return input_pos_; + } + /** * Prefill helper. Run multiple inferences as needed depending on the length * of the prompt and method's input length. Returns the position in the output @@ -586,6 +590,7 @@ class StaticAttentionIOManager { executorch::runtime::Span tokens, executorch::runtime::Span input_buffer, executorch::runtime::Method& method) { + ET_LOG(Info, "Prefilling at position %zu", input_pos_); size_t input_len = input_buffer.size(); auto& masks = get_mask(input_buffer.size()); for (auto& pair : masks) { @@ -621,6 +626,7 @@ class StaticAttentionIOManager { executorch::runtime::Method& method, std::function& sample, std::function& token_callback) { + ET_LOG(Info, "Decoding at position %zu", input_pos_); set_input(method, 0, input_buffer.data()); auto& masks = get_mask(input_buffer.size()); for (auto& pair : masks) { @@ -661,6 +667,10 @@ class StaticAttentionIOManager { size_t window_size, size_t n_verifications, std::unordered_map> suffix_caches) { + ET_LOG( + Info, + "Decoding with lookahead and verification at position %zu", + input_pos_); set_input(method, 0, input_buffer.data()); size_t input_len = input_buffer.size(); diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index 0fbdd1936ef..8d4d37e0e93 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -269,7 +269,7 @@ def replace_kv_cache_with_quantized_kv_cache(module): executorch_package_path = executorch.__path__[-1] libs = list( glob.glob( - f"{executorch_package_path}/**/libquantized_ops_aot_lib.*", + f"{executorch_package_path}/**/*quantized_ops_aot_lib.*", recursive=True, ) ) diff --git a/examples/models/llama/source_transformation/quantize.py b/examples/models/llama/source_transformation/quantize.py index fed36c39081..835972b7f3e 100644 --- a/examples/models/llama/source_transformation/quantize.py +++ b/examples/models/llama/source_transformation/quantize.py @@ -116,7 +116,6 @@ def quantize( # noqa C901 assert len(matches) == 1, f"Expected 1 match for pattern but got {len(matches)}" bitwidth = int(matches[0][0]) - from torchao.dtypes import PackedLinearInt8DynamicActivationIntxWeightLayout from torchao.quantization.granularity import PerAxis, PerGroup from torchao.quantization.quant_api import ( Int8DynamicActivationIntxWeightConfig, @@ -136,7 +135,7 @@ def quantize( # noqa C901 PerAxis(0) if group_size == 0 else PerGroup(group_size) ), weight_mapping_type=MappingType.SYMMETRIC, - layout=PackedLinearInt8DynamicActivationIntxWeightLayout(), + intx_packing_format="opaque_torchao_auto", ), ) model = unwrap_tensor_subclass(model) @@ -148,10 +147,21 @@ def quantize( # noqa C901 # TODO: Default value for group size for 8da4w. Need this here for refactor, will clean this up. group_size = 128 - from torchao.quantization import int8_dynamic_activation_int4_weight, quantize_ + from torchao.quantization import ( + Int8DynamicActivationIntxWeightConfig, + quantize_, + ) + from torchao.quantization.granularity import PerGroup from torchao.utils import unwrap_tensor_subclass - quantize_(model, int8_dynamic_activation_int4_weight(group_size=group_size)) + quantize_( + model, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, + weight_granularity=PerGroup(group_size), + ), + ) + model = unwrap_tensor_subclass(model) # TODO: deal with checkpoint / computation dtype decoupling. @@ -159,11 +169,19 @@ def quantize( # noqa C901 if verbose: print("quantized model:", model) return model - elif qmode == "vulkan_4w": - from executorch.backends.vulkan._passes import VkInt4WeightOnlyQuantizer + elif qmode == "4w": + from torchao.quantization.granularity import PerGroup + from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ + from torchao.utils import unwrap_tensor_subclass q_group_size = 256 if group_size is None else group_size - model = VkInt4WeightOnlyQuantizer(groupsize=q_group_size).quantize(model) + q_config = IntxWeightOnlyConfig( + # pyre-ignore[16] + weight_dtype=torch.int4, + granularity=PerGroup(q_group_size), + ) + quantize_(model, q_config) + model = unwrap_tensor_subclass(model) return model else: @@ -736,9 +754,9 @@ def get_quant_embedding_transform( dtype_override: Optional[DType] = None, ): if embedding_quantize.startswith("torchao:"): - from torchao.experimental.quant_api import ( + from torchao.prototype.quantization.embedding.api import ( EmbeddingQuantizer, - SharedEmbeddingQuantizer, + TiedEmbeddingQuantizer, ) from torchao.quantization.granularity import PerAxis, PerGroup from torchao.quantization.quant_api import MappingType @@ -772,7 +790,7 @@ def _torchao_embedding_quantizer(model): use_fallback=False, ).quantize(model) else: - SharedEmbeddingQuantizer( + TiedEmbeddingQuantizer( weight_dtype=weight_dtype, granularity=granularity, mapping_type=mapping_type, diff --git a/examples/models/llama/source_transformation/spin_quant.py b/examples/models/llama/source_transformation/spin_quant.py index e07b78dc6ee..c8dcee94617 100644 --- a/examples/models/llama/source_transformation/spin_quant.py +++ b/examples/models/llama/source_transformation/spin_quant.py @@ -14,7 +14,7 @@ import torch.nn.functional as F -from executorch.examples.models.llama.llama_transformer import FeedForward +from executorch.examples.models.llama.feed_forward import FeedForward from torch import nn diff --git a/examples/models/llama/source_transformation/vulkan_rope.py b/examples/models/llama/source_transformation/vulkan_rope.py deleted file mode 100644 index cdaf6f0baa7..00000000000 --- a/examples/models/llama/source_transformation/vulkan_rope.py +++ /dev/null @@ -1,36 +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. - -import executorch.backends.vulkan.custom_ops_lib # noqa -import torch - -from executorch.examples.models.llama.rope import RotaryEmbedding - - -class VkRotaryEmbedding(torch.nn.Module): - def __init__(self): - super().__init__() - - def forward( - self, - xq: torch.Tensor, - xk: torch.Tensor, - freqs_cos: torch.Tensor, - freqs_sin: torch.Tensor, - ): - xq_out, xk_out = torch.ops.et_vk.apply_rotary_emb(xq, xk, freqs_cos, freqs_sin) - return xq_out, xk_out - - -def replace_with_vulkan_rotary_emb(module: torch.nn.Module): - for name, child in module.named_children(): - if isinstance(child, RotaryEmbedding): - new_module = VkRotaryEmbedding() - setattr(module, name, new_module) - else: - replace_with_vulkan_rotary_emb(child) - - return module diff --git a/examples/models/llama/static_attention.py b/examples/models/llama/static_attention.py index e3859b98210..b42371dc090 100644 --- a/examples/models/llama/static_attention.py +++ b/examples/models/llama/static_attention.py @@ -242,7 +242,8 @@ def __init__( config: ModelArgs, input_len: int, cache_lens: Union[int, List[int]], - dtype=torch.float32, + batch_size: int = 1, + dtype: torch.dtype = torch.float32, style: str = "shift_pointer", mask_val: float = float("-inf"), ): @@ -266,7 +267,10 @@ def __init__( if split_mha: self.k_caches = { StaticKVCache.calculate_cache_key(layer_id, head_id): torch.zeros( - 1, cache_lens[layer_id], none_throws(config.head_dim), dtype=dtype + batch_size, + cache_lens[layer_id], + none_throws(config.head_dim), + dtype=dtype, ) for layer_id in range(config.n_layers) for head_id in range(none_throws(config.n_kv_heads)) @@ -274,7 +278,10 @@ def __init__( } self.v_caches = { StaticKVCache.calculate_cache_key(layer_id, head_id): torch.zeros( - 1, cache_lens[layer_id], none_throws(config.head_dim), dtype=dtype + batch_size, + cache_lens[layer_id], + none_throws(config.head_dim), + dtype=dtype, ) for layer_id in range(config.n_layers) for head_id in range(none_throws(config.n_kv_heads)) @@ -283,7 +290,7 @@ def __init__( else: self.k_caches = { StaticKVCache.calculate_cache_key(layer_id, 0): torch.zeros( - 1, + batch_size, none_throws(config.n_kv_heads), cache_lens[layer_id], none_throws(config.head_dim), @@ -293,7 +300,7 @@ def __init__( } self.v_caches = { StaticKVCache.calculate_cache_key(layer_id, 0): torch.zeros( - 1, + batch_size, none_throws(config.n_kv_heads), cache_lens[layer_id], none_throws(config.head_dim), @@ -323,7 +330,7 @@ def reset(self): def prefill( self, model: Callable[..., Any], - tokens: List[int], + tokens: Union[List[int], torch.Tensor], ) -> torch.Tensor: if self.cache_full: raise RuntimeError("KV cache is full.") @@ -336,10 +343,13 @@ def prefill( ) ) + if isinstance(tokens, list): + tokens = torch.tensor([tokens], dtype=torch.int32) + logits = None all_logits = None - for i in range(0, len(tokens), self.input_len): - logits = self._run_once(model, tokens[i : i + self.input_len])[0] + for i in range(0, tokens.size(1), self.input_len): + logits = self._run_once(model, tokens[:, i : i + self.input_len])[0] if self.config.generate_full_logits: if all_logits is None: all_logits = logits @@ -347,7 +357,7 @@ def prefill( all_logits = torch.cat([all_logits, logits], dim=1) if self.config.generate_full_logits: - return all_logits[:, : len(tokens), :] + return all_logits[:, : tokens.size(1), :] return logits @@ -510,15 +520,16 @@ def lookahead_decode( # noqa: C901 def _run_once( self, model: Callable[..., Any], - tokens: List[int], + tokens: Union[List[int], torch.Tensor], non_padded_len: Optional[int] = None, freqs_cos_override: Optional[torch.Tensor] = None, freqs_sin_override: Optional[torch.Tensor] = None, ): - n_tokens = len(tokens) + if isinstance(tokens, list): + tokens = torch.tensor([tokens], dtype=torch.int32) + n_tokens = tokens.size(1) if n_tokens < self.input_len: - tokens += [0] * (self.input_len - n_tokens) - tokens = torch.tensor([tokens], dtype=torch.int32) # pyre-ignore[9] + tokens = F.pad(tokens, (0, self.input_len - n_tokens)) if freqs_cos_override is None: freqs_cos_override = self.freqs_cos[self.pos : self.pos + self.input_len] if freqs_sin_override is None: @@ -549,7 +560,7 @@ def _update_states(self, attn_updates, update_pos, update_len): style=self.style, update_pos=update_pos, update_len=update_len, - ) + ).detach() for cache_id, update in v_cache_updates.items(): self.v_caches[cache_id] = StaticKVCache.apply_update( self.v_caches[cache_id], @@ -558,7 +569,7 @@ def _update_states(self, attn_updates, update_pos, update_len): style=self.style, update_pos=update_pos, update_len=update_len, - ) + ).detach() self.pos += update_len def _get_lookahead_decoding_mask( @@ -658,7 +669,12 @@ class StaticAttention(Attention): """ def __init__( - self, config: ModelArgs, layer_id: int, rope: Rope, split_mha: bool = True + self, + config: ModelArgs, + layer_id: int, + rope: Rope, + split_mha: bool = True, + **kwargs: Any, ): super().__init__() self.n_heads = config.n_heads @@ -676,6 +692,7 @@ def __init__( self.qk_norm_before_rope = config.qk_norm_before_rope self.split_mha = split_mha self.use_conv2d = False + self.enable_qnn_masked_softmax = kwargs.get("enable_qnn_masked_softmax", False) if self.split_mha: self.wqs = nn.ModuleList( @@ -759,7 +776,7 @@ def forward( bsz, seq_len, dim = x.shape if self.use_conv2d: - x = x.reshape(bsz, seq_len, 1, dim).transpose(1, 3) + x = x.reshape(bsz, -1, 1, dim).transpose(1, 3) new_qs = [wq(x) for wq in self.wqs] new_ks = [wk(x) for wk in self.wks] @@ -768,9 +785,7 @@ def forward( if self.use_conv2d: def from_conv2ds(ts): - return [ - t.reshape(bsz, self.head_dim, seq_len).transpose(1, 2) for t in ts - ] + return [t.reshape(bsz, self.head_dim, -1).transpose(1, 2) for t in ts] new_qs = from_conv2ds(new_qs) new_ks = from_conv2ds(new_ks) @@ -800,9 +815,11 @@ def from_conv2ds(ts): if self.use_conv2d: y = ( - self.wo(y.reshape(bsz, seq_len, 1, -1).transpose(1, 3)) + self.wo( + y.reshape(bsz, -1, 1, self.n_heads * self.head_dim).transpose(1, 3) + ) .transpose(1, 3) - .reshape(bsz, seq_len, -1) + .reshape(bsz, -1, self.dim) ) else: y = self.wo(y) @@ -857,7 +874,14 @@ def _forward_sha( kv_idx = i // self.n_heads_per_kv_group attn = new_qs[i] @ all_ks[kv_idx].transpose(-2, -1) attn = attn * self.inv_scale - attn = attn + mask + if self.enable_qnn_masked_softmax: + attn_min = torch.amin(attn, dim=-1, keepdim=True) + minus_value = -20 + attn = torch.where( + mask == 0, attn, attn_min + minus_value + ) # prye-ignore + else: + attn = attn + mask attn = F.softmax(attn, dim=-1) heads.append(attn @ all_vs[kv_idx]) @@ -1031,5 +1055,5 @@ def transfer_weight(linear, conv2d): @register_attention("static_mha") class StaticAttentionMHA(StaticAttention): - def __init__(self, config: ModelArgs, layer_id: int, rope: Rope): - super().__init__(config, layer_id, rope, split_mha=False) + def __init__(self, config: ModelArgs, layer_id: int, rope: Rope, **kwargs: Any): + super().__init__(config, layer_id, rope, split_mha=False, **kwargs) diff --git a/examples/models/llama/tests/test_ring_attention.py b/examples/models/llama/tests/test_ring_attention.py index df0d0733033..ae440e00e47 100644 --- a/examples/models/llama/tests/test_ring_attention.py +++ b/examples/models/llama/tests/test_ring_attention.py @@ -163,10 +163,17 @@ def test_single_token_processing( ) # Check that outputs are the same - self.assertTrue( - torch.allclose(baseline_out, ring_out, rtol=1e-7, atol=1e-7), - f"Outputs differ at position {pos}", - ) + if kv_cache_type == KVCacheType.REGULAR: + self.assertTrue( + torch.allclose(baseline_out, ring_out, rtol=1e-7, atol=1e-7), + f"Outputs differ at position {pos}", + ) + else: + # For quantized kv cache we need bigger margin + self.assertTrue( + torch.allclose(baseline_out, ring_out, rtol=1e-6, atol=1e-6), + f"Outputs differ at position {pos}", + ) def test_single_token_processing_quantized(self): """Test single token processing with QuantizedKVCache.""" diff --git a/examples/models/llama/tests/test_static_attention.py b/examples/models/llama/tests/test_static_attention.py index 2461732db5a..8786c70da11 100644 --- a/examples/models/llama/tests/test_static_attention.py +++ b/examples/models/llama/tests/test_static_attention.py @@ -195,7 +195,7 @@ def test_with_style(style): test_with_style("shift_pointer") test_with_style("smart_mask") - def _get_test_transformers(self, config, attention_type="static"): + def _get_test_transformers(self, config, attention_type="static", use_conv2d=False): mha_transformer = construct_transformer(config).eval() config = copy.copy(config) @@ -207,6 +207,8 @@ def _get_test_transformers(self, config, attention_type="static"): ): static_layer.attention.load_weights_from_attention_mha(mha_layer.attention) static_layer.attention.adopt_hf_rope() + if use_conv2d: + static_layer.linear_to_conv2d() config.use_hf_rope = True return mha_transformer, static_transformer, config @@ -220,7 +222,8 @@ def test_within_transformer(self): n_layers=4, vocab_size=128, ) - x = torch.randint(config.vocab_size, (1, config.max_seq_len)) + batch_size = 3 + x = torch.randint(config.vocab_size, (batch_size, config.max_seq_len)) n_chunks = 3 chunk_len = config.max_seq_len // n_chunks cache_len = config.max_seq_len - chunk_len @@ -235,13 +238,13 @@ def test(style, attention_type): expected = mha_transformer(x) mgr = StaticAttentionIOManager( - static_config, chunk_len, cache_len, style=style + static_config, chunk_len, cache_len, style=style, batch_size=batch_size ) ys = [] for i in range(n_chunks): y_i = mgr.prefill( static_transformer, - x[0][i * chunk_len : (i + 1) * chunk_len].tolist(), + x[:, i * chunk_len : (i + 1) * chunk_len], ) ys.append(y_i) @@ -300,3 +303,59 @@ def test_lookahead_decode(self): ngram_caches=ngram_caches, ) self.assertEqual(lookahead_output[: len(ref_output)], ref_output) + + def test_batched_export_with_backprop(self): + config = ModelArgs( + dim=64, + n_heads=4, + n_kv_heads=2, + max_seq_len=128, + n_layers=4, + vocab_size=128, + generate_full_logits=True, + ) + _, static_transformer, static_config = self._get_test_transformers(config) + batch_size = 4 + input_len = 32 + cache_len = static_config.max_seq_len - input_len + mgr = StaticAttentionIOManager( + static_config, input_len, cache_len, batch_size=batch_size + ) + example_inputs = ( + torch.zeros(batch_size, input_len), + { + "masks": mgr.masks, + "freqs_cos_override": mgr.freqs_cos[:input_len], + "freqs_sin_override": mgr.freqs_sin[:input_len], + "in_cache_state": (mgr.k_caches, mgr.v_caches), + }, + ) + batched_gm = torch.export.export(static_transformer, example_inputs).module() + + # Test backprop + for _ in range(10): + x = torch.randint(config.vocab_size, (batch_size, input_len)) + y = mgr.prefill(batched_gm, x) + loss = torch.nn.functional.cross_entropy( + y, torch.rand(batch_size, input_len, config.vocab_size) + ) + loss.backward() + mgr.reset() + + # Test loading state dict into a non batched graph for inference + mgr = StaticAttentionIOManager( + static_config, input_len, cache_len, batch_size=1 + ) + example_inputs = ( + torch.zeros(1, input_len), + { + "masks": mgr.masks, + "freqs_cos_override": mgr.freqs_cos[:input_len], + "freqs_sin_override": mgr.freqs_sin[:input_len], + "in_cache_state": (mgr.k_caches, mgr.v_caches), + }, + ) + non_batched_gm = torch.export.export( + static_transformer, example_inputs + ).module() + non_batched_gm.load_state_dict(batched_gm.state_dict()) diff --git a/examples/models/llama3_2_vision/TARGETS b/examples/models/llama3_2_vision/TARGETS index 133fd3a6839..b5ceb1da8ae 100644 --- a/examples/models/llama3_2_vision/TARGETS +++ b/examples/models/llama3_2_vision/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "multimodal_lib", srcs = [ "__init__.py", diff --git a/examples/models/llama3_2_vision/text_decoder/TARGETS b/examples/models/llama3_2_vision/text_decoder/TARGETS index e87b567f2bc..a0d3f5f7e51 100644 --- a/examples/models/llama3_2_vision/text_decoder/TARGETS +++ b/examples/models/llama3_2_vision/text_decoder/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "model", srcs = [ "model.py", diff --git a/examples/models/llama3_2_vision/vision_encoder/TARGETS b/examples/models/llama3_2_vision/vision_encoder/TARGETS index 82717a56d08..e910388f268 100644 --- a/examples/models/llama3_2_vision/vision_encoder/TARGETS +++ b/examples/models/llama3_2_vision/vision_encoder/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "model", srcs = [ "__init__.py", diff --git a/examples/models/llava/CMakeLists.txt b/examples/models/llava/CMakeLists.txt index 500265cbb77..cf9d54ad3ec 100644 --- a/examples/models/llava/CMakeLists.txt +++ b/examples/models/llava/CMakeLists.txt @@ -21,7 +21,6 @@ project(llava) # Duplicating options as root CMakeLists.txt option(EXECUTORCH_BUILD_KERNELS_OPTIMIZED "Build the optimized kernels" OFF) - include(CMakeDependentOption) # # pthreadpool: build pthreadpool library. Disable on unsupported platforms diff --git a/examples/models/llava/README.md b/examples/models/llava/README.md index 86d522862f0..e7b8ba523fd 100644 --- a/examples/models/llava/README.md +++ b/examples/models/llava/README.md @@ -86,9 +86,9 @@ to for full instructions on building the Android LLAMA Demo App. #### iOS -We can run LLAVA using the LLAMA Demo Apps. Please refer to [this -tutorial](https://github.com/pytorch/executorch/tree/main/examples/demo-apps/apple_ios/LLaMA) -to for full instructions on building the iOS LLAMA Demo App. +We can run LLAVA using the etLLM Demo Apps. Please refer to [this +tutorial](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/apple) +to for full instructions on building the iOS etLLM Demo App. ### Running LLaVA diff --git a/examples/models/llava/export_llava.py b/examples/models/llava/export_llava.py index d95bd7fb054..f903e0f2ecf 100644 --- a/examples/models/llava/export_llava.py +++ b/examples/models/llava/export_llava.py @@ -226,11 +226,11 @@ def export_all(llava_model: LlavaModel): { "image_encoder": image_encoder_ep, "token_embedding": token_embedding_ep, - "text_model": text_model_ep, + "text_decoder": text_model_ep, }, partitioner={ "image_encoder": [XnnpackPartitioner()], - "text_model": [ + "text_decoder": [ # First partition the DQLinear nodes, then partition the rest of the nodes, # to avoid multiple DQLinear nodes in the same partition, # to avoid holding multiple unpacked and packed weight buffers in memory, @@ -242,6 +242,7 @@ def export_all(llava_model: LlavaModel): XnnpackPartitioner(), ], }, + constant_methods={"get_max_seq_len": llava_model.max_seq_len}, compile_config=EdgeCompileConfig(_check_ir_validity=False), ) @@ -254,7 +255,7 @@ def export_all(llava_model: LlavaModel): memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), sym_shape_eval_pass={ "image_encoder": ConstraintBasedSymShapeEvalPass(), - "text_model": ConstraintBasedSymShapeEvalPass(), + "text_decoder": ConstraintBasedSymShapeEvalPass(), "token_embedding": HintBasedSymShapeEvalPass(), }, ) diff --git a/examples/models/llava/runner/CMakeLists.txt b/examples/models/llava/runner/CMakeLists.txt index ce7c0968d9e..88ad8590ee5 100644 --- a/examples/models/llava/runner/CMakeLists.txt +++ b/examples/models/llava/runner/CMakeLists.txt @@ -26,9 +26,7 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) set(_common_include_directories ${EXECUTORCH_ROOT}/..) # build llava_runner library -set(_llava_runner__srcs - "${CMAKE_CURRENT_SOURCE_DIR}/llava_runner.cpp" -) +set(_llava_runner__srcs "${CMAKE_CURRENT_SOURCE_DIR}/llava_runner.cpp") if(NOT TARGET extension_llm_runner) message( @@ -40,8 +38,9 @@ endif() add_library(llava_runner STATIC ${_llava_runner__srcs}) target_include_directories(llava_runner PRIVATE ${_common_include_directories}) -set(llava_runner_deps executorch_core extension_data_loader extension_llm_runner - extension_module extension_tensor extension_flat_tensor +set(llava_runner_deps + executorch_core extension_data_loader extension_llm_runner extension_module + extension_tensor extension_flat_tensor ) target_link_libraries(llava_runner PUBLIC ${llava_runner_deps}) diff --git a/examples/models/llava/runner/llava_runner.h b/examples/models/llava/runner/llava_runner.h index 184522c2cf1..62df890b46d 100644 --- a/examples/models/llava/runner/llava_runner.h +++ b/examples/models/llava/runner/llava_runner.h @@ -42,7 +42,7 @@ class ET_EXPERIMENTAL LlavaRunner { const float temperature = 0.8f) : temperature_(temperature), module_(std::make_unique(model_path, Module::LoadMode::File)), - io_manager_(std::make_unique()), + io_manager_(std::make_unique(*module_)), tokenizer_path_(tokenizer_path) { ET_LOG( Info, diff --git a/examples/models/llava/runner/llava_text_decoder_runner.h b/examples/models/llava/runner/llava_text_decoder_runner.h index 09b8e82d49d..cfa92e0c253 100644 --- a/examples/models/llava/runner/llava_text_decoder_runner.h +++ b/examples/models/llava/runner/llava_text_decoder_runner.h @@ -89,7 +89,7 @@ class ET_EXPERIMENTAL LlavaTextDecoderRunner } inline static const std::string kTokenEmbeddingMethod = "token_embedding"; - inline static const std::string kTextModelMethod = "text_model"; + inline static const std::string kTextModelMethod = "text_decoder"; }; } // namespace example diff --git a/examples/models/llava/test/test_llava.py b/examples/models/llava/test/test_llava.py index 05cfd5b1497..def9eaa02bd 100644 --- a/examples/models/llava/test/test_llava.py +++ b/examples/models/llava/test/test_llava.py @@ -96,7 +96,7 @@ def test_llava_export(self): "token_embedding", (prompt_before_image,) )[0] llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_before_img), ) @@ -107,7 +107,7 @@ def test_llava_export(self): # pte prefill image pte_embeds_img = llava_module.run_method("image_encoder", (resized,))[0] llava_module.run_method( - "text_model", + "text_decoder", ( torch.tensor([start_pos], dtype=torch.int64), pte_embeds_img, @@ -122,7 +122,7 @@ def test_llava_export(self): "token_embedding", (prompt_after_image,) )[0] pte_prefill_after_img = llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_after_img), )[0] @@ -139,7 +139,7 @@ def test_llava_export(self): "token_embedding", (torch.tensor([[new_tokens[i]]], dtype=torch.int64),) )[0] logits = llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos + i], dtype=torch.int64), token_embeds), )[0] new_tokens.append(torch.argmax(logits).item()) diff --git a/examples/models/llava/test/test_pte.py b/examples/models/llava/test/test_pte.py index f12d72f854b..1f4aaa9938c 100644 --- a/examples/models/llava/test/test_pte.py +++ b/examples/models/llava/test/test_pte.py @@ -47,7 +47,7 @@ def main(): "token_embedding", (prompt_before_image,) )[0] pte_prefill_before_img = llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_before_img), )[0] print(pte_prefill_before_img) @@ -60,7 +60,7 @@ def main(): logging.warning("Image encoder finished") logging.warning("Image token prefill started") pte_prefill_img = llava_module.run_method( - "text_model", + "text_decoder", ( torch.tensor([start_pos], dtype=torch.int64), pte_embeds_img, @@ -77,7 +77,7 @@ def main(): "token_embedding", (prompt_after_image,) )[0] pte_prefill_after_img = llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos], dtype=torch.int64), pte_embeds_after_img), )[0] logging.warning("Text token prefill finished") @@ -91,7 +91,7 @@ def main(): "token_embedding", (torch.tensor([[new_tokens[i]]], dtype=torch.int64),) )[0] logits = llava_module.run_method( - "text_model", + "text_decoder", (torch.tensor([start_pos + i], dtype=torch.int64), token_embeds), )[0] new_tokens.append(torch.argmax(logits[..., -1, :]).item()) diff --git a/examples/models/mobilenet_v2/model.py b/examples/models/mobilenet_v2/model.py index f15178ac71b..32e82197e46 100644 --- a/examples/models/mobilenet_v2/model.py +++ b/examples/models/mobilenet_v2/model.py @@ -15,7 +15,8 @@ class MV2Model(EagerModelBase): - def __init__(self): + def __init__(self, use_real_input=True): + self.use_real_input = use_real_input pass def get_eager_model(self) -> torch.nn.Module: @@ -26,7 +27,37 @@ def get_eager_model(self) -> torch.nn.Module: def get_example_inputs(self): tensor_size = (1, 3, 224, 224) - return (torch.randn(tensor_size),) + input_batch = (torch.randn(tensor_size),) + if self.use_real_input: + logging.info("Loaded real input image dog.jpg") + import urllib + + url, filename = ( + "https://github.com/pytorch/hub/raw/master/images/dog.jpg", + "dog.jpg", + ) + try: + urllib.URLopener().retrieve(url, filename) + except: + urllib.request.urlretrieve(url, filename) + from PIL import Image + from torchvision import transforms + + input_image = Image.open(filename) + preprocess = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), + ] + ) + input_tensor = preprocess(input_image) + input_batch = input_tensor.unsqueeze(0) + input_batch = (input_batch,) + return input_batch class MV2UntrainedModel(EagerModelBase): diff --git a/examples/models/moshi/mimi/install_requirements.sh b/examples/models/moshi/mimi/install_requirements.sh index 828bcac0abb..23e65465f56 100755 --- a/examples/models/moshi/mimi/install_requirements.sh +++ b/examples/models/moshi/mimi/install_requirements.sh @@ -7,6 +7,8 @@ set -x +conda install "ffmpeg<8" +pip install torchcodec==0.7.0.dev20250906 --extra-index-url https://download.pytorch.org/whl/nightly/cpu pip install moshi==0.2.4 pip install bitsandbytes soundfile # Run llama2/install requirements for torchao deps diff --git a/examples/models/phi-3-mini/CMakeLists.txt b/examples/models/phi-3-mini/CMakeLists.txt index 4e55d4f9cb0..3c7ed6a4acb 100644 --- a/examples/models/phi-3-mini/CMakeLists.txt +++ b/examples/models/phi-3-mini/CMakeLists.txt @@ -31,7 +31,7 @@ set(BUILD_TESTING OFF) if(NOT TARGET extension_llm_runner) message( FATAL_ERROR - "ExecuTorch must be installed with EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER enabled." + "ExecuTorch must be installed with EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER enabled." ) endif() diff --git a/examples/models/phi-3-mini/export_phi-3-mini.py b/examples/models/phi-3-mini/export_phi-3-mini.py index d1239d9769d..017c15f783e 100644 --- a/examples/models/phi-3-mini/export_phi-3-mini.py +++ b/examples/models/phi-3-mini/export_phi-3-mini.py @@ -23,7 +23,7 @@ from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.passes import MemoryPlanningPass from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass -from torch.export import export_for_training +from torch.export import export as torch_export from torch.nn.attention import SDPBackend from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -108,7 +108,7 @@ def export(args) -> None: gm(*example_inputs) gm = convert_pt2e(gm) DuplicateDynamicQuantChainPass()(gm) - exported_program = export_for_training( + exported_program = torch_export( gm, example_inputs, dynamic_shapes=dynamic_shapes, strict=False ) diff --git a/examples/models/phi_4_mini/convert_weights.py b/examples/models/phi_4_mini/convert_weights.py index 3d91747f468..01b7302ed2e 100644 --- a/examples/models/phi_4_mini/convert_weights.py +++ b/examples/models/phi_4_mini/convert_weights.py @@ -1,8 +1,8 @@ import argparse -import os from typing import Dict import torch +from executorch.examples.models.checkpoint import load_checkpoint_from_pytorch_model from torchtune.models.convert_weights import get_mapped_key @@ -87,10 +87,8 @@ def phi_4_tune_to_meta(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.T Convert a state dict from torchtune's format to Meta's format. This function doesn't handle any sharding or splitting of state dicts. It follows the state_dict IN -> state_dict OUT pattern. - Args: state_dict (Dict[str, torch.Tensor]): State dict in torchtune's format. - Returns: Dict[str, torch.Tensor]: State dict in Meta's format. """ @@ -105,14 +103,15 @@ def phi_4_tune_to_meta(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.T converted_state_dict["output.weight"] = converted_state_dict[ "tok_embeddings.weight" ] - return converted_state_dict def convert_weights(input_dir_or_checkpoint: str, output_file: str) -> None: - # If input_dir_or_checkpoint is a directory downloaded from HF, FullModelHFCheckpointer is used to extract the state dict - # If input_dir_or_checkpoint is a checkpoint (from eager model model), it is loaded directly - if os.path.isdir(input_dir_or_checkpoint): + try: + sd = load_checkpoint_from_pytorch_model(input_dir_or_checkpoint) + print("Converting checkpoint...") + sd = phi_4_hf_to_meta(sd) + except FileNotFoundError: checkpointer = FullModelHFCheckpointer( checkpoint_dir=input_dir_or_checkpoint, checkpoint_files=[ @@ -127,11 +126,6 @@ def convert_weights(input_dir_or_checkpoint: str, output_file: str) -> None: sd = sd["model"] print("Converting checkpoint...") sd = phi_4_tune_to_meta(sd) - else: - print("Loading checkpoint from file...") - sd = torch.load(input_dir_or_checkpoint, map_location="cpu", weights_only=True) - print("Converting checkpoint...") - sd = phi_4_hf_to_meta(sd) print("Saving checkpoint...") torch.save(sd, output_file) diff --git a/examples/models/qwen2_5/README.md b/examples/models/qwen2_5/README.md index c58807b46cb..11c92f4bd73 100644 --- a/examples/models/qwen2_5/README.md +++ b/examples/models/qwen2_5/README.md @@ -6,9 +6,10 @@ Qwen 2.5 is the latest iteration of the Qwen series of large language models (LL Qwen 2.5 uses the same example code as Llama, while the checkpoint, model params, and tokenizer are different. Please see the [Llama README page](../llama/README.md) for details. All commands for exporting and running Llama on various backends should also be applicable to Qwen 2.5, by swapping the following args: + ``` -base.model_class="qwen2_5" -base.params="examples/models/qwen2_5/config/1_5b_config.json" +base.model_class=[qwen2_5_0_5b, qwen2_5_1_5b] +base.params=[examples/models/qwen2_5/config/0_5b_config.json, examples/models/qwen2_5/config/1_5b_config.json] base.checkpoint= ``` @@ -34,7 +35,7 @@ QWEN_CHECKPOINT=path/to/checkpoint.pth python -m extension.llm.export.export_llm \ --config examples/models/qwen2_5/config/qwen2_5_xnnpack_q8da4w.yaml \ - +base.model_class="qwen2_5" \ + +base.model_class="qwen2_5_1_5b" \ +base.checkpoint="${QWEN_CHECKPOINT:?}" \ +base.params="examples/models/qwen2_5/config/1_5b_config.json" \ +export.output_name="qwen2_5-1_5b.pte" \ diff --git a/examples/models/qwen2_5/config/qwen2_5_xnnpack_q8da4w.yaml b/examples/models/qwen2_5/config/qwen2_5_xnnpack_q8da4w.yaml index 0e5c6f7624e..77b34527026 100644 --- a/examples/models/qwen2_5/config/qwen2_5_xnnpack_q8da4w.yaml +++ b/examples/models/qwen2_5/config/qwen2_5_xnnpack_q8da4w.yaml @@ -8,4 +8,7 @@ model: backend: xnnpack: - enabled: True \ No newline at end of file + enabled: True + +quantization: + qmode: 8da4w diff --git a/examples/models/qwen3/README.md b/examples/models/qwen3/README.md index 3d5a6cb1ea9..c3d960adfe0 100644 --- a/examples/models/qwen3/README.md +++ b/examples/models/qwen3/README.md @@ -45,7 +45,7 @@ python -m extension.llm.export.export_llm \ ### Example run With ExecuTorch pybindings: ``` -python -m examples.models.llama.runner.native +python -m examples.models.llama.runner.native \ --model qwen3_0_6b \ --pte qwen3_0_6b.pte \ --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json \ @@ -59,9 +59,9 @@ python -m examples.models.llama.runner.native With ExecuTorch's sample c++ runner (see the Llama README's [Step 3: Run on your computer to validate](../llama/README.md#step-3-run-on-your-computer-to-validate) to build the runner): ``` -cmake-out/examples/models/llama/llama_main - --model_path qwen3_0_6b.pte - --tokenizer_path ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json +cmake-out/examples/models/llama/llama_main \ + --model_path qwen3_0_6b.pte \ + --tokenizer_path ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/a9c98e602b9d36d2a2f7ba1eb0f5f31e4e8e5143/tokenizer.json \ --prompt="Who is the president of the US?" ``` diff --git a/examples/models/qwen3/convert_weights.py b/examples/models/qwen3/convert_weights.py index 6d5254906fb..404fd4cbe88 100644 --- a/examples/models/qwen3/convert_weights.py +++ b/examples/models/qwen3/convert_weights.py @@ -5,6 +5,7 @@ from typing import Dict import torch +from executorch.examples.models.checkpoint import load_checkpoint_from_pytorch_model from safetensors.torch import load_file from torchtune.models.convert_weights import get_mapped_key @@ -80,19 +81,34 @@ def load_checkpoint_from_safetensors(input_dir: str) -> Dict: tensor = shard_to_weights[shard][weight_name] merged_state_dict[weight_name] = tensor return merged_state_dict - else: - # Single checkpoint. - state_dict = load_file(os.path.join(input_dir, "model.safetensors")) - return state_dict + + # Single checkpoint. + model_path = os.path.join(input_dir, "model.safetensors") + if os.path.exists(model_path): + return load_file(os.path.join(input_dir, "model.safetensors")) + + raise FileNotFoundError(f"Could not find safetensors checkpoint in {input_dir}") def load_checkpoint(input_dir: str) -> Dict: - pytorch_path = os.path.join(input_dir, "pytorch_model.bin") - if os.path.exists(pytorch_path): - print("Loading checkpoint from PyTorch .bin file") - return torch.load(pytorch_path, map_location="cpu", weights_only=True) - print("Loading checkpoint from safetensors directory") - return load_checkpoint_from_safetensors(input_dir) + try: + print("Loading checkpoint from pytorch_model directory") + state_dict = load_checkpoint_from_pytorch_model(input_dir) + return state_dict + except FileNotFoundError: + print( + "Could not find pytorch_model checkpoints in directory, trying safetensors" + ) + pass + + try: + print("Loading checkpoint from safetensors directory") + state_dict = load_checkpoint_from_safetensors(input_dir) + return state_dict + except FileNotFoundError: + pass + + raise FileNotFoundError(f"Could not find checkpoint in {input_dir}") def convert_weights(input_dir: str, output_file: str) -> None: diff --git a/examples/models/smollm3/3b_config.json b/examples/models/smollm3/3b_config.json new file mode 100644 index 00000000000..c44f0727919 --- /dev/null +++ b/examples/models/smollm3/3b_config.json @@ -0,0 +1,15 @@ +{ + "dim": 2048, + "ffn_dim_multiplier": 1, + "hidden_dim": 11008, + "n_heads": 16, + "n_kv_heads": 4, + "n_layers": 36, + "norm_eps": 1e-06, + "rope_theta": 5000000.0, + "use_scaled_rope": false, + "vocab_size": 128256, + "use_hf_rope": false, + "no_rope_layer_interval": 4, + "attention_qkv_bias": false + } diff --git a/examples/models/smollm3/__init__.py b/examples/models/smollm3/__init__.py new file mode 100644 index 00000000000..627cbc631f0 --- /dev/null +++ b/examples/models/smollm3/__init__.py @@ -0,0 +1,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 executorch.examples.models.llama.model import Llama2Model +from executorch.examples.models.smollm3.convert_weights import convert_weights + + +class SmolLM3Model(Llama2Model): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + +__all__ = [ + "SmolLM3Model", + "convert_weights", +] diff --git a/examples/models/smollm3/convert_weights.py b/examples/models/smollm3/convert_weights.py new file mode 100644 index 00000000000..df51674c3b0 --- /dev/null +++ b/examples/models/smollm3/convert_weights.py @@ -0,0 +1,112 @@ +import argparse +import json +import os +from typing import Dict + +import torch + +from safetensors.torch import load_file + +from torchtune.models.convert_weights import get_mapped_key + + +_SMOLLM_TO_META = { + "model.embed_tokens.weight": "tok_embeddings.weight", + "model.norm.weight": "norm.weight", + "model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight", + "model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight", + "model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight", + "model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight", + "model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight", + "model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight", + "model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight", + "model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight", + "model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight", +} + + +def smollm_to_meta(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """ + Convert a state dict from torchtune's format to Meta's format. This function + doesn't handle any sharding or splitting of state dicts. It follows the + state_dict IN -> state_dict OUT pattern. + + Args: + state_dict (Dict[str, torch.Tensor]): State dict in torchtune's format. + + Returns: + Dict[str, torch.Tensor]: State dict in Meta's format. + """ + converted_state_dict = {} + for key, value in state_dict.items(): + new_key = get_mapped_key(key, _SMOLLM_TO_META) + converted_state_dict[new_key] = value + converted_state_dict["output.weight"] = converted_state_dict[ + "tok_embeddings.weight" + ] + + return converted_state_dict + + +def load_checkpoint_from_safetensors(input_dir: str) -> Dict: + index_path = os.path.join(input_dir, "model.safetensors.index.json") + if os.path.exists(index_path): + # Sharded checkpoint. + with open(index_path, "r") as f: + index = json.load(f) + weight_map = index["weight_map"] + checkpoint_shards = sorted(set(weight_map.values())) + + # Load all the shards into memory + shard_to_weights = {} + for shard in checkpoint_shards: + shard_to_weights[shard] = load_file(os.path.join(input_dir, shard)) + + # Merge tensors into consolidated state dict. + merged_state_dict = {} + for weight_name, shard in weight_map.items(): + tensor = shard_to_weights[shard][weight_name] + merged_state_dict[weight_name] = tensor + return merged_state_dict + else: + # Single checkpoint. + state_dict = load_file(os.path.join(input_dir, "model.safetensors")) + return state_dict + + +def load_checkpoint(input_dir: str) -> Dict: + pytorch_path = os.path.join(input_dir, "pytorch_model.bin") + if os.path.exists(pytorch_path): + print("Loading checkpoint from PyTorch .bin file") + return torch.load(pytorch_path, map_location="cpu", weights_only=True) + print("Loading checkpoint from safetensors directory") + return load_checkpoint_from_safetensors(input_dir) + + +def convert_weights(input_dir: str, output_file: str) -> None: + print("Loading checkpoint...") + sd = load_checkpoint(input_dir) + print("Converting checkpoint...") + sd = smollm_to_meta(sd) + print("Saving checkpoint...") + torch.save(sd, output_file) + print("Done.") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert SmolLM weights to Meta format." + ) + parser.add_argument( + "input_dir", + type=str, + help="Path to directory containing checkpoint files", + ) + parser.add_argument("output", type=str, help="Path to the output checkpoint") + + args = parser.parse_args() + convert_weights(args.input_dir, args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/models/test/test_export.py b/examples/models/test/test_export.py index 306f54c0e89..53b6278bd6c 100644 --- a/examples/models/test/test_export.py +++ b/examples/models/test/test_export.py @@ -29,9 +29,7 @@ def collect_executorch_and_eager_outputs( Returns a tuple containing the outputs of the eager mode model and the executorch mode model. """ eager_model = eager_model.eval() - model = torch.export.export_for_training( - eager_model, example_inputs, strict=True - ).module() + model = torch.export.export(eager_model, example_inputs, strict=True).module() edge_model = export_to_edge(model, example_inputs) executorch_prog = edge_model.to_executorch() diff --git a/examples/models/voxtral/CMakeLists.txt b/examples/models/voxtral/CMakeLists.txt new file mode 100644 index 00000000000..85c6a13e0ff --- /dev/null +++ b/examples/models/voxtral/CMakeLists.txt @@ -0,0 +1,102 @@ +# 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. + +# +# Simple CMake build system for voxtral runner. +# +cmake_minimum_required(VERSION 3.24) +project(voxtral) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) + +include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) + +if(CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") + set(CMAKE_TOOLCHAIN_IOS ON) +else() + set(CMAKE_TOOLCHAIN_IOS OFF) +endif() + +if(NOT CMAKE_CXX_STANDARD) + set(CMAKE_CXX_STANDARD 17) +endif() + +# Let files say "include " +set(_common_include_directories ${EXECUTORCH_ROOT}/..) + +# Need this for gflags for some reason +set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../third-party/gflags) +find_package(gflags REQUIRED) + +# Find `executorch` libraries, same as for gflags +list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../..) +find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH) +executorch_target_link_options_shared_lib(executorch) + +set(link_libraries executorch gflags) +set(_srcs multimodal.cpp) + +list( + APPEND + link_libraries + optimized_native_cpu_ops_lib + quantized_ops_lib + custom_ops + cpublas + eigen_blas +) +executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib) +executorch_target_link_options_shared_lib(quantized_ops_lib) +executorch_target_link_options_shared_lib(custom_ops) + +# XNNPACK +if(TARGET xnnpack_backend) + set(xnnpack_backend_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod) + if(TARGET kleidiai) + list(APPEND xnnpack_backend_libs kleidiai) + endif() + list(APPEND link_libraries ${xnnpack_backend_libs}) + executorch_target_link_options_shared_lib(xnnpack_backend) +endif() + +# Add LLM runner and extension module +if(NOT TARGET extension_llm_runner) + message( + FATAL_ERROR + "ExecuTorch must be installed with EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER enabled." + ) +endif() + +# Needed for cpuinfo where it uses android specific log lib +if(ANDROID) + list(APPEND link_libraries log) +endif() + +# Add the required ExecuTorch extensions for multimodal LLM runner +list( + APPEND + link_libraries + extension_llm_runner + extension_module + extension_data_loader + extension_tensor + extension_flat_tensor +) + +# Add tokenizers +list(APPEND link_libraries tokenizers::tokenizers) + +add_executable(voxtral_runner ${_srcs}) +if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") + target_link_options_gc_sections(voxtral_runner) + if(NOT APPLE) + target_link_options(voxtral_runner PRIVATE "LINKER:-s") + endif() +endif() + +target_include_directories(voxtral_runner PUBLIC ${_common_include_directories}) +target_link_libraries(voxtral_runner PUBLIC ${link_libraries}) +target_compile_options(voxtral_runner PUBLIC ${_common_compile_options}) diff --git a/examples/models/voxtral/README.md b/examples/models/voxtral/README.md new file mode 100644 index 00000000000..5bc675e0615 --- /dev/null +++ b/examples/models/voxtral/README.md @@ -0,0 +1,105 @@ +# Summary + +This example demonstrates how to export and run Mistral's [Voxtral](https://huggingface.co/mistralai/Voxtral-Mini-3B-2507) audio multimodal model locally on ExecuTorch. + +# Exporting the model +To export the model, we use [Optimum ExecuTorch](https://github.com/huggingface/optimum-executorch), a repo that enables exporting models straight from the source - from HuggingFace's Transformers repo. + +## Setting up Optimum ExecuTorch +Install through pip package: +``` +pip install optimum-executorch +``` + +Or install from source: +``` +git clone https://github.com/huggingface/optimum-executorch.git +cd optimum-executorch +python install_dev.py +``` + +We are currently working on a Transformers pin bump for Optimum. In the meantime, manually override the Transformers dep to the earliest compatible version. +``` +pip install git+https://github.com/huggingface/transformers@6121e9e46c4fc4e5c91d9f927aef5490691850cf#egg=transformers +``` + +## Using the export CLI +We export Voxtral using the Optimum CLI, which will export `model.pte` to the `voxtral` output directory: +``` +optimum-cli export executorch \ + --model "mistralai/Voxtral-Mini-3B-2507" \ + --task "multimodal-text-to-text" \ + --recipe "xnnpack" \ + --use_custom_sdpa \ + --use_custom_kv_cache \ + --qlinear 8da4w \ + --qembedding 4w \ + --output_dir="voxtral" +``` + +This exports Voxtral with XNNPack backend acceleration and 4-bit weight/8-bit activation linear quantization. + +# Running the model +To run the model, we will use the Voxtral runner, which utilizes ExecuTorch's MultiModal runner API. +The Voxtral runner will do the following things: + +- Audio Input: + - Option A: Pass the raw audio tensor into exported preprocessor to produce a mel spectrogram tensor. + - Option B: If starting directly with an already processed audio input tensor, format the inputs to the multimodal runner (metadata tokens, audio tokens, text tokens, etc.). +- Feed the formatted inputs to the multimodal modal runner. + + +# [Option A] Exporting the audio preprocessor +The exported model takes in a mel spectrogram input tensor as its audio inputs. +We provide a simple way to transform raw audio data into a mel spectrogram by exporting a version of Voxtral's audio preprocessor used directly by Transformers. + +``` +python -m executorch.extension.audio.mel_spectrogram --feature_size 128 --output_file voxtral_preprocessor.pte +``` + +## Building the multimodal runner +``` +# Build and install ExecuTorch +cmake --preset llm -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=cmake-out -DEXECUTORCH_ENABLE_LOGGING=ON && cmake --build cmake-out -j16 --target install --config Release + +# Build and install Voxtral runner +cmake -DCMAKE_INSTALL_PREFIX=cmake-out -DBUILD_TESTING=OFF -DCMAKE_BUILD_TYPE=Release -Bcmake-out/examples/models/voxtral examples/models/voxtral && cmake --build cmake-out/examples/models/voxtral -j16 --config Release +``` + +## Running the model +You can download the `tekken.json` tokenizer from [Voxtral's HuggingFace repo](https://huggingface.co/mistralai/Voxtral-Mini-3B-2507). +``` +./cmake-out/examples/models/voxtral/voxtral_runner \ + --model_path path/to/model.pte \ + --tokenizer_path path/to/tekken.json \ + --prompt "What can you tell me about this audio?" \ + --audio_path path/to/audio_input.bin \ + --processor_path path/to/voxtral_preprocessor.pte # If you're passing raw audio file in audio_path +``` + +Example output: +``` +The speaker in this audio seems to be talking about their concerns about a device called the model or maybe they're just talking about the model in general. They mention that the model was trained with the speaker for inference, which suggests that + the model was trained based on the speaker's data or instructions. They also mention that the volume is quite small, which could imply that the speaker is trying to control the volume of the model's output, likely because they are concerned about how loud the model's responses might +PyTorchObserver {"prompt_tokens":388,"generated_tokens":99,"model_load_start_ms":0,"model_load_end_ms":0,"inference_start_ms":1756351346381,"inference_end_ms":1756351362602,"prompt_eval_end_ms":1756351351435,"first_token_ms":1756351351435,"aggregate_sampling_time_ms":99,"SCALING_FACTOR_UNITS_PER_SECOND":1000} +I 00:00:24.036773 executorch:stats.h:104] Prompt Tokens: 388 Generated Tokens: 99 +I 00:00:24.036800 executorch:stats.h:110] Model Load Time: 0.000000 (seconds) +I 00:00:24.036805 executorch:stats.h:117] Total inference time: 16.221000 (seconds) Rate: 6.103200 (tokens/second) +I 00:00:24.036815 executorch:stats.h:127] Prompt evaluation: 5.054000 (seconds) Rate: 76.770875 (tokens/second) +I 00:00:24.036819 executorch:stats.h:136] Generated 99 tokens: 11.167000 (seconds) Rate: 8.865407 (tokens/second) +I 00:00:24.036822 executorch:stats.h:147] Time to first generated token: 5.054000 (seconds) +I 00:00:24.036828 executorch:stats.h:153] Sampling time over 487 tokens: 0.099000 (seconds) +``` + +You can easily produce an `.bin` for the audio input in Python like this: +``` +# t = some torch.Tensor +with open("tensor.bin", "wb") as f: + f.write(t.numpy().tobytes()) +``` + +You can also produce raw audio file as follows (for Option A): + +``` +ffmpeg -i audio.mp3 -f f32le -acodec pcm_f32le -ar 16000 audio_input.bin +``` diff --git a/examples/models/voxtral/multimodal.cpp b/examples/models/voxtral/multimodal.cpp new file mode 100644 index 00000000000..17013df96e1 --- /dev/null +++ b/examples/models/voxtral/multimodal.cpp @@ -0,0 +1,355 @@ +/* + * 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 + +#if defined(ET_USE_THREADPOOL) +#include +#include +#endif + +DEFINE_string( + model_path, + "multimodal.pte", + "Model serialized in flatbuffer format."); + +DEFINE_string(tokenizer_path, "tekken.json", "Tokenizer stuff."); + +DEFINE_string(prompt, "What is happening in this audio?", "Text prompt."); + +DEFINE_string(audio_path, "", "Path to input audio file."); + +DEFINE_string( + processor_path, + "", + "Path to processor .pte file for raw audio processing."); + +DEFINE_double( + temperature, + 0.8f, + "Temperature; Default is 0.8f. 0 = greedy argmax sampling (deterministic). Lower temperature = more deterministic"); + +DEFINE_int32( + cpu_threads, + -1, + "Number of CPU threads for inference. Defaults to -1, which implies we'll use a heuristic to derive the # of performant cores for a specific device."); + +DEFINE_bool(warmup, false, "Whether to run a warmup run."); + +namespace { + +using ::executorch::extension::from_blob; +using ::executorch::extension::Module; +using ::executorch::extension::llm::Image; +using ::executorch::extension::llm::make_image_input; +using ::executorch::extension::llm::make_text_input; +using ::executorch::extension::llm::MultimodalInput; +using ::executorch::runtime::EValue; + +bool ends_with(const std::string& str, const std::string& suffix) { + return str.size() >= suffix.size() && + str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; +} + +/** + * @brief Loads preprocessed audio data from a binary file + * + * Reads mel spectrogram features that have been pre-computed and saved as a + * binary file. The audio data is expected to be stored as float values in + * binary format, typically saved using: + * with open("tensor.bin", "wb") as f: + * f.write(t.numpy().tobytes()) + * + * @param audio_path Path to the binary audio file (.bin) + * @return MultimodalInput containing the loaded audio data + */ +MultimodalInput loadPreprocessedAudio(const std::string& audio_path) { + std::ifstream f(audio_path, std::ios::binary | std::ios::ate); + if (!f.is_open()) { + ET_LOG(Error, "Failed to open audio file: %s", audio_path.c_str()); + throw std::runtime_error("Failed to open audio file"); + } + + std::size_t n_floats = f.tellg() / sizeof(float); + f.seekg(0, std::ios::beg); + + int32_t n_bins = 128; + int32_t n_frames = 3000; + + int32_t batch_size = ceil( + n_floats / + (n_bins * n_frames)); // Batch in increments of n_frames, rounding up. + + ET_LOG(Info, "audio_data len = %zu", n_floats); + + // Create Audio multimodal input + auto audio = std::make_unique<::executorch::extension::llm::Audio>(); + audio->batch_size = batch_size; + audio->n_bins = n_bins; + audio->n_frames = n_frames; + audio->data.resize(n_floats * sizeof(float)); + f.read(reinterpret_cast(audio->data.data()), n_floats * sizeof(float)); + f.close(); + return ::executorch::extension::llm::make_audio_input(std::move(*audio)); +} + +/** + * @brief Loads a .bin file into a tensor and processes it using a .pte + * processor + * + * This function loads raw audio data from a .bin file (similar to + * loadPreprocessedAudio), creates a tensor from it, and then passes it through + * a processor module loaded from a .pte file to generate processed audio + * features. + * + * @param audio_path Path to the .bin audio file + * @param processor_path Path to the .pte processor file + * @return MultimodalInput containing the processed audio data + * @throws std::runtime_error if file loading or processing fails + */ +MultimodalInput processRawAudioFile( + const std::string& audio_path, + const std::string& processor_path) { + if (processor_path.empty()) { + ET_LOG(Error, "Processor path is required for raw audio processing"); + throw std::runtime_error( + "Processor path is required for raw audio processing"); + } + + // Load the audio processor .pte. + std::unique_ptr processor_module; + try { + processor_module = + std::make_unique(processor_path, Module::LoadMode::File); + auto load_error = processor_module->load(); + if (load_error != ::executorch::runtime::Error::Ok) { + ET_LOG( + Error, + "Failed to load processor module from: %s", + processor_path.c_str()); + throw std::runtime_error("Failed to load processor module"); + } + } catch (const std::exception& e) { + ET_LOG(Error, "Exception while loading processor module: %s", e.what()); + throw std::runtime_error("Exception while loading processor module"); + } + + // Load the audio data from file. + std::ifstream f(audio_path, std::ios::binary | std::ios::ate); + if (!f.is_open()) { + ET_LOG(Error, "Failed to open audio file: %s", audio_path.c_str()); + throw std::runtime_error("Failed to open audio file"); + } + + std::size_t n_floats = f.tellg() / sizeof(float); + f.seekg(0, std::ios::beg); + + std::vector audio_data(n_floats); + f.read( + reinterpret_cast(audio_data.data()), + audio_data.size() * sizeof(float)); + f.close(); + + ET_LOG( + Info, "Loaded .bin file: %s, %zu floats", audio_path.c_str(), n_floats); + + // Execute the processor + std::vector tensor_shape = { + static_cast(audio_data.size())}; + auto input_tensor = from_blob( + audio_data.data(), tensor_shape, ::executorch::aten::ScalarType::Float); + + ET_LOG(Info, "Processing audio through processor module..."); + auto result = processor_module->execute("forward", input_tensor); + if (!result.ok()) { + ET_LOG(Error, "Failed to execute processor's forward method"); + throw std::runtime_error("Failed to execute processor forward method"); + } + + auto outputs = result.get(); + if (outputs.empty()) { + ET_LOG(Error, "Processor returned no outputs"); + throw std::runtime_error("Processor returned no outputs"); + } + + // Extract processed audio features + const auto& processed_tensor = outputs[0].toTensor(); + const float* processed_data = processed_tensor.const_data_ptr(); + const auto& sizes = processed_tensor.sizes(); + + ET_LOG( + Info, + "Processed audio tensor shape: [%d, %d, %d]", + static_cast(sizes[0]), + static_cast(sizes[1]), + static_cast(sizes[2])); + + // Create Audio multimodal input from processed features + auto processed_audio = + std::make_unique<::executorch::extension::llm::Audio>(); + processed_audio->batch_size = + static_cast(sizes[0]); // Note: batching for s > 30 doesn't work + // yet, so this will just be = 1. + processed_audio->n_bins = static_cast(sizes[1]); + processed_audio->n_frames = + static_cast(sizes[2]); // And this will just be = 3000. + + size_t total_elements = processed_audio->batch_size * + processed_audio->n_bins * processed_audio->n_frames; + processed_audio->data.resize(total_elements * sizeof(float)); + std::memcpy( + processed_audio->data.data(), + processed_data, + total_elements * sizeof(float)); + + ET_LOG( + Info, + "Created processed Audio: batch_size=%d, n_bins=%d, n_frames=%d", + processed_audio->batch_size, + processed_audio->n_bins, + processed_audio->n_frames); + + return ::executorch::extension::llm::make_audio_input( + std::move(*processed_audio)); +} + +/** + * @brief Processes audio files for multimodal input + * + * Dispatches audio file processing based on file extension and processor + * availability: + * - .bin files with processor: Loads raw audio from .bin and processes through + * processor + * - .bin files without processor: Loads preprocessed mel spectrogram features + * directly + * + * @param audio_path Path to the audio file (.bin) + * @param processor_path Path to the processor .pte file (optional) + * @return MultimodalInput containing the processed audio data + * @throws std::runtime_error if file format is unsupported or processing fails + */ +MultimodalInput processAudioFile( + const std::string& audio_path, + const std::string& processor_path = "") { + if (ends_with(audio_path, ".bin")) { + if (!processor_path.empty()) { + // Process raw audio from .bin file through the processor + return processRawAudioFile(audio_path, processor_path); + } else { + // Load preprocessed audio stored as a binary file (existing behavior) + return loadPreprocessedAudio(audio_path); + } + } else { + ET_LOG( + Error, + "Unsupported audio file format: %s (only .bin files are supported)", + audio_path.c_str()); + throw std::runtime_error("Unsupported audio file format"); + } +} + +} // namespace + +int32_t main(int32_t argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + + const char* model_path = FLAGS_model_path.c_str(); + + const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); + const char* prompt = FLAGS_prompt.c_str(); + const char* audio_path = FLAGS_audio_path.c_str(); + const char* processor_path = FLAGS_processor_path.c_str(); + float temperature = FLAGS_temperature; + int32_t cpu_threads = FLAGS_cpu_threads; + bool warmup = FLAGS_warmup; + +#if defined(ET_USE_THREADPOOL) + uint32_t num_performant_cores = cpu_threads == -1 + ? ::executorch::extension::cpuinfo::get_num_performant_cores() + : static_cast(cpu_threads); + ET_LOG( + Info, "Resetting threadpool with num threads = %d", num_performant_cores); + if (num_performant_cores > 0) { + ::executorch::extension::threadpool::get_threadpool() + ->_unsafe_reset_threadpool(num_performant_cores); + } +#endif + + // Load tokenizer + std::unique_ptr<::tokenizers::Tokenizer> tokenizer = + ::executorch::extension::llm::load_tokenizer(tokenizer_path); + if (tokenizer == nullptr) { + ET_LOG(Error, "Failed to load tokenizer from: %s", tokenizer_path); + return 1; + } + + // Create multimodal runner + std::unique_ptr<::executorch::extension::llm::MultimodalRunner> runner = + ::executorch::extension::llm::create_multimodal_runner( + model_path, std::move(tokenizer)); + if (runner == nullptr) { + ET_LOG(Error, "Failed to create multimodal runner"); + return 1; + } + + // Load runner + auto load_error = runner->load(); + if (load_error != ::executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load multimodal runner"); + return 1; + } + + // Prepare inputs + std::vector inputs = { + make_text_input("[INST][BEGIN_AUDIO]"), + processAudioFile(audio_path, processor_path), + make_text_input(std::string(prompt) + "[/INST]"), + }; + + ::executorch::extension::llm::GenerationConfig config; + config.max_new_tokens = 100; + config.temperature = temperature; + + // Run warmup if requested + if (warmup) { + ET_LOG(Info, "Running warmup..."); + auto warmup_error = runner->generate(inputs, config); + if (warmup_error != ::executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to run warmup"); + return 1; + } + runner->reset(); + } + + // Generate + ET_LOG(Info, "Starting generation..."); + auto error = runner->generate(inputs, config); + if (error != ::executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to generate with multimodal runner"); + return 1; + } + + printf("\n"); + return 0; +} diff --git a/examples/models/yolo12/CMakeLists.txt b/examples/models/yolo12/CMakeLists.txt index 9c92e5eaeae..60b11685bdf 100644 --- a/examples/models/yolo12/CMakeLists.txt +++ b/examples/models/yolo12/CMakeLists.txt @@ -37,7 +37,6 @@ set(link_libraries gflags) list(APPEND link_libraries portable_ops_lib portable_kernels) executorch_target_link_options_shared_lib(portable_ops_lib) - if(USE_XNNPACK_BACKEND) set(xnnpack_backend_libs xnnpack_backend XNNPACK microkernels-prod) list(APPEND link_libraries ${xnnpack_backend_libs}) @@ -49,9 +48,10 @@ if(USE_OPENVINO_BACKEND) target_include_directories( openvino_backend - INTERFACE ${CMAKE_CURRENT_BINARY_DIR}/../../include - ${CMAKE_CURRENT_BINARY_DIR}/../../include/executorch/runtime/core/portable_type/c10 - ${CMAKE_CURRENT_BINARY_DIR}/../../lib + INTERFACE + ${CMAKE_CURRENT_BINARY_DIR}/../../include + ${CMAKE_CURRENT_BINARY_DIR}/../../include/executorch/runtime/core/portable_type/c10 + ${CMAKE_CURRENT_BINARY_DIR}/../../lib ) list(APPEND link_libraries openvino_backend) executorch_target_link_options_shared_lib(openvino_backend) @@ -72,14 +72,13 @@ set(PROJECT_SOURCES ) add_executable(Yolo12DetectionDemo ${PROJECT_SOURCES}) -target_link_libraries(Yolo12DetectionDemo PUBLIC - ${link_libraries} - ${OpenCV_LIBS} - executorch_core - extension_module - extension_tensor +target_link_libraries( + Yolo12DetectionDemo PUBLIC ${link_libraries} ${OpenCV_LIBS} executorch_core + extension_module extension_tensor ) find_package(Threads REQUIRED) target_link_libraries(Yolo12DetectionDemo PRIVATE Threads::Threads) -target_include_directories(Yolo12DetectionDemo PUBLIC ${_common_include_directories}) \ No newline at end of file +target_include_directories( + Yolo12DetectionDemo PUBLIC ${_common_include_directories} +) diff --git a/examples/nxp/README.md b/examples/nxp/README.md index 66ca0785b4c..4d9831c73f8 100644 --- a/examples/nxp/README.md +++ b/examples/nxp/README.md @@ -13,7 +13,7 @@ $ examples/nxp/setup.sh 2. Now run the `aot_neutron_compile.py` example with the `cifar10` model ```commandline $ python -m examples.nxp.aot_neutron_compile --quantize \ - --delegate --neutron_converter_flavor SDK_25_03 -m cifar10 + --delegate --neutron_converter_flavor SDK_25_06 -m cifar10 ``` 3. It will generate you `cifar10_nxp_delegate.pte` file which can be used with the MXUXpresso SDK `cifarnet_example` project, presented [here](https://mcuxpresso.nxp.com/mcuxsdk/latest/html/middleware/eiq/executorch/docs/nxp/topics/example_applications.html#how-to-build-and-run-executorch-cifarnet-example). diff --git a/examples/nxp/aot_neutron_compile.py b/examples/nxp/aot_neutron_compile.py index 5c0634697d0..cdaf913a9ea 100644 --- a/examples/nxp/aot_neutron_compile.py +++ b/examples/nxp/aot_neutron_compile.py @@ -15,7 +15,6 @@ import executorch.kernels.quantized # noqa F401 import torch - from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import ( RemoveIOQuantOpsPass, ) @@ -24,19 +23,19 @@ from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.examples.models import MODEL_NAME_TO_MODEL from executorch.examples.models.model_factory import EagerModelFactory - from executorch.exir import ( EdgeCompileConfig, ExecutorchBackendConfig, to_edge_transform_and_lower, ) from executorch.extension.export_util import save_pte_program - from torch.export import export from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from .experimental.cifar_net.cifar_net import CifarNet, test_cifarnet_model +from .models.mobilenet_v2 import MobilenetV2 + FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) @@ -87,7 +86,7 @@ def get_model_and_inputs_from_name(model_name: str): logging.warning( "Using a model from examples/models not all of these are currently supported" ) - model, example_inputs, _ = EagerModelFactory.create_model( + model, example_inputs, _, _ = EagerModelFactory.create_model( *MODEL_NAME_TO_MODEL[model_name] ) else: @@ -100,6 +99,7 @@ def get_model_and_inputs_from_name(model_name: str): models = { "cifar10": CifarNet, + "mobilenetv2": MobilenetV2, } @@ -164,7 +164,7 @@ def _get_batch_size(data): "-c", "--neutron_converter_flavor", required=False, - default="SDK_25_03", + default="SDK_25_06", help="Flavor of installed neutron-converter module. Neutron-converter module named " "'neutron_converter_SDK_24_12' has flavor 'SDK_24_12'.", ) @@ -224,9 +224,7 @@ def _get_batch_size(data): model = model.eval() # 2. Export the model to ATEN - exported_program = torch.export.export_for_training( - model, example_inputs, strict=True - ) + exported_program = torch.export.export(model, example_inputs, strict=True) module = exported_program.module() diff --git a/examples/nxp/experimental/cifar_net/cifar_net.pth b/examples/nxp/experimental/cifar_net/cifar_net.pth index 6dc4efde21d..63c49bf494b 100644 Binary files a/examples/nxp/experimental/cifar_net/cifar_net.pth and b/examples/nxp/experimental/cifar_net/cifar_net.pth differ diff --git a/examples/nxp/experimental/cifar_net/cifar_net.py b/examples/nxp/experimental/cifar_net/cifar_net.py index 1378d00cf12..8d057c1ca34 100644 --- a/examples/nxp/experimental/cifar_net/cifar_net.py +++ b/examples/nxp/experimental/cifar_net/cifar_net.py @@ -1,4 +1,4 @@ -# Copyright 2024 NXP +# Copyright 2024-2025 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -57,7 +57,7 @@ class CifarNetModel(nn.Module): def __init__(self): super().__init__() - self.conv1 = nn.Conv2d(8, 32, 5) + self.conv1 = nn.Conv2d(3, 32, 5) self.conv2 = nn.Conv2d(32, 32, 5) self.conv3 = nn.Conv2d(32, 64, 5) self.pool1 = nn.MaxPool2d(2, 2) @@ -66,10 +66,7 @@ def __init__(self): self.softmax = nn.Softmax(1) def forward(self, x): - - # Neutron Backend does not yet have passses for automated padding if number of channels does not - # fit to Neutron constrains (#channels == #MAC units). So define the model explicitly tailored for Neutron-C-64. - x = F.pad(x, (2, 2, 2, 2, 0, 5)) + x = F.pad(x, (2, 2, 2, 2)) x = self.conv1(x) x = self.pool1(x) diff --git a/examples/nxp/models/mobilenet_v2.py b/examples/nxp/models/mobilenet_v2.py new file mode 100644 index 00000000000..ccda4155d39 --- /dev/null +++ b/examples/nxp/models/mobilenet_v2.py @@ -0,0 +1,114 @@ +# Copyright 2025 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 itertools +from typing import Iterator + +import torch +import torchvision + +from executorch.examples.models.mobilenet_v2 import MV2Model +from torch.utils.data import DataLoader +from torchvision import transforms + + +class MobilenetV2(MV2Model): + + def get_calibration_inputs( + self, batch_size: int = 1 + ) -> Iterator[tuple[torch.Tensor]]: + """ + Returns an iterator for the Imagenette validation dataset, downloading it if necessary. + + Args: + batch_size (int): The batch size for the iterator. + + Returns: + iterator: An iterator that yields batches of images from the Imagnetette validation dataset. + """ + dataloader = self.get_dataset(batch_size) + + # Return the iterator + dataloader_iterable = itertools.starmap( + lambda data, label: (data,), iter(dataloader) + ) + + # We want approximately 500 samples + batch_count = 500 // batch_size + return itertools.islice(dataloader_iterable, batch_count) + + def get_dataset(self, batch_size): + # Define data transformations + data_transforms = transforms.Compose( + [ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] + ), # ImageNet stats + ] + ) + + dataset = torchvision.datasets.Imagenette( + root="./data", split="val", transform=data_transforms, download=True + ) + dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + num_workers=1, + ) + return dataloader + + +def gather_samples_per_class_from_dataloader( + dataloader, num_samples_per_class=10 +) -> list[tuple]: + """ + Gathers a specified number of samples for each class from a DataLoader. + + Args: + dataloader (DataLoader): The PyTorch DataLoader object. + num_samples_per_class (int): The number of samples to gather for each class. Defaults to 10. + + Returns: + samples: A list of (sample, label) tuples. + """ + + if not isinstance(dataloader, DataLoader): + raise TypeError("dataloader must be a torch.utils.data.DataLoader object") + if not isinstance(num_samples_per_class, int) or num_samples_per_class <= 0: + raise ValueError("num_samples_per_class must be a positive integer") + + labels = sorted( + set([label for _, label in dataloader.dataset]) + ) # Get unique labels from the dataset + samples_per_label = {label: [] for label in labels} # Initialize dictionary + + for sample, label in dataloader: + label = label.item() + if len(samples_per_label[label]) < num_samples_per_class: + samples_per_label[label].append((sample, label)) + + samples = [] + + for label in labels: + samples.extend(samples_per_label[label]) + + return samples + + +def generate_input_samples_file(): + model = MobilenetV2() + dataloader = model.get_dataset(batch_size=1) + samples = gather_samples_per_class_from_dataloader( + dataloader, num_samples_per_class=2 + ) + + torch.save(samples, "calibration_data.pt") + + +if __name__ == "__main__": + generate_input_samples_file() diff --git a/examples/nxp/run_aot_example.sh b/examples/nxp/run_aot_example.sh index 1710490f6d7..fa8c318778d 100755 --- a/examples/nxp/run_aot_example.sh +++ b/examples/nxp/run_aot_example.sh @@ -7,11 +7,12 @@ set -eux SCRIPT_DIR=$(dirname $(readlink -fm $0)) EXECUTORCH_DIR=$(dirname $(dirname $SCRIPT_DIR)) +MODEL=${1:-"cifar10"} cd $EXECUTORCH_DIR # Run the AoT example python -m examples.nxp.aot_neutron_compile --quantize \ - --delegate --neutron_converter_flavor SDK_25_03 -m cifar10 + --delegate --neutron_converter_flavor SDK_25_06 -m ${MODEL} # verify file exists -test -f cifar10_nxp_delegate.pte +test -f ${MODEL}_nxp_delegate.pte diff --git a/examples/nxp/setup.sh b/examples/nxp/setup.sh index 1a050a79c19..74038220e74 100755 --- a/examples/nxp/setup.sh +++ b/examples/nxp/setup.sh @@ -7,4 +7,4 @@ set -u # Install neutron-converter -pip install --extra-index-url https://eiq.nxp.com/repository neutron_converter_SDK_25_03 +pip install --index-url https://eiq.nxp.com/repository neutron_converter_SDK_25_06 diff --git a/examples/portable/custom_ops/CMakeLists.txt b/examples/portable/custom_ops/CMakeLists.txt index c4a00e47991..4188554af79 100644 --- a/examples/portable/custom_ops/CMakeLists.txt +++ b/examples/portable/custom_ops/CMakeLists.txt @@ -59,15 +59,9 @@ option( # ------------------------------- OPTIONS END -------------------------------- # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../../executorch_srcs.cmake" -) - -extract_sources(${EXECUTORCH_SRCS_FILE}) - -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() # Generate C++ bindings to register kernels into both PyTorch (for AOT) and # Executorch (for runtime). diff --git a/examples/portable/executor_runner/executor_runner.cpp b/examples/portable/executor_runner/executor_runner.cpp index 434b4783bac..4f4208a5b53 100644 --- a/examples/portable/executor_runner/executor_runner.cpp +++ b/examples/portable/executor_runner/executor_runner.cpp @@ -1,7 +1,7 @@ /* * Copyright (c) Meta Platforms, Inc. and affiliates. - * Copyright 2024-2025 Arm Limited and/or its affiliates. * All rights reserved. + * Copyright 2024-2025 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. @@ -18,6 +18,7 @@ * all fp32 tensors. */ +#include #include #include @@ -49,6 +50,16 @@ DEFINE_string( model_path, "model.pte", "Model serialized in flatbuffer format."); +DEFINE_string(inputs, "", "Comma-separated list of input files"); +DEFINE_string( + output_file, + "", + "Base name of output file. If not empty output will be written to the file(s)."); + +DEFINE_bool( + print_all_output, + false, + "Prints all output. By default only first and last 100 elements are printed."); DEFINE_uint32(num_executions, 1, "Number of times to run the model."); #ifdef ET_EVENT_TRACER_ENABLED DEFINE_string(etdump_path, "model.etdump", "Write ETDump data to this path."); @@ -58,6 +69,8 @@ DEFINE_int32( -1, "Number of CPU threads for inference. Defaults to -1, which implies we'll use a heuristic to derive the # of performant cores for a specific device."); +using executorch::aten::ScalarType; +using executorch::aten::Tensor; using executorch::extension::FileDataLoader; using executorch::runtime::Error; using executorch::runtime::EValue; @@ -70,6 +83,8 @@ using executorch::runtime::MethodMeta; using executorch::runtime::Program; using executorch::runtime::Result; using executorch::runtime::Span; +using executorch::runtime::Tag; +using executorch::runtime::TensorInfo; /// Helper to manage resources for ETDump generation class EventTraceManager { @@ -156,6 +171,31 @@ int main(int argc, char** argv) { "FileDataLoader::from() failed: 0x%" PRIx32, (uint32_t)loader.error()); + std::vector inputs_storage; + std::vector> input_buffers; + + std::stringstream list_of_input_files(FLAGS_inputs); + std::string token; + + while (std::getline(list_of_input_files, token, ',')) { + std::ifstream input_file_handle(token, std::ios::binary | std::ios::ate); + if (!input_file_handle) { + ET_LOG(Error, "Failed to open input file: %s\n", token.c_str()); + return 1; + } + + std::streamsize file_size = input_file_handle.tellg(); + input_file_handle.seekg(0, std::ios::beg); + + inputs_storage.emplace_back(file_size, '\0'); + if (!input_file_handle.read(&inputs_storage.back()[0], file_size)) { + ET_LOG(Error, "Failed to read input file: %s\n", token.c_str()); + return 1; + } + + input_buffers.emplace_back(&inputs_storage.back()[0], file_size); + } + // Parse the program file. This is immutable, and can also be reused between // multiple execution invocations across multiple threads. Result program = Program::load(&loader.get()); @@ -254,7 +294,8 @@ int main(int argc, char** argv) { // Run the model. for (uint32_t i = 0; i < FLAGS_num_executions; i++) { ET_LOG(Debug, "Preparing inputs."); - // Allocate input tensors and set all of their elements to 1. The `inputs` + // Allocate input tensors and set all of their elements to 1 or to the + // contents of input_buffers if available. The `inputs` // variable owns the allocated memory and must live past the last call to // `execute()`. // @@ -262,7 +303,8 @@ int main(int argc, char** argv) { // because inputs whose space gets reused by memory planning (if // any such inputs exist) will not be preserved for the next // execution. - auto inputs = executorch::extension::prepare_input_tensors(*method); + auto inputs = executorch::extension::prepare_input_tensors( + *method, {}, input_buffers); ET_CHECK_MSG( inputs.ok(), "Could not prepare inputs: 0x%" PRIx32, @@ -295,10 +337,66 @@ int main(int argc, char** argv) { ET_LOG(Info, "%zu outputs: ", outputs.size()); Error status = method->get_outputs(outputs.data(), outputs.size()); ET_CHECK(status == Error::Ok); - // Print the first and last 100 elements of long lists of scalars. - std::cout << executorch::extension::evalue_edge_items(100); - for (int i = 0; i < outputs.size(); ++i) { - std::cout << "Output " << i << ": " << outputs[i] << std::endl; + + if (FLAGS_output_file.size() > 0) { + for (int i = 0; i < outputs.size(); ++i) { + if (outputs[i].isTensor()) { + Tensor tensor = outputs[i].toTensor(); + + char out_filename[255]; + snprintf(out_filename, 255, "%s-%d.bin", FLAGS_output_file.c_str(), i); + ET_LOG(Info, "Writing output to file: %s", out_filename); + FILE* out_file = fopen(out_filename, "wb"); + fwrite(tensor.const_data_ptr(), 1, tensor.nbytes(), out_file); + fclose(out_file); + } + } + } + + if (FLAGS_print_all_output) { + for (int i = 0; i < outputs.size(); ++i) { + if (outputs[i].isTensor()) { + Tensor tensor = outputs[i].toTensor(); + + for (int j = 0; j < tensor.numel(); ++j) { + if (tensor.scalar_type() == ScalarType::Int) { + printf( + "Output[%d][%d]: (int) %d\n", + i, + j, + tensor.const_data_ptr()[j]); + } else if (tensor.scalar_type() == ScalarType::Float) { + printf( + "Output[%d][%d]: (float) %f\n", + i, + j, + tensor.const_data_ptr()[j]); + } else if (tensor.scalar_type() == ScalarType::Char) { + printf( + "Output[%d][%d]: (char) %d\n", + i, + j, + tensor.const_data_ptr()[j]); + } else if (tensor.scalar_type() == ScalarType::Bool) { + printf( + "Output[%d][%d]: (bool) %s (0x%x)\n", + i, + j, + tensor.const_data_ptr()[j] ? "true " : "false", + tensor.const_data_ptr()[j]); + } + } + } else { + printf("Output[%d]: Not Tensor\n", i); + } + } + } else { + // Print the first and last 100 elements of long lists of scalars. + std::cout << executorch::extension::evalue_edge_items(100); + + for (int i = 0; i < outputs.size(); ++i) { + std::cout << "OutputX " << i << ": " << outputs[i] << std::endl; + } } if (tracer.get_event_tracer()) { diff --git a/examples/portable/scripts/export_and_delegate.py b/examples/portable/scripts/export_and_delegate.py index 1c2adf67688..509916959b0 100644 --- a/examples/portable/scripts/export_and_delegate.py +++ b/examples/portable/scripts/export_and_delegate.py @@ -61,7 +61,7 @@ def export_composite_module_with_lower_graph(): m_compile_spec = m.get_compile_spec() # pre-autograd export. eventually this will become torch.export - m = torch.export.export_for_training(m, m_inputs, strict=True).module() + m = torch.export.export(m, m_inputs, strict=True).module() edge = export_to_edge(m, m_inputs) logging.info(f"Exported graph:\n{edge.exported_program().graph}") @@ -84,7 +84,7 @@ def forward(self, *args): m = CompositeModule() m = m.eval() # pre-autograd export. eventually this will become torch.export - m = torch.export.export_for_training(m, m_inputs, strict=True).module() + m = torch.export.export(m, m_inputs, strict=True).module() composited_edge = export_to_edge(m, m_inputs) # The graph module is still runnerable @@ -134,7 +134,7 @@ def get_example_inputs(self): m = Model() m_inputs = m.get_example_inputs() # pre-autograd export. eventually this will become torch.export - m = torch.export.export_for_training(m, m_inputs, strict=True).module() + m = torch.export.export(m, m_inputs, strict=True).module() edge = export_to_edge(m, m_inputs) logging.info(f"Exported graph:\n{edge.exported_program().graph}") @@ -171,7 +171,7 @@ def export_and_lower_the_whole_graph(): m_inputs = m.get_example_inputs() # pre-autograd export. eventually this will become torch.export - m = torch.export.export_for_training(m, m_inputs, strict=True).module() + m = torch.export.export(m, m_inputs, strict=True).module() edge = export_to_edge(m, m_inputs) logging.info(f"Exported graph:\n{edge.exported_program().graph}") diff --git a/examples/qualcomm/CMakeLists.txt b/examples/qualcomm/CMakeLists.txt index d33d666b9c0..19190b6f794 100644 --- a/examples/qualcomm/CMakeLists.txt +++ b/examples/qualcomm/CMakeLists.txt @@ -35,13 +35,9 @@ find_package(gflags REQUIRED) set(_common_compile_options -Wno-deprecated-declarations -fPIC) # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../executorch_srcs.cmake" -) -extract_sources(${EXECUTORCH_SRCS_FILE}) -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() get_filename_component( EXECUTORCH_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE diff --git a/examples/qualcomm/README.md b/examples/qualcomm/README.md index d33d4d2c52b..355209f43a7 100644 --- a/examples/qualcomm/README.md +++ b/examples/qualcomm/README.md @@ -2,17 +2,17 @@ This directory contains examples for some AI models. -We have seperated the example scripts into the following subfolders, please refer to [README.md](../../backends/qualcomm/README.md) for the example scripts' directory structure: +We have separated the example scripts into the following subfolders, please refer to [README.md](../../backends/qualcomm/README.md) for the example scripts' directory structure: 1. executor_runner: This folder contains a general executor runner capable of running most of the models. As a rule of thumb, if a model does not have its own customized runner, execute the model using [executor_runner](executor_runner/qnn_executor_runner.cpp). On the other hand, if a model has its own runner, such as [llama](oss_scripts/llama/qnn_llama_runner.cpp), use the customized runner to execute the model. Customized runner should be located under the same folder as the model's python script. 2. oss_scripts: OSS stands for Open Source Software. This folder contains python scripts for open source models. Some models under this folder might also have their own customized runner. For example, [llama](oss_scripts/llama/qnn_llama_runner.cpp) contains not only the python scripts to prepare the model but also a customized runner for executing the model. -3. qaihub_scripts: QAIHub stands for [Qualcomm AI Hub](https://aihub.qualcomm.com/). On QAIHub, users can find pre-compiled context binaries, a format used by QNN to save its models. This provides users with a new option for model deployment. Different from oss_scripts & scripts, which the example scripts are converting a model from nn.Module to ExecuTorch .pte files, qaihub_scripts provides example scripts for converting pre-compiled context binaries to ExecuTorch .pte files. Additionaly, users can find customized example runners specific to the QAIHub models for execution. For example [qaihub_llama2_7b](qaihub_scripts/llama/llama2/qaihub_llama2_7b.py) is a script converting context binaries to ExecuTorch .pte files, and [qaihub_llama2_7b_runner](qaihub_scripts/llama/llama2/qaihub_llama2_7b_runner.cpp) is a customized example runner to execute llama2 .pte files. Please be aware that context-binaries downloaded from QAIHub are tied to a specific QNN SDK version. +3. qaihub_scripts: QAIHub stands for [Qualcomm AI Hub](https://aihub.qualcomm.com/). On QAIHub, users can find pre-compiled context binaries, a format used by QNN to save its models. This provides users with a new option for model deployment. Different from oss_scripts & scripts, which the example scripts are converting a model from nn.Module to ExecuTorch .pte files, qaihub_scripts provides example scripts for converting pre-compiled context binaries to ExecuTorch .pte files. Additionally, users can find customized example runners specific to the QAIHub models for execution. For example [qaihub_llama2_7b](qaihub_scripts/llama/llama2/qaihub_llama2_7b.py) is a script converting context binaries to ExecuTorch .pte files, and [qaihub_llama2_7b_runner](qaihub_scripts/llama/llama2/qaihub_llama2_7b_runner.cpp) is a customized example runner to execute llama2 .pte files. Please be aware that context-binaries downloaded from QAIHub are tied to a specific QNN SDK version. Before executing the scripts and runner, please ensure that you are using the QNN SDK version that is matching the context binary. Please refer to [Check context binary version](#check-context-binary-version) for tutorial on how to check the QNN Version for a context binary. -4. scripts: This folder contains scripts to build models provided by Executorch. +4. scripts: This folder contains scripts to build models provided by ExecuTorch. @@ -22,7 +22,7 @@ Here are some general information and limitations. ## Prerequisite -Please finish tutorial [Setting up executorch](https://pytorch.org/executorch/main/getting-started-setup). +Please finish tutorial [Setting up ExecuTorch](https://pytorch.org/executorch/main/getting-started-setup). Please finish [setup QNN backend](../../docs/source/backends-qualcomm.md). @@ -44,7 +44,7 @@ commands in the `SimpleADB` class inside [utils.py](utils.py). ## Please use python xxx.py --help for information of each examples. -Some CLI examples here. Please adjust according to your environment. If you want to export the model without running it, please add `-compile_only` to the command.: +Some CLI examples here. Please adjust according to your environment. If you want to export the model without running it, please add `--compile_only` to the command.: #### First switch to following folder ```bash @@ -83,6 +83,32 @@ If you run into the following error, that means the ${QNN_SDK_ROOT} that you are ``` Error: Failed to get context binary info. ``` +## Model Structure +This section outlines the essential APIs and utilities provided to streamline the process of model conversion, deployment, and evaluation on Qualcomm hardware using ExecuTorch. + +1. `build_executorch_binary()`: + + build_executorch_binary is a high-level API used to convert a PyTorch model into a Qualcomm-compatible .pte binary format. This function streamlines the process of quantization, transformation, optimization, and export, enabling users to efficiently deploy models on Qualcomm hardware. + +2. `SimpleADB`: + + SimpleADB is a Python class that provides a simplified interface for interacting with Android devices. It allows users to execute ADB commands, retrieve device information, and manage files on the device. + +3. `get_imagenet_dataset`: + + If the model requires ImageNet, this function can be used to load the dataset and apply the necessary preprocessing steps to prepare it for inference or quantization calibration. + +4. `topk_accuracy`: + + Calculates the Top-K accuracy for classification models, used to evaluate model performance. + +5. `parse_skip_delegation_node`: + + Parses command-line arguments to identify node IDs or operation types that should be skipped during model conversion. + +6. `make_output_dir`: + + Creates a clean directory for storing model outputs or intermediate results. If the directory already exists, it will be deleted and recreated to ensure a consistent environment for each run. ## Additional Dependency @@ -95,7 +121,7 @@ pip install scikit-learn pandas ## Limitation -1. QNN 2.24 is used for all examples. Newer or older QNN might work, +1. QNN 2.28 is used for all examples. Newer or older QNN might work, but the performance and accuracy number can differ. 2. The mobilebert example is on QNN HTP fp16, which is only supported by a limited diff --git a/examples/qualcomm/TARGETS b/examples/qualcomm/TARGETS index 43ca4db6be5..600bc330457 100644 --- a/examples/qualcomm/TARGETS +++ b/examples/qualcomm/TARGETS @@ -1,7 +1,7 @@ # Any targets that should be shared between fbcode and xplat must be defined in # targets.bzl. This file can contain fbcode-only targets. -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") @@ -33,7 +33,7 @@ runtime.command_alias( exe = ":export_example", ) -python_library( +runtime.python_library( name = "utils", srcs = ["utils.py"], deps = [ diff --git a/examples/qualcomm/custom_op/custom_ops_1.py b/examples/qualcomm/custom_op/custom_ops_1.py index 4a865197584..28aabfbb342 100644 --- a/examples/qualcomm/custom_op/custom_ops_1.py +++ b/examples/qualcomm/custom_op/custom_ops_1.py @@ -102,15 +102,6 @@ def annotate_custom(gm: torch.fx.GraphModule) -> None: ) -def create_device_inputs(example_inputs): - input_list = "" - for idx, _ in enumerate(example_inputs): - input_name = f"input_0_{idx}.raw" - input_list += input_name + " " - input_list = input_list.strip() + "\n" - return input_list - - def _run(cmd, cwd=None): subprocess.run(cmd, stdout=sys.stdout, cwd=cwd, check=True) @@ -189,12 +180,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - quant_dtype = QuantDtype.use_8a8w if args.use_fp16: quant_dtype = None @@ -204,7 +189,6 @@ def main(args): sample_input = (torch.ones(1, 32, 28, 28),) workspace = f"/data/local/tmp/executorch/{pte_filename}" - input_list = create_device_inputs(sample_input) soc_info = _soc_info_table[getattr(QcomChipset, args.model)] op_package_options, op_package_paths = prepare_op_package( @@ -237,8 +221,7 @@ def main(args): if args.enable_x86_64: input_list_filename = "input_list.txt" - input_list = f"{args.artifact}/{input_list}" - generate_inputs(args.artifact, input_list_filename, sample_input, input_list) + generate_inputs(args.artifact, input_list_filename, sample_input) qnn_sdk = os.getenv("QNN_SDK_ROOT") assert qnn_sdk, "QNN_SDK_ROOT was not found in environment variable" target = "x86_64-linux-clang" @@ -276,7 +259,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=sample_input, input_list=input_list, files=op_package_paths) + adb.push(inputs=sample_input, files=op_package_paths) adb.execute() adb.pull(output_path=args.artifact) @@ -342,6 +325,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) diff --git a/examples/qualcomm/executor_runner/qnn_executor_runner.cpp b/examples/qualcomm/executor_runner/qnn_executor_runner.cpp index 83478bd8e68..26e70c90f38 100644 --- a/examples/qualcomm/executor_runner/qnn_executor_runner.cpp +++ b/examples/qualcomm/executor_runner/qnn_executor_runner.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include #include #include @@ -33,7 +35,6 @@ #include #include #include - static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4 MB DEFINE_string( @@ -83,12 +84,38 @@ DEFINE_int32( 20000000, // 20MB "Size of the debug buffer in bytes to allocate for intermediate outputs and program outputs logging."); +DEFINE_string( + performance_output_path, + "inference_speed.txt", + "Records inference speed. For CI purpose."); + +DEFINE_int32( + log_level, + 0, + "Log level between 1-5, higher is more verbose. " + "This is a runtime option and will override the log level set during AOT. " + "Refer to QnnExecuTorchLogLevel under qc_compiler_spec.fbs for more info."); +DEFINE_int32( + htp_performance_mode, + 0, + "HTP Performance mode between 0-8. " + "This is a runtime option and will override the performance mode set during AOT. " + "Refer to QnnExecuTorchHtpPerformanceMode under qc_compiler_spec.fbs for more info."); +DEFINE_int32( + profile_level, + 0, + "Profile level between 0-2. " + "Level 3(Optrace) must be turned on during AOT and cannot be enabled during runtime. " + "This is a runtime option and will override the profile level set during AOT. " + "Refer to QnnExecuTorchProfileLevel under qc_compiler_spec.fbs for more info."); + using executorch::aten::Tensor; using executorch::aten::TensorImpl; using executorch::etdump::ETDumpGen; using executorch::etdump::ETDumpResult; using executorch::extension::FileDataLoader; using executorch::extension::prepare_input_tensors; +using executorch::runtime::BackendOption; using executorch::runtime::Error; using executorch::runtime::EValue; using executorch::runtime::EventTracerDebugLogLevel; @@ -151,6 +178,40 @@ int main(int argc, char** argv) { return 1; } + // Set runtime options + executorch::runtime::BackendOptions<3> backend_options; + if (!gflags::GetCommandLineFlagInfoOrDie("log_level").is_default) { + ET_LOG(Info, "Setting runtime log level: %d", FLAGS_log_level); + ET_CHECK_MSG( + backend_options.set_option(QNN_RUNTIME_LOG_LEVEL, FLAGS_log_level) == + Error::Ok, + "Failed to set backend options: %s", + QNN_RUNTIME_LOG_LEVEL); + } + if (!gflags::GetCommandLineFlagInfoOrDie("htp_performance_mode").is_default) { + ET_LOG( + Info, + "Setting runtime performance mode: %d", + FLAGS_htp_performance_mode); + ET_CHECK_MSG( + backend_options.set_option( + QNN_RUNTIME_HTP_PERFORMANCE_MODE, FLAGS_htp_performance_mode) == + Error::Ok, + "Failed to set backend options: %s", + QNN_RUNTIME_HTP_PERFORMANCE_MODE); + } + if (!gflags::GetCommandLineFlagInfoOrDie("profile_level").is_default) { + ET_LOG(Info, "Setting runtime profile level: %d", FLAGS_profile_level); + ET_CHECK_MSG( + backend_options.set_option( + QNN_RUNTIME_PROFILE_LEVEL, FLAGS_profile_level) == Error::Ok, + "Failed to set backend options: %s", + QNN_RUNTIME_PROFILE_LEVEL); + } + ET_CHECK_MSG( + set_option(QNN_BACKEND, backend_options.view()) == Error::Ok, + "Failed to set runtime options."); + // Create a loader to get the data of the program file. There are other // 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. @@ -483,10 +544,20 @@ int main(int argc, char** argv) { } ET_LOG( Info, - "%d inference took %f ms, avg %f ms", + "Total %d inference took %f ms, avg %f ms", inference_index, elapsed_time, elapsed_time / inference_index); + + // Save avg inference time for CI + std::ofstream outfile(FLAGS_performance_output_path.c_str()); + if (outfile.is_open()) { + double avg_time = elapsed_time / inference_index; + outfile << avg_time; + outfile.close(); + } else { + ET_CHECK_MSG(false, "Error saving the inference speed file"); + } } else { // if no input is provided, fill the inputs with default values auto inputs = prepare_input_tensors(*method); diff --git a/examples/qualcomm/oss_scripts/README.md b/examples/qualcomm/oss_scripts/README.md new file mode 100644 index 00000000000..b68024d5fbf --- /dev/null +++ b/examples/qualcomm/oss_scripts/README.md @@ -0,0 +1,141 @@ +# Usage Guide for Models Provided by ExecuTorch + +This guide provides examples and instructions for open source models. Some models under this folder might also have their own customized runner. + +## Model categories +The following models can be categorized based on their primary use cases. + +1. Language Model: + - albert + - bert + - distilbert + - eurobert + - llama + - roberta + +2. Vision Model: + - conv_former + - cvt + - deit + - dino_v2 + - dit + - efficientnet + - efficientSAM + - esrgan + - fastvit + - fbnet + - focalnet + - gMLP_image_classification + - mobilevit1 + - mobilevit_v2 + - pvt + - regnet + - retinanet + - squeezenet + - ssd300_vgg16 + - swin_transformer + +## Prerequisite +Please follow another [README](../README.md) first to set up environment. + +## Model running +Some models require specific datasets. Please download them in advance and place them in the appropriate folders. + +Detailed instructions for each model are provided below. +If you want to export the model without running it, please add `--compile_only` to the command. + +1. `albert`,`bert`,`distilbert`, `eurobert`, `roberta`: + - Required Dataset : wikisent2 + + download [dataset](https://www.kaggle.com/datasets/mikeortman/wikipedia-sentences) first, and place it in a valid folder. + ```bash + python albert.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d path/to/wikisent2 + +2. `conv_former`,`cvt`,`deit`,`dino_v2`,`efficientnet`,`fbnet`, `focalnet`, `gMLP_image_classification`, `mobilevit1`,`mobilevit_v2`, `pvt`, `squeezenet`, `swin_transformer` : + - Required Dataset : ImageNet + + Download [dataset](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000) first, and place it in a valid folder. + ```bash + python SCRIPT_NAME.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d path/to/ImageNet + +3. `dit`: + + ```bash + python dit.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} +4. `esrgan`: + - Required Dataset: B100 + + Will be downloaded automatically if -d is specified. Alternatively, you can provide your own dataset using `--hr_ref_dir` and `--lr_dir`. + + - Required OSS Repo: Real-ESRGAN + + Clone [OSS Repo](https://github.com/ai-forever/Real-ESRGAN) first, and place it in a valid folder. + + ```bash + python esrgan.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} --oss_repo path/to/Real-ESRGAN + +5. `fastvit`: + - Required Dataset: ImageNet + + Download [dataset](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000) first, and place it in a valid folder. + + - Required OSS Repo: ml-fastvit + + Clone [OSS Repo](https://github.com/apple/ml-fastvit) first, and place it in a valid folder. + + - Pretrained weight: + + Download [pretrained weight](https://docs-assets.developer.apple.com/ml-research/models/fastvit/image_classification_distilled_models/fastvit_s12_reparam.pth.tar) first, and place it in a valid folder(should be fastvit_s12_reparam.pth.tar). + ```bash + python fastvit.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} --oss_repo path/to/ml-fastvit -p path/to/pretrained_weight -d path/to/ImageNet + +6. `regnet`: + - Required Dataset: ImageNet + + Download [dataset](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000) first, and place it in a valid folder. + - Weights: regnet_y_400mf, regnet_x_400mf + + use `--weights` to specify which regent weights/model to execute. + ```bash + python regnet.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d path/to/ImageNet --weights + +7. `retinanet`: + - Required Dataset: COCO + + Download [val2017](http://images.cocodataset.org/zips/val2017.zip) and [annotations](http://images.cocodataset.org/annotations/annotations_trainval2017.zip) first, and place it in a valid folder. + + ```bash + python retinanet.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d path/to/PATH/TO/COCO #(which contains 'val_2017' & 'annotations') + +8. `ssd300_vgg16`: + - Required OSS Repo: + + Clone [OSS Repo](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Object-Detection) first, and place it in a valid folder. + + - Pretrained weight: + + Download [pretrained weight](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Object-Detection) first, and place it in a valid folder.(checkpoint_ssd300.pth.tar) + + - Required Dataset: VOCSegmentation + download [VOC 2007](https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Object-Detection?tab=readme-ov-file#download) first, and place it in a valid folder. + ```bash + python ssd300_vgg16.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} --oss_repo path/to/a-PyTorch-Tutorial-to-Object-Detection -p path/to/pretrained_weight + +9. `llama`: + For llama, please check [README](llama/README.md) under llama folder for more details. + +10. `efficientSAM`: + For efficientSAM, please get access to efficientSAM folder. + - Pretrained weight: + + Download [EfficientSAM-S](https://github.com/yformer/EfficientSAM/blob/main/weights/efficient_sam_vits.pt.zip) or [EfficientSAM-Ti](https://github.com/yformer/EfficientSAM/blob/main/weights/efficient_sam_vitt.pt) first, and place it in a valid folder. + - Required Dataset: ImageNet + + Download [dataset](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000) first, and place it in a valid folder. + + - Required OSS Repo: + + Clone [OSS Repo](https://github.com/yformer/EfficientSAM) first, and place it in a valid folder. + ```bash + python efficientSAM.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} --oss_repo path/to/EfficientSAM -p path/to/pretrained_weight -d path/to/ImageNet + \ No newline at end of file diff --git a/examples/qualcomm/oss_scripts/albert.py b/examples/qualcomm/oss_scripts/albert.py index 6af554655f1..1935798d536 100644 --- a/examples/qualcomm/oss_scripts/albert.py +++ b/examples/qualcomm/oss_scripts/albert.py @@ -51,7 +51,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_masked_language_model_dataset( + inputs, targets = get_masked_language_model_dataset( args.dataset, tokenizer, data_size ) @@ -94,7 +94,7 @@ def main(args): make_output_dir(output_data_folder) # accuracy analysis - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() adb.pull(output_path=args.artifact) # since the original nn.Module could not perform well on this task either @@ -152,6 +152,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/bert.py b/examples/qualcomm/oss_scripts/bert.py index 96c7826d89c..31171c3b689 100644 --- a/examples/qualcomm/oss_scripts/bert.py +++ b/examples/qualcomm/oss_scripts/bert.py @@ -13,7 +13,6 @@ import evaluate import numpy as np import torch - from executorch.backends.qualcomm._passes.qnn_pass_manager import ( get_capture_program_passes, ) @@ -50,7 +49,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_masked_language_model_dataset( + inputs, targets = get_masked_language_model_dataset( args.dataset, tokenizer, data_size ) module = AutoModelForMaskedLM.from_pretrained( @@ -92,7 +91,7 @@ def main(args): make_output_dir(output_data_folder) # accuracy analysis - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() adb.pull(output_path=args.artifact) goldens, predictions = [], [] @@ -139,6 +138,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/conv_former.py b/examples/qualcomm/oss_scripts/conv_former.py index 8ce16abcc87..70304568a50 100644 --- a/examples/qualcomm/oss_scripts/conv_former.py +++ b/examples/qualcomm/oss_scripts/conv_former.py @@ -5,7 +5,9 @@ # LICENSE file in the root directory of this source tree. import json +import logging import os + import sys from multiprocessing.connection import Client @@ -37,17 +39,14 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - if args.compile_only: + if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] + logging.warning( + "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." + ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -85,7 +84,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -132,10 +131,12 @@ def main(args): "for https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000)" ), type=str, - required=True, + required=False, ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/cvt.py b/examples/qualcomm/oss_scripts/cvt.py index eefbb6f2259..d7418c549a9 100644 --- a/examples/qualcomm/oss_scripts/cvt.py +++ b/examples/qualcomm/oss_scripts/cvt.py @@ -7,6 +7,7 @@ import json import logging import os + import types from multiprocessing.connection import Client @@ -93,12 +94,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -106,7 +101,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -146,7 +141,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -198,6 +193,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/deit.py b/examples/qualcomm/oss_scripts/deit.py index e0719dfffb9..cffbba3ff30 100644 --- a/examples/qualcomm/oss_scripts/deit.py +++ b/examples/qualcomm/oss_scripts/deit.py @@ -8,6 +8,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -55,7 +56,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(height, width), @@ -96,7 +97,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -147,6 +148,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/dino_v2.py b/examples/qualcomm/oss_scripts/dino_v2.py index db0981248e9..3263a9f13aa 100644 --- a/examples/qualcomm/oss_scripts/dino_v2.py +++ b/examples/qualcomm/oss_scripts/dino_v2.py @@ -6,6 +6,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -42,14 +43,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - img_size, data_num = 224, 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -85,7 +80,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -137,6 +132,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/distilbert.py b/examples/qualcomm/oss_scripts/distilbert.py index 2863a653200..6ddaf8b18d4 100644 --- a/examples/qualcomm/oss_scripts/distilbert.py +++ b/examples/qualcomm/oss_scripts/distilbert.py @@ -50,7 +50,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_masked_language_model_dataset( + inputs, targets = get_masked_language_model_dataset( args.dataset, tokenizer, data_size ) module = AutoModelForMaskedLM.from_pretrained( @@ -92,7 +92,7 @@ def main(args): make_output_dir(output_data_folder) # accuracy analysis - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() adb.pull(output_path=args.artifact) goldens, predictions = [], [] @@ -139,6 +139,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/dit.py b/examples/qualcomm/oss_scripts/dit.py index 1dc4cebee75..68d3500487a 100644 --- a/examples/qualcomm/oss_scripts/dit.py +++ b/examples/qualcomm/oss_scripts/dit.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -37,7 +38,7 @@ def get_rvlcdip_dataset(data_size): ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] for index, data in enumerate(dataset): if index >= data_size: break @@ -47,9 +48,8 @@ def get_rvlcdip_dataset(data_size): ) inputs.append((feature["pixel_values"],)) targets.append(torch.tensor(target)) - input_list += f"input_{index}_0.raw\n" - return inputs, targets, input_list + return inputs, targets def main(args): @@ -58,11 +58,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) data_num = 160 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -70,7 +65,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_rvlcdip_dataset(data_num) + inputs, targets = get_rvlcdip_dataset(data_num) module = ( AutoModelForImageClassification.from_pretrained( @@ -112,7 +107,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -151,6 +146,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py b/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py index 8b7c1dc3dd3..5e5f3b0d235 100644 --- a/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py +++ b/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py @@ -97,19 +97,13 @@ def get_dataset(dataset_path, data_size=1): dataloader = DataLoader(dataset) # prepare input data - inputs, input_list = [], "" + inputs = [] for index, data in enumerate(dataloader): if index >= data_size: break inputs.append(tuple(data)) - num_feature = len(data) - for idx, _ in enumerate(data): - input_name = f"input_{index}_{idx}.raw" - input_list += input_name + " " if idx < num_feature - 1 else input_name - input_list = input_list + "\n" - - return inputs, input_list + return inputs def source_transform( @@ -226,7 +220,7 @@ def main(args): os.makedirs(args.artifact, exist_ok=True) data_size = 1 - inputs, input_list = get_dataset(args.dataset, data_size) + inputs = get_dataset(args.dataset, data_size) assert args.pretrained_weight, "Checkpoint params can't be empty" # Get the EfficientSAM model. @@ -271,7 +265,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -317,6 +311,7 @@ def post_process(): ) parser.add_argument( + "-p", "--pretrained_weight", help="Path to ESAM checkpoint, such as ./efficient_sam_vitt.pt or ./efficient_sam_vits.pt.zip", type=str, @@ -343,6 +338,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/efficientnet.py b/examples/qualcomm/oss_scripts/efficientnet.py index b11ad7abc47..3430977022e 100644 --- a/examples/qualcomm/oss_scripts/efficientnet.py +++ b/examples/qualcomm/oss_scripts/efficientnet.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -31,12 +32,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +39,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -82,7 +77,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -135,6 +130,8 @@ def main(args): ) args = parser.parse_args() + args.validate(args) + try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/esrgan.py b/examples/qualcomm/oss_scripts/esrgan.py index a5f027f79a6..82da1064d6f 100644 --- a/examples/qualcomm/oss_scripts/esrgan.py +++ b/examples/qualcomm/oss_scripts/esrgan.py @@ -6,6 +6,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -45,17 +46,11 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - dataset = get_dataset( args.hr_ref_dir, args.lr_dir, args.default_dataset, args.artifact ) - inputs, targets, input_list = dataset.lr, dataset.hr, dataset.get_input_list() + inputs, targets = dataset.lr, dataset.hr pte_filename = "esrgan_qnn" instance = get_instance(args.oss_repo) @@ -83,7 +78,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -170,6 +165,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/eurobert.py b/examples/qualcomm/oss_scripts/eurobert.py index 97e70428e01..46dea4584e7 100644 --- a/examples/qualcomm/oss_scripts/eurobert.py +++ b/examples/qualcomm/oss_scripts/eurobert.py @@ -43,12 +43,6 @@ def main(args): os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - module_id = "EuroBERT/EuroBERT-210m" tokenizer = AutoTokenizer.from_pretrained(module_id) model = AutoModelForMaskedLM.from_pretrained( @@ -88,7 +82,7 @@ def replace_rms_norm_with_native_rms_norm(module: torch.nn.Module): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_masked_language_model_dataset( + inputs, targets = get_masked_language_model_dataset( args.dataset, tokenizer, data_size ) @@ -130,7 +124,7 @@ def replace_rms_norm_with_native_rms_norm(module: torch.nn.Module): make_output_dir(output_data_folder) # accuracy analysis - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() adb.pull(output_path=args.artifact) goldens, predictions = [], [] @@ -177,6 +171,7 @@ def replace_rms_norm_with_native_rms_norm(module: torch.nn.Module): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/fastvit.py b/examples/qualcomm/oss_scripts/fastvit.py index ee062735fbd..854f3ad6e27 100644 --- a/examples/qualcomm/oss_scripts/fastvit.py +++ b/examples/qualcomm/oss_scripts/fastvit.py @@ -6,6 +6,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -65,14 +66,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -146,7 +141,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -218,6 +213,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/fbnet.py b/examples/qualcomm/oss_scripts/fbnet.py index 67fe2fba380..4ef802b6233 100755 --- a/examples/qualcomm/oss_scripts/fbnet.py +++ b/examples/qualcomm/oss_scripts/fbnet.py @@ -23,19 +23,13 @@ def main(args): - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) instance = timm.create_model("fbnetc_100", pretrained=True).eval() data_num = 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(299, 299), @@ -65,7 +59,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -129,6 +123,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/focalnet.py b/examples/qualcomm/oss_scripts/focalnet.py index 377d49a3a18..421c4537c74 100644 --- a/examples/qualcomm/oss_scripts/focalnet.py +++ b/examples/qualcomm/oss_scripts/focalnet.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -31,12 +32,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +39,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -82,7 +77,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -135,6 +130,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/gMLP_image_classification.py b/examples/qualcomm/oss_scripts/gMLP_image_classification.py index 1dffa6831b4..e684335c623 100644 --- a/examples/qualcomm/oss_scripts/gMLP_image_classification.py +++ b/examples/qualcomm/oss_scripts/gMLP_image_classification.py @@ -7,6 +7,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -31,14 +32,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(224, 224), @@ -73,7 +68,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -125,6 +120,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/llama/CMakeLists.txt b/examples/qualcomm/oss_scripts/llama/CMakeLists.txt index bf83a456bca..612f898028c 100644 --- a/examples/qualcomm/oss_scripts/llama/CMakeLists.txt +++ b/examples/qualcomm/oss_scripts/llama/CMakeLists.txt @@ -28,6 +28,7 @@ list( ${CMAKE_CURRENT_LIST_DIR}/qnn_llama_runner.cpp ${CMAKE_CURRENT_LIST_DIR}/runner/runner.cpp ${CMAKE_CURRENT_LIST_DIR}/runner/runner.h + ${CMAKE_CURRENT_LIST_DIR}/runner/cache_utils.h ${CMAKE_CURRENT_LIST_DIR}/runner/decoder_runner.cpp ${CMAKE_CURRENT_LIST_DIR}/runner/decoder_runner.h ${CMAKE_CURRENT_LIST_DIR}/runner/prompt_processor.cpp @@ -42,6 +43,8 @@ list( ${CMAKE_CURRENT_LIST_DIR}/runner/rpc_mem.h ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.cpp ${CMAKE_CURRENT_LIST_DIR}/runner/kv_manager.h + ${EXECUTORCH_SOURCE_DIR}/examples/models/llama/runner/runner.cpp + ${EXECUTORCH_SOURCE_DIR}/examples/models/llama/runner/runner.h ) list(APPEND _llama_runner__srcs) diff --git a/examples/qualcomm/oss_scripts/llama/README.md b/examples/qualcomm/oss_scripts/llama/README.md index cbfd1b46a06..1be94ec04d6 100644 --- a/examples/qualcomm/oss_scripts/llama/README.md +++ b/examples/qualcomm/oss_scripts/llama/README.md @@ -5,7 +5,13 @@ This file provides you the instructions to run LLM Decoder model with different 1. LLAMA2 Stories 110M 2. LLAMA3.2 1B 3. LLAMA3.2 3B - 4. QWEN2.5 0.5B + 4. Gemma3 1B + 5. Phi4-mini-instruct + 6. QWEN2.5 0.5B / 1.5B + 7. QWEN3 0.6B / 1.7B + 8. SmolLM2 135M + 9. SmolLM3 3B + We offer the following modes to execute the model: @@ -54,18 +60,73 @@ Follow the [instructions](https://www.llama.com/) to download models. At the end of this step, users should have the following files ready: `consolidated.00.pth`, `params.json`, and `tokenizer.model`. -### Step3: Run default examples using hybrid mode. +### Step3: Run default examples using hybrid mode for smaller models and kv mode for larger models. #### LLAMA2 ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint stories110M.pt --params params.json --tokenizer_model tokenizer.model --tokenizer_bin tokenizer.bin --llama_model stories110m --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "Once upon a time" +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint stories110M.pt --params params.json --tokenizer_model tokenizer.model --tokenizer_bin tokenizer.bin --decoder_model stories110m --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "Once upon a time" ``` -#### LLAMA3.2 +#### LLAMA3.2 1B Instruct +Default example using kv mode. +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2-1b_instruct --model_mode hybrid --prefill_ar_len 128 --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### LLAMA3.2 3B Instruct +Default example using kv mode. +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2-3b_instruct --model_mode kv --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### Gemma3 1B +Default example using hybrid mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma3-1b --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### Phi4-mini-instruct +Default example using kv mode. +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --decoder_model phi_4_mini --model_mode kv --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### QWEN2.5 0.5B +Default example using hybrid mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model qwen2_5-0_5b --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### QWEN2.5 1.5B +Default example using kv mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen2_5-1_5b --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### QWEN3 0.6B +Default example using hybrid mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model qwen3-0_6b --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### QWEN3 1.7B +Default example using hybrid mode +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen3-1_7b --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### SmolLM2 Default example using hybrid mode. ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --decoder_model smollm2_135m --model_mode hybrid --prefill_ar_len 128 --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 +``` + +#### SmolLM3 +Default example using kv mode. +```bash +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --decoder_model smollm3-3b --model_mode kv --max_seq_len 1024 --prompt "I would like to learn python, could you teach me with a simple example?" --tasks wikitext --limit 1 ``` + ### KV Cache update mechanism We have two distinct mechanisms for updating the key-value (KV) cache, which can be selected at runtime. Shift Pointer and Smart Mask. @@ -118,13 +179,13 @@ We have two distinct mechanisms for updating the key-value (KV) cache, which can #### Compile Only If you would like to compile the model only, we have provided the flag `--compile_only`. Taking LLAMA3.2 as an example: ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --compile_only +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --compile_only ``` #### Pre Generated PTE On the other hand, if you already have a pre-compiled .pte model, you can perform inference by providing the flag `--pre_gen_pte` and specifying the folder that contains the .pte model. Taking LLAMA3.2 as an example: ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --pre_gen_pte ${FOLDER_TO_PRE_GEN_PTE} +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --pre_gen_pte ${FOLDER_TO_PRE_GEN_PTE} ``` #### KV Cache Updater @@ -132,7 +193,7 @@ python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL You can select the KV Cache update mechanism at runtime by setting the `KV_UPDATER` variable to either "shift_pointer" or "smart_mask". By default, it is set to "smart_mask". `KV_UPDATER` = "shift_pointer" ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --kv_updator ${KV_UPDATER} +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2 --model_mode hybrid --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --kv_updator ${KV_UPDATER} ``` #### Lookahead Decoding Mode @@ -145,14 +206,9 @@ You can choose the lookahead mode to enhance decoding speed. To use this mode, y For more details, please refer to the paper ["Break the Sequential Dependency of LLM Inference Using Lookahead Decoding"](https://arxiv.org/abs/2402.02057) ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode lookahead --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --ngram 3 --window 2 --gcap 2 +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --decoder_model llama3_2 --model_mode lookahead --prefill_ar_len 32 --max_seq_len 128 --prompt "what is 1+1" --ngram 3 --window 2 --gcap 2 ``` -#### Masked Softmax - -You can enable MaskedSoftmax feature by providing the flag `--enable_masked_softmax`. It is designed to optimize the LLMs accuracy and performance executed on HTP backend. MaskedSoftmax is used to replace the Softmax(Add(In, Mask)) structure in attention block in LLMs during backend optimization. For more details, please refer to QNN documents. -Note that it is only supported starting from QNN 2.35. - #### Perplexity Evaluation This script supports perplexity evaluation and is capable of assessing perplexity scores across 3 phases: prepare_pt2e(CPU FP), convert_pt2e(CPU QDQ), QNN on device. @@ -160,18 +216,18 @@ To evaluate the perplexity across all 3 phases, users should provide the `--eval For example, using the Qwen model and 1 wikitext sample as the evaluation task, users can assess all 3 phases perplexity score in a single run by including the appropriate configuration: ```bash -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --ptq 16a8w --decoder_model qwen2_5 --eval_perplexity --tasks wikitext --limit 1 +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen2_5-0_5b --eval_perplexity --tasks wikitext --limit 1 ``` For the example script above, 1 wikitext sample is used to evaluate all 3 phases. However, there are cases where a user may want to use one sample for quantization calibration and multiple samples for perplexity evaluation. In this case, the process should be split into two runs. In the 1st run, the model is compiled using one sample. In the 2nd run, the user can provide a different configuration for QNN device execution. Example: ```bash # 1st run to compile with --limit 1 -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --ptq 16a8w --decoder_model qwen2_5 --eval_perplexity --tasks wikitext --limit 1 --compile_only +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen2_5-0_5b --eval_perplexity --tasks wikitext --limit 1 --compile_only ``` ```bash # 2nd run to perform QNN device execution with --limit 3 -python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --ptq 16a8w --decoder_model qwen2_5 --eval_perplexity --tasks wikitext --limit 3 --pre_gen_pte ${PATH_TO_ARTIFACT_IN_1ST_RUN} --quant_attrs_path ${PATH_TO_ARTIFACT_IN_1ST_RUN}/kv_llama_qnn_quant_attrs.json +python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --prompt "What is 1+1?" --temperature 0 --model_mode kv --max_seq_len 1024 --decoder_model qwen2_5-0_5b --eval_perplexity --tasks wikitext --limit 3 --pre_gen_pte ${PATH_TO_ARTIFACT_IN_1ST_RUN} --quant_attrs_path ${PATH_TO_ARTIFACT_IN_1ST_RUN}/kv_llama_qnn_quant_attrs.json ``` #### Tasks quantization calibration diff --git a/examples/qualcomm/oss_scripts/llama/TARGETS b/examples/qualcomm/oss_scripts/llama/TARGETS index 09a2948f3a0..10462595c56 100644 --- a/examples/qualcomm/oss_scripts/llama/TARGETS +++ b/examples/qualcomm/oss_scripts/llama/TARGETS @@ -1,11 +1,11 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "static_llama", srcs = [ "model/static_llama.py", @@ -15,10 +15,30 @@ python_library( ], ) -python_library( +runtime.python_library( + name = "decoder_utils", + srcs = [ + "decoder_utils.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/llama:eval_library", + ], +) + +runtime.python_library( + name = "decoder_constants", + srcs = [ + "decoder_constants.py", + ], +) + +runtime.python_library( name = "llama_lib", - srcs = ["llama.py"], + srcs = ["__init__.py", "llama.py"], deps = [ + ":decoder_constants", + ":decoder_utils", "//executorch/examples/models/llama:source_transformation", "//caffe2:torch", "//executorch/backends/qualcomm/partition:partition", @@ -27,6 +47,7 @@ python_library( "//executorch/devtools:lib", "//executorch/examples/models:models", "//executorch/examples/models/llama:hf_download", + "//executorch/examples/qualcomm/oss_scripts/llama:range_setting_pt2e", "//executorch/examples/qualcomm/oss_scripts/llama:static_llama", "//executorch/examples/qualcomm:utils", "//executorch/extension/export_util:export_util", @@ -35,7 +56,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "range_setting_pt2e", srcs = [ "range_setting_pt2e.py", @@ -53,7 +74,6 @@ python_binary( ], deps = [ ":llama_lib", - "//executorch/examples/qualcomm/oss_scripts/llama:range_setting_pt2e", ], ) diff --git a/examples/qualcomm/oss_scripts/llama/__init__.py b/examples/qualcomm/oss_scripts/llama/__init__.py new file mode 100644 index 00000000000..99b0739f1f0 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/__init__.py @@ -0,0 +1,516 @@ +# 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 os +from abc import ABC +from dataclasses import dataclass +from enum import Enum +from functools import partial +from typing import Callable, Dict, Tuple, Type + +import torch +from executorch.backends.qualcomm.quantizer.custom_annotation import ( + annotate_down_proj, + annotate_kv_8bit, + annotate_output_16a8w, + annotate_qkv_proj_sha, + StaticLLMQuantConfig, +) +from executorch.backends.qualcomm.quantizer.qconfig import ( + get_ptq_per_channel_quant_config, +) +from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + +from executorch.examples.models.gemma3 import convert_weights as convert_gemma3_weights +from executorch.examples.models.phi_4_mini import ( + convert_weights as convert_phi_4_mini_weights, +) +from executorch.examples.models.qwen2_5 import ( + convert_weights as convert_qwen2_5_weights, +) +from executorch.examples.models.qwen3 import convert_weights as convert_qwen3_weights +from executorch.examples.models.smollm2 import ( + convert_weights as convert_smollm2_weights, +) +from executorch.examples.models.smollm3 import ( + convert_weights as convert_smollm3_weights, +) + +from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import ( + DECODER_MODEL_VERSION, +) +from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import ( + MultiScopeAwareLlamaModel, +) +from tabulate import tabulate +from torchao.quantization.pt2e import MinMaxObserver + + +BASE_DIR = os.path.dirname(__file__) + + +LLM_VARIANT_ARCHS = { + "gemma3-1b": MultiScopeAwareLlamaModel, +} +annotate_wqkv_sha = partial( + annotate_qkv_proj_sha, + qkv_tags={ + StaticLLMQuantConfig.wq_sha, + StaticLLMQuantConfig.wk_sha, + StaticLLMQuantConfig.wv_sha, + }, +) +annotate_wv_sha = partial(annotate_qkv_proj_sha, qkv_tags={StaticLLMQuantConfig.wv_sha}) + + +@dataclass(init=False, frozen=True) +class LLMModelConfig(ABC): + """ + Base class for all LLM Models, including native Llama and Hugging Face Models. + This class stores configs for each supported LLM model, including quant config + + repo_id: Hugging Face Repo ID. + params_path: Path to model's config.json. If the corresponding .json has not yet exist, please create one. + convert_weights: Used to convert Hugging Face weights parameters to Static Decoder's parameter naming. + decoder_model_version: This is to determine the chat template to use during runtime(qnn_llama_runner). + transform_weight: Set to true to change Hugging Face weight to improve the performance of RoPE in HTP backend. + instruct_model: True if the model uses chat templates. Check Hugging Face model card to ensure the model uses chat templates. + num_sharding: Specify the number of splits by inserting the fallback custom op. The graph will be split evenly by layers. + ptq: Set to true to perform PTQ quantization. Support 16a16w, 16a8w, 16a4w, 16a4w_block, 8a8w. + group_size: Group size used in block quantization for weight quantization. Will only be used when ptq = 16a4w_block + masked_softmax: The MaskedSoftmax feature is designed to optimize the LLMs accuracy and performance executed on HTP backend. + MaskedSoftmax is used to replace the Softmax(Add(In, Mask)) structure in attention block in LLMs during backend optimization. + For more details, please refer to QNN documents. Note that it is only supported starting from QNN 2.35. + seq_mse_candidates: Number of steps to sequentially search for optimum scales for quantized parameters which will minimize + the MSE of activation value between floating point golden & fake quantization. + r1: Enable SpinQuant R1 quantization optimization. + r2: Enable SpinQuant R2 quantization optimization. + r3: Enable SpinQuant R3 quantization optimization. + custom_annotation: Custom annotation to use when setting quant configs for the model. + """ + + repo_id: str + params_path: str + convert_weights: Callable + # TODO: Replace decoder_model_version with chat_template.jinja + decoder_model_version: str + transform_weight: bool + instruct_model: bool + num_sharding: int + ptq: QuantDtype + group_size: int + masked_softmax: bool + seq_mse_candidates: int + r1: bool + r2: bool + r3: bool + custom_annotation: Tuple + + def __str__(self): # noqa: C901 + """ + Visualize the current LLMModelConfig settings in a readable table format. + + This method helps users quickly inspect key configuration, + skipping internal or irrelevant attributes and formatting complex types + like functions, enums, and partials for clarity. + + Returns: + str: A table showing the current config for LLM models. + """ + + def format_value(v): + if isinstance(v, partial): + func_name = ( + v.func.__name__ if hasattr(v.func, "__name__") else str(v.func) + ) + return f"partial({func_name})" + elif isinstance(v, Callable): + return v.__name__ if hasattr(v, "__name__") else str(v) + elif isinstance(v, Enum): + return f"{v.__class__.__name__}.{v.name}" + elif isinstance(v, (tuple, list)): + return "(" + ", ".join(format_value(i) for i in v) + ")" + elif isinstance(v, (str, int, float, bool)): + return v + else: + return f"<{v.__class__.__name__}>" + + attrs = {} + for k in dir(self): + if k.startswith("_") or k in {"convert_weights", "params_path"}: + continue + try: + v = getattr(self, k) + if k in {"get_kv_io_bit_width", "get_logits_output_bit_width"}: + v = v() + except Exception: + v = f"Warning: failed to retrieve config for '{k}'" + if isinstance(v, (str, int, float, bool, tuple, list, Callable)): + attrs[k] = format_value(v) + table = [(k, v) for k, v in attrs.items()] + return tabulate(table, headers=["Config", "Value"], tablefmt="grid") + + def get_kv_io_bit_width(self) -> int: + if self.ptq is None: + return 32 + elif ( + self.ptq == QuantDtype.use_8a8w + or annotate_kv_8bit in self.custom_annotation + ): + return 8 + else: + # If quantized but not 8a8w or mix_quantization, it has to be 16bit kv io. + return 16 + + def get_logits_output_bit_width(self) -> int: + # We use 16bit logits for all quant config + return 32 if self.ptq is None else 16 + + +SUPPORTED_LLM_MODELS: Dict[str, LLMModelConfig] = {} + + +def register_llm_model(name: str): + def decorator(cls: Type[LLMModelConfig]): + cls.decoder_model_version = DECODER_MODEL_VERSION[name] + SUPPORTED_LLM_MODELS[name.lower()] = cls() + return cls() + + return decorator + + +@register_llm_model("stories260k") +@dataclass(init=False, frozen=True) +class LlamaStories260K(LLMModelConfig): + repo_id = None + params_path = None + convert_weights = None + transform_weight = True + instruct_model = False + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w + group_size = None + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + quantization_config_wv_sha_8a4w = get_ptq_per_channel_quant_config( + act_dtype=torch.uint8, + weight_dtype=torch.int4, + act_observer=MinMaxObserver, + act_symmetric=True, + ) + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + partial(annotate_wv_sha, quantization_config=quantization_config_wv_sha_8a4w), + ) + + +@register_llm_model("stories110m") +@dataclass(init=False, frozen=True) +class LlamaStories110M(LLMModelConfig): + repo_id = None + params_path = None + convert_weights = None + transform_weight = True + instruct_model = False + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w + group_size = None + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + quantization_config_wv_sha_8a4w = get_ptq_per_channel_quant_config( + act_dtype=torch.uint8, + weight_dtype=torch.int4, + act_observer=MinMaxObserver, + act_symmetric=True, + ) + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + partial(annotate_wv_sha, quantization_config=quantization_config_wv_sha_8a4w), + ) + + +@register_llm_model("llama3_2-1b_instruct") +@dataclass(init=False, frozen=True) +class Llama3_2_1B_Instruct(LLMModelConfig): + repo_id = None + params_path = None + convert_weights = None + transform_weight = True + # The Llama3_2 enabled should be instruct, however, Llama's tokenizer does not provide utility to apply chat template. + instruct_model = False + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 32 + masked_softmax = False + seq_mse_candidates = 1000 + r1 = False + r2 = False + r3 = False + quantization_config_down_proj_16a8w = get_ptq_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + partial( + annotate_down_proj, quantization_config=quantization_config_down_proj_16a8w + ), + ) + + +@register_llm_model("llama3_2-3b_instruct") +@dataclass(init=False, frozen=True) +class Llama3_2_3B_Instruct(LLMModelConfig): + repo_id = None + params_path = None + convert_weights = None + transform_weight = True + # The Llama3_2 enabled should be instruct, however, Llama's tokenizer does not provide utility to apply chat template. + instruct_model = False + + num_sharding = 4 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 32 + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + ) + + +@register_llm_model("gemma3-1b") +@dataclass(init=False, frozen=True) +class Gemma3(LLMModelConfig): + repo_id: str = "google/gemma-3-1b-it" + params_path: str = os.path.join( + BASE_DIR, "../../../models/gemma3/config/1b_config.json" + ) + convert_weights = convert_gemma3_weights + transform_weight = False + instruct_model = True + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 64 + masked_softmax = True + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + quantization_config_wv_sha_16a8w = get_ptq_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) + custom_annotation = ( + annotate_kv_8bit, + partial(annotate_wv_sha, quantization_config=quantization_config_wv_sha_16a8w), + ) + + +@register_llm_model("phi_4_mini") +@dataclass(init=False, frozen=True) +class Phi4Mini(LLMModelConfig): + repo_id: str = "microsoft/Phi-4-mini-instruct" + params_path: str = os.path.join( + BASE_DIR, "../../../models/phi_4_mini/config/config.json" + ) + convert_weights = convert_phi_4_mini_weights + transform_weight = False + instruct_model = True + + num_sharding = 8 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 16 + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + quantization_config_wv_sha_8a4w = get_ptq_per_channel_quant_config( + act_dtype=torch.uint8, + weight_dtype=torch.int4, + act_observer=MinMaxObserver, + act_symmetric=True, + ) + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + partial(annotate_wv_sha, quantization_config=quantization_config_wv_sha_8a4w), + ) + + +@register_llm_model("qwen2_5-0_5b") +@dataclass(init=False, frozen=True) +class Qwen2_5_0_5B(LLMModelConfig): + repo_id: str = "Qwen/Qwen2.5-0.5B" + params_path: str = os.path.join( + BASE_DIR, "../../../models/qwen2_5/config/0_5b_config.json" + ) + convert_weights = convert_qwen2_5_weights + transform_weight = False + instruct_model = False + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 16 + masked_softmax = True + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = True + custom_annotation = () + + +@register_llm_model("qwen2_5-1_5b") +@dataclass(init=False, frozen=True) +class Qwen2_5_1_5B(LLMModelConfig): + repo_id: str = "Qwen/Qwen2.5-1.5B" + params_path: str = os.path.join( + BASE_DIR, "../../../models/qwen2_5/config/1_5b_config.json" + ) + convert_weights = convert_qwen2_5_weights + transform_weight = False + instruct_model = False + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 16 + masked_softmax = True + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = True + custom_annotation = (annotate_output_16a8w,) + + +@register_llm_model("qwen3-0_6b") +@dataclass(init=False, frozen=True) +class Qwen3_0_6B(LLMModelConfig): + repo_id: str = "Qwen/Qwen3-0.6B" + params_path: str = os.path.join( + BASE_DIR, "../../../models/qwen3/config/0_6b_config.json" + ) + convert_weights = convert_qwen3_weights + transform_weight = False + instruct_model = True + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 32 + masked_softmax = True + seq_mse_candidates = 1000 + r1 = False + r2 = False + r3 = False + quantization_config_down_proj_16a8w = get_ptq_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) + custom_annotation = ( + partial( + annotate_down_proj, quantization_config=quantization_config_down_proj_16a8w + ), + ) + + +@register_llm_model("qwen3-1_7b") +@dataclass(init=False, frozen=True) +class Qwen3_1_7B(LLMModelConfig): + repo_id: str = "Qwen/Qwen3-1.7B" + params_path: str = os.path.join( + BASE_DIR, "../../../models/qwen3/config/1_7b_config.json" + ) + convert_weights = convert_qwen3_weights + transform_weight = False + instruct_model = True + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 16 + masked_softmax = True + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = True + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + ) + + +@register_llm_model("smollm2_135m") +@dataclass(init=False, frozen=True) +class Smollm2_135M(LLMModelConfig): + repo_id: str = "HuggingFaceTB/SmolLM2-135M-Instruct" + params_path: str = os.path.join( + BASE_DIR, "../../../models/smollm2/135M_config.json" + ) + convert_weights = convert_smollm2_weights + transform_weight = True + instruct_model = True + + num_sharding = 1 + # quant config + ptq = QuantDtype.use_16a8w + group_size = None + masked_softmax = False + seq_mse_candidates = 0 + r1 = False + r2 = False + r3 = False + custom_annotation = () + + +@register_llm_model("smollm3-3b") +@dataclass(init=False, frozen=True) +class Smollm3_3B(LLMModelConfig): + repo_id: str = "HuggingFaceTB/SmolLM3-3B" + params_path: str = os.path.join(BASE_DIR, "../../../models/smollm3/3b_config.json") + convert_weights = convert_smollm3_weights + transform_weight = False + instruct_model = True + + num_sharding = 4 + # quant config + ptq = QuantDtype.use_16a4w_block + group_size = 32 + masked_softmax = True + r1 = False + r2 = False + r3 = False + quantization_config_wqkv_sha_16a8w = get_ptq_per_channel_quant_config( + torch.uint16, weight_dtype=torch.int8, act_observer=MinMaxObserver + ) + custom_annotation = ( + annotate_kv_8bit, + annotate_output_16a8w, + partial( + annotate_wqkv_sha, quantization_config=quantization_config_wqkv_sha_16a8w + ), + ) diff --git a/examples/qualcomm/oss_scripts/llama/decoder_constants.py b/examples/qualcomm/oss_scripts/llama/decoder_constants.py index cf5aa02a357..ac96770b889 100644 --- a/examples/qualcomm/oss_scripts/llama/decoder_constants.py +++ b/examples/qualcomm/oss_scripts/llama/decoder_constants.py @@ -4,17 +4,24 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -HUGGING_FACE_REPO_IDS = {"qwen2_5": "Qwen/Qwen2.5-0.5B"} - EVAL_MODE = { "kv": 0, "hybrid": 1, "lookahead": 2, } +# The dict's value is mainly for runner to decide what special tokens are required to wrap the prompt. DECODER_MODEL_VERSION = { "stories260k": "llama2", "stories110m": "llama2", - "llama3_2": "llama3", - "qwen2_5": "qwen2_5", + "gemma3-1b": "gemma3", + "phi_4_mini": "phi_4_mini", + "llama3_2-1b_instruct": "llama3", + "llama3_2-3b_instruct": "llama3", + "qwen2_5-0_5b": "qwen2_5", + "qwen2_5-1_5b": "qwen2_5", + "qwen3-0_6b": "qwen3", + "qwen3-1_7b": "qwen3", + "smollm2_135m": "smollm2_135m", + "smollm3-3b": "smollm3", } diff --git a/examples/qualcomm/oss_scripts/llama/decoder_utils.py b/examples/qualcomm/oss_scripts/llama/decoder_utils.py index eba8c375468..76cf85c6e9c 100644 --- a/examples/qualcomm/oss_scripts/llama/decoder_utils.py +++ b/examples/qualcomm/oss_scripts/llama/decoder_utils.py @@ -5,7 +5,6 @@ # LICENSE file in the root directory of this source tree. import getpass -import json import logging import os from typing import Callable, Optional, Union @@ -13,12 +12,13 @@ import numpy as np import torch +from executorch.backends.qualcomm._passes import SeqMSE from executorch.examples.models.llama.evaluate.eager_eval import EagerEvalWrapper - from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import ( DECODER_MODEL_VERSION, EVAL_MODE, ) +from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import AttentionMask from executorch.examples.qualcomm.utils import make_output_dir, SimpleADB from executorch.exir._serialize._program import deserialize_pte_binary @@ -34,6 +34,16 @@ ) +INFERENCE_REGISTRY = {} + + +def register_inference(use_kv_cache: bool): + def decorator(func): + INFERENCE_REGISTRY[use_kv_cache] = func + + return decorator + + class GraphModuleCalibrationWrapper(EagerEvalWrapper): """ A wrapper class for calibration @@ -45,14 +55,16 @@ def __init__( tokenizer: Union[ SentencePieceTokenizer, TiktokenTokenizer, HuggingFaceTokenizer ], - max_seq_length: Optional[int], + max_seq_length: int, ar_len: int, use_kv_cache: bool, get_example_inputs: Callable, kv_updater: Callable, use_i64_token: bool, + seq_mse_candidates: int, ): # n seq len = n-1 cache len, so we len(inps) = n-1 during _model_call + assert max_seq_length is not None, "max_seq_length must be provided" super().__init__( model=model, tokenizer=tokenizer, max_seq_length=max_seq_length - 1 ) @@ -63,32 +75,28 @@ def __init__( self.max_seq_length = max_seq_length self.kv_updater = kv_updater self.use_i64_token = use_i64_token + self.seq_mse_candidates = seq_mse_candidates def _model_call(self, inps): all_logits = None + kwargs = {} if self._use_kv_cache: - all_logits = kv_inference( - self.get_example_inputs, - inps, - self._model, - self._tokenizer, - self.ar_len, - self.max_seq_length, - kv_updater=self.kv_updater, - use_i64_token=self.use_i64_token, - collect_logits=True, - ) - else: - all_logits = prefill_inference( - self.get_example_inputs, - inps, - self._model, - self._tokenizer, - self.ar_len, - self.max_seq_length, - use_i64_token=self.use_i64_token, - collect_logits=True, - ) + kwargs["ar_len"] = self.ar_len + kwargs["kv_updater"] = self.kv_updater + kwargs["seq_mse_candidates"] = self.seq_mse_candidates + + all_logits = INFERENCE_REGISTRY[self._use_kv_cache]( + self.get_example_inputs, + inps, + self._model, + self._tokenizer, + max_seq_len=self.max_seq_length, + use_i64_token=self.use_i64_token, + collect_logits=True, + **kwargs, + ) + # one shot is enough for seq mse + self.seq_mse_candidates = 0 return all_logits @@ -117,12 +125,33 @@ def __init__( # Retrieve vocab_size from get_metadata under static_llama that is passed to edge manager self.output_vocab_size = None pte_max_seq_len = None + self.logits_scale = None + self.logits_zero_point = None + self.kv_io_bit_width = 32 for method in program.execution_plan: # Don't use tokenizer.n_words, the numbers are off once calling get_tokenizer() if method.name == "get_vocab_size": + # pyre-ignore self.output_vocab_size = method.values[0].val.int_val if method.name == "get_max_seq_len": + # pyre-ignore pte_max_seq_len = method.values[0].val.int_val + if method.name == "get_logits_scale": + self.logits_scale = method.values[0].val.double_val + if method.name == "get_logits_zero_point": + self.logits_zero_point = method.values[0].val.int_val + if method.name == "get_kv_io_bit_width": + self.kv_io_bit_width = method.values[0].val.int_val + + # FP has no scale/zero_point, use following values, which is equivalent to not performing dequantize. + if self.kv_io_bit_width == 32: + self.logits_scale = 1 + self.logits_zero_point = 0 + elif self.logits_scale is None or self.logits_zero_point is None: + raise RuntimeError( + "Unable to find scale/offset. The .pte file might be deprecated. Please generate a new .pte file" + ) + assert self.output_vocab_size is not None, "Couldn't find the vocab size" assert pte_max_seq_len is not None, "Couldn't find the max_seq_len from pte" if pte_max_seq_len != max_seq_length: @@ -135,11 +164,6 @@ def __init__( ) max_seq_length = pte_max_seq_len self.max_seq_length = max_seq_length - - assert ( - args.quant_attrs_path is not None - ), "Please provide path to quant_attrs json file" - self.quant_attrs = json.load(open(args.quant_attrs_path)) self.runtime_tokenizer_path = runtime_tokenizer_path self.output_dir = args.artifact @@ -155,8 +179,9 @@ def __init__( soc_model=args.model, runner="examples/qualcomm/oss_scripts/llama/qnn_llama_runner", ) - self.adb.push(inputs=[], input_list="", files=[self.runtime_tokenizer_path]) + self.adb.push(inputs=[], files=[self.runtime_tokenizer_path]) # n seq len = n-1 cache len, so we len(inps) = n-1 during _model_call + # pyre-ignore super().__init__(None, tokenizer, max_seq_length - 1) def _model_call(self, inps): @@ -189,7 +214,7 @@ def _model_call(self, inps): ] ) - self.adb.push(inputs=[], input_list="", files=[input_file_name], init_env=False) + self.adb.push(inputs=[], files=[input_file_name], init_env=False) self.adb.execute(custom_runner_cmd=runner_cmd) output_data_folder = f"{self.output_dir}/outputs" make_output_dir(output_data_folder) @@ -203,8 +228,8 @@ def post_process(): ) ) output_tensor = ( - output_tensor.to(torch.float32) - self.quant_attrs["zero_point"] - ) * self.quant_attrs["scale"] + output_tensor.to(torch.float32) - self.logits_zero_point + ) * self.logits_scale output_tensor_list.append(output_tensor) # simple_eval will run multiple rounds, use last run for inference speed @@ -216,40 +241,58 @@ def post_process(): def smart_mask_updater( - ar_len, atten_mask, pos, k_caches, v_caches, new_k_caches, new_v_caches + n_updates: int, + atten_mask: AttentionMask, + pos, + k_caches, + v_caches, + new_k_caches, + new_v_caches, ): - # Update the KV cache input for the next inference when the position exceeds the autoregressive length. - if pos >= ar_len: + # ar_len is unused in smart mask + max_cache_len = k_caches[0].size(-1) + if pos + n_updates <= max_cache_len: for i, k_cache in enumerate(k_caches): - k_cache[:, :, pos - ar_len] = new_k_caches[i][:, :, 0] + k_cache[:, :, pos : pos + n_updates] = new_k_caches[i][:, :, :n_updates] for i, v_cache in enumerate(v_caches): - v_cache[:, pos - ar_len, :] = new_v_caches[i][:, 0, :] - atten_mask[:, :, pos - ar_len] = 0 + v_cache[:, pos : pos + n_updates, :] = new_v_caches[i][:, :n_updates, :] + atten_mask.smart_mask_update(pos, n_updates) + pos += n_updates - pos += 1 - return (atten_mask, pos, k_caches, v_caches) + return pos, k_caches, v_caches def shift_pointer_updater( - ar_len, atten_mask, pos, k_caches, v_caches, new_k_caches, new_v_caches + n_updates: int, + atten_mask: AttentionMask, + pos, + k_caches, + v_caches, + new_k_caches, + new_v_caches, ): - # Update the KV cache input for the next inference when the position exceeds the autoregressive length. - if pos >= ar_len: + max_cache_len = k_caches[0].size(-1) + if pos + n_updates <= max_cache_len: k_caches = [ - torch.cat([k_cache[:, :, 1:], new_k_caches[i][:, :, :1]], dim=-1) + torch.cat( + [k_cache[:, :, n_updates:], new_k_caches[i][:, :, :n_updates]], dim=-1 + ) for i, k_cache in enumerate(k_caches) ] v_caches = [ - torch.cat([v_cache[:, 1:, :], new_v_caches[i][:, :1, :]], dim=1) + torch.cat( + [v_cache[:, n_updates:, :], new_v_caches[i][:, :n_updates, :]], dim=1 + ) for i, v_cache in enumerate(v_caches) ] - atten_mask[:, :, -pos - 1] = 0 + atten_mask.shift_pointer_update(pos, n_updates) + pos += n_updates - pos += 1 - return (atten_mask, pos, k_caches, v_caches) + return pos, k_caches, v_caches +@register_inference(use_kv_cache=True) def kv_inference( get_example_inputs, prompt: Union[str, list], @@ -260,80 +303,144 @@ def kv_inference( kv_updater=smart_mask_updater, use_i64_token=False, collect_logits=False, + seq_mse_candidates=0, ): _, atten_mask, _, k_caches, v_caches = get_example_inputs(use_kv_cache=True) # TODO: change criteria & support batch inputs if necessary all_pos = torch.arange(0, max_seq_len, 1, dtype=torch.int32).unsqueeze(0) - token_list, result_logits = [], [] + prompt_token_list, total_token_list, result_logits = [], [], [] if isinstance(prompt, str): # Llama2 tokenizer has no special tokens if isinstance(tokenizer, (SentencePieceTokenizer, HuggingFaceTokenizer)): - token_list = tokenizer.encode(prompt, bos=True, eos=False) + prompt_token_list = tokenizer.encode(prompt, bos=True, eos=False) elif isinstance(tokenizer, TiktokenTokenizer): - token_list = tokenizer.encode( + prompt_token_list = tokenizer.encode( prompt, bos=True, eos=False, allowed_special="all" ) else: raise RuntimeError("Unknown tokenizer") else: - token_list = prompt.flatten().tolist() - pos = len(token_list) if len(token_list) < ar_len else ar_len + # pyre-ignore + prompt_token_list = prompt.flatten().tolist() + total_token_list = prompt_token_list dtype = torch.int64 if use_i64_token else torch.int32 with torch.no_grad(): - while token_list[-1] != tokenizer.eos_id and pos < max_seq_len: - tmp_token_list = torch.tensor( - token_list[pos - ar_len : pos], dtype=dtype - ).reshape(1, -1) - tmp_pos = all_pos[:, pos - ar_len : pos] - tmp_atten_mask = atten_mask - if pos < ar_len: - tmp_token_list = torch.cat( - [ - torch.zeros((1, ar_len - pos), dtype=dtype), - torch.tensor(token_list, dtype=dtype).reshape(1, -1), - ], - dim=1, - ) - tmp_pos = torch.cat( - [ - torch.zeros((1, ar_len - pos), dtype=torch.int32), - all_pos[:, :pos], - ], - dim=1, - ) - tmp_atten_mask = torch.cat( - [ - torch.ones(1, ar_len, max_seq_len - pos) * -255.0, - atten_mask[:, :, -pos:], - ], - dim=-1, - ) + # Phase 1: Prefill the prompt in ar_len chunks. + num_prompt_tokens = len(prompt_token_list) + pos = 0 # Tracks how many prompt tokens have been processed. + while pos < num_prompt_tokens: + chunk_start_idx = pos + # Take a chunk of prompt tokens, up to ar_len length. + chunk_end_idx = min(num_prompt_tokens, pos + ar_len) + actual_chunk_tokens = prompt_token_list[chunk_start_idx:chunk_end_idx] + num_tokens_in_chunk = len(actual_chunk_tokens) + + # Prepare tmp_token_list (padded with zeros). + tmp_token_list = torch.zeros((1, ar_len), dtype=dtype) + tmp_token_list[0, :num_tokens_in_chunk] = torch.tensor( + actual_chunk_tokens, dtype=dtype + ) + + # Prepare tmp_pos (padded with zeros). + tmp_pos = torch.zeros((1, ar_len), dtype=torch.int32) + tmp_pos[0, :num_tokens_in_chunk] = all_pos[ + 0, + pos : pos + num_tokens_in_chunk, + ] + + # Run inference. + logits, new_k_caches, new_v_caches = module( + tmp_token_list, + *atten_mask, + tmp_pos, + *k_caches, + *v_caches, + ) + if collect_logits: + result_logits.append(logits[:, :num_tokens_in_chunk]) + + # We should have enough calibration data when generating last token if task was specified + if seq_mse_candidates != 0 and pos == num_prompt_tokens - 1: + with SeqMSE(module, seq_mse_candidates): + module( + tmp_token_list, + *atten_mask, + tmp_pos, + *k_caches, + *v_caches, + ) + + # Update the pos, KV cache and attention mask. + pos, k_caches, v_caches = kv_updater( + num_tokens_in_chunk, + atten_mask, + pos, + k_caches, + v_caches, + new_k_caches, + new_v_caches, + ) + # Append the last run logits to the total_token_list. + total_token_list.append( + torch.argmax(logits[:, num_tokens_in_chunk - 1], dim=-1).item() + ) + + # Phase 2: Generate tokens until the EOS token is generated or max_seq_len is reached. + # When run on wikitext for ppl evaluation, this while-loop is not expected to run. + max_cache_len = max_seq_len - ar_len + num_tokens = len(total_token_list) + while total_token_list[-1] != tokenizer.eos_id and num_tokens < max_seq_len: + chunk_start_idx = min(pos, max_cache_len) + # Take a chunk of generated tokens, up to ar_len length. + chunk_end_idx = num_tokens + actual_chunk_tokens = total_token_list[chunk_start_idx:chunk_end_idx] + num_tokens_in_chunk = len(actual_chunk_tokens) + + # Prepare tmp_token_list (padded with zeros). + tmp_token_list = torch.zeros((1, ar_len), dtype=dtype) + tmp_token_list[0, :num_tokens_in_chunk] = torch.tensor( + actual_chunk_tokens, dtype=dtype + ) + + # Prepare tmp_pos (padded with zeros). + tmp_pos = torch.zeros((1, ar_len), dtype=torch.int32) + tmp_pos[0, :num_tokens_in_chunk] = all_pos[0, chunk_start_idx:chunk_end_idx] logits, new_k_caches, new_v_caches = module( tmp_token_list, - tmp_atten_mask, + *atten_mask, tmp_pos, *k_caches, *v_caches, ) if collect_logits: - result_logits.append(logits) - atten_mask, pos, k_caches, v_caches = kv_updater( - ar_len, atten_mask, pos, k_caches, v_caches, new_k_caches, new_v_caches + result_logits.append(logits[:, :num_tokens_in_chunk]) + + pos, k_caches, v_caches = kv_updater( + 1, + atten_mask, + pos, + k_caches, + v_caches, + new_k_caches, + new_v_caches, ) - if pos > len(token_list): - token_list.append(torch.argmax(logits[:, -1], dim=-1).item()) + total_token_list.append( + torch.argmax(logits[:, num_tokens_in_chunk - 1], dim=-1).item() + ) + num_tokens = len(total_token_list) - logging.info(f"kv inference result:\n{tokenizer.decode(token_list)}") + logging.info(f"kv inference result:\n{tokenizer.decode(total_token_list)}") if collect_logits: result_logits = torch.cat(result_logits, dim=1) return result_logits +@register_inference(use_kv_cache=False) def prefill_inference( get_example_inputs, prompt: Union[str, list], @@ -360,6 +467,7 @@ def prefill_inference( else: raise RuntimeError("Unknown tokenizer") else: + # pyre-ignore token_list = prompt.flatten().tolist() pos = len(token_list) @@ -376,61 +484,61 @@ def prefill_inference( ], dim=1, ) - results = module( - tmp_token_list, - atten_mask, - ) + results = module(tmp_token_list, *atten_mask) if len(results) == 3: - logits, new_k_caches, new_v_caches = results + logits, _, _ = results elif len(results) == 1: logits = results - logits = torch.argmax(logits[:, pos - 1], dim=-1).item() - token_list.append(logits) + token = torch.argmax(logits[:, pos - 1], dim=-1).item() + token_list.append(token) if collect_logits: - result_logits.append(logits) + result_logits = logits[:, :pos] pos += 1 logging.info(f"prefill inference result:\n{tokenizer.decode(token_list)}") - if collect_logits: - result_logits = torch.cat(result_logits, dim=1) return result_logits def graph_module_inference( - args, - use_kv_cache, + use_kv_cache: bool, get_example_inputs: Callable, module: torch.fx.GraphModule, tokenizer, ar_len=1, max_seq_len=512, kv_updater=smart_mask_updater, + prompt=None, + tasks=None, + tasks_limit=1, + num_fewshot=None, use_i64_token=False, - event_name: str = None, + event_name: Optional[str] = None, + seq_mse_candidates: int = 0, ): - if args.tasks is None: + """ + This function supports model execution from static nn.Module decoder model + all the way to edge program. + Users could choose to provide either the prompt or tasks for execution but not both. + """ + # Checks 1 and only 1 is provided. + assert (tasks is None) != ( + prompt is None + ), "Please provide either tasks or prompt - not both or neither" + if tasks is None: + kwargs = {} if use_kv_cache: - kv_inference( - get_example_inputs, - args.prompt[0], - module, - tokenizer, - ar_len, - max_seq_len, - kv_updater=kv_updater, - use_i64_token=use_i64_token, - collect_logits=False, - ) - else: - prefill_inference( - get_example_inputs, - args.prompt[0], - module, - tokenizer, - max_seq_len, - use_i64_token, - collect_logits=False, - ) + kwargs["ar_len"] = ar_len + kwargs["kv_updater"] = kv_updater + INFERENCE_REGISTRY[use_kv_cache]( + get_example_inputs, + prompt, + module, + tokenizer, + max_seq_len=max_seq_len, + use_i64_token=use_i64_token, + collect_logits=False, + **kwargs, + ) else: calibration_wrapper = GraphModuleCalibrationWrapper( model=module, @@ -441,14 +549,30 @@ def graph_module_inference( get_example_inputs=get_example_inputs, kv_updater=kv_updater, use_i64_token=use_i64_token, + seq_mse_candidates=seq_mse_candidates, ) # Evaluate the model with torch.no_grad(): eval_results = simple_evaluate( model=calibration_wrapper, - tasks=args.tasks, - limit=args.limit, + tasks=tasks, + num_fewshot=num_fewshot, + limit=tasks_limit, ) logging.info(f"Perplexity evaluation summary for {event_name}") for task, res in eval_results["results"].items(): logging.info(f"{task}: {res}") + + +def apply_prompt_template( + chat_template: Callable, prompt: str, system_prompt: str = None +): + messages = [{"role": "user", "content": prompt}] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + + template_prompt = chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + logging.info(f"Prompt after applying template: {template_prompt}") + return template_prompt diff --git a/examples/qualcomm/oss_scripts/llama/eval_llama_qnn.py b/examples/qualcomm/oss_scripts/llama/eval_llama_qnn.py index 00c36a59582..b25e0cbdc7d 100644 --- a/examples/qualcomm/oss_scripts/llama/eval_llama_qnn.py +++ b/examples/qualcomm/oss_scripts/llama/eval_llama_qnn.py @@ -10,12 +10,14 @@ import logging import sys import types +from functools import partial import torch - from executorch.backends.qualcomm.quantizer.custom_annotation import ( - annotate_linear_16a8w_in_affine_layer, - annotate_matmul_16a8w, + annotate_kv_8bit, + annotate_output_16a8w, + annotate_qkv_proj_sha, + StaticLLMQuantConfig, ) from executorch.backends.qualcomm.quantizer.observers.per_channel_param_observer import ( @@ -23,6 +25,7 @@ ) from executorch.backends.qualcomm.quantizer.qconfig import ( _derived_bias_quant_spec, + get_ptq_per_channel_quant_config, QuantizationConfig, ) @@ -51,11 +54,11 @@ set_scales, WrappedLlamaModel, ) - from lm_eval.evaluator import simple_evaluate from pytorch_tokenizers import get_tokenizer from torchao.prototype.spinquant import apply_spinquant +from torchao.quantization.pt2e import MinMaxObserver from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import QuantizationSpec @@ -207,11 +210,22 @@ def gen_eval_wrapper(model_name, args): if args.ptq is not None: quant_dtype = getattr(QuantDtype, f"use_{args.ptq}") - custom_annotations = (annotate_matmul_16a8w,) + quantization_config_wv_sha_8a4w = get_ptq_per_channel_quant_config( + act_dtype=torch.uint8, + weight_dtype=torch.int4, + act_observer=MinMaxObserver, + act_symmetric=True, + ) + custom_annotations = ( + annotate_kv_8bit, + partial( + annotate_qkv_proj_sha, + qkv_tags={StaticLLMQuantConfig.wv_sha}, + quantization_config=quantization_config_wv_sha_8a4w, + ), + ) if args.llama_model == "stories110m": - custom_annotations = custom_annotations + ( - annotate_linear_16a8w_in_affine_layer, - ) + custom_annotations = custom_annotations + (annotate_output_16a8w,) quantizer = make_custom_quantizer( quant_dtype, args.range_setting, custom_annotations, args.quant_linear_only diff --git a/examples/qualcomm/oss_scripts/llama/llama.py b/examples/qualcomm/oss_scripts/llama/llama.py index b37dc75dc39..273829d214e 100755 --- a/examples/qualcomm/oss_scripts/llama/llama.py +++ b/examples/qualcomm/oss_scripts/llama/llama.py @@ -32,8 +32,6 @@ from executorch.backends.qualcomm.builders.utils import is_graph_output from executorch.backends.qualcomm.quantizer.custom_annotation import ( - annotate_linear_16a8w_in_affine_layer, - annotate_matmul_16a8w, annotate_prefill_kv_output, ) @@ -43,7 +41,6 @@ from executorch.backends.qualcomm.utils.constants import ( QCOM_PASS_ACTIVATE_KEY, QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY, - QCOM_QUANT_ATTRS_MAP, ) from executorch.backends.qualcomm.utils.utils import ( convert_linear_to_conv2d, @@ -64,12 +61,14 @@ from executorch.examples.models.llama.source_transformation.quantize import ( get_quant_embedding_transform, ) -from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import ( - DECODER_MODEL_VERSION, - EVAL_MODE, - HUGGING_FACE_REPO_IDS, +from executorch.examples.qualcomm.oss_scripts.llama import ( + LLM_VARIANT_ARCHS, + LLMModelConfig, + SUPPORTED_LLM_MODELS, ) +from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import EVAL_MODE from executorch.examples.qualcomm.oss_scripts.llama.decoder_utils import ( + apply_prompt_template, graph_module_inference, QnnRunnerEvalWrapper, shift_pointer_updater, @@ -116,6 +115,8 @@ FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) logging.getLogger().setLevel(logging.INFO) +# Avoid the error message "Could not initialize NNPACK! Reason: Unsupported hardware." +torch.backends.nnpack.set_flags(False) def next_power_of_two(n): @@ -125,24 +126,26 @@ def next_power_of_two(n): class SingleLlama: - def __init__(self, decoder_model, pte_filename) -> None: + def __init__( + self, decoder_model, decoder_model_config: LLMModelConfig, pte_filename + ) -> None: super().__init__() self.decoder_model = decoder_model + self.decoder_model_config = decoder_model_config self.passes_job = get_capture_program_passes() self.dep_table = get_passes_dependency_for_capture_program() - self.quant_attrs = None self.quant_dtype = None self.llama_meta = self.decoder_model.get_metadata() self.has_quant_io = False self.pte_filename = pte_filename - if self.llama_meta["get_use_kv_cache"]: - tokens, atten_mask, pos_ids, k_caches, v_caches = self.get_example_inputs( - use_kv_cache=True - ) - self.inputs = (tokens, atten_mask, pos_ids, *k_caches, *v_caches) - else: - tokens, atten_mask = self.get_example_inputs(use_kv_cache=False) - self.inputs = (tokens, atten_mask) + inputs = self.get_example_inputs(self.llama_meta["get_use_kv_cache"]) + self.inputs = ( + inputs[0], # tokens + *inputs[1], # attn_mask + *((inputs[2],) if self.llama_meta["get_use_kv_cache"] else []), # pos_ids + *(inputs[3] if self.llama_meta["get_use_kv_cache"] else []), # k_caches + *(inputs[4] if self.llama_meta["get_use_kv_cache"] else []), # v_caches + ) self.llama_graph_module = decoder_model self.io_shape = { # logit output @@ -219,25 +222,31 @@ def quantize( tokenizer, custom_annotations=(), scales_state_dict=None, + chat_template=None, ): self.quant_dtype = quant_dtype quantizer = make_custom_quantizer( quant_dtype, args.range_setting, custom_annotations ) - self.has_quant_io = True fx_graph_module = None - with torch.no_grad(): fx_graph_module = torch.export.export( self.llama_graph_module, self.inputs, strict=True ).module() if quant_dtype == QuantDtype.use_16a4w_block: + if self.decoder_model_config.group_size is None: + raise ValueError( + "Group size is required when use quant_dtype 16a4w_block" + ) conv_nodes = [ n for n in fx_graph_module.graph.nodes if "conv" in n.name ] - block_size_map = {n.name: (1, 64, 1, 1) for n in conv_nodes} + block_size_map = { + n.name: (1, self.decoder_model_config.group_size, 1, 1) + for n in conv_nodes + } quantizer.set_block_size_map(block_size_map) fx_graph_module = prepare_pt2e(fx_graph_module, quantizer) @@ -245,8 +254,32 @@ def quantize( logging.info("Quantizing the model...") # Calibration + if args.tasks is not None: + graph_module_inference( + use_kv_cache=self.llama_meta["get_use_kv_cache"], + get_example_inputs=self.get_example_inputs, + module=fx_graph_module, + tokenizer=tokenizer, + ar_len=self.llama_meta["get_ar_len"], + max_seq_len=self.llama_meta["get_max_seq_len"], + kv_updater=args.kv_updater, + tasks=args.tasks, + tasks_limit=args.limit, + num_fewshot=args.num_fewshot, + use_i64_token=args.embedding_quantize is not None, + event_name="prepare_pt2e_tasks", + seq_mse_candidates=self.decoder_model_config.seq_mse_candidates, + ) + + # Check user's prompt, helps calibrate special token + prompt = ( + args.prompt[0] + if chat_template is None + else apply_prompt_template( + chat_template, args.prompt[0], args.system_prompt + ) + ) graph_module_inference( - args=args, use_kv_cache=self.llama_meta["get_use_kv_cache"], get_example_inputs=self.get_example_inputs, module=fx_graph_module, @@ -254,10 +287,10 @@ def quantize( ar_len=self.llama_meta["get_ar_len"], max_seq_len=self.llama_meta["get_max_seq_len"], kv_updater=args.kv_updater, + prompt=prompt, use_i64_token=args.embedding_quantize is not None, - event_name="prepare_pt2e", + event_name="prepare_pt2e_prompt", ) - if scales_state_dict: set_scales( fx_graph_module, scales_state_dict, self.llama_graph_module.head_dim @@ -265,11 +298,34 @@ def quantize( self.llama_graph_module = convert_pt2e(fx_graph_module) - logging.info("Verifying the QDQ model...") - if args.eval_perplexity: - # Check qdq cpu results + if args.verbose: + logging.info("Verifying the QDQ model...") + # qdq cpu ppl evaluation is time consuming, only enable when eval_perplexity + if args.eval_perplexity: + # Check qdq cpu results + graph_module_inference( + use_kv_cache=self.llama_meta["get_use_kv_cache"], + get_example_inputs=self.get_example_inputs, + module=self.llama_graph_module, + tokenizer=tokenizer, + ar_len=self.llama_meta["get_ar_len"], + max_seq_len=self.llama_meta["get_max_seq_len"], + kv_updater=args.kv_updater, + tasks=args.tasks, + tasks_limit=args.limit, + num_fewshot=args.num_fewshot, + use_i64_token=args.embedding_quantize is not None, + event_name="convert_pt2e_tasks", + ) + # Check user's prompt + prompt = ( + args.prompt[0] + if chat_template is None + else apply_prompt_template( + chat_template, args.prompt[0], args.system_prompt + ) + ) graph_module_inference( - args=args, use_kv_cache=self.llama_meta["get_use_kv_cache"], get_example_inputs=self.get_example_inputs, module=self.llama_graph_module, @@ -277,16 +333,32 @@ def quantize( ar_len=self.llama_meta["get_ar_len"], max_seq_len=self.llama_meta["get_max_seq_len"], kv_updater=args.kv_updater, + prompt=prompt, use_i64_token=args.embedding_quantize is not None, - event_name="convert_pt2e", + event_name="convert_pt2e_prompt", ) + def save_logits_quant_attrs(self): + for node in self.llama_graph_module.graph.nodes: + if node.op == "output": + for output_node in node.args[0]: + if ( + output_node.target + == torch.ops.quantized_decomposed.dequantize_per_tensor.default + ): + source_node = output_node.args[0].args[0] + if source_node.meta["val"].size() in self.io_shape: + self.llama_meta["get_logits_scale"] = output_node.args[1] + self.llama_meta["get_logits_zero_point"] = output_node.args[ + 2 + ] + break + def lowering_modules( self, work_space, use_fp16=False, soc_model=QcomChipset.SM8650, - num_sharding=1, shared_buffer=False, verbose=False, ): @@ -304,7 +376,8 @@ def lowering_modules( with torch.no_grad(): # backend option backend_options = generate_htp_compiler_spec( - use_fp16=use_fp16, use_multi_contexts=num_sharding > 1 + use_fp16=use_fp16, + use_multi_contexts=self.decoder_model_config.num_sharding > 1, ) compiler_specs = generate_qnn_executorch_compiler_spec( soc_model=soc_model, @@ -312,6 +385,8 @@ def lowering_modules( shared_buffer=shared_buffer, ) skip_node_op_set = {"llama.fallback.default"} + + self.save_logits_quant_attrs() edge_prog_mgr = to_edge_transform_and_lower_to_qnn( self.llama_graph_module, self.inputs, @@ -322,13 +397,7 @@ def lowering_modules( skip_node_op_set=skip_node_op_set, ) - for n in edge_prog_mgr.exported_program().graph.nodes: - if n.op == "output": - for node, output_encoding in n.meta[QCOM_QUANT_ATTRS_MAP].items(): - if node.meta["val"].size() in self.io_shape: - self.quant_attrs = output_encoding - - if num_sharding > 1: + if self.decoder_model_config.num_sharding > 1: update_spill_fill_size(edge_prog_mgr.exported_program()) if verbose: @@ -341,31 +410,40 @@ def lowering_modules( def get_example_inputs(self, use_kv_cache=True): return self.decoder_model.get_example_inputs(use_kv_cache) - def get_quant_attrs(self): - return self.quant_attrs - -def compile(args, pte_filename, tokenizer): +def compile( + args, + decoder_model_config: LLMModelConfig, + pte_filename: str, + tokenizer, + chat_template, +): os.makedirs(args.artifact, exist_ok=True) start_ts = time.time() kv_config, prefill_config = None, None if args.params: - with open(args.params) as f: - kv_config = ModelArgs(**json.load(f)) - elif args.decoder_model == "qwen2_5": - from importlib.resources import files - - data_dir = files("executorch").joinpath("examples/models/qwen2_5/config") - config_file = data_dir.joinpath("0_5b_config.json") - kv_config = ModelArgs(**json.loads(config_file.read_text())) + params_path = args.params + else: + params_path = decoder_model_config.params_path + with open(params_path) as f: + kv_config = ModelArgs(**json.load(f)) # TODO: support batch inputs if necessary kv_config.max_batch_size = 1 kv_config.max_seq_len = args.max_seq_len kv_config.use_kv_cache = True - kv_config.enable_masked_softmax = args.enable_masked_softmax - kv_config.enable_r3 = args.r3 + kv_config.enable_r3 = decoder_model_config.r3 + kv_config.kv_io_bit_width = decoder_model_config.get_kv_io_bit_width() + + if decoder_model_config.masked_softmax: + if is_qnn_sdk_version_less_than("2.35"): + logging.warning( + f"Masked softmax is supported after QNN SDK 2.35. Given sdk version {get_sdk_build_id()} is lower the target version. Disabling the feature." + ) + kv_config.enable_masked_softmax = False + else: + kv_config.enable_masked_softmax = True prefill_config = copy.copy(kv_config) prefill_config.use_kv_cache = ( @@ -374,39 +452,60 @@ def compile(args, pte_filename, tokenizer): llama_instance_list = [] use_i64_token = args.embedding_quantize is not None + extra_kwargs = {} + if args.decoder_model == "gemma3-1b": + from transformers import Gemma3Config + + hf_config = Gemma3Config.from_pretrained(decoder_model_config.repo_id) + extra_kwargs["layer_types"] = hf_config.text_config.layer_types + extra_kwargs["rope_local_base_freq"] = ( + hf_config.text_config.rope_local_base_freq + ) + extra_kwargs["sliding_window"] = hf_config.sliding_window + with torch.device("meta"): if args.model_mode == "kv": llama_instance_list.append( - LlamaModel( + LLM_VARIANT_ARCHS.get(args.decoder_model, LlamaModel)( kv_config, ar_len=1, output_new_cache_only=True, output_cache=True, use_i64_token=use_i64_token, + **extra_kwargs, ) ) elif args.model_mode == "hybrid": llama_instance_list.append( - LlamaModel( + LLM_VARIANT_ARCHS.get(args.decoder_model, LlamaModel)( kv_config, ar_len=1, output_new_cache_only=True, output_cache=True, use_i64_token=use_i64_token, + **extra_kwargs, ) ) llama_instance_list.append( - LlamaModel( + LLM_VARIANT_ARCHS.get(args.decoder_model, LlamaModel)( prefill_config, ar_len=args.prefill_ar_len, output_new_cache_only=True, output_cache=True, use_i64_token=use_i64_token, + **extra_kwargs, ) ) elif args.model_mode == "lookahead": + # TODO: Lookahead decoding is not yet supported for gemma3-1b. + # This will be implemented once the model architecture and KV update logic are adapted. + if args.decoder_model == "gemma3-1b": + raise NotImplementedError( + "gemma3-1b does not currently support lookahead decoding." + ) + llama_instance_list.append( - LlamaModel( + LLM_VARIANT_ARCHS.get(args.decoder_model, LlamaModel)( kv_config, # To get better performance, we round up to the nearest power of 2. ar_len=next_power_of_two( @@ -415,31 +514,38 @@ def compile(args, pte_filename, tokenizer): output_new_cache_only=True, output_cache=True, use_i64_token=use_i64_token, + **extra_kwargs, ) ) llama_instance_list.append( - LlamaModel( + LLM_VARIANT_ARCHS.get(args.decoder_model, LlamaModel)( prefill_config, ar_len=args.prefill_ar_len, output_new_cache_only=True, output_cache=True, use_i64_token=use_i64_token, + **extra_kwargs, ) ) else: raise RuntimeError(f"Unknown model_mode: {args.model_mode}.") if args.checkpoint is None: # HF models - model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] - if args.decoder_model == "qwen2_5": - from executorch.examples.models.qwen2_5 import ( # pyre-ignore[21] - convert_weights, - ) - - checkpoint = download_and_convert_hf_checkpoint(model_id, convert_weights) + checkpoint = download_and_convert_hf_checkpoint( + decoder_model_config.repo_id, + decoder_model_config.convert_weights.__func__, + ) state_dict = torch.load( checkpoint, weights_only=True, map_location="cpu", mmap=True ) + if args.decoder_model == "gemma3-1b": + for k, v in state_dict.items(): + if "norm" not in k: + continue + # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + state_dict[k] = v.float() + torch.ones(v.shape, dtype=torch.float32) + else: state_dict = torch.load( args.checkpoint, weights_only=True, map_location="cpu", mmap=True @@ -451,6 +557,7 @@ def compile(args, pte_filename, tokenizer): if args.decoder_model == "stories260k": state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()} + if decoder_model_config.transform_weight: # Change to HuggingFace weight to improve the performance of RoPE in HTP backend. def permute(w, heads): dim_0 = w.size(0) @@ -482,7 +589,7 @@ def permute(w, heads): end_load_ts = time.time() logging.info(f"Time for loading checkpoint: {end_load_ts - start_ts}") - if args.spinquant: + if decoder_model_config.r1 or decoder_model_config.r2: config = types.SimpleNamespace( dim=prefill_config.dim, head_dim=prefill_config.dim // prefill_config.n_heads, @@ -495,8 +602,8 @@ def permute(w, heads): # Currently this script is on CPU: run with CUDA_VISIBLE_DEVICES=-1 apply_spinquant( model, - use_r1=True, - use_r2=False, + use_r1=decoder_model_config.r1, + use_r2=decoder_model_config.r2, use_r4=False, pretrained_rotation_path=None, qkv_split=True, @@ -516,15 +623,19 @@ def permute(w, heads): model.to(torch.float) ar_len, model.ar_len = model.ar_len, model.max_seq_len tokens, atten_mask = model.get_example_inputs(use_kv_cache=False) - atten_mask.to(torch.float) + atten_mask.mask.to(torch.float) wrapped_model = WrappedLlamaModel( - model, atten_mask, model.use_kv_cache, args.max_seq_len, args.device + model, + atten_mask.mask, + model.use_kv_cache, + args.max_seq_len, + args.device, ) act_bits, weight_bits = { - "8a8w": (8, 8), - "16a4w": (16, 4), - "16a4w_block": (16, 4), - }[args.ptq] + QuantDtype.use_8a8w: (8, 8), + QuantDtype.use_16a4w: (16, 4), + QuantDtype.use_16a4w_block: (16, 4), + }[decoder_model_config.ptq] scales_state_dict = compute_scales( wrapped_model, tokens, weight_bits, act_bits, 1600 ) @@ -539,22 +650,26 @@ def permute(w, heads): layer.feed_forward.prepare_feedfoward_conv() use_fp16 = True + # "io_type" here refers to logits output and "kv_type" refers to kv_cache input/output. fixed_point_type = {"kv_type": torch.float32, "io_type": torch.float32} - if args.ptq: - use_fp16 = False - fixed_point_type["kv_type"] = torch.uint8 - if args.ptq == "8a8w": - fixed_point_type["io_type"] = torch.uint8 - elif args.ptq in ("16a4w", "16a4w_block", "16a8w"): + if decoder_model_config.ptq: + if decoder_model_config.get_kv_io_bit_width() == 8: + fixed_point_type["kv_type"] = torch.uint8 + elif decoder_model_config.get_kv_io_bit_width() == 16: + fixed_point_type["kv_type"] = torch.uint16 + else: + raise RuntimeError( + f"Unknown kv io bit width {decoder_model_config.get_kv_io_bit_width()}" + ) + + if decoder_model_config.get_logits_output_bit_width() == 16: fixed_point_type["io_type"] = torch.uint16 else: - assert args.ptq in [ - "8a8w", - "16a4w", - "16a4w_block", - "16a8w", - ], f"No support for quant type {args.ptq}. Support 8a8w, 16a4w and 16a4w_block." - quant_dtype = getattr(QuantDtype, f"use_{args.ptq}") + raise RuntimeError( + f"Unknown logits io bit width {decoder_model_config.get_logits_output_bit_width()}" + ) + + quant_dtype = decoder_model_config.ptq if args.dtype_override is not None: dtype_override = DType[args.dtype_override] @@ -570,26 +685,16 @@ def permute(w, heads): )(llama_instance_list[i]) llama_instance_list[i] = convert_linear_to_conv2d(llama_instance_list[i]) llama_instance_list[i] = SingleLlama( - llama_instance_list[i].eval(), pte_filename + llama_instance_list[i].eval(), decoder_model_config, pte_filename ) if args.embedding_quantize: llama_instance_list[i].passes_job[I64toI32][ QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY ]["skip_node"] = {"tokens"} - if args.ptq: + if decoder_model_config.ptq: start_quantize_ts = time.time() - custom_annotations = ( - # For qwen2.5, skip annotate_conv can improve result. - partial( - annotate_matmul_16a8w, - annotate_conv=args.ptq != "16a8w", - ), - ) - if args.decoder_model == {"stories110m", "stories260k"}: - custom_annotations = custom_annotations + ( - annotate_linear_16a8w_in_affine_layer, - ) + custom_annotations = decoder_model_config.custom_annotation kv_quant_attrs = {} for i, llama_instance in enumerate(llama_instance_list): llama_instance.quantize( @@ -598,6 +703,7 @@ def permute(w, heads): tokenizer=tokenizer, custom_annotations=custom_annotations, scales_state_dict=scales_state_dict, + chat_template=chat_template, ) # If hybrid and lookahead mode, we store kv output quant_attrs and apply to prefill output quant_attrs later if i == 0 and args.model_mode in ["hybrid", "lookahead"]: @@ -618,16 +724,57 @@ def permute(w, heads): llama_instance.passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ "get_quant_io_dtype_fn" ] = partial(llama_instance._tag_ios, fixed_point_type=fixed_point_type) + + # force overriding frozen parameters here for model quantizing under seq mse scenario + # this will make weight sharing work properly + if decoder_model_config.seq_mse_candidates != 0 and args.model_mode in [ + "hybrid", + "lookahead", + ]: + decode, prompt = [ + instance.llama_graph_module for instance in llama_instance_list + ] + override_nodes = { + str(node.meta["nn_module_stack"].values()): node + for node in prompt.graph.nodes + if node.target == torch.ops.aten.conv2d.default + } + indices_map = { + # (affine_tensor, group_size, scales, zero_points, dtype, min, max) + torch.ops.torchao.dequantize_affine: [0, 2, 3], + # (per_channel_tensor, scales, zero_points, dim, dtype, min, max) + torch.ops.quantized_decomposed.dequantize_per_channel.default: [ + 0, + 1, + 2, + ], + # should not need to worry about per-tensor case + } + for node in decode.graph.nodes: + if node.target == torch.ops.aten.conv2d.default: + if target_node := override_nodes.get( + str(node.meta["nn_module_stack"].values()) + ): + # arguments of conv: (input, weight, bias) + for i, dq_node in enumerate(node.args[1:]): + for index in indices_map[dq_node.target]: + setattr( + prompt, + target_node.args[i + 1].args[index].target, + getattr(decode, dq_node.args[index].target), + ) + else: + raise RuntimeError("failed to override quantization attribute") + end_quantize_ts = time.time() logging.info(f"Time for quantizing: {end_quantize_ts - start_quantize_ts}") start_lowering_ts = time.time() - quant_attrs = None - if args.num_sharding > 1: + if decoder_model_config.num_sharding > 1: for llama_instance in llama_instance_list: SplitGraph, setting = model_sharding.get_split_graph_pass( llama_instance.llama_meta["get_n_layers"], - shares=args.num_sharding, + shares=decoder_model_config.num_sharding, ) llama_instance.passes_job[SplitGraph] = setting llama_instance.dep_table[SplitGraph] = [FoldQDQ] @@ -638,17 +785,15 @@ def permute(w, heads): args.artifact, use_fp16=use_fp16, soc_model=get_soc_to_chipset_map()[args.model], - num_sharding=args.num_sharding, shared_buffer=args.shared_buffer, ) - quant_attrs = llama_instance_list[0].get_quant_attrs() elif args.model_mode in ["hybrid", "lookahead"]: sample_inputs_list = [ llama_instace.inputs for llama_instace in llama_instance_list ] backend_options = generate_htp_compiler_spec( use_fp16=use_fp16, - use_multi_contexts=args.num_sharding > 1, + use_multi_contexts=decoder_model_config.num_sharding > 1, use_weight_sharing=not args.enable_x86_64, # x86 emulator does not support weight sharing ) graph_names = ["kv_forward", "prefill_forward"] @@ -661,6 +806,8 @@ def permute(w, heads): ) for graph_name in graph_names ] + + llama_instance_list[1].save_logits_quant_attrs() edge_prog_mgr = to_edge_transform_and_lower_to_qnn( { graph_name: instance.llama_graph_module @@ -685,13 +832,8 @@ def permute(w, heads): }, skip_node_op_set={"llama.fallback.default"}, ) - for n in list(edge_prog_mgr._edge_programs.values())[0].graph.nodes: - if n.op == "output": - for node, output_encoding in n.meta[QCOM_QUANT_ATTRS_MAP].items(): - if node.meta["val"].size() in llama_instance_list[0].io_shape: - quant_attrs = output_encoding - if args.num_sharding > 1: + if decoder_model_config.num_sharding > 1: # TODO: add arg parser of spill_fill_size since weight-sharing based # context binaries cannot be opened in x86 host pass @@ -717,33 +859,17 @@ def permute(w, heads): exec_prog_mgr.write_to_file(file) end_lowering_ts = time.time() - - quant_attrs_path = ( - f"{args.artifact}/{pte_filename}_quant_attrs.json" - if args.quant_attrs_path is None - else args.quant_attrs_path - ) - if quant_attrs: - json.dump( - { - "scale": quant_attrs["scale"], - "zero_point": quant_attrs["zero_point"], - }, - open(quant_attrs_path, "w"), - ) - else: - logging.warning("Quant attributes of the logit is None.") - if args.quant_attrs_path is None: - args.quant_attrs_path = quant_attrs_path - logging.info(f"Time for compiling: {end_lowering_ts - start_lowering_ts}") -def inference(args, pte_filename, runtime_tokenizer_path, tokenizer): +def inference( + args, + decoder_model_config: LLMModelConfig, + pte_filename, + runtime_tokenizer_path, + tokenizer, +): assert args.model_mode in EVAL_MODE, f"Unknown model_mode: {args.model_mode}." - assert ( - args.decoder_model in DECODER_MODEL_VERSION - ), f"Unknown decoder_model: {args.decoder_model}." pte_path = ( f"{args.pre_gen_pte}/{pte_filename}.pte" @@ -803,12 +929,20 @@ def post_process(): seq_len = args.max_seq_len multi_prompts = " ".join([f'--prompt "{prompt}"' for prompt in args.prompt]) + lookahead_args = " ".join( + [ + f"--window {args.window}", + f"--gcap {args.gcap}", + f"--ngram {args.ngram}", + ] + ) runner_args = " ".join( [ multi_prompts, f"--eval_mode {EVAL_MODE[args.model_mode]}", f"--temperature {args.temperature}", f"--system_prompt '{args.system_prompt}'", + lookahead_args if args.model_mode == "lookahead" else "", ] ) @@ -829,7 +963,7 @@ def post_process(): [ f"export LD_LIBRARY_PATH={qnn_sdk}/lib/{target}/:{args.build_folder}/lib &&", f"./{args.build_folder}/examples/qualcomm/oss_scripts/llama/qnn_llama_runner", - f"--decoder_model_version {DECODER_MODEL_VERSION[args.decoder_model]}", + f"--decoder_model_version {decoder_model_config.decoder_model_version}", f"--tokenizer_path {runtime_tokenizer_path}", f"--model_path {pte_path}", f"--seq_len {seq_len}", @@ -851,16 +985,13 @@ def post_process(): [ f"cd {workspace} &&", f"./qnn_llama_runner", - f"--decoder_model_version {DECODER_MODEL_VERSION[args.decoder_model]}", + f"--decoder_model_version {decoder_model_config.decoder_model_version}", f"--tokenizer_path {os.path.basename(runtime_tokenizer_path)}", f"--model_path {pte_filename}.pte", f"--seq_len {seq_len}", "--output_path outputs/outputs.txt", f"--performance_output_path {performance_output_path}", f"--kv_updater {'SmartMask' if args.kv_updater == smart_mask_updater else 'ShiftPointer'}", - f"--window {args.window}", - f"--gcap {args.gcap}", - f"--ngram {args.ngram}", runner_args, ] ) @@ -877,10 +1008,10 @@ def post_process(): runner=f"examples/qualcomm/oss_scripts/llama/qnn_llama_runner", ) # No pregen inputs, input_list is not required - adb.push(inputs=[], input_list="", files=[runtime_tokenizer_path]) + adb.push(inputs=[], files=[runtime_tokenizer_path]) adb.execute(custom_runner_cmd=runner_cmd) - adb.pull(output_path=args.artifact, callback=post_process) + if args.ip and args.port != -1: inference_speed = 0 with open( @@ -934,13 +1065,6 @@ def _build_tasks_parser(parser): help="Number of examples in few-shot context", ) - parser.add_argument( - "--quant_attrs_path", - help="A json file holding logit's quant_attributes. This file is generated after model compilation, stored under the artifacts. This file is required when eval_perplexity is enabled", - type=str, - required=False, - ) - return parser @@ -955,17 +1079,10 @@ def _build_parser(): type=str, ) - parser.add_argument( - "-P", - "--ptq", - help="If specified, will do PTQ quantization. default is 16bits activation and 4bits weight. Support 8a8w, 16a4w and 16a4w_block, 16a8w.", - type=str, - ) - parser.add_argument( "--decoder_model", - choices=["stories260k", "stories110m", "llama3_2", "qwen2_5"], - help="The Llama model to export. Current available options are: [stories260k, stories110m, llama3_2, qwen2_5]", + choices=list(SUPPORTED_LLM_MODELS.keys()), + help=f"The llm model to export. Current available options are: { SUPPORTED_LLM_MODELS.keys()}", required=True, ) @@ -1034,17 +1151,10 @@ def _build_parser(): type=str, ) - parser.add_argument( - "--num_sharding", - type=int, - default=1, - help="Specify the number of splits by inserting the fallback custom op. The graph will be split evenly by layers.", - ) - parser.add_argument( "--model_mode", help="Export and inference kv mode, hybrid mode, or lookahead decoding mode", - default="kv", + default="hybrid", choices=["kv", "hybrid", "lookahead"], type=str, ) @@ -1106,25 +1216,6 @@ def _build_parser(): type=str, ) - parser.add_argument( - "--spinquant", - help="Apply SpinQuant (R1+R2) to the model. Uses random Hadamard matrices for rotations", - action="store_true", - ) - - parser.add_argument( - "--enable_masked_softmax", - help="The MaskedSoftmax feature is designed to optimize the LLMs accuracy and performance executed on HTP backend. Note that it is only supported starting from QNN 2.35.", - action="store_true", - ) - - parser.add_argument( - "--r3", - help="Enable SpinQuant R3 quantization optimization. Please notice enable R3 could possibly cause performance drop.", - action="store_true", - default=False, - ) - parser.add_argument("-v", "--verbose", action="store_true") return parser @@ -1137,6 +1228,11 @@ def export_llama(args) -> None: raise RuntimeError("Eval device perplexity is only supported for KV mode") if args.eval_perplexity and args.tasks is None: raise RuntimeError("Please provide --tasks to eval perplexity") + assert ( + args.decoder_model in SUPPORTED_LLM_MODELS + ), f"Unknown decoder_model: {args.decoder_model}." + decoder_model_config = SUPPORTED_LLM_MODELS[args.decoder_model] + logging.info(f"*** {args.decoder_model} ***\n%s", str(decoder_model_config)) if args.model_mode == "kv": pte_filename = "kv_llama_qnn" @@ -1161,6 +1257,7 @@ def export_llama(args) -> None: tokenizer = None runtime_tokenizer_path = "" + chat_template = None if args.decoder_model in {"stories110m", "stories260k"}: tokenizer = get_tokenizer(args.tokenizer_model) assert isinstance( @@ -1170,19 +1267,36 @@ def export_llama(args) -> None: args.tokenizer_bin is not None ), "Please provide tokenizer_bin for stories." runtime_tokenizer_path = args.tokenizer_bin - elif args.decoder_model == "llama3_2": + elif "llama3_2" in args.decoder_model: tokenizer = get_tokenizer(args.tokenizer_model) assert isinstance( tokenizer, TiktokenTokenizer ), f"Wrong tokenizer provided for llama3_2." runtime_tokenizer_path = args.tokenizer_model - elif args.decoder_model == "qwen2_5": - model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] + elif args.decoder_model in SUPPORTED_LLM_MODELS: + model_id = decoder_model_config.repo_id tokenizer = AutoTokenizer.from_pretrained(model_id) - runtime_tokenizer_path = tokenizer.save_pretrained(args.artifact)[-1] - tokenizer = get_tokenizer(runtime_tokenizer_path) - else: - raise RuntimeError(f"Unknown decoder_model: {args.decoder_model}.") + chat_template = ( + tokenizer.apply_chat_template + if hasattr(tokenizer, "apply_chat_template") + and decoder_model_config.instruct_model + else None + ) + tokenizer_artifacts = tokenizer.save_pretrained(args.artifact) + tokenizer_config = tokenizer_artifacts[0] + runtime_tokenizer_path = tokenizer_artifacts[-1] + tokenizer = get_tokenizer(runtime_tokenizer_path, tokenizer_config) + + # TODO: Remove this once error is resolved. + if args.decoder_model == "phi_4_mini": + with open(runtime_tokenizer_path, "r+") as file: + data = json.load(file) + # TODO: Encountered the following error during runtime, so switched behavior for now. + # Error: libc++abi: terminating due to uncaught exception of type std::runtime_error: invert=true is not supported for Split PreTokenizer. Only invert=false is supported. + data["pre_tokenizer"]["pretokenizers"][-2]["invert"] = False + file.seek(0) + json.dump(data, file, indent=4) + file.truncate() if args.kv_updater == "smart_mask": args.shared_buffer = True @@ -1192,18 +1306,15 @@ def export_llama(args) -> None: else: raise RuntimeError(f"Using an unknown kv update {args.kv_updater}") - if args.enable_masked_softmax and is_qnn_sdk_version_less_than("2.35"): - raise RuntimeError( - f"Masked softmax is supported after QNN SDK 2.35. Given sdk version {get_sdk_build_id()} is lower the target version" - ) - if args.pre_gen_pte: - inference(args, pte_filename, runtime_tokenizer_path, tokenizer) + inference( + args, decoder_model_config, pte_filename, runtime_tokenizer_path, tokenizer + ) print(f"Finish the running pre_gen_pte from {args.pre_gen_pte}") return if args.compile_only: - compile(args, pte_filename, tokenizer) + compile(args, decoder_model_config, pte_filename, tokenizer, chat_template) if args.ip and args.port != -1: pte_path = f"{args.artifact}/{pte_filename}.pte" @@ -1219,8 +1330,10 @@ def export_llama(args) -> None: print(f"Finish compile_only and save to {args.artifact}") return - compile(args, pte_filename, tokenizer) - inference(args, pte_filename, runtime_tokenizer_path, tokenizer) + compile(args, decoder_model_config, pte_filename, tokenizer, chat_template) + inference( + args, decoder_model_config, pte_filename, runtime_tokenizer_path, tokenizer + ) def main(): diff --git a/examples/qualcomm/oss_scripts/llama/masking_utils.py b/examples/qualcomm/oss_scripts/llama/masking_utils.py new file mode 100644 index 00000000000..8d9d9ead154 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/masking_utils.py @@ -0,0 +1,324 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from abc import ABC, abstractmethod +from typing import List, Union + +import torch + + +def create_causal_attn_mask(max_batch_size: int, ar_len: int, max_seq_len: int): + """ + Creating a causal attention mask (ar_len: 5, max_seq_len: 15) + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ● + + ● = activate (can attend), ○ = inactivate (masked) + """ + mask = torch.full((ar_len, ar_len), -255.0) + mask_cond = torch.arange(ar_len) + mask.masked_fill_(mask_cond.view(1, ar_len) <= mask_cond.view(ar_len, 1), 0) + + if max_seq_len != ar_len: + mask = torch.cat( + [ + torch.ones(ar_len, max_seq_len - ar_len) * -255.0, + mask, + ], + dim=-1, + ) + mask = mask[None, :, :].expand(max_batch_size, ar_len, max_seq_len) + return mask + + +def create_sliding_window_attn_mask( + max_batch_size: int, ar_len: int, max_seq_len: int, sliding_window: int +): + """ + Creating a sliding_window attention mask (ar_len: 5, max_seq_len: 15, sliding_window: 3) + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + + ● = activate (can attend), ○ = inactivate (masked) + """ + mask = torch.full((ar_len, ar_len), -255.0) + mask_cond = torch.arange(ar_len) + mask.masked_fill_( + (mask_cond.view(1, ar_len) <= mask_cond.view(ar_len, 1)) + & (mask_cond.view(ar_len, 1) - mask_cond.view(1, ar_len) < sliding_window), + 0, + ) + + if max_seq_len != ar_len: + mask = torch.cat( + [ + torch.ones(ar_len, max_seq_len - ar_len) * -255.0, + mask, + ], + dim=-1, + ) + mask = mask[None, :, :].expand(max_batch_size, ar_len, max_seq_len) + return mask + + +class BaseAttentionMask(ABC): + def __init__(self, max_batch_size: int, ar_len: int, max_seq_len: int): + """ + Base class for attention masks used in autoregressive or hybrid attention mechanisms. + + Args: + max_batch_size (int): Maximum batch size supported. + ar_len (int): Length of the autoregressive sequence. + max_seq_len (int): Maximum sequence length. + """ + self.max_batch_size = max_batch_size + self.ar_len = ar_len + self.max_seq_len = max_seq_len + + @property + @abstractmethod + def mask(self) -> torch.Tensor: + """ + Attention mask tensor that must be initialized by child classes. + """ + pass + + @abstractmethod + def smart_mask_update(self, pos, n_updates): + """ + Update the attention mask by smart mask update method after model forward. + + Args: + pos (int): Current position in the sequence. + n_updates (int): Number of new tokens to update. + """ + pass + + @abstractmethod + def shift_pointer_update(self, pos, n_updates): + """ + Update the attention mask by shift pointer update method after model forward. + + Args: + pos (int): Current position in the sequence. + n_updates (int): Number of tokens to shift. + """ + pass + + +class CausalAttentionMask(BaseAttentionMask): + def __init__(self, max_batch_size: int, ar_len: int, max_seq_len: int): + super().__init__(max_batch_size, ar_len, max_seq_len) + self._mask = create_causal_attn_mask(max_batch_size, ar_len, max_seq_len) + + @property + def mask(self): + return self._mask + + def smart_mask_update(self, pos, n_updates): + """ + Smart Mask mechanism for attention mask updating + + Initial mask(5x15) layout (before any updates): + Each row represents a query token in the autoregressive context. + ● = activate (can attend), ○ = inactivate (masked) + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ● + + After 1st update (e.g., pos=0, n_updates=5, sliding_window=3): + Newly added tokens are unmasked (set to 0). + + 0 ● ● ● ● ● ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ● ● ● ● ● ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ● ● ● ● ● ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ● ● ● ● ● ○ ○ ○ ○ ○ ● ● ● ● ○ + 4 ● ● ● ● ● ○ ○ ○ ○ ○ ● ● ● ● ● + + After 2nd update (e.g., pos=5, n_updates=5): + + 0 ● ● ● ● ● ● ● ● ● ● ● ○ ○ ○ ○ + 1 ● ● ● ● ● ● ● ● ● ● ● ● ○ ○ ○ + 2 ● ● ● ● ● ● ● ● ● ● ● ● ● ○ ○ + 3 ● ● ● ● ● ● ● ● ● ● ● ● ● ● ○ + 4 ● ● ● ● ● ● ● ● ● ● ● ● ● ● ● + """ + start_pos = pos + end_pos = pos + n_updates + self.mask[:, :, start_pos:end_pos] = 0 + + def shift_pointer_update(self, pos, n_updates): + """ + Shift Pointer mechanism for attention mask updating + + Initial mask(5x15) layout (before any updates): + Each row represents a query token in the autoregressive context. + ● = activate (can attend), ○ = inactivate (masked) + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ● ● + + After 1st update (e.g., pos=0, n_updates=5, sliding_window=3): + Newly added tokens are unmasked (set to 0). + + 0 ○ ○ ○ ○ ○ ● ● ● ● ● ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ● ● ● ● ● ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ● ● ● ● ● ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ● ● ● ● ● ● ● ● ● ○ + 4 ○ ○ ○ ○ ○ ● ● ● ● ● ● ● ● ● ● + + After 2nd update (e.g., pos=5, n_updates=5): + + 0 ● ● ● ● ● ● ● ● ● ● ● ○ ○ ○ ○ + 1 ● ● ● ● ● ● ● ● ● ● ● ● ○ ○ ○ + 2 ● ● ● ● ● ● ● ● ● ● ● ● ● ○ ○ + 3 ● ● ● ● ● ● ● ● ● ● ● ● ● ● ○ + 4 ● ● ● ● ● ● ● ● ● ● ● ● ● ● ● + """ + start_pos = -pos - n_updates - self.ar_len + end_pos = -pos - self.ar_len + self.mask[:, :, start_pos:end_pos] = 0 + + +class SlidingWindowAttentionMask(BaseAttentionMask): + def __init__( + self, + max_batch_size: int, + ar_len: int, + max_seq_len: int, + sliding_window: int, + ): + super().__init__(max_batch_size, ar_len, max_seq_len) + self._mask = create_sliding_window_attn_mask( + max_batch_size, ar_len, max_seq_len, sliding_window + ) + self.sliding_window = sliding_window + + @property + def mask(self): + return self._mask + + def smart_mask_update(self, pos, n_updates): + """ + Smart Mask mechanism for attention mask updating + + Initial mask(5x15) layout (before any updates): + Each row represents a query token in the autoregressive context. + ● = activate (can attend), ○ = inactivate (masked) + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + + After 1st update (e.g., pos=0, n_updates=5, sliding_window=3): + Newly added tokens are unmasked (set to 0). + Earlier tokens lose access to older cache due to sliding window limits. + + 0 ○ ○ ○ ● ● ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ● ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + + After 2nd update (e.g., pos=5, n_updates=5): + Sliding window shifts again, masking older positions and activate new postion. + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + """ + start_pos = pos + end_pos = pos + n_updates + # Unmask the same range in the sliding window mask + self.mask[:, :, start_pos:end_pos] = 0 + + for i in range(self.ar_len): + # Calculate how many cached tokens are still avalible for this row + avalible_cache_len = self.sliding_window - (i + 1) + + # If the current position exceeds available cache, mask the overflow + if end_pos > avalible_cache_len: + # Mask tokens that are no longer within the sliding window + # TODO: [Optional]: it can be optimized by computing the exact start index + self.mask[:, i, : end_pos - avalible_cache_len] = -255.0 + + def shift_pointer_update(self, pos, n_updates): + """ + Shift Pointer mechanism for attention mask updating + + Initial mask(5x15) layout (before any updates): + Each row represents a query token in the autoregressive context. + ● = activate (can attend), ○ = inactivate (masked) + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + + After 1st update (e.g., pos=0, n_updates=5, sliding_window=3): + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + + After 2nd update (e.g., pos=5, n_updates=5): + + 0 ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ ○ + 1 ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ ○ + 2 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ ○ + 3 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● ○ + 4 ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ○ ● ● ● + """ + + start_pos = -pos - n_updates - self.ar_len + end_pos = -pos - self.ar_len + self.mask[:, :, start_pos:end_pos] = 0 + + for i in range(self.ar_len): + avalible_cache_len = self.sliding_window - (i + 1) + if abs(start_pos + self.ar_len) > avalible_cache_len: + self.mask[ + :, + i, + start_pos : start_pos + + abs(start_pos + self.ar_len) + - avalible_cache_len, + ] = -255.0 + + +class AttentionMask: + def __init__(self, masks: Union[BaseAttentionMask, List[BaseAttentionMask]]): + self.masks = masks if isinstance(masks, list) else [masks] + + def smart_mask_update(self, pos, n_updates): + for mask in self.masks: + mask.smart_mask_update(pos, n_updates) + + def shift_pointer_update(self, pos, n_updates): + for mask in self.masks: + mask.shift_pointer_update(pos, n_updates) + + def __iter__(self): + return iter([mask.mask for mask in self.masks]) diff --git a/examples/qualcomm/oss_scripts/llama/model/static_llama.py b/examples/qualcomm/oss_scripts/llama/model/static_llama.py index 192f23de302..8dcfced95fb 100755 --- a/examples/qualcomm/oss_scripts/llama/model/static_llama.py +++ b/examples/qualcomm/oss_scripts/llama/model/static_llama.py @@ -13,9 +13,16 @@ import scipy import torch import torch.nn as nn -import torch.nn.functional as F from executorch.examples.models.llama.model_args import ModelArgs -from executorch.examples.models.llama.rope import precompute_freqs_cis +from executorch.examples.models.llama.rope import ( + hf_precompute_freqs_cis, + precompute_freqs_cis, +) +from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import ( + AttentionMask, + CausalAttentionMask, + SlidingWindowAttentionMask, +) def apply_rotary_emb_single( @@ -36,8 +43,26 @@ def apply_rotary_emb_single( return x_out +def apply_partial_rotary_emb_single( + x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor +) -> torch.Tensor: + + if x.dim() == 4: + freqs_cos = freqs_cos[None, :, None, :] + freqs_sin = freqs_sin[None, :, None, :] + + rotary_dim = freqs_cos.shape[-1] * 2 + + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x_r, x_i = x_rot[..., : x_rot.shape[-1] // 2], x_rot[..., x_rot.shape[-1] // 2 :] + x_out_r = x_r * freqs_cos - x_i * freqs_sin + x_out_i = x_r * freqs_sin + x_i * freqs_cos + x_rotated = torch.cat([x_out_r, x_out_i], dim=-1) + return torch.cat([x_rotated, x_pass], dim=-1) + + class LlamaAttention(nn.Module): - def __init__(self, config: ModelArgs, output_new_cache_only=False): + def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): super().__init__() self.config = config self.dim = config.dim @@ -48,6 +73,23 @@ def __init__(self, config: ModelArgs, output_new_cache_only=False): self.max_seq_len = config.max_seq_len self.output_new_cache_only = output_new_cache_only self.enable_masked_softmax = getattr(config, "enable_masked_softmax", False) + self.use_qk_norm = config.use_qk_norm + self.qk_norm_before_rope = config.qk_norm_before_rope + self.use_rope = ( + config.no_rope_layer_interval + and (layer_idx + 1) % config.no_rope_layer_interval + ) + + if self.use_qk_norm: + q_norm_dim = self.head_dim + k_norm_dim = self.head_dim + self.q_norm_fn = torch.nn.RMSNorm(q_norm_dim, eps=config.norm_eps) + self.k_norm_fn = torch.nn.RMSNorm(k_norm_dim, eps=config.norm_eps) + + if config.partial_rotary_factor < 1: + self.apply_rope_emb = apply_partial_rotary_emb_single + else: + self.apply_rope_emb = apply_rotary_emb_single self.wq = nn.Linear( self.dim, @@ -70,7 +112,7 @@ def __init__(self, config: ModelArgs, output_new_cache_only=False): self.scale = float(self.head_dim) ** 0.5 - if config.enable_r3: + if getattr(config, "enable_r3", False): self.register_buffer( "r3_weight", torch.tensor( @@ -151,7 +193,7 @@ def prepare_sha(self): ) self.wo_sha.weight.data.copy_(self.wo.weight[:, :, None, None]) - def forward_sha( + def forward_sha( # noqa: C901 self, hidden_states: torch.Tensor, freqs_cos: torch.Tensor, @@ -184,14 +226,26 @@ def forward_sha( .reshape(bsz, seq_len, self.head_dim) for wv_sha in self.wv_sha ] + for i in range(len(q)): - q[i] = apply_rotary_emb_single(q[i], freqs_cos, freqs_sin) - if self.config.enable_r3: - q[i] = torch.matmul(q[i], self.r3_weight.T) + if self.use_qk_norm and self.qk_norm_before_rope: + q[i] = self.q_norm_fn(q[i]) + if self.use_rope: + q[i] = self.apply_rope_emb(q[i], freqs_cos, freqs_sin) + if self.use_qk_norm and not self.qk_norm_before_rope: + q[i] = self.q_norm_fn(q[i]) + if getattr(self.config, "enable_r3", False): + q[i] = torch.matmul(q[i], self.r3_weight) + for i in range(len(k)): - k[i] = apply_rotary_emb_single(k[i], freqs_cos, freqs_sin) - if self.config.enable_r3: - k[i] = torch.matmul(k[i], self.r3_weight.T) + if self.use_qk_norm and self.qk_norm_before_rope: + k[i] = self.k_norm_fn(k[i]) + if self.use_rope: + k[i] = self.apply_rope_emb(k[i], freqs_cos, freqs_sin) + if self.use_qk_norm and not self.qk_norm_before_rope: + k[i] = self.k_norm_fn(k[i]) + if getattr(self.config, "enable_r3", False): + k[i] = torch.matmul(k[i], self.r3_weight) k[i] = k[i].transpose(1, 2) output_y = [] @@ -249,8 +303,18 @@ def forward( k = k.view(bsz, seq_len, self.n_kv_heads, self.head_dim) v = v.view(bsz, seq_len, self.n_kv_heads, self.head_dim) - q = apply_rotary_emb_single(q, freqs_cos, freqs_sin) - k = apply_rotary_emb_single(k, freqs_cos, freqs_sin).permute(0, 2, 3, 1) + if self.use_qk_norm and self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) + + if self.use_rope: + q = self.apply_rope_emb(q, freqs_cos, freqs_sin) + k = self.apply_rope_emb(k, freqs_cos, freqs_sin) + k = k.permute(0, 2, 3, 1) + + if self.use_qk_norm and not self.qk_norm_before_rope: + q = self.q_norm_fn(q) + k = self.k_norm_fn(k) output_kh, output_vh, output_y = [], [], [] kh, vh = [], [] @@ -306,6 +370,7 @@ def __init__(self, args: ModelArgs): self.w1 = nn.Linear(self.dim, self.hidden_dim, bias=False) self.w2 = nn.Linear(self.hidden_dim, self.dim, bias=False) self.w3 = nn.Linear(self.dim, self.hidden_dim, bias=False) + self.act_fn = args.act_fn.get_function() def prepare_feedfoward_conv(self): self.w1_conv = nn.Conv2d(self.dim, self.hidden_dim, 1, bias=False) @@ -327,25 +392,37 @@ def forward_feedfoward_conv(self, x): bsz, _, _ = x.size() x = torch.reshape(x, (bsz, -1, 1, self.dim)) x = x.transpose(1, 3) # Transpose right before and after Conv - x = self.w2_conv(F.silu(self.w1_conv(x)) * self.w3_conv(x)) + x = self.w2_conv(self.act_fn(self.w1_conv(x)) * self.w3_conv(x)) x = x.transpose(1, 3) x = torch.reshape(x, (bsz, -1, self.dim)) return x def forward(self, x): - return self.w2(F.silu(self.w1(x)) * self.w3(x)) + return self.w2(self.act_fn(self.w1(x)) * self.w3(x)) class LlamaDecoderLayer(nn.Module): - def __init__(self, config: ModelArgs, output_new_cache_only=False): + def __init__(self, layer_idx: int, config: ModelArgs, output_new_cache_only=False): super().__init__() self.dim = config.dim self.attention = LlamaAttention( - config=config, output_new_cache_only=output_new_cache_only + layer_idx=layer_idx, + config=config, + output_new_cache_only=output_new_cache_only, ) self.feed_forward = FeedForward(config) self.attention_norm = torch.nn.RMSNorm(config.dim, eps=config.norm_eps) self.ffn_norm = torch.nn.RMSNorm(config.dim, eps=config.norm_eps) + self.post_attention_norm = ( + torch.nn.RMSNorm(config.dim, eps=config.norm_eps) + if config.post_attention_norm + else None + ) + self.post_ffn_norm = ( + torch.nn.RMSNorm(config.dim, eps=config.norm_eps) + if config.post_ffn_norm + else None + ) def forward( self, @@ -364,8 +441,13 @@ def forward( k_caches=k_caches, v_caches=v_caches, ) + if self.post_attention_norm: + h = self.post_attention_norm(h) h = x + h - output = h + self.feed_forward(self.ffn_norm(h)) + out = self.feed_forward(self.ffn_norm(h)) + if self.post_ffn_norm: + out = self.post_ffn_norm(out) + output = h + out return output, k_cache, v_cache @@ -377,6 +459,7 @@ def __init__( output_new_cache_only=True, output_cache=True, use_i64_token=False, + **kwargs, ): super().__init__() self.dim = config.dim @@ -389,27 +472,39 @@ def __init__( self.vocab_size = config.vocab_size self.rope_freq_base = config.rope_freq_base self.use_kv_cache = config.use_kv_cache + self.embedding_scale_factor = config.embedding_scale_factor self.ar_len = ar_len self.output_new_cache_only = output_new_cache_only self.use_i64_token = use_i64_token self.output_cache = output_cache + self.kv_io_bit_width = config.kv_io_bit_width self.layers = nn.ModuleList( [ - LlamaDecoderLayer(config, self.output_new_cache_only) - for _ in range(config.n_layers) + LlamaDecoderLayer(i, config, self.output_new_cache_only) + for i in range(config.n_layers) ] ) self.norm = torch.nn.RMSNorm(config.dim, eps=config.norm_eps) self.output = nn.Linear(config.dim, config.vocab_size, bias=False) self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim) - freqs_cos, freqs_sin = precompute_freqs_cis( - config.head_dim, - config.max_seq_len, - config.rope_freq_base, - config.use_scaled_rope, - config.rope_scale_factor, - ) + if config.use_hf_rope: + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_seq_len, + config.rope_freq_base, + config.partial_rotary_factor, + ) + freqs_cos = freqs_cos[:, : freqs_cos.shape[-1] // 2] + freqs_sin = freqs_sin[:, : freqs_sin.shape[-1] // 2] + else: + freqs_cos, freqs_sin = precompute_freqs_cis( + config.head_dim, + config.max_seq_len, + config.rope_freq_base, + config.use_scaled_rope, + config.rope_scale_factor, + ) self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False) @@ -447,7 +542,7 @@ def forward( self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin ) - hidden_states = self.tok_embeddings(tokens) + hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) for ind, decoder_layer in enumerate(self.layers): k_caches = None v_caches = None @@ -479,22 +574,8 @@ def get_example_inputs(self, use_kv_cache=True): tokens = torch.randint( self.vocab_size, (self.max_batch_size, self.ar_len), dtype=dtype ) - - atten_mask = torch.full((self.ar_len, self.ar_len), torch.tensor(-255.0)) - mask_cond = torch.arange(atten_mask.size(-1)) - atten_mask.masked_fill_( - mask_cond < (mask_cond + 1).view(atten_mask.size(-1), 1), 0 - ) - if self.max_seq_len != self.ar_len: - atten_mask = torch.cat( - [ - torch.ones(self.ar_len, self.max_seq_len - self.ar_len) * -255.0, - atten_mask, - ], - dim=-1, - ) - atten_mask = atten_mask[None, :, :].expand( - self.max_batch_size, self.ar_len, self.max_seq_len + atten_mask = AttentionMask( + CausalAttentionMask(self.max_batch_size, self.ar_len, self.max_seq_len) ) if use_kv_cache: pos_ids = torch.zeros((self.max_batch_size, self.ar_len), dtype=torch.int32) @@ -546,4 +627,132 @@ def get_metadata(self): "get_n_layers": self.n_layers, "get_vocab_size": self.vocab_size, "get_use_kv_cache": self.use_kv_cache, + "get_kv_io_bit_width": self.kv_io_bit_width, } + + +class MultiScopeAwareLlamaModel(LlamaModel): + def __init__( + self, + config: ModelArgs, + ar_len=1, + output_new_cache_only=True, + output_cache=True, + use_i64_token=False, + **kwargs, + ): + super().__init__( + config=config, + ar_len=ar_len, + output_new_cache_only=output_new_cache_only, + output_cache=output_cache, + use_i64_token=use_i64_token, + ) + + for key in ["layer_types", "sliding_window", "rope_local_base_freq"]: + assert key in kwargs, f"Missing required argument: '{key}' in kwargs" + + # Get attention type for each layer + self.layer_types = kwargs["layer_types"] + # Get sliding window size (used in local/global attention) + self.sliding_window = kwargs["sliding_window"] + # Get local freq base for sliding attention + rope_freq_base = kwargs["rope_local_base_freq"] + + local_freqs_cos, local_freqs_sin = hf_precompute_freqs_cis( + config.head_dim, + config.max_seq_len, + rope_freq_base, + config.partial_rotary_factor, + ) + local_freqs_cos = local_freqs_cos[:, : local_freqs_cos.shape[-1] // 2] + local_freqs_sin = local_freqs_sin[:, : local_freqs_sin.shape[-1] // 2] + self.register_buffer("local_freqs_cos", local_freqs_cos, persistent=False) + self.register_buffer("local_freqs_sin", local_freqs_sin, persistent=False) + + def forward( + self, + tokens: torch.Tensor, + atten_mask: torch.Tensor, + window_atten_mask: torch.Tensor, + input_pos: Optional[torch.Tensor] = None, + *args, + ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + + output_k_cache = [] + output_v_cache = [] + # following tensors should be invariant across batches + freqs_cos = ( + self.freqs_cos[input_pos][0] if self.use_kv_cache else self.freqs_cos + ) + freqs_sin = ( + self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin + ) + local_freqs_cos = ( + self.local_freqs_cos[input_pos][0] + if self.use_kv_cache + else self.local_freqs_cos + ) + local_freqs_sin = ( + self.local_freqs_sin[input_pos][0] + if self.use_kv_cache + else self.local_freqs_sin + ) + + hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens) + for ind, decoder_layer in enumerate(self.layers): + k_caches = None + v_caches = None + if self.use_kv_cache: + offset_k = ind * self.n_kv_heads + offset_v = self.n_layers * self.n_kv_heads + offset_k + k_caches = args[offset_k : offset_k + self.n_kv_heads] + v_caches = args[offset_v : offset_v + self.n_kv_heads] + + if self.layer_types[ind] == "sliding_attention": + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=local_freqs_cos, + freqs_sin=local_freqs_sin, + atten_mask=window_atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + else: + hidden_states, k, v = decoder_layer( + hidden_states, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + atten_mask=atten_mask, + k_caches=k_caches, + v_caches=v_caches, + ) + + output_k_cache.extend(k) + output_v_cache.extend(v) + + hidden_states = self.norm(hidden_states) + logits = self.output(hidden_states) + if self.output_cache: + return logits, output_k_cache, output_v_cache + return logits + + def get_example_inputs(self, use_kv_cache=True): + inputs = list(super().get_example_inputs(use_kv_cache=use_kv_cache)) + causal_mask = CausalAttentionMask( + self.max_batch_size, self.ar_len, self.max_seq_len + ) + sliding_window_mask = SlidingWindowAttentionMask( + self.max_batch_size, + self.ar_len, + self.max_seq_len, + sliding_window=self.sliding_window, + ) + # Don't reverse the order of attention mask + inputs[1] = AttentionMask([causal_mask, sliding_window_mask]) + return tuple(inputs) + + def get_metadata(self): + meta_data = super().get_metadata() + meta_data["get_sliding_window"] = self.sliding_window + return meta_data diff --git a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp index 7004b793661..71eaea2b8d6 100644 --- a/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp +++ b/examples/qualcomm/oss_scripts/llama/qnn_llama_runner.cpp @@ -9,13 +9,15 @@ /** * @file * - * This tool can run Llama2 110M, Llama3.2 1B / 3B, Qwen2.5 0.5B with Qualcomm - * AI Engine Direct. + * This tool can run Llama2 110M, Llama3.2 1B / 3B, Gemma3 1B, + * phi4-mini-instruct, Qwen2.5 0.5B / 1.5B, Qwen3 0.6B / 1.7B, SmolLM2 135M, + * SmolLM3 3B with Qualcomm AI Engine Direct. * */ #include #include +#include #include #include #include @@ -61,7 +63,7 @@ DEFINE_int32( "Total number of tokens to generate (prompt + output)."); DEFINE_int32( eval_mode, - 0, + 1, "0: TokenGenerator(kv) / 1: HybridMode (prefill+kv) / 2: Lookahead Decoding"); DEFINE_string( kv_updater, @@ -101,7 +103,6 @@ std::string get_formatted_prompt( std::string formatted_prompt; switch (decoder_model_version) { case example::DecoderModelVersion::kLlama2: - case example::DecoderModelVersion::kQwen2_5: formatted_prompt.append(prompt); break; case example::DecoderModelVersion::kLlama3: @@ -116,6 +117,61 @@ std::string get_formatted_prompt( formatted_prompt.append( "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"); break; + case example::DecoderModelVersion::kGemma3: + formatted_prompt.append("user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("\n"); + formatted_prompt.append("model\n"); + if (!system_prompt.empty()) { + formatted_prompt.append(system_prompt); + formatted_prompt.append("\n"); + } + break; + case example::DecoderModelVersion::kPhi4: + if (!system_prompt.empty()) { + formatted_prompt.append("<|system|>"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("<|end|>"); + } + formatted_prompt.append("<|user|>"); + formatted_prompt.append(prompt); + formatted_prompt.append("<|end|><|assistant|>"); + break; + case example::DecoderModelVersion::kQwen2_5: + formatted_prompt.append(prompt); + break; + case example::DecoderModelVersion::kQwen3: + formatted_prompt.append("<|im_start|>user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("<|im_end|>\n"); + if (!system_prompt.empty()) { + formatted_prompt.append("<|im_start|>system\n"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("<|im_end|>\n"); + } + formatted_prompt.append("<|im_start|>assistant"); + break; + case example::DecoderModelVersion::kSmollm2_135m: + if (!system_prompt.empty()) { + formatted_prompt.append("<|im_start|>system\n"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("<|im_end|>\n"); + } + formatted_prompt.append("<|im_start|>user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("<|im_end|>\n\n"); + break; + case example::DecoderModelVersion::kSmollm3: + if (!system_prompt.empty()) { + formatted_prompt.append("<|im_start|>system\n"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("\n\n"); + } + formatted_prompt.append("<|im_start|>user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("<|im_end|>\n"); + formatted_prompt.append("<|im_start|>assistant\n"); + break; default: ET_CHECK_MSG(false, "unsupported llama version"); break; @@ -123,24 +179,16 @@ std::string get_formatted_prompt( return formatted_prompt; } -int main(int argc, char** argv) { - std::vector prompts = CollectPrompts(argc, argv); - gflags::ParseCommandLineFlags(&argc, &argv, true); - if (!gflags::GetCommandLineFlagInfoOrDie("prompt").is_default && - !gflags::GetCommandLineFlagInfoOrDie("tokenized_prompt").is_default) { - ET_CHECK_MSG(false, "Only provide prompt or tokenized_input but not both."); - } - if (!gflags::GetCommandLineFlagInfoOrDie("dump_logits_path").is_default && - FLAGS_eval_mode != 0) { - ET_CHECK_MSG( - false, "Only TokenGenerator(kv) mode is supported to dump all logits."); - } - +template +void start_runner( + std::unique_ptr module, + std::vector& prompts) { bool use_tokenized_prompt = gflags::GetCommandLineFlagInfoOrDie("tokenized_prompt").is_default ? false : true; // create llama runner - example::Runner runner( + example::Runner runner( + std::move(module), FLAGS_decoder_model_version.c_str(), FLAGS_model_path.c_str(), FLAGS_tokenizer_path.c_str(), @@ -161,13 +209,17 @@ int main(int argc, char** argv) { buf.push_back(c); } }; - + executorch::extension::llm::GenerationConfig config{ + true, + -1, + false, + FLAGS_seq_len, + static_cast(FLAGS_temperature), + 0, + 0}; if (use_tokenized_prompt) { - runner.generate( - FLAGS_tokenized_prompt.c_str(), - use_tokenized_prompt, - FLAGS_seq_len, - callback); + runner.generate_from_prompt_or_file( + FLAGS_tokenized_prompt.c_str(), use_tokenized_prompt, config, callback); } else { // generate tokens & store inference output for (int i = 0; i < FLAGS_num_iters; i++) { @@ -175,16 +227,51 @@ int main(int argc, char** argv) { std::string formatted_prompt; formatted_prompt = get_formatted_prompt( prompt, FLAGS_system_prompt, decoder_model_version.get()); - runner.generate( - formatted_prompt.c_str(), - use_tokenized_prompt, - FLAGS_seq_len, - callback); + runner.generate_from_prompt_or_file( + formatted_prompt.c_str(), use_tokenized_prompt, config, callback); } } } fout.write(buf.data(), buf.size()); fout.close(); +} + +int main(int argc, char** argv) { + std::vector prompts = CollectPrompts(argc, argv); + gflags::ParseCommandLineFlags(&argc, &argv, true); + if (!gflags::GetCommandLineFlagInfoOrDie("prompt").is_default && + !gflags::GetCommandLineFlagInfoOrDie("tokenized_prompt").is_default) { + ET_CHECK_MSG(false, "Only provide prompt or tokenized_input but not both."); + } + if (!gflags::GetCommandLineFlagInfoOrDie("dump_logits_path").is_default && + FLAGS_eval_mode != 0) { + ET_CHECK_MSG( + false, "Only TokenGenerator(kv) mode is supported to dump all logits."); + } + + std::unique_ptr module = + std::make_unique( + FLAGS_model_path.c_str(), + executorch::extension::Module::LoadMode::MmapUseMlockIgnoreErrors); + // Using 8bit as default since this meta is introduced with 16bit kv io + // support and older models only have 8bit kv io. + example::KvBitWidth kv_bitwidth = example::KvBitWidth::kWidth8; + if (module->method_names()->count("get_kv_io_bit_width") > 0) { + kv_bitwidth = static_cast( + module->get("get_kv_io_bit_width").get().toScalar().to()); + } + + if (kv_bitwidth == example::KvBitWidth::kWidth8) { + start_runner(std::move(module), prompts); + } else if (kv_bitwidth == example::KvBitWidth::kWidth16) { + start_runner(std::move(module), prompts); + } else { + ET_CHECK_MSG( + false, + "Unsupported kv bitwidth: %ld", + static_cast(kv_bitwidth)); + } + return 0; } diff --git a/examples/qualcomm/oss_scripts/llama/range_setting_pt2e.py b/examples/qualcomm/oss_scripts/llama/range_setting_pt2e.py index 4ef3e8cfe94..d5c68606f60 100644 --- a/examples/qualcomm/oss_scripts/llama/range_setting_pt2e.py +++ b/examples/qualcomm/oss_scripts/llama/range_setting_pt2e.py @@ -210,12 +210,13 @@ def make_custom_quantizer( per_channel_linear=True, act_observer=MinMaxObserver, ) - if range_setting in ("mse_weight_only", "mse_with_act_loss", "na"): - if range_setting == "na": - observer = PerChannelMinMaxObserver - elif range_setting == "mse_weight_only": + if range_setting in ("mse_weight_only", "mse_with_act_loss"): + assert ( + quant_dtype != QuantDtype.use_16a4w_block + ), "Range setting only supported for per-channel quantization" + if range_setting == "mse_weight_only": observer = PerChannelMSEObserver.with_args( - **{"steps": 200, "use_mse": True} + **{"steps": 1600, "use_mse": True} ) else: observer = PerChannelFixedQParamsObserver.with_args(**{"eps": 2**-12}) diff --git a/examples/qualcomm/oss_scripts/llama/runner/cache_utils.h b/examples/qualcomm/oss_scripts/llama/runner/cache_utils.h new file mode 100644 index 00000000000..11eeb52d9b4 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llama/runner/cache_utils.h @@ -0,0 +1,15 @@ +/* + * 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. + */ + +#pragma once + +enum CacheMode { + StaticCahce = 0, + // For models with global/local attention architecture (e.g., Gemma3), + HybridCache, +}; diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp index b563049eb8d..a049b54abb6 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp @@ -9,34 +9,35 @@ #include #include namespace example { -KVManager::KVManager(KVManagerMode kv_updater, Metadata metadata) +template +KVManager::KVManager(KVManagerMode kv_updater, Metadata metadata) : kv_updater_(kv_updater), metadata_(metadata) { k_cache_.resize( - metadata_.num_layers, std::vector(metadata_.num_heads)); + metadata_.num_layers, std::vector>(metadata_.num_heads)); v_cache_.resize( - metadata_.num_layers, std::vector(metadata_.num_heads)); + metadata_.num_layers, std::vector>(metadata_.num_heads)); // Calculate cache size switch (kv_updater_) { case KVManagerMode::SMART_MASK: { size_t cache_in_bytes = metadata_.num_layers * metadata_.num_heads * - metadata_.head_dim * metadata_.max_cache_len * sizeof(uint8_t); + metadata_.head_dim * metadata_.max_cache_len * sizeof(T); size_t cache_out_bytes = metadata_.num_layers * metadata_.num_heads * - metadata_.head_dim * metadata_.max_ar_len * sizeof(uint8_t); + metadata_.head_dim * metadata_.max_ar_len * sizeof(T); total_cache_size_ = 2 * (cache_in_bytes + cache_out_bytes); break; } case KVManagerMode::SHIFT_POINTER: { size_t k_cache_in_bytes = metadata_.num_layers * metadata_.num_heads * - (metadata_.head_dim + 1) * metadata_.max_cache_len * sizeof(uint8_t); + (metadata_.head_dim + 1) * metadata_.max_cache_len * sizeof(T); size_t k_cache_out_bytes = metadata_.num_layers * metadata_.num_heads * - metadata_.head_dim * metadata_.max_ar_len * sizeof(uint8_t); + metadata_.head_dim * metadata_.max_ar_len * sizeof(T); // Use the same memory for input and output of value cache in shift // pointer mode. Note that using context length to prevent exceeding the // range when the AR-N model updates the last block in shift pointer // mode. size_t v_cache_bytes = metadata_.num_layers * (metadata_.num_heads + 1) * - metadata_.head_dim * metadata_.context_len * sizeof(uint8_t); + metadata_.head_dim * metadata_.context_len * sizeof(T); total_cache_size_ = k_cache_in_bytes + k_cache_out_bytes + v_cache_bytes; break; } @@ -45,7 +46,8 @@ KVManager::KVManager(KVManagerMode kv_updater, Metadata metadata) } }; -void KVManager::init_attention_mask( +template +void KVManager::init_attention_mask( uint16_t* attention_mask, const std::vector& attention_map, int32_t ar_len, @@ -114,7 +116,86 @@ void KVManager::init_attention_mask( } } -void KVManager::update_attention_mask( +template +void KVManager::init_attention_mask( + uint16_t* attention_mask, + const std::vector& attention_map, + int32_t ar_len, + int32_t n_past, + int32_t sliding_window) { + ET_CHECK_MSG( + attention_map.size() <= ar_len, + "The size of attention_map (%zu) doesn't match with ar_len (%d)", + attention_map.size(), + ar_len); + uint16_t neg_val = 0; + uint16_t pos_val = 65535; + // Clear the attention mask + std::fill_n(attention_mask, ar_len * metadata_.context_len, neg_val); + + // SMART_MASK requires special handling of attention mask + switch (kv_updater_) { + case KVManagerMode::SMART_MASK: { + uint16_t* past_ptr = attention_mask; + uint16_t* new_ptr = attention_mask + (metadata_.context_len - ar_len); + // All inputs will necessarily attend to n_past and itself + for (int i = 0; i < ar_len; i++) { + // Iterate across ar_len + if (attention_map[i] < 0) { + // If negative, attend to only past tokens + std::fill_n(past_ptr, n_past, pos_val); + } else { + // If positive, copy attention map from (relative to 0th input) parent + // Parent token index + const int32_t pidx = attention_map[i]; + uint16_t* parent_ptr = attention_mask + pidx * metadata_.context_len; + std::memcpy( + past_ptr, parent_ptr, metadata_.context_len * sizeof(uint16_t)); + } + // Attend to itself + new_ptr[i] = pos_val; + + // mask by limitation of sliding_window + int32_t avalible_context_len = sliding_window - (i + 1) - n_past; + if (n_past > avalible_context_len) { + std::fill_n(past_ptr, n_past - avalible_context_len, neg_val); + } + + past_ptr += metadata_.context_len; + new_ptr += metadata_.context_len; + } + break; + } + case KVManagerMode::SHIFT_POINTER: { + // Only fill in ar_len. Rest will be padding + const size_t attn_row_start = metadata_.context_len - n_past - ar_len; + for (int i = 0; i < ar_len; i++) { + uint16_t* cur_ptr = + attention_mask + i * metadata_.context_len + attn_row_start; + // Attend to itself + cur_ptr[n_past + i] = pos_val; + if (attention_map[i] < 0) { + // If negative, attend to only past tokens + std::fill_n(cur_ptr, n_past, pos_val); + } else { + // If positive, copy attention map from (relative to 0th input) parent + // Parent token index + const int32_t pidx = attention_map[i]; + uint16_t* parent_ptr = + attention_mask + pidx * metadata_.context_len + attn_row_start; + std::memcpy( + cur_ptr, parent_ptr, (n_past + pidx + 1) * sizeof(uint16_t)); + } + } + break; + } + default: + break; + } +} + +template +void KVManager::update_attention_mask( uint16_t* attention_mask, int32_t ar_len, int32_t n_past, @@ -132,12 +213,47 @@ void KVManager::update_attention_mask( } } -void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { +template +void KVManager::update_attention_mask( + uint16_t* attention_mask, + int32_t ar_len, + int32_t n_past, + int32_t n_update, + int32_t sliding_window) { + uint16_t pos_val = 65535; + uint16_t neg_val = 0; + uint16_t* cur_ptr = attention_mask; + if (kv_updater_ == KVManagerMode::SMART_MASK) + cur_ptr += n_past; + if (kv_updater_ == KVManagerMode::SHIFT_POINTER) + cur_ptr += metadata_.context_len - n_past - ar_len - n_update; + + for (int i = 0; i < ar_len; i++) { + std::fill_n(cur_ptr, n_update, pos_val); + int32_t avalible_cache_len = sliding_window - (i + 1); + if (kv_updater_ == KVManagerMode::SMART_MASK) { + if (n_past + n_update > avalible_cache_len) { + std::fill_n( + cur_ptr - n_past, n_past + n_update - avalible_cache_len, neg_val); + } + } else if (kv_updater_ == KVManagerMode::SHIFT_POINTER) { + if (std::abs(n_past + ar_len) > avalible_cache_len) { + int32_t n_invalid = n_past - avalible_cache_len; + std::fill_n( + cur_ptr, std::abs(n_past + ar_len) - avalible_cache_len, neg_val); + } + + cur_ptr += metadata_.context_len; + } + } +} + +template +void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { cur_ar_len_ = ar_len; const size_t max_in_cache_block_in_bytes = - metadata_.max_cache_len * sizeof(uint8_t); - const size_t max_out_cache_block_in_bytes = - metadata_.max_ar_len * sizeof(uint8_t); + metadata_.max_cache_len * sizeof(T); + const size_t max_out_cache_block_in_bytes = metadata_.max_ar_len * sizeof(T); switch (kv_updater_) { case KVManagerMode::SMART_MASK: { @@ -148,14 +264,14 @@ void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { for (int layer = 0; layer < metadata_.num_layers; ++layer) { for (int head = 0; head < metadata_.num_heads; ++head) { // Allocate buffer for key cache and value cache - uint8_t* single_layer_k_cache_in = reinterpret_cast( - buffer_manager->allocate(cache_in_bytes)); - uint8_t* single_layer_k_cache_out = reinterpret_cast( - buffer_manager->allocate(cache_out_bytes)); - uint8_t* single_layer_v_cache_in = reinterpret_cast( - buffer_manager->allocate(cache_in_bytes)); - uint8_t* single_layer_v_cache_out = reinterpret_cast( - buffer_manager->allocate(cache_out_bytes)); + T* single_layer_k_cache_in = + reinterpret_cast(buffer_manager->allocate(cache_in_bytes)); + T* single_layer_k_cache_out = + reinterpret_cast(buffer_manager->allocate(cache_out_bytes)); + T* single_layer_v_cache_in = + reinterpret_cast(buffer_manager->allocate(cache_in_bytes)); + T* single_layer_v_cache_out = + reinterpret_cast(buffer_manager->allocate(cache_out_bytes)); k_cache_[layer][head].buffer = single_layer_k_cache_in; k_cache_[layer][head].output_buffer = single_layer_k_cache_out; @@ -171,20 +287,20 @@ void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { const size_t k_cache_out_size_in_bytes = metadata_.num_heads * metadata_.head_dim * max_out_cache_block_in_bytes; const size_t v_cache_size_in_bytes = (metadata_.num_heads + 1) * - metadata_.head_dim * metadata_.context_len * sizeof(uint8_t); + metadata_.head_dim * metadata_.context_len * sizeof(T); const int32_t single_head_size_in = metadata_.head_dim * metadata_.max_cache_len; const int32_t single_head_size_out = metadata_.head_dim * metadata_.max_ar_len; for (int layer = 0; layer < metadata_.num_layers; ++layer) { // Allocate buffer for key cache and value cache - uint8_t* single_layer_k_cache_in = reinterpret_cast( + T* single_layer_k_cache_in = reinterpret_cast( buffer_manager->allocate(k_cache_in_size_in_bytes)); - uint8_t* single_layer_k_cache_out = reinterpret_cast( + T* single_layer_k_cache_out = reinterpret_cast( buffer_manager->allocate(k_cache_out_size_in_bytes)); // Note that using context length to prevent exceeding the range when // the AR-N model updates the last block in shift pointer mode. - uint8_t* single_layer_v_cache = reinterpret_cast( + T* single_layer_v_cache = reinterpret_cast( buffer_manager->allocate(v_cache_size_in_bytes)); for (int head = 0; head < metadata_.num_heads; ++head) { k_cache_[layer][head].buffer = single_layer_k_cache_in + @@ -211,7 +327,8 @@ void KVManager::init_cache(IMemAlloc* buffer_manager, int32_t ar_len) { } } -void KVManager::rearrange_cache(int32_t ar_len_dst) { +template +void KVManager::rearrange_cache(int32_t ar_len_dst) { // Don't need to rearrange if cur_ar_len_ is equal to target ar_len if (cur_ar_len_ == ar_len_dst) return; @@ -225,15 +342,16 @@ void KVManager::rearrange_cache(int32_t ar_len_dst) { cur_ar_len_ = ar_len_dst; } -void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { +template +void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { // The output of key cache doesn't need to rearrange for both of SMART_MASK // and SHIFT_POINTER const int32_t src_cache_num = (cur_ar_len_ == metadata_.context_len) ? metadata_.context_len : metadata_.context_len - cur_ar_len_; const int32_t dst_cache_num = metadata_.context_len - ar_len_dst; - uint8_t* k_cache_in_read_ptr = k_cache.buffer; - uint8_t* k_cache_in_write_ptr = k_cache.buffer; + T* k_cache_in_read_ptr = k_cache.buffer; + T* k_cache_in_write_ptr = k_cache.buffer; if (src_cache_num > dst_cache_num) { if (kv_updater_ == KVManagerMode::SHIFT_POINTER) { @@ -263,7 +381,8 @@ void KVManager::rearrange_key(KVCache& k_cache, int32_t ar_len_dst) { } } -void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) { +template +void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) { // The input and output of the value cache don't need to rearrange for both // SMART_MASK and SHIFT_POINTER. However, the input pointer of the value cache // needs to be reset by ar_len_dst in SHIFT_POINTER mode. The output pointer @@ -276,7 +395,8 @@ void KVManager::rearrange_value(KVCache& v_cache, int32_t ar_len_dst) { } } -bool KVManager::update_cache_tensor( +template +bool KVManager::update_cache_tensor( std::vector>>& k_cache_in, std::vector>>& @@ -313,7 +433,8 @@ bool KVManager::update_cache_tensor( return updated; } -void KVManager::update_cache( +template +void KVManager::update_cache( int32_t ar_len, int32_t n_past, int32_t n_update, @@ -331,14 +452,15 @@ void KVManager::update_cache( } } -void KVManager::update_key( - KVCache& k_cache, +template +void KVManager::update_key( + KVCache& k_cache, int32_t n_past, int32_t n_update, const std::vector& selected) { - uint8_t* write_ptr = k_cache.buffer; - uint8_t* read_ptr = k_cache.output_buffer; - const int32_t copy_size = n_update * sizeof(uint8_t); + T* write_ptr = k_cache.buffer; + T* read_ptr = k_cache.output_buffer; + const int32_t copy_size = n_update * sizeof(T); const int32_t iter_size = (cur_ar_len_ == metadata_.context_len) ? metadata_.context_len : metadata_.context_len - cur_ar_len_; @@ -374,14 +496,15 @@ void KVManager::update_key( } } -void KVManager::update_value( - KVCache& v_cache, +template +void KVManager::update_value( + KVCache& v_cache, int32_t n_past, int32_t n_update, const std::vector& selected) { - uint8_t* write_ptr = v_cache.buffer; - uint8_t* read_ptr = v_cache.output_buffer; - const int32_t copy_size = n_update * metadata_.head_dim * sizeof(uint8_t); + T* write_ptr = v_cache.buffer; + T* read_ptr = v_cache.output_buffer; + const int32_t copy_size = n_update * metadata_.head_dim * sizeof(T); const int32_t past_size = n_past * metadata_.head_dim; if (kv_updater_ == KVManagerMode::SMART_MASK) @@ -403,7 +526,7 @@ void KVManager::update_value( auto wp = write_ptr, rp = read_ptr; for (auto sel : selected) { if (sel) { - std::memcpy(wp, rp, metadata_.head_dim * sizeof(uint8_t)); + std::memcpy(wp, rp, metadata_.head_dim * sizeof(T)); wp += metadata_.head_dim; update_times--; if (update_times == 0) @@ -414,4 +537,8 @@ void KVManager::update_value( } } +// Explicit instantiations +template class KVManager; +template class KVManager; + } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h index e1a756d1215..af9cf49a34f 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h +++ b/examples/qualcomm/oss_scripts/llama/runner/kv_manager.h @@ -15,9 +15,10 @@ namespace example { // Structure to hold key-value cache buffers +template struct KVCache { - uint8_t* buffer; - uint8_t* output_buffer; + T* buffer; + T* output_buffer; }; // Enumeration for key-value manager modes @@ -26,6 +27,7 @@ enum KVManagerMode { SMART_MASK = 0x0, SHIFT_POINTER = 0x1 }; * @class KVManager * @brief Class for kv cache update, rearrangement, and buffer allocatation. */ +template class KVManager { public: struct Metadata { @@ -76,6 +78,31 @@ class KVManager { int32_t ar_len, int32_t n_past); + /** + * @brief Initialize attention mask based on kv manager mode, and attention + * map. + * For example, + * ar_len = 4, CL = 6, n_past = 0, + * attention map: {-1, 0, 1, 2} and SMART_MASK. + * Attention_mask will be: + * [ 0 0 65535 0 0 0 ] + * [ 0 0 65535 65535 0 0 ] + * [ 0 0 65535 65535 65535 0 ] + * [ 0 0 65535 65535 65535 65535 ] + * @param attention_mask Pointer to the attention mask array to be + * initialized. + * @param attention_map Vector containing the attention map values. The shape + * of attention map should be [ar_len]. + * @param ar_len Length of input tokens. + * @param n_past Number of past elements in the cache. + */ + void init_attention_mask( + uint16_t* attention_mask, + const std::vector& attention_map, + int32_t ar_len, + int32_t n_past, + int32_t sliding_window); + /** * @brief Update attention mask based on kv manager mode, and n_update. * @param attention_mask Pointer to the attention mask array to be @@ -90,6 +117,23 @@ class KVManager { int32_t n_past, int32_t n_update); + /** + * @brief Update attention mask based on kv manager mode, and n_update. + * @param attention_mask Pointer to the attention mask array to be + * initialized. + * @param ar_len Length of input tokens. + * @param n_past Number of past elements in the cache. + * @param n_update Number of elements to be updated. + * @param sliding_window Length of sliding window for sliding window attention + * mask + */ + void update_attention_mask( + uint16_t* attention_mask, + int32_t ar_len, + int32_t n_past, + int32_t n_update, + int32_t sliding_window); + /** * @brief Reset the data pointer of the I/O cache tensor based on number of * past cache, kv manager mode, current ar length and KV cache data pointer @@ -128,10 +172,10 @@ class KVManager { int32_t n_update, const std::vector& selected); - const std::vector>& get_k_cache_() const { + const std::vector>>& get_k_cache_() const { return k_cache_; } - const std::vector>& get_v_cache_() const { + const std::vector>>& get_v_cache_() const { return v_cache_; } @@ -141,15 +185,15 @@ class KVManager { private: // Helper functions to rearrange and update key and value caches - void rearrange_key(KVCache& k_cache, int32_t ar_len_dst); - void rearrange_value(KVCache& v_cache, int32_t ar_len_dst); + void rearrange_key(KVCache& k_cache, int32_t ar_len_dst); + void rearrange_value(KVCache& v_cache, int32_t ar_len_dst); void update_key( - KVCache& k_cache, + KVCache& k_cache, int32_t n_past, int32_t n_update, const std::vector& selected); void update_value( - KVCache& v_cache, + KVCache& v_cache, int32_t n_past, int32_t n_update, const std::vector& selected); @@ -162,7 +206,7 @@ class KVManager { // Store start pointer of k and v cache for input and output // input: layer -> head -> head_dim * max_cache_len // output: layer -> head -> head_dim * max_ar_len - std::vector> k_cache_; - std::vector> v_cache_; + std::vector>> k_cache_; + std::vector>> v_cache_; }; } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp b/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp index 9b5030c461c..1692caa2756 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp @@ -13,28 +13,31 @@ using executorch::runtime::Result; namespace example { -void LhdTokenGenerator::prepare_io( +template +void LhdTokenGenerator::prepare_io( std::vector input_tokens, std::vector input_pos) { for (int i = 0; i < metadata_.ar_len; i++) { if (i < input_tokens.size()) { // Prepare pos data - input_pos_.data[i] = input_pos[i]; + this->input_pos_.data[i] = input_pos[i]; // Support CPU 4-bit embedding, which requires int64 input. // However, for QNN embedding, only int32 input is needed. // Therefore, we need to cast to the correct type to write the data. if (metadata_.use_int64_token) { - input_toks_.data[i] = input_tokens[i]; + this->input_toks_.data[i] = input_tokens[i]; } else { - int32_t* input_toks_ptr = reinterpret_cast(input_toks_.data); + int32_t* input_toks_ptr = + reinterpret_cast(this->input_toks_.data); input_toks_ptr[i] = static_cast(input_tokens[i]); } } } } -void LhdTokenGenerator::init_attention_mask(int32_t n_past) { +template +void LhdTokenGenerator::init_attention_mask(int32_t n_past) { std::vector attention_map; attention_map.reserve(metadata_.ar_len); // Initialize attention mask with current position @@ -56,11 +59,12 @@ void LhdTokenGenerator::init_attention_mask(int32_t n_past) { } } - kv_manager_->init_attention_mask( - attention_mask_.data, attention_map, metadata_.ar_len, n_past); + this->kv_manager_->init_attention_mask( + this->attention_mask_.data, attention_map, metadata_.ar_len, n_past); } -void LhdTokenGenerator::init_lookahead_branch( +template +void LhdTokenGenerator::init_lookahead_branch( const std::vector& tokens) { for (int i = 0; i < metadata_.ngram - 1; ++i) { for (int j = 0; j < metadata_.window; ++j) { @@ -77,7 +81,8 @@ void LhdTokenGenerator::init_lookahead_branch( is_lhd_branch_initialized_ = true; } -void LhdTokenGenerator::init_verification_branch(uint64_t cur_token) { +template +void LhdTokenGenerator::init_verification_branch(uint64_t cur_token) { const int g_cur = ngrams_pool_.cnt[cur_token]; v_branch_.resize(g_cur); @@ -101,7 +106,8 @@ void LhdTokenGenerator::init_verification_branch(uint64_t cur_token) { } } -void LhdTokenGenerator::update_ngrams_pool() { +template +void LhdTokenGenerator::update_ngrams_pool() { std::vector ngram(metadata_.ngram - 1); // n-gram pool generation for (int f = 0; f < metadata_.window; ++f) { @@ -154,7 +160,8 @@ void LhdTokenGenerator::update_ngrams_pool() { } } -void LhdTokenGenerator::update_lookahead_branch( +template +void LhdTokenGenerator::update_lookahead_branch( const executorch::aten::Tensor& logits_tensor) { for (int i = 0; i < metadata_.window; i++) { lhd_branch_prev_[i] = lhd_branch_[0][i]; @@ -168,11 +175,12 @@ void LhdTokenGenerator::update_lookahead_branch( for (int i = 0; i < metadata_.window; i++) { size_t sample_idx = (metadata_.ngram - 2) * metadata_.window + i; lhd_branch_[metadata_.ngram - 2][i] = - decoder_runner_->logits_to_token(logits_tensor, sample_idx); + this->decoder_runner_->logits_to_token(logits_tensor, sample_idx); } } -Result LhdTokenGenerator::generate( +template +Result LhdTokenGenerator::generate( std::vector tokens, int64_t start_pos, int32_t seq_len, @@ -197,7 +205,7 @@ Result LhdTokenGenerator::generate( input_pos.reserve(metadata_.ar_len); // Rearrange KV cache first and initialize the input and output of KV cache - kv_manager_->rearrange_cache(metadata_.ar_len); + this->kv_manager_->rearrange_cache(metadata_.ar_len); // Initialize attention mask with pos init_attention_mask(pos); @@ -210,10 +218,11 @@ Result LhdTokenGenerator::generate( // Initialize the output of the module ET_CHECK_MSG( - decoder_runner_->set_outputs(method_name_, output_tensors_) == + this->decoder_runner_->set_outputs( + this->method_name_, this->output_tensors_) == executorch::runtime::Error::Ok, "Failed to set output tensor for module %s", - method_name_.c_str()); + this->method_name_.c_str()); // Generate tokens while (pos < seq_len - 1) { @@ -252,25 +261,27 @@ Result LhdTokenGenerator::generate( prepare_io(input_tokens, input_pos); // Only update data pointer of the cache to the tensor for SHIFT_POINTER // mode - bool updated = kv_manager_->update_cache_tensor( - k_cache_in_, - k_cache_out_, - v_cache_in_, - v_cache_out_, + bool updated = this->kv_manager_->update_cache_tensor( + this->k_cache_in_, + this->k_cache_out_, + this->v_cache_in_, + this->v_cache_out_, metadata_.ar_len, pos); // Only update the output of module for SHIFT_POINTER mode if (updated) { // Update the output of the module ET_CHECK_MSG( - decoder_runner_->set_outputs(method_name_, output_tensors_) == + this->decoder_runner_->set_outputs( + this->method_name_, this->output_tensors_) == executorch::runtime::Error::Ok, "Failed to set output tensor for module %s", - method_name_.c_str()); + this->method_name_.c_str()); } // Run inference - auto logits_res = decoder_runner_->step(method_name_, inputs_); + auto logits_res = + this->decoder_runner_->step(this->method_name_, this->inputs_); ET_CHECK_OK_OR_RETURN_ERROR(logits_res.error()); executorch::aten::Tensor& logits_tensor = logits_res.get(); prev_pos = pos; @@ -313,18 +324,19 @@ Result LhdTokenGenerator::generate( prev_token = cur_token; // sampler from logits all - stats_->on_sampling_begin(); - cur_token = decoder_runner_->logits_to_token(logits_tensor, sample_idx); - stats_->on_sampling_end(); + this->stats_->on_sampling_begin(); + cur_token = + this->decoder_runner_->logits_to_token(logits_tensor, sample_idx); + this->stats_->on_sampling_end(); result_tokens.push_back(cur_token); pos++; // print the token as string, decode it with the Tokenizer object token_callback( - ET_UNWRAP_TOKENIZER(tokenizer_->decode(prev_token, cur_token))); + ET_UNWRAP_TOKENIZER(this->tokenizer_->decode(prev_token, cur_token))); // data-dependent terminating condition: we have n_eos_ number of EOS - if (eos_ids_->count(cur_token) > 0) { + if (this->eos_ids_->count(cur_token) > 0) { printf("\n"); ET_LOG(Info, "\nReached to the end of generation"); break; @@ -360,14 +372,15 @@ Result LhdTokenGenerator::generate( } // Update KV Cache with the output results int32_t n_update = pos - prev_pos; - kv_manager_->update_cache(metadata_.ar_len, prev_pos, n_update, selected); + this->kv_manager_->update_cache( + metadata_.ar_len, prev_pos, n_update, selected); // Update attention mask with current position - kv_manager_->update_attention_mask( - attention_mask_.data, metadata_.ar_len, prev_pos, n_update); + this->kv_manager_->update_attention_mask( + this->attention_mask_.data, metadata_.ar_len, prev_pos, n_update); // data-dependent terminating condition: we have n_eos_ number of EOS - if (eos_ids_->count(cur_token) > 0) { + if (this->eos_ids_->count(cur_token) > 0) { printf("\n"); ET_LOG(Info, "\nReached to the end of generation"); break; @@ -381,4 +394,9 @@ Result LhdTokenGenerator::generate( return pos - start_pos; } + +// Explicit instantiations +template class LhdTokenGenerator; +template class LhdTokenGenerator; + } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.h b/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.h index fde50972f06..fe5e4b49230 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.h +++ b/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.h @@ -15,7 +15,8 @@ namespace example { * @brief Class for generating the token using decoder and key-value manager * with lookahead decoding. */ -class LhdTokenGenerator : public TokenGenerator { +template +class LhdTokenGenerator : public TokenGenerator { public: struct Metadata { int32_t context_len; @@ -27,28 +28,30 @@ class LhdTokenGenerator : public TokenGenerator { int32_t ngram; int32_t window; int32_t gcap; + int sliding_window; }; LhdTokenGenerator( tokenizers::Tokenizer* tokenizer, DecoderRunner* decoder_runner, - KVManager* kv_manager, + KVManager* kv_manager, const std::string& forward_name, std::unique_ptr>&& eos_ids, Metadata metadata, executorch::llm::Stats* stats) - : TokenGenerator( + : TokenGenerator( tokenizer, decoder_runner, kv_manager, forward_name, std::move(eos_ids), - TokenGenerator::Metadata{ + typename TokenGenerator::Metadata{ metadata.context_len, metadata.num_heads, metadata.num_layers, metadata.ar_len, metadata.vocab_size, - metadata.use_int64_token}, + metadata.use_int64_token, + metadata.sliding_window}, stats), metadata_(metadata), lhd_branch_(metadata.ngram - 1, std::vector(metadata.window)), diff --git a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp index 8794a1651da..73da764b584 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp @@ -14,9 +14,11 @@ using executorch::runtime::Result; using executorch::runtime::TensorInfo; namespace example { -PromptProcessor::PromptProcessor( + +template +PromptProcessor::PromptProcessor( DecoderRunner* decoder_runner, - KVManager* kv_manager, + KVManager* kv_manager, const std::string& method_name, Metadata metadata) : decoder_runner_(decoder_runner), @@ -33,17 +35,35 @@ PromptProcessor::PromptProcessor( input_pos_.size = 0; else input_pos_.size = metadata_.ar_len * sizeof(int32_t); - attention_mask_.size = - metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + + switch (metadata_.cache_mode) { + case CacheMode::StaticCahce: + attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + window_attention_mask_.size = 0; + break; + case CacheMode::HybridCache: + attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + window_attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + break; + default: + ET_CHECK_MSG(false, "Unsupported llama cache mode"); + break; + } + logits_.size = metadata_.ar_len * metadata_.vocab_size * sizeof(uint16_t); }; -void PromptProcessor::init_io( +template +void PromptProcessor::init_io( IMemAlloc* buffer_manager, Result method_meta) { + size_t idx = 0; input_tensors_.reserve(method_meta->num_inputs()); output_tensors_.reserve(method_meta->num_outputs()); // [I]: input_tokens - Result input_toks = method_meta->input_tensor_meta(0); + Result input_toks = method_meta->input_tensor_meta(idx++); input_toks_.data = reinterpret_cast(buffer_manager->allocate(input_toks_.size)); input_toks_.tensor = std::make_unique( @@ -57,7 +77,7 @@ void PromptProcessor::init_io( input_toks_.data, input_toks_.size, input_toks.get()); // [I]: attention_mask - Result attention_mask = method_meta->input_tensor_meta(1); + Result attention_mask = method_meta->input_tensor_meta(idx++); attention_mask_.data = reinterpret_cast( buffer_manager->allocate(attention_mask_.size)); attention_mask_.tensor = std::make_unique( @@ -71,9 +91,30 @@ void PromptProcessor::init_io( buffer_manager->add_memory_info( attention_mask_.data, attention_mask_.size, attention_mask.get()); + // [I]: sliding window attention_mask + if (metadata_.cache_mode == CacheMode::HybridCache) { + Result window_attention_mask = + method_meta->input_tensor_meta(idx++); + window_attention_mask_.data = reinterpret_cast( + buffer_manager->allocate(window_attention_mask_.size)); + window_attention_mask_.tensor = std::make_unique( + window_attention_mask->scalar_type(), + window_attention_mask->sizes().size(), + const_cast( + window_attention_mask->sizes().data()), + window_attention_mask_.data, + const_cast( + window_attention_mask->dim_order().data())); + input_tensors_.emplace_back(window_attention_mask_.tensor.get()); + buffer_manager->add_memory_info( + window_attention_mask_.data, + window_attention_mask_.size, + window_attention_mask.get()); + } + if (!is_bert()) { // [I]: input_pos - Result input_pos = method_meta->input_tensor_meta(2); + Result input_pos = method_meta->input_tensor_meta(idx++); input_pos_.data = reinterpret_cast(buffer_manager->allocate(input_pos_.size)); input_pos_.tensor = std::make_unique( @@ -87,18 +128,18 @@ void PromptProcessor::init_io( input_pos_.data, input_pos_.size, input_pos.get()); // [I] kv_cache - int index = 3; // bypass input_tokens, atten_mask, input_pos + size_t index = idx; // bypass input_tokens, atten_mask, input_pos for (int cache_group = 0; cache_group < 2; ++cache_group) { std::vector>>& cache = (cache_group == 0 ? k_cache_in_ : v_cache_in_); - std::vector> cache_ptrs = (cache_group == 0) + std::vector>> cache_ptrs = (cache_group == 0) ? kv_manager_->get_k_cache_() : kv_manager_->get_v_cache_(); for (int layer = 0; layer < metadata_.num_layers; ++layer) { for (int head = 0; head < metadata_.num_heads; ++head, ++index) { Result kv_cache = method_meta->input_tensor_meta(index); - uint8_t* cache_ptr = cache_ptrs[layer][head].buffer; + T* cache_ptr = cache_ptrs[layer][head].buffer; cache[layer].emplace_back(std::make_unique( kv_cache->scalar_type(), @@ -129,17 +170,17 @@ void PromptProcessor::init_io( buffer_manager->add_memory_info(logits_.data, logits_.size, logits.get()); // [O] kv_cache - int index = 1; + size_t index = 1; for (int cache_group = 0; cache_group < 2; ++cache_group) { std::vector>>& cache = (cache_group == 0 ? k_cache_out_ : v_cache_out_); - std::vector> cache_ptrs = (cache_group == 0) + std::vector>> cache_ptrs = (cache_group == 0) ? kv_manager_->get_k_cache_() : kv_manager_->get_v_cache_(); for (int layer = 0; layer < metadata_.num_layers; ++layer) { for (int head = 0; head < metadata_.num_heads; ++head, ++index) { Result kv_cache = method_meta->output_tensor_meta(index); - uint8_t* cache_ptr = cache_ptrs[layer][head].output_buffer; + T* cache_ptr = cache_ptrs[layer][head].output_buffer; cache[layer].emplace_back(std::make_unique( kv_cache->scalar_type(), kv_cache->sizes().size(), @@ -160,11 +201,13 @@ void PromptProcessor::init_io( } } -const std::vector& PromptProcessor::get_all_logits() { +template +const std::vector& PromptProcessor::get_all_logits() { return prompt_all_logits_; } -void PromptProcessor::prepare_io( +template +void PromptProcessor::prepare_io( const std::vector& prompt_tokens, int64_t prompt_pos, int64_t start_pos) { @@ -189,7 +232,8 @@ void PromptProcessor::prepare_io( } } -Result PromptProcessor::prefill( +template +Result PromptProcessor::prefill( std::vector prompt_tokens, int64_t start_pos, bool dump_logits) { @@ -227,6 +271,16 @@ Result PromptProcessor::prefill( // Initialize attention mask with current position kv_manager_->init_attention_mask( attention_mask_.data, attention_map, metadata_.ar_len, pos); + // Initialize window attention mask with current position + if (metadata_.cache_mode == CacheMode::HybridCache) { + kv_manager_->init_attention_mask( + window_attention_mask_.data, + attention_map, + metadata_.ar_len, + pos, + metadata_.sliding_window); + } + // Initialize the output of the module ET_CHECK_MSG( decoder_runner_->set_outputs(method_name_, output_tensors_) == @@ -268,9 +322,18 @@ Result PromptProcessor::prefill( } // Update KV Cache with the output results kv_manager_->update_cache(metadata_.ar_len, pos, n_update, {}); + // Update attention mask with current position kv_manager_->update_attention_mask( attention_mask_.data, metadata_.ar_len, pos, n_update); + if (metadata_.cache_mode == CacheMode::HybridCache) { + kv_manager_->update_attention_mask( + window_attention_mask_.data, + metadata_.ar_len, + pos, + n_update, + metadata_.sliding_window); + } prompt_pos += metadata_.ar_len; pos += metadata_.ar_len; } @@ -281,4 +344,8 @@ Result PromptProcessor::prefill( return cur_token; } +// Explicit instantiations +template class PromptProcessor; +template class PromptProcessor; + } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h index 244e26577e9..a3dd2079461 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h +++ b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h @@ -7,6 +7,7 @@ */ #pragma once +#include #include #include #include @@ -19,6 +20,7 @@ namespace example { * @class PromptProcessor * @brief Class for processing prompts using decoder and key-value manager. */ +template class PromptProcessor { public: struct Metadata { @@ -28,10 +30,12 @@ class PromptProcessor { int32_t ar_len; int32_t vocab_size; bool use_int64_token; + int sliding_window; + CacheMode cache_mode; }; PromptProcessor( DecoderRunner* decoder_runner, - KVManager* kv_manager, + KVManager* kv_manager, const std::string& method_name, Metadata metadata); @@ -71,8 +75,13 @@ class PromptProcessor { * @return Total I/O size in bytes. */ inline const size_t total_prompt_processor_io_size_in_bytes() const { - return input_toks_.size + input_pos_.size + attention_mask_.size + - logits_.size; + if (metadata_.cache_mode == CacheMode::HybridCache) { + return input_toks_.size + input_pos_.size + attention_mask_.size + + window_attention_mask_.size + logits_.size; + } else { + return input_toks_.size + input_pos_.size + attention_mask_.size + + logits_.size; + } } private: @@ -92,7 +101,7 @@ class PromptProcessor { int64_t prompt_pos, int64_t start_pos); DecoderRunner* decoder_runner_; - KVManager* kv_manager_; + KVManager* kv_manager_; std::string method_name_; // metadata @@ -102,6 +111,7 @@ class PromptProcessor { TensorStruct input_toks_; TensorStruct input_pos_; TensorStruct attention_mask_; + TensorStruct window_attention_mask_; TensorStruct logits_; // layer -> head -> TensorImpl diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp index 9c61863bc9d..0c9be4d441d 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp @@ -9,6 +9,7 @@ // A llama 3.2 runner that includes preprocessing and post processing // logic. The module takes in a string as input and emits a string as output. +#include #include #include #include @@ -21,7 +22,6 @@ #include #include #include - #include #include @@ -59,7 +59,7 @@ void print_performance_report( outfile << num_tok; outfile.close(); } else { - ET_CHECK_MSG(false, "Error saving the inference speed file"); + ET_LOG(Error, "Error saving the inference speed file"); } } @@ -84,14 +84,9 @@ void save_logits( } // namespace -std::unique_ptr<::tokenizers::Tokenizer> load_llama_tokenizer( - const std::string& tokenizer_path, - Version version) { - auto special_tokens = get_special_tokens(version); - return llm::load_tokenizer(tokenizer_path, std::move(special_tokens)); -} - -Runner::Runner( +template +Runner::Runner( + std::unique_ptr module, const std::string& decoder_model_version, const std::string& model_path, const std::string& tokenizer_path, @@ -104,7 +99,8 @@ Runner::Runner( const int window, const int gcap, std::unique_ptr tokenizer) - : ngram_(ngram), + : module_(std::move(module)), + ngram_(ngram), window_(window), gcap_(gcap), tokenizer_path_(tokenizer_path), @@ -113,8 +109,6 @@ Runner::Runner( temperature_(temperature), eval_mode_(static_cast(eval_mode)), tokenizer_(std::move(tokenizer)) { - module_ = std::make_unique( - model_path, Module::LoadMode::MmapUseMlockIgnoreErrors); stats_.reset(); if (kv_updater == "SmartMask") { kv_updater_ = KVManagerMode::SMART_MASK; @@ -128,8 +122,19 @@ Runner::Runner( decoder_model_version_ = DecoderModelVersion::kLlama2; } else if (decoder_model_version == "llama3") { decoder_model_version_ = DecoderModelVersion::kLlama3; + } else if (decoder_model_version == "gemma3") { + decoder_model_version_ = DecoderModelVersion::kGemma3; + cache_mode_ = CacheMode::HybridCache; + } else if (decoder_model_version == "phi_4_mini") { + decoder_model_version_ = DecoderModelVersion::kPhi4; } else if (decoder_model_version == "qwen2_5") { decoder_model_version_ = DecoderModelVersion::kQwen2_5; + } else if (decoder_model_version == "qwen3") { + decoder_model_version_ = DecoderModelVersion::kQwen3; + } else if (decoder_model_version == "smollm2_135m") { + decoder_model_version_ = DecoderModelVersion::kSmollm2_135m; + } else if (decoder_model_version == "smollm3") { + decoder_model_version_ = DecoderModelVersion::kSmollm3; } else { ET_CHECK_MSG(false, "Unsupported Decoder Model"); } @@ -140,12 +145,14 @@ Runner::Runner( ET_LOG(Info, "kv updater=%s", kv_updater.c_str()); } -bool Runner::is_loaded() const { +template +bool Runner::is_loaded() const { return module_->is_loaded() && tokenizer_ && decoder_runner_ && prompt_processor_ && token_generator_ && kv_manager_ && buffer_manager_; } -Error Runner::load() { +template +Error Runner::load() { if (is_loaded()) { return Error::Ok; } @@ -169,19 +176,34 @@ Error Runner::load() { ET_CHECK_MSG(false, "Unsupported llama evaluation mode"); break; } - - tokenizer_ = load_llama_tokenizer(tokenizer_path_, Version::Default); - if (tokenizer_ == nullptr) { - ET_LOG(Error, "Failed to load tokenizer with %s", tokenizer_path_.c_str()); - return Error::Internal; + auto eos_ids = std::make_unique>(); + if (tokenizer_ != nullptr) { + eos_ids->insert(tokenizer_->encode("<|eot_id|>", 0, 0).get()[0]); + eos_ids->insert(tokenizer_->encode("<|eot|>", 0, 0).get()[0]); + eos_ids->insert(tokenizer_->encode("<|end_of_text|>", 0, 0).get()[0]); + } else { + tokenizer_ = + example::load_llama_tokenizer(tokenizer_path_, Version::Default); + if (tokenizer_ == nullptr) { + ET_LOG( + Error, "Failed to load tokenizer with %s", tokenizer_path_.c_str()); + return Error::Internal; + } + eos_ids->insert(tokenizer_->eos_tok()); } - - auto eos_ids = std::make_unique>( - std::unordered_set{tokenizer_->eos_tok()}); - if (decoder_model_version_ == DecoderModelVersion::kLlama3) { eos_ids->insert(tokenizer_->encode("<|eot_id|>", 0, 0).get()[0]); + } else if (decoder_model_version_ == DecoderModelVersion::kPhi4) { + eos_ids->insert(tokenizer_->encode("<|end|>", 0, 0).get()[0]); + } else if ( + decoder_model_version_ == DecoderModelVersion::kQwen3 || + decoder_model_version_ == DecoderModelVersion::kSmollm2_135m || + decoder_model_version_ == DecoderModelVersion::kSmollm3) { + eos_ids->insert(tokenizer_->encode("<|im_end|>", 0, 0).get()[0]); + } else if (decoder_model_version_ == DecoderModelVersion::kGemma3) { + eos_ids->insert(tokenizer_->encode("", 0, 0).get()[0]); } + // Try avoid getMetadataHelper as it is time consuming. Result method_meta = module_->method_meta(token_generator_method_name); @@ -195,10 +217,10 @@ Error Runner::load() { ET_CHECK_OK_OR_RETURN_ERROR(decoder_runner_->load(method_names)); ET_LOG(Info, "Reading metadata from model"); - // retrieve any method meta, can be either prefill or kv int64_t num_layers = ET_UNWRAP(module_->get("get_n_layers")).toScalar().to(); + ET_CHECK_MSG(num_layers != -1, "Could not retrieve num layers"); // k_cache: [1, head_dim, seq_len] int64_t head_dim = method_meta->output_tensor_meta(1)->sizes()[1]; @@ -233,9 +255,16 @@ Error Runner::load() { std::min(token_generator_ar_len, prompt_processor_ar_len); max_ar_len = std::max(token_generator_ar_len, prompt_processor_ar_len); - kv_manager_ = std::make_unique( + // Load the sliding window size if the model supports it. + // This is used to configure the attention mask for models with window + // attention + int32_t sliding_window = context_len_; + if (module_->method_names()->count("get_sliding_window") > 0) { + sliding_window = ET_UNWRAP(module_->get("get_sliding_window")).toInt(); + } + kv_manager_ = std::make_unique>( kv_updater_, - KVManager::Metadata{ + typename KVManager::Metadata{ context_len_, head_dim, max_ar_len, @@ -243,25 +272,33 @@ Error Runner::load() { num_heads, num_layers}); - prompt_processor_ = std::make_unique( + prompt_processor_ = std::make_unique>( decoder_runner_.get(), kv_manager_.get(), prompt_processor_method_name, - PromptProcessor::Metadata{ + typename PromptProcessor::Metadata{ context_len_, num_heads, num_layers, prompt_processor_ar_len, vocab_size, - use_int64_token}); + use_int64_token, + sliding_window, + cache_mode_}); if (eval_mode_ == EvalMode::kLookaheadDecoding) { - token_generator_ = std::make_unique( + // TODO: sliding window attention will be supported in future. + if (sliding_window < context_len_) { + ET_CHECK_MSG( + false, + "Lookahead decoding (eval_mode == 2) is not yet supported for sliding window attention."); + } + token_generator_ = std::make_unique>( tokenizer_.get(), decoder_runner_.get(), kv_manager_.get(), token_generator_method_name, std::move(eos_ids), - LhdTokenGenerator::Metadata{ + typename LhdTokenGenerator::Metadata{ context_len_, num_heads, num_layers, @@ -270,22 +307,25 @@ Error Runner::load() { use_int64_token, ngram_, window_, - gcap_}, + gcap_, + sliding_window}, &stats_); } else { - token_generator_ = std::make_unique( + token_generator_ = std::make_unique>( tokenizer_.get(), decoder_runner_.get(), kv_manager_.get(), token_generator_method_name, std::move(eos_ids), - TokenGenerator::Metadata{ + typename TokenGenerator::Metadata{ context_len_, num_heads, num_layers, token_generator_ar_len, vocab_size, - use_int64_token}, + use_int64_token, + sliding_window, + cache_mode_}, &stats_); } @@ -308,14 +348,34 @@ Error Runner::load() { return Error::Ok; } -Error Runner::generate( +template +Error Runner::generate( + const std::string& prompt, + const llm::GenerationConfig& config, + std::function token_callback, + std::function stats_callback) { + return generate_from_pos(prompt, 0, config, token_callback, stats_callback); +} + +template +Error Runner::generate_from_pos( + const std::string& prompt, + int64_t start_pos, + const llm::GenerationConfig& config, + std::function token_callback, + std::function stats_callback) { + // TODO: currently only support start_pos == 0 + return generate_from_prompt_or_file( + prompt, false, config, token_callback, stats_callback); +} + +template +Error Runner::generate_from_prompt_or_file( const std::string& prompt, bool tokenized_prompt, - int32_t seq_len, + const llm::GenerationConfig& config, std::function token_callback, - std::function stats_callback, - bool echo, - bool warming) { + std::function stats_callback) { ET_CHECK_MSG(!prompt.empty(), "prompt cannot be null"); if (!is_loaded()) { stats_.model_load_start_ms = time_in_ms(); @@ -324,6 +384,7 @@ Error Runner::generate( } stats_.inference_start_ms = time_in_ms(); + int32_t seq_len = config.seq_len; seq_len = (seq_len > 0 && seq_len <= context_len_) ? seq_len : context_len_; int32_t n_bos = (cur_pos_ == 0) ? 1 : 0; @@ -362,7 +423,7 @@ Error Runner::generate( "sequence length exceeded - please increase the seq_len value"); // Prompt Processor first - if (token_callback) { + if (token_callback && config.echo) { token_callback(prompt); } bool dump_logits = dump_logits_path_.empty() ? false : true; @@ -414,7 +475,8 @@ Error Runner::generate( return Error::Ok; } -Result Runner::get_decoder_model_version() { +template +Result Runner::get_decoder_model_version() { if (!is_loaded()) { stats_.model_load_start_ms = time_in_ms(); ET_CHECK_OK_OR_RETURN_ERROR(load()); @@ -423,4 +485,8 @@ Result Runner::get_decoder_model_version() { return decoder_model_version_; } +// Explicit instantiations +template class Runner; +template class Runner; + } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.h b/examples/qualcomm/oss_scripts/llama/runner/runner.h index 1205bcb0eed..30fba71ecef 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/runner.h +++ b/examples/qualcomm/oss_scripts/llama/runner/runner.h @@ -16,11 +16,13 @@ #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -30,36 +32,60 @@ namespace example { enum DecoderModelVersion { kLlama2 = 0, kLlama3, + kGemma3, + kPhi4, kQwen2_5, + kQwen3, + kSmollm2_135m, + kSmollm3 }; -class Runner { + +enum KvBitWidth { + kWidth8 = 8, + kWidth16 = 16, +}; + +template +class Runner : public executorch::extension::llm::IRunner { public: explicit Runner( + std::unique_ptr module, const std::string& decoder_model, const std::string& model_path, const std::string& tokenizer_path, const std::string& performance_output_path, const std::string& dump_logits_path, const float temperature = 0.8f, - const int eval_mode = EvalMode::kKVCached, + const int eval_mode = EvalMode::kHybrid, const std::string& kv_updater = "SmartMask", const int ngram = 0, const int window = 0, const int gcap = 0, std::unique_ptr tokenizer = nullptr); - bool is_loaded() const; - executorch::runtime::Error load(); + bool is_loaded() const override; + executorch::runtime::Error load() override; // TODO: Support echo and warming executorch::runtime::Error generate( + const std::string& prompt, + const executorch::extension::llm::GenerationConfig& config, + std::function token_callback = {}, + std::function stats_callback = {}) + override; + executorch::runtime::Error generate_from_pos( + const std::string& prompt, + int64_t start_pos, + const executorch::extension::llm::GenerationConfig& config, + std::function token_callback = {}, + std::function stats_callback = {}) + override; + executorch::runtime::Error generate_from_prompt_or_file( const std::string& prompt, bool tokenized_prompt, - int32_t seq_len, + const executorch::extension::llm::GenerationConfig& config, std::function token_callback = {}, - std::function stats_callback = {}, - bool echo = true, - bool warming = false); - void stop() {}; + std::function stats_callback = {}); + void stop() override {}; executorch::runtime::Result get_decoder_model_version(); private: @@ -76,6 +102,10 @@ class Runner { int ngram_{0}; int window_{0}; int gcap_{0}; + + // Defaults to StaticCahce, indicating that the model does not use a + // global/local architecture. + CacheMode cache_mode_{CacheMode::StaticCahce}; int64_t cur_pos_{0}; std::string tokenizer_path_; @@ -83,14 +113,15 @@ class Runner { std::string dump_logits_path_; float temperature_; EvalMode eval_mode_; + DecoderModelVersion decoder_model_version_; KVManagerMode kv_updater_; std::unique_ptr buffer_manager_; - std::unique_ptr kv_manager_; + std::unique_ptr> kv_manager_; std::unique_ptr tokenizer_; std::unique_ptr decoder_runner_; - std::unique_ptr prompt_processor_; - std::unique_ptr token_generator_; + std::unique_ptr> prompt_processor_; + std::unique_ptr> token_generator_; // stats executorch::llm::Stats stats_; diff --git a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp index bacff94f594..6775c08bd87 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp @@ -14,10 +14,11 @@ using executorch::runtime::Result; using executorch::runtime::TensorInfo; namespace example { -TokenGenerator::TokenGenerator( +template +TokenGenerator::TokenGenerator( tokenizers::Tokenizer* tokenizer, DecoderRunner* decoder_runner, - KVManager* kv_manager, + KVManager* kv_manager, const std::string& method_name, std::unique_ptr>&& eos_ids, Metadata metadata, @@ -39,15 +40,35 @@ TokenGenerator::TokenGenerator( input_pos_.size = metadata_.ar_len * sizeof(int32_t); attention_mask_.size = metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + + switch (metadata_.cache_mode) { + case CacheMode::StaticCahce: + attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + window_attention_mask_.size = 0; + break; + case CacheMode::HybridCache: + attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + window_attention_mask_.size = + metadata_.ar_len * metadata_.context_len * sizeof(uint16_t); + break; + default: + ET_CHECK_MSG(false, "Unsupported llama cache mode"); + break; + } + logits_.size = metadata_.ar_len * metadata_.vocab_size * sizeof(uint16_t); } -void TokenGenerator::init_io( +template +void TokenGenerator::init_io( IMemAlloc* buffer_manager, Result method_meta) { + size_t idx = 0; input_tensors_.reserve(method_meta->num_inputs()); output_tensors_.reserve(method_meta->num_outputs()); // [I]: input_tokens - Result input_toks = method_meta->input_tensor_meta(0); + Result input_toks = method_meta->input_tensor_meta(idx++); input_toks_.data = reinterpret_cast(buffer_manager->allocate(input_toks_.size)); input_toks_.tensor = std::make_unique( @@ -61,7 +82,7 @@ void TokenGenerator::init_io( input_toks_.data, input_toks_.size, input_toks.get()); // [I]: attention_mask - Result attention_mask = method_meta->input_tensor_meta(1); + Result attention_mask = method_meta->input_tensor_meta(idx++); attention_mask_.data = reinterpret_cast( buffer_manager->allocate(attention_mask_.size)); attention_mask_.tensor = std::make_unique( @@ -75,8 +96,29 @@ void TokenGenerator::init_io( buffer_manager->add_memory_info( attention_mask_.data, attention_mask_.size, attention_mask.get()); + // [I]: sliding window attention_mask + if (metadata_.cache_mode == CacheMode::HybridCache) { + Result window_attention_mask = + method_meta->input_tensor_meta(idx++); + window_attention_mask_.data = reinterpret_cast( + buffer_manager->allocate(window_attention_mask_.size)); + window_attention_mask_.tensor = std::make_unique( + window_attention_mask->scalar_type(), + window_attention_mask->sizes().size(), + const_cast( + window_attention_mask->sizes().data()), + window_attention_mask_.data, + const_cast( + window_attention_mask->dim_order().data())); + input_tensors_.emplace_back(window_attention_mask_.tensor.get()); + buffer_manager->add_memory_info( + window_attention_mask_.data, + window_attention_mask_.size, + window_attention_mask.get()); + } + // [I]: input_pos - Result input_pos = method_meta->input_tensor_meta(2); + Result input_pos = method_meta->input_tensor_meta(idx++); input_pos_.data = reinterpret_cast(buffer_manager->allocate(input_pos_.size)); input_pos_.tensor = std::make_unique( @@ -90,18 +132,18 @@ void TokenGenerator::init_io( input_pos_.data, input_pos_.size, input_pos.get()); // [I] kv_cache - int index = 3; // bypass input_tokens, atten_mask, input_pos + size_t index = idx; // bypass input_tokens, atten_mask, input_pos for (int cache_group = 0; cache_group < 2; ++cache_group) { std::vector>>& cache = (cache_group == 0 ? k_cache_in_ : v_cache_in_); - std::vector> cache_ptrs = (cache_group == 0) + std::vector>> cache_ptrs = (cache_group == 0) ? kv_manager_->get_k_cache_() : kv_manager_->get_v_cache_(); for (int layer = 0; layer < metadata_.num_layers; ++layer) { for (int head = 0; head < metadata_.num_heads; ++head, ++index) { Result kv_cache = method_meta->input_tensor_meta(index); - uint8_t* cache_ptr = cache_ptrs[layer][head].buffer; + T* cache_ptr = cache_ptrs[layer][head].buffer; cache[layer].emplace_back(std::make_unique( kv_cache->scalar_type(), @@ -135,13 +177,13 @@ void TokenGenerator::init_io( for (int cache_group = 0; cache_group < 2; ++cache_group) { std::vector>>& cache = (cache_group == 0 ? k_cache_out_ : v_cache_out_); - std::vector> cache_ptrs = (cache_group == 0) + std::vector>> cache_ptrs = (cache_group == 0) ? kv_manager_->get_k_cache_() : kv_manager_->get_v_cache_(); for (int layer = 0; layer < metadata_.num_layers; ++layer) { for (int head = 0; head < metadata_.num_heads; ++head, ++index) { Result kv_cache = method_meta->output_tensor_meta(index); - uint8_t* cache_ptr = cache_ptrs[layer][head].output_buffer; + T* cache_ptr = cache_ptrs[layer][head].output_buffer; cache[layer].emplace_back(std::make_unique( kv_cache->scalar_type(), kv_cache->sizes().size(), @@ -162,12 +204,14 @@ void TokenGenerator::init_io( } } -const std::vector& TokenGenerator::get_all_logits() { +template +const std::vector& TokenGenerator::get_all_logits() { return token_all_logits_; } // This function only considers the case where token_generator_ar_len equals 1. -void TokenGenerator::prepare_io(uint64_t cur_token, int64_t start_pos) { +template +void TokenGenerator::prepare_io(uint64_t cur_token, int64_t start_pos) { // update input_tok *input_toks_.data = metadata_.use_int64_token ? cur_token : static_cast(cur_token); @@ -175,7 +219,8 @@ void TokenGenerator::prepare_io(uint64_t cur_token, int64_t start_pos) { *input_pos_.data = static_cast(start_pos); } -Result TokenGenerator::generate( +template +Result TokenGenerator::generate( std::vector tokens, int64_t start_pos, int32_t seq_len, @@ -192,9 +237,20 @@ Result TokenGenerator::generate( kv_manager_->rearrange_cache(metadata_.ar_len); std::vector attention_map(metadata_.ar_len); std::iota(attention_map.begin(), attention_map.end(), -1); + // Initialize attention mask with current position kv_manager_->init_attention_mask( attention_mask_.data, attention_map, metadata_.ar_len, pos); + // Initialize window attention mask with current position + if (metadata_.cache_mode == CacheMode::HybridCache) { + kv_manager_->init_attention_mask( + window_attention_mask_.data, + attention_map, + metadata_.ar_len, + pos, + metadata_.sliding_window); + } + // Initialize the output of the module ET_CHECK_MSG( decoder_runner_->set_outputs(method_name_, output_tensors_) == @@ -246,6 +302,14 @@ Result TokenGenerator::generate( // Update attention mask with current position kv_manager_->update_attention_mask( attention_mask_.data, metadata_.ar_len, pos, metadata_.ar_len); + if (metadata_.cache_mode == CacheMode::HybridCache) { + kv_manager_->update_attention_mask( + window_attention_mask_.data, + metadata_.ar_len, + pos, + metadata_.ar_len, + metadata_.sliding_window); + } pos++; // print the token as string, decode it with the Tokenizer object @@ -261,4 +325,8 @@ Result TokenGenerator::generate( } return pos - start_pos; } +// Explicit instantiations +template class TokenGenerator; +template class TokenGenerator; + } // namespace example diff --git a/examples/qualcomm/oss_scripts/llama/runner/token_generator.h b/examples/qualcomm/oss_scripts/llama/runner/token_generator.h index f76340d4d87..9f0198f3040 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/token_generator.h +++ b/examples/qualcomm/oss_scripts/llama/runner/token_generator.h @@ -7,6 +7,7 @@ */ #pragma once +#include #include #include #include @@ -20,6 +21,7 @@ namespace example { * @class TokenGenerator * @brief Class for generating the token using decoder and key-value manager. */ +template class TokenGenerator { public: struct Metadata { @@ -29,11 +31,13 @@ class TokenGenerator { int32_t ar_len; int32_t vocab_size; bool use_int64_token; + int sliding_window; + CacheMode cache_mode; }; TokenGenerator( tokenizers::Tokenizer* tokenizer, DecoderRunner* decoder_runner, - KVManager* kv_manager, + KVManager* kv_manager, const std::string& method_name, std::unique_ptr>&& eos_ids, Metadata metadata, @@ -72,14 +76,19 @@ class TokenGenerator { std::function token_callback, bool dump_logits); inline const size_t total_token_generator_io_size_in_bytes() const { - return input_toks_.size + input_pos_.size + attention_mask_.size + - logits_.size; + if (metadata_.cache_mode == CacheMode::HybridCache) { + return input_toks_.size + input_pos_.size + attention_mask_.size + + window_attention_mask_.size + logits_.size; + } else { + return input_toks_.size + input_pos_.size + attention_mask_.size + + logits_.size; + } } protected: tokenizers::Tokenizer* tokenizer_; DecoderRunner* decoder_runner_; - KVManager* kv_manager_; + KVManager* kv_manager_; std::string method_name_; std::unique_ptr> eos_ids_; @@ -87,6 +96,7 @@ class TokenGenerator { TensorStruct input_toks_; TensorStruct input_pos_; TensorStruct attention_mask_; + TensorStruct window_attention_mask_; TensorStruct logits_; // layer -> head -> TensorImpl diff --git a/examples/qualcomm/oss_scripts/llama/targets.bzl b/examples/qualcomm/oss_scripts/llama/targets.bzl index b70c8fd2f33..062edf7594c 100644 --- a/examples/qualcomm/oss_scripts/llama/targets.bzl +++ b/examples/qualcomm/oss_scripts/llama/targets.bzl @@ -29,6 +29,7 @@ def define_common_targets(): exported_deps = [ "//executorch/extension/module:module", "//executorch/extension/llm/sampler:sampler", + "//executorch/examples/models/llama/runner:runner", "//executorch/examples/models/llama/tokenizer:tiktoken", "//executorch/extension/evalue_util:print_evalue", "//executorch/backends/qualcomm/runtime:runtime", diff --git a/examples/qualcomm/oss_scripts/llm_utils/README.md b/examples/qualcomm/oss_scripts/llm_utils/README.md new file mode 100644 index 00000000000..a713d6df5de --- /dev/null +++ b/examples/qualcomm/oss_scripts/llm_utils/README.md @@ -0,0 +1,78 @@ +## Tutorial to run [eval_decoder_model_qnn.py](./eval_decoder_model_qnn.py) +This script, [`eval_decoder_model_qnn.py`](./eval_decoder_model_qnn.py), is designed to evaluate large language models (LLMs) from transformers that have been compiled into ExecuTorch Portable Executable (PTE) format for execution on Qualcomm devices. It leverages the `lm-evaluation-harness` library to perform various NLP evaluation tasks. + +> ⚠️ **Important:** Note that this script runs PTE files generated specifically for Hugging Face Transformers, such as [qwen2_5.py](../qwen2_5/qwen2_5.py), rather than [the static LLaMA version](../llama/llama.py). + +### Features: + +* Evaluates ExecuTorch PTE models on Qualcomm devices (requires ADB setup and QNN SDK). +* Integrates with `lm-evaluation-harness` for standardized LLM evaluation tasks. + +### Prerequisites + +Before running this script, ensure you have the following: + +1. **Setup ExecuTorch** Follow the [tutorial](https://pytorch.org/executorch/main/getting-started-setup) to set up ExecuTorch. +2. **Setup QNN ExecuTorch** Follow the [tutorial](https://pytorch.org/executorch/main/backends-qualcomm) to build Qualcomm AI Engine Direct Backend. +3. **`lm-evaluation-harness`:** Install the `lm-evaluation-harness` library. +4. **PTE Model:** A pre-exported ExecuTorch PTE model of your decoder LLM. +5. **Tokenizer:** A tokenizer in a format supported by `pytorch-tokenizers` (e.g., SentencePiece `.json` or `.model`, Tiktoken). + +### How to Use + +The script evaluates the model by running the PTE file on a connected Qualcomm device. + +#### Command Line Arguments: + +* `-a`, `--artifact`: (Required) Path for storing generated artifacts by this example. +* `--tokenizer_path`: (Required) Path to your tokenizer file (e.g., `tokenizer.model` or `tokenizer.json`). +* `--pte`: (Required) Path to the ExecuTorch Portable Executable (`.pte`) model file. +* `--logits_quant_attr_path`: (Optional) Path to a JSON file containing quantization attributes. This is needed if your PTE model uses tag quant I/O for logits and requires de-quantization before evaluation. +* `--max_seq_len`: (Optional, default: 128) Maximum sequence length the model can process. +* `--tasks`: (Optional, default: `["wikitext"]`) A list of `lm-evaluation-harness` tasks to evaluate. You can specify multiple tasks separated by spaces (e.g., `--tasks wikitext piqa`). +* `--limit`: (Optional) Number of samples to evaluate per task. If not set, all samples will be evaluated. +* `--num_fewshot`: (Optional) Number of examples to use in few-shot context for evaluation. +* `--model`: (Required for QNN execution) The SoC model name (e.g., `SM8550`, `SM8650`). +* `--device`: (Required for QNN execution) The ADB device ID. +* `--host`: (Required for QNN execution) The ADB host ID (usually `localhost`). +* `--build_folder`: (Optional, default: `build-android`) The build folder for ExecuTorch artifacts, relative to the current directory. + +#### Example Usage: + +```bash +python examples/qualcomm/oss_scripts/llm_utils/eval_decoder_model_qnn.py \ + --artifact ./eval_output \ + --tokenizer_path /path/to/your/tokenizer.model \ + --pte /path/to/your/model.pte \ + --model SM8550 \ + --device YOUR_DEVICE_ID \ + --host localhost \ + --tasks wikitext \ + --limit 1 \ + --max_seq_len 512 +``` + +Replace `/path/to/your/tokenizer.model`, `/path/to/your/model.pte`, and `YOUR_DEVICE_ID` with your actual paths and device ID. + +If your model's logits are quantized and require de-quantization: + +```bash +python examples/qualcomm/oss_scripts/llm_utils/eval_decoder_model_qnn.py \ + --artifact ./eval_output \ + --tokenizer_path /path/to/your/tokenizer.model \ + --pte /path/to/your/model.pte \ + --logits_quant_attr_path /path/to/your/logits_quant_attrs.json \ + --model SM8550 \ + --device YOUR_DEVICE_ID \ + --host localhost \ + --tasks wikitext \ + --limit 1 \ + --max_seq_len 512 +``` + +### Output + +The script will print the evaluation results for each specified task to the console, similar to the `lm-evaluation-harness` output format. For example: + +``` +wikitext: {'word_perplexity': ..., 'byte_perplexity': ..., 'bits_per_byte': ...} diff --git a/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py new file mode 100644 index 00000000000..f59dc548c44 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llm_utils/decoder_model_wrapper.py @@ -0,0 +1,160 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import math +from typing import Optional + +import scipy +import torch +import transformers +from transformers import GenerationConfig, PretrainedConfig + +from transformers.integrations.executorch import TorchExportableModuleForDecoderOnlyLM + +TRANSFORMERS_VERSION = "4.53.1" + + +def save_config_to_constant_methods( + config: PretrainedConfig, + generation_config: Optional[GenerationConfig] = None, + **kwargs, +): + # Initialize metadata with values from model config + metadata = { + "get_bos_id": getattr(config, "bos_token_id", None), + "get_eos_id": getattr(config, "eos_token_id", None), + "get_vocab_size": getattr(config, "vocab_size", None), + "get_max_seq_len": getattr(config, "max_position_embeddings", None), + "use_kv_cache": getattr(generation_config, "use_cache", None), + "use_sdpa_with_kv_cache": False, + } + + # Safely access fields from generation_config if it exists + if generation_config is not None: + # Check for cache_config and its attributes + cache_config = getattr(generation_config, "cache_config", None) + if cache_config is not None: + max_seq_len = getattr(cache_config, "max_cache_len", None) + if max_seq_len is not None: + metadata["get_max_seq_len"] = max_seq_len + + # Combine with any additional kwargs and filter out None values + return {k: v for k, v in {**metadata, **kwargs}.items() if v is not None} + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +@torch._dynamo.assume_constant_result +def get_transposed_hadamard_matrix(head_dim): + r3_weight = torch.tensor( + scipy.linalg.hadamard(head_dim, dtype=float) / math.sqrt(head_dim), + dtype=torch.float32, + ) + return r3_weight.transpose(0, 1) + + +def _qnn_attention( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + if getattr(module.config, "enable_spinquant_r3", False): + r3_weight = get_transposed_hadamard_matrix(module.head_dim) + query = torch.matmul(query, r3_weight) + key = torch.matmul(key, r3_weight) + + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = torch.nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query.dtype) + attn_weights = torch.nn.functional.dropout( + attn_weights, p=dropout, training=module.training + ) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +def _qnn_attention_mask( + batch_size: int, + cache_position: torch.Tensor, + kv_length: int, + **kwargs, +): + kv_arange = torch.arange(kv_length, device=cache_position.device) + reshaped_cache_position = cache_position.view(-1, 1) + + # Simplest and most efficient way to obtain a causal mask + causal_mask = kv_arange <= reshaped_cache_position + atten_mask = torch.full((causal_mask.shape[0], kv_length), torch.tensor(-65504.0)) + atten_mask = atten_mask.masked_fill(causal_mask, 0) + atten_mask = atten_mask[None, None, :, :].expand(batch_size, -1, -1, -1) + + return atten_mask + + +class QnnCausalLMExportableModule(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.config = model.config + self._metadata = save_config_to_constant_methods( + model.config, model.generation_config + ) + logging.info(f"Metadata to be recorded in PTE: {self._metadata}") + self.exportable_module = TorchExportableModuleForDecoderOnlyLM( + self.model, + max_batch_size=1, + max_cache_len=self._metadata.get("get_max_seq_len"), + ) + self._register_attention_mask_for_4_53(self.exportable_module) + + def _register_attention_mask_for_4_53(self, exportable_module: torch.nn.Module): + if transformers.__version__ >= TRANSFORMERS_VERSION: + from transformers.masking_utils import AttentionMaskInterface + from transformers.modeling_utils import AttentionInterface + + AttentionInterface.register("qnn_attention", _qnn_attention) + AttentionMaskInterface.register("qnn_attention", _qnn_attention_mask) + exportable_module.model.model.config._attn_implementation = "qnn_attention" + self._metadata.update({"use_sdpa_with_kv_cache": False}) + + def get_example_inputs(self): + example_input_ids = torch.tensor([[1]], dtype=torch.long) + example_cache_position = torch.tensor([0], dtype=torch.long) + return (example_input_ids, example_cache_position) + + def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor): + return self.exportable_module(input_ids, cache_position) + + def get_metadata(self): + return self._metadata diff --git a/examples/qualcomm/oss_scripts/llm_utils/eval_decoder_model_qnn.py b/examples/qualcomm/oss_scripts/llm_utils/eval_decoder_model_qnn.py new file mode 100644 index 00000000000..49cdc192f22 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llm_utils/eval_decoder_model_qnn.py @@ -0,0 +1,293 @@ +# 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 argparse +import json +import os + +from typing import Optional, Union + +import numpy as np + +import torch +from executorch.examples.models.llama.evaluate.eager_eval import EagerEvalWrapper + +from executorch.examples.qualcomm.utils import ( + make_output_dir, + setup_common_args_and_variables, +) + +from lm_eval.evaluator import simple_evaluate +from pytorch_tokenizers import get_tokenizer +from pytorch_tokenizers.llama2c import Llama2cTokenizer as SentencePieceTokenizer +from pytorch_tokenizers.tiktoken import TiktokenTokenizer as Tiktoken + + +def create_device_inputs(example_inputs): + # TODO: support batch inputs if necessary + input_list = "" + inputs = [] + + for index, data in enumerate(example_inputs): + inputs.append(data) + input_list += " ".join([f"input_{index}_{i}.raw" for i in range(len(data))]) + input_list += "\n" + return inputs, input_list + + +class GraphModuleCalibrationWrapper(EagerEvalWrapper): + """ + A wrapper class for calibration + """ + + def __init__( + self, + model: torch.fx.GraphModule, + tokenizer: Union[SentencePieceTokenizer, Tiktoken], + max_seq_length: Optional[int] = None, + use_kv_cache: bool = False, + generate_full_logits: bool = False, + enable_dynamic_shape: bool = True, + ): + super().__init__( + model=model, tokenizer=tokenizer, max_seq_length=max_seq_length + ) + self._model = model.to(self.device) + self._use_kv_cache = use_kv_cache + self._generate_full_logits = generate_full_logits + self._enable_dynamic_shape = enable_dynamic_shape + + def _model_call(self, inps): + if self._use_kv_cache: + if not self._enable_dynamic_shape: + # graph module exported without dynamic shape won't work with a different shape. + # And we have to do single token prefill here. + result_logits = [] + for pos in range(inps.shape[-1]): + pos_tensor = torch.tensor([pos], dtype=torch.int64) + logits = self._model(inps[:, pos : pos + 1], pos_tensor) + result_logits.append(logits) + if self._generate_full_logits: + return torch.cat(result_logits, dim=1) + else: + return torch.stack(result_logits, dim=1) + else: + pos_tensor = torch.tensor([0], dtype=torch.int64, device=self.device) + # Batch process the whole sequence. + logits = self._model(inps[:, : self._max_seq_length], pos_tensor) + return logits + + else: + return self._model(inps) + + def _model_generate(self, context, max_length, eos_token_id): + raise Exception("unimplemented") + + +class QNNRunnerEvalWrapper(EagerEvalWrapper): + """ + A wrapper class for ExecuTorch Runtime integration with the + lm-evaluation-harness library. + """ + + def __init__( + self, + model: str, + tokenizer: Union[SentencePieceTokenizer, Tiktoken], + soc_model: str, + device: str, + host: str, + max_seq_length: Optional[int] = None, + output_dir: str = ".", + quant_attrs=None, + build_folder: str = "build-android", + ): + super().__init__(None, tokenizer, max_seq_length) + import getpass + + from executorch.examples.qualcomm.utils import SimpleADB + + self._model = model + self.output_dir = output_dir + self.quant_attrs = quant_attrs + workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/meta_llama" + self.adb = SimpleADB( + qnn_sdk=os.getenv("QNN_SDK_ROOT"), + build_path=build_folder, + pte_path=model, + workspace=workspace, + device_id=device, + host_id=host, + soc_model=soc_model, + ) + self.adb.push() + + def _model_call(self, inps): + # Given inps (tokens), return the logits from a single + # forward call + + # Example: + # inps: Tensor of shape (1, N) + # logits: Tensor of shape (1, N, vocab_size) + result_logits = [] + inputs = [] + + for pos in range(self._max_seq_length): + pos_tensor = torch.tensor([pos], dtype=torch.int64) + inputs.append([inps[:, pos : pos + 1], pos_tensor]) + + inputs, input_list = create_device_inputs(inputs) + self.adb.push(inputs=inputs, input_list=input_list, init_env=False) + self.adb.execute() + output_data_folder = f"{self.output_dir}/outputs" + make_output_dir(output_data_folder) + + def post_process(): + for f in sorted( + os.listdir(output_data_folder), key=lambda f: int(f.split("_")[1]) + ): + output_tensor = None + if self.quant_attrs: + output_tensor = torch.from_numpy( + np.fromfile( + os.path.join(output_data_folder, f), dtype=np.uint16 + ).reshape(1, 1, -1) + ) + output_tensor = ( + output_tensor.to(torch.float32) - self.quant_attrs["zero_point"] + ) * self.quant_attrs["scale"] + else: + output_tensor = torch.from_numpy( + np.fromfile( + os.path.join(output_data_folder, f), dtype=np.float32 + ).reshape(1, 1, -1) + ) + + result_logits.append(output_tensor) + + self.adb.pull(output_path=self.output_dir, callback=post_process) + return torch.cat(result_logits, dim=1) + + +def gen_eval_wrapper( + args: argparse.ArgumentParser, +): + """ + Generates a wrapper interface around the provided model and tokenizer for + the lm-evaluation-harness library. + + Returns: + eval_wrapper (LM): A wrapper interface for the lm-evaluation-harness library. + """ + tokenizer = get_tokenizer(args.tokenizer_path) + + # ExecuTorch Binary Evaluation + if (model := args.pte) is not None: # pyre-ignore + assert args.device is not None, "please specify the device to execute pte" + quant_attrs = None + if args.logits_quant_attr_path is not None: + quant_attrs = json.load(open(f"{args.logits_quant_attr_path}")) + return QNNRunnerEvalWrapper( + model=model, + tokenizer=tokenizer, + soc_model=args.model, + device=args.device, + host=args.host, + max_seq_length=args.max_seq_len - 1, + output_dir=args.artifact, + quant_attrs=quant_attrs, + build_folder=args.build_folder, + ) + else: + raise RuntimeError("Currently only support evaluate pte on device") + + +def eval_llama( + args: argparse.ArgumentParser, +) -> None: + + # Generate the eval wrapper + eval_wrapper = gen_eval_wrapper(args) + + # Evaluate the model + with torch.no_grad(): + eval_results = simple_evaluate( + model=eval_wrapper, + tasks=args.tasks, + num_fewshot=args.num_fewshot, + limit=args.limit, + ) + + for task, res in eval_results["results"].items(): + print(f"{task}: {res}") + + +def main() -> None: + seed = 42 + torch.manual_seed(seed) + parser = setup_common_args_and_variables() + + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts by this example.", + type=str, + ) + + parser.add_argument( + "--tokenizer_path", + help="path to tokenizer.json.", + type=str, + ) + + parser.add_argument( + "--pte", + type=str, + default=None, + help="[For ExecuTorch] Path to the ExecuTorch model being evaluated. If provided, don't go through the export flow", + ) + parser.add_argument( + "--logits_quant_attr_path", + type=str, + default=None, + help="For the pte with tag quant io, it needs to be dequantize and compute ppl.", + ) + + parser.add_argument( + "--max_seq_len", + type=int, + default=128, + help="This refers to maximum number of tokens that the model can process & consider at once to generate predictions/responses.", + ) + parser.add_argument( + "--tasks", + nargs="+", + type=str, + default=["wikitext"], + help="list of lm-eluther tasks to evaluate usage: --tasks task1 task2", + ) + parser.add_argument( + "--limit", + type=int, + default=None, + help="number of samples to evalulate. If not set, evaluate all samples", + ) + parser.add_argument( + "--num_fewshot", + type=int, + default=None, + metavar="N", + help="Number of examples in few-shot context", + ) + + args = parser.parse_args() + + eval_llama(args) + + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py new file mode 100644 index 00000000000..2a2cd250a40 --- /dev/null +++ b/examples/qualcomm/oss_scripts/llm_utils/qnn_decoder_model_manager.py @@ -0,0 +1,331 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from functools import partial +from typing import Callable, List + +import torch +from executorch.backends.qualcomm._passes import TagQuantIO +from executorch.backends.qualcomm._passes.build_quant_io import BuildQuantIo +from executorch.backends.qualcomm._passes.qnn_pass_manager import ( + get_capture_program_passes, +) +from executorch.backends.qualcomm.builders.utils import is_graph_output +from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype +from executorch.backends.qualcomm.utils.constants import ( + QCOM_PASS_ACTIVATE_KEY, + QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY, + QCOM_QUANT_ATTRS_MAP, +) +from executorch.backends.qualcomm.utils.utils import ( + generate_htp_compiler_spec, + generate_qnn_executorch_compiler_spec, + get_soc_to_chipset_map, + to_edge_transform_and_lower_to_qnn, +) +from executorch.devtools.backend_debug import print_delegation_info +from executorch.examples.qualcomm.oss_scripts.llm_utils.decoder_model_wrapper import ( + QnnCausalLMExportableModule, +) +from executorch.examples.qualcomm.utils import make_quantizer +from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass +from pytorch_tokenizers import get_tokenizer +from torchao.quantization.pt2e import MinMaxObserver +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from transformers import AutoConfig, AutoModelForCausalLM, GenerationConfig + + +FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" +logging.basicConfig(level=logging.INFO, format=FORMAT) + +HUGGING_FACE_REPO_IDS = { + "qwen2.5_0.5B": "Qwen/Qwen2.5-0.5B", + "qwen2.5_1.5B_instruct": "Qwen/Qwen2.5-1.5B-Instruct", + "qwen2.5_0.5B_instruct": "Qwen/Qwen2.5-0.5B-Instruct", +} + + +def get_qnn_llm_edge_manager(model_name, max_seq_len=128, enable_spinquant_r3=True): + model_id = HUGGING_FACE_REPO_IDS[model_name] + config = AutoConfig.from_pretrained(model_id) + device = "cpu" + batch_size = 1 + dtype = "float32" + cache_implementation = "static" + attn_implementation = "eager" + + # Set configs + config.max_seq_len = max_seq_len + config.ar_len = 1 # kv mode + config.max_batch_size = batch_size + config.enable_spinquant_r3 = enable_spinquant_r3 + + # Some config has head_dim provided that is different from equation below(e.g., qwen3) + if not hasattr(config, "head_dim"): + config.head_dim = config.hidden_size // config.num_attention_heads + + model = AutoModelForCausalLM.from_pretrained( + model_id, + device_map=device, + torch_dtype=dtype, + config=config, + attn_implementation=attn_implementation, + generation_config=GenerationConfig( + use_cache=True, + cache_implementation=cache_implementation, + max_length=max_seq_len, + cache_config={ + "batch_size": batch_size, + "max_cache_len": max_seq_len, + }, + ), + ).eval() + model_wrapper = QnnCausalLMExportableModule(model) + + return QnnLLMEdgeManager(model_name, model_wrapper, config) + + +class QnnLLMEdgeManager: + def __init__(self, model_name, model_wrapper, config, verbose=True) -> None: + self.model_name = model_name + self.model_wrapper = model_wrapper + self.graph_module = model_wrapper + self.config = config + self.verbose = verbose + self.use_fp16 = True + self.passes_job = get_capture_program_passes() + self.edge_prog_mgr = None + self.logits_quant_attrs = None + + def source_transform( + self, transforms: List[Callable[[torch.nn.Module], torch.nn.Module]] + ) -> "QnnLLMEdgeManager": + """ + Apply source transforms to the model. The transforms are callables that + takes nn.Module as input and returns nn.Module. + Args: + transforms (List[Callable[[torch.nn.Module], torch.nn.Module]]): A + list of source transforms. + """ + for transform in transforms: + self.graph_module = transform(self.graph_module) + + if self.verbose: + logging.info(f"Applied source transforms: {transforms}") + logging.info(f"Model after source transforms: {self.graph_module}") + return self + + def _tag_ios(self, node, fixed_point_type, config): + # shape of k caches and v caches + kv_cache_shape = { + # single head, kv input + (config.head_dim, config.max_seq_len), + (config.max_seq_len, config.head_dim), + # single head, kv output + (config.head_dim, config.ar_len), + (config.ar_len, config.head_dim), + } + + logit_out_shape = { + ( + config.max_batch_size, + config.ar_len, + config.vocab_size, + ) + } + + quant_io_type = None + + if node.op == "placeholder": + if ( + len(users := list(node.users)) == 1 + and users[0].meta["val"].size()[-2:] in kv_cache_shape + ): + quant_io_type = fixed_point_type["kv_type"] + if is_graph_output(node): + if node.meta["val"].size()[-2:] in kv_cache_shape: + quant_io_type = fixed_point_type["kv_type"] + elif node.meta["val"].size() in logit_out_shape: + quant_io_type = fixed_point_type["io_type"] + + return quant_io_type + + def export(self): + with torch.no_grad(): + self.graph_module = torch.export.export( + self.graph_module, + args=self.model_wrapper.get_example_inputs(), + strict=True, + ).module() + + def pt2e_calibrate( + self, + calibration_tasks, + calibration_limit, + calibration_seq_length, + calibration_data, + tokenizer_path, + ): + try: + from executorch.examples.qualcomm.oss_scripts.llm_utils.eval_decoder_model_qnn import ( + GraphModuleCalibrationWrapper, + ) + from lm_eval.evaluator import simple_evaluate + except ImportError: + raise ImportError( + "Please install the llm eval dependency via examples/models/llama/install_requirements.sh" + ) + + tokenizer = get_tokenizer(tokenizer_path) + logging.info( + f"Calibrating with tasks: {calibration_tasks}, limit: {calibration_limit}, calibration_data: {calibration_data}, tokenizer_path: {tokenizer_path}, seq_length: {self.config.max_seq_len}" + ) + + def calibrate_template( + module: torch.fx.GraphModule, tokenizer, prompts: str, max_len: int + ): + # TODO: change criteria & support batch inputs if necessary + pos = 0 + token_list = tokenizer.encode(prompts, bos=True, eos=False) + + with torch.no_grad(): + while token_list[-1] != tokenizer.eos_id and pos < max_len: + cur_pos = torch.tensor([pos], dtype=torch.long) + logits = module(torch.full((1, 1), token_list[pos]), cur_pos) + pos += 1 + if pos >= len(token_list): + token_list.append(torch.argmax(logits, dim=-1).item()) + logging.info( + f"Result of LLM with static cache:\n {tokenizer.decode(token_list)} \n\n\n" + ) + + calibrate_template( + module=self.graph_module, + tokenizer=tokenizer, + prompts=calibration_data, + max_len=calibration_seq_length, + ) + if calibration_tasks is not None and calibration_limit is not None: + eval_wrapper = GraphModuleCalibrationWrapper( + model=self.graph_module, + tokenizer=tokenizer, + max_seq_length=calibration_seq_length, + use_kv_cache=True, + generate_full_logits=True, + enable_dynamic_shape=False, + ) + + # Evaluate the model + with torch.no_grad(): + eval_results = simple_evaluate( + model=eval_wrapper, + tasks=calibration_tasks, + limit=calibration_limit, + ) + + for task, res in eval_results["results"].items(): + print(f"{task}: {res}") + logging.info("Calibration finish...") + + def pt2e_quantize( + self, + quant_dtype, + fixed_point_type, + calibration_tasks, + calibration_limit, + calibration_data, + tokenizer_path, + ): + self.export() + + quantizer = make_quantizer( + quant_dtype=quant_dtype, + per_channel_linear=True, + per_channel_conv=True, + act_observer=MinMaxObserver, + ) + if quant_dtype == QuantDtype.use_16a4w_block: + + def extract_linear_nodes(graph): + linear_nodes = [] + for node in graph.nodes: + if node.target == torch.ops.aten.linear.default: + linear_nodes.append(node) # linear node + linear_nodes.append(node.args[1]) # weight node + return linear_nodes + + linear_nodes = extract_linear_nodes(self.graph_module.graph) + block_size_map = {n.name: (1, 16) for n in linear_nodes} + quantizer.set_block_size_map(block_size_map) + self.graph_module = prepare_pt2e(self.graph_module, quantizer) + self.pt2e_calibrate( + calibration_tasks, + calibration_limit, + self.config.max_seq_len, + calibration_data, + tokenizer_path, + ) + self.graph_module = convert_pt2e(self.graph_module) + + self.passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True + self.passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ + "get_quant_io_dtype_fn" + ] = partial( + self._tag_ios, fixed_point_type=fixed_point_type, config=self.config + ) + self.use_fp16 = False + + def to_edge_transform_and_lower_to_qnn( + self, soc_model, skip_node_id_set, skip_node_op_set + ): + backend_options = generate_htp_compiler_spec(use_fp16=self.use_fp16) + compiler_spec = generate_qnn_executorch_compiler_spec( + soc_model=get_soc_to_chipset_map()[soc_model], + backend_options=backend_options, + ) + with torch.no_grad(): + self.edge_prog_mgr = to_edge_transform_and_lower_to_qnn( + self.graph_module, + self.model_wrapper.get_example_inputs(), + compiler_spec, + constant_methods=self.model_wrapper.get_metadata(), + passes_job=self.passes_job, + skip_node_id_set=skip_node_id_set, + skip_node_op_set=skip_node_op_set, + convert_linear_to_conv2d=True, + ) + + print_delegation_info(self.edge_prog_mgr.exported_program().graph_module) + if not self.use_fp16: + logit_out_shape = { + ( + self.config.max_batch_size, + self.config.ar_len, + self.config.vocab_size, + ) + } + for n in self.edge_prog_mgr.exported_program().graph.nodes: + if n.op == "output": + for node, output_encoding in n.meta[QCOM_QUANT_ATTRS_MAP].items(): + if node.meta["val"].size() in logit_out_shape: + self.logits_quant_attrs = output_encoding + + def get_logits_quant_attrs(self): + return self.logits_quant_attrs + + def to_executorch(self, artifact, pte_filename): + executorch_config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, + ), + passes=[BuildQuantIo()], + ) + exec_prog_mgr = self.edge_prog_mgr.to_executorch(config=executorch_config) + with open(f"{artifact}/{pte_filename}.pte", "wb") as file: + exec_prog_mgr.write_to_file(file) + logging.info(f"Saved exported program to {artifact}/{pte_filename}.pte") diff --git a/examples/qualcomm/oss_scripts/mobilevit_v1.py b/examples/qualcomm/oss_scripts/mobilevit_v1.py index 99b7160f669..f3f5f8e86aa 100644 --- a/examples/qualcomm/oss_scripts/mobilevit_v1.py +++ b/examples/qualcomm/oss_scripts/mobilevit_v1.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -36,7 +37,7 @@ def get_data_loader(): ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] data_loader = get_data_loader() feature_extractor = MobileViTFeatureExtractor.from_pretrained( "apple/mobilevit-xx-small" @@ -49,9 +50,8 @@ def get_data_loader(): feature = feature_extractor(images=image, return_tensors="pt") inputs.append((feature["pixel_values"],)) targets.append(torch.tensor(target)) - input_list += f"input_{index}_0.raw\n" - return inputs, targets, input_list + return inputs, targets def main(args): @@ -60,12 +60,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -73,7 +67,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, ) @@ -110,7 +104,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -163,6 +157,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/mobilevit_v2.py b/examples/qualcomm/oss_scripts/mobilevit_v2.py index 70a233a7988..64042efa212 100644 --- a/examples/qualcomm/oss_scripts/mobilevit_v2.py +++ b/examples/qualcomm/oss_scripts/mobilevit_v2.py @@ -7,6 +7,7 @@ import json import logging import os + import warnings from multiprocessing.connection import Client @@ -37,7 +38,7 @@ def get_data_loader(): ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] data_loader = get_data_loader() feature_extractor = MobileViTFeatureExtractor.from_pretrained( "apple/mobilevit-xx-small" @@ -50,9 +51,8 @@ def get_data_loader(): feature = feature_extractor(images=image, return_tensors="pt") inputs.append((feature["pixel_values"],)) targets.append(torch.tensor(target)) - input_list += f"input_{index}_0.raw\n" - return inputs, targets, input_list + return inputs, targets def main(args): @@ -66,12 +66,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 256, 256),)] @@ -79,7 +73,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, ) @@ -118,7 +112,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -170,6 +164,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/moshi/CMakeLists.txt b/examples/qualcomm/oss_scripts/moshi/CMakeLists.txt index 70356e54906..0853866c50b 100644 --- a/examples/qualcomm/oss_scripts/moshi/CMakeLists.txt +++ b/examples/qualcomm/oss_scripts/moshi/CMakeLists.txt @@ -7,28 +7,26 @@ set(_qnn_mimi_decoder_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/qnn_mimi_decoder_runner.cpp ${CMAKE_CURRENT_LIST_DIR}/runner/runner.cpp - ${CMAKE_CURRENT_LIST_DIR}/runner/runner.h + ${CMAKE_CURRENT_LIST_DIR}/runner/runner.h ) # build mimi decoder runner add_executable(qnn_mimi_decoder_runner ${_qnn_mimi_decoder_runner__srcs}) target_include_directories( - qnn_mimi_decoder_runner PUBLIC ${_common_include_directories} + qnn_mimi_decoder_runner PUBLIC ${_common_include_directories} ) target_link_libraries( - qnn_mimi_decoder_runner - qnn_executorch_backend - executorch_core - extension_module - extension_data_loader - extension_flat_tensor - gflags + qnn_mimi_decoder_runner + qnn_executorch_backend + executorch_core + extension_module + extension_data_loader + extension_flat_tensor + gflags ) -target_compile_options( - qnn_llama_runner PUBLIC ${_common_compile_options} -) +target_compile_options(qnn_llama_runner PUBLIC ${_common_compile_options}) set_target_properties( - qnn_mimi_decoder_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" + qnn_mimi_decoder_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" ) diff --git a/examples/qualcomm/oss_scripts/moshi/mimi.py b/examples/qualcomm/oss_scripts/moshi/mimi.py index 1dba9bc8da1..0679b649d9f 100644 --- a/examples/qualcomm/oss_scripts/moshi/mimi.py +++ b/examples/qualcomm/oss_scripts/moshi/mimi.py @@ -176,9 +176,7 @@ def forward(self, x): ) -def inference_mimi_encoder( - args, encoder_inputs, encoder_input_list, encoder_pte_filename -): +def inference_mimi_encoder(args, encoder_inputs, encoder_pte_filename): adb = SimpleADB( qnn_sdk=os.getenv("QNN_SDK_ROOT"), build_path=f"{args.build_folder}", @@ -189,7 +187,7 @@ def inference_mimi_encoder( soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=encoder_inputs, input_list=encoder_input_list) + adb.push(inputs=encoder_inputs) adb.execute() # collect output data @@ -210,7 +208,7 @@ def inference_mimi_encoder( def export_mimi_encoder( args, orig_mimi, sample_pcm, pcm_chunk_size, skip_node_id_set, skip_node_op_set ): - encoder_inputs, encoder_input_list = [], "" + encoder_inputs = [] count = 0 cpu_encoded_results = [] logging.info("streaming encoding...") @@ -219,7 +217,6 @@ def export_mimi_encoder( chunk = sample_pcm[..., start_idx:end_idx] # Preparing QNN inputs encoder_inputs.append((chunk,)) - encoder_input_list += f"input_{count}_0.raw\n" count += 1 # Performing cpu encoding for golden codes = orig_mimi.encode(chunk) @@ -244,7 +241,6 @@ def export_mimi_encoder( qnn_encoded_results = inference_mimi_encoder( args, encoder_inputs, - encoder_input_list, encoder_pte_filename, ) else: @@ -260,7 +256,6 @@ def export_mimi_encoder( qnn_encoded_results = inference_mimi_encoder( args, encoder_inputs, - encoder_input_list, encoder_pte_filename, ) @@ -367,7 +362,7 @@ def inference_static_mimi_decoder( shared_buffer=args.shared_buffer, runner="examples/qualcomm/oss_scripts/moshi/qnn_mimi_decoder_runner", ) - adb.push(inputs=encoded_results, input_list=encoded_results_list) + adb.push(inputs=encoded_results) adb.execute(custom_runner_cmd=runner_cmd) # collect output data diff --git a/examples/qualcomm/oss_scripts/pvt.py b/examples/qualcomm/oss_scripts/pvt.py index fd2dee56e2f..0f850d0376e 100644 --- a/examples/qualcomm/oss_scripts/pvt.py +++ b/examples/qualcomm/oss_scripts/pvt.py @@ -7,10 +7,10 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np - import torch from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.examples.qualcomm.utils import ( @@ -31,12 +31,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +38,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -83,7 +77,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -135,6 +129,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/qwen2_5/qwen2_5.py b/examples/qualcomm/oss_scripts/qwen2_5/qwen2_5.py new file mode 100644 index 00000000000..5a08e1126f1 --- /dev/null +++ b/examples/qualcomm/oss_scripts/qwen2_5/qwen2_5.py @@ -0,0 +1,268 @@ +# 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 getpass +import json +import logging +import os +import subprocess +from multiprocessing.connection import Client + +import torch + +from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + +from executorch.examples.qualcomm.oss_scripts.llm_utils.qnn_decoder_model_manager import ( + get_qnn_llm_edge_manager, + HUGGING_FACE_REPO_IDS, +) + +from executorch.examples.qualcomm.utils import ( + make_output_dir, + parse_skip_delegation_node, + setup_common_args_and_variables, + SimpleADB, +) + +from transformers import AutoTokenizer + +FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" +logging.basicConfig(level=logging.INFO, format=FORMAT) +logging.getLogger().setLevel(logging.INFO) + +PTE_FILENAME = "qwen_qnn_q16" + + +def compile(args): # noqa: C901 + skip_node_id_set, skip_node_op_set = parse_skip_delegation_node(args) + + # ensure the working directory exist. + os.makedirs(args.artifact, exist_ok=True) + + manager = get_qnn_llm_edge_manager( + args.decoder_model, args.max_seq_len, args.enable_spinquant_r3 + ) + + fixed_point_type = {} + if args.ptq: + if args.ptq == "8a8w": + fixed_point_type["io_type"] = torch.uint8 + fixed_point_type["kv_type"] = torch.uint8 + elif args.ptq in ( + "16a8w", + "16a4w", + "16a4w_block", + "16a16w", + ): + fixed_point_type["io_type"] = torch.uint16 + fixed_point_type["kv_type"] = torch.uint16 + else: + raise ValueError( + f"No support for quant type {args.ptq}. Support 8a8w, 16a8w, 16a4w and 16a4w_block." + ) + quant_dtype = getattr(QuantDtype, f"use_{args.ptq}") + model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] + tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1] + + manager.pt2e_quantize( + quant_dtype, + fixed_point_type, + args.calibration_tasks, + args.calibration_limit, + args.prompt, + tokenizer_json_path, + ) + + manager.to_edge_transform_and_lower_to_qnn( + args.model, skip_node_id_set, skip_node_op_set + ) + if args.ptq: + logits_quant_attrs = manager.get_logits_quant_attrs() + json.dump( + { + "scale": logits_quant_attrs["scale"], + "zero_point": logits_quant_attrs["zero_point"], + }, + open(f"{args.artifact}/{PTE_FILENAME}_quant_attrs.txt", "w"), + ) + + manager.to_executorch(args.artifact, PTE_FILENAME) + + +def inference(args): + workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/{PTE_FILENAME}" + pte_path = f"{args.artifact}/{PTE_FILENAME}.pte" + # collect output data + output_data_folder = f"{args.artifact}/outputs" + make_output_dir(output_data_folder) + outputs = [] + + def post_process(): + with open(f"{args.artifact}/outputs/result.txt", "r") as f: + outputs.append(f.read()) + + model_id = HUGGING_FACE_REPO_IDS[args.decoder_model] + tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer_json_path = tokenizer.save_pretrained(args.artifact)[-1] + seq_len = args.max_seq_len + if args.enable_x86_64: + # x86 emulator is intended for CI and not performance. Check only the first few tokens. + seq_len = min(seq_len, 16) + + qnn_sdk = os.getenv("QNN_SDK_ROOT") + target = "x86_64-linux-clang" + runner_cmd = " ".join( + [ + f"export LD_LIBRARY_PATH={qnn_sdk}/lib/{target}/:{args.build_folder}/lib &&", + f"{args.build_folder}/examples/models/llama/llama_main", + f'--prompt "{args.prompt}"', + f"--tokenizer_path {tokenizer_json_path}", + f"--model_path {pte_path}", + f"--seq_len {seq_len}", + "--temperature 0", + f" > {output_data_folder}/result.txt", + ] + ) + subprocess.run( + runner_cmd, + shell=True, + executable="/bin/bash", + capture_output=True, + ) + post_process() + else: + runner_cmd = " ".join( + [ + f"cd {workspace} &&", + "./llama_main", + f'--prompt "{args.prompt}"', + "--tokenizer_path tokenizer.json", + f"--model_path {PTE_FILENAME}.pte", + f"--seq_len {seq_len}", + "--temperature 0", + " > outputs/result.txt", + ] + ) + adb = SimpleADB( + qnn_sdk=os.getenv("QNN_SDK_ROOT"), + build_path=f"{args.build_folder}", + pte_path=pte_path, + workspace=workspace, + device_id=args.device, + host_id=args.host, + soc_model=args.model, + runner="examples/models/llama/llama_main", + ) + # No pregen inputs, input_list is not required + adb.push(inputs=[], input_list="", files=[tokenizer_json_path]) + adb.execute(custom_runner_cmd=runner_cmd) + + adb.pull(output_path=args.artifact, callback=post_process) + + if args.ip and args.port != -1: + with Client((args.ip, args.port)) as conn: + conn.send( + json.dumps( + { + "result": outputs, + } + ) + ) + else: + for idx, output in enumerate(outputs): + logging.info(f"Results[{idx}]:\n{output}") + + +def main(args): + if args.compile_only and args.pre_gen_pte: + raise RuntimeError("Cannot set both compile_only and pre_gen_pte as true") + + if args.compile_only: + compile(args) + elif args.pre_gen_pte: + inference(args) + else: + compile(args) + inference(args) + + +if __name__ == "__main__": + parser = setup_common_args_and_variables() + + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts by this example.", + default="qwen2_5", + type=str, + ) + + parser.add_argument( + "-P", + "--ptq", + choices=["8a8w", "16a8w", "16a4w", "16a4w_block"], + help="If specified, will do PTQ quantization.", + type=str, + ) + + parser.add_argument( + "--pre_gen_pte", + help="Run the pre-generated Qwen in the given directory.", + type=str, + ) + + parser.add_argument( + "--prompt", + help="User prompts for Qwen.", + required=True, + type=str, + ) + + parser.add_argument( + "--decoder_model", + choices=["qwen2.5_0.5B", "qwen2.5_0.5B_instruct", "qwen2.5_1.5B_instruct"], + help="The Qwen model to export. Current available options are: [qwen2.5_0.5B, qwen2.5_0.5B_instruct, qwen2.5_1.5B_instruct]", + required=True, + ) + + parser.add_argument( + "--max_seq_len", + help="This refers to maximum number of tokens that the model can process & consider at once to generate predictions/responses.", + default=128, + type=int, + ) + parser.add_argument( + "--calibration_tasks", + nargs="+", + type=str, + default=None, + help="Tasks for GPTQ calibration from lm_eval", + ) + parser.add_argument( + "--calibration_limit", + type=int, + default=None, + help="number of samples used for calibration from lm_eval", + ) + parser.add_argument( + "--enable_spinquant_r3", + action="store_true", + help="Specify to enable spin quant R3", + ) + + try: + args = parser.parse_args() + args.validate(args) + if args.artifact is None: + args.artifact = args.decoder_model + main(args) + except Exception as e: + if args.ip and args.port != -1: + with Client((args.ip, args.port)) as conn: + conn.send(json.dumps({"Error": str(e)})) + else: + raise Exception(e) diff --git a/examples/qualcomm/oss_scripts/regnet.py b/examples/qualcomm/oss_scripts/regnet.py index 01b6bb9937e..617859df1af 100644 --- a/examples/qualcomm/oss_scripts/regnet.py +++ b/examples/qualcomm/oss_scripts/regnet.py @@ -6,6 +6,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -34,14 +35,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -81,7 +76,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -140,6 +135,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/retinanet.py b/examples/qualcomm/oss_scripts/retinanet.py index 229b35e3f8f..adb1241486d 100644 --- a/examples/qualcomm/oss_scripts/retinanet.py +++ b/examples/qualcomm/oss_scripts/retinanet.py @@ -6,6 +6,7 @@ import json import os + import sys from multiprocessing.connection import Client @@ -103,17 +104,16 @@ def resize_bbox(self, bbox, orig_shape): dataset = COCODataset(dataset_root=dataset_dir) test_loader = torch.utils.data.DataLoader(dataset=dataset, shuffle=True) - inputs, input_list = [], "" + inputs = [] bboxes, targets = [], [] for index, (img, boxes, labels) in enumerate(test_loader): if index >= data_size: break inputs.append((img,)) - input_list += f"input_{index}_0.raw\n" bboxes.append(boxes) targets.append(labels) - return inputs, input_list, bboxes, targets, dataset.label_names + return inputs, bboxes, targets, dataset.label_names def calculate_precision( @@ -214,19 +214,13 @@ def main(args): # ensure the working directory exist os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - model = get_instance() # retrieve dataset data_num = 100 # 91 classes appear in COCO dataset n_classes, n_coord_of_bbox = 91, 4 - inputs, input_list, bboxes, targets, label_names = get_dataset( + inputs, bboxes, targets, label_names = get_dataset( data_size=data_num, dataset_dir=args.dataset ) pte_filename = "retinanet_qnn" @@ -255,7 +249,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -320,6 +314,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/roberta.py b/examples/qualcomm/oss_scripts/roberta.py index cd70edc5dec..dc4a89d69b5 100644 --- a/examples/qualcomm/oss_scripts/roberta.py +++ b/examples/qualcomm/oss_scripts/roberta.py @@ -55,7 +55,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_masked_language_model_dataset( + inputs, targets = get_masked_language_model_dataset( args.dataset, tokenizer, data_size ) @@ -109,7 +109,7 @@ def main(args): sample_input["attention_mask"] = sample_input["attention_mask"].to(torch.float32) sample_input = tuple(sample_input.values()) golden = model(*sample_input)[0] - adb.push(inputs=[sample_input], input_list="input_0_0.raw input_0_1.raw\n") + adb.push(inputs=[sample_input]) adb.execute() adb.pull(output_path=args.artifact) @@ -121,7 +121,7 @@ def main(args): print(f"QNN output: {tokenizer.batch_decode(predictions.argmax(axis=2))}") # accuracy analysis - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() adb.pull(output_path=args.artifact) goldens, predictions = [], [] @@ -167,6 +167,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/squeezenet.py b/examples/qualcomm/oss_scripts/squeezenet.py index 9e486e94c07..e5d89d8f159 100644 --- a/examples/qualcomm/oss_scripts/squeezenet.py +++ b/examples/qualcomm/oss_scripts/squeezenet.py @@ -6,6 +6,7 @@ import json import os + from multiprocessing.connection import Client import numpy as np @@ -29,14 +30,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -72,7 +67,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -125,6 +120,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/ssd300_vgg16.py b/examples/qualcomm/oss_scripts/ssd300_vgg16.py index 2db51cd5c48..b9c30936a28 100644 --- a/examples/qualcomm/oss_scripts/ssd300_vgg16.py +++ b/examples/qualcomm/oss_scripts/ssd300_vgg16.py @@ -6,6 +6,7 @@ import json import os + import sys from multiprocessing.connection import Client from pprint import PrettyPrinter @@ -88,7 +89,7 @@ def get_dataset(data_size, dataset_dir, download): test_dataset, shuffle=True, collate_fn=test_dataset.collate_fn ) - inputs, input_list = [], "" + inputs = [] true_boxes = [] true_labels = [] true_difficulties = [] @@ -96,12 +97,11 @@ def get_dataset(data_size, dataset_dir, download): if index >= data_size: break inputs.append((images,)) - input_list += f"input_{index}_0.raw\n" true_boxes.extend(boxes) true_labels.extend(labels) true_difficulties.extend(difficulties) - return inputs, input_list, true_boxes, true_labels, true_difficulties + return inputs, true_boxes, true_labels, true_difficulties def SSD300VGG16(pretrained_weight_model): @@ -126,14 +126,8 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 - inputs, input_list, true_boxes, true_labels, true_difficulties = get_dataset( + inputs, true_boxes, true_labels, true_difficulties = get_dataset( data_size=data_num, dataset_dir=args.artifact, download=args.download ) @@ -165,7 +159,7 @@ def main(args): host_id=args.host, soc_model=args.model, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -275,6 +269,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/swin_transformer.py b/examples/qualcomm/oss_scripts/swin_transformer.py index 3c62eba45cd..1ee4a392add 100644 --- a/examples/qualcomm/oss_scripts/swin_transformer.py +++ b/examples/qualcomm/oss_scripts/swin_transformer.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -81,12 +82,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -94,7 +89,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -135,7 +130,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -188,6 +183,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/t5/CMakeLists.txt b/examples/qualcomm/oss_scripts/t5/CMakeLists.txt index 70fb613bb22..1bbec379341 100644 --- a/examples/qualcomm/oss_scripts/t5/CMakeLists.txt +++ b/examples/qualcomm/oss_scripts/t5/CMakeLists.txt @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - # preprocess qnn runner src files for t5 set(_qnn_t5_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/qnn_t5_runner.cpp @@ -19,10 +18,7 @@ set(_qnn_t5_runner__srcs # build qnn t5 runner add_executable(qnn_t5_runner ${_qnn_t5_runner__srcs}) -target_include_directories( - qnn_t5_runner PUBLIC ${_common_include_directories} -) - +target_include_directories(qnn_t5_runner PUBLIC ${_common_include_directories}) target_link_libraries( qnn_t5_runner @@ -37,9 +33,7 @@ target_link_libraries( tokenizers::tokenizers ) -target_compile_options( - qnn_t5_runner PUBLIC ${_common_compile_options} -) +target_compile_options(qnn_t5_runner PUBLIC ${_common_compile_options}) set_target_properties( - qnn_t5_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" + qnn_t5_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" ) diff --git a/examples/qualcomm/oss_scripts/t5/t5.py b/examples/qualcomm/oss_scripts/t5/t5.py index 1b8ea1b1665..093572f032a 100644 --- a/examples/qualcomm/oss_scripts/t5/t5.py +++ b/examples/qualcomm/oss_scripts/t5/t5.py @@ -11,7 +11,6 @@ from multiprocessing.connection import Client import torch - from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset from executorch.backends.qualcomm.utils.utils import ( @@ -207,19 +206,13 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_size = 100 max_hidden_seq_length = 384 max_cache_length = 512 tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-small").eval() - inputs, targets, input_list = get_seq2seq_dataset_from_squad_csv( + inputs, targets = get_seq2seq_dataset_from_squad_csv( args.dataset, tokenizer, data_size, @@ -307,7 +300,6 @@ def post_process(): ) adb.push( inputs=inputs, - input_list=input_list, files=[spiece_model], ) adb.execute(custom_runner_cmd=runner_cmd) @@ -351,6 +343,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/oss_scripts/whisper/CMakeLists.txt b/examples/qualcomm/oss_scripts/whisper/CMakeLists.txt index 5845575cba2..8f7d0f9a9be 100644 --- a/examples/qualcomm/oss_scripts/whisper/CMakeLists.txt +++ b/examples/qualcomm/oss_scripts/whisper/CMakeLists.txt @@ -4,7 +4,6 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - # preprocess qnn runner src files for whisper set(_qnn_whisper_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/qnn_whisper_runner.cpp @@ -20,10 +19,9 @@ set(_qnn_whisper_runner__srcs # build qnn whisper runner add_executable(qnn_whisper_runner ${_qnn_whisper_runner__srcs}) target_include_directories( - qnn_whisper_runner PUBLIC ${_common_include_directories} + qnn_whisper_runner PUBLIC ${_common_include_directories} ) - target_link_libraries( qnn_whisper_runner qnn_executorch_backend @@ -37,9 +35,7 @@ target_link_libraries( tokenizers::tokenizers ) -target_compile_options( - qnn_whisper_runner PUBLIC ${_common_compile_options} -) +target_compile_options(qnn_whisper_runner PUBLIC ${_common_compile_options}) set_target_properties( - qnn_whisper_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" + qnn_whisper_runner PROPERTIES LINK_FLAGS "-Wl,-rpath='$ORIGIN'" ) diff --git a/examples/qualcomm/oss_scripts/whisper/TARGETS b/examples/qualcomm/oss_scripts/whisper/TARGETS new file mode 100644 index 00000000000..7bd33c3db3f --- /dev/null +++ b/examples/qualcomm/oss_scripts/whisper/TARGETS @@ -0,0 +1,48 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") +load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_library( + name = "whisper_model_lib", + srcs = [ + "whisper_model.py", + ], + deps = [ + "//caffe2:torch", + "fbsource//third-party/pypi/transformers:transformers", + ], +) + +runtime.python_library( + name = "whisper_lib", + srcs = ["whisper.py"], + deps = [ + ":whisper_model_lib", + "//caffe2:torch", + "//executorch/backends/qualcomm/_passes:passes", + "//executorch/backends/qualcomm/partition:partition", + "//executorch/backends/qualcomm/quantizer:quantizer", + "//executorch/backends/qualcomm/serialization:serialization", + "//executorch/backends/qualcomm/utils:utils", + "//executorch/devtools/backend_debug:delegation_info", + "//executorch/examples/qualcomm:utils", + "//executorch/exir/capture:config", + "//executorch/exir/passes:memory_planning_pass", + "fbsource//third-party/pypi/datasets:datasets", + "fbsource//third-party/pypi/librosa:librosa", + "fbsource//third-party/pypi/soundfile:soundfile", + "fbsource//third-party/pypi/torchmetrics:torchmetrics", + "fbsource//third-party/pypi/transformers:transformers", + ], +) + +python_binary( + name = "whisper", + main_module = "executorch.examples.qualcomm.oss_scripts.whisper.whisper", + deps = [ + ":whisper_lib", + ], +) diff --git a/examples/qualcomm/oss_scripts/whisper/runner/runner.cpp b/examples/qualcomm/oss_scripts/whisper/runner/runner.cpp index 8cd75f433f7..c98326778bf 100644 --- a/examples/qualcomm/oss_scripts/whisper/runner/runner.cpp +++ b/examples/qualcomm/oss_scripts/whisper/runner/runner.cpp @@ -112,7 +112,11 @@ uint64_t Runner::logits_to_token( const executorch::aten::Tensor& logits_tensor) { return sampler_->sample(logits_tensor.data_ptr()); } - +/** + * @param inputs: A vector containing one element: a vector of bytes that + * encodes a float tensor in little-endian byte order. + * + */ Error Runner::transcribe( int32_t seq_len, std::vector>& inputs, diff --git a/examples/qualcomm/oss_scripts/whisper/targets.bzl b/examples/qualcomm/oss_scripts/whisper/targets.bzl new file mode 100644 index 00000000000..48f0174f392 --- /dev/null +++ b/examples/qualcomm/oss_scripts/whisper/targets.bzl @@ -0,0 +1,60 @@ +load( + "@fbsource//tools/build_defs:default_platform_defs.bzl", + "ANDROID", +) +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_oss_build_kwargs", "runtime") +load("@fbsource//xplat/executorch/backends/qualcomm/qnn_version.bzl", "get_qnn_library_version") + +def define_common_targets(): + runtime.cxx_library( + name = "runner_lib", + srcs = glob( + [ + "runner/*.cpp", + ], + ), + exported_headers = glob([ + "runner/*.h", + ]), + compiler_flags = [ + "-Wno-global-constructors", + "-Wunused-command-line-argument", + ], + deps = [ + "//executorch/extension/llm/runner:stats", + "//executorch/kernels/quantized:generated_lib", + "fbsource//third-party/qualcomm/qnn/qnn-{0}:api".format(get_qnn_library_version()), + ], + exported_deps = [ + "//executorch/extension/module:module", + "//executorch/extension/llm/sampler:sampler", + "//executorch/extension/tensor:tensor", + "//pytorch/tokenizers:hf_tokenizer", + "//executorch/extension/evalue_util:print_evalue", + "//executorch/backends/qualcomm/runtime:runtime", + ], + external_deps = [ + "gflags", + ], + platforms = [ANDROID], + **get_oss_build_kwargs() + ) + + runtime.cxx_binary( + name = "qnn_whisper_runner", + srcs = [ + "qnn_whisper_runner.cpp", + ], + compiler_flags = [ + "-Wno-global-constructors", + ], + deps = [ + ":runner_lib", + "//executorch/extension/threadpool:threadpool", + ], + external_deps = [ + "gflags", + ], + platforms = [ANDROID], + **get_oss_build_kwargs() + ) diff --git a/examples/qualcomm/oss_scripts/whisper/whisper.py b/examples/qualcomm/oss_scripts/whisper/whisper.py index a9f666e5f54..3985849411e 100644 --- a/examples/qualcomm/oss_scripts/whisper/whisper.py +++ b/examples/qualcomm/oss_scripts/whisper/whisper.py @@ -3,6 +3,10 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +# TODO: reenable pyre after fixing the issues +# pyre-ignore-all-errors + import getpass import json import logging @@ -71,7 +75,7 @@ def get_dataset(data_size): processor = AutoProcessor.from_pretrained("openai/whisper-tiny") # prepare input data - inputs, target, input_list = [], [], "" + inputs, target = [], [] for index, data in enumerate(dataset): if index >= data_size: break @@ -84,9 +88,8 @@ def get_dataset(data_size): ).input_features inputs.append((feature,)) target.append(data["text"]) - input_list += f"input_{index}_0.raw\n" - return inputs, input_list, target + return inputs, target def calibrate( @@ -331,11 +334,6 @@ def compile_whisper(args, inputs): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) tokenizer = AutoTokenizer.from_pretrained("openai/whisper-tiny") module = ( AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-tiny") @@ -362,7 +360,7 @@ def compile_whisper(args, inputs): ) -def inference_whisper(args, inputs, input_list, target): +def inference_whisper(args, inputs, target): workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/whisper" tokenizer = AutoTokenizer.from_pretrained("openai/whisper-tiny") tokenizer_json = tokenizer.save_pretrained(args.artifact)[-1] @@ -432,7 +430,7 @@ def post_process(): runner="examples/qualcomm/oss_scripts/whisper/qnn_whisper_runner", ) # No pregen inputs, input_list is not required - adb.push(inputs=inputs, input_list=input_list, files=[tokenizer_json]) + adb.push(inputs=inputs, files=[tokenizer_json]) adb.execute(custom_runner_cmd=runner_cmd) adb.pull(output_path=args.artifact, callback=post_process) @@ -479,6 +477,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) if args.compile_only and args.pre_gen_pte: exit("Cannot set both compile_only and pre_gen_pte as true") @@ -490,10 +489,10 @@ def post_process(): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, input_list, target = get_dataset(data_num) + inputs, target = get_dataset(data_num) if args.pre_gen_pte: - inference_whisper(args, inputs, input_list, target) + inference_whisper(args, inputs, target) exit(f"Finish the running pre_gen_pte from {args.pre_gen_pte}") if args.compile_only: @@ -502,7 +501,7 @@ def post_process(): try: compile_whisper(args, inputs) - inference_whisper(args, inputs, input_list, target) + inference_whisper(args, inputs, target) except Exception as e: if args.ip and args.port != -1: with Client((args.ip, args.port)) as conn: diff --git a/examples/qualcomm/qaihub_scripts/llama/CMakeLists.txt b/examples/qualcomm/qaihub_scripts/llama/CMakeLists.txt index 4088d685ec0..b42ceef6eae 100644 --- a/examples/qualcomm/qaihub_scripts/llama/CMakeLists.txt +++ b/examples/qualcomm/qaihub_scripts/llama/CMakeLists.txt @@ -30,8 +30,7 @@ list(PREPEND _qaihub_llama2_7b_runner__srcs add_executable(qaihub_llama2_7b_runner ${_qaihub_llama2_7b_runner__srcs}) target_include_directories( - qaihub_llama2_7b_runner - PUBLIC ${_common_include_directories} + qaihub_llama2_7b_runner PUBLIC ${_common_include_directories} ) target_link_libraries( qaihub_llama2_7b_runner @@ -64,9 +63,7 @@ list(APPEND _common_compile_options -DQAIHUB_LLAMA3_RUNNER) # build qaihub llama3 8b runner add_executable(qaihub_llama3_8b_runner ${_qaihub_llama3_8b_runner__srcs}) target_include_directories( - qaihub_llama3_8b_runner - PUBLIC - ${_common_include_directories} + qaihub_llama3_8b_runner PUBLIC ${_common_include_directories} ) target_link_libraries( diff --git a/examples/qualcomm/qaihub_scripts/stable_diffusion/qaihub_stable_diffusion.py b/examples/qualcomm/qaihub_scripts/stable_diffusion/qaihub_stable_diffusion.py index 8e56ce11e2e..7905dfa9a7e 100644 --- a/examples/qualcomm/qaihub_scripts/stable_diffusion/qaihub_stable_diffusion.py +++ b/examples/qualcomm/qaihub_scripts/stable_diffusion/qaihub_stable_diffusion.py @@ -258,13 +258,11 @@ def inference(args, compiler_specs, pte_files): ) input_unet = () - input_list_unet = "" - for i, t in enumerate(scheduler.timesteps): + for t in scheduler.timesteps: time_emb = get_quant_data( encoding, get_time_embedding(t, time_embedding), "unet", 1 ) - input_list_unet += f"input_{i}_0.raw\n" input_unet = input_unet + (time_emb,) qnn_executor_runner_args = [ @@ -333,7 +331,7 @@ def inference(args, compiler_specs, pte_files): files.append(os.path.join(args.artifact, "latents.raw")) if not args.skip_push: - adb.push(inputs=input_unet, input_list=input_list_unet, files=files) + adb.push(inputs=input_unet, files=files) adb.execute(custom_runner_cmd=qnn_executor_runner_args) output_image = [] diff --git a/examples/qualcomm/qaihub_scripts/utils/export.py b/examples/qualcomm/qaihub_scripts/utils/export.py index 2ee1968dd82..ff364ab986e 100644 --- a/examples/qualcomm/qaihub_scripts/utils/export.py +++ b/examples/qualcomm/qaihub_scripts/utils/export.py @@ -126,9 +126,8 @@ def get_tensor(io_info, tensors, logger, checking_output=False): return [get_ones_tensor(t, logger) for t in io_info] # list of tensors to be returned - ret_tensors, ret_list = [], [] + ret_tensors = [] for i, info in enumerate(io_info): - ret_list.append(f"input_0_{i}.raw") if list(tensors[i].shape) != info["shape"]: logger.error( f"tensor '{info['name']}' shape mismatch: " @@ -145,7 +144,7 @@ def get_tensor(io_info, tensors, logger, checking_output=False): # try quant / dequant for given tensor if possible ret_tensors.append(get_tensor_with_encoding(tensors[i], info, logger)) ) - return [ret_tensors], " ".join(ret_list) + return [ret_tensors] def to_context_binary( @@ -297,7 +296,7 @@ def execute(args): # check if inputs are valid, fallback to ones tensor if any logger.info("generating input data") - inputs, input_list = get_tensor(graph_info["inputs"], user_inputs, logger) + inputs = get_tensor(graph_info["inputs"], user_inputs, logger) logger.info("preparing ADB connection") # leverage SimpleADB for e2e inference @@ -313,7 +312,7 @@ def execute(args): ) logger.info("pushing QNN libraries & other artifacts") - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) logger.info("starting inference") adb.execute() diff --git a/examples/qualcomm/scripts/README.md b/examples/qualcomm/scripts/README.md new file mode 100644 index 00000000000..da9d10ba663 --- /dev/null +++ b/examples/qualcomm/scripts/README.md @@ -0,0 +1,68 @@ +# Usage Guide for Models Provided by ExecuTorch + +This guide provides examples and instructions for deploying and executing models using the ExecuTorch runtime on Qualcomm platforms. + +## Model categories +The following models in two folders can be categorized based on their primary use cases. + +1. Language Model: + - mobilebert_fine_tune + +2. Speech Model: + - wav2letter + +3. Vision Model: + - edsr + - inception_v3 + - inception_v4 + - mobilenet_v2 + - mobilenet_v3 + - torchvision_vit + +## Prerequisite +Please follow another [README](../README.md) first to set up environment. + +## Model running +Some models require specific datasets. Please download them in advance and place them in the appropriate folders. + +Detailed instructions for each model are provided below. +If you want to export the model without running it, please add `--compile_only` to the command. + +1. `deeplab_v3`: + - Required Dataset : VOCSegmentation + Will be downloaded automatically if `-d` is specified + + ```bash + python deeplab_v3.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} + +2. `edsr`: + - Required Dataset : DIV2K + + Will be downloaded automatically if -d is specified. Alternatively, you can provide your own dataset using `--hr_ref_dir` and `--lr_dir`. + + ```bash + pip install piq + python edsr.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d + +3. `inception_v3`, `inception_v4`, `mobilenet_v2`,`mobilenet_v3`, `torchvision_vit`: + - Required Dataset : ImageNet + + Download [dataset](https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000) first, and place it in a valid folder. + + ```bash + python SCRIPT_NAME.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -d path/to/ImageNet + +4. `mobilebert_fine_tune`: + - You can specify the pretrained weight using `-p `, if no pretrained weights are provided, using `--num_epochs` to set number of epochs to train the model. + - `-F --use_fp16`: If specified, the model will run in FP16 mode and the PTQ will be ignored. + + ```bash + python mobilebert_fine_tune.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} --num_epochs + +5. `wav2letter`: + - Pretrained weight : + + for torchaudio.models.Wav2Letter version, please download at [here](https://github.com/nipponjo/wav2letter-ctc-pytorch/tree/main?tab=readme-ov-file#wav2letter-ctc-pytorch), and place it in a valid folder. + ```bash + python wav2letter.py -m ${SOC_MODEL} -b path/to/build-android/ -s ${DEVICE_SERIAL} -p path/to/pretrained_weight + \ No newline at end of file diff --git a/examples/qualcomm/scripts/deeplab_v3.py b/examples/qualcomm/scripts/deeplab_v3.py index cb64d904919..164dc312fae 100755 --- a/examples/qualcomm/scripts/deeplab_v3.py +++ b/examples/qualcomm/scripts/deeplab_v3.py @@ -50,16 +50,15 @@ def get_dataset(data_size, dataset_dir, download): # prepare input data random.shuffle(dataset) - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] for index, data in enumerate(dataset): if index >= data_size: break image, target = data inputs.append((image.unsqueeze(0),)) targets.append(np.array(target.resize(input_size))) - input_list += f"input_{index}_0.raw\n" - return inputs, targets, input_list + return inputs, targets def main(args): @@ -68,12 +67,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -81,7 +74,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_dataset( + inputs, targets = get_dataset( data_size=data_num, dataset_dir=args.artifact, download=args.download ) @@ -113,7 +106,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -200,6 +193,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/edsr.py b/examples/qualcomm/scripts/edsr.py index 222c04ed1b1..0601ff9e0b7 100755 --- a/examples/qualcomm/scripts/edsr.py +++ b/examples/qualcomm/scripts/edsr.py @@ -7,6 +7,7 @@ import json import logging import os + import re from multiprocessing.connection import Client @@ -57,12 +58,6 @@ def _resize_img(self, file: str, scale: int): with Image.open(file) as img: return to_tensor(img.resize(tuple(self.input_size * scale))).unsqueeze(0) - def get_input_list(self): - input_list = "" - for i in range(len(self.lr)): - input_list += f"input_{i}_0.raw\n" - return input_list - def get_b100( dataset_dir: str, @@ -107,12 +102,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - instance = EdsrModel() if args.ci: inputs = instance.get_example_inputs() @@ -124,7 +113,7 @@ def main(args): args.hr_ref_dir, args.lr_dir, args.default_dataset, args.artifact ) - inputs, targets, input_list = dataset.lr, dataset.hr, dataset.get_input_list() + inputs, targets = dataset.lr, dataset.hr pte_filename = "edsr_qnn_q8" build_executorch_binary( @@ -152,7 +141,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -235,6 +224,7 @@ def post_process(): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/export_example.py b/examples/qualcomm/scripts/export_example.py index 44335ebb32c..9b5ace12327 100644 --- a/examples/qualcomm/scripts/export_example.py +++ b/examples/qualcomm/scripts/export_example.py @@ -1,6 +1,5 @@ # pyre-ignore-all-errors import argparse -import copy import torch from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer @@ -10,10 +9,8 @@ get_soc_to_chipset_map, to_edge_transform_and_lower_to_qnn, ) -from executorch.devtools import generate_etrecord from executorch.examples.models import MODEL_NAME_TO_MODEL from executorch.examples.models.model_factory import EagerModelFactory -from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.extension.export_util.utils import save_pte_program from torchao.quantization.pt2e.quantize_pt2e import ( @@ -107,19 +104,14 @@ def main() -> None: backend_options=backend_options, ) delegated_program = to_edge_transform_and_lower_to_qnn( - m, example_inputs, compile_spec + m, example_inputs, compile_spec, generate_etrecord=args.generate_etrecord ) - # this is needed for the ETRecord as lowering modifies the graph in-place - edge_copy = copy.deepcopy(delegated_program) - - executorch_program = delegated_program.to_executorch( - config=ExecutorchBackendConfig(extract_delegate_segments=False) - ) + executorch_program = delegated_program.to_executorch() if args.generate_etrecord: etrecord_path = args.output_folder + "etrecord.bin" - generate_etrecord(etrecord_path, edge_copy, executorch_program) + executorch_program.get_etrecord().save(etrecord_path) save_pte_program(executorch_program, args.model_name, args.output_folder) diff --git a/examples/qualcomm/scripts/inception_v3.py b/examples/qualcomm/scripts/inception_v3.py index 6cfb44adcf7..fbd693d7100 100755 --- a/examples/qualcomm/scripts/inception_v3.py +++ b/examples/qualcomm/scripts/inception_v3.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -31,12 +32,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +39,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -77,7 +72,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -130,6 +125,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/inception_v4.py b/examples/qualcomm/scripts/inception_v4.py index 92de33f8cba..a13804d8868 100755 --- a/examples/qualcomm/scripts/inception_v4.py +++ b/examples/qualcomm/scripts/inception_v4.py @@ -7,10 +7,10 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np - import torch from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype from executorch.examples.models.inception_v4 import InceptionV4Model @@ -31,12 +31,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +38,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(299, 299), @@ -76,7 +70,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -129,6 +123,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/mobilebert_fine_tune.py b/examples/qualcomm/scripts/mobilebert_fine_tune.py index bd0b6dfbcf2..8eff2c9f2eb 100755 --- a/examples/qualcomm/scripts/mobilebert_fine_tune.py +++ b/examples/qualcomm/scripts/mobilebert_fine_tune.py @@ -65,10 +65,10 @@ def accuracy_per_class(preds, goldens, labels): def get_dataset(data_val): # prepare input data - inputs, input_list = [], "" + inputs = [] # max_position_embeddings defaults to 512 position_ids = torch.arange(512).expand((1, -1)).to(torch.int32) - for index, data in enumerate(data_val): + for data in data_val: data = [d.to(torch.int32) for d in data] # input_ids, attention_mask, token_type_ids, position_ids inputs.append( @@ -78,12 +78,8 @@ def get_dataset(data_val): position_ids[:, : data[0].shape[1]], ) ) - input_text = " ".join( - [f"input_{index}_{i}.raw" for i in range(len(inputs[-1]))] - ) - input_list += f"{input_text}\n" - return inputs, input_list + return inputs def get_fine_tuned_mobilebert(artifacts_dir, pretrained_weight, batch_size): @@ -228,17 +224,11 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - batch_size, pte_filename = 1, "ptq_mb_qnn" model, data_val, labels = get_fine_tuned_mobilebert( args.artifact, args.pretrained_weight, batch_size ) - inputs, input_list = get_dataset(data_val) + inputs = get_dataset(data_val) try: quant_dtype = getattr(QuantDtype, f"use_{args.ptq}") @@ -303,7 +293,7 @@ def calibrator(gm): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -381,6 +371,7 @@ def calibrator(gm): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/mobilenet_v2.py b/examples/qualcomm/scripts/mobilenet_v2.py index 1b153431741..e3c5d16e545 100755 --- a/examples/qualcomm/scripts/mobilenet_v2.py +++ b/examples/qualcomm/scripts/mobilenet_v2.py @@ -31,12 +31,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -44,7 +38,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -77,7 +71,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -130,6 +124,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/mobilenet_v3.py b/examples/qualcomm/scripts/mobilenet_v3.py index e34125bbfca..02fa8c93269 100644 --- a/examples/qualcomm/scripts/mobilenet_v3.py +++ b/examples/qualcomm/scripts/mobilenet_v3.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -30,12 +31,6 @@ def main(args): # ensure the working directory exist. os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - data_num = 100 if args.ci: inputs = [(torch.rand(1, 3, 224, 224),)] @@ -43,7 +38,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -75,7 +70,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -128,6 +123,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/torchvision_vit.py b/examples/qualcomm/scripts/torchvision_vit.py index 428863daf4b..0417ffdc673 100755 --- a/examples/qualcomm/scripts/torchvision_vit.py +++ b/examples/qualcomm/scripts/torchvision_vit.py @@ -7,6 +7,7 @@ import json import logging import os + from multiprocessing.connection import Client import numpy as np @@ -35,7 +36,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_imagenet_dataset( + inputs, targets = get_imagenet_dataset( dataset_path=f"{args.dataset}", data_size=data_num, image_shape=(256, 256), @@ -67,7 +68,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -118,6 +119,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/scripts/wav2letter.py b/examples/qualcomm/scripts/wav2letter.py index e5b97a8241e..980a1dd9acc 100644 --- a/examples/qualcomm/scripts/wav2letter.py +++ b/examples/qualcomm/scripts/wav2letter.py @@ -7,6 +7,7 @@ import json import logging import os + import sys from multiprocessing.connection import Client @@ -66,17 +67,16 @@ def collate_fun(batch): collate_fn=lambda x: collate_fun(x), ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] for wave, label in data_loader: for index in range(data_size): # reshape input tensor to NCHW inputs.append((wave[index].reshape(1, 1, -1, 1),)) targets.append(label[index]) - input_list += f"input_{index}_0.raw\n" # here we only take first batch, i.e. 'data_size' tensors break - return inputs, targets, input_list + return inputs, targets def eval_metric(pred, target_str): @@ -102,12 +102,6 @@ def main(args): # ensure the working directory exist os.makedirs(args.artifact, exist_ok=True) - if not args.compile_only and args.device is None: - raise RuntimeError( - "device serial is required if not compile only. " - "Please specify a device serial by -s/--device argument." - ) - instance = Wav2LetterModel() # target labels " abcdefghijklmnopqrstuvwxyz'*" instance.vocab_size = 29 @@ -140,9 +134,7 @@ def main(args): "This option is for CI to verify the export flow. It uses random input and will result in poor accuracy." ) else: - inputs, targets, input_list = get_dataset( - data_size=data_num, artifact_dir=args.artifact - ) + inputs, targets = get_dataset(data_size=data_num, artifact_dir=args.artifact) pte_filename = "w2l_qnn" build_executorch_binary( model, @@ -169,7 +161,7 @@ def main(args): soc_model=args.model, shared_buffer=args.shared_buffer, ) - adb.push(inputs=inputs, input_list=input_list) + adb.push(inputs=inputs) adb.execute() # collect output data @@ -228,6 +220,7 @@ def main(args): ) args = parser.parse_args() + args.validate(args) try: main(args) except Exception as e: diff --git a/examples/qualcomm/test_qualcomm.sh b/examples/qualcomm/test_qualcomm.sh index 19d3d798418..51a563863f3 100644 --- a/examples/qualcomm/test_qualcomm.sh +++ b/examples/qualcomm/test_qualcomm.sh @@ -15,7 +15,7 @@ cmake_install_executorch_qnn_lib() { echo "Installing libexecutorch.a, libqnn_executorch_backend.a" rm -rf cmake-out - retry cmake -DBUCK2="$BUCK" \ + retry cmake \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_QNN=ON \ @@ -55,11 +55,5 @@ then PYTHON_EXECUTABLE=python3 fi -if [[ -z $BUCK ]]; -then - BUCK=buck2 -fi - - cmake_install_executorch_qnn_lib test_cmake_qualcomm diff --git a/examples/qualcomm/util_scripts/cli.py b/examples/qualcomm/util_scripts/cli.py index e4c4c5dcaf8..5745e248808 100644 --- a/examples/qualcomm/util_scripts/cli.py +++ b/examples/qualcomm/util_scripts/cli.py @@ -229,7 +229,7 @@ def execute(args): # load input files logger.info("loading user inputs") - user_inputs, input_list = [], "" + user_inputs = [] with open(args.input_list, "r") as f: for line in f.read().split("\n")[:-1]: inputs, input_names = [], "" @@ -237,7 +237,6 @@ def execute(args): input_names += f"{Path(data).stem}.raw " inputs.append(torch.load(data, weights_only=True)) user_inputs.append(inputs) - input_list += input_names.strip() + "\n" logger.info("retrieving graph I/O") # setup compiler spec dedicated to QNN HTP backend @@ -263,7 +262,7 @@ def execute(args): ) logger.info("pushing QNN libraries & other artifacts") - adb.push(inputs=user_inputs, input_list=input_list) + adb.push(inputs=user_inputs) logger.info("starting inference") adb.execute() diff --git a/examples/qualcomm/util_scripts/gen_etrecord.py b/examples/qualcomm/util_scripts/gen_etrecord.py index 305a6054735..7c1ced1e032 100644 --- a/examples/qualcomm/util_scripts/gen_etrecord.py +++ b/examples/qualcomm/util_scripts/gen_etrecord.py @@ -1,4 +1,3 @@ -import copy import os import torch @@ -10,7 +9,7 @@ QcomChipset, to_edge_transform_and_lower_to_qnn, ) -from executorch.devtools import generate_etrecord, Inspector +from executorch.devtools import Inspector from executorch.devtools.inspector._inspector_utils import TimeScale from executorch.examples.qualcomm.utils import ( make_quantizer, @@ -46,11 +45,9 @@ def main(args): module=converted, inputs=sample_input, compiler_specs=compiler_specs, + generate_etrecord=True, ) - # for inspector API - edge_copy = copy.deepcopy(edge_prog_mgr) - # store pte file exec_prog = edge_prog_mgr.to_executorch() with open(f"{pte_filename}.pte", "wb") as f: @@ -65,13 +62,12 @@ def main(args): device_id=args.device, soc_model=args.model, ) - input_list = "input_0_0.raw input_0_1.raw\n" - adb.push(inputs=[sample_input], input_list=input_list) + adb.push(inputs=[sample_input]) adb.execute() # pull etdump back and display the statistics adb.pull_etdump(".") - generate_etrecord("etrecord.bin", edge_copy, exec_prog) + exec_prog.get_etrecord().save("etrecord.bin") inspector = Inspector( etdump_path="etdump.etdp", etrecord="etrecord.bin", diff --git a/examples/qualcomm/utils.py b/examples/qualcomm/utils.py index c12cb582961..e43821bda64 100755 --- a/examples/qualcomm/utils.py +++ b/examples/qualcomm/utils.py @@ -104,16 +104,22 @@ def __init__( self.expected_output_shape = expected_output_shape self.extra_cmds = "" - def _adb(self, cmd): + def _adb(self, cmd, output_callback: Optional[Callable[[str], None]] = None): if not self.host_id: cmds = ["adb", "-s", self.device_id] else: cmds = ["adb", "-H", self.host_id, "-s", self.device_id] cmds.extend(cmd) - subprocess.run( - cmds, stdout=subprocess.DEVNULL if self.error_only else sys.stdout - ) + if output_callback: + result = subprocess.run( + cmds, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True + ) + output_callback(result) + else: + subprocess.run( + cmds, stdout=subprocess.DEVNULL if self.error_only else sys.stdout + ) def push(self, inputs=None, input_list=None, files=None, init_env=True): artifacts = [] @@ -140,7 +146,7 @@ def push(self, inputs=None, input_list=None, files=None, init_env=True): f"{self.qnn_sdk}/lib/aarch64-android/libQnnModelDlc.so", ] input_list_file, input_files = generate_inputs( - self.working_dir, self.input_list_filename, inputs, input_list + self.working_dir, self.input_list_filename, inputs ) if input_list_file is not None: @@ -173,7 +179,12 @@ def push(self, inputs=None, input_list=None, files=None, init_env=True): for file_name in files: self._adb(["push", file_name, self.workspace]) - def execute(self, custom_runner_cmd=None, method_index=0): + def execute( + self, + custom_runner_cmd=None, + method_index=0, + output_callback: Optional[Callable[[str], None]] = None, + ): self._adb(["shell", f"mkdir -p {self.output_folder}"]) # run the delegation if custom_runner_cmd is None: @@ -205,8 +216,9 @@ def execute(self, custom_runner_cmd=None, method_index=0): ) else: qnn_executor_runner_cmds = custom_runner_cmd - - self._adb(["shell", f"{qnn_executor_runner_cmds}"]) + self._adb( + ["shell", f"{qnn_executor_runner_cmds}"], output_callback=output_callback + ) def pull(self, output_path, callback=None): self._adb(["pull", "-a", self.output_folder, output_path]) @@ -574,7 +586,7 @@ def get_data_loader(): ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] data_loader = get_data_loader() for index, data in enumerate(data_loader): if index >= data_size: @@ -582,9 +594,8 @@ def get_data_loader(): feature, target = data inputs.append((feature,)) targets.append(target) - input_list += f"input_{index}_0.raw\n" - return inputs, targets, input_list + return inputs, targets def get_masked_language_model_dataset(dataset_path, tokenizer, data_size, shuffle=True): @@ -624,10 +635,9 @@ def __len__(self): ) # prepare input data - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] data_loader = get_data_loader() - for _, data in enumerate(data_loader): - index = len(inputs) + for data in data_loader: if len(inputs) >= data_size: break input_ids = data[0] @@ -639,9 +649,8 @@ def __len__(self): continue inputs.append((input_ids, attention_mask)) targets.append(target) - input_list += f"input_{index}_0.raw input_{index}_1.raw\n" - return inputs, targets, input_list + return inputs, targets def get_seq2seq_dataset_from_squad_csv( # noqa: C901 @@ -720,9 +729,9 @@ def __getitem__(self, idx): dataset, batch_size=1, shuffle=shuffle, collate_fn=collator ) - inputs, targets, input_list = [], [], "" + inputs, targets = [], [] data_loader = get_data_loader(max_hidden_seq_length) - for idx, batch in enumerate(data_loader): + for batch in data_loader: if len(inputs) >= data_size: break input_ids = batch["input_ids"] @@ -741,9 +750,8 @@ def __getitem__(self, idx): ) ) targets.append(labels) - input_list += f"input_{idx}_0.raw input_{idx}_1.raw input_{idx}_2.raw\n" - return inputs, targets, input_list + return inputs, targets def setup_common_args_and_variables(): @@ -860,12 +868,31 @@ def setup_common_args_and_variables(): default=False, ) + parser.add_argument( + "--seed", + help="Set the seed for generating random numbers in both torch and random.", + type=int, + ) + # QNN_SDK_ROOT might also be an argument, but it is used in various places. # So maybe it's fine to just use the environment. if "QNN_SDK_ROOT" not in os.environ: raise RuntimeError("Environment variable QNN_SDK_ROOT must be set") print(f"QNN_SDK_ROOT={os.getenv('QNN_SDK_ROOT')}") + def validate(args): + if not args.compile_only and args.device is None: + raise RuntimeError( + "device serial is required if not compile only. " + "Please specify a device serial by -s/--device argument." + ) + if args.seed: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + random.seed(args.seed) + + parser.set_defaults(validate=validate) + return parser @@ -884,25 +911,28 @@ def parse_skip_delegation_node(args): return skip_node_id_set, skip_node_op_set -def generate_inputs(dest_path: str, file_name: str, inputs=None, input_list=None): +def generate_inputs(dest_path: str, file_name: str, inputs=None): input_list_file = None input_files = [] - # Prepare input list - if input_list is not None: - input_list_file = f"{dest_path}/{file_name}" - with open(input_list_file, "w") as f: - f.write(input_list) - f.flush() - # Prepare input data if inputs is not None: - for idx, data in enumerate(inputs): - for i, d in enumerate(data): - file_name = f"{dest_path}/input_{idx}_{i}.raw" - if not isinstance(d, torch.Tensor): - d = torch.tensor(d) - d.detach().numpy().tofile(file_name) - input_files.append(file_name) + input_list_file = f"{dest_path}/{file_name}" + with open(input_list_file, "w") as f: + for idx, data in enumerate(inputs): + for i, d in enumerate(data): + # transform torch.Tensor to raw file + file_name = f"input_{idx}_{i}.raw" + file_path = f"{dest_path}/{file_name}" + if not isinstance(d, torch.Tensor): + d = torch.tensor(d) + d.detach().numpy().tofile(file_path) + input_files.append(file_path) + + # prepare input_list + if i > 0: + f.write(" ") + f.write(file_name) + f.write("\n") return input_list_file, input_files diff --git a/examples/rpi/build_firmware_pico.sh b/examples/rpi/build_firmware_pico.sh new file mode 100755 index 00000000000..8880a834777 --- /dev/null +++ b/examples/rpi/build_firmware_pico.sh @@ -0,0 +1,98 @@ +#!/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. + +#!/bin/bash +# build_firmware_pico.sh +# Simple script to cross-compile ExecuTorch and build Pico2 firmware with optional model input + +set -e + +# Paths +ROOT_DIR="$(cd "$(dirname "$0")/../.." && pwd)" # examples/rpi/ -> root dir +PICO2_DIR="${ROOT_DIR}/examples/rpi/pico2" +BUILD_DIR="${PICO2_DIR}/build" +EXECUTORCH_BUILD_DIR="${ROOT_DIR}/cmake-out" + +# Default model +DEFAULT_MODEL="default_model.pte" + +usage() { + echo "Usage: $0 [--clean] [--model=path/to/model.pte]" + echo " --clean Clean build directories" + echo " --model=FILE Specify model file to embed (relative to pico2/)" + exit 1 +} + +# Parse args +MODEL_INPUT="" +CLEAN_BUILD=0 + +for arg in "$@"; do + case $arg in + --clean) + CLEAN_BUILD=1 + shift + ;; + --model=*) + MODEL_INPUT="${arg#*=}" + shift + ;; + *) + usage + ;; + esac +done + +# Clean if requested +if [ $CLEAN_BUILD -eq 1 ]; then + echo "Cleaning build directories..." + rm -rf "${EXECUTORCH_BUILD_DIR}" "${BUILD_DIR}" + echo "Clean complete." + exit 0 +fi + +# Step 1: Cross compile ExecuTorch from root dir +echo "Cross compiling ExecuTorch baremetal ARM..." + +cmake -B "${EXECUTORCH_BUILD_DIR}" \ + -DCMAKE_TOOLCHAIN_FILE="${ROOT_DIR}/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake" \ + -DTARGET_CPU=cortex-m0plus \ + -DEXECUTORCH_BUILD_ARM_BAREMETAL=ON \ + -DEXECUTORCH_PAL_DEFAULT=minimal \ + -DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON \ + -DCMAKE_BUILD_TYPE=MinSizeRel \ + -DEXECUTORCH_ENABLE_LOGGING=OFF \ + -DEXECUTORCH_SELECT_ALL_OPS=OFF \ + -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \ + -DCMAKE_INSTALL_PREFIX="${EXECUTORCH_BUILD_DIR}" \ + "${ROOT_DIR}" + +cmake --build "${EXECUTORCH_BUILD_DIR}" --target install -j$(nproc) + +echo "ExecuTorch cross compile complete." + +# Step 2: Build firmware for Pico2 with model input + +cd "${PICO2_DIR}" + +if [ -n "$MODEL_INPUT" ]; then + # Use specified model + if [ ! -f "${MODEL_INPUT}" ]; then + echo "Error: Specified model file '${MODEL_INPUT}' not found in pico2 directory." + exit 1 + fi + echo "Building firmware with model: ${MODEL_INPUT}" + cmake -B "${BUILD_DIR}" -DPICO_BOARD=pico2 -DINPUT_MODEL="./${MODEL_INPUT}" -DCMAKE_BUILD_TYPE=Release +else + # Use default model + echo "Building firmware with default model: ${DEFAULT_MODEL}" + cmake -B "${BUILD_DIR}" -DPICO_BOARD=pico2 -DINPUT_MODEL="./${DEFAULT_MODEL}" -DCMAKE_BUILD_TYPE=Release +fi + +cmake --build "${BUILD_DIR}" -j$(nproc) + +echo "Firmware build complete. Output in ${BUILD_DIR}, Binary: executorch_pico.uf2" diff --git a/examples/rpi/pico2/CMakeLists.txt b/examples/rpi/pico2/CMakeLists.txt new file mode 100644 index 00000000000..34d9833e6c0 --- /dev/null +++ b/examples/rpi/pico2/CMakeLists.txt @@ -0,0 +1,143 @@ +# 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. + +cmake_minimum_required(VERSION 3.13) + +# Include FetchContent for Pico SDK +include(FetchContent) + +FetchContent_Declare( + pico_sdk + GIT_REPOSITORY https://github.com/raspberrypi/pico-sdk.git + GIT_TAG 2.0.0 +) + +FetchContent_MakeAvailable(pico_sdk) +set(PICO_SDK_PATH ${pico_sdk_SOURCE_DIR}) +include(${PICO_SDK_PATH}/pico_sdk_init.cmake) + +project(executorch_pico C CXX ASM) +pico_sdk_init() + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + +# Configure input model file +set(INPUT_MODEL + "default_model.pte" + CACHE STRING "Input PTE model file to embed" +) + +# Convert relative paths to absolute +if(NOT IS_ABSOLUTE "${INPUT_MODEL}") + set(INPUT_MODEL "${CMAKE_CURRENT_SOURCE_DIR}/${INPUT_MODEL}") +endif() + +message(STATUS "Using model file: ${INPUT_MODEL}") + +# Validate model file exists +if(NOT EXISTS "${INPUT_MODEL}") + message(FATAL_ERROR "Model file '${INPUT_MODEL}' does not exist.") +endif() + +# Create model_pte.c if it doesn't exist +set(MODEL_PTE_C "${CMAKE_CURRENT_SOURCE_DIR}/model_pte.c") +if(NOT EXISTS "${MODEL_PTE_C}") + file( + WRITE "${MODEL_PTE_C}" + "#include \"model_pte.h\"\n" + "\n" + "const uint8_t model_pte[] __attribute__((aligned(8))) = {\n" + " // Model data will be injected here\n" + "};\n" + "\n" + "const unsigned int model_pte_len = sizeof(model_pte);\n" + ) + message(STATUS "Created initial model_pte.c") +endif() + +# Use a stamp file instead of declaring model_pte.c as OUTPUT +set(MODEL_STAMP "${CMAKE_CURRENT_BINARY_DIR}/model_pte.stamp") + +add_custom_command( + OUTPUT ${MODEL_STAMP} + COMMAND python3 ${EXECUTORCH_ROOT}/executorch/examples/rpi/pte_to_array.py + --model ${INPUT_MODEL} --file ${MODEL_PTE_C} + COMMAND ${CMAKE_COMMAND} -E touch ${MODEL_STAMP} + DEPENDS ${INPUT_MODEL} + ${EXECUTORCH_ROOT}/executorch/examples/rpi/pte_to_array.py + COMMENT "Injecting PTE data from '${INPUT_MODEL}' into model_pte.c" + VERBATIM +) + +# Target to ensure model injection happens +add_custom_target(inject_model_data ALL DEPENDS ${MODEL_STAMP}) + +# Check if source files exist +if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/main.cpp") + message(FATAL_ERROR "main.cpp not found in ${CMAKE_CURRENT_SOURCE_DIR}") +endif() + +# Create the executable +add_executable( + executorch_pico main.cpp ${MODEL_PTE_C} # Use the full path to be explicit +) + +# Ensure model data is injected before building +add_dependencies(executorch_pico inject_model_data) + +# Set correct target properties for Pico2/RP2350 +set_target_properties( + executorch_pico PROPERTIES SUFFIX ".elf" OUTPUT_NAME "executorch_pico" +) + +# Configure Pico-specific settings +pico_enable_stdio_usb(executorch_pico 1) +pico_enable_stdio_uart(executorch_pico 0) + +# Set correct flags for Pico2 (RP2350) +target_compile_options( + executorch_pico PRIVATE -mcpu=cortex-m33 -mfloat-abi=soft -mthumb +) + +target_include_directories( + executorch_pico + PRIVATE ${EXECUTORCH_ROOT} ${EXECUTORCH_ROOT}/executorch/third-party + ${EXECUTORCH_ROOT}/executorch/runtime/core/portable_type/c10 +) + +target_compile_definitions( + executorch_pico + PRIVATE C10_USING_CUSTOM_GENERATED_MACROS EXECUTORCH_ENABLE_LOGGING=OFF + EXECUTORCH_PAL_DEFAULT=minimal +) + +# Optimization flags +target_compile_options( + executorch_pico PRIVATE -Os -ffunction-sections -fdata-sections +) + +set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -Wl,--gc-sections") + +set(BAREMETAL_BUILD_DIR ${EXECUTORCH_ROOT}/executorch/cmake-out/) + +# Link ExecuTorch and Pico libraries +target_link_libraries( + executorch_pico + PRIVATE ${BAREMETAL_BUILD_DIR}/lib/libexecutorch.a + ${BAREMETAL_BUILD_DIR}/lib/libexecutorch_core.a + -Wl,--whole-archive + ${BAREMETAL_BUILD_DIR}/lib/libportable_ops_lib.a + -Wl,--no-whole-archive + ${BAREMETAL_BUILD_DIR}/lib/libportable_kernels.a + pico_stdlib + pico_stdio_usb +) + +# Only add extra outputs if the target builds successfully +if(TARGET executorch_pico) + pico_add_extra_outputs(executorch_pico) +endif() diff --git a/examples/rpi/pico2/README.md b/examples/rpi/pico2/README.md new file mode 100644 index 00000000000..8c617cc5aa3 --- /dev/null +++ b/examples/rpi/pico2/README.md @@ -0,0 +1,45 @@ +# Overview +This document outlines the steps required to run a simple Add Module on the Pico2 microcontroller using executorch. + +## (Pre-requisistes) Prepare the Environment for Arm +1. See for instructions on setting up the environment for Arm. +2. Make sure you have the toolchain configured correctly. +```bash +which arm-none-eabi-gcc +--> return something like executorch/examples/arm/ethos-u-scratch/arm-gnu-toolchain-13.3.rel1-x86_64-arm-none-eabi/bin/arm-none-eabi-gcc +``` + +## Build Pico2 Firmware with Executorch + +This step involves two sub steps + +1. Cross Compile Executorch for Arm Cortex M, Pico2 target +2. Build the firmware with the input model provided (If not provided, it will use the default_model.pte) + +Use the following command to build the firmware: +``` bash +executorch/examples/rpi/build_firmware_pico.sh --model= +``` + +### Flash Firmware +Hold the BOOTSEL button on the Pico2 and connect it to your computer; it will mount as RPI-RP2. Copy the executorch_pico.uf2 file to this drive. + +### Verify Execution +Check that the Pico2's LED blinks 10 times at 500 ms interval to confirm successful firmware execution. +The Pico2's LED should blink 10 times at 500 ms intervals, indicating successful firmware execution. If connected via serial, you should see: + +```bash +Method loaded [forward] +Output: 13.000000, 136.000000, 24.000000, 131.000000 +``` + +### Debugging via Serial Terminal +On macOS or Linux, open a serial terminal with: +```bash +screen /dev/tty.usbmodem1101 115200 +``` + +Replace /dev/tty.usbmodem1101 with your device path. This terminal shows program logs and errors. If +the LED blinks 10 times at 100 ms intervals, your program hit an error state—check the logs here. + +These steps complete running the simple model on Pico2 using ExecuTorch. diff --git a/examples/rpi/pico2/main.cpp b/examples/rpi/pico2/main.cpp new file mode 100644 index 00000000000..afec859863c --- /dev/null +++ b/examples/rpi/pico2/main.cpp @@ -0,0 +1,181 @@ +/* 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. + */ + +// Model data +#include "model_pte.h" + +// Pico includes +#include "pico/stdio_usb.h" +#include "pico/stdlib.h" + +// Executorch includes +#include +#include +#include +#include +#include +#include + +// Std c++ includes +#include + +using namespace executorch::runtime; +using executorch::aten::Tensor; +using executorch::aten::TensorImpl; +using ScalarType = executorch::runtime::etensor::ScalarType; +using executorch::runtime::runtime_init; + +// Define GPIO pins for indicators +const uint INDICATOR_PIN_1 = 25; // Onboard LED +const uint INDICATOR_PIN_2 = 22; // External LED +const uint INDICATOR_PIN_3 = 23; // Onboard LED + +void init_gpio_pins() { + gpio_init(INDICATOR_PIN_1); + gpio_set_dir(INDICATOR_PIN_1, GPIO_OUT); + gpio_init(INDICATOR_PIN_2); + gpio_set_dir(INDICATOR_PIN_2, GPIO_OUT); + gpio_init(INDICATOR_PIN_3); + gpio_set_dir(INDICATOR_PIN_3, GPIO_OUT); +} + +void wait_for_usb() { + const int kMaxRetryCount = 10; + int retry_usb_count = 0; + while (!stdio_usb_connected() && retry_usb_count++ < kMaxRetryCount) { + printf("Retry again! USB not connected \n"); + sleep_ms(1000); + } +} + +// Helper function to blink an indicator pin on the pico board a given number of +// times +void blink_indicator(uint pin, int times, int delay_ms = 100) { + for (int i = 0; i < times; ++i) { + gpio_put(pin, 1); + sleep_ms(delay_ms); + gpio_put(pin, 0); + sleep_ms(delay_ms); + } +} + +bool load_and_prepare_model( + std::unique_ptr& program_ptr, + std::unique_ptr& method_ptr, + MemoryManager& memory_manager) { + executorch::extension::BufferDataLoader loader(model_pte, model_pte_len); + auto program_result = Program::load(&loader); + if (!program_result.ok()) { + printf("Failed to load model: %d\n", (int)program_result.error()); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + program_ptr = std::make_unique(std::move(*program_result)); + auto method_name_result = program_ptr->get_method_name(0); + if (!method_name_result.ok()) { + printf("Failed to get method name: %d\n", (int)method_name_result.error()); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + auto method_result = + program_ptr->load_method(*method_name_result, &memory_manager); + if (!method_result.ok()) { + printf("Failed to load method: %d\n", (int)method_result.error()); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + method_ptr = std::make_unique(std::move(*method_result)); + printf("Method loaded [%s]\n", *method_name_result); + return true; +} + +bool validate_and_set_inputs(Method& method) { + float input_data_0[4] = {4.0, 109.0, 13.0, 123.0}; + float input_data_1[4] = {9.0, 27.0, 11.0, 8.0}; + TensorImpl::SizesType sizes[1] = {4}; + TensorImpl::DimOrderType dim_order[] = {0}; + TensorImpl impl0(ScalarType::Float, 1, sizes, input_data_0, dim_order); + TensorImpl impl1(ScalarType::Float, 1, sizes, input_data_1, dim_order); + Tensor input0(&impl0); + Tensor input1(&impl1); + if (method.set_input(input0, 0) != Error::Ok) { + printf("Failed to set input0\n"); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + if (method.set_input(input1, 1) != Error::Ok) { + printf("Failed to set input1\n"); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + return true; +} + +bool run_inference(Method& method) { + if (!validate_and_set_inputs(method)) { + return false; // Input validation or setting failed + } + + if (method.execute() != Error::Ok) { + printf("Failed to execute\n"); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + const EValue& output = method.get_output(0); + if (output.isTensor()) { + const float* out_data = output.toTensor().const_data_ptr(); + printf( + "Output: %f, %f, %f, %f\n", + out_data[0], + out_data[1], + out_data[2], + out_data[3]); + } else { + printf("Output is not a tensor!\n"); + blink_indicator(INDICATOR_PIN_1, 10); + return false; + } + return true; +} + +int executor_runner() { + init_gpio_pins(); + stdio_init_all(); + sleep_ms(1000); + + wait_for_usb(); + runtime_init(); + + static uint8_t method_allocator_pool[32 * 1024]; // 32KB + static uint8_t activation_pool[64 * 1024]; // 64KB + MemoryAllocator method_allocator( + sizeof(method_allocator_pool), method_allocator_pool); + method_allocator.enable_profiling("method allocator"); + Span memory_planned_buffers[1]{ + {activation_pool, sizeof(activation_pool)}}; + HierarchicalAllocator planned_memory({memory_planned_buffers, 1}); + MemoryManager memory_manager(&method_allocator, &planned_memory); + + std::unique_ptr program_ptr; + std::unique_ptr method_ptr; + if (!load_and_prepare_model(program_ptr, method_ptr, memory_manager)) { + printf("Failed to load and prepare model\n"); + return 1; + } + if (!run_inference(*method_ptr)) { + printf("Failed to run inference\n"); + return 1; + } + + // If everything went well, it will blink the indicator pin + blink_indicator(INDICATOR_PIN_1, 10, 500); + return 0; +} + +int main() { + return executor_runner(); +} diff --git a/examples/rpi/pico2/model_pte.c b/examples/rpi/pico2/model_pte.c new file mode 100644 index 00000000000..3208c787380 --- /dev/null +++ b/examples/rpi/pico2/model_pte.c @@ -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. + */ + +#include "model_pte.h" + +/* + * The contents of this file will be auto generated by the cmake command via the helper script (examples/rpi/pte_to_array.py) + * The file then gets compiled into the binary that can be flashed onto pico2 board. + */ +const uint8_t model_pte[] __attribute__((aligned(8))) = { +}; + +const unsigned int model_pte_len = sizeof(model_pte); diff --git a/examples/rpi/pico2/model_pte.h b/examples/rpi/pico2/model_pte.h new file mode 100644 index 00000000000..3146a279baa --- /dev/null +++ b/examples/rpi/pico2/model_pte.h @@ -0,0 +1,21 @@ +/* + * 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 + +#ifdef __cplusplus +extern "C" { +#endif + +extern const uint8_t model_pte[] __attribute__((aligned(8))); +extern const unsigned int model_pte_len; + +#ifdef __cplusplus +} +#endif diff --git a/examples/rpi/pte_to_array.py b/examples/rpi/pte_to_array.py new file mode 100644 index 00000000000..e69df603e39 --- /dev/null +++ b/examples/rpi/pte_to_array.py @@ -0,0 +1,207 @@ +# 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. + +""" +Automated PTE to Array Injector for ExecuTorch +Reads a PTE file and either: +1. Updates model_pte[] array in a C file (when --file points to .c file) +2. Dumps hex array to text file (when --file points to .txt or is not provided) + +Supports multiple array declaration patterns via regex map. +""" + +import argparse +import os +import re +import sys + + +# Define regex patterns for different array declarations +# Each pattern should have 3 capture groups: +# 1. Declaration part (before opening brace) +# 2. Array contents (to be replaced) +# 3. Closing part (closing brace and semicolon) +ARRAY_PATTERNS = { + "const_uint8_model_pte": r"(const\s+uint8_t\s+model_pte\[\].*?=\s*\{)(.*?)(\};)", + "network_model_sec": r'(__attribute__\(\(section\("network_model_sec"\).*?\)\)\s+\w+\s+\w+\[\].*?=\s*\{)(.*?)(\};)', + # Add more patterns here as needed +} + + +def generate_hex_array(pte_file_path): + """Generate formatted hex array contents from PTE file.""" + with open(pte_file_path, "rb") as f: + data = f.read() + + hex_bytes = [f"0x{b:02x}" for b in data] + lines = [] + bytes_per_line = 12 # Format with 12 bytes per line for readability + + for i in range(0, len(hex_bytes), bytes_per_line): + line = ", ".join(hex_bytes[i : i + bytes_per_line]) + lines.append(f" {line}") + + return ",\n".join(lines) + + +def dump_to_text_file(output_file, hex_array_contents): + """Dump hex array contents to a text file.""" + with open(output_file, "w") as f: + f.write(hex_array_contents) + + byte_count = len(hex_array_contents.split(",")) + print(f"Successfully wrote {byte_count} bytes to {output_file}") + + +def find_and_replace_array(content, hex_array_contents): + """Try each regex pattern until one matches, then replace array contents.""" + for pattern_name, pattern in ARRAY_PATTERNS.items(): + match = re.search(pattern, content, re.DOTALL) + if match: + print(f"Found array using pattern: '{pattern_name}'") + # Replace the array contents (keeping the declaration and closing brace) + new_content = ( + content[: match.start(2)] + + f"\n{hex_array_contents}\n" + + content[match.start(3) :] + ) + return new_content, pattern_name + # No pattern matched + return None, None + + +def update_model_file(c_file_path, hex_array_contents): + """Update the model array in the C file with new hex contents.""" + with open(c_file_path, "r") as f: + content = f.read() + + new_content, matched_pattern = find_and_replace_array(content, hex_array_contents) + if new_content is None: + available_patterns = ", ".join(ARRAY_PATTERNS.keys()) + raise ValueError( + f"Could not find any model array pattern in {c_file_path}.\n" + f"Available patterns: {available_patterns}\n" + "Add a new pattern to ARRAY_PATTERNS if needed." + ) + + with open(c_file_path, "w") as f: + f.write(new_content) + + byte_count = len(hex_array_contents.split(",")) + print( + f"Successfully updated {c_file_path} with {byte_count} bytes using pattern '{matched_pattern}'" + ) + + +def add_custom_pattern(name, pattern): + """Add a custom regex pattern to the pattern map.""" + ARRAY_PATTERNS[name] = pattern + print(f"Added custom pattern '{name}': {pattern}") + + +def list_patterns(): + """List all available regex patterns.""" + print("Available regex patterns:") + for name, pattern in ARRAY_PATTERNS.items(): + print(f" {name}: {pattern}") + + +def main(): + parser = argparse.ArgumentParser( + description="Generate hex array from PTE file and either inject into C file or dump to text file", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=( + "Examples:\n" + " # Inject into C file (auto-detects array pattern):\n" + " python3 pte_to_array.py --model model.pte --file model_pte.c\n\n" + " # Dump to text file (default):\n" + " python3 pte_to_array.py --model model.pte\n" + " python3 pte_to_array.py --model model.pte --file my_array.txt\n\n" + " # List available patterns:\n" + " python3 pte_to_array.py --list-patterns\n\n" + " # Add custom pattern:\n" + ' python3 pte_to_array.py --model model.pte --file model.c --add-pattern "my_pattern" "regex_here"\n' + ), + ) + + parser.add_argument( + "--model", + "-m", + help="Path to the input PTE model file", + ) + + parser.add_argument( + "--file", + "-f", + default="model_array.txt", + help="Output file: .c file to inject array into, or .txt file to dump array (default: model_array.txt)", + ) + + parser.add_argument( + "--list-patterns", + action="store_true", + help="List all available regex patterns and exit", + ) + + parser.add_argument( + "--add-pattern", + nargs=2, + metavar=("NAME", "REGEX"), + help="Add a custom regex pattern: --add-pattern 'name' 'regex_pattern'", + ) + + args = parser.parse_args() + + # Handle list patterns command + if args.list_patterns: + list_patterns() + sys.exit(0) + + # Add custom pattern if provided + if args.add_pattern: + add_custom_pattern(args.add_pattern[0], args.add_pattern[1]) + + # Validate required arguments + if not args.model: + print("Error: --model argument is required", file=sys.stderr) + sys.exit(1) + + # Validate input PTE file + if not os.path.exists(args.model): + print(f"Error: PTE file '{args.model}' does not exist", file=sys.stderr) + sys.exit(1) + + try: + # Generate hex array from PTE file + print(f"Reading PTE file: {args.model}") + hex_array = generate_hex_array(args.model) + + # Determine action based on file extension + file_ext = os.path.splitext(args.file)[1].lower() + + if file_ext == ".c": + # Inject into C file + if not os.path.exists(args.file): + print(f"Error: C file '{args.file}' does not exist", file=sys.stderr) + sys.exit(1) + + print(f"Injecting into C file: {args.file}") + update_model_file(args.file, hex_array) + + else: + # Dump to text file (default behavior) + print(f"Dumping hex array to: {args.file}") + dump_to_text_file(args.file, hex_array) + + sys.exit(0) + + except Exception as exc: + print(f"Error: {exc}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/samsung/README.md b/examples/samsung/README.md new file mode 100644 index 00000000000..426dce40670 --- /dev/null +++ b/examples/samsung/README.md @@ -0,0 +1,60 @@ +# 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). + +## Environment +We set up `PYTHONPATH` because it's easier to develop and import executorch Python APIs. +Users might also build and install executorch package as usual python package. +```bash +export PYTHONPATH=${EXECUTORCH_ROOT}/.. +``` + +Note: Since we set `PYTHONPATH`, we may have issue with finding `program.fbs` and `scalar_type.fbs` +when we export a model. A workaround is that we copy them to directory `${EXECUTORCH_ROOT}/exir/_serialize/`. +We can find the files in `${EXECUTORCH_ROOT}/schema` or in +`${EXECUTORCH_ROOT}/pip-out/lib.linux-x86_64-cpython-310/executorch/exir/_serialize`. + +## Device +Prepare an android phone with samsung exynos chip. Use `adb` to connect with the mobile phone. + +Check the chip's version, when lower the model, set the corresponding chip version. + +## Lowering + +Before running an example, please copy python artifacts to target directory `PYTHON_TARGET_DIR`. +If you use `build.sh` script to compile samsung backend, please skip the copy step. + +Set up `PYTHON_TARGET_DIR` to `${EXECUTORCH_ROOT}/backends/samsung/python`. +```bash +cp -rf ${EXECUTORCH_ROOT}/build_x86_64/backends/samsung/Py*.so ${PYTHON_TARGET_DIR} +cp -rf ${EXYNOS_AI_LITECORE_PATH}/python/snc_py_api*.so ${PYTHON_TARGET_DIR} +``` + +Take `EXECUTORCH_ROOT` as work directory and here is an example for ic3. +```bash +python -m executorch.examples.samsung.aot_compiler --chipset E9955 -m ic3 --output_dir ic3_artifact +``` + +## 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} +``` + +#### Step 2: Indicate dynamic linkers and execute model +```bash +adb push ./ic3_exynos_fp32.pte ${DEVICE_DIR} +adb shell "cd ${DEVICE_DIR} \ + && export LD_LIBRARY_PATH=${DEVICE_DIR} \ + && ./enn_executor_runner -model ./ic3_exynos_fp32.pte -input ./ic3_input_0.bin --output_path ." +``` + +`enn_executor_runner` has more usages, please refer to the help message. diff --git a/examples/samsung/aot_compiler.py b/examples/samsung/aot_compiler.py new file mode 100644 index 00000000000..771d74e0b45 --- /dev/null +++ b/examples/samsung/aot_compiler.py @@ -0,0 +1,90 @@ +# 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. + +import argparse +import logging + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.utils.export_utils import ( + to_edge_transform_and_lower_to_enn, +) +from executorch.examples.samsung.utils import save_tensors +from executorch.exir import ExecutorchBackendConfig +from executorch.extension.export_util.utils import save_pte_program + +from ..models import MODEL_NAME_TO_MODEL +from ..models.model_factory import EagerModelFactory + +FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" +logging.basicConfig(level=logging.INFO, format=FORMAT) + +SUPPORT_MODEL_NAMES = [ + "mv2", + "ic3", + "ic4", + "resnet18", + "resnet50", + "mv3", + "edsr", + "dl3", + "vit", + "w2l", +] + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "-c", + "--chipset", + required=True, + help="Samsung chipset, i.e. E9955, etc", + type=str, + ) + parser.add_argument( + "-m", + "--model_name", + required=True, + help=f"Model name. Valid ones: {SUPPORT_MODEL_NAMES}", + ) + parser.add_argument("-o", "--output_dir", default=".", help="output directory") + + args = parser.parse_args() + + if args.model_name not in SUPPORT_MODEL_NAMES: + raise RuntimeError( + f"Model {args.model_name} is not a valid name. or not support yet. " + "In the near future, more example models will be supported. Currently, " + f"Available models are {SUPPORT_MODEL_NAMES}." + ) + + model, example_inputs, dynamic_shapes, _ = EagerModelFactory.create_model( + *MODEL_NAME_TO_MODEL[args.model_name] + ) + assert ( + dynamic_shapes is None + ), "enn backend doesn't support dynamic shapes currently." + + model = model.eval() + outputs = model(*example_inputs) + + compile_specs = [gen_samsung_backend_compile_spec(args.chipset)] + edge = to_edge_transform_and_lower_to_enn( + model, example_inputs, compile_specs=compile_specs + ) + + exec_prog = edge.to_executorch( + config=ExecutorchBackendConfig(extract_delegate_segments=True) + ) + + model_name = f"{args.model_name}_exynos_fp32" + save_pte_program(exec_prog, model_name, args.output_dir) + + save_tensors(example_inputs, f"{args.model_name}_input", args.output_dir) + save_tensors(outputs, f"{args.model_name}_output", args.output_dir) diff --git a/examples/samsung/executor_runner/enn_executor_runner.cpp b/examples/samsung/executor_runner/enn_executor_runner.cpp new file mode 100644 index 00000000000..5facd6301f3 --- /dev/null +++ b/examples/samsung/executor_runner/enn_executor_runner.cpp @@ -0,0 +1,277 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * 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. + * + */ + +/** + * @file + * + * This tool can run ExecuTorch model files with Enn runtime. + * It assumes all inputs and output are fp32, please give a list for input + * files. And Enn backends is going to inference, and output results. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +static uint8_t method_allocator_pool[4 * 1024U * 1024U]; // 4 MB + +DEFINE_string(model, "model.pte", "Model serialized in flatbuffer format."); +DEFINE_string( + input, + "", + "Input file path, support multiple inputs: input_1 input_2 ..."); + +DEFINE_string(output_path, "", "Output Execution results to target directory."); + +using namespace torch::executor; +using torch::executor::util::FileDataLoader; + +std::vector split(std::string str, char delimiter = ' ') { + std::vector result; + std::stringstream ss(str); + std::string temp; + while (std::getline(ss, temp, delimiter)) { + if (!temp.empty()) { + result.push_back(temp); + } + } + return result; +} + +class DataReader { + public: + typedef std::vector data_t; + + DataReader(size_t size) : data_set_(size) {} + + void read(const std::string file_path) { + ET_CHECK(index_ < data_set_.size()); + data_t& data = data_set_[index_]; + std::ifstream input_file(file_path.c_str(), std::ios::binary); + ET_CHECK(input_file.is_open()); + input_file.seekg(0, std::ios::end); + data.resize(input_file.tellg()); + input_file.seekg(0); + input_file.read(reinterpret_cast(data.data()), data.size()); + input_file.close(); + ++index_; + } + + void* get(int32_t index) { + ET_CHECK(index < data_set_.size()); + return data_set_[index].data(); + } + + size_t nbytes(int32_t index) { + ET_CHECK(index < data_set_.size()); + return data_set_[index].size(); + } + + ~DataReader() = default; + + private: + std::vector data_set_; + int32_t index_ = 0; +}; + +void saveOutput(const exec_aten::Tensor& tensor, int32_t output_index) { + if (FLAGS_output_path.empty()) { + return; + } + auto output_file_name = + FLAGS_output_path + "/output_" + std::to_string(output_index) + ".bin"; + std::ofstream fout(output_file_name.c_str(), std::ios::binary); + ET_CHECK_MSG( + fout.is_open(), + "Directory or have no visit permission: %s", + FLAGS_output_path.c_str()); + fout.write(tensor.const_data_ptr(), tensor.nbytes()); + fout.close(); +} + +int main(int argc, char** argv) { + runtime_init(); + + gflags::ParseCommandLineFlags(&argc, &argv, true); + if (argc != 1) { + std::string msg = "Extra commandline args:"; + for (int i = 1 /* skip argv[0] (program name) */; i < argc; i++) { + msg += std::string(" ") + argv[i]; + } + ET_LOG(Error, "%s", msg.c_str()); + return 1; + } + + // Create a loader to get the data of the program file. There are other + // 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); + ET_CHECK_MSG( + loader.ok(), + "FileDataLoader::from() failed: 0x%" PRIx32, + (uint32_t)loader.error()); + + // Parse the program file. This is immutable, and can also be reused between + // multiple execution invocations across multiple threads. + Result program = Program::load(&loader.get()); + if (!program.ok()) { + ET_LOG(Error, "Failed to parse model file %s", model_path); + return 1; + } + ET_LOG(Info, "Model file %s is loaded.", model_path); + + // Use the first method in the program. + const char* method_name = nullptr; + { + const auto method_name_result = program->get_method_name(0); + ET_CHECK_MSG(method_name_result.ok(), "Program has no methods"); + method_name = *method_name_result; + } + ET_LOG(Info, "Using method %s", method_name); + + // MethodMeta describes the memory requirements of the method. + Result method_meta = program->method_meta(method_name); + ET_CHECK_MSG( + method_meta.ok(), + "Failed to get method_meta for %s: 0x%" PRIx32, + method_name, + (uint32_t)method_meta.error()); + + // + // The runtime does not use malloc/new; it allocates all memory using the + // MemoryManger provided by the client. Clients are responsible for allocating + // the memory ahead of time, or providing MemoryAllocator subclasses that can + // do it dynamically. + // + + // The method allocator is used to allocate all dynamic C++ metadata/objects + // used to represent the loaded method. This allocator is only used during + // loading a method of the program, which will return an error if there was + // not enough memory. + // + // The amount of memory required depends on the loaded method and the runtime + // code itself. The amount of memory here is usually determined by running the + // method and seeing how much memory is actually used, though it's possible to + // subclass MemoryAllocator so that it calls malloc() under the hood (see + // MallocMemoryAllocator). + // + // In this example we use a statically allocated memory pool. + MemoryAllocator method_allocator{ + MemoryAllocator(sizeof(method_allocator_pool), method_allocator_pool)}; + + // The memory-planned buffers will back the mutable tensors used by the + // method. The sizes of these buffers were determined ahead of time during the + // memory-planning pasees. + // + // Each buffer typically corresponds to a different hardware memory bank. Most + // mobile environments will only have a single buffer. Some embedded + // environments may have more than one for, e.g., slow/large DRAM and + // fast/small SRAM, or for memory associated with particular cores. + std::vector> planned_buffers; // Owns the memory + std::vector> planned_spans; // Passed to the allocator + size_t num_memory_planned_buffers = method_meta->num_memory_planned_buffers(); + for (size_t id = 0; id < num_memory_planned_buffers; ++id) { + // .get() will always succeed because id < num_memory_planned_buffers. + size_t buffer_size = + static_cast(method_meta->memory_planned_buffer_size(id).get()); + ET_LOG(Info, "Setting up planned buffer %zu, size %zu.", id, buffer_size); + planned_buffers.push_back(std::make_unique(buffer_size)); + planned_spans.push_back({planned_buffers.back().get(), buffer_size}); + } + HierarchicalAllocator planned_memory( + {planned_spans.data(), planned_spans.size()}); + + // Assemble all of the allocators into the MemoryManager that the Executor + // will use. + MemoryManager memory_manager(&method_allocator, &planned_memory); + + // + // Load the method from the program, using the provided allocators. Running + // the method can mutate the memory-planned buffers, so the method should only + // be used by a single thread at at time, but it can be reused. + // + + Result method = program->load_method(method_name, &memory_manager); + ET_CHECK_MSG( + method.ok(), + "Loading of method %s failed with status 0x%" PRIx32, + method_name, + (uint32_t)method.error()); + + auto input_files = split(FLAGS_input); + ET_CHECK_MSG( + input_files.size() == method->inputs_size(), + "Please check the number of given input binary files"); + DataReader input_data_reader(input_files.size()); + for (const auto& input_file : input_files) { + input_data_reader.read(input_file); + } + + for (int input_index = 0; input_index < method->inputs_size(); + ++input_index) { + MethodMeta method_meta = method->method_meta(); + Result tensor_meta = method_meta.input_tensor_meta(input_index); + ET_CHECK_MSG( + input_data_reader.nbytes(input_index) == tensor_meta->nbytes(), + "Given inputs size is invalid"); + TensorImpl impl = TensorImpl( + tensor_meta->scalar_type(), + tensor_meta->sizes().size(), + const_cast(tensor_meta->sizes().data()), + input_data_reader.get(input_index), + const_cast(tensor_meta->dim_order().data())); + Error ret = method->set_input(Tensor(&impl), input_index); + ET_CHECK_MSG(ret == Error::Ok, "Failed to set input tensor: %d", ret); + } + // Allocate input tensors and set all of their elements to 1. The `inputs` + // variable owns the allocated memory and must live past the last call to + // `execute()`. + // auto inputs = util::prepare_input_tensors(*method); + + // Run the model. + ET_LOG(Info, "Start inference."); + auto start = std::chrono::high_resolution_clock::now(); + Error status = method->execute(); + auto end = std::chrono::high_resolution_clock::now(); + double elapse = + std::chrono::duration_cast(end - start) + .count() / + 1000.0; + ET_CHECK_MSG( + status == Error::Ok, + "Execution of method %s failed with status 0x%" PRIx32, + method_name, + static_cast(status)); + ET_LOG(Info, "End with elapsed time(ms): %f", elapse); + + // Get the outputs. + std::vector outputs(method->outputs_size()); + status = method->get_outputs(outputs.data(), outputs.size()); + ET_CHECK(status == Error::Ok); + + for (size_t output_index = 0; output_index < method->outputs_size(); + ++output_index) { + auto output_tensor = outputs[output_index].toTensor(); + // Save the results to given directory in order. + saveOutput(output_tensor, output_index); + } + + return 0; +} diff --git a/examples/samsung/scripts/mobilebert_finetune.py b/examples/samsung/scripts/mobilebert_finetune.py new file mode 100644 index 00000000000..e0c11df246a --- /dev/null +++ b/examples/samsung/scripts/mobilebert_finetune.py @@ -0,0 +1,261 @@ +# 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. + +import argparse +import os +from pathlib import Path + +import torch + +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) +from executorch.backends.samsung.utils.export_utils import ( + to_edge_transform_and_lower_to_enn, +) +from executorch.examples.samsung.utils import save_tensors +from executorch.exir import ExecutorchBackendConfig +from executorch.extension.export_util.utils import save_pte_program +from transformers import AutoTokenizer, MobileBertForSequenceClassification + + +# Output from pretrained model exceeds the representation scale of half-float. +# Finetune bert model on specific task and make output more reasonable for hardware. +# Here is an example. +class MobileBertFinetune: + def __init__(self): + self.tokenizer = self.load_tokenizer() + + def load_tokenizer(self): + return AutoTokenizer.from_pretrained("google/mobilebert-uncased") + + def get_example_inputs(self): + encoding = self.tokenizer.encode_plus( + "Hello, my dog is cute", + add_special_tokens=True, + max_length=128, + return_token_type_ids=False, + return_attention_mask=True, + truncation=True, + return_tensors="pt", + padding="max_length", + ) + + return ( + encoding["input_ids"], + encoding["attention_mask"].to(torch.float32), + ) + + def build_loader_from_dataset(self, dataset, batch_size, usage="train"): + """ + :param data: Provide dataset in pandas table type. The header names should be ['text', 'label'], + and label range from 0 (include) to total number of classification kinds (not include). + For example: + index text label + 0 despite its title , punch drunk love is never heavy handed 1 + 1 at once half baked and overheated 0 + 2 this is a shameless sham, ... 0 + ... + :param batch_size: Size of data fetch in one batch. + :param usage: The type of dataset which is used to build dataloader, like train, val. + :return: dataloader + """ + from torch.utils.data import ( + DataLoader, + RandomSampler, + SequentialSampler, + TensorDataset, + ) + + encoded_dataset = self.tokenizer.batch_encode_plus( + dataset.text.values.tolist(), + return_attention_mask=True, + truncation=True, + padding="max_length", + max_length=128, + return_tensors="pt", + ) + + labels = torch.tensor(dataset.label.values.tolist()) + + tensor_dataset = TensorDataset( + encoded_dataset["input_ids"], encoded_dataset["attention_mask"], labels + ) + data_loader = None + if usage == "train": + data_loader = DataLoader( + tensor_dataset, + sampler=RandomSampler(tensor_dataset), + batch_size=batch_size, + ) + elif usage == "val": + data_loader = DataLoader( + tensor_dataset, + sampler=SequentialSampler(tensor_dataset), + batch_size=batch_size, + drop_last=True, + ) + else: + raise NotImplementedError( + f"Unsupported `{usage}` dataset for building dataloader." + ) + + return data_loader + + def get_finetune_mobilebert(self, artifacts_dir): + # 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. + from io import BytesIO + + import pandas as pd + import requests + + from tqdm import tqdm + from transformers import get_linear_schedule_with_warmup + + # sentiment classification + train_url = "https://raw.githubusercontent.com/clairett/pytorch-sentiment-classification/refs/heads/master/data/SST2/train.tsv" + content = requests.get(train_url, allow_redirects=True).content + train_data = pd.read_csv( + BytesIO(content), delimiter="\t", header=None, names=["text", "label"] + ) + labels_set = train_data.label.unique() + + train_data_loader = self.build_loader_from_dataset( + train_data, batch_size=64, usage="train" + ) + + val_url = "https://raw.githubusercontent.com/clairett/pytorch-sentiment-classification/refs/heads/master/data/SST2/test.tsv" + content = requests.get(val_url, allow_redirects=True).content + val_data = pd.read_csv( + BytesIO(content), delimiter="\t", header=None, names=["text", "label"] + ) + val_data_loader = self.build_loader_from_dataset( + val_data, batch_size=64, usage="val" + ) + + artifacts_dir = artifacts_dir if artifacts_dir is not None else "./mobilebert" + need_finetune = True + os.makedirs(artifacts_dir, exist_ok=True) + pretrained_required_files = ["config.json", "model.safetensors"] + path = Path(artifacts_dir) + if (path / pretrained_required_files[0]).exists() and ( + path / pretrained_required_files[1] + ).exists(): + need_finetune = False + + # get pre-trained mobilebert + model = MobileBertForSequenceClassification.from_pretrained( + "google/mobilebert-uncased" if need_finetune else artifacts_dir, + num_labels=len(labels_set), + return_dict=False, + ) + + if not need_finetune: + return model.eval(), val_data_loader + + num_epochs = 5 + num_train_steps = len(train_data_loader) * num_epochs + + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5) + scheduler = get_linear_schedule_with_warmup( + optimizer, num_warmup_steps=0, num_training_steps=num_train_steps + ) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + for epoch in range(1, num_epochs + 1): + print(f"Epoch {epoch}") + model.train() + for batch in tqdm(train_data_loader): + texts, attention_mask, labels = batch + texts = texts.to(device) + labels = labels.to(device) + + loss = model(texts, attention_mask=attention_mask, labels=labels)[0] + # backward + optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) + optimizer.step() + # update learning rate + scheduler.step() + model.to("cpu") + + model.save_pretrained(artifacts_dir) + + return model.eval(), val_data_loader + + def validate(self, model, val_data_loader): + model.eval() + total_loss = 0 + correct = 0 + total = 0 + with torch.no_grad(): + for batch in val_data_loader: + texts, attention_mask, labels = batch + + loss, output = model( + texts, attention_mask=attention_mask, labels=labels + ) + total_loss += loss.item() + predictions = torch.argmax(output, dim=1) + correct += (predictions == labels).sum().item() + total += labels.size(0) + + val_loss, val_accuracy = total_loss / len(val_data_loader), correct / total + print(f"Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "-c", + "--chipset", + required=True, + help="Samsung chipset, i.e. E9955, etc", + type=str, + ) + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts by this example.", + default="./mobilebert", + type=str, + ) + parser.add_argument( + "--dump", + default=False, + action="store_true", + help=("Whether to dump all outputs. If not set, we only dump pte."), + ) + args = parser.parse_args() + # ensure the working directory exist. + os.makedirs(args.artifact, exist_ok=True) + + mobilebert_finetune = MobileBertFinetune() + model, val_dataset = mobilebert_finetune.get_finetune_mobilebert(args.artifact) + mobilebert_finetune.validate(model, val_dataset) + + example_inputs = mobilebert_finetune.get_example_inputs() + output = model(*example_inputs) + + compile_specs = [gen_samsung_backend_compile_spec(args.chipset)] + edge = to_edge_transform_and_lower_to_enn( + model, example_inputs, compile_specs=compile_specs + ) + model_name = "mobilebert_exynos_fp32" + exec_prog = edge.to_executorch( + config=ExecutorchBackendConfig(extract_delegate_segments=True) + ) + save_pte_program(exec_prog, model_name, args.artifact) + + if args.dump: + # Expect example inputs are tuple, including input ids and attn mask + save_tensors(example_inputs, prefix="float_input", artifact_dir=args.artifact) + save_tensors(output, prefix="float_output", artifact_dir=args.artifact) diff --git a/examples/samsung/utils.py b/examples/samsung/utils.py new file mode 100644 index 00000000000..4f2cc56086c --- /dev/null +++ b/examples/samsung/utils.py @@ -0,0 +1,20 @@ +import collections +import logging +import os + +import torch + + +def save_tensors(tensors, prefix, artifact_dir): + if isinstance(tensors, tuple): + for index, output in enumerate(tensors): + save_path = prefix + "_" + str(index) + ".bin" + output.detach().numpy().tofile(os.path.join(artifact_dir, save_path)) + elif isinstance(tensors, torch.Tensor): + tensors.detach().numpy().tofile(os.path.join(artifact_dir, prefix + ".bin")) + elif isinstance(tensors, collections.OrderedDict): + for index, output in enumerate(tensors.values()): + save_path = prefix + "_" + str(index) + ".bin" + output.detach().numpy().tofile(os.path.join(artifact_dir, save_path)) + else: + logging.warning("Unsupported type (", type(tensors), ") skip saving tensor. ") diff --git a/examples/selective_build/README.md b/examples/selective_build/README.md index 888020ce770..c6c8dc1ba57 100644 --- a/examples/selective_build/README.md +++ b/examples/selective_build/README.md @@ -1,23 +1,137 @@ # Selective Build Examples -To optimize binary size of ExecuTorch runtime, selective build can be used. This folder contains examples to select only the operators needed for ExecuTorch build. This example will demonstrate the CMake build. +To optimize binary size of ExecuTorch runtime, selective build can be used. This folder contains examples to select only the operators needed for ExecuTorch build. -## How to run +These examples showcase two flows - the simple way, using CMake options to configure the framework build, and an advanced flow - showcasing user-defined kernel targets including custom operators. Prerequisite: finish the [setting up wiki](https://pytorch.org/executorch/main/getting-started-setup). -Run: +## Example 1 - Basic Flow -```bash -cd executorch -bash examples/selective_build/test_selective_build.sh cmake +This example showcases using CMake options to control which operators are included. This approach should be preferred when not using +custom operators or additional kernel libraries beyond the standard kernels provided by ExecuTorch. + +The code under the basic/ directory builds a minimal model runner binary which links to a selective kernel target. To build the +example with operators needed for the MobileNetV2 model, run the following commands: +``` +# From the executorch directory +python -m examples.portable.scripts.export --model_name="mv2" # Create a PTE file for MobileNetV2 +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 +./selective_build_test --model_path="../../mv2.pte" # Run the model with the selective kernel library +``` + +### CMake Options + +The example commands above show use of the EXECUTORCH_SELECT_OPS_MODEL option to select operators used in a PTE file, but there are +several ways to provide the operator list. The options can be passed to CMake in the same way (during configuration) and are mutually +exclusive, meaning that only one of these options should be chosen. + + * `EXECUTORCH_SELECT_OPS_MODEL`: Select operators used in a .PTE file. Takes a path to the file. + * `EXECUTORCH_SELECT_OPS_YAML`: Provide a list of operators from a .yml file, typically generated with the `codegen/tools/gen_oplist.py` script. See this script for usage information. + * `EXECUTORCH_SELECT_OPS_LIST`: Provide a comma-separated list of operators to include. An example is included below. + +Example operator list specification (passed as a CLI arg to the CMake configure command): +``` +-DEXECUTORCH_SELECT_OPS_LIST="aten::convolution.out,\ +aten::_native_batch_norm_legit_no_training.out,aten::hardtanh.out,aten::add.out,\ +aten::mean.out,aten::view_copy.out,aten::permute_copy.out,aten::addmm.out,\ +aten,aten::clone.out" +``` + +#### DType-Selective Build + +To further reduce binary size, ExecuTorch can specialize the individual operators for only the dtypes (data types) used. For example, if +the model only calls add with 32-bit floating point tensors, it can drop parts of the code that handle integer tensors or other floating point types. This option is controlled by passing `-DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON` to CMake. It is only supported in conjunction +with the `EXECUTORCH_SELECT_OPS_MODEL` option and is not yet supported for other modes. It is recommended to enable this option when using `EXECUTORCH_SELECT_OPS_MODEL` as it provides significant size savings on top of the kernel selective build. + +### How it Works + +The CMake options described above are read by ExecuTorch framework build, which is referenced via `add_subdirectory` in basic/CMakeLists.txt. These options reflect in the `executorch_kernels` CMake target, which is linked against the example binary. + +```cmake +# basic/CMakeLists.txt +target_link_libraries( + selective_build_test + PRIVATE executorch_core extension_evalue_util extension_runner_util + gflags::gflags executorch_kernels +) +``` + +To use selective build in a user CMake project, take the following steps: + * Reference the executorch framework via `add_subdirectory`. + * Add `executorch_kernels` as a dependency (via target_link_libraries). + * Set CMake options at build time or in user CMake code. + +To use the CMake-build framework libraries from outside of the CMake ecosystem, link against libexecutorch_selected_kernels. + +## Example 2 - Advanced Flow for Custom Ops and Kernel Libraries + +This example showcases defined a custom kernel target. This option can be used when defining custom operators or integrating with +kernel libraries not part of the standard ExecuTorch build. + +The code under the advanced/ directory builds a minimal model runner binary which links to a user-defined kernel library target. To run a model with a simple custom operator, run the following commands: +``` +# From the executorch directory +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 +./selective_build_test --model_path="../../custom_ops_1.pte" # Run the model with the selective kernel library +``` + +### CMake Options + +The CMake logic in `advanced/CMakeLists.txt` respects the CMake options described in the basic flow, as well as the following options: + + * `EXECUTORCH_EXAMPLE_USE_CUSTOM_OPS`: Build and link some simple custom operators. + * `EXECUTORCH_EXAMPLE_SELECT_ALL_OPS`: Build a kernel target with all available operators. + +### How it Works + +The build logic in `advanced/CMakeLists.txt` uses the `gen_selected_ops`, `generate_bindings_for_kernels`, and `gen_operators_lib` CMake functions to define an operator target. See [Kernel Library Selective Build](https://docs.pytorch.org/executorch/main/kernel-library-selective-build.html) for more information on selective build. + +```cmake +gen_selected_ops( + LIB_NAME + "select_build_lib" + OPS_SCHEMA_YAML + "${_custom_ops_yaml}" + ROOT_OPS + "${EXECUTORCH_SELECT_OPS_LIST}" + INCLUDE_ALL_OPS + "${EXECUTORCH_SELECT_ALL_OPS}" + OPS_FROM_MODEL + "${EXECUTORCH_SELECT_OPS_MODEL}" + DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}" +) + +generate_bindings_for_kernels( + LIB_NAME + "select_build_lib" + FUNCTIONS_YAML + ${EXECUTORCH_ROOT}/kernels/portable/functions.yaml + CUSTOM_OPS_YAML + "${_custom_ops_yaml}" + DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}" +) + +gen_operators_lib( + LIB_NAME + "select_build_lib" + KERNEL_LIBS + ${_kernel_lib} + DEPS + executorch_core + DTYPE_SELECTIVE_BUILD + "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}" +) ``` -Check out `CMakeLists.txt` for demo of selective build APIs: -1. `SELECT_ALL_OPS`: Select all ops from the dependency kernel libraries, register all of them into ExecuTorch runtime. -2. `SELECT_OPS_LIST`: Only select operators from a list. -3. `SELECT_OPS_YAML`: Only select operators from a yaml file. -4. `SELECT_OPS_FROM_MODEL`: Only select operators from a from an exported model pte. -5. `DTYPE_SELECTIVE_BUILD`: Enable rebuild of `portable_kernels` to use dtype selection. Currently only supported for `SELECTED_OPS_FROM_MODEL` API and `portable_kernels` lib. +To link against this target, the top-level binary target declares a dependency on `select_build_lib`, which is the library name defined by the above function invocations. To use outside of the CMake ecosystem, link against libselect_build_lib. -Other configs: -- `MAX_KERNEL_NUM=N`: Only allocate memory for N operators. +See `test_selective_build.sh` for additional build examples. diff --git a/examples/selective_build/CMakeLists.txt b/examples/selective_build/advanced/CMakeLists.txt similarity index 66% rename from examples/selective_build/CMakeLists.txt rename to examples/selective_build/advanced/CMakeLists.txt index 3cc5e759ac4..65ebb50bcac 100644 --- a/examples/selective_build/CMakeLists.txt +++ b/examples/selective_build/advanced/CMakeLists.txt @@ -18,7 +18,7 @@ cmake_minimum_required(VERSION 3.19) project(selective_build_example) -set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) set(TORCH_ROOT ${EXECUTORCH_ROOT}/third-party/pytorch) include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) @@ -37,62 +37,37 @@ set(_common_compile_options -Wno-deprecated-declarations -fPIC -ffunction-sections -fdata-sections ) -# Let files say "include ". -set(_common_include_directories ${EXECUTORCH_ROOT}/..) - -find_package(executorch CONFIG REQUIRED) -find_package( - gflags REQUIRED PATHS ${CMAKE_CURRENT_BINARY_DIR}/../../third-party -) - -target_include_directories( - executorch_core INTERFACE ${_common_include_directories} -) +add_subdirectory(${EXECUTORCH_ROOT} ${CMAKE_CURRENT_BINARY_DIR}/executorch) # ------------------------------ OPTIONS BEGIN ------------------------------- -# Option to register ops from yaml file -option(EXECUTORCH_SELECT_OPS_YAML "Register all the ops from a given yaml file" - OFF -) - -# Option to register op list -option(EXECUTORCH_SELECT_OPS_LIST "Register a list of ops, separated by comma" - OFF -) - # Selective build options. -option(EXECUTORCH_SELECT_ALL_OPS +option(EXECUTORCH_EXAMPLE_SELECT_ALL_OPS "Whether to register all ops defined in portable kernel library." OFF ) -# Option to enable parsing ops and dtypes directly from model pte file -option(EXECUTORCH_SELECT_OPS_FROM_MODEL - "Enable op selection from pte during build." OFF +option(EXECUTORCH_EXAMPLE_USE_CUSTOM_OPS + "Whether to include custom ops in the example." OFF ) -# Option to enable dtype selective build. Note: must be using selective build -# model API. -option(EXECUTORCH_DTYPE_SELECTIVE_BUILD "Enable dtype selective build." OFF) -# ------------------------------- OPTIONS END -------------------------------- - +# Note that the following options are defined by the core framework and are also +# used by this example when defining a custom operator target: # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. -# -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../executorch_srcs.cmake" -) - -extract_sources(${EXECUTORCH_SRCS_FILE}) +# EXECUTORCH_SELECT_OPS_YAML EXECUTORCH_SELECT_OPS_LIST +# EXECUTORCH_SELECT_OPS_MODEL EXECUTORCH_DTYPE_SELECTIVE_BUILD -include(${EXECUTORCH_SRCS_FILE}) +# ------------------------------- OPTIONS END -------------------------------- # -# select_build_lib: C++ library to register selected ops in custom kernel -# library +# The `__srcs` lists are defined by executorch_load_build_variables. # +executorch_load_build_variables() + +# For advanced use cases, we can define a custom operator target. This is useful +# when using custom operators. set(_kernel_lib) -if(EXECUTORCH_SELECT_OPS_YAML) + +if(EXECUTORCH_EXAMPLE_USE_CUSTOM_OPS) set(_custom_ops_yaml ${EXECUTORCH_ROOT}/examples/portable/custom_ops/custom_ops.yaml ) @@ -122,7 +97,7 @@ gen_selected_ops( INCLUDE_ALL_OPS "${EXECUTORCH_SELECT_ALL_OPS}" OPS_FROM_MODEL - "${EXECUTORCH_SELECT_OPS_FROM_MODEL}" + "${EXECUTORCH_SELECT_OPS_MODEL}" DTYPE_SELECTIVE_BUILD "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}" ) @@ -149,6 +124,9 @@ gen_operators_lib( "${EXECUTORCH_DTYPE_SELECTIVE_BUILD}" ) +executorch_target_link_options_shared_lib(select_build_lib) +set(selected_kernel_target select_build_lib) + list(TRANSFORM _executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/") # @@ -160,8 +138,8 @@ if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") target_link_options_gc_sections(selective_build_test) endif() target_link_libraries( - selective_build_test PRIVATE executorch_core extension_evalue_util - extension_runner_util gflags select_build_lib + selective_build_test + PRIVATE executorch_core extension_evalue_util extension_runner_util + gflags::gflags ${selected_kernel_target} ) -executorch_target_link_options_shared_lib(select_build_lib) target_compile_options(selective_build_test PUBLIC ${_common_compile_options}) diff --git a/examples/selective_build/basic/CMakeLists.txt b/examples/selective_build/basic/CMakeLists.txt new file mode 100644 index 00000000000..3cc68ad53b6 --- /dev/null +++ b/examples/selective_build/basic/CMakeLists.txt @@ -0,0 +1,77 @@ +# 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. + +# +# Simple CMake build system for selective build demo. +# +# ### Editing this file ### +# +# This file should be formatted with +# ~~~ +# cmake-format -i CMakeLists.txt +# ~~~ +# It should also be cmake-lint clean. +# +cmake_minimum_required(VERSION 3.19) +project(selective_build_example) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) +set(TORCH_ROOT ${EXECUTORCH_ROOT}/third-party/pytorch) + +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_CXX_STANDARD) + set(CMAKE_CXX_STANDARD 17) + # Can't set to 11 due to executor_runner.cpp make_unique +endif() + +set(_common_compile_options -Wno-deprecated-declarations -fPIC + -ffunction-sections -fdata-sections +) + +add_subdirectory(${EXECUTORCH_ROOT} ${CMAKE_CURRENT_BINARY_DIR}/executorch) + +# ------------------------------ OPTIONS BEGIN ------------------------------- + +# The following options are defined by the core framework and are also used in +# the generated kernel target. +# +# EXECUTORCH_SELECT_OPS_YAML EXECUTORCH_SELECT_OPS_LIST +# EXECUTORCH_SELECT_OPS_MODEL EXECUTORCH_DTYPE_SELECTIVE_BUILD + +# ------------------------------- OPTIONS END -------------------------------- + +# +# The `__srcs` lists are defined by executorch_load_build_variables. +# +executorch_load_build_variables() + +# For most use cases, we can configure the ExecuTorch kernel library build using +# the EXECUTORCH_SELECT_OPS_* variables. This will reflect in the +# executorch_kernels target, which includes the configured kernel libraries, +# including selective build, where supported. + +list(TRANSFORM _executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/") + +# +# selective_build_test: test binary to allow different operator libraries to +# link to +# +add_executable(selective_build_test ${_executor_runner__srcs}) +if(NOT CMAKE_BUILD_TYPE STREQUAL "Debug") + target_link_options_gc_sections(selective_build_test) +endif() +target_link_libraries( + selective_build_test + PRIVATE executorch_core extension_evalue_util extension_runner_util + gflags::gflags executorch_kernels +) +target_compile_options(selective_build_test PUBLIC ${_common_compile_options}) diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh old mode 100644 new mode 100755 index eb9038b7cdd..942df7eae04 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -85,35 +85,11 @@ test_buck2_select_ops_from_yaml() { } # CMake examples; test in OSS. Check the README for more information. -test_cmake_select_all_ops() { - echo "Exporting MobilenetV3" - ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv3" - - local example_dir=examples/selective_build - local build_dir=cmake-out/${example_dir} - rm -rf ${build_dir} - retry cmake -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_SELECT_ALL_OPS=ON \ - -DCMAKE_INSTALL_PREFIX=cmake-out \ - -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ - -B${build_dir} \ - ${example_dir} - - echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release - - echo 'Running selective build test' - ${build_dir}/selective_build_test --model_path="./mv3.pte" - - echo "Removing mv3.pte" - rm "./mv3.pte" -} - test_cmake_select_ops_in_list() { echo "Exporting MobilenetV2" ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="mv2" - local example_dir=examples/selective_build + local example_dir=examples/selective_build/basic local build_dir=cmake-out/${example_dir} # set MAX_KERNEL_NUM=22: 19 primops, add, mul rm -rf ${build_dir} @@ -141,11 +117,12 @@ aten,aten::clone.out" \ test_cmake_select_ops_in_yaml() { echo "Exporting custom_op_1" ${PYTHON_EXECUTABLE} -m examples.portable.custom_ops.custom_ops_1 - local example_dir=examples/selective_build + local example_dir=examples/selective_build/advanced local build_dir=cmake-out/${example_dir} rm -rf ${build_dir} retry cmake -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_SELECT_OPS_YAML=ON \ + -DEXECUTORCH_EXAMPLE_USE_CUSTOM_OPS=ON \ + -DEXECUTORCH_EXAMPLE_DEFINE_CUSTOM_TARGET=ON \ -DCMAKE_INSTALL_PREFIX=cmake-out \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -B${build_dir} \ @@ -166,11 +143,11 @@ test_cmake_select_ops_in_model() { local model_export_name="${model_name}.pte" echo "Exporting ${model_name}" ${PYTHON_EXECUTABLE} -m examples.portable.scripts.export --model_name="${model_name}" - local example_dir=examples/selective_build + local example_dir=examples/selective_build/basic local build_dir=cmake-out/${example_dir} rm -rf ${build_dir} retry cmake -DCMAKE_BUILD_TYPE="$CMAKE_BUILD_TYPE" \ - -DEXECUTORCH_SELECT_OPS_FROM_MODEL="./${model_export_name}" \ + -DEXECUTORCH_SELECT_OPS_MODEL="./${model_export_name}" \ -DEXECUTORCH_DTYPE_SELECTIVE_BUILD=ON \ -DEXECUTORCH_OPTIMIZE_SIZE=ON \ -DCMAKE_INSTALL_PREFIX=cmake-out \ @@ -206,7 +183,6 @@ fi if [[ $1 == "cmake" ]]; then cmake_install_executorch_lib $CMAKE_BUILD_TYPE - test_cmake_select_all_ops test_cmake_select_ops_in_list test_cmake_select_ops_in_yaml test_cmake_select_ops_in_model diff --git a/examples/vulkan/README.md b/examples/vulkan/README.md new file mode 100644 index 00000000000..71fdd0e4183 --- /dev/null +++ b/examples/vulkan/README.md @@ -0,0 +1,80 @@ +# Vulkan Delegate Export Examples + +This directory contains scripts for exporting models with the Vulkan delegate in ExecuTorch. Vulkan delegation allows you to run your models on devices with Vulkan-capable GPUs, potentially providing significant performance improvements over CPU execution. + +## Scripts + +- `export.py`: Basic export script for models to use with Vulkan delegate +- `aot_compiler.py`: Advanced export script with quantization support + +## Usage + +### Basic Export + +```bash +python -m executorch.examples.vulkan.export -m -o +``` + +### Export with Quantization (Experimental) + +```bash +python -m executorch.examples.vulkan.aot_compiler -m -q -o +``` + +### Dynamic Shape Support + +```bash +python -m executorch.examples.vulkan.export -m -d -o +``` + +### Additional Options + +- `-s/--strict`: Export with strict mode (default: True) +- `-a/--segment_alignment`: Specify segment alignment in hex (default: 0x1000) +- `-e/--external_constants`: Save constants in external .ptd file (default: False) +- `-r/--etrecord`: Generate and save an ETRecord to the given file location + +## Examples + +```bash +# Export MobileNetV2 with Vulkan delegate +python -m executorch.examples.vulkan.export -m mobilenet_v2 -o ./exported_models + +# Export MobileNetV3 with quantization +python -m executorch.examples.vulkan.aot_compiler -m mobilenet_v3 -q -o ./exported_models + +# Export with dynamic shapes +python -m executorch.examples.vulkan.export -m mobilenet_v2 -d -o ./exported_models + +# Export with ETRecord for debugging +python -m executorch.examples.vulkan.export -m mobilenet_v2 -r ./records/mobilenet_record.etrecord -o ./exported_models +``` + +## Supported Operations + +The Vulkan delegate supports various operations including: + +- Basic arithmetic (add, subtract, multiply, divide) +- Activations (ReLU, Sigmoid, Tanh, etc.) +- Convolutions (Conv1d, Conv2d, ConvTranspose2d) +- Pooling operations (MaxPool2d, AvgPool2d) +- Linear/Fully connected layers +- BatchNorm, GroupNorm +- Various tensor operations (cat, reshape, permute, etc.) + +For a complete list of supported operations, refer to the Vulkan delegate implementation in the ExecuTorch codebase. + +## Debugging and Optimization + +If you encounter issues with Vulkan delegation: + +1. Use `-r/--etrecord` to generate an ETRecord for debugging +2. Check if your operations are supported by the Vulkan delegate +3. Ensure your Vulkan drivers are up to date +4. Try using the export script with `--strict False` if strict mode causes issues + +## Requirements + +- Vulkan runtime libraries (libvulkan.so.1) +- A Vulkan-capable GPU with appropriate drivers +- PyTorch with Vulkan support diff --git a/examples/vulkan/__init__.py b/examples/vulkan/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/examples/vulkan/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/examples/vulkan/export.py b/examples/vulkan/export.py new file mode 100644 index 00000000000..4d85d83c862 --- /dev/null +++ b/examples/vulkan/export.py @@ -0,0 +1,263 @@ +# 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. + +# Example script for exporting models to flatbuffer with the Vulkan delegate + +# pyre-unsafe + +import argparse +import logging + +import backends.vulkan.test.utils as test_utils + +import torch + +from executorch.backends.transforms.convert_dtype_pass import I64toI32 +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.devtools import BundledProgram +from executorch.devtools.bundled_program.config import MethodTestCase, MethodTestSuite +from executorch.devtools.bundled_program.serialize import ( + serialize_from_bundled_program_to_flatbuffer, +) +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.extension.export_util.utils import save_pte_program +from executorch.extension.pytree import tree_flatten +from torch.export import export + +from ..models import MODEL_NAME_TO_MODEL +from ..models.model_factory import EagerModelFactory + +FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" +logging.basicConfig(level=logging.INFO, format=FORMAT) + + +def main() -> None: + logger = logging.getLogger("") + logger.setLevel(logging.INFO) + + parser = argparse.ArgumentParser() + parser.add_argument( + "-m", + "--model_name", + required=True, + help=f"provide a model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}", + ) + + parser.add_argument( + "-fp16", + "--force_fp16", + action=argparse.BooleanOptionalAction, + default=False, + help="Force fp32 tensors to be converted to fp16 internally. Input/s outputs " + "will be converted to/from fp32 when entering/exiting the delegate. Default is " + "False", + ) + + parser.add_argument( + "-s", + "--strict", + action=argparse.BooleanOptionalAction, + default=True, + help="whether to export with strict mode. Default is True", + ) + + parser.add_argument( + "-a", + "--segment_alignment", + required=False, + help="specify segment alignment in hex. Default is 0x1000. Use 0x4000 for iOS", + ) + + parser.add_argument( + "-e", + "--external_constants", + action=argparse.BooleanOptionalAction, + default=False, + help="Save constants in external .ptd file. Default is False", + ) + + parser.add_argument( + "-d", + "--dynamic", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable dynamic shape support. Default is False", + ) + + parser.add_argument( + "-r", + "--etrecord", + required=False, + default="", + help="Generate and save an ETRecord to the given file location", + ) + + parser.add_argument("-o", "--output_dir", default=".", help="output directory") + + parser.add_argument( + "-b", + "--bundled", + action=argparse.BooleanOptionalAction, + default=False, + help="Export as bundled program (.bpte) instead of regular program (.pte). Default is False", + ) + + parser.add_argument( + "-t", + "--test", + action=argparse.BooleanOptionalAction, + default=False, + help="Execute lower_module_and_test_output to validate the model. Default is False", + ) + + args = parser.parse_args() + + if args.model_name not in MODEL_NAME_TO_MODEL: + raise RuntimeError( + f"Model {args.model_name} is not a valid name. " + f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}." + ) + + model, example_inputs, _, dynamic_shapes = EagerModelFactory.create_model( + *MODEL_NAME_TO_MODEL[args.model_name] + ) + + # Prepare model + model.eval() + + # Setup compile options + compile_options = {} + if args.dynamic or dynamic_shapes is not None: + compile_options["require_dynamic_shapes"] = True + if args.force_fp16: + compile_options["force_fp16"] = True + + # Configure Edge compilation + edge_compile_config = EdgeCompileConfig( + _skip_dim_order=False, # Proper handling for Vulkan memory format + ) + + logging.info(f"Exporting model {args.model_name} with Vulkan delegate") + + # Export the model using torch.export + if dynamic_shapes is not None: + program = export( + model, example_inputs, dynamic_shapes=dynamic_shapes, strict=args.strict + ) + else: + program = export(model, example_inputs, strict=args.strict) + + # Transform and lower with Vulkan partitioner + edge_program = to_edge_transform_and_lower( + program, + compile_config=edge_compile_config, + transform_passes=[ + I64toI32(edge_compile_config._skip_dim_order), + ], + partitioner=[VulkanPartitioner(compile_options)], + generate_etrecord=args.etrecord, + ) + + logging.info( + f"Exported and lowered graph:\n{edge_program.exported_program().graph}" + ) + + # Configure backend options + backend_config = ExecutorchBackendConfig(external_constants=args.external_constants) + if args.segment_alignment is not None: + backend_config.segment_alignment = int(args.segment_alignment, 16) + + # Create executorch program + exec_prog = edge_program.to_executorch(config=backend_config) + + # Save ETRecord if requested + if args.etrecord: + exec_prog.get_etrecord().save(args.etrecord) + logging.info(f"Saved ETRecord to {args.etrecord}") + + # Save the program + output_filename = f"{args.model_name}_vulkan" + + atol = 1e-4 + rtol = 1e-4 + + # If forcing fp16, then numerical divergence is expected + if args.force_fp16: + atol = 2e-2 + rtol = 1e-1 + + # Test the model if --test flag is provided + if args.test: + test_result = test_utils.run_and_check_output( + reference_model=model, + executorch_program=exec_prog, + sample_inputs=example_inputs, + atol=atol, + rtol=rtol, + ) + + if test_result: + logging.info( + "✓ Model test PASSED - outputs match reference within tolerance" + ) + else: + logging.error("✗ Model test FAILED - outputs do not match reference") + raise RuntimeError( + "Model validation failed: ExecutorTorch outputs do not match reference model outputs" + ) + + if args.bundled: + # Create bundled program + logging.info("Creating bundled program with test cases") + + # Generate expected outputs by running the model + expected_outputs = [model(*example_inputs)] + + # Flatten sample inputs to match expected format + inputs_flattened, _ = tree_flatten(example_inputs) + + # Create test suite with the sample inputs and expected outputs + test_suites = [ + MethodTestSuite( + method_name="forward", + test_cases=[ + MethodTestCase( + inputs=inputs_flattened, + expected_outputs=expected_outputs, + ) + ], + ) + ] + + # Create bundled program + bp = BundledProgram(exec_prog, test_suites) + + # Serialize to flatbuffer + bp_buffer = serialize_from_bundled_program_to_flatbuffer(bp) + + # Save bundled program + bundled_output_path = f"{args.output_dir}/{output_filename}.bpte" + with open(bundled_output_path, "wb") as file: + file.write(bp_buffer) + + logging.info( + f"Bundled program exported and saved as {output_filename}.bpte in {args.output_dir}" + ) + else: + # Save regular program + save_pte_program(exec_prog, output_filename, args.output_dir) + logging.info( + f"Model exported and saved as {output_filename}.pte in {args.output_dir}" + ) + + +if __name__ == "__main__": + with torch.no_grad(): + main() # pragma: no cover diff --git a/examples/xnnpack/aot_compiler.py b/examples/xnnpack/aot_compiler.py index 79496c82a58..81eeb75c72c 100644 --- a/examples/xnnpack/aot_compiler.py +++ b/examples/xnnpack/aot_compiler.py @@ -9,12 +9,10 @@ # Example script for exporting simple models to flatbuffer import argparse -import copy import logging import torch from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner -from executorch.devtools import generate_etrecord from executorch.exir import ( EdgeCompileConfig, ExecutorchBackendConfig, @@ -60,6 +58,7 @@ "-r", "--etrecord", required=False, + default="", help="Generate and save an ETRecord to the given file location", ) parser.add_argument("-o", "--output_dir", default=".", help="output directory") @@ -87,14 +86,14 @@ model = model.eval() # pre-autograd export. eventually this will become torch.export - ep = torch.export.export_for_training(model, example_inputs, strict=False) + ep = torch.export.export(model, example_inputs, strict=False) model = ep.module() if args.quantize: logging.info("Quantizing Model...") # TODO(T165162973): This pass shall eventually be folded into quantizer model = quantize(model, example_inputs, quant_type) - ep = torch.export.export_for_training(model, example_inputs, strict=False) + ep = torch.export.export(model, example_inputs, strict=False) edge = to_edge_transform_and_lower( ep, @@ -103,18 +102,16 @@ _check_ir_validity=False if args.quantize else True, _skip_dim_order=True, # TODO(T182187531): enable dim order in xnnpack ), + generate_etrecord=args.etrecord, ) logging.info(f"Exported and lowered graph:\n{edge.exported_program().graph}") - # this is needed for the ETRecord as lowering modifies the graph in-place - edge_copy = copy.deepcopy(edge) - exec_prog = edge.to_executorch( config=ExecutorchBackendConfig(extract_delegate_segments=False) ) - if args.etrecord is not None: - generate_etrecord(args.etrecord, edge_copy, exec_prog) + if args.etrecord: + exec_prog.get_etrecord().save(args.etrecord) logging.info(f"Saved ETRecord to {args.etrecord}") quant_tag = "q8" if args.quantize else "fp32" diff --git a/examples/xnnpack/quantization/example.py b/examples/xnnpack/quantization/example.py index 93831ab8252..f5425155008 100644 --- a/examples/xnnpack/quantization/example.py +++ b/examples/xnnpack/quantization/example.py @@ -60,9 +60,7 @@ def verify_xnnpack_quantizer_matching_fx_quant_model(model_name, model, example_ m = model # 1. pytorch 2.0 export quantization flow (recommended/default flow) - m = torch.export.export_for_training( - m, copy.deepcopy(example_inputs), strict=True - ).module() + m = torch.export.export(m, copy.deepcopy(example_inputs), strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config(is_per_channel=True) quantizer.set_global(quantization_config) @@ -179,9 +177,7 @@ def main() -> None: model = model.eval() # pre-autograd export. eventually this will become torch.export - model = torch.export.export_for_training( - model, example_inputs, strict=True - ).module() + model = torch.export.export(model, example_inputs, strict=True).module() start = time.perf_counter() quantized_model = quantize(model, example_inputs) end = time.perf_counter() diff --git a/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake b/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake index 8e1a478a6cb..ef58f9b4e8d 100644 --- a/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake +++ b/examples/zephyr/x86_64-linux-arm-zephyr-eabi-gcc.cmake @@ -97,10 +97,6 @@ add_compile_options( # -Wall -Wextra -Wcast-align -Wdouble-promotion -Wformat # -Wmissing-field-initializers -Wnull-dereference -Wredundant-decls -Wshadow # -Wswitch -Wswitch-default -Wunused -Wno-redundant-decls - -Wno-stringop-overread - -Wno-error=format= - -Wno-error=maybe-uninitialized - -Wno-error=deprecated-declarations - -Wno-error=shift-count-overflow - -Wno-psabi + -Wno-stringop-overread -Wno-error=format= -Wno-error=maybe-uninitialized + -Wno-error=deprecated-declarations -Wno-error=shift-count-overflow -Wno-psabi ) diff --git a/exir/TARGETS b/exir/TARGETS index cda57de7f80..853d5e199ba 100644 --- a/exir/TARGETS +++ b/exir/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "tracer", srcs = [ "tracer.py", @@ -20,7 +20,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "graph", srcs = [ "graph.py", @@ -31,7 +31,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "graph_module", srcs = [ "graph_module.py", @@ -41,14 +41,14 @@ python_library( ], ) -python_library( +runtime.python_library( name = "scalar_type", srcs = [ "scalar_type.py", ], ) -python_library( +runtime.python_library( name = "schema", srcs = [ "schema.py", @@ -59,14 +59,14 @@ python_library( ], ) -python_library( +runtime.python_library( name = "version", srcs = [ "version.py", ], ) -python_library( +runtime.python_library( name = "tensor", srcs = [ "tensor.py", @@ -79,7 +79,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory", srcs = [ "memory.py", @@ -92,7 +92,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "control_flow", srcs = [ "control_flow.py", @@ -105,7 +105,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "delegate", srcs = [ "delegate.py", @@ -116,7 +116,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "lowered_backend_module", srcs = [ "lowered_backend_module.py", @@ -135,7 +135,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "lib", srcs = [ "__init__.py", @@ -150,7 +150,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory_planning", srcs = [ "memory_planning.py", @@ -167,7 +167,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "common", srcs = [ "common.py", @@ -177,7 +177,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "error", srcs = [ "error.py", @@ -185,7 +185,7 @@ python_library( deps = [], ) -python_library( +runtime.python_library( name = "types", srcs = [ "types.py", @@ -196,7 +196,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "wrap", srcs = [ "wrap.py", @@ -208,7 +208,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "print_program", srcs = [ "print_program.py", @@ -220,14 +220,14 @@ python_library( ], ) -python_library( +runtime.python_library( name = "dynamic_shape", srcs = [ "dynamic_shape.py", ], ) -python_library( +runtime.python_library( name = "pass_base", srcs = [ "pass_base.py", @@ -241,7 +241,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "pass_manager", srcs = [ "pass_manager.py", @@ -253,7 +253,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "sym_util", srcs = ["sym_util.py"], deps = [ @@ -262,7 +262,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "dim_order_utils", srcs = ["dim_order_utils.py"], deps = [ @@ -270,7 +270,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "_warnings", srcs = ["_warnings.py"], deps = [ @@ -278,7 +278,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "debug_handle_utils", srcs = ["debug_handle_utils.py"], deps = [ diff --git a/exir/_serialize/_program.py b/exir/_serialize/_program.py index 0994156ae50..448a3afb90c 100644 --- a/exir/_serialize/_program.py +++ b/exir/_serialize/_program.py @@ -146,6 +146,8 @@ class _ExtendedHeader: + 8 # Segment base offset + 8 + # Segment data size + + 8 ) # Instance attributes. @dataclass will turn these into ctor args. @@ -155,6 +157,9 @@ class _ExtendedHeader: # Offset to the start of the first segment, or zero if there # are no segments. segment_base_offset: int + # Size of the segment data, in bytes, or zero if there are no segments, or + # if the this field isn't populated in the PTE file. + segment_data_size: int # The magic bytes read from or to be written to the binary header. magic: bytes = EXPECTED_MAGIC @@ -189,6 +194,7 @@ def from_bytes(data: bytes) -> "_ExtendedHeader": segment_base_offset=int.from_bytes( data[16:24], byteorder=_HEADER_BYTEORDER ), + segment_data_size=int.from_bytes(data[24:32], byteorder=_HEADER_BYTEORDER), ) def is_valid(self) -> bool: @@ -220,6 +226,9 @@ def to_bytes(self) -> bytes: # uint64_t: Offset to the start of the first segment, or zero if # there are no segments. + self.segment_base_offset.to_bytes(8, byteorder=_HEADER_BYTEORDER) + # uint64_t: size of the segment data, or zero if there are no + # segments. + + self.segment_data_size.to_bytes(8, byteorder=_HEADER_BYTEORDER) ) return data @@ -512,7 +521,9 @@ def serialize_pte_binary( # Construct and pad the extended header. header_data: bytes = _ExtendedHeader( - program_size=program_size, segment_base_offset=segment_base_offset + program_size=program_size, + segment_base_offset=segment_base_offset, + segment_data_size=len(segments_data), ).to_bytes() header_data = pad_to(header_data, padded_header_length) diff --git a/exir/_serialize/test/test_program.py b/exir/_serialize/test/test_program.py index c67849dd28d..18803da05b6 100644 --- a/exir/_serialize/test/test_program.py +++ b/exir/_serialize/test/test_program.py @@ -191,6 +191,8 @@ def constant_segment_with_tensor_alignment( # the end of the file. self.assertGreaterEqual(eh.segment_base_offset, eh.program_size) self.assertLess(eh.segment_base_offset, len(pte_data)) + # Segment data_size should be non-zero since there are segments. + self.assertGreater(eh.segment_data_size, 0) # Peek inside the actual flatbuffer data to see the segments. program_with_segments = _json_to_program(_program_flatbuffer_to_json(pte_data)) @@ -232,6 +234,8 @@ def constant_segment_with_tensor_alignment( # Check segment data. offsets = subsegment_offsets.offsets segment_data: bytes = pte_data[eh.segment_base_offset :] + # Check segment data size. + self.assertEqual(len(segment_data), eh.segment_data_size) # tensor[1]: padding. self.assertEqual( @@ -514,6 +518,8 @@ def test_round_trip_with_segments(self) -> None: # the end of the file. self.assertGreaterEqual(eh.segment_base_offset, eh.program_size) self.assertLess(eh.segment_base_offset, len(pte_data)) + # Segment data size should be non-zero since there are segments. + self.assertGreater(eh.segment_data_size, 0) # Peek inside the actual flatbuffer data to see the segments. Note that # this also implicity tests the case where we try parsing the entire @@ -566,6 +572,8 @@ def test_round_trip_with_segments(self) -> None: # Now that we've shown that the base offset is correct, slice off the # front so that all segment offsets are relative to zero. segment_data: bytes = pte_data[segment_base_offset:] + # Check segment data size. + self.assertEqual(len(segment_data), eh.segment_data_size) # End of the first segment. It's much smaller than the alignment, # so we know that it's followed by zeros. @@ -729,6 +737,8 @@ def test_constant_delegate_and_named_data_segments(self) -> None: # the end of the file. self.assertGreaterEqual(eh.segment_base_offset, eh.program_size) self.assertLess(eh.segment_base_offset, len(pte_data)) + # Segment data size should be non-zero since there are segments. + self.assertGreater(eh.segment_data_size, 0) # Peek inside the actual flatbuffer data to see the segments. program_with_segments = _json_to_program(_program_flatbuffer_to_json(pte_data)) @@ -811,6 +821,8 @@ def test_constant_delegate_and_named_data_segments(self) -> None: # Now that we've shown that the base offset is correct, slice off the # front so that all segment offsets are relative to zero. segment_data: bytes = pte_data[segment_base_offset:] + # Check segment data size. + self.assertEqual(len(segment_data), eh.segment_data_size) # Check segment[0] for constants. offsets = subsegment_offsets.offsets @@ -925,6 +937,8 @@ def test_named_data_segments(self) -> None: # the end of the file. self.assertGreaterEqual(eh.segment_base_offset, eh.program_size) self.assertLess(eh.segment_base_offset, len(pte_data)) + # Segment data size should be non-zero since there are segments. + self.assertGreater(eh.segment_data_size, 0) # Peek inside the actual flatbuffer data to see the named data segments. program_with_segments = _json_to_program(_program_flatbuffer_to_json(pte_data)) @@ -958,6 +972,9 @@ def test_named_data_segments(self) -> None: # Check the pte data for buffer values. segment_data: bytes = pte_data[eh.segment_base_offset :] + # Check segment data size. + self.assertEqual(len(segment_data), eh.segment_data_size) + self.assertEqual( segment_data[ segment_table[0].offset : segment_table[0].offset @@ -985,6 +1002,7 @@ def test_named_data_segments(self) -> None: # the example data. EXAMPLE_PROGRAM_SIZE: int = 0x1122112233443344 EXAMPLE_SEGMENT_BASE_OFFSET: int = 0x5566556677887788 +EXAMPLE_SEGMENT_DATA_SIZE: int = 0x5544554433223322 # This data is intentionally fragile. If the header layout or magic changes, # this test must change too. The layout of the header is a contract, not an # implementation detail. @@ -992,11 +1010,13 @@ def test_named_data_segments(self) -> None: # Magic bytes b"eh00" # uint32_t header size (little endian) - + b"\x18\x00\x00\x00" + + b"\x20\x00\x00\x00" # uint64_t program size + b"\x44\x33\x44\x33\x22\x11\x22\x11" # uint64_t segment base offset + b"\x88\x77\x88\x77\x66\x55\x66\x55" + # uint64_t segment data size + + b"\x22\x33\x22\x33\x44\x55\x44\x55" ) @@ -1005,6 +1025,7 @@ def test_to_bytes(self) -> None: eh = _ExtendedHeader( program_size=EXAMPLE_PROGRAM_SIZE, segment_base_offset=EXAMPLE_SEGMENT_BASE_OFFSET, + segment_data_size=EXAMPLE_SEGMENT_DATA_SIZE, ) self.assertTrue(eh.is_valid()) self.assertEqual(eh.to_bytes(), EXAMPLE_HEADER_DATA) @@ -1013,6 +1034,7 @@ def test_to_bytes_with_non_defaults(self) -> None: eh = _ExtendedHeader( program_size=EXAMPLE_PROGRAM_SIZE, segment_base_offset=EXAMPLE_SEGMENT_BASE_OFFSET, + segment_data_size=EXAMPLE_SEGMENT_DATA_SIZE, # Override the default magic and length, to demonstrate that this # does not affect the serialized header. magic=b"ABCD", @@ -1036,6 +1058,7 @@ def test_from_bytes_valid(self) -> None: self.assertEqual(eh.length, _ExtendedHeader.EXPECTED_LENGTH) self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) def test_from_bytes_with_more_data_than_necessary(self) -> None: # Pass in more data than necessary to parse the header. @@ -1049,6 +1072,7 @@ def test_from_bytes_with_more_data_than_necessary(self) -> None: self.assertEqual(eh.length, _ExtendedHeader.EXPECTED_LENGTH) self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) def test_from_bytes_larger_than_needed_header_size_field(self) -> None: # Simulate a backwards-compatibility situation. Parse a header @@ -1059,11 +1083,13 @@ def test_from_bytes_larger_than_needed_header_size_field(self) -> None: # Magic bytes b"eh00" # uint32_t header size (little endian) - + b"\x1c\x00\x00\x00" # Longer than expected + + b"\x21\x00\x00\x00" # Longer than expected # uint64_t program size + b"\x44\x33\x44\x33\x22\x11\x22\x11" # uint64_t segment base offset + b"\x88\x77\x88\x77\x66\x55\x66\x55" + # uint64_t segment data size + + b"\x22\x33\x22\x33\x44\x55\x44\x55" # uint32_t new field (ignored) + b"\xff\xee\xff\xee" ) @@ -1075,9 +1101,10 @@ def test_from_bytes_larger_than_needed_header_size_field(self) -> None: self.assertTrue(eh.is_valid()) self.assertEqual(eh.magic, _ExtendedHeader.EXPECTED_MAGIC) - self.assertEqual(eh.length, 28) + self.assertEqual(eh.length, 33) self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) def test_from_bytes_not_enough_data_fails(self) -> None: # Parsing a truncated prefix should fail. @@ -1090,11 +1117,13 @@ def test_from_bytes_invalid_magic(self) -> None: # Magic bytes b"ABCD" # Invalid # uint32_t header size (little endian) - + b"\x18\x00\x00\x00" + + b"\x20\x00\x00\x00" # uint64_t program size + b"\x44\x33\x44\x33\x22\x11\x22\x11" # uint64_t segment base offset + b"\x88\x77\x88\x77\x66\x55\x66\x55" + # uint64_t segment data size + + b"\x22\x33\x22\x33\x44\x55\x44\x55" ) # Parse the serialized extended header. @@ -1109,6 +1138,7 @@ def test_from_bytes_invalid_magic(self) -> None: self.assertEqual(eh.length, _ExtendedHeader.EXPECTED_LENGTH) self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) def test_from_bytes_invalid_length(self) -> None: # An invalid serialized header @@ -1121,6 +1151,8 @@ def test_from_bytes_invalid_length(self) -> None: + b"\x44\x33\x44\x33\x22\x11\x22\x11" # uint64_t segment base offset + b"\x88\x77\x88\x77\x66\x55\x66\x55" + # uint64_t segment data size + + b"\x22\x33\x22\x33\x44\x55\x44\x55" ) # Parse the serialized extended header. @@ -1135,3 +1167,4 @@ def test_from_bytes_invalid_length(self) -> None: self.assertEqual(eh.length, 16) self.assertEqual(eh.program_size, EXAMPLE_PROGRAM_SIZE) self.assertEqual(eh.segment_base_offset, EXAMPLE_SEGMENT_BASE_OFFSET) + self.assertEqual(eh.segment_data_size, EXAMPLE_SEGMENT_DATA_SIZE) diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index 6bb2df3dfdb..dd8d97d66ac 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -126,6 +126,10 @@ def to_backend( lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map } + if preprocess_result._delegate_info_meta is not None: + lowered_module.meta["_delegate_info_meta"] = ( + preprocess_result._delegate_info_meta + ) return lowered_module raise NotImplementedError(f"Backend {backend_id} was not found.") @@ -610,6 +614,11 @@ def lower_all_submodules_to_backend( lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map, } + if preprocess_result._delegate_info_meta is not None: + assert lowered_module.meta is not None + lowered_module.meta["_delegate_info_meta"] = ( + preprocess_result._delegate_info_meta + ) is_submodule = call_submodule_node.meta["is_submodule"] toplevel_input_specs_to_delete = call_submodule_node.meta[ "toplevel_input_specs_to_delete" diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index 513ae7c64b3..6999dadb9f9 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -7,7 +7,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from executorch.exir._serialize._named_data_store import NamedDataStoreOutput @@ -32,6 +32,11 @@ class PreprocessResult: # but retrieveable by delegates via the NamedDataMap at runtime. data_store_output: Optional[NamedDataStoreOutput] = None + # Optional delegate-specific information that will be added to the + # lowered_module.meta field in the graph, but not directly serialized + # into the PTE file. + _delegate_info_meta: Optional[Any] = None + """ How to create a backend (for example, BackendWithCompilerDemo): diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py index 961bd741205..50e77e0a884 100644 --- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py +++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py @@ -61,10 +61,14 @@ def duplicate_constant_node( new_input_specs = [] old_signature = exported_program.graph_signature copied_nodes = set() - for idx, node in enumerate(exported_program.graph.nodes): + + placeholder_idx = -1 + for node in exported_program.graph.nodes: if node.op != "placeholder": continue - old_input_spec = old_signature.input_specs[idx] + + placeholder_idx += 1 + old_input_spec = old_signature.input_specs[placeholder_idx] old_input_spec_copy = copy.deepcopy(old_input_spec) if node == to_be_copied[0]: constant_or_attribute_node = node diff --git a/exir/backend/test/TARGETS b/exir/backend/test/TARGETS index 18c0d84eacb..f5913826c17 100644 --- a/exir/backend/test/TARGETS +++ b/exir/backend/test/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") -python_library( +runtime.python_library( name = "backend_with_compiler_demo", srcs = [ "backend_with_compiler_demo.py", @@ -20,7 +20,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "backend_with_delegate_mapping_demo", srcs = [ "backend_with_delegate_mapping_demo.py", @@ -38,7 +38,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "backend_with_named_data_map", srcs = [ "backend_with_named_data_map.py", @@ -94,7 +94,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "qnn_backend_demo", srcs = [ "qnn_backend_demo.py", @@ -109,7 +109,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "hta_partitioner_demo", srcs = [ "hta_partitioner_demo.py", @@ -128,7 +128,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "op_partitioner_demo", srcs = [ "op_partitioner_demo.py", @@ -149,7 +149,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "example_backends", deps = [ ":backend_with_compiler_demo", @@ -219,7 +219,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "backend_with_preprocess_all_demo", srcs = [ "backend_with_preprocess_all_demo.py" diff --git a/exir/backend/test/backend_with_compiler_demo.py b/exir/backend/test/backend_with_compiler_demo.py index b419db153ee..aa60fa2154c 100644 --- a/exir/backend/test/backend_with_compiler_demo.py +++ b/exir/backend/test/backend_with_compiler_demo.py @@ -138,4 +138,5 @@ def preprocess( encoding="utf8", ), debug_handle_map=debug_handle_map, + _delegate_info_meta="test", ) diff --git a/exir/backend/test/demos/rpc/ExecutorBackend.cpp b/exir/backend/test/demos/rpc/ExecutorBackend.cpp index 7632e4ad33c..977c548b1a9 100644 --- a/exir/backend/test/demos/rpc/ExecutorBackend.cpp +++ b/exir/backend/test/demos/rpc/ExecutorBackend.cpp @@ -188,7 +188,7 @@ class ExecutorBackend final : public ::executorch::runtime::BackendInterface { Error execute( ET_UNUSED BackendExecutionContext& context, DelegateHandle* handle, - EValue** args) const override { + Span args) const override { Method* client_method = static_cast(handle); auto num_inputs = client_method->inputs_size(); Error status = Error::Ok; diff --git a/exir/backend/test/demos/rpc/executor_backend_partitioner.py b/exir/backend/test/demos/rpc/executor_backend_partitioner.py index 563d587cfb8..ac8d79482b0 100644 --- a/exir/backend/test/demos/rpc/executor_backend_partitioner.py +++ b/exir/backend/test/demos/rpc/executor_backend_partitioner.py @@ -8,6 +8,8 @@ from typing import final import torch +import torch.fx + from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import ( generate_pattern_op_partitions, ) @@ -65,8 +67,9 @@ def partition(self, edge_exported_program: ExportedProgram) -> PartitionResult: partition_tags[delegation_tag] = self.delegation_spec # Tag the delegate submodules - if node.args[0].op == "get_attr": - node.args[0].meta["delegation_tag"] = delegation_tag + arg0 = node.args[0] + if isinstance(arg0, torch.fx.Node) and arg0.op == "get_attr": + arg0.meta["delegation_tag"] = delegation_tag return PartitionResult( tagged_exported_program=edge_exported_program, diff --git a/exir/backend/test/test_backends.py b/exir/backend/test/test_backends.py index 544b97bb53c..42e8ef16bd7 100644 --- a/exir/backend/test/test_backends.py +++ b/exir/backend/test/test_backends.py @@ -266,7 +266,6 @@ def forward(self, a, x, b): executorch_module = _load_for_executorch_from_buffer(buff) - # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(model_inputs) model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) ref_output = add_mul_module(*model_inputs) @@ -319,16 +318,17 @@ def forward(self, x): ) buff = exec_prog.buffer - + executorch_module = _load_for_executorch_from_buffer(buff) # This line should raise an exception like # RuntimeError: failed with error 0x12 - _load_for_executorch_from_buffer(buff) + inputs_flattened, _ = tree_flatten(model_inputs) + executorch_module.run_method("forward", tuple(inputs_flattened)) @vary_segments def test_backend_with_compiler_out_of_range(self, extract_delegate_segments: bool): with self.assertRaisesRegex( RuntimeError, - "loading method forward failed with error 0x12", + "Failed to execute method forward, error: 0x12", ): self.run_model_in_unsupported_backend( extract_delegate_segments=extract_delegate_segments @@ -876,7 +876,6 @@ def forward(self, a, x, b): self.assertEqual(counter, 2) executorch_module = _load_for_executorch_from_buffer(executorch_prog.buffer) - # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(inputs) model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) ref_output = m(*inputs) diff --git a/exir/backend/test/test_backends_lifted.py b/exir/backend/test/test_backends_lifted.py index b6aea7f8bb3..ef3b502bfd3 100644 --- a/exir/backend/test/test_backends_lifted.py +++ b/exir/backend/test/test_backends_lifted.py @@ -297,7 +297,6 @@ def forward(self, a, x, b): executorch_module = _load_for_executorch_from_buffer(buff) - # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(model_inputs) model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) ref_output = add_mul_module(*model_inputs) @@ -347,15 +346,17 @@ def forward(self, x): buff = exec_prog.buffer + executorch_module = _load_for_executorch_from_buffer(buff) # This line should raise an exception like # RuntimeError: failed with error 0x12 - _load_for_executorch_from_buffer(buff) + inputs_flattened, _ = tree_flatten(model_inputs) + executorch_module.run_method("forward", tuple(inputs_flattened)) @vary_segments def test_backend_with_compiler_out_of_range(self, extract_delegate_segments: bool): with self.assertRaisesRegex( RuntimeError, - "loading method forward failed with error 0x12", + "Failed to execute method forward, error: 0x12", ): self.run_model_in_unsupported_backend( extract_delegate_segments=extract_delegate_segments @@ -887,7 +888,6 @@ def forward(self, a, x, b): self.assertEqual(counter, 2) executorch_module = _load_for_executorch_from_buffer(executorch_prog.buffer) - # pyre-fixme[16]: Module `pytree` has no attribute `tree_flatten`. inputs_flattened, _ = tree_flatten(inputs) model_output = executorch_module.run_method("forward", tuple(inputs_flattened)) ref_output = m(*inputs) @@ -1264,3 +1264,63 @@ def forward(self, x: List[torch.Tensor]): gm = to_edge(export(ComposedM(), inputs, strict=True)) gm.exported_program().module()(*inputs) + + def test_delegate_info_full_delegate(self): + """ + Test that _delegate_info_meta from BackendWithCompilerDemo ends up in the call_delegate node metadata + when using full delegation (to_backend directly). + """ + + class SinModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.sin(x) + + sin_module = SinModule() + model_inputs = (torch.ones(1),) + edgeir_m = to_edge(export(sin_module, model_inputs, strict=True)) + max_value = model_inputs[0].shape[0] + compile_specs = [CompileSpec("max_value", bytes([max_value]))] + lowered_sin_module = to_backend( + "BackendWithCompilerDemo", edgeir_m.exported_program(), compile_specs + ) + + # Check that the lowered module has _delegate_info_meta in its meta + self.assertIn("_delegate_info_meta", lowered_sin_module.meta.keys()) + self.assertEqual(lowered_sin_module.meta["_delegate_info_meta"], "test") + + def test_delegate_info_partitioner(self): + """ + Test that _delegate_info_meta from BackendWithCompilerDemo ends up in the call_delegate node metadata + when using partitioner-based delegation. + """ + + class SinModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return torch.sin(x) + + sin_module = SinModule() + model_inputs = (torch.ones(1),) + max_value = model_inputs[0].shape[0] + + partitioner = AllNodePartitioner( + "BackendWithCompilerDemo", [CompileSpec("max_value", bytes([max_value]))] + ) + + edgeir_m = to_edge(export(sin_module, model_inputs, strict=True)) + lowered_m = edgeir_m.to_backend(partitioner) + + # Check that the lowered submodule has _delegate_info_meta in its meta + lowered_submodules = get_lowered_submodules( + lowered_m.exported_program().graph_module + ) + self.assertEqual(len(lowered_submodules), 1) + + lowered_module = lowered_submodules[0][1] + self.assertIn("_delegate_info_meta", lowered_module.meta) + self.assertEqual(lowered_module.meta["_delegate_info_meta"], "test") diff --git a/exir/backend/test/test_backends_nested.py b/exir/backend/test/test_backends_nested.py index 3313e2a8204..5751706959b 100644 --- a/exir/backend/test/test_backends_nested.py +++ b/exir/backend/test/test_backends_nested.py @@ -197,8 +197,11 @@ def _partition_graph_module( and node.target is torch.ops.higher_order.cond ): # Tag the arguments that take in the submodules to cond - node.args[1].meta["delegation_tag"] = delegation_tag - node.args[2].meta["delegation_tag"] = delegation_tag + arg1, arg2 = node.args[1], node.args[2] + if isinstance(arg1, torch.fx.Node): + arg1.meta["delegation_tag"] = delegation_tag + if isinstance(arg2, torch.fx.Node): + arg2.meta["delegation_tag"] = delegation_tag node.meta["delegation_tag"] = delegation_tag partition_tags[delegation_tag] = self.delegation_spec return partition_tags diff --git a/exir/backend/test/test_compatibility.py b/exir/backend/test/test_compatibility.py index bcda1d36516..4bde3d40b2c 100644 --- a/exir/backend/test/test_compatibility.py +++ b/exir/backend/test/test_compatibility.py @@ -62,12 +62,13 @@ def forward(self, x): ) ) + executorch_module = _load_for_executorch_from_buffer(buff) # Throw runtime error with error code 0x30, meaning delegate is incompatible. with self.assertRaisesRegex( RuntimeError, - "loading method forward failed with error 0x30", + "Failed to execute method forward, error: 0x30", ): - executorch_module = _load_for_executorch_from_buffer(buff) + executorch_module.run_method("forward") def test_compatibility_in_runtime_edge_program_manager(self): class SinModule(torch.nn.Module): @@ -108,9 +109,10 @@ def forward(self, x): ) ) + executorch_module = _load_for_executorch_from_buffer(buff) # Throw runtime error with error code 0x30, meaning delegate is incompatible. with self.assertRaisesRegex( RuntimeError, - "loading method forward failed with error 0x30", + "Failed to execute method forward, error: 0x30", ): - executorch_module = _load_for_executorch_from_buffer(buff) + executorch_module.run_method("forward") diff --git a/exir/backend/test/test_partitioner.py b/exir/backend/test/test_partitioner.py index d369a914fac..dedcfe52966 100644 --- a/exir/backend/test/test_partitioner.py +++ b/exir/backend/test/test_partitioner.py @@ -40,7 +40,7 @@ ) from executorch.extension.pytree import tree_flatten from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param -from torch.export import export, export_for_training +from torch.export import export from torch.fx.passes.operator_support import any_chain @@ -76,7 +76,7 @@ def partition( mlp = MLP() example_inputs = mlp.get_random_inputs() - model = export_for_training(mlp, example_inputs, strict=True).module() + model = export(mlp, example_inputs, strict=True).module() aten = export(model, example_inputs, strict=True) spec_key = "path" spec_value = "/a/b/c/d" @@ -137,7 +137,7 @@ def partition( mlp = MLP() example_inputs = mlp.get_random_inputs() - model = export_for_training(mlp, example_inputs, strict=True).module() + model = export(mlp, example_inputs, strict=True).module() aten = export(model, example_inputs, strict=True) edge = exir.to_edge(aten) @@ -177,7 +177,7 @@ def partition( mlp = MLP() example_inputs = mlp.get_random_inputs() - model = export_for_training(mlp, example_inputs, strict=True).module() + model = export(mlp, example_inputs, strict=True).module() edge = exir.to_edge(export(model, example_inputs, strict=True)) with self.assertRaisesRegex( @@ -229,9 +229,7 @@ def partition( partition_tags=partition_tags, ) - model = export_for_training( - self.AddConst(), (torch.ones(2, 2),), strict=True - ).module() + model = export(self.AddConst(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True)) delegated = edge.to_backend(PartitionerNoTagData()) @@ -310,9 +308,7 @@ def partition( partition_tags=partition_tags, ) - model = export_for_training( - self.AddConst(), (torch.ones(2, 2),), strict=True - ).module() + model = export(self.AddConst(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True)) delegated = edge.to_backend(PartitionerTagData()) @@ -387,9 +383,7 @@ def partition( partition_tags=partition_tags, ) - model = export_for_training( - self.AddConst(), (torch.ones(2, 2),), strict=True - ).module() + model = export(self.AddConst(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True)) delegated = edge.to_backend(PartitionerTagData()) @@ -477,9 +471,7 @@ def partition( ) inputs = (torch.ones(2, 2),) - model = export_for_training( - ReuseConstData(), (torch.ones(2, 2),), strict=True - ).module() + model = export(ReuseConstData(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True)) exec_prog = edge.to_backend(PartitionerTagData()).to_executorch() executorch_module = _load_for_executorch_from_buffer(exec_prog.buffer) @@ -539,9 +531,7 @@ def partition( partition_tags=partition_tags, ) - model = export_for_training( - ReuseConstData(), (torch.ones(2, 2),), strict=True - ).module() + model = export(ReuseConstData(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge(export(model, (torch.ones(2, 2),), strict=True)) with self.assertRaises(RuntimeError) as error: _ = edge.to_backend(PartitionerTagData()) diff --git a/exir/backend/test/test_passes.py b/exir/backend/test/test_passes.py index 1cdf494fa01..32fb75e90f9 100644 --- a/exir/backend/test/test_passes.py +++ b/exir/backend/test/test_passes.py @@ -12,7 +12,7 @@ duplicate_constant_node, ) from torch._export.utils import is_buffer -from torch.export import export_for_training +from torch.export import export from torch.testing import FileCheck @@ -28,9 +28,7 @@ def forward(self, x): z = x - self.const return y, z - model = export_for_training( - ReuseConstData(), (torch.ones(2, 2),), strict=True - ).module() + model = export(ReuseConstData(), (torch.ones(2, 2),), strict=True).module() edge = exir.to_edge( torch.export.export(model, (torch.ones(2, 2),), strict=True) ) diff --git a/exir/capture/TARGETS b/exir/capture/TARGETS index 65e03c70bc9..983cdd210c9 100644 --- a/exir/capture/TARGETS +++ b/exir/capture/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "__init__.py", @@ -13,7 +13,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "capture", srcs = [ "_capture.py", @@ -29,7 +29,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "config", srcs = [ "_config.py", diff --git a/exir/dialects/TARGETS b/exir/dialects/TARGETS index c1f677071c5..39f0a9d0f1b 100644 --- a/exir/dialects/TARGETS +++ b/exir/dialects/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "_ops.py", diff --git a/exir/dialects/backend/TARGETS b/exir/dialects/backend/TARGETS index ed3d46a6b38..1bfd78a9ae1 100644 --- a/exir/dialects/backend/TARGETS +++ b/exir/dialects/backend/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "_ops.py", diff --git a/exir/dialects/edge/TARGETS b/exir/dialects/edge/TARGETS index 3a941d0ecee..4fcf1948b5f 100644 --- a/exir/dialects/edge/TARGETS +++ b/exir/dialects/edge/TARGETS @@ -1,9 +1,9 @@ load("@fbcode_macros//build_defs:export_files.bzl", "export_file") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "_ops.py", diff --git a/exir/dialects/edge/_ops.py b/exir/dialects/edge/_ops.py index 1915bf3d318..206d4d196de 100644 --- a/exir/dialects/edge/_ops.py +++ b/exir/dialects/edge/_ops.py @@ -1,12 +1,14 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 importlib import resources from typing import Any, Dict, List, Optional, Set, Union -import pkg_resources +import executorch.exir.dialects.edge as edge_package import torch @@ -166,9 +168,8 @@ def __getitem__(self, arg_name: str) -> Set[torch.dtype]: def _load_edge_dialect_info() -> Dict[str, Dict[str, Any]]: # pyre-ignore yaml = YAML(typ="safe") - edge_dialect_yaml_info = yaml.load( - pkg_resources.resource_string(__name__, "edge.yaml").decode("utf8") - ) + edge_yaml = resources.read_text(edge_package, "edge.yaml", encoding="utf-8") + edge_dialect_yaml_info = yaml.load(edge_yaml) if edge_dialect_yaml_info: return { edge_op_yaml_info["inherits"]: edge_op_yaml_info diff --git a/exir/dialects/edge/arg/TARGETS b/exir/dialects/edge/arg/TARGETS index 7c499fa4bfa..79705a06357 100644 --- a/exir/dialects/edge/arg/TARGETS +++ b/exir/dialects/edge/arg/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("ai_infra_mobile_platform") -python_library( +runtime.python_library( name = "lib", srcs = [ "model.py", diff --git a/exir/dialects/edge/dtype/TARGETS b/exir/dialects/edge/dtype/TARGETS index 6beb5635c5b..9fbec419812 100644 --- a/exir/dialects/edge/dtype/TARGETS +++ b/exir/dialects/edge/dtype/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("ai_infra_mobile_platform") -python_library( +runtime.python_library( name = "lib", srcs = [ "runner.py", diff --git a/exir/dialects/edge/dtype/supported.py b/exir/dialects/edge/dtype/supported.py index 89c269de1d0..810966220af 100644 --- a/exir/dialects/edge/dtype/supported.py +++ b/exir/dialects/edge/dtype/supported.py @@ -23,6 +23,7 @@ torch.int32: "Int", torch.int64: "Long", torch.float16: "Half", + torch.bfloat16: "BFloat16", torch.float: "Float", torch.double: "Double", torch.uint16: "UInt16", diff --git a/exir/dialects/edge/edge.yaml b/exir/dialects/edge/edge.yaml index 6970f7b00cc..2dac0741007 100644 --- a/exir/dialects/edge/edge.yaml +++ b/exir/dialects/edge/edge.yaml @@ -7,7 +7,7 @@ namespace: edge inherits: aten::_log_softmax type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - self: T0 __ret_0: T0 @@ -16,7 +16,7 @@ namespace: edge inherits: aten::_softmax type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - self: T0 __ret_0: T0 @@ -26,12 +26,12 @@ inherits: aten::_to_copy type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -72,7 +72,7 @@ namespace: edge inherits: aten::abs type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -82,7 +82,7 @@ inherits: aten::acos type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -95,7 +95,7 @@ inherits: aten::acosh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -117,7 +117,7 @@ T7: [Double] T8: [Float] T9: [Float, Int] - T10: [Half] + T10: [Half, BFloat16] T11: [Int] T12: [Long] T13: [Short] @@ -254,9 +254,9 @@ type_alias: T0: [Bool] T1: [Bool, Byte] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T4: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T4: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T5: [Bool, Byte, Char, Int, Long, Short] T6: [Bool, Byte, Char, Int, Short] T7: [Bool, Byte, Char, Short] @@ -269,7 +269,7 @@ T14: [Double] T15: [Float] T16: [Float, Int] - T17: [Half] + T17: [Half, BFloat16] T18: [Int] T19: [Long] T20: [Short] @@ -598,7 +598,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -896,7 +896,7 @@ namespace: edge inherits: aten::alias_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -905,7 +905,7 @@ namespace: edge inherits: aten::amax type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -914,7 +914,7 @@ namespace: edge inherits: aten::amin type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -924,7 +924,7 @@ inherits: aten::any type_alias: T0: [Bool] - T1: [Bool, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] type_constraint: - self: T1 @@ -941,7 +941,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -981,7 +981,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -2071,7 +2071,7 @@ namespace: edge inherits: aten::argmax type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] T1: [Long] type_constraint: - self: T0 @@ -2081,7 +2081,7 @@ namespace: edge inherits: aten::argmin type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] T1: [Long] type_constraint: - self: T0 @@ -2091,7 +2091,7 @@ namespace: edge inherits: aten::as_strided_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -2101,7 +2101,7 @@ inherits: aten::asin type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2114,7 +2114,7 @@ inherits: aten::asinh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2127,7 +2127,7 @@ inherits: aten::atan type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2140,7 +2140,7 @@ inherits: aten::atanh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2152,7 +2152,7 @@ namespace: edge inherits: aten::avg_pool2d type_alias: - T0: [Double, Float, Half, Long] + T0: [Double, Float, Half, BFloat16, Long] type_constraint: - self: T0 __ret_0: T0 @@ -2479,7 +2479,7 @@ namespace: edge inherits: aten::bmm type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 mat2: T0 @@ -2489,7 +2489,7 @@ namespace: edge inherits: aten::cat type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - tensors: T0 __ret_0: T0 @@ -2498,7 +2498,7 @@ namespace: edge inherits: aten::ceil type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -2516,7 +2516,7 @@ T6: [Char] T7: [Double] T8: [Float] - T9: [Half] + T9: [Half, BFloat16] T10: [Int] T11: [Long] T12: [Short] @@ -2775,7 +2775,7 @@ namespace: edge inherits: aten::clone type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -2790,7 +2790,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -2831,12 +2831,12 @@ namespace: edge inherits: aten::convolution type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T1: [Byte] T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -2879,12 +2879,12 @@ inherits: aten::copy type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -2926,7 +2926,7 @@ inherits: aten::cos type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2939,7 +2939,7 @@ inherits: aten::cosh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -2951,12 +2951,12 @@ namespace: edge inherits: aten::cumsum type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T1: [Byte] T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -2990,7 +2990,7 @@ namespace: edge inherits: aten::detach_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -3006,7 +3006,7 @@ T4: [Char] T5: [Double] T6: [Float] - T7: [Half] + T7: [Half, BFloat16] T8: [Int] T9: [Long] T10: [Short] @@ -3057,16 +3057,16 @@ inherits: aten::div.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Bool, Byte, Char, Float, Int, Long, Short] - T4: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T4: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T5: [Byte] T6: [Char] T7: [Double] T8: [Float] T9: [Float, UInt16] - T10: [Half] + T10: [Half, BFloat16] T11: [Int] T12: [Long] T13: [Short] @@ -3138,9 +3138,9 @@ inherits: aten::div.Tensor_mode type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -3151,7 +3151,7 @@ T11: [Char, Short] T12: [Double] T13: [Float] - T14: [Half] + T14: [Half, BFloat16] T15: [Int] T16: [Long] T17: [Short] @@ -3225,7 +3225,7 @@ T1: [Bool, Float, Int] T2: [Double] T3: [Float] - T4: [Half] + T4: [Half, BFloat16] T5: [Int] type_constraint: - self: T2 @@ -3643,7 +3643,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Int, Long] T8: [Long] @@ -3685,7 +3685,7 @@ namespace: edge inherits: aten::empty.memory_format type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - dtype: T0 __ret_0: T0 @@ -3695,13 +3695,13 @@ inherits: aten::eq.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Bool, Float, Int] T3: [Byte] T4: [Char] T5: [Double] T6: [Float] - T7: [Half] + T7: [Half, BFloat16] T8: [Int] T9: [Long] T10: [Short] @@ -3752,7 +3752,7 @@ inherits: aten::erf type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -3765,7 +3765,7 @@ inherits: aten::exp type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -3777,7 +3777,7 @@ namespace: edge inherits: aten::expand_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -3792,7 +3792,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -3834,12 +3834,12 @@ inherits: aten::fill.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -3880,7 +3880,7 @@ namespace: edge inherits: aten::floor type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -3890,9 +3890,9 @@ inherits: aten::floor_divide type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -3903,7 +3903,7 @@ T11: [Char, Short] T12: [Double] T13: [Float] - T14: [Half] + T14: [Half, BFloat16] T15: [Int] T16: [Long] T17: [Short] @@ -3981,7 +3981,7 @@ T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -4022,9 +4022,9 @@ inherits: aten::fmod.Tensor type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -4035,7 +4035,7 @@ T11: [Char, Short] T12: [Double] T13: [Float] - T14: [Half] + T14: [Half, BFloat16] T15: [Int] T16: [Long] T17: [Short] @@ -4111,7 +4111,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -4153,13 +4153,13 @@ inherits: aten::full_like type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Bool, Float, Int] T3: [Byte] T4: [Char] T5: [Double] T6: [Float] - T7: [Half] + T7: [Half, BFloat16] T8: [Int] T9: [Long] T10: [Short] @@ -4691,14 +4691,14 @@ inherits: aten::ge.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Bool, Float, Int] T4: [Byte] T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -4745,14 +4745,14 @@ inherits: aten::ge.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -4823,7 +4823,7 @@ namespace: edge inherits: aten::gelu type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - self: T0 __ret_0: T0 @@ -4832,7 +4832,7 @@ namespace: edge inherits: aten::glu type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - self: T0 __ret_0: T0 @@ -4842,14 +4842,14 @@ inherits: aten::gt.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Bool, Float, Int] T4: [Byte] T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -4896,14 +4896,14 @@ inherits: aten::gt.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -4980,7 +4980,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -5188,7 +5188,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -5231,7 +5231,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -5282,7 +5282,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Int, Long] T8: [Long] @@ -5325,7 +5325,7 @@ inherits: aten::isinf type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T1 __ret_0: T0 @@ -5335,7 +5335,7 @@ inherits: aten::isnan type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T1 __ret_0: T0 @@ -5345,14 +5345,14 @@ inherits: aten::le.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Bool, Float, Int] T4: [Byte] T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5399,14 +5399,14 @@ inherits: aten::le.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5480,7 +5480,7 @@ T0: [Bool, Float, Int] T1: [Double] T2: [Float] - T3: [Half] + T3: [Half, BFloat16] type_constraint: - self: T1 negative_slope: T0 @@ -5496,7 +5496,7 @@ namespace: edge inherits: aten::lift_fresh_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -5506,7 +5506,7 @@ inherits: aten::log type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -5519,14 +5519,14 @@ inherits: aten::logical_and type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5598,7 +5598,7 @@ inherits: aten::logical_not type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T1 __ret_0: T0 @@ -5608,14 +5608,14 @@ inherits: aten::logical_or type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5687,14 +5687,14 @@ inherits: aten::logical_xor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5766,7 +5766,7 @@ inherits: aten::logit type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -5779,14 +5779,14 @@ inherits: aten::lt.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Bool, Float, Int] T4: [Byte] T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5833,14 +5833,14 @@ inherits: aten::lt.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half] + T6: [Double, Float, Half, BFloat16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -5917,7 +5917,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -5968,7 +5968,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -6009,7 +6009,7 @@ T1: [Char] T2: [Double] T3: [Float] - T4: [Half] + T4: [Half, BFloat16] T5: [Int] T6: [Long] T7: [Short] @@ -6043,10 +6043,10 @@ namespace: edge inherits: aten::mean.dim type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T1: [Double] T2: [Float] - T3: [Half] + T3: [Half, BFloat16] type_constraint: - self: T0 dtype: T1 @@ -6067,7 +6067,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -6106,9 +6106,9 @@ type_alias: T0: [Bool] T1: [Bool, Byte] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T4: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T4: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T5: [Bool, Byte, Char, Int, Long, Short] T6: [Bool, Byte, Char, Int, Short] T7: [Bool, Byte, Char, Short] @@ -6119,7 +6119,7 @@ T12: [Char, Short] T13: [Double] T14: [Float] - T15: [Half] + T15: [Half, BFloat16] T16: [Int] T17: [Long] T18: [Short] @@ -6192,7 +6192,7 @@ namespace: edge inherits: aten::mm type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 mat2: T0 @@ -6211,7 +6211,7 @@ T6: [Char] T7: [Double] T8: [Float] - T9: [Half] + T9: [Half, BFloat16] T10: [Int] T11: [Long] T12: [Short] @@ -6259,9 +6259,9 @@ inherits: aten::mul.Tensor type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -6273,7 +6273,7 @@ T12: [Char, Short] T13: [Double] T14: [Float] - T15: [Half] + T15: [Half, BFloat16] T16: [Int] T17: [Long] T18: [Short] @@ -6346,7 +6346,7 @@ namespace: edge inherits: aten::native_layer_norm type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - input: T0 weight: T0 @@ -6360,13 +6360,13 @@ inherits: aten::ne.Scalar type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Bool, Float, Int] T3: [Byte] T4: [Char] T5: [Double] T6: [Float] - T7: [Half] + T7: [Half, BFloat16] T8: [Int] T9: [Long] T10: [Short] @@ -6417,14 +6417,14 @@ inherits: aten::ne.Tensor type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T3: [Byte] T4: [Char] T5: [Double] - T6: [Double, Float, Half, UInt16] + T6: [Double, Float, Half, BFloat16, UInt16] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -6495,7 +6495,7 @@ namespace: edge inherits: aten::neg type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -6504,7 +6504,7 @@ namespace: edge inherits: aten::nonzero type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] T1: [Long] type_constraint: - self: T0 @@ -6514,7 +6514,7 @@ namespace: edge inherits: aten::ones type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - dtype: T0 __ret_0: T0 @@ -6523,7 +6523,7 @@ namespace: edge inherits: aten::permute_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -6532,7 +6532,7 @@ namespace: edge inherits: aten::pixel_shuffle type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -6550,7 +6550,7 @@ T6: [Char] T7: [Double] T8: [Float] - T9: [Half] + T9: [Half, BFloat16] T10: [Int] T11: [Long] T12: [Short] @@ -6598,9 +6598,9 @@ inherits: aten::pow.Tensor_Tensor type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -6611,7 +6611,7 @@ T11: [Char, Short] T12: [Double] T13: [Float] - T14: [Half] + T14: [Half, BFloat16] T15: [Int] T16: [Long] T17: [Short] @@ -6682,7 +6682,7 @@ inherits: aten::reciprocal type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -6694,7 +6694,7 @@ namespace: edge inherits: aten::relu type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -6711,7 +6711,7 @@ T5: [Char] T6: [Double] T7: [Float] - T8: [Half] + T8: [Half, BFloat16] T9: [Int] T10: [Long] T11: [Short] @@ -6752,9 +6752,9 @@ inherits: aten::remainder.Tensor type_alias: T0: [Bool, Byte] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Bool, Byte, Char, Int, Long, Short] T5: [Bool, Byte, Char, Int, Short] T6: [Bool, Byte, Char, Short] @@ -6765,7 +6765,7 @@ T11: [Char, Short] T12: [Double] T13: [Float] - T14: [Half] + T14: [Half, BFloat16] T15: [Int] T16: [Long] T17: [Short] @@ -6835,7 +6835,7 @@ namespace: edge inherits: aten::repeat type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -6844,7 +6844,7 @@ namespace: edge inherits: aten::round type_alias: - T0: [Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -6854,7 +6854,7 @@ inherits: aten::rsqrt type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -6872,7 +6872,7 @@ T3: [Double] T4: [Float] T5: [Float, Int] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -6989,7 +6989,7 @@ T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -7035,7 +7035,7 @@ T2: [Char] T3: [Double] T4: [Float] - T5: [Half] + T5: [Half, BFloat16] T6: [Int] T7: [Long] T8: [Short] @@ -7081,7 +7081,7 @@ namespace: edge inherits: aten::select_copy.int type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7091,12 +7091,12 @@ inherits: aten::select_scatter type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -7138,7 +7138,7 @@ inherits: aten::sigmoid type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7150,7 +7150,7 @@ namespace: edge inherits: aten::sign type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -7160,7 +7160,7 @@ inherits: aten::sin type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7173,7 +7173,7 @@ inherits: aten::sinh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7185,7 +7185,7 @@ namespace: edge inherits: aten::slice_copy.Tensor type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7195,12 +7195,12 @@ inherits: aten::slice_scatter type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -7241,7 +7241,7 @@ namespace: edge inherits: aten::split_copy.Tensor type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7250,7 +7250,7 @@ namespace: edge inherits: aten::split_with_sizes_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7260,7 +7260,7 @@ inherits: aten::sqrt type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7272,7 +7272,7 @@ namespace: edge inherits: aten::squeeze_copy.dim type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7281,7 +7281,7 @@ namespace: edge inherits: aten::squeeze_copy.dims type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7290,7 +7290,7 @@ namespace: edge inherits: aten::stack type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - tensors: T0 __ret_0: T0 @@ -7305,7 +7305,7 @@ T3: [Double] T4: [Float] T5: [Float, Int] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -7417,9 +7417,9 @@ inherits: aten::sub.Tensor type_alias: T0: [Byte] - T1: [Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T2: [Byte, Char, Float, Half, Int, Long, Short, UInt16] - T3: [Byte, Char, Half, Int, Long, Short, UInt16] + T1: [Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T2: [Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T4: [Byte, Char, Int, Long, Short] T5: [Byte, Char, Int, Short] T6: [Byte, Char, Short] @@ -7429,7 +7429,7 @@ T10: [Double] T11: [Float] T12: [Float, Int] - T13: [Half] + T13: [Half, BFloat16] T14: [Int] T15: [Long] T16: [Short] @@ -7717,12 +7717,12 @@ inherits: aten::sum.dim_IntList type_alias: T0: [Bool] - T1: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T1: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] T2: [Byte] T3: [Char] T4: [Double] T5: [Float] - T6: [Half] + T6: [Half, BFloat16] T7: [Int] T8: [Long] T9: [Short] @@ -7759,7 +7759,7 @@ namespace: edge inherits: aten::t_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7769,7 +7769,7 @@ inherits: aten::tan type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7782,7 +7782,7 @@ inherits: aten::tanh type_alias: T0: [Bool, Byte, Char, Float, Int, Long, Short, UInt16] - T1: [Double, Half] + T1: [Double, Half, BFloat16] T2: [Float] type_constraint: - self: T0 @@ -7794,7 +7794,7 @@ namespace: edge inherits: aten::transpose_copy.int type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7803,7 +7803,7 @@ namespace: edge inherits: aten::tril type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short] type_constraint: - self: T0 __ret_0: T0 @@ -7812,7 +7812,7 @@ namespace: edge inherits: aten::unbind_copy.int type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7821,7 +7821,7 @@ namespace: edge inherits: aten::unsqueeze_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7830,7 +7830,7 @@ namespace: edge inherits: aten::var.dim type_alias: - T0: [Double, Float, Half] + T0: [Double, Float, Half, BFloat16] type_constraint: - self: T0 __ret_0: T0 @@ -7839,7 +7839,7 @@ namespace: edge inherits: aten::view_copy type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - self: T0 __ret_0: T0 @@ -7850,9 +7850,9 @@ type_alias: T0: [Bool] T1: [Bool, Byte] - T2: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] - T3: [Bool, Byte, Char, Float, Half, Int, Long, Short, UInt16] - T4: [Bool, Byte, Char, Half, Int, Long, Short, UInt16] + T2: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] + T3: [Bool, Byte, Char, Float, Half, BFloat16, Int, Long, Short, UInt16] + T4: [Bool, Byte, Char, Half, BFloat16, Int, Long, Short, UInt16] T5: [Bool, Byte, Char, Int, Long, Short] T6: [Bool, Byte, Char, Int, Short] T7: [Bool, Byte, Char, Short] @@ -7863,7 +7863,7 @@ T12: [Char, Short] T13: [Double] T14: [Float] - T15: [Half] + T15: [Half, BFloat16] T16: [Int] T17: [Long] T18: [Short] @@ -8382,7 +8382,7 @@ namespace: edge inherits: aten::zeros type_alias: - T0: [Bool, Byte, Char, Double, Float, Half, Int, Long, Short, UInt16] + T0: [Bool, Byte, Char, Double, Float, Half, BFloat16, Int, Long, Short, UInt16] type_constraint: - dtype: T0 __ret_0: T0 diff --git a/exir/dialects/edge/op/TARGETS b/exir/dialects/edge/op/TARGETS index ed4cda67bad..c4f5b5c4d70 100644 --- a/exir/dialects/edge/op/TARGETS +++ b/exir/dialects/edge/op/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("ai_infra_mobile_platform") -python_library( +runtime.python_library( name = "lib", srcs = [ "api.py", diff --git a/exir/dialects/edge/spec/TARGETS b/exir/dialects/edge/spec/TARGETS index 9433415b993..ce5c5970e14 100644 --- a/exir/dialects/edge/spec/TARGETS +++ b/exir/dialects/edge/spec/TARGETS @@ -1,9 +1,9 @@ load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("ai_infra_mobile_platform") -python_library( +runtime.python_library( name = "lib", srcs = [ "gen.py", diff --git a/exir/dialects/edge/test/test_edge_ops.py b/exir/dialects/edge/test/test_edge_ops.py index f621ba1bebb..de0aa205100 100644 --- a/exir/dialects/edge/test/test_edge_ops.py +++ b/exir/dialects/edge/test/test_edge_ops.py @@ -108,14 +108,16 @@ def test_edge_argument_dtype_constraints(self) -> None: if isinstance(arg.type, torch.TensorType): self.assertTrue(isinstance(arg.allowed_types, set)) self.assertEqual( - arg.allowed_types, {torch.float16, torch.float32, torch.float64} + arg.allowed_types, + {torch.float16, torch.float32, torch.float64, torch.bfloat16}, ) for ret in returns: if isinstance(ret.type, torch.TensorType): self.assertTrue(isinstance(ret.allowed_types, set)) self.assertEqual( - ret.allowed_types, {torch.float16, torch.float32, torch.float64} + ret.allowed_types, + {torch.float16, torch.float32, torch.float64, torch.bfloat16}, ) def test_allowed_dtype_set(self) -> None: diff --git a/exir/emit/TARGETS b/exir/emit/TARGETS index b97368c8312..0adee10e82e 100644 --- a/exir/emit/TARGETS +++ b/exir/emit/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "__init__.py", @@ -12,7 +12,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "emit", srcs = [ "_emit_program.py", diff --git a/exir/emit/_emit_program.py b/exir/emit/_emit_program.py index cb849dde11a..0618871bd40 100644 --- a/exir/emit/_emit_program.py +++ b/exir/emit/_emit_program.py @@ -45,6 +45,11 @@ class EmitterOutput: str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]] ] + # This dictionary maps the method name to the corresponding dict which + # contains the mapping of the instruction ids to the number of outputs + # generated by each instruction. + instruction_id_to_num_outs_map: Dict[str, Dict[int, int]] + mutable_data: Optional[List[Buffer]] # Constants are optionally stored in external files. @@ -148,6 +153,7 @@ def emit_program( plans = [] debug_handle_map = {} method_to_delegate_debug_id_map = {} + instruction_id_to_num_outs_map = {} program_state = _ProgramState() # emit each entry point in order according to name. @@ -176,6 +182,7 @@ def emit_program( method_to_delegate_debug_id_map[name] = ( emitter.instr_id_to_delegate_debug_id_map ) + instruction_id_to_num_outs_map[name] = emitter.instruction_id_to_num_outs_map training_metadata = _get_training_metadata(methods) if len(training_metadata) > 0: @@ -188,6 +195,7 @@ def emit_program( return EmitterOutput( debug_handle_map=debug_handle_map, method_to_delegate_debug_id_map=method_to_delegate_debug_id_map, + instruction_id_to_num_outs_map=instruction_id_to_num_outs_map, program=Program( version=EXECUTORCH_SCHEMA_VERSION, execution_plan=plans, diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index 5ee8ca56091..6995f9f73a9 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -253,6 +253,7 @@ def __init__( self.concrete_output_ids: List[_AbstractValue] = [] self.debug_handle_map: Dict[int, Union[int, List[int]]] = {} + self.instruction_id_to_num_outs_map: Dict[int, int] = {} self.instr_id_to_delegate_debug_id_map: Dict[ int, Dict[str, Union[str, _DelegateDebugIdentifierMap]] ] = {} @@ -1003,7 +1004,11 @@ def _add_debug_handle( and node.meta.get("debug_handle") is not None ): debug_handle_list.append(node.meta.get("debug_handle")) + output_node = lowered_module.original_module.graph.output_node() + outputs = output_node.args[0] + num_outputs = len(outputs) if isinstance(outputs, (list, tuple)) else 1 self.debug_handle_map[emitter_id] = debug_handle_list + self.instruction_id_to_num_outs_map[emitter_id] = num_outputs # Debug handle for this node is the emitter_id which is essentially the index of the # instruction in the chain. self.node.meta["debug_handle"] = emitter_id @@ -1030,7 +1035,7 @@ def _add_delegate_map( code, module hierarchy etc. """ delegate_map = {} - if hasattr(lowered_module, "meta"): + if lowered_module.meta is not None: delegate_map = lowered_module.meta.get("debug_handle_map", {}) self.instr_id_to_delegate_debug_id_map[delegate_instruction_id] = { diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index 7d0da7170c6..43b5fcfa99b 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -66,7 +66,7 @@ from functorch.experimental import control_flow from torch import nn -from torch.export import Dim, export, export_for_training +from torch.export import Dim, export from torch.export.experimental import _export_forward_backward @@ -1751,8 +1751,8 @@ def forward(self, x): module_1(*example_inputs) module_2(*example_inputs) - ep1 = export_for_training(module_1, example_inputs, strict=True) - ep2 = export_for_training(module_2, example_inputs, strict=True) + ep1 = export(module_1, example_inputs, strict=True) + ep2 = export(module_2, example_inputs, strict=True) edge_program_manager = exir.to_edge( {"forward1": ep1, "forward2": ep2}, @@ -1794,7 +1794,6 @@ def forward(self, input, label): net = TrainingNet(Net()) # Captures the forward graph. The graph will look similar to the model definition now. - # Will move to export_for_training soon which is the api planned to be supported in the long term. ep = export( net, (torch.randn(1, 2), torch.ones(1, dtype=torch.int64)), strict=True ) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index e1dd7cb4079..61414990703 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -66,6 +66,7 @@ class LoweredBackendModule(torch.nn.Module): _named_data_store_output: Optional[ NamedDataStoreOutput ] # Named Data serialized by the backend + meta: Optional[Dict[str, Any]] # Metadata for the lowered module def __init__( self, @@ -81,6 +82,7 @@ def __init__( self._processed_bytes = processed_bytes self._compile_specs = compile_specs self._named_data_store_output = named_data_store_output + self.meta = None # pyre-ignore def __deepcopy__(self, memo: Optional[Dict[int, Any]]) -> "LoweredBackendModule": @@ -109,7 +111,6 @@ def __deepcopy__(self, memo: Optional[Dict[int, Any]]) -> "LoweredBackendModule" compile_specs=copy.deepcopy(self._compile_specs, memo), named_data_store_output=self._named_data_store_output, ) - # pyre-fixme[16]: `LoweredBackendModule` has no attribute `meta`. res.meta = copy.copy(getattr(self, "meta", {})) return res diff --git a/exir/operator/TARGETS b/exir/operator/TARGETS index f169fbfc8a7..1a8f82a495a 100644 --- a/exir/operator/TARGETS +++ b/exir/operator/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") -python_library( +runtime.python_library( name = "convert", srcs = [ "convert.py", @@ -14,7 +14,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "manip", srcs = [ "manip.py", @@ -25,7 +25,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "util", srcs = [ "util.py", diff --git a/exir/passes/TARGETS b/exir/passes/TARGETS index 0a1f5117f20..45e505cdc7d 100644 --- a/exir/passes/TARGETS +++ b/exir/passes/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "__init__.py", @@ -48,7 +48,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "pass_registry", srcs = [ "pass_registry.py", @@ -59,7 +59,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "external_constants_pass", srcs = [ "external_constants_pass.py", @@ -69,7 +69,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "reinplace_pass", srcs = [ "reinplace.py", @@ -80,7 +80,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "insert_write_back_for_buffers_pass", srcs = [ "insert_write_back_for_buffers_pass.py", @@ -90,7 +90,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "weights_to_outputs_pass", srcs = [ "weights_to_outputs_pass.py", @@ -100,7 +100,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "const_prop_pass", srcs = [ "const_prop_pass.py", @@ -112,7 +112,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "scalar_to_tensor_pass", srcs = [ "scalar_to_tensor_pass.py", @@ -123,7 +123,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "constant_prop_pass", srcs = [ "constant_prop_pass.py", @@ -137,7 +137,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "remove_graph_asserts_pass", srcs = [ "remove_graph_asserts_pass.py", @@ -147,7 +147,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "sym_to_tensor_pass", srcs = [ "sym_to_tensor_pass.py", @@ -158,7 +158,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "quant_fusion_pass", srcs = [ "_quant_patterns_and_replacements.py", @@ -174,7 +174,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "quantize_io_pass", srcs = [ "quantize_io_pass.py", @@ -187,7 +187,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory_planning_pass", srcs = [ "memory_planning_pass.py", @@ -203,7 +203,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "remove_noop_pass", srcs = [ "remove_noop_pass.py", @@ -215,7 +215,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "prune_empty_tensor_pass", srcs = [ "prune_empty_tensors_pass.py", @@ -227,7 +227,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "remove_mixed_type_operators", srcs = [ "remove_mixed_type_operators.py", @@ -238,7 +238,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "spec_prop_pass", srcs = [ "spec_prop_pass.py", @@ -251,7 +251,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "init_mutable_pass", srcs = [ "init_mutable_pass.py", @@ -261,7 +261,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "sym_shape_eval_pass", srcs = [ "sym_shape_eval_pass.py", @@ -276,7 +276,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_broken_ops_with_function_ops_pass", srcs = [ "replace_broken_ops_with_function_ops_pass.py", @@ -287,7 +287,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_sym_size_op_pass", srcs = [ "replace_sym_size_op_pass.py", @@ -298,7 +298,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_edge_with_backend_pass", srcs = [ "replace_edge_with_backend_pass.py", @@ -310,7 +310,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "normalize_transpose_pass", srcs = [ "normalize_transpose_pass.py", @@ -321,7 +321,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_aten_with_edge_pass", srcs = [ "replace_aten_with_edge_pass.py", @@ -335,7 +335,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "debug_handle_generator_pass", srcs = [ "debug_handle_generator_pass.py", @@ -348,7 +348,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "prim_ops_py_registry", srcs = ["executorch_prim_ops_registry.py"], deps = [ @@ -357,7 +357,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "dim_order_ops_registry", srcs = ["dim_order_ops_registry.py"], deps = [ @@ -367,7 +367,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "memory_format_ops_pass", srcs = [ "memory_format_ops_pass.py", @@ -381,7 +381,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "normalize_view_copy_base_pass", srcs = [ "normalize_view_copy_base_pass.py", @@ -392,7 +392,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "replace_view_copy_with_view_pass", srcs = [ "replace_view_copy_with_view_pass.py", @@ -405,7 +405,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "remove_unused_parameters_pass", srcs = [ "remove_unused_parameters_pass.py", diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index fc93aa1b0ca..7daa3a247e8 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -295,7 +295,11 @@ def create_constant_nodes_and_return_specs( return name_to_spec_dict -def _update_output_node_and_specs(exported_program: ExportedProgram) -> None: +# add _skip_dim_order to ensure the introduced correct clone node for different dim order schema +# TODO(gasoonjia): only relying on _clone_dim_order once we remove _skip_dim_order option in the EdgeCompileConfig +def _update_output_node_and_specs( + exported_program: ExportedProgram, _skip_dim_order: bool +) -> None: """ Update the output node and output specs in the exported program. In case a constant node is used as output, we replace it with a clone of the constant node. @@ -307,15 +311,19 @@ def _update_output_node_and_specs(exported_program: ExportedProgram) -> None: output_specs = exported_program.graph_signature.output_specs assert len(output_nodes) == len(output_specs) + clone_op = ( + exir_ops.edge.aten.clone.default + if _skip_dim_order + else exir_ops.edge.dim_order_ops._clone_dim_order.default + ) + for i in range(len(output_specs)): out_node = output_nodes[i] if out_node not in updated_constant_placeholders: continue with exported_program.graph.inserting_after(out_node): - new_node = exported_program.graph.call_function( - exir_ops.edge.aten.clone.default, (out_node,) - ) + new_node = exported_program.graph.call_function(clone_op, (out_node,)) assert "val" in out_node.meta new_node.meta["val"] = out_node.meta["val"] output_nodes[i] = new_node @@ -329,6 +337,7 @@ def _update_output_node_and_specs(exported_program: ExportedProgram) -> None: def constant_prop_pass( exported_program: ExportedProgram, custom_skip_targets: Optional[set[EdgeOpOverload]] = None, + _skip_dim_order: bool = True, ) -> ExportedProgram: """ This pass is for constant propagation for Exported Program with lifted parameters, @@ -376,7 +385,7 @@ def constant_prop_pass( new_input_specs.append(name_to_spec_dict[node.name]) exported_program.graph_signature.input_specs = new_input_specs - _update_output_node_and_specs(exported_program) + _update_output_node_and_specs(exported_program, _skip_dim_order=_skip_dim_order) # Cleanup the graph. exported_program.graph.eliminate_dead_code() diff --git a/exir/passes/dim_order_ops_registry.py b/exir/passes/dim_order_ops_registry.py index f3fc009f109..7a5dff387c1 100644 --- a/exir/passes/dim_order_ops_registry.py +++ b/exir/passes/dim_order_ops_registry.py @@ -28,6 +28,14 @@ "_empty_dim_order.out(int[] size, *, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "_clone_dim_order(Tensor self, *, bool non_blocking=False, int[]? dim_order=None) -> Tensor" +) + +lib.define( + "_clone_dim_order.out(Tensor self, *, bool non_blocking=False, int[]? dim_order=None, Tensor(a!) out) -> Tensor(a!)" +) + def _op_impl(target, *args, **kwargs): kwargs["memory_format"] = get_memory_format(kwargs.get("dim_order", None)) @@ -57,12 +65,23 @@ def _empty_dim_order_out_impl(*args, **kwargs): return _op_impl(torch.ops.aten.empty.out, *args, **kwargs) +@impl(lib, "_clone_dim_order", "CompositeImplicitAutograd") +def _clone_dim_order_impl(*args, **kwargs): + return _op_impl(torch.ops.aten.clone.default, *args, **kwargs) + + +@impl(lib, "_clone_dim_order.out", "CompositeImplicitAutograd") +def _clone_dim_order_out_impl(*args, **kwargs): + return _op_impl(torch.ops.aten.clone.out, *args, **kwargs) + + """ Defines a map of edge ops to the corresponding dim_order ops for quick lookup """ DimOrderOpsMap = { exir_ops.edge.aten._to_copy.default: exir_ops.edge.dim_order_ops._to_dim_order_copy.default, exir_ops.edge.aten.empty.memory_format: exir_ops.edge.dim_order_ops._empty_dim_order.default, + exir_ops.edge.aten.clone.default: exir_ops.edge.dim_order_ops._clone_dim_order.default, } """ diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py index d9bba4635ff..ba1d85adb4f 100644 --- a/exir/passes/external_constants_pass.py +++ b/exir/passes/external_constants_pass.py @@ -27,6 +27,7 @@ def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool: def external_constants_pass( gm: GraphModule, + gen_tag_fn: Optional[Callable[[torch.fx.Node], Optional[str]]] = None, ) -> PassResult: """ Move all non-lifted constants to external file. @@ -42,7 +43,10 @@ def external_constants_pass( if (node.op == "placeholder") and ("_lifted_tensor" not in node.name): spec = node.meta.get("spec") if isinstance(spec, TensorSpec) and spec.const: - node.meta["constant_tag"] = "_default_external_constant" + if gen_tag_fn is not None: + node.meta["constant_tag"] = gen_tag_fn(node) + else: + node.meta["constant_tag"] = "_default_external_constant" mutated = True return PassResult(gm, mutated) @@ -88,31 +92,22 @@ def external_mutable_weights_pass( return PassResult(gm, mutated) -def delegate_external_constants_pass( - gm: GraphModule, - ep: ExportedProgram, - gen_tag_fn: Optional[Callable[[torch.fx.Node], str]] = None, +# Note: this pass must be run on an unlifted graph, e.g. ep.module(), +# and not on a lifted graph, e.g. ep.graph_module. +# This is using 'get_attr' to tag constants, which only appears in +# unlifted graphs. +def delegate_external_constants_pass_unlifted( + module: torch.nn.Module, + gen_tag_fn: Optional[Callable[[torch.fx.Node], Optional[str]]] = None, ) -> PassResult: - """ - Tag external constants before to_backend. - - Note: this pass must be run after run_decompositions(), as tags on - constants are removed then. - - Args: - gm: GraphModule to tag. - ep: ExportedProgram, to distinguish if a node is a constant. - gen_tag_fn: node -> str callable indicating the tag for the node. - Returns: - PassResult: The resulting gm, and if it was mutated or not. - """ mutated = False - for module in gm.modules(): - if not isinstance(module, torch.fx.GraphModule): + for m in module.modules(): + if not isinstance(m, torch.fx.GraphModule): continue - for node in module.graph.nodes: - if node.op == "placeholder" and is_param_node(ep, node): + for node in m.graph.nodes: + if node.op == "get_attr": if gen_tag_fn is not None: - node.meta["delegate_constant_tag"] = gen_tag_fn(node) + node.meta.setdefault("custom", {}) + node.meta["custom"]["delegate_constant_tag"] = gen_tag_fn(node) mutated = True - return PassResult(gm, mutated) + return PassResult(module, mutated) diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py index db597b2b233..421f30960b6 100644 --- a/exir/passes/memory_format_ops_pass.py +++ b/exir/passes/memory_format_ops_pass.py @@ -19,9 +19,6 @@ logger = logging.getLogger(__file__) logger.setLevel(logging.INFO) -# TODO - these passes are too specialized on a single to_copy op. -# We should be able to replace (or revert) any of the dim_order ops in the future. - class MemoryFormatOpsPass(ExportPass): """ @@ -43,7 +40,7 @@ def call_operator(self, op, args, kwargs, meta): # new kwargs with dim_order, and no memory_format for the new op nkwargs = dict(copy.deepcopy(kwargs)) # orig kwargs are immutable - # get the "to" memory format for the EdgeOp + # get the target memory format for the EdgeOp mem_format = nkwargs.pop("memory_format", torch.contiguous_format) # can always get the shape, assuming rank is specialized diff --git a/exir/passes/remove_mixed_type_operators.py b/exir/passes/remove_mixed_type_operators.py index d0e48a277c0..86a71354337 100644 --- a/exir/passes/remove_mixed_type_operators.py +++ b/exir/passes/remove_mixed_type_operators.py @@ -23,6 +23,7 @@ def call_operator(self, op, args, kwargs, meta: NodeMetadata): # noqa: C901 promotion_type_allow_list = { torch.ops.aten.add.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, torch.ops.aten.mul.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + torch.ops.aten.sub.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, # The correct promotion for div depends on the mode! If there is no mode, # it's INT_TO_FLOAT, otherwise it's default. torch.ops.aten.div.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, diff --git a/exir/program/TARGETS b/exir/program/TARGETS index cb28d563623..b113d8ab2bc 100644 --- a/exir/program/TARGETS +++ b/exir/program/TARGETS @@ -1,9 +1,9 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "__init__.py", @@ -14,7 +14,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "program", srcs = [ "_program.py", @@ -50,7 +50,7 @@ python_library( ] + (["//executorch/exir/program/fb:logger"] if not runtime.is_oss else []) ) -python_library( +runtime.python_library( name = "fake_program", srcs = [ "_fake_program.py", diff --git a/exir/program/_program.py b/exir/program/_program.py index 63b49d9860d..f3d9eef9221 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -22,6 +22,7 @@ ) from executorch.exir._serialize._serialize import serialize_for_executorch from executorch.exir._serialize.data_serializer import DataSerializer +from executorch.exir._warnings import experimental from executorch.exir.backend.backend_api import ( MethodProgramsPartitionerSpec, to_backend, @@ -109,8 +110,8 @@ # Define a stub decorator that does nothing def et_logger(api_name: str) -> Callable[[Any], Any]: def decorator(func: Callable[..., Any]) -> Callable[..., Any]: - def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: - return func(self, *args, **kwargs) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) return wrapper @@ -240,8 +241,29 @@ def _transform( isinstance(p, (list, Verifier)) for p in passes ), f"Expected all passes to be of PassType, not list or Verifier. Use override_verifiers kwarg instead. Got: {list(passes)}" - pm = PassManager(list(passes)) - res = pm(self.graph_module) + return _transform_with_pass_manager( + self, PassManager(list(passes)), override_verifiers + ) + + +def _transform_with_pass_manager( + self, + pass_manager: PassManager, + override_verifiers: None | list[Type[Verifier]] = None, +) -> "ExportedProgram": + """ + Transforms the program using the provided pass_manager. + + Args: + self: The ExportedProgram instance to transform + pass_manager: An instance of PassManager to apply transformations. + override_verifiers: Optional list of verifier classes to use instead of the default verifiers. + This is needed if the transforms yields illegal graph that the default verifier cannot handle. + + Returns: + ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made + """ + res = pass_manager(self.graph_module) transformed_gm = res.graph_module if res is not None else self.graph_module assert transformed_gm is not None @@ -605,6 +627,24 @@ def buffer(self) -> bytes: self._buffer = bytes(self._get_pte_data()) return self._buffer + @property + @experimental("This API is experimental and subject to change without notice.") + def data_files(self) -> Dict[str, bytes]: + """Returns the data files as a dictionary of filename to byte data. + + Returns: + Dict[str, bytes]: Dictionary mapping data filenames (e.g., .ptd files) to + their serialized byte content. + Returns empty dict if no data files are available. + """ + if self._pte_data is None: + self._get_pte_data() # This populates _tensor_data + + if self._tensor_data is None: + return {} + + return {filename: bytes(cord) for filename, cord in self._tensor_data.items()} + @property def program(self) -> Program: return self._get_emitter_output().program @@ -613,7 +653,7 @@ def program(self) -> Program: def debug_handle_map(self) -> Dict[int, Union[int, List[int]]]: if self._emitter_output: return self._emitter_output.debug_handle_map - return {} + return self._get_emitter_output().debug_handle_map @property def delegate_map( @@ -621,7 +661,15 @@ def delegate_map( ) -> Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]]: if self._emitter_output: return self._emitter_output.method_to_delegate_debug_id_map - return {} + return self._get_emitter_output().method_to_delegate_debug_id_map + + @property + def instruction_id_to_num_outs_map( + self, + ) -> Dict[str, Dict[int, Union[int, List[int]]]]: + if self._emitter_output: + return self._emitter_output.instruction_id_to_num_outs_map + return self._get_emitter_output().instruction_id_to_num_outs_map @property def graph_module(self) -> torch.fx.GraphModule: @@ -787,11 +835,13 @@ def edge_to_executorch_passes( Get the pre memory planning passes based on the method name, if the pass is not in the dict, use the default pass. """ passes: List[PassType] = [ - *config.passes, SpecPropPass(), # ExecuTorch backend ops are unable to handle unbacked symints. So after # this pass, passes cannot be Interpreter-based, because it will fail if # there exists an unbacked symint operation. + *config.passes, + # config.passes may contain external_constants_pass. This pass has to + # run after SpecPropPass, which populates tensor names. EdgeToBackendOpsPass(), RemoveGraphAssertsPass(), ] + pre_memory_planning_passes(config, name) @@ -1230,7 +1280,7 @@ def collect_named_data_store_outputs( def to_edge_transform_and_lower( # noqa: C901 programs: Union[ExportedProgram, Dict[str, ExportedProgram]], transform_passes: Optional[ - Union[Sequence[PassType], Dict[str, Sequence[PassType]]] + Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager] ] = None, partitioner: Optional[ Union[List[Partitioner], Dict[str, List[Partitioner]]] @@ -1259,11 +1309,15 @@ def to_edge_transform_and_lower( # noqa: C901 to their corresponding ExportedPrograms. If only a single ExportedProgram is provided it will be assigned the name "forward". - transform_passes: The passes can either be a list of passes, or a dictionary - mapping method names to lists of passes. If it is just a list of passes, all methods - in the given EdgeProgramManager will be transformed with the provided passes. If it - is a dictionary, only method names specified in the dictionary will be transformed - with their corresponding passes. + transform_passes: The transform_passes can be one of: + 1) a list of passes - + all methods in the given EdgeProgramManager will be transformed with the provided passes. + 2) a dictionary - + only method names specified in the dictionary will be transformed + with their corresponding passes + 3) an instance of a PassManager - + all methods in the given EdgeProgramManager will be + transformed with the given PassManager instance. partitioner: The partitioner can either be a Partitioner subclass instance, or a dictionary mapping method names to Partitioner subclass instance. If it is a @@ -1351,6 +1405,7 @@ def to_edge( programs: Union[ExportedProgram, Dict[str, ExportedProgram]], constant_methods: Optional[Dict[str, Any]] = None, compile_config: Optional[EdgeCompileConfig] = None, + generate_etrecord: bool = False, ) -> "EdgeProgramManager": """ :func:`to_edge` constructs an EdgeProgramManager from a set of exported programs in @@ -1363,6 +1418,8 @@ def to_edge( compile_config: An optional argument used to provide greater control over the transformation to edge dialect process. + generate_etrecord: An optional argument used to generate an etrecord for debugging purposes. Default is False. + Returns: EdgeProgramManager """ @@ -1416,7 +1473,14 @@ def to_edge( logging.info(f"Input program {name} is not in Edge dialect.") raise e - return EdgeProgramManager(edge_programs, constant_methods, config) + epm = EdgeProgramManager(edge_programs, constant_methods, config) + if generate_etrecord: + etrecord = _create_empty_etrecord() + etrecord.add_exported_program(aten_programs) + etrecord.add_edge_dialect_program(copy.deepcopy(epm)) + epm._etrecord = etrecord + + return epm class EdgeProgramManager: @@ -1493,19 +1557,23 @@ def exported_program(self, method_name: str = "forward") -> ExportedProgram: @et_logger("transform") def transform( self, - passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]]], + passes: Union[Sequence[PassType], Dict[str, Sequence[PassType]], PassManager], compile_config: Optional[EdgeCompileConfig] = None, ) -> "EdgeProgramManager": """ Transforms the program according to the provided passes. Args: - passes: The passes can either be a list of passes, or a - dictionary mapping method names to lists of passes. If it is - just a list of passes, all methods in the given EdgeProgramManager - will be transformed with the provided passes. If it is a - dictionary, only method names specified in the dictionary will be - transformed with their corresponding passes. + passes: This param can be one of: + 1) a list of passes - + all methods in the given EdgeProgramManager + will be transformed with the provided passes. + 2) a dictionary mapping method names to lists of passes - + only method names specified in the dictionary will be + transformed with their corresponding passes. + 3) a PassManager instance - + all methods in the given EdgeProgramManager will be + transformed with the given PassManager instance. compile_config: Compile config to use for veriy the correctness of model graph after each pass. If not specified, the compile config of the calling EdgeProgramManager will be used. It will be used in as compile @@ -1515,29 +1583,52 @@ def transform( EdgeProgramManager: A copy of the calling EdgeProgramManager with the transformations applied. """ + compile_config = compile_config or self.compile_config new_programs: Dict[str, ExportedProgram] = {} + + # Cast passes parameter upfront. + passes_seq: Optional[Sequence[PassType]] = None + passes_dict: Optional[Dict[str, Sequence[PassType]]] = None + pass_manager: Optional[PassManager] = None + + if isinstance(passes, Sequence): + passes_seq = passes if isinstance(passes, dict): - for name, program in self._edge_programs.items(): - if name in passes.keys(): - new_programs[name] = _transform(program, *passes[name]) - EXIREdgeDialectVerifier(edge_compile_config=compile_config)( - new_programs[name].graph_module - ) - else: - new_programs[name] = copy.deepcopy(program) + passes_dict = passes + if isinstance(passes, PassManager): + pass_manager = passes - else: # apply passes to every method - for name, program in self._edge_programs.items(): - new_programs[name] = _transform(program, *passes) - EXIREdgeDialectVerifier(edge_compile_config=compile_config)( - new_programs[name].graph_module - ) + for name, program in self._edge_programs.items(): + # If the method name is enforced, but not matched, we skip transformation. + if ( + isinstance(passes, dict) + and passes_dict + and name not in passes_dict.keys() + ): + new_programs[name] = copy.deepcopy(program) + continue + + # Depending on the passes parameter, call the corresponding transform function. + if passes_seq is not None: + new_programs[name] = _transform(program, *passes_seq) + elif passes_dict is not None: + new_programs[name] = _transform(program, *passes_dict[name]) + elif pass_manager is not None: + new_programs[name] = _transform_with_pass_manager(program, pass_manager) + + # Verify the correctness of model graph after each transformation. + EXIREdgeDialectVerifier(edge_compile_config=compile_config)( + new_programs[name].graph_module + ) - return EdgeProgramManager( + epm = EdgeProgramManager( new_programs, copy.deepcopy(self._config_methods), compile_config ) + epm._etrecord = self._etrecord + return epm + @et_logger("to_backend") def to_backend( self, @@ -1580,12 +1671,15 @@ def to_backend( new_edge_programs = to_backend(method_to_programs_and_partitioners) config = EdgeCompileConfig(_check_ir_validity=False) - return EdgeProgramManager( + epm = EdgeProgramManager( new_edge_programs, copy.deepcopy(self._config_methods), config, ) + epm._etrecord = self._etrecord + return epm + @et_logger("to_executorch") def to_executorch( self, @@ -1793,6 +1887,12 @@ def delegate_map( ) -> Dict[str, Dict[int, Dict[str, Union[str, _DelegateDebugIdentifierMap]]]]: return self._emitter_output.method_to_delegate_debug_id_map + @property + def instruction_id_to_num_outs_map( + self, + ) -> Dict[str, Dict[int, Union[int, List[int]]]]: + return self._emitter_output.instruction_id_to_num_outs_map + @property def executorch_program(self) -> Program: """ @@ -1843,7 +1943,9 @@ def write_tensor_data_to_file(self, outdir) -> None: """ assert self._tensor_data is not None for filename, cord in self._tensor_data.items(): - with open(os.path.join(outdir, f"{filename}.ptd"), "wb") as f: + if not filename.endswith(".ptd"): + filename += ".ptd" + with open(os.path.join(outdir, f"{filename}"), "wb") as f: logging.info(f"Writing data file to {filename}") cord.write_to_file(f) diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index da5647936aa..2e788ef5c74 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -37,6 +37,7 @@ from torch._export.verifier import Verifier from torch.export import Dim, export, ExportedProgram from torch.export._trace import _export +from torch.fx.passes.infra.pass_manager import PassManager from torch.library import impl, Library from torch.nn import functional as F @@ -470,6 +471,30 @@ def test_transform_dict_api(self): torch.ones(1) + 1, # x + 1 ) + def test_transform_pass_manager_api(self): + edge_manager = to_edge(get_exported_programs(), get_config_methods()) + + pm = PassManager() + pm.add_pass(AddToMulPassEdge()) + + transformed_edge = edge_manager.transform(pm) + + x = torch.ones(1) * 2 + y = torch.ones(1) * 3 + + # x * y + x -> x * y * x + self.assertEqual( + transformed_edge.exported_program("forward").module()(x, y), x * y * x + ) + + # x + 1 -> x * 1 + self.assertEqual( + transformed_edge.exported_program("foo").module()( + x, + ), + x * 1, + ) + def test_edge_to_backend_replaces_subgraph(self): edge_manager: EdgeProgramManager = to_edge( get_exported_programs(), get_config_methods() diff --git a/exir/serde/TARGETS b/exir/serde/TARGETS index ec3db22aac6..4cc442a0fca 100644 --- a/exir/serde/TARGETS +++ b/exir/serde/TARGETS @@ -1,8 +1,8 @@ -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") -python_library( +runtime.python_library( # @autodeps-skip for some reason autodeps thinks this target # needs to depend on exir:lib which it doesn't. name = "serialize", diff --git a/exir/serde/schema_check.py b/exir/serde/schema_check.py index acf6fa20677..9f95341f487 100644 --- a/exir/serde/schema_check.py +++ b/exir/serde/schema_check.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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. @@ -205,15 +206,16 @@ class _Commit: def update_schema(): - import importlib.resources + import importlib.resources as _resources + import executorch.exir.serde as serde_package - if importlib.resources.is_resource(__package__, "schema.yaml"): - content = importlib.resources.read_text(__package__, "schema.yaml") + if _resources.is_resource(serde_package, "schema.yaml"): + content = _resources.read_text(serde_package, "schema.yaml") match = re.search("checksum<<([A-Fa-f0-9]{64})>>", content) _check(match is not None, "checksum not found in schema.yaml") assert match is not None checksum_base = match.group(1) - from yaml import load, Loader + from yaml import Loader, load dst = load(content, Loader=Loader) assert isinstance(dst, dict) diff --git a/exir/tests/TARGETS b/exir/tests/TARGETS index 27256b3b3b4..855793d53a5 100644 --- a/exir/tests/TARGETS +++ b/exir/tests/TARGETS @@ -1,10 +1,10 @@ load("@fbcode_macros//build_defs:cpp_library.bzl", "cpp_library") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") -python_library( +runtime.python_library( name = "lib", srcs = [ "common.py", @@ -15,7 +15,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "models", srcs = [ "models.py", @@ -27,7 +27,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "control_flow_models", srcs = [ "control_flow_models.py", @@ -37,7 +37,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "dynamic_shape_models", srcs = [ "dynamic_shape_models.py", @@ -47,7 +47,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "transformer", srcs = [ "transformer.py", @@ -57,7 +57,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "asr_joiner", srcs = [ "asr_joiner.py", @@ -419,7 +419,7 @@ python_unittest( ], ) -python_library( +runtime.python_library( name = "test_memory_format_ops_pass_utils", srcs = [ "test_memory_format_ops_pass_utils.py", diff --git a/exir/tests/test_arg_validator.py b/exir/tests/test_arg_validator.py index ede8b224329..5580881ab01 100644 --- a/exir/tests/test_arg_validator.py +++ b/exir/tests/test_arg_validator.py @@ -37,16 +37,15 @@ def forward(self, x): self.assertEqual(len(validator.violating_ops), 0) def test_edge_dialect_fails(self) -> None: - # torch.bfloat16 is not supported by edge::aten::_log_softmax + # torch.complex64 is not supported by edge::aten::add class M(torch.nn.Module): def __init__(self): super().__init__() - self.m = torch.nn.LogSoftmax(dim=1) def forward(self, x): - return self.m(x) + return x + x - inputs = (torch.randn(1, 3, 100, 100).to(dtype=torch.bfloat16),) + inputs = (torch.randn(1, 3, 100, 100).to(dtype=torch.complex64),) egm = ( to_edge( export(M(), inputs, strict=True), @@ -61,12 +60,13 @@ def forward(self, x): key: EdgeOpOverload = next(iter(validator.violating_ops)) self.assertEqual( key.name(), - ops.edge.aten._log_softmax.default.name(), + ops.edge.aten.add.Tensor.name(), ) self.assertDictEqual( validator.violating_ops[key][0], { - "self": torch.bfloat16, - "__ret_0": torch.bfloat16, + "self": torch.complex64, + "other": torch.complex64, + "__ret_0": torch.complex64, }, ) diff --git a/exir/tests/test_extract_io_quant_params.py b/exir/tests/test_extract_io_quant_params.py index 84da01c673d..ec018e8ae68 100644 --- a/exir/tests/test_extract_io_quant_params.py +++ b/exir/tests/test_extract_io_quant_params.py @@ -41,7 +41,7 @@ def setUp(self): operator_config = get_symmetric_quantization_config() self.quantizer.set_global(operator_config) - exported = torch.export.export_for_training( + exported = torch.export.export( self.mod, copy.deepcopy(self.example_inputs), strict=True, @@ -54,7 +54,7 @@ def setUp(self): converted = convert_pt2e(prepared) # Export again with quant parameters - final_export = torch.export.export_for_training( + final_export = torch.export.export( converted, self.example_inputs, strict=True, diff --git a/exir/tests/test_memory_format_ops_pass.py b/exir/tests/test_memory_format_ops_pass.py index 84cd0faa485..2384f6123a9 100644 --- a/exir/tests/test_memory_format_ops_pass.py +++ b/exir/tests/test_memory_format_ops_pass.py @@ -27,7 +27,10 @@ AmbiguousDimOrderError, MemoryFormatOpsPassTestUtils, MemoryFormatTestSet, + PropagateToCloneChannelsLastModule, PropagateToCopyChannalsLastModule, + SimpleCloneChannelsLastModule, + SimpleCloneContiguousModule, SimpleEmptyChannelLastModule, SimpleEmptyContiguoustModule, SimpleToCopyChannelsLastModule, @@ -91,6 +94,36 @@ def test_op_empty_replacement_contiguous(self) -> None: ), ) + def test_op_clone_replacement_contiguous(self) -> None: + model = SimpleCloneContiguousModule() + MemoryFormatOpsPassTestUtils.memory_format_test_runner( + self, + MemoryFormatTestSet( + module=model.eval(), + op=torch.ops.aten.clone.default, + sample_input=( + torch.randn((3, 4, 5, 6)).to(memory_format=torch.channels_last), + ), + target_memory_format=torch.contiguous_format, + _load_for_executorch_from_buffer=_load_for_executorch_from_buffer, + ), + ) + + def test_op_clone_replacement_channels_last(self) -> None: + model = SimpleCloneChannelsLastModule() + MemoryFormatOpsPassTestUtils.memory_format_test_runner( + self, + MemoryFormatTestSet( + module=model.eval(), + op=torch.ops.aten.clone.default, + sample_input=( + torch.randn((3, 4, 5, 6)).to(memory_format=torch.contiguous_format), + ), + target_memory_format=torch.channels_last, + _load_for_executorch_from_buffer=_load_for_executorch_from_buffer, + ), + ) + def test_op_dim_order_update(self) -> None: MemoryFormatOpsPassTestUtils.memory_format_test_runner( self, @@ -128,6 +161,25 @@ def test_op_dim_order_propagation(self) -> None: check_unambiguous_dim_order=True, ) + def test_op_clone_dim_order_propagation(self) -> None: + MemoryFormatOpsPassTestUtils.memory_format_test_runner( + self, + MemoryFormatTestSet( + module=PropagateToCloneChannelsLastModule().eval(), + op=torch.ops.aten.clone.default, + sample_input=( + torch.rand_like( + torch.zeros([2, 2, 2, 2]), + dtype=torch.float32, + memory_format=torch.contiguous_format, + ), + ), + target_memory_format=torch.channels_last, + _load_for_executorch_from_buffer=_load_for_executorch_from_buffer, + ), + check_unambiguous_dim_order=True, + ) + def test_op_dim_order_propagation_ambiguous(self) -> None: try: MemoryFormatOpsPassTestUtils.memory_format_test_runner( diff --git a/exir/tests/test_memory_format_ops_pass_utils.py b/exir/tests/test_memory_format_ops_pass_utils.py index 6daf38b187f..f5a786c6f74 100644 --- a/exir/tests/test_memory_format_ops_pass_utils.py +++ b/exir/tests/test_memory_format_ops_pass_utils.py @@ -38,6 +38,10 @@ "torch.ops.aten.empty.memory_format", "executorch_exir_dialects_edge__ops_dim_order_ops__empty_dim_order_default", ), + torch.ops.aten.clone.default: ( + "torch.ops.aten.clone.default", + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", + ), } @@ -70,6 +74,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(dtype=torch.double, memory_format=torch.channels_last) +class SimpleCloneContiguousModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.clone(memory_format=torch.contiguous_format) + + +class SimpleCloneChannelsLastModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.clone(memory_format=torch.channels_last) + + class SimpleEmptyContiguoustModule(torch.nn.Module): def __init__(self): super().__init__() @@ -102,6 +122,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return t1 * t2 +class PropagateToCloneChannelsLastModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + t1 = x.clone(memory_format=torch.channels_last) + t2 = t1 + t1 + return t1 * t2 + + class AmbiguousDimOrderError(RuntimeError): pass diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 70a4c88e3b6..716b808b087 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -159,6 +159,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return Module Add = make_module(lambda x, y: (x + y) + x) + Sub = make_module(lambda x, y: (x - y) - x) Mult = make_module(lambda x, y: x * y) Minimum = make_module(torch.minimum) DivWithoutMode = make_module(torch.div) @@ -177,6 +178,12 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: 2, ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, ), + ( + Sub, + exir_ops.edge.aten.sub.Tensor, + 2, + ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), ( Mult, exir_ops.edge.aten.mul.Tensor, @@ -1185,7 +1192,7 @@ def forward(self) -> torch.Tensor: ) edge._edge_programs["forward"] = constant_prop_pass( - edge.exported_program("forward") + edge.exported_program("forward"), _skip_dim_order=False ) # Check (c_lifted_tensor_*) nodes are all replaced by _prop_tensor_constant. @@ -1333,9 +1340,7 @@ def forward(self, query, key, value): value = torch.randn(32, 32, 32, 32) # Capture the model - m = torch.export.export_for_training( - M(32), (query, key, value), strict=True - ).module() + m = torch.export.export(M(32), (query, key, value), strict=True).module() # 8w16a quantization from torchao.quantization.pt2e import MinMaxObserver, PerChannelMinMaxObserver @@ -1608,9 +1613,7 @@ def quantize_model( m_eager: torch.nn.Module, example_inputs: Tuple[torch.Tensor] ) -> Tuple[EdgeProgramManager, int, int]: # program capture - m = torch.export.export_for_training( - m_eager, example_inputs, strict=True - ).module() + m = torch.export.export(m_eager, example_inputs, strict=True).module() quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config() diff --git a/exir/tests/test_quantization.py b/exir/tests/test_quantization.py index c7bcdeeeb5c..7cfa926dae3 100644 --- a/exir/tests/test_quantization.py +++ b/exir/tests/test_quantization.py @@ -51,7 +51,7 @@ def test_resnet(self) -> None: m = torchvision.models.resnet18().eval() m_copy = copy.deepcopy(m) # program capture - m = torch.export.export_for_training( + m = torch.export.export( m, copy.deepcopy(example_inputs), strict=True ).module() diff --git a/exir/tests/test_quantize_io_pass.py b/exir/tests/test_quantize_io_pass.py index f670594616a..1dde08ce15b 100644 --- a/exir/tests/test_quantize_io_pass.py +++ b/exir/tests/test_quantize_io_pass.py @@ -38,15 +38,13 @@ def _quantize(self, mod, example_inputs): quantizer = XNNPACKQuantizer() operator_config = get_symmetric_quantization_config() quantizer.set_global(operator_config) - m = torch.export.export_for_training( + m = torch.export.export( mod, copy.deepcopy(example_inputs), strict=True ).module() m = prepare_pt2e(m, quantizer) _ = m(*example_inputs) m = convert_pt2e(m) - exported_program = torch.export.export_for_training( - m, example_inputs, strict=True - ) + exported_program = torch.export.export(m, example_inputs, strict=True) return exported_program def _check_count(self, op, count, epm): diff --git a/exir/tests/test_remove_unused_parameters_pass.py b/exir/tests/test_remove_unused_parameters_pass.py index b7a63b80d82..8eacf692c20 100644 --- a/exir/tests/test_remove_unused_parameters_pass.py +++ b/exir/tests/test_remove_unused_parameters_pass.py @@ -196,7 +196,7 @@ def _test_pass_e2e( self.assertEqual(1, len(runtime_outputs)) self.assertTrue( - torch.allclose(runtime_outputs[0], eager_outputs, atol=2e-6), + torch.allclose(runtime_outputs[0], eager_outputs, atol=1e-5), "Values out of tolerance.\n" + f" Strict: {strict}, ToEdge: {use_to_edge}, Delegate: {delegate}.\n" + f" Eager: {eager_outputs}.\n" diff --git a/exir/tests/test_tracer.py b/exir/tests/test_tracer.py index 594e760ab32..22e01f33332 100644 --- a/exir/tests/test_tracer.py +++ b/exir/tests/test_tracer.py @@ -307,76 +307,6 @@ def f(x: torch.Tensor, y: List[torch.Tensor]) -> Dict[str, torch.Tensor]: self.assertEqual(prog(*inp), f(*inp)) - def test_aot_buffer_mutation(self) -> None: - class Module(torch.nn.Module): - def __init__(self): - super().__init__() - self.register_buffer( - "_bin_num_examples", - torch.empty([42]).fill_( - 0.0, - ), - ) - - def forward(self, x, y, z): - self._bin_num_examples.index_copy_( - dim=0, - index=y, - source=z, - ) - self._bin_num_examples.index_add_( - dim=0, index=torch.arange(4), source=x - ) - return self._bin_num_examples - 1, x * z - - model = Module() - example_inputs = ( - torch.randn(4, requires_grad=True), - torch.tensor(0), - torch.tensor(3.14), - ) - - with self.assertRaisesRegex( - RuntimeError, - "Found a graph input that requires gradients, and received a mutation.", - ): - _ = exir.capture( - model, - example_inputs, - exir.CaptureConfig( - enable_aot=True, - ), - ) - - # Note that model._bin_num_examples is mutated during exir.capture - # We need to create a new_model - new_model = Module() - example_inputs = ( - torch.randn(4), - torch.tensor(0), - torch.tensor(3.14), - ) - - ep = exir.capture( - new_model, - example_inputs, - exir.CaptureConfig( - enable_aot=True, - ), - ) - - test_inputs = ( - torch.randn(4), - torch.tensor(0), - torch.tensor(2.1), - ) - graph_outputs = ep(*test_inputs) - eager_outputs = Module()(*test_inputs) - self.assertEqual(len(graph_outputs), 2) - self.assertEqual(len(eager_outputs), 2) - self.assertTrue(torch.allclose(graph_outputs[0], eager_outputs[0])) - self.assertTrue(torch.allclose(graph_outputs[1], eager_outputs[1])) - def test_assume_constant_by_default_prop(self) -> None: def foo(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: if x.shape[0] > 3: diff --git a/exir/tests/test_verification.py b/exir/tests/test_verification.py index f18e9d74b75..90073216b2d 100644 --- a/exir/tests/test_verification.py +++ b/exir/tests/test_verification.py @@ -252,14 +252,21 @@ def forward(self, x): self.assertTrue(verifier.is_valid(egm)) def test_edge_sad_with_edge_ops(self) -> None: - # log_softmax only takes float or double Tensor - m = torch.nn.LogSoftmax(dim=1) + # add operation does not support complex64 dtype + class TestModel(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x + x + + m = TestModel() with self.assertRaises(SpecViolationError): _ = ( to_edge( export( m, - (torch.randn(1, 3, 100, 100).to(dtype=torch.bfloat16),), + (torch.randn(1, 3, 100, 100).to(dtype=torch.complex64),), strict=True, ) ) diff --git a/exir/tracer.py b/exir/tracer.py index a8a3a16be3b..0a761e6a308 100644 --- a/exir/tracer.py +++ b/exir/tracer.py @@ -48,6 +48,7 @@ from torch._decomp import get_decompositions from torch._dynamo.guards import Guard from torch._functorch.eager_transforms import _maybe_unwrap_functional_tensor + from torch.export import default_decompositions from torch.func import functionalize from torch.fx.operator_schemas import normalize_function diff --git a/exir/verification/TARGETS b/exir/verification/TARGETS index 092b48658df..8362d87a238 100644 --- a/exir/verification/TARGETS +++ b/exir/verification/TARGETS @@ -1,5 +1,5 @@ load("@fbcode_macros//build_defs:cpp_python_extension.bzl", "cpp_python_extension") -load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") @@ -17,7 +17,7 @@ cpp_python_extension( ], ) -python_library( +runtime.python_library( name = "interpreter", srcs = [ "interpreter.py", @@ -32,7 +32,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "devhtml", srcs = [ "dev_html.py", @@ -42,7 +42,7 @@ python_library( ], ) -python_library( +runtime.python_library( name = "verifier", srcs = [ "arg_validator.py", diff --git a/export/TARGETS b/export/TARGETS index 816a3a1a289..ae41393d883 100644 --- a/export/TARGETS +++ b/export/TARGETS @@ -110,3 +110,16 @@ runtime.python_library( "types.py", ], ) + +runtime.python_library( + name = "target_recipes", + srcs = [ + "target_recipes.py", + ], + deps = [ + "fbsource//third-party/pypi/coremltools:coremltools", + "//executorch/export:recipe", + "//executorch/backends/xnnpack/recipes:xnnpack_recipes", + "//executorch/backends/apple/coreml:coreml_recipes", + ] +) diff --git a/export/__init__.py b/export/__init__.py index d5f3826ab90..a7b165185de 100644 --- a/export/__init__.py +++ b/export/__init__.py @@ -15,12 +15,19 @@ """ from .export import export, ExportSession -from .recipe import ExportRecipe, LoweringRecipe, QuantizationRecipe, RecipeType +from .recipe import ( + AOQuantizationConfig, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) from .recipe_provider import BackendRecipeProvider from .recipe_registry import recipe_registry from .types import StageType __all__ = [ + "AOQuantizationConfig", "StageType", "ExportRecipe", "LoweringRecipe", diff --git a/export/export.py b/export/export.py index e5c3b793ccd..86a932d153c 100644 --- a/export/export.py +++ b/export/export.py @@ -44,6 +44,7 @@ def export( dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None, constant_methods: Optional[Union[Dict[str, Callable]]] = None, artifact_dir: Optional[str] = None, + generate_etrecord: bool = False, ) -> "ExportSession": """ Create and configure an ExportSession with the given parameters. @@ -61,6 +62,7 @@ def export( dynamic_shapes: Optional dynamic shape specifications constant_methods: Optional dictionary of constant methods artifact_dir: Optional directory to store artifacts + generate_etrecord: Optional flag to generate an etrecord Returns: A configured ExportSession instance with the export process completed if requested @@ -73,6 +75,7 @@ def export( dynamic_shapes=dynamic_shapes, constant_methods=constant_methods, artifact_dir=artifact_dir, + generate_etrecord=generate_etrecord, ) session.export() @@ -104,6 +107,7 @@ def __init__( dynamic_shapes: Optional[Union[Any, Dict[str, Any]]] = None, constant_methods: Optional[Union[Dict[str, Callable]]] = None, artifact_dir: Optional[str] = None, + generate_etrecord: Optional[bool] = False, ) -> None: """ Initialize the ExportSession with model, inputs, and recipe. @@ -118,6 +122,7 @@ def __init__( dynamic_shapes: Optional dynamic shape specifications constant_methods: Optional dictionary of constant methods artifact_dir: Optional directory to store artifacts + generate_etrecord: Optional flag to generate an etrecord """ # Standardize model to dictionary format self._model = model if isinstance(model, dict) else {"forward": model} @@ -165,6 +170,7 @@ def __init__( "export_recipe": self._export_recipe, "session_name": name, "artifact_dir": artifact_dir, + "generate_etrecord": generate_etrecord, } self._stage_to_artifacts: Dict[StageType, PipelineArtifact] = {} @@ -189,12 +195,12 @@ def _build_stages(self, stages: List[StageType]) -> Dict[StageType, Stage]: elif stage_type == StageType.QUANTIZE: stage = QuantizeStage(self._quant_recipe) elif stage_type == StageType.TORCH_EXPORT: - pre_edge_passes = None - if self._export_recipe.pre_edge_transform_passes is not None: - pre_edge_passes = list( - self._export_recipe.pre_edge_transform_passes + aten_transform_passes = None + if self._export_recipe.aten_transform_passes is not None: + aten_transform_passes = list( + self._export_recipe.aten_transform_passes ) - stage = TorchExportStage(pre_edge_passes) + stage = TorchExportStage(aten_transform_passes) elif stage_type == StageType.TO_EDGE_TRANSFORM_AND_LOWER: stage = EdgeTransformAndLowerStage.from_recipe(self._lowering_recipe) elif stage_type == StageType.TO_EDGE: @@ -446,10 +452,37 @@ def print_delegation_info(self) -> None: """ Print delegation information for the exported program. """ - delegation_info = self._run_context.get("delegation_info", None) + lowering_stage = list( + set(self._pipeline_stages) + & {StageType.TO_EDGE_TRANSFORM_AND_LOWER, StageType.TO_BACKEND} + ) + if not lowering_stage: + RuntimeError( + "No delegation info available, atleast one of the lowering stages should be present" + ) + + stage_artifact = self._stage_to_artifacts.get(lowering_stage[0]) + if stage_artifact is None: + RuntimeError("No delegation info available, run the lowering stage first") + + # pyre-ignore + delegation_info = stage_artifact.get_context("delegation_info", None) if delegation_info: - logging.info(delegation_info.get_summary()) + print(delegation_info.get_summary()) df = delegation_info.get_operator_delegation_dataframe() - logging.info(tabulate(df, headers="keys", tablefmt="fancy_grid")) + print(tabulate(df, headers="keys", tablefmt="fancy_grid")) else: - logging.info("No delegation info available") + print("No delegation info available") + + # Use Any instead of ETRecord as return type to avoid static dependency on etrecord + def get_etrecord(self) -> Any: + """ + Get the etrecord from the ExecuTorchProgramManager. + + Returns: + The etrecord in the ExecuTorchProgramManager + + Raises: + RuntimeError: If the ExecuTorchManager is unavailable, or etrecord is not available in the ExecuTorchProgramManager + """ + return self.get_executorch_program_manager().get_etrecord() diff --git a/export/recipe.py b/export/recipe.py index 8f7251cd419..18f4b8aebb9 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -3,10 +3,14 @@ # # 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 from abc import ABCMeta, abstractmethod from dataclasses import dataclass from enum import Enum, EnumMeta -from typing import List, Optional, Sequence +from typing import Callable, List, Optional + +import torch +from executorch.exir import ExportedProgram from executorch.exir._warnings import experimental @@ -64,6 +68,20 @@ class Mode(str, Enum): RELEASE = "release" +@dataclass +class AOQuantizationConfig: + """ + Configuration for torchao quantization with optional filter function. + + Attributes: + ao_base_config: The AOBaseConfig for quantization + filter_fn: Optional filter function to selectively apply quantization + """ + + ao_base_config: AOBaseConfig + filter_fn: Optional[Callable[[torch.nn.Module, str], bool]] = None + + @dataclass class QuantizationRecipe: """ @@ -73,11 +91,12 @@ class QuantizationRecipe: Attributes: quantizers: Optional list of quantizers for model quantization - ao_base_config: Optional list of AO base configurations + ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair + AOBaseConfig with optional filter functions """ quantizers: Optional[List[Quantizer]] = None - ao_base_config: Optional[List[AOBaseConfig]] = None + ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None def get_quantizers(self) -> Optional[List[Quantizer]]: """ @@ -99,12 +118,15 @@ class LoweringRecipe: Attributes: partitioners: Optional list of partitioners for model partitioning - edge_transform_passes: Optional sequence of transformation passes to apply + edge_transform_passes: Optional list of callables that take (method_name: str, exported_program: ExportedProgram) as arguments + and return a list of passes (PassType) to be executed during lowering stages. edge_compile_config: Optional edge compilation configuration """ partitioners: Optional[List[Partitioner]] = None - edge_transform_passes: Optional[Sequence[PassType]] = None + edge_transform_passes: ( + None | List[Callable[[str, ExportedProgram], List[PassType]]] + ) = None # pyre-ignore[11]: Type not defined edge_compile_config: Optional[EdgeCompileConfig] = None @@ -123,8 +145,8 @@ class ExportRecipe: Attributes: name: Optional name for the recipe quantization_recipe: Optional quantization recipe for model quantization - pre_edge_transform_passes: Optional function to apply transformation passes - before edge lowering + aten_transform_passes: Optional list of functions to apply transformation passes to the program before edge lowering. + These callables are invoked to modify and return the transformed program. lowering_recipe: Optional lowering recipe for model lowering and partitioning executorch_backend_config: Optional backend configuration for ExecuTorch pipeline_stages: Optional list of stages to execute, defaults to a standard pipeline. @@ -133,7 +155,9 @@ class ExportRecipe: name: Optional[str] = None quantization_recipe: Optional[QuantizationRecipe] = None - pre_edge_transform_passes: Optional[Sequence[PassType]] = None + aten_transform_passes: Optional[ + List[Callable[[str, ExportedProgram], ExportedProgram]] + ] = None lowering_recipe: Optional[LoweringRecipe] = None # pyre-ignore[11]: Type not defined executorch_backend_config: Optional[ExecutorchBackendConfig] = None @@ -167,3 +191,129 @@ def get_recipe(cls, recipe: "RecipeType", **kwargs) -> "ExportRecipe": f"Supported: {[r.value for r in supported]}" ) return export_recipe + + @classmethod + def combine( + cls, recipes: List["ExportRecipe"], recipe_name: Optional[str] = None + ) -> "ExportRecipe": + """ + Combine multiple ExportRecipe objects into a single recipe. + + Args: + recipes: List of ExportRecipe objects to combine + recipe_name: Optional name for the combined recipe + + Returns: + A new ExportRecipe that combines all input recipes + + Example: + recipe1 = ExportRecipe.get_recipe(CoreMLRecipeType.FP32) + recipe2 = ExportRecipe.get_recipe(XNNPackRecipeType.FP32) + combined_recipe = ExportRecipe.combine( + [recipe1, recipe2], + recipe_name="multi_backend_coreml_xnnpack_fp32" + ) + """ + if not recipes: + raise ValueError("Recipes cannot be empty") + + if len(recipes) == 1: + return recipes[0] + + return cls._combine_recipes(recipes, recipe_name) + + @classmethod + def _combine_recipes( # noqa: C901 + 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 + """ + # Extract components from individual recipes + all_partitioners = [] + all_quantizers = [] + all_ao_quantization_configs = [] + all_pre_edge_passes = [] + all_transform_passes = [] + combined_backend_config = None + + 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: + all_partitioners.extend(recipe.lowering_recipe.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 + ) + + # 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 + ), + ) + + # Create combined lowering recipe + combined_lowering_recipe = None + if all_partitioners or all_transform_passes: + edge_compile_config = None + for recipe in backend_recipes: + if ( + recipe.lowering_recipe + and recipe.lowering_recipe.edge_compile_config + ): + edge_compile_config = recipe.lowering_recipe.edge_compile_config + break + + combined_lowering_recipe = LoweringRecipe( + partitioners=all_partitioners if all_partitioners else None, + edge_transform_passes=( + all_transform_passes if all_transform_passes else None + ), + edge_compile_config=edge_compile_config or EdgeCompileConfig(), + ) + + recipe_name = recipe_name or "_".join( + [r.name for r in backend_recipes if r.name is not None] + ) + return cls( + name=recipe_name, + quantization_recipe=combined_quantization_recipe, + aten_transform_passes=all_pre_edge_passes, + lowering_recipe=combined_lowering_recipe, + executorch_backend_config=combined_backend_config, + ) diff --git a/export/stages.py b/export/stages.py index dd22155e929..323b327bfa4 100644 --- a/export/stages.py +++ b/export/stages.py @@ -4,23 +4,27 @@ # 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 logging from abc import ABC, abstractmethod -from typing import Any, Callable, Dict, List, Optional, Sequence +from collections import defaultdict +from typing import Any, Callable, Dict, List, Optional import torch from executorch.devtools.backend_debug import get_delegation_info -from executorch.exir import EdgeCompileConfig +from executorch.exir import EdgeCompileConfig, ExportedProgram from executorch.exir.backend.backend_api import validation_disabled from executorch.exir.program import to_edge, to_edge_transform_and_lower -from executorch.exir.program._program import _transform from executorch.export.recipe import LoweringRecipe, QuantizationRecipe from executorch.export.types import StageType from torch import nn from torch._export.pass_base import PassType from torchao.quantization import quantize_ from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e -from torchao.quantization.pt2e.quantizer import ComposableQuantizer +from torchao.quantization.pt2e.quantizer import ( + ComposableQuantizer, + Quantizer as TorchAOPT2EQuantizer, +) from torchao.utils import unwrap_tensor_subclass @@ -103,10 +107,12 @@ class TorchExportStage(Stage): def __init__( self, - pre_edge_transform_passes: Optional[List[PassType]] = None, + aten_transform_passes: Optional[ + List[Callable[[str, ExportedProgram], ExportedProgram]] + ] = None, ) -> None: super().__init__() - self._pre_edge_transform_passes = pre_edge_transform_passes + self._aten_transform_passes = aten_transform_passes @property def stage_type(self) -> str: @@ -145,9 +151,13 @@ def run(self, artifact: PipelineArtifact) -> None: ) # Apply pre-edge transform passes if available - for pass_ in self._pre_edge_transform_passes or []: - exported_programs[method_name] = _transform( - exported_programs[method_name], pass_ + for pass_ in self._aten_transform_passes or []: + if not callable(pass_): + raise ValueError( + "Aten transform passes must be a callable that can transform and return an exported program" + ) + exported_programs[method_name] = pass_( + method_name, exported_programs[method_name] ) self._artifact = artifact.copy_with_new_data(exported_programs) @@ -161,7 +171,9 @@ class EdgeTransformAndLowerStage(Stage): def __init__( self, partitioners: Optional[List[Any]] = None, - transform_passes: Optional[Sequence[Callable[[Any], Optional[Any]]]] = None, + transform_passes: ( + None | List[Callable[[str, ExportedProgram], List[PassType]]] + ) = None, compile_config: Optional[Any] = None, ) -> None: self._partitioners = partitioners @@ -199,14 +211,33 @@ def run(self, artifact: PipelineArtifact) -> None: """ exported_programs = artifact.data constant_methods = artifact.get_context("constant_methods") + generate_etrecord = artifact.get_context("generate_etrecord", False) + + # per method transform passes + transform_passes = defaultdict(list) + for method_name, ep in exported_programs.items(): + # Resolve transform passes from callable + for pass_ in self._transform_passes or []: + if not callable(pass_): + raise ValueError( + "Transform passes must be a callable that resolves to a list of passes" + ) + passes = pass_(method_name, ep) + if isinstance(passes, list): + transform_passes[method_name].extend(passes) + else: + raise ValueError( + "Transform passes must be a callable that resolves to a list of passes" + ) with validation_disabled(): edge_program_manager = to_edge_transform_and_lower( exported_programs, partitioner=self._partitioners, - transform_passes=self._transform_passes, + transform_passes=transform_passes, constant_methods=constant_methods, compile_config=self._compile_config, + generate_etrecord=generate_etrecord, ) delegation_info = get_delegation_info( @@ -287,7 +318,7 @@ def run(self, artifact: PipelineArtifact) -> None: """ if ( not self._quantization_recipe - or not self._quantization_recipe.ao_base_config + or not self._quantization_recipe.ao_quantization_configs ): logging.info( "Quantization recipe is invalid to run SourceTransform, returning original artifact" @@ -298,15 +329,19 @@ def run(self, artifact: PipelineArtifact) -> None: assert isinstance(artifact.data, dict) # Store the original models - self._transformed_models = artifact.data + self._transformed_models = copy.deepcopy(artifact.data) # Apply torchao quantize_ to each model - for method_name, model in artifact.data.items(): + for _, model in self._transformed_models.items(): # pyre-ignore - for config in self._quantization_recipe.ao_base_config: - quantize_(model, config) - unwrap_tensor_subclass(model) - self._transformed_models[method_name] = model + if len(self._quantization_recipe.ao_quantization_configs) > 1: + raise ValueError( + "AO quantization configs cannot be reliably composed together, multiple quantization configs are disallowed for source transform at this point" + ) + + ao_config = self._quantization_recipe.ao_quantization_configs[0] + quantize_(model, ao_config.ao_base_config, ao_config.filter_fn) + unwrap_tensor_subclass(model) self._artifact = artifact.copy_with_new_data(self._transformed_models) @@ -331,6 +366,36 @@ def valid_predecessor_stages(self) -> List["StageType"]: def can_start_pipeline(self) -> bool: return True + def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]): + torch_ao_quantizers = [] + torchao_pt2e_quantizers = [] + + for quantizer in quantizers: + if isinstance(quantizer, TorchAOPT2EQuantizer): + torchao_pt2e_quantizers.append(quantizer) + else: + # torch.ao quantizer support will soon be deprecated, remove this once CoreML moves to torchao quantizer + logging.warning( + f"torch.ao quantizer {quantizer} is deprecated, consider moving to torchao quantizer" + ) + torch_ao_quantizers.append(quantizer) + + if torch_ao_quantizers and torchao_pt2e_quantizers: + raise ValueError("Mixed quantizer types are not supported") + if len(torch_ao_quantizers) > 1: + raise ValueError( + "Multiple quantizers of torch.ao.quantization.quantizer not supported" + ) + + if torch_ao_quantizers: + # prepare_pt2e has backward compat with torch.ao quantizer + return torch_ao_quantizers[0] + elif torchao_pt2e_quantizers: + # Multiple torchao quantizers - use ComposableQuantizer + return ComposableQuantizer(torchao_pt2e_quantizers) + else: + raise ValueError("No quantizers detected") + def run(self, artifact: PipelineArtifact) -> None: if not self._quantization_recipe or not self._quantization_recipe.quantizers: logging.info( @@ -355,11 +420,10 @@ def run(self, artifact: PipelineArtifact) -> None: inputs = example_inputs[method_name][0] captured_graph = torch.export.export(model, inputs, strict=True).module() - composed_quantizer = ComposableQuantizer( - # pyre-ignore - self._quantization_recipe.quantizers + quantizer = self._get_quantizer_for_prepare_pt2e( + self._quantization_recipe.quantizers # pyre-ignore ) - prepared_model = prepare_pt2e(captured_graph, composed_quantizer) + prepared_model = prepare_pt2e(captured_graph, quantizer) for calibration_input in example_inputs[method_name]: prepared_model(*calibration_input) @@ -418,6 +482,7 @@ def run(self, artifact: PipelineArtifact) -> None: exported_programs, constant_methods=constant_methods, compile_config=self._edge_compile_config, + generate_etrecord=artifact.get_context("generate_etrecord", False), ) self._artifact = artifact.copy_with_new_data(edge_program_manager) @@ -431,7 +496,9 @@ class ToBackendStage(Stage): def __init__( self, partitioners: Optional[List[Any]] = None, - transform_passes: Optional[Sequence[Callable[[Any], Optional[Any]]]] = None, + transform_passes: ( + None | List[Callable[[str, ExportedProgram], List[PassType]]] + ) = None, ) -> None: super().__init__() self._partitioners = partitioners @@ -473,11 +540,24 @@ def run(self, artifact: PipelineArtifact) -> None: if edge_program_manager is None: raise RuntimeError("Edge program manager is not set.") - # Apply transform passes if available - if self._transform_passes: - edge_program_manager = edge_program_manager.transform( - self._transform_passes - ) + # per method transform passes + transform_passes = defaultdict(list) + for method_name in edge_program_manager.methods: + # Resolve transform passes if it's a callable + ep = edge_program_manager.exported_program(method_name) + for pass_ in self._transform_passes or []: + if not callable(pass_): + raise ValueError( + "Transform passes must be a callable that resolves to a list of passes" + ) + passes = pass_(method_name, ep) + if isinstance(passes, list): + transform_passes[method_name].extend(passes) + else: + raise ValueError("Transform passes must return list of passes") + + # Apply transform passes + edge_program_manager = edge_program_manager.transform(transform_passes) # Apply partitioners if available if self._partitioners is not None and len(self._partitioners) > 0: diff --git a/export/target_recipes.py b/export/target_recipes.py new file mode 100644 index 00000000000..0a5ae9ce754 --- /dev/null +++ b/export/target_recipes.py @@ -0,0 +1,117 @@ +# 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. + +""" +Target-specific recipe functions for simplified multi-backend deployment. + +This module provides platform-specific functions that abstract away backend +selection and combine multiple backends optimally for target hardware. +""" + +import sys +from typing import Dict, List + +if sys.platform != "win32": + import coremltools as ct + from executorch.backends.apple.coreml.recipes import CoreMLRecipeType + +# pyre-ignore +from executorch.backends.xnnpack.recipes import XNNPackRecipeType +from executorch.export.recipe import ExportRecipe, RecipeType + + +## IOS Target configs +# The following list of recipes are not exhaustive for CoreML; refer to CoreMLRecipeType for more detailed recipes. +IOS_CONFIGS: Dict[str, List[RecipeType]] = ( + { + # pyre-ignore + "ios-arm64-coreml-fp32": [CoreMLRecipeType.FP32, XNNPackRecipeType.FP32], + # pyre-ignore + "ios-arm64-coreml-fp16": [CoreMLRecipeType.FP16], + # pyre-ignore + "ios-arm64-coreml-int8": [CoreMLRecipeType.PT2E_INT8_STATIC], + } + if sys.platform != "win32" + else {} +) + + +def _create_target_recipe( + target_config: str, recipes: List[RecipeType], **kwargs +) -> ExportRecipe: + """ + Create a combined recipe for a target. + + Args: + target: Human-readable hardware configuration name + recipes: List of backend recipe types to combine + **kwargs: Additional parameters - each backend will use what it needs + + Returns: + Combined ExportRecipe for the hardware configuration + """ + if not recipes: + raise ValueError(f"No backends configured for: {target_config}") + + # Create individual backend recipes + backend_recipes = [] + for recipe_type in recipes: + try: + backend_recipe = ExportRecipe.get_recipe(recipe_type, **kwargs) + backend_recipes.append(backend_recipe) + except Exception as e: + raise ValueError( + f"Failed to create {recipe_type.value} recipe for {target_config}: {e}" + ) from e + + # Combine into single recipe + if len(backend_recipes) == 1: + return backend_recipes[0] + + return ExportRecipe.combine(backend_recipes, recipe_name=target_config) + + +# IOS Recipe +def get_ios_recipe( + target_config: str = "ios-arm64-coreml-fp16", **kwargs +) -> ExportRecipe: + """ + Get iOS-optimized recipe for specified hardware configuration. + + Supported configurations: + - 'ios-arm64-coreml-fp32': CoreML + XNNPACK fallback (FP32) + - 'ios-arm64-coreml-fp16': CoreML fp16 recipe + - 'ios-arm64-coreml-int8': CoreML INT8 quantization recipe + + Args: + target_config: iOS configuration string + **kwargs: Additional parameters for backend recipes + + Returns: + ExportRecipe configured for iOS deployment + + Raises: + ValueError: If target configuration is not supported + + Example: + recipe = get_ios_recipe('ios-arm64-coreml-int8') + session = export(model, recipe, example_inputs) + """ + if target_config not in IOS_CONFIGS: + supported = list(IOS_CONFIGS.keys()) + raise ValueError( + f"Unsupported iOS configuration: '{target_config}'. " + f"Supported: {supported}" + ) + + kwargs = kwargs or {} + + if target_config == "ios-arm64-coreml-int8": + if "minimum_deployment_target" not in kwargs: + kwargs["minimum_deployment_target"] = ct.target.iOS17 + + backend_recipes = IOS_CONFIGS[target_config] + return _create_target_recipe(target_config, backend_recipes, **kwargs) diff --git a/export/tests/TARGETS b/export/tests/TARGETS index 068c3436b6a..71f28b64df7 100644 --- a/export/tests/TARGETS +++ b/export/tests/TARGETS @@ -31,3 +31,17 @@ runtime.python_test( "//executorch/runtime:runtime", ] ) + +runtime.python_test( + name = "test_target_recipes", + srcs = [ + "test_target_recipes.py", + ], + deps = [ + "//executorch/export:lib", + "//executorch/export:target_recipes", + "//executorch/runtime:runtime", + "//executorch/backends/xnnpack/recipes:xnnpack_recipes", + "//executorch/backends/apple/coreml:coreml_recipes", + ] +) diff --git a/export/tests/test_export_session.py b/export/tests/test_export_session.py index 92aeebb7304..fcec1b7a59a 100644 --- a/export/tests/test_export_session.py +++ b/export/tests/test_export_session.py @@ -12,7 +12,11 @@ import torch from executorch.export import ExportRecipe, ExportSession -from executorch.export.recipe import LoweringRecipe, QuantizationRecipe +from executorch.export.recipe import ( + AOQuantizationConfig, + LoweringRecipe, + QuantizationRecipe, +) from executorch.export.stages import PipelineArtifact from executorch.export.types import StageType @@ -20,7 +24,7 @@ class SimpleTestModel(torch.nn.Module): def __init__(self) -> None: super().__init__() - self.linear = torch.nn.Linear(10, 5) + self.linear: torch.nn.Module = torch.nn.Linear(10, 5) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(x) @@ -184,6 +188,7 @@ def test_context_propagation_through_pipeline(self) -> None: "export_recipe", "session_name", "artifact_dir", + "generate_etrecord", } self.assertEqual(set(session._run_context.keys()), expected_context_keys) self.assertEqual(session._run_context["session_name"], "test_session") @@ -448,7 +453,7 @@ def test_pipeline_building_with_all_recipes(self) -> None: """Test pipeline building with quantization and lowering recipes.""" # Create comprehensive recipes quant_recipe = QuantizationRecipe( - ao_base_config=[Mock()], + ao_quantization_configs=[AOQuantizationConfig(Mock())], quantizers=[Mock()], ) lowering_recipe = LoweringRecipe( diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index 2b3e533723a..4e8144bd487 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -7,29 +7,29 @@ # pyre-strict import unittest -from unittest.mock import Mock, patch +from unittest.mock import Mock, patch, PropertyMock import torch from executorch.exir.program import EdgeProgramManager, ExecutorchProgramManager -from executorch.export import QuantizationRecipe +from executorch.export import AOQuantizationConfig, QuantizationRecipe, StageType from executorch.export.stages import ( EdgeTransformAndLowerStage, ExecutorchStage, PipelineArtifact, QuantizeStage, SourceTransformStage, - StageType, ToBackendStage, ToEdgeStage, TorchExportStage, ) from torch.export import ExportedProgram +from torchao.quantization.pt2e.quantizer import Quantizer as TorchAOPT2EQuantizer class SimpleTestModel(torch.nn.Module): def __init__(self) -> None: super().__init__() - self.linear = torch.nn.Linear(10, 5) + self.linear: torch.nn.Module = torch.nn.Linear(10, 5) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(x) @@ -99,6 +99,66 @@ def test_get_artifacts_before_run(self) -> None: stage.get_artifacts() self.assertIn("Stage: TorchExportStage not executed", str(cm.exception)) + @patch("torch.export.export") + def test_export_stage_with_aten_transform_passes( + self, mock_torch_export: Mock + ) -> None: + """Test TorchExportStage with aten_transform_passes.""" + mock_exported_program = Mock(spec=ExportedProgram) + mock_transformed_program = Mock(spec=ExportedProgram) + mock_torch_export.return_value = mock_exported_program + + # Create a mock aten transform pass that we can verify + mock_aten_transform_pass = Mock() + mock_aten_transform_pass.return_value = mock_transformed_program + aten_transform_passes = [mock_aten_transform_pass] + + stage = TorchExportStage(aten_transform_passes=aten_transform_passes) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + + stage.run(artifact) + + # Verify torch.export.export was called + mock_torch_export.assert_called_once_with( + self.model, + self.example_inputs[0], + dynamic_shapes=None, + strict=True, + ) + + # Verify the aten transform pass was called with correct parameters + mock_aten_transform_pass.assert_called_once_with( + "forward", mock_exported_program + ) + + # Verify artifacts contain the transformed program + result_artifact = stage.get_artifacts() + self.assertIn("forward", result_artifact.data) + self.assertEqual(result_artifact.data["forward"], mock_transformed_program) + + @patch("torch.export.export") + def test_export_stage_invalid_aten_transform_pass( + self, mock_torch_export: Mock + ) -> None: + """Test TorchExportStage with invalid aten_transform_pass (not callable).""" + mock_exported_program = Mock(spec=ExportedProgram) + mock_torch_export.return_value = mock_exported_program + + # Use a non-callable object as transform pass + invalid_transform_pass = "not_callable" + aten_transform_passes = [invalid_transform_pass] + + # pyre-ignore + stage = TorchExportStage(aten_transform_passes=aten_transform_passes) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + + with self.assertRaises(ValueError) as cm: + stage.run(artifact) + self.assertIn( + "Aten transform passes must be a callable that can transform and return an exported program", + str(cm.exception), + ) + class TestEdgeTransformAndLowerStage(unittest.TestCase): def setUp(self) -> None: @@ -106,12 +166,32 @@ def setUp(self) -> None: self.exported_programs = {"forward": self.mock_exported_program} self.context = {"constant_methods": None} - def test_run_with_partitioners_and_config(self) -> None: + @patch("executorch.export.stages.to_edge_transform_and_lower") + @patch("executorch.export.stages.get_delegation_info") + def test_run_with_partitioners_and_config( + self, mock_get_delegation_info: Mock, mock_to_edge_transform_and_lower: Mock + ) -> None: """Test execution with partitioners and compile config""" + mock_delegation_info = {"delegation": "info"} + mock_get_delegation_info.return_value = mock_delegation_info + mock_partitioners = [Mock()] - mock_transform_passes = [Mock()] mock_compile_config = Mock() + # Create a mock transform pass callable that we can verify + mock_transform_pass = Mock() + mock_pass1 = Mock() + mock_pass2 = Mock() + mock_transform_pass.return_value = [mock_pass1, mock_pass2] + mock_transform_passes = [mock_transform_pass] + + mock_edge_program_manager = Mock(spec=EdgeProgramManager) + mock_exported_program = Mock() + mock_graph_module = Mock() + mock_exported_program.graph_module = mock_graph_module + mock_edge_program_manager.exported_program.return_value = mock_exported_program + mock_to_edge_transform_and_lower.return_value = mock_edge_program_manager + stage = EdgeTransformAndLowerStage( partitioners=mock_partitioners, transform_passes=mock_transform_passes, @@ -124,6 +204,33 @@ def test_run_with_partitioners_and_config(self) -> None: self.assertEqual(stage._transform_passes, mock_transform_passes) self.assertEqual(stage._compile_config, mock_compile_config) + # Test the run method + artifact = PipelineArtifact(data=self.exported_programs, context=self.context) + stage.run(artifact) + + # Verify the transform pass callable was called with correct parameters + mock_transform_pass.assert_called_once_with( + "forward", self.mock_exported_program + ) + + # Verify to_edge_transform_and_lower was called with the expected structure + expected_transform_passes = {"forward": [mock_pass1, mock_pass2]} + mock_to_edge_transform_and_lower.assert_called_once_with( + self.exported_programs, + partitioner=mock_partitioners, + transform_passes=expected_transform_passes, + constant_methods=None, + compile_config=mock_compile_config, + generate_etrecord=False, + ) + + # Verify artifacts are set correctly + result_artifact = stage.get_artifacts() + self.assertEqual(result_artifact.data, mock_edge_program_manager) + self.assertEqual( + result_artifact.get_context("delegation_info"), mock_delegation_info + ) + class TestExecutorchStage(unittest.TestCase): def setUp(self) -> None: @@ -163,7 +270,7 @@ def setUp(self) -> None: def test_source_transform_stage_no_quantization(self) -> None: mock_recipe = Mock(spec=QuantizationRecipe) - mock_recipe.ao_base_config = None + mock_recipe.ao_quantization_configs = None stage = SourceTransformStage(mock_recipe) artifact = PipelineArtifact(data=self.models_dict, context={}) @@ -174,12 +281,18 @@ def test_source_transform_stage_no_quantization(self) -> None: @patch("executorch.export.stages.quantize_") @patch("executorch.export.stages.unwrap_tensor_subclass") - def test_run_with_ao_base_config( + def test_run_with_ao_quantization_configs( self, mock_unwrap: Mock, mock_quantize: Mock ) -> None: - mock_config = Mock() + from torchao.core.config import AOBaseConfig + + mock_config = Mock(spec=AOBaseConfig) + mock_filter_fn = Mock() + mock_ao_config: AOQuantizationConfig = AOQuantizationConfig( + ao_base_config=mock_config, filter_fn=mock_filter_fn + ) mock_recipe = Mock(spec=QuantizationRecipe) - mock_recipe.ao_base_config = [mock_config] + mock_recipe.ao_quantization_configs = [mock_ao_config] stage = SourceTransformStage(mock_recipe) @@ -187,11 +300,26 @@ def test_run_with_ao_base_config( artifact = PipelineArtifact(data=models_dict, context={}) stage.run(artifact) - # Verify quantize_ was called with the model and config - mock_quantize.assert_called_once_with(self.model, mock_config) + # Verify quantize_ was called once (with the copied model, not the original) + self.assertEqual(mock_quantize.call_count, 1) + # Verify the config and filter_fn arguments are correct + call_args = mock_quantize.call_args[0] + self.assertNotEqual(self.model, call_args[0]) + self.assertEqual(call_args[1], mock_config) + self.assertEqual(call_args[2], mock_filter_fn) + + # Verify unwrap_tensor_subclass was called once (with the copied model) + self.assertEqual(mock_unwrap.call_count, 1) - # Verify unwrap_tensor_subclass was called with the model - mock_unwrap.assert_called_once_with(self.model) + # Verify that the original models_dict is unchanged + self.assertEqual(models_dict, {"forward": self.model}) + + # Verify that the result artifact data contains valid models + result_artifact = stage.get_artifacts() + self.assertIn("forward", result_artifact.data) + self.assertIsNotNone(result_artifact.data["forward"]) + # verify the result model is NOT the same object as the original + self.assertIsNot(result_artifact.data["forward"], self.model) class TestQuantizeStage(unittest.TestCase): @@ -201,6 +329,21 @@ def setUp(self) -> None: self.example_inputs = [(torch.randn(2, 10),)] self.context = {"example_inputs": {"forward": self.example_inputs}} + @staticmethod + def create_dummy_quantizer() -> TorchAOPT2EQuantizer: + + class DummyQuantizer(TorchAOPT2EQuantizer): + def __init__(self): + pass + + def annotate(self, model): + return model + + def validate(self, model): + pass + + return DummyQuantizer() + def test_run_no_quantizers(self) -> None: """Test execution with no quantizers.""" mock_recipe = Mock(spec=QuantizationRecipe) @@ -224,7 +367,7 @@ def test_run_with_quantizers( mock_convert_pt2e: Mock, ) -> None: """Test execution with quantizers""" - mock_quantizer = Mock() + mock_quantizer = self.create_dummy_quantizer() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] stage = QuantizeStage(mock_recipe) @@ -270,6 +413,10 @@ def test_run_with_quantizers( self.assertIn("forward", result_artifact.data) self.assertEqual(result_artifact.data["forward"], mock_quantized_model) + # Verify that the original model in the input artifact is unchanged + self.assertEqual(artifact.data["forward"], self.model) + self.assertIsNot(result_artifact.data["forward"], self.model) + def test_run_empty_example_inputs(self) -> None: """Test error when example inputs list is empty.""" mock_quantizer = Mock() @@ -285,6 +432,35 @@ def test_run_empty_example_inputs(self) -> None: "Example inputs for method forward not found or empty", str(cm.exception) ) + @patch("executorch.export.stages.ComposableQuantizer") + def test_get_quantizer_for_prepare_pt2e( + self, mock_composable_quantizer: Mock + ) -> None: + """Test _get_quantizer_for_prepare_pt2e method with different quantizer scenarios.""" + mock_recipe = Mock(spec=QuantizationRecipe) + stage = QuantizeStage(mock_recipe) + + # Test empty quantizers list - should raise ValueError + with self.assertRaises(ValueError) as cm: + stage._get_quantizer_for_prepare_pt2e([]) + self.assertIn("No quantizers detected", str(cm.exception)) + + # Test ComposableQuantizer path with multiple torchao quantizers + # Create instances of dummy quantizers using the reusable method + quantizer1 = self.create_dummy_quantizer() + quantizer2 = self.create_dummy_quantizer() + + # Set up ComposableQuantizer mock + mock_composed_quantizer = Mock() + mock_composable_quantizer.return_value = mock_composed_quantizer + + # Call the method with multiple torchao quantizers + result = stage._get_quantizer_for_prepare_pt2e([quantizer1, quantizer2]) + + # Verify ComposableQuantizer was called with the quantizers + mock_composable_quantizer.assert_called_once_with([quantizer1, quantizer2]) + self.assertEqual(result, mock_composed_quantizer) + class TestToEdgeStage(unittest.TestCase): def setUp(self) -> None: @@ -307,6 +483,7 @@ def test_run_success(self, mock_to_edge: Mock) -> None: self.exported_programs, constant_methods=None, compile_config=mock_config, + generate_etrecord=False, ) # Verify artifacts are set correctly @@ -329,7 +506,10 @@ def test_run_success_no_transforms_or_partitioners( mock_exported_program = Mock() mock_graph_module = Mock() mock_exported_program.graph_module = mock_graph_module + + self.mock_edge_manager.transform.return_value = self.mock_edge_manager self.mock_edge_manager.exported_program.return_value = mock_exported_program + self.mock_edge_manager.methods = {"forward"} stage = ToBackendStage() artifact = PipelineArtifact(data=self.mock_edge_manager, context=self.context) @@ -358,9 +538,21 @@ def test_run_with_partitioners_and_passes( mock_edge_program_manager = Mock(spec=EdgeProgramManager) mock_edge_program_manager.transform.return_value = mock_edge_program_manager mock_edge_program_manager.to_backend.return_value = mock_edge_program_manager + mock_edge_program_manager.exported_program.return_value = mock_exported_program + + # Use PropertyMock for the methods property + methods_property_mock = PropertyMock(return_value={"forward"}) + type(mock_edge_program_manager).methods = methods_property_mock mock_partitioner = Mock() - mock_transform_passes = [Mock(), Mock()] + + # Create a mock transform pass callable that we can verify + mock_transform_pass = Mock() + mock_pass1 = Mock() + mock_pass2 = Mock() + mock_transform_pass.return_value = [mock_pass1, mock_pass2] + mock_transform_passes = [mock_transform_pass] + stage = ToBackendStage( partitioners=[mock_partitioner], transform_passes=mock_transform_passes ) @@ -369,10 +561,19 @@ def test_run_with_partitioners_and_passes( ) stage.run(artifact) - # Verify transform and to_backend called correctly + # Verify that the methods property was accessed + methods_property_mock.assert_called_once() + + # Verify the transform pass callable was called with correct parameters + mock_transform_pass.assert_called_once_with("forward", mock_exported_program) + + # Verify transform was called with the expected structure + expected_transform_passes = {"forward": [mock_pass1, mock_pass2]} mock_edge_program_manager.transform.assert_called_once_with( - mock_transform_passes + expected_transform_passes ) + + # Verify to_backend called correctly mock_edge_program_manager.to_backend.assert_called_once_with(mock_partitioner) # Verify artifacts contain the backend manager diff --git a/export/tests/test_target_recipes.py b/export/tests/test_target_recipes.py new file mode 100644 index 00000000000..7a2a7c87342 --- /dev/null +++ b/export/tests/test_target_recipes.py @@ -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. + +# pyre-strict + +import logging +import sys +import unittest + +import torch +from executorch.backends.xnnpack.recipes.xnnpack_recipe_provider import ( + XNNPACKRecipeProvider, +) +from executorch.export import export, recipe_registry +from executorch.export.target_recipes import get_ios_recipe +from executorch.runtime import Runtime + +if sys.platform != "win32": + from executorch.backends.apple.coreml.recipes import ( # pyre-ignore + CoreMLRecipeProvider, + ) + + +class TestTargetRecipes(unittest.TestCase): + """Test target recipes.""" + + def setUp(self) -> None: + torch._dynamo.reset() + super().setUp() + recipe_registry.register_backend_recipe_provider(XNNPACKRecipeProvider()) + if sys.platform != "win32": + # pyre-ignore + recipe_registry.register_backend_recipe_provider(CoreMLRecipeProvider()) + + def tearDown(self) -> None: + super().tearDown() + + @unittest.skipIf(sys.platform == "win32", "Core ML is not available on Windows.") + def test_ios_fp32_recipe_with_xnnpack_fallback(self) -> None: + # Linear ops skipped by coreml but handled by xnnpack + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(4, 4) + self.linear2 = torch.nn.Linear(4, 2) + + def forward(self, x, y): + a = self.linear1(x) + b = a + y + c = b - x + result = self.linear2(c) + return result + + model = Model() + model.eval() + + example_inputs = [(torch.randn(2, 4), torch.randn(2, 4))] + + # Export using multi-backend target recipe with CoreML configured to skip linear operations + recipe = get_ios_recipe( + "ios-arm64-coreml-fp32", + skip_ops_for_coreml_delegation=["aten.linear.default"], + ) + + # Export the model + session = export( + model=model, example_inputs=example_inputs, export_recipe=recipe + ) + + # Verify we can create executable + executorch_program = session.get_executorch_program() + # session.print_delegation_info() + + self.assertIsNotNone( + executorch_program, "ExecuTorch program should not be None" + ) + + # Assert there is an execution plan + self.assertTrue(len(executorch_program.execution_plan) == 1) + + # Check number of partitions created + self.assertTrue(len(executorch_program.execution_plan[0].delegates) == 3) + + # First delegate backend is Xnnpack + self.assertEqual( + executorch_program.execution_plan[0].delegates[0].id, + "XnnpackBackend", + ) + + # Second delegate backend is CoreML + self.assertEqual( + executorch_program.execution_plan[0].delegates[1].id, + "CoreMLBackend", + ) + + # Third delegate backend is Xnnpack + self.assertEqual( + executorch_program.execution_plan[0].delegates[2].id, + "XnnpackBackend", + ) + + et_runtime: Runtime = Runtime.get() + backend_registry = et_runtime.backend_registry + logging.info( + f"backends registered: {et_runtime.backend_registry.registered_backend_names}" + ) + if backend_registry.is_available( + "CoreMLBackend" + ) and backend_registry.is_available("XnnpackBackend"): + logging.info("Running with CoreML and XNNPACK backends") + et_output = session.run_method("forward", example_inputs[0]) + logging.info(f"et output {et_output}") + + @unittest.skipIf(sys.platform == "win32", "Core ML is not available on Windows.") + def test_ios_quant_recipes(self) -> None: + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear1 = torch.nn.Linear(4, 4) + self.linear2 = torch.nn.Linear(4, 2) + + def forward(self, x, y): + a = self.linear1(x) + b = a + y + c = b - x + result = self.linear2(c) + return result + + model = Model() + model.eval() + + example_inputs = [(torch.randn(2, 4), torch.randn(2, 4))] + + for recipe in [ + get_ios_recipe("ios-arm64-coreml-fp16"), + get_ios_recipe("ios-arm64-coreml-int8"), + ]: + # Export the model + session = export( + model=model, example_inputs=example_inputs, export_recipe=recipe + ) + + # Verify we can create executable + executorch_program = session.get_executorch_program() + session.print_delegation_info() + + self.assertIsNotNone( + executorch_program, "ExecuTorch program should not be None" + ) + + # Assert there is an execution plan + self.assertTrue(len(executorch_program.execution_plan) == 1) + + # Check number of partitions created + self.assertTrue(len(executorch_program.execution_plan[0].delegates) == 1) + + # Delegate backend is CoreML + self.assertEqual( + executorch_program.execution_plan[0].delegates[0].id, + "CoreMLBackend", + ) + + # Check number of instructions + instructions = executorch_program.execution_plan[0].chains[0].instructions + self.assertIsNotNone(instructions) + self.assertEqual(len(instructions), 1) + + et_runtime: Runtime = Runtime.get() + backend_registry = et_runtime.backend_registry + logging.info( + f"backends registered: {et_runtime.backend_registry.registered_backend_names}" + ) + if backend_registry.is_available("CoreMLBackend"): + et_output = session.run_method("forward", example_inputs[0]) + logging.info(f"et output {et_output}") diff --git a/extension/android/BUCK b/extension/android/BUCK index 191e6ce4714..b02003fdc34 100644 --- a/extension/android/BUCK +++ b/extension/android/BUCK @@ -10,6 +10,7 @@ non_fbcode_target(_kind = fb_android_library, "executorch_android/src/main/java/org/pytorch/executorch/DType.java", "executorch_android/src/main/java/org/pytorch/executorch/EValue.java", "executorch_android/src/main/java/org/pytorch/executorch/ExecuTorchRuntime.java", + "executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java", "executorch_android/src/main/java/org/pytorch/executorch/MethodMetadata.java", "executorch_android/src/main/java/org/pytorch/executorch/Module.java", "executorch_android/src/main/java/org/pytorch/executorch/Tensor.java", diff --git a/extension/android/CMakeLists.txt b/extension/android/CMakeLists.txt index feb4fdb6bb9..be6715f93d5 100644 --- a/extension/android/CMakeLists.txt +++ b/extension/android/CMakeLists.txt @@ -23,19 +23,19 @@ if(NOT ANDROID_PLATFORM) set(ANDROID_PLATFORM android-30) endif() -# We need to download fbjni library from maven, and use its "prefab" library -# and headers, and link executorch library against that fbjni library. -# We don't know which NDK is used to compile fbjni, and we need to link our -# executorch library to the version which Android APK links against for runtime -# to ensure the libc++ dependencies are consistent. -# WARNING # -# Users need to use the SAME fbjni version here and in app gradle dependency -# for runtime compatibility! +# We need to download fbjni library from maven, and use its "prefab" library and +# headers, and link executorch library against that fbjni library. We don't know +# which NDK is used to compile fbjni, and we need to link our executorch library +# to the version which Android APK links against for runtime to ensure the +# libc++ dependencies are consistent. WARNING # Users need to use the SAME fbjni +# version here and in app gradle dependency for runtime compatibility! if(NOT FBJNI_VERSION) set(FBJNI_VERSION 0.5.1) endif() -set(FBJNI_AAR_URL https://repo1.maven.org/maven2/com/facebook/fbjni/fbjni/${FBJNI_VERSION}/fbjni-${FBJNI_VERSION}.aar) +set(FBJNI_AAR_URL + https://repo1.maven.org/maven2/com/facebook/fbjni/fbjni/${FBJNI_VERSION}/fbjni-${FBJNI_VERSION}.aar +) set(FBJNI_DOWNLOAD_PATH ${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/fbjni.aar) if(NOT EXISTS "${FBJNI_DOWNLOAD_PATH}") @@ -43,25 +43,36 @@ if(NOT EXISTS "${FBJNI_DOWNLOAD_PATH}") endif() add_custom_command( - OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" - COMMAND unzip -o ${FBJNI_DOWNLOAD_PATH} -d ${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni + OUTPUT + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" + COMMAND unzip -o ${FBJNI_DOWNLOAD_PATH} -d + ${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni DEPENDS "${FBJNI_DOWNLOAD_PATH}" ) add_custom_target( fbjni_prefab - DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" + DEPENDS + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" ) add_library(fbjni SHARED IMPORTED) add_dependencies(fbjni fbjni_prefab) -set_target_properties(fbjni PROPERTIES - IMPORTED_LOCATION "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" +set_target_properties( + fbjni + PROPERTIES + IMPORTED_LOCATION + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/libs/android.${ANDROID_ABI}/libfbjni.so" ) executorch_target_link_options_shared_lib(executorch) -add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp) +add_library( + executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp + jni/jni_helper.cpp +) set(link_libraries) list( @@ -78,21 +89,14 @@ list( ) if(EXECUTORCH_ANDROID_PROFILING) - list( - APPEND - link_libraries - etdump - flatccrt + list(APPEND link_libraries etdump flatccrt) + target_compile_definitions( + executorch_jni PUBLIC EXECUTORCH_ANDROID_PROFILING=1 ) - target_compile_definitions(executorch_jni PUBLIC EXECUTORCH_ANDROID_PROFILING=1) endif() if(TARGET optimized_native_cpu_ops_lib) - list( - APPEND - link_libraries - optimized_native_cpu_ops_lib - ) + list(APPEND link_libraries optimized_native_cpu_ops_lib) executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib) else() list(APPEND link_libraries portable_ops_lib portable_kernels) @@ -110,7 +114,15 @@ endif() if(TARGET xnnpack_backend) executorch_target_link_options_shared_lib(xnnpack_backend) - list(APPEND link_libraries xnnpack_backend XNNPACK pthreadpool cpuinfo xnnpack-microkernels-prod) + list( + APPEND + link_libraries + xnnpack_backend + XNNPACK + pthreadpool + cpuinfo + xnnpack-microkernels-prod + ) if(TARGET kleidiai) list(APPEND link_libraries kleidiai) endif() @@ -149,7 +161,9 @@ endif() if(EXECUTORCH_BUILD_EXTENSION_TRAINING) target_sources(executorch_jni PRIVATE jni/jni_layer_training.cpp jni/log.cpp) list(APPEND link_libraries extension_training) - target_compile_definitions(executorch_jni PUBLIC EXECUTORCH_BUILD_EXTENSION_TRAINING=1) + target_compile_definitions( + executorch_jni PUBLIC EXECUTORCH_BUILD_EXTENSION_TRAINING=1 + ) endif() if(EXECUTORCH_BUILD_LLAMA_JNI) @@ -166,33 +180,70 @@ if(EXECUTORCH_BUILD_LLAMA_JNI) ${CMAKE_CURRENT_BINARY_DIR}/../../examples/models/llama/runner ) + target_sources( + executorch_jni + PRIVATE ${EXECUTORCH_ROOT}/extension/llm/runner/llm_runner_helper.cpp + ) + + target_include_directories( + executorch_jni PRIVATE ${EXECUTORCH_ROOT}/extension/llm/runner + ) + + if(QNN_SDK_ROOT) + target_sources( + executorch_jni + PRIVATE + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/runner.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/decoder_runner.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/token_generator.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/lhd_token_generator.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/rpc_mem.cpp + ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner/kv_manager.cpp + ) + + target_include_directories( + executorch_jni + PRIVATE ${EXECUTORCH_ROOT}/examples/qualcomm/oss_scripts/llama/runner + ) + target_compile_definitions(executorch_jni PRIVATE EXECUTORCH_BUILD_QNN=1) + endif() + if(NEURON_BUFFER_ALLOCATOR_LIB) - target_sources( - executorch_jni PRIVATE - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/mtk_llama_runner.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/MultiModelLoader.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/rotary_embedding.cpp - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/token_embedding.cpp + target_sources( + executorch_jni + PRIVATE + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/mtk_llama_runner.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/LlamaModelChunk.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/LlamaRuntime.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/ModelChunk.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/MultiModelLoader.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/mask_builder.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/rotary_embedding.cpp + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner/llm_helper/token_embedding.cpp ) target_include_directories( - executorch_jni PRIVATE - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/ - ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner + executorch_jni + PRIVATE ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/ + ${EXECUTORCH_ROOT}/examples/mediatek/executor_runner/llama_runner ) add_library(libneuron_buffer_allocator SHARED IMPORTED) - set_property(TARGET libneuron_buffer_allocator PROPERTY IMPORTED_LOCATION ${NEURON_BUFFER_ALLOCATOR_LIB}) + set_property( + TARGET libneuron_buffer_allocator PROPERTY IMPORTED_LOCATION + ${NEURON_BUFFER_ALLOCATOR_LIB} + ) list(APPEND link_libraries neuron_backend libneuron_buffer_allocator) - target_compile_definitions(executorch_jni PRIVATE EXECUTORCH_BUILD_MEDIATEK=1) + target_compile_definitions( + executorch_jni PRIVATE EXECUTORCH_BUILD_MEDIATEK=1 + ) endif() endif() target_include_directories( - executorch_jni PRIVATE ${_common_include_directories} - "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" + executorch_jni + PRIVATE + ${_common_include_directories} + "${CMAKE_CURRENT_BINARY_DIR}/third-party/fbjni/prefab/modules/fbjni/include/" ) target_compile_options(executorch_jni PUBLIC ${_common_compile_options}) diff --git a/extension/android/executorch_android/build.gradle b/extension/android/executorch_android/build.gradle index 330dfc83479..7d91cfd1194 100644 --- a/extension/android/executorch_android/build.gradle +++ b/extension/android/executorch_android/build.gradle @@ -6,12 +6,16 @@ * LICENSE file in the root directory of this source tree. */ + plugins { id "com.android.library" version "8.9.0" id "com.vanniktech.maven.publish" version "0.31.0" alias(libs.plugins.jetbrains.kotlin.android) } +def qnnVersion = System.properties['qnnVersion'] +def execuTorchVersion = System.properties['execuTorchVersion'] + android { namespace = "org.pytorch.executorch" compileSdk = 34 @@ -56,13 +60,16 @@ dependencies { androidTestImplementation 'commons-io:commons-io:2.4' androidTestImplementation 'org.json:json:20250107' androidTestImplementation 'org.jetbrains.kotlin:kotlin-test:1.9.23' + if (qnnVersion) { + implementation "com.qualcomm.qti:qnn-runtime:$qnnVersion" + } } mavenPublishing { publishToMavenCentral() signAllPublications() - coordinates("org.pytorch", "executorch-android", "0.7.0-SNAPSHOT") + coordinates("org.pytorch", "executorch-android" + (qnnVersion ? "-qnn" : ""), execuTorchVersion ? execuTorchVersion : "0.7.0-SNAPSHOT") pom { name = "ExecuTorch Android" diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java new file mode 100644 index 00000000000..e180efcdfca --- /dev/null +++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/ExecutorchRuntimeException.java @@ -0,0 +1,157 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +package org.pytorch.executorch; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +public class ExecutorchRuntimeException extends RuntimeException { + // Error code constants - keep in sync with runtime/core/error.h + // System errors + public static final int OK = 0x00; + public static final int INTERNAL = 0x01; + public static final int INVALID_STATE = 0x02; + public static final int END_OF_METHOD = 0x03; + + // Logical errors + public static final int NOT_SUPPORTED = 0x10; + public static final int NOT_IMPLEMENTED = 0x11; + public static final int INVALID_ARGUMENT = 0x12; + public static final int INVALID_TYPE = 0x13; + public static final int OPERATOR_MISSING = 0x14; + public static final int REGISTRATION_EXCEEDING_MAX_KERNELS = 0x15; + public static final int REGISTRATION_ALREADY_REGISTERED = 0x16; + + // Resource errors + public static final int NOT_FOUND = 0x20; + public static final int MEMORY_ALLOCATION_FAILED = 0x21; + public static final int ACCESS_FAILED = 0x22; + public static final int INVALID_PROGRAM = 0x23; + public static final int INVALID_EXTERNAL_DATA = 0x24; + public static final int OUT_OF_RESOURCES = 0x25; + + // Delegate errors + public static final int DELEGATE_INVALID_COMPATIBILITY = 0x30; + public static final int DELEGATE_MEMORY_ALLOCATION_FAILED = 0x31; + public static final int DELEGATE_INVALID_HANDLE = 0x32; + + private static final Map ERROR_CODE_MESSAGES; + + static { + Map map = new HashMap<>(); + + // System errors + map.put(OK, "Operation successful"); + map.put(INTERNAL, "Internal error"); + map.put(INVALID_STATE, "Invalid state"); + map.put(END_OF_METHOD, "End of method reached"); + // Logical errors + map.put(NOT_SUPPORTED, "Operation not supported"); + map.put(NOT_IMPLEMENTED, "Operation not implemented"); + map.put(INVALID_ARGUMENT, "Invalid argument"); + map.put(INVALID_TYPE, "Invalid type"); + map.put(OPERATOR_MISSING, "Operator missing"); + map.put(REGISTRATION_EXCEEDING_MAX_KERNELS, "Exceeded max kernels"); + map.put(REGISTRATION_ALREADY_REGISTERED, "Kernel already registered"); + // Resource errors + map.put(NOT_FOUND, "Resource not found"); + map.put(MEMORY_ALLOCATION_FAILED, "Memory allocation failed"); + map.put(ACCESS_FAILED, "Access failed"); + map.put(INVALID_PROGRAM, "Invalid program"); + map.put(INVALID_EXTERNAL_DATA, "Invalid external data"); + map.put(OUT_OF_RESOURCES, "Out of resources"); + // Delegate errors + map.put(DELEGATE_INVALID_COMPATIBILITY, "Delegate invalid compatibility"); + map.put(DELEGATE_MEMORY_ALLOCATION_FAILED, "Delegate memory allocation failed"); + map.put(DELEGATE_INVALID_HANDLE, "Delegate invalid handle"); + ERROR_CODE_MESSAGES = Collections.unmodifiableMap(map); + } + + static class ErrorHelper { + // Reusable StringBuilder instance + private static final StringBuilder sb = new StringBuilder(); + + static String formatMessage(int errorCode, String details) { + synchronized (sb) { + sb.setLength(0); // Clear the StringBuilder before use + + String baseMessage = ERROR_CODE_MESSAGES.get(errorCode); + if (baseMessage == null) { + baseMessage = "Unknown error code 0x" + Integer.toHexString(errorCode); + } + + sb.append("[Executorch Error 0x") + .append(Integer.toHexString(errorCode)) + .append("] ") + .append(baseMessage) + .append(": ") + .append(details) + .append("\nDetailed Logs:\n"); + + try { + String[] logEntries = readLogBuffer(); // JNI call + formatLogEntries(sb, logEntries); + } catch (Exception e) { + sb.append("Failed to retrieve detailed logs: ").append(e.getMessage()); + } + + return sb.toString(); + } + } + + // Native JNI method declaration + private static native String[] readLogBuffer(); + + // Append log entries to the provided StringBuilder + private static void formatLogEntries(StringBuilder sb, String[] logEntries) { + if (logEntries == null || logEntries.length == 0) { + sb.append("No detailed logs available."); + return; + } + for (String entry : logEntries) { + sb.append(entry).append("\n"); + } + } + } + + private final int errorCode; + + public ExecutorchRuntimeException(int errorCode, String details) { + super(ErrorHelper.formatMessage(errorCode, details)); + this.errorCode = errorCode; + } + + public int getErrorCode() { + return errorCode; + } + + // Idiomatic Java exception for invalid arguments. + public static class ExecutorchInvalidArgumentException extends IllegalArgumentException { + private final int errorCode = INVALID_ARGUMENT; + + public ExecutorchInvalidArgumentException(String details) { + super(ErrorHelper.formatMessage(INVALID_ARGUMENT, details)); + } + + public int getErrorCode() { + return errorCode; + } + } + + // Factory method to create an exception of the appropriate subclass. + public static RuntimeException makeExecutorchException(int errorCode, String details) { + switch (errorCode) { + case INVALID_ARGUMENT: + return new ExecutorchInvalidArgumentException(details); + default: + return new ExecutorchRuntimeException(errorCode, details); + } + } +} diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java index b014ceb75d8..289df5defd9 100644 --- a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java +++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java @@ -125,9 +125,7 @@ public int generate(String prompt, LlmCallback llmCallback, boolean echo) { * @param llmCallback callback object to receive results * @param echo indicate whether to echo the input prompt or not (text completion vs chat) */ - public int generate(String prompt, int seqLen, LlmCallback llmCallback, boolean echo) { - return generate(null, 0, 0, 0, prompt, seqLen, llmCallback, echo); - } + public native int generate(String prompt, int seqLen, LlmCallback llmCallback, boolean echo); /** * Start generating tokens from the module. @@ -154,8 +152,7 @@ public int generate(String prompt, LlmGenerationConfig config, LlmCallback llmCa * @param llmCallback callback object to receive results. * @param echo indicate whether to echo the input prompt or not (text completion vs chat) */ - @DoNotStrip - public native int generate( + public int generate( int[] image, int width, int height, @@ -163,7 +160,11 @@ public native int generate( String prompt, int seqLen, LlmCallback llmCallback, - boolean echo); + boolean echo) { + prefillPrompt(prompt); + prefillImages(image, width, height, channels); + return generate("", llmCallback, echo); + } /** * Prefill an LLaVA Module with the given images input. @@ -172,56 +173,47 @@ public native int generate( * @param width Input image width * @param height Input image height * @param channels Input image number of channels - * @param startPos The starting position in KV cache of the input in the LLM. - * @return The updated starting position in KV cache of the input in the LLM. + * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer + * exposed to user. * @throws RuntimeException if the prefill failed */ - public long prefillImages(int[] image, int width, int height, int channels, long startPos) { - long[] nativeResult = prefillImagesNative(image, width, height, channels, startPos); - if (nativeResult[0] != 0) { - throw new RuntimeException("Prefill failed with error code: " + nativeResult[0]); + @Deprecated + public long prefillImages(int[] image, int width, int height, int channels) { + int nativeResult = appendImagesInput(image, width, height, channels); + if (nativeResult != 0) { + throw new RuntimeException("Prefill failed with error code: " + nativeResult); } - return nativeResult[1]; + return 0; } - // returns a tuple of (status, updated startPos) - private native long[] prefillImagesNative( - int[] image, int width, int height, int channels, long startPos); + private native int appendImagesInput(int[] image, int width, int height, int channels); /** * Prefill an LLaVA Module with the given text input. * * @param prompt The text prompt to LLaVA. - * @param startPos The starting position in KV cache of the input in the LLM. It's passed as - * reference and will be updated inside this function. - * @param bos The number of BOS (begin of sequence) token. - * @param eos The number of EOS (end of sequence) token. - * @return The updated starting position in KV cache of the input in the LLM. + * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer + * exposed to user. * @throws RuntimeException if the prefill failed */ - public long prefillPrompt(String prompt, long startPos, int bos, int eos) { - long[] nativeResult = prefillPromptNative(prompt, startPos, bos, eos); - if (nativeResult[0] != 0) { - throw new RuntimeException("Prefill failed with error code: " + nativeResult[0]); + @Deprecated + public long prefillPrompt(String prompt) { + int nativeResult = appendTextInput(prompt); + if (nativeResult != 0) { + throw new RuntimeException("Prefill failed with error code: " + nativeResult); } - return nativeResult[1]; + return 0; } - // returns a tuple of (status, updated startPos) - private native long[] prefillPromptNative(String prompt, long startPos, int bos, int eos); + // returns status + private native int appendTextInput(String prompt); /** - * Generate tokens from the given prompt, starting from the given position. + * Reset the context of the LLM. This will clear the KV cache and reset the state of the LLM. * - * @param prompt The text prompt to LLaVA. - * @param seqLen The total sequence length, including the prompt tokens and new tokens. - * @param startPos The starting position in KV cache of the input in the LLM. - * @param callback callback object to receive results. - * @param echo indicate whether to echo the input prompt or not. - * @return The error code. + *

The startPos will be reset to 0. */ - public native int generateFromPos( - String prompt, int seqLen, long startPos, LlmCallback callback, boolean echo); + public native void resetContext(); /** Stop current generate() before it finishes. */ @DoNotStrip diff --git a/extension/android/jni/BUCK b/extension/android/jni/BUCK index 2a903da3e33..0ba39a71666 100644 --- a/extension/android/jni/BUCK +++ b/extension/android/jni/BUCK @@ -28,7 +28,7 @@ non_fbcode_target(_kind = executorch_generated_lib, non_fbcode_target(_kind = fb_android_cxx_library, name = "executorch_jni", - srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"], + srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_helper.cpp"], allow_jni_merging = False, compiler_flags = ET_JNI_COMPILER_FLAGS, soname = "libexecutorch.$(ext)", @@ -49,7 +49,7 @@ non_fbcode_target(_kind = fb_android_cxx_library, non_fbcode_target(_kind = fb_android_cxx_library, name = "executorch_jni_full", - srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp"], + srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_helper.cpp"], allow_jni_merging = False, compiler_flags = ET_JNI_COMPILER_FLAGS, soname = "libexecutorch.$(ext)", @@ -71,7 +71,7 @@ non_fbcode_target(_kind = fb_android_cxx_library, non_fbcode_target(_kind = fb_android_cxx_library, name = "executorch_training_jni", - srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp"], + srcs = ["jni_layer.cpp", "log.cpp", "jni_layer_runtime.cpp", "jni_layer_training.cpp", "jni_helper.cpp"], allow_jni_merging = False, compiler_flags = ET_JNI_COMPILER_FLAGS + [ "-DEXECUTORCH_BUILD_EXTENSION_TRAINING", @@ -98,11 +98,7 @@ non_fbcode_target(_kind = fb_android_cxx_library, non_fbcode_target(_kind = fb_android_cxx_library, name = "executorch_llama_jni", - srcs = [ - "jni_layer.cpp", - "jni_layer_llama.cpp", - "jni_layer_runtime.cpp", - ], + srcs = ["jni_layer.cpp", "jni_layer_llama.cpp", "jni_layer_runtime.cpp", "jni_helper.cpp"], allow_jni_merging = False, compiler_flags = ET_JNI_COMPILER_FLAGS + [ "-DEXECUTORCH_BUILD_LLAMA_JNI", @@ -145,9 +141,14 @@ runtime.export_file( name = "jni_layer_runtime.cpp", ) +runtime.export_file( + name = "jni_helper.cpp", +) + runtime.cxx_library( name = "jni_headers", exported_headers = [ "jni_layer_constants.h", + "jni_helper.h", ] ) diff --git a/extension/android/jni/jni_helper.cpp b/extension/android/jni/jni_helper.cpp new file mode 100644 index 00000000000..a8fb2aeddcf --- /dev/null +++ b/extension/android/jni/jni_helper.cpp @@ -0,0 +1,34 @@ +/* + * 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 "jni_helper.h" + +namespace executorch::jni_helper { + +void throwExecutorchException(uint32_t errorCode, const std::string& details) { + // Get the current JNI environment + auto env = facebook::jni::Environment::current(); + + // Find the Java ExecutorchRuntimeException class + static auto exceptionClass = facebook::jni::findClassLocal( + "org/pytorch/executorch/ExecutorchRuntimeException"); + + // Find the static factory method: makeExecutorchException(int, String) + static auto makeExceptionMethod = exceptionClass->getStaticMethod< + facebook::jni::local_ref( + int, facebook::jni::alias_ref)>( + "makeExecutorchException", + "(ILjava/lang/String;)Lorg/pytorch/executorch/ExecutorchRuntimeException;"); + + auto jDetails = facebook::jni::make_jstring(details); + // Call the factory method to create the exception object + auto exception = makeExceptionMethod(exceptionClass, errorCode, jDetails); + facebook::jni::throwNewJavaException(exception.get()); +} + +} // namespace executorch::jni_helper diff --git a/extension/android/jni/jni_helper.h b/extension/android/jni/jni_helper.h new file mode 100644 index 00000000000..996d75581d3 --- /dev/null +++ b/extension/android/jni/jni_helper.h @@ -0,0 +1,26 @@ +/* + * 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 executorch::jni_helper { + +/** + * Throws a Java ExecutorchRuntimeException corresponding to the given error + * code and details. Uses the Java factory method + * ExecutorchRuntimeException.makeExecutorchException(int, String). + * + * @param errorCode The error code from the C++ Executorch runtime. + * @param details Additional details to include in the exception message. + */ +void throwExecutorchException(uint32_t errorCode, const std::string& details); + +} // namespace executorch::jni_helper diff --git a/extension/android/jni/jni_layer.cpp b/extension/android/jni/jni_layer.cpp index 7111a0bc6bc..e7ef6e62c74 100644 --- a/extension/android/jni/jni_layer.cpp +++ b/extension/android/jni/jni_layer.cpp @@ -6,7 +6,9 @@ * LICENSE file in the root directory of this source tree. */ +#include #include + #include #include #include @@ -55,14 +57,14 @@ class TensorHybrid : public facebook::jni::HybridClass { // Java wrapper currently only supports contiguous tensors. const auto scalarType = tensor.scalar_type(); - + int jdtype = scalar_type_to_java_dtype.at(scalarType); if (scalar_type_to_java_dtype.count(scalarType) == 0) { - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "executorch::aten::Tensor scalar type %d is not supported on java side", - scalarType); + std::stringstream ss; + ss << "executorch::aten::Tensor scalar [java] type: " << jdtype + << " is not supported on java side"; + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str().c_str()); } - int jdtype = scalar_type_to_java_dtype.at(scalarType); const auto& tensor_shape = tensor.sizes(); std::vector tensor_shape_vec; @@ -124,19 +126,19 @@ class TensorHybrid : public facebook::jni::HybridClass { } JNIEnv* jni = facebook::jni::Environment::current(); if (java_dtype_to_scalar_type.count(jdtype) == 0) { - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "Unknown Tensor jdtype %d", - jdtype); + std::stringstream ss; + ss << "Unknown Tensor jdtype: [" << jdtype << "]"; + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str().c_str()); } ScalarType scalar_type = java_dtype_to_scalar_type.at(jdtype); const auto dataCapacity = jni->GetDirectBufferCapacity(jbuffer.get()); if (dataCapacity != numel) { - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "Tensor dimensions(elements number:%d inconsistent with buffer capacity(%d)", - numel, - dataCapacity); + std::stringstream ss; + ss << "Tensor dimensions(elements number: " << numel + << "inconsistent with buffer capacity " << dataCapacity << "]"; + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str().c_str()); } return from_blob( jni->GetDirectBufferAddress(jbuffer.get()), shape_vec, scalar_type); @@ -194,10 +196,11 @@ class JEValue : public facebook::jni::JavaClass { return jMethodTensor( JEValue::javaClassStatic(), facebook::jni::make_jstring(str)); } - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "Unsupported EValue type: %d", - evalue.tag); + std::stringstream ss; + ss << "Unknown EValue type: [" << static_cast(evalue.tag) << "]"; + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str().c_str()); + return {}; } static TensorPtr JEValueToTensorImpl( @@ -213,10 +216,11 @@ class JEValue : public facebook::jni::JavaClass { auto jtensor = jMethodGetTensor(JEValue); return TensorHybrid::newTensorFromJTensor(jtensor); } - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "Unknown EValue typeCode %d", - typeCode); + std::stringstream ss; + ss << "Unknown EValue typeCode: " << typeCode; + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str().c_str()); + return {}; } }; @@ -296,13 +300,24 @@ class ExecuTorchJni : public facebook::jni::HybridClass { jinputs) { // If no inputs is given, it will run with sample inputs (ones) if (jinputs->size() == 0) { - if (module_->load_method(method) != Error::Ok) { + auto result = module_->load_method(method); + if (result != Error::Ok) { + // Format hex string + std::stringstream ss; + ss << "Cannot get method names [Native Error: 0x" << std::hex + << std::uppercase << static_cast(result) << "]"; + + jni_helper::throwExecutorchException( + static_cast(result), ss.str()); return {}; } auto&& underlying_method = module_->methods_[method].method; auto&& buf = prepare_input_tensors(*underlying_method); - auto result = underlying_method->execute(); + result = underlying_method->execute(); if (result != Error::Ok) { + jni_helper::throwExecutorchException( + static_cast(result), + "Execution failed for method: " + method); return {}; } facebook::jni::local_ref> jresult = @@ -356,11 +371,9 @@ class ExecuTorchJni : public facebook::jni::HybridClass { #endif if (!result.ok()) { - facebook::jni::throwNewJavaException( - "java/lang/Exception", - "Execution of method %s failed with status 0x%" PRIx32, - method.c_str(), - static_cast(result.error())); + jni_helper::throwExecutorchException( + static_cast(result.error()), + "Execution failed for method: " + method); return {}; } @@ -438,9 +451,15 @@ class ExecuTorchJni : public facebook::jni::HybridClass { facebook::jni::local_ref> getMethods() { const auto& names_result = module_->method_names(); if (!names_result.ok()) { - facebook::jni::throwNewJavaException( - facebook::jni::gJavaLangIllegalArgumentException, - "Cannot get load module"); + // Format hex string + std::stringstream ss; + ss << "Cannot get load module [Native Error: 0x" << std::hex + << std::uppercase << static_cast(names_result.error()) + << "]"; + + jni_helper::throwExecutorchException( + static_cast(Error::InvalidArgument), ss.str()); + return {}; } const auto& methods = names_result.get(); facebook::jni::local_ref> ret = diff --git a/extension/android/jni/jni_layer_llama.cpp b/extension/android/jni/jni_layer_llama.cpp index 48bc62141a2..23686f01ee7 100644 --- a/extension/android/jni/jni_layer_llama.cpp +++ b/extension/android/jni/jni_layer_llama.cpp @@ -13,10 +13,12 @@ #include #include -#include -#include #include #include +#include +#include +#include +#include #include #include #include @@ -29,6 +31,10 @@ #include #include +#if defined(EXECUTORCH_BUILD_QNN) +#include +#endif + #if defined(EXECUTORCH_BUILD_MEDIATEK) #include #endif @@ -115,7 +121,9 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { float temperature_ = 0.0f; int model_type_category_; std::unique_ptr runner_; - std::unique_ptr multi_modal_runner_; + std::unique_ptr + multi_modal_runner_; + std::vector prefill_inputs_; public: constexpr static auto kJavaDescriptor = @@ -124,6 +132,7 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { constexpr static int MODEL_TYPE_CATEGORY_LLM = 1; constexpr static int MODEL_TYPE_CATEGORY_MULTIMODAL = 2; constexpr static int MODEL_TYPE_MEDIATEK_LLAMA = 3; + constexpr static int MODEL_TYPE_QNN_LLAMA = 4; static facebook::jni::local_ref initHybrid( facebook::jni::alias_ref, @@ -160,20 +169,33 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { model_type_category_ = model_type_category; if (model_type_category == MODEL_TYPE_CATEGORY_MULTIMODAL) { - multi_modal_runner_ = std::make_unique( + multi_modal_runner_ = llm::create_multimodal_runner( model_path->toStdString().c_str(), - tokenizer_path->toStdString().c_str(), - temperature); + llm::load_tokenizer(tokenizer_path->toStdString())); } else if (model_type_category == MODEL_TYPE_CATEGORY_LLM) { std::optional data_path_str = data_path ? std::optional{data_path->toStdString()} : std::nullopt; - // TODO(larryliu0820): Use the API in text_llm_runner.h to create the - // runner. - runner_ = example::create_llama_runner( + runner_ = executorch::extension::llm::create_text_llm_runner( model_path->toStdString(), - tokenizer_path->toStdString(), + llm::load_tokenizer(tokenizer_path->toStdString()), data_path_str); +#if defined(EXECUTORCH_BUILD_QNN) + } else if (model_type_category == MODEL_TYPE_QNN_LLAMA) { + std::unique_ptr module = std::make_unique< + executorch::extension::Module>( + model_path->toStdString().c_str(), + executorch::extension::Module::LoadMode::MmapUseMlockIgnoreErrors); + std::string decoder_model = "llama3"; // use llama3 for now + runner_ = std::make_unique>( // QNN runner + std::move(module), + decoder_model.c_str(), + model_path->toStdString().c_str(), + tokenizer_path->toStdString().c_str(), + data_path->toStdString().c_str(), + ""); + model_type_category_ = MODEL_TYPE_CATEGORY_LLM; +#endif #if defined(EXECUTORCH_BUILD_MEDIATEK) } else if (model_type_category == MODEL_TYPE_MEDIATEK_LLAMA) { runner_ = std::make_unique( @@ -186,34 +208,26 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { } jint generate( - facebook::jni::alias_ref image, - jint width, - jint height, - jint channels, facebook::jni::alias_ref prompt, jint seq_len, facebook::jni::alias_ref callback, jboolean echo) { if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) { - auto image_size = image->size(); - std::vector images; - if (image_size != 0) { - std::vector image_data_jint(image_size); - std::vector image_data(image_size); - image->getRegion(0, image_size, image_data_jint.data()); - for (int i = 0; i < image_size; i++) { - image_data[i] = image_data_jint[i]; - } - llm::Image image_runner{image_data, width, height, channels}; - images.push_back(image_runner); + std::vector inputs = prefill_inputs_; + prefill_inputs_.clear(); + if (!prompt->toStdString().empty()) { + inputs.emplace_back(llm::MultimodalInput{prompt->toStdString()}); } + executorch::extension::llm::GenerationConfig config{ + .echo = static_cast(echo), + .seq_len = seq_len, + .temperature = temperature_, + }; multi_modal_runner_->generate( - std::move(images), - prompt->toStdString(), - seq_len, - [callback](std::string result) { callback->onResult(result); }, - [callback](const llm::Stats& result) { callback->onStats(result); }, - echo); + std::move(inputs), + config, + [callback](const std::string& result) { callback->onResult(result); }, + [callback](const llm::Stats& result) { callback->onStats(result); }); } else if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) { executorch::extension::llm::GenerationConfig config{ .echo = static_cast(echo), @@ -229,50 +243,24 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { return 0; } - // Returns a tuple of (error, start_pos) + // Returns status_code // Contract is valid within an AAR (JNI + corresponding Java code) - // If the first element is not Error::Ok, the other element is undefined. - facebook::jni::local_ref prefill_prompt( - facebook::jni::alias_ref prompt, - jlong start_pos, - jint bos, - jint eos) { - facebook::jni::local_ref tuple_result = - facebook::jni::make_long_array(2); - if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) { - tuple_result->pin()[0] = static_cast(Error::NotSupported); - return tuple_result; - } - - auto&& result = multi_modal_runner_->prefill_prompt( - prompt->toStdString(), start_pos, bos, eos); - tuple_result->pin()[0] = static_cast(Error::Ok); - if (result.ok()) { - tuple_result->pin()[1] = static_cast(start_pos); - } - return tuple_result; + jint append_text_input(facebook::jni::alias_ref prompt) { + prefill_inputs_.emplace_back(llm::MultimodalInput{prompt->toStdString()}); + return 0; } - // Returns a tuple of (error, start_pos) - // Contract is valid within an AAR (JNI + corresponding Java code) - // If the first element is not Error::Ok, the other element is undefined. - - facebook::jni::local_ref prefill_images( + // Returns status_code + jint append_images_input( facebook::jni::alias_ref image, jint width, jint height, - jint channels, - jlong start_pos) { - facebook::jni::local_ref tuple_result = - facebook::jni::make_long_array(2); - - if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) { - tuple_result->pin()[0] = static_cast(Error::NotSupported); - return tuple_result; + jint channels) { + std::vector images; + if (image == nullptr) { + return static_cast(Error::EndOfMethod); } - auto image_size = image->size(); - std::vector images; if (image_size != 0) { std::vector image_data_jint(image_size); std::vector image_data(image_size); @@ -281,43 +269,11 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { image_data[i] = image_data_jint[i]; } llm::Image image_runner{image_data, width, height, channels}; - images.push_back(image_runner); + prefill_inputs_.emplace_back( + llm::MultimodalInput{std::move(image_runner)}); } - // TODO(hsz): make start_pos a reference and update it here - jint result = static_cast( - multi_modal_runner_->prefill_images(images, start_pos)); - tuple_result->pin()[0] = result; - tuple_result->pin()[1] = static_cast(start_pos); - return tuple_result; - } - jint generate_from_pos( - facebook::jni::alias_ref prompt, - jint seq_len, - jlong start_pos, - facebook::jni::alias_ref callback, - jboolean echo) { - if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) { - return static_cast(multi_modal_runner_->generate_from_pos( - prompt->toStdString(), - seq_len, - start_pos, - [callback](const std::string& result) { callback->onResult(result); }, - [callback](const llm::Stats& stats) { callback->onStats(stats); }, - echo)); - } else if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) { - executorch::extension::llm::GenerationConfig config{ - .echo = static_cast(echo), - .seq_len = seq_len, - .temperature = temperature_, - }; - return static_cast(runner_->generate_from_pos( - prompt->toStdString(), - start_pos, - config, - [callback](std::string result) { callback->onResult(result); }, - [callback](const llm::Stats& stats) { callback->onStats(stats); })); - } + return 0; } void stop() { @@ -328,6 +284,15 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { } } + void reset_context() { + if (runner_ != nullptr) { + runner_->reset(); + } + if (multi_modal_runner_ != nullptr) { + multi_modal_runner_->reset(); + } + } + jint load() { if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) { return static_cast(multi_modal_runner_->load()); @@ -344,11 +309,10 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { makeNativeMethod("stop", ExecuTorchLlmJni::stop), makeNativeMethod("load", ExecuTorchLlmJni::load), makeNativeMethod( - "prefillImagesNative", ExecuTorchLlmJni::prefill_images), - makeNativeMethod( - "prefillPromptNative", ExecuTorchLlmJni::prefill_prompt), + "appendImagesInput", ExecuTorchLlmJni::append_images_input), makeNativeMethod( - "generateFromPos", ExecuTorchLlmJni::generate_from_pos), + "appendTextInput", ExecuTorchLlmJni::append_text_input), + makeNativeMethod("resetContext", ExecuTorchLlmJni::reset_context), }); } }; diff --git a/extension/android/jni/selective_jni.buck.bzl b/extension/android/jni/selective_jni.buck.bzl index d557606b7d1..8e20f903ca9 100644 --- a/extension/android/jni/selective_jni.buck.bzl +++ b/extension/android/jni/selective_jni.buck.bzl @@ -10,6 +10,7 @@ def selective_jni_target(name, deps, srcs = [], soname = "libexecutorch.$(ext)") srcs = [ "//xplat/executorch/extension/android/jni:jni_layer.cpp", "//xplat/executorch/extension/android/jni:jni_layer_runtime.cpp", + "//xplat/executorch/extension/android/jni:jni_helper.cpp", ] + srcs, allow_jni_merging = False, compiler_flags = ET_JNI_COMPILER_FLAGS, diff --git a/extension/apple/CMakeLists.txt b/extension/apple/CMakeLists.txt index 0e978073aa2..180c13777be 100644 --- a/extension/apple/CMakeLists.txt +++ b/extension/apple/CMakeLists.txt @@ -20,36 +20,28 @@ endif() add_library(extension_apple) -file(GLOB OBJC_SOURCES - ExecuTorch/Exported/*.m - ExecuTorch/Exported/*.mm - ExecuTorch/Internal/*.m - ExecuTorch/Internal/*.mm +file(GLOB OBJC_SOURCES ExecuTorch/Exported/*.m ExecuTorch/Exported/*.mm + ExecuTorch/Internal/*.m ExecuTorch/Internal/*.mm ) -file(GLOB SWIFT_SOURCES - ExecuTorch/Exported/*.swift -) +file(GLOB SWIFT_SOURCES ExecuTorch/Exported/*.swift) -target_sources(extension_apple PRIVATE - ${OBJC_SOURCES} - ${SWIFT_SOURCES} -) +target_sources(extension_apple PRIVATE ${OBJC_SOURCES} ${SWIFT_SOURCES}) -target_include_directories(extension_apple +target_include_directories( + extension_apple PUBLIC ExecuTorch/Exported PRIVATE ExecuTorch/Internal ) find_library(FOUNDATION_FRAMEWORK Foundation) -target_link_libraries(extension_apple - PRIVATE executorch ${FOUNDATION_FRAMEWORK} +target_link_libraries( + extension_apple PRIVATE executorch ${FOUNDATION_FRAMEWORK} ) -set_source_files_properties(${OBJC_SOURCES} PROPERTIES COMPILE_FLAGS - "-fobjc-arc" - "-fno-exceptions" - "-fno-rtti" +set_source_files_properties( + ${OBJC_SOURCES} PROPERTIES COMPILE_FLAGS "-fobjc-arc" "-fno-exceptions" + "-fno-rtti" ) set(MODULE_MAP_DIR ${CMAKE_CURRENT_BINARY_DIR}/module) @@ -57,30 +49,36 @@ set(MODULE_MAP_FILE ${MODULE_MAP_DIR}/module.modulemap) configure_file( "${CMAKE_CURRENT_SOURCE_DIR}/ExecuTorch/Exported/ExecuTorch.h" - "${MODULE_MAP_DIR}/ExecuTorch.h" - COPYONLY + "${MODULE_MAP_DIR}/ExecuTorch.h" COPYONLY ) file(MAKE_DIRECTORY ${MODULE_MAP_DIR}) -file(WRITE ${MODULE_MAP_FILE} -"module ExecuTorch { +file( + WRITE ${MODULE_MAP_FILE} + "module ExecuTorch { umbrella header \"ExecuTorch.h\" export * } -") +" +) -set(SWIFT_CLANG_INTEROP_FLAGS "-Xcc -fmodule-map-file=${MODULE_MAP_FILE} -I ${MODULE_MAP_DIR}") +set(SWIFT_CLANG_INTEROP_FLAGS + "-Xcc -fmodule-map-file=${MODULE_MAP_FILE} -I ${MODULE_MAP_DIR}" +) set(SWIFT_REMAP_FLAGS "-debug-prefix-map ${PROJECT_SOURCE_DIR}=/executorch") -set_target_properties(extension_apple PROPERTIES - Swift_MODULE_NAME "ExecuTorch" - Swift_FLAGS "${SWIFT_CLANG_INTEROP_FLAGS} ${SWIFT_REMAP_FLAGS}" - XCODE_ATTRIBUTE_SWIFT_MODULE_NAME "ExecuTorch" - XCODE_ATTRIBUTE_BUILD_LIBRARY_FOR_DISTRIBUTION "YES" - XCODE_ATTRIBUTE_OTHER_SWIFT_FLAGS "${SWIFT_CLANG_INTEROP_FLAGS} ${SWIFT_REMAP_FLAGS}" +set_target_properties( + extension_apple + PROPERTIES Swift_MODULE_NAME "ExecuTorch" + Swift_FLAGS "${SWIFT_CLANG_INTEROP_FLAGS} ${SWIFT_REMAP_FLAGS}" + XCODE_ATTRIBUTE_SWIFT_MODULE_NAME "ExecuTorch" + XCODE_ATTRIBUTE_BUILD_LIBRARY_FOR_DISTRIBUTION "YES" + XCODE_ATTRIBUTE_OTHER_SWIFT_FLAGS + "${SWIFT_CLANG_INTEROP_FLAGS} ${SWIFT_REMAP_FLAGS}" ) add_custom_command( - TARGET extension_apple POST_BUILD + TARGET extension_apple + POST_BUILD COMMAND ${CMAKE_COMMAND} -E rm -rf ${MODULE_MAP_DIR} ) diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift b/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift index 599a990b64c..11b20000ee1 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift +++ b/extension/apple/ExecuTorch/Exported/ExecuTorch+Module.swift @@ -176,3 +176,123 @@ public extension Module { try execute("forward") } } + +@available(*, deprecated, message: "This API is experimental.") +public extension Module { + /// Sets a single input value for a method at the specified index. + /// + /// - Parameters: + /// - value: The input as a `ValueConvertible`. + /// - method: The method name. + /// - index: Zero-based input index. + /// - Throws: If setting the input fails. + func setInput(_ value: ValueConvertible, for method: String, at index: Int) throws { + try __setInput(value.asValue(), forMethod: method, at: index) + } + + /// Sets a single input value for a method at index 0. + /// + /// - Parameters: + /// - value: The input as a `ValueConvertible`. + /// - method: The method name. + /// - Throws: If setting the input fails. + func setInput(_ value: ValueConvertible, for method: String) throws { + try setInput(value, for: method, at: 0) + } + + /// Sets a single input value for the "forward" method at the specified index. + /// + /// - Parameters: + /// - value: The input as a `ValueConvertible`. + /// - index: Zero-based input index. + /// - Throws: If setting the input fails. + func setInput(_ value: ValueConvertible, at index: Int) throws { + try setInput(value, for: "forward", at: index) + } + + /// Sets the first input value (index 0) for the "forward" method. + /// + /// - Parameter value: The input as a `ValueConvertible`. + /// - Throws: If setting the input fails. + func setInput(_ value: ValueConvertible) throws { + try setInput(value, for: "forward", at: 0) + } + + /// Sets all input values for a method. + /// + /// - Parameters: + /// - values: The inputs as an array of `ValueConvertible`. + /// - method: The method name. + /// - Throws: If setting the inputs fails. + func setInputs(_ values: [ValueConvertible], for method: String) throws { + try __setInputs(values.map { $0.asValue() }, forMethod: method) + } + + /// Sets all input values for the "forward" method. + /// + /// - Parameter values: The inputs as an array of `ValueConvertible`. + /// - Throws: If setting the inputs fails. + func setInputs(_ values: [ValueConvertible]) throws { + try setInputs(values, for: "forward") + } + + /// Sets all input values for a method using variadic arguments. + /// + /// - Parameters: + /// - values: The inputs as a variadic list of `ValueConvertible`. + /// - method: The method name. + /// - Throws: If setting the inputs fails. + func setInputs(_ values: ValueConvertible..., for method: String) throws { + try setInputs(values, for: method) + } + + /// Sets all input values for the "forward" method using variadic arguments. + /// + /// - Parameter values: The inputs as a variadic list of `ValueConvertible`. + /// - Throws: If setting the inputs fails. + func setInputs(_ values: ValueConvertible...) throws { + try setInputs(values, for: "forward") + } + + /// Sets the output location for a method at the specified index. + /// + /// Only tensor outputs are supported. The provided value must wrap a tensor + /// with compatible shape and data type for the method’s output slot. + /// + /// - Parameters: + /// - value: The output buffer as a `ValueConvertible` (tensor). + /// - method: The method name. + /// - index: Zero-based output index. + /// - Throws: If setting the output fails. + func setOutput(_ value: ValueConvertible, for method: String, at index: Int) throws { + try __setOutput(value.asValue(), forMethod: method, at: index) + } + + /// Sets the output location for a method at index 0. + /// + /// - Parameters: + /// - value: The output buffer as a `ValueConvertible` (tensor). + /// - method: The method name. + /// - Throws: If setting the output fails. + func setOutput(_ value: ValueConvertible, for method: String) throws { + try setOutput(value, for: method, at: 0) + } + + /// Sets the output location for the "forward" method at the specified index. + /// + /// - Parameters: + /// - value: The output buffer as a `ValueConvertible` (tensor). + /// - index: Zero-based output index. + /// - Throws: If setting the output fails. + func setOutput(_ value: ValueConvertible, at index: Int) throws { + try setOutput(value, for: "forward", at: index) + } + + /// Sets the first output location (index 0) for the "forward" method. + /// + /// - Parameter value: The output buffer as a `ValueConvertible` (tensor). + /// - Throws: If setting the output fails. + func setOutput(_ value: ValueConvertible) throws { + try setOutput(value, for: "forward", at: 0) + } +} diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h index 0eafcca8cc7..51d99780da3 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.h @@ -124,16 +124,41 @@ __attribute__((deprecated("This API is experimental."))) @interface ExecuTorchModule : NSObject /** - * Initializes a module with a file path and a specified load mode. + * Initializes a module with a file path, data path and a specified load mode. * * @param filePath A string representing the path to the ExecuTorch program file. - * @param loadMode A value from ExecuTorchModuleLoadMode that determines the file loading behavior. + * @param dataFilePath A string representing the path to a .ptd file with + * external tensors and external data. + * @param loadMode A value from ExecuTorchModuleLoadMode that determines the + * file loading behavior. * @return An initialized ExecuTorchModule instance. */ - (instancetype)initWithFilePath:(NSString *)filePath + dataFilePath:(NSString *)dataFilePath loadMode:(ExecuTorchModuleLoadMode)loadMode NS_DESIGNATED_INITIALIZER; +/** + * Initializes a module with a file path, data path and a specified load mode. + * + * @param filePath A string representing the path to the ExecuTorch program file. + * @param dataFilePath A string representing the path to a .ptd file with + * external tensors and external data. + * @return An initialized ExecuTorchModule instance. + */ +- (instancetype)initWithFilePath:(NSString *)filePath + dataFilePath:(NSString *)dataFilePath; + +/** + * Initializes a module with a file path and a specified load mode. + * + * @param filePath A string representing the path to the ExecuTorch program file. + * @param loadMode A value from ExecuTorchModuleLoadMode that determines the file loading behavior. + * @return An initialized ExecuTorchModule instance. + */ +- (instancetype)initWithFilePath:(NSString *)filePath + loadMode:(ExecuTorchModuleLoadMode)loadMode; + /** * Initializes a module with a file path using the default load mode (File mode). * @@ -187,6 +212,14 @@ __attribute__((deprecated("This API is experimental."))) */ - (BOOL)isMethodLoaded:(NSString *)methodName NS_SWIFT_NAME(isLoaded(_:)); +/** + * Unloads a method and releases its native resources and planned buffers. + * + * @param methodName The method to unload. + * @return YES if the method was unloaded; NO if it was not loaded at all. + */ +- (BOOL)unloadMethod:(NSString *)methodName NS_SWIFT_NAME(unload(_:)); + /** * Retrieves the set of method names available in the loaded program. * @@ -358,6 +391,145 @@ __attribute__((deprecated("This API is experimental."))) NS_SWIFT_UNAVAILABLE("") NS_RETURNS_RETAINED; +/** + * Sets a single input value for the "forward" method at index 0. + * + * @param value The input value. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInput:(ExecuTorchValue *)value + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets a single input value for the "forward" method at the specified index. + * + * @param value The input value. + * @param index Zero-based input index. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInput:(ExecuTorchValue *)value + atIndex:(NSInteger)index + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets a single input value for the specified method at index 0. + * + * @param value The input value. + * @param methodName The method name. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets a single input value for the specified method at the given index. + * + * The module retains the provided value to keep its backing storage alive + * until the value is overwritten or the module is deallocated. + * + * @param value The input value. + * @param methodName The method name. + * @param index Zero-based input index. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + atIndex:(NSInteger)index + error:(NSError **)error NS_REFINED_FOR_SWIFT; + +/** + * Sets all input values for the "forward" method. + * + * The number and types of values must match the method’s declared inputs. + * + * @param values The input values, one per declared input. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInputs:(NSArray *)values + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets all input values for the specified method. + * + * The module retains the provided values to keep their backing storage alive + * until the values are overwritten or the module is deallocated. + * + * @param values The input values, one per declared input. + * @param methodName The method name. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setInputs:(NSArray *)values + forMethod:(NSString *)methodName + error:(NSError **)error NS_REFINED_FOR_SWIFT; + +/** + * Sets the output buffer for the "forward" method at index 0. + * + * Only tensor outputs are supported. The provided value must wrap a tensor + * compatible with the method’s output slot. + * + * @param value The output buffer (must wrap a tensor). + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setOutput:(ExecuTorchValue *)value + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets the output buffer for the "forward" method at the specified index. + * + * Only tensor outputs are supported. The provided value must wrap a tensor + * compatible with the method’s output slot. + * + * @param value The output buffer (must wrap a tensor). + * @param index Zero-based output index. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setOutput:(ExecuTorchValue *)value + atIndex:(NSInteger)index + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets the output buffer for the specified method at index 0. + * + * Only tensor outputs are supported. The provided value must wrap a tensor + * compatible with the method’s output slot. + * + * @param value The output buffer (must wrap a tensor). + * @param methodName The method name. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setOutput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + error:(NSError **)error NS_SWIFT_UNAVAILABLE(""); + +/** + * Sets the output buffer for the specified method at the given index. + * + * The module retains the provided value to keep its backing storage alive + * until the value is overwritten or the module is deallocated. + * Only tensor outputs are supported. + * + * @param value The output buffer (must wrap a tensor). + * @param methodName The method name. + * @param index Zero-based output index. + * @param error On failure, set to an NSError describing the issue. + * @return YES on success; NO otherwise. + */ +- (BOOL)setOutput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + atIndex:(NSInteger)index + error:(NSError **)error NS_REFINED_FOR_SWIFT; + + (instancetype)new NS_UNAVAILABLE; - (instancetype)init NS_UNAVAILABLE; diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm index 7b0b15c00d0..ce58f2fb21a 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchModule.mm @@ -245,22 +245,43 @@ - (nullable instancetype)initWithMethodMetadata:(const MethodMeta &)methodMeta @implementation ExecuTorchModule { std::unique_ptr _module; + NSMutableDictionary *> *_inputs; + NSMutableDictionary *> *_outputs; } - (instancetype)initWithFilePath:(NSString *)filePath + dataFilePath:(NSString *)dataFilePath loadMode:(ExecuTorchModuleLoadMode)loadMode { self = [super init]; if (self) { _module = std::make_unique( filePath.UTF8String, + dataFilePath.UTF8String, static_cast(loadMode) ); + _inputs = [NSMutableDictionary new]; + _outputs = [NSMutableDictionary new]; } return self; } +- (instancetype)initWithFilePath:(NSString *)filePath + dataFilePath:(NSString *)dataFilePath { + return [self initWithFilePath:filePath + dataFilePath:dataFilePath + loadMode:ExecuTorchModuleLoadModeFile]; +} + +- (instancetype)initWithFilePath:(NSString *)filePath + loadMode:(ExecuTorchModuleLoadMode)loadMode { + return [self initWithFilePath:filePath + dataFilePath:@"" + loadMode:loadMode]; +} - (instancetype)initWithFilePath:(NSString *)filePath { - return [self initWithFilePath:filePath loadMode:ExecuTorchModuleLoadModeFile]; + return [self initWithFilePath:filePath + dataFilePath:@"" + loadMode:ExecuTorchModuleLoadModeFile]; } - (BOOL)loadWithVerification:(ExecuTorchVerification)verification @@ -300,6 +321,13 @@ - (BOOL)isMethodLoaded:(NSString *)methodName { return _module->is_method_loaded(methodName.UTF8String); } +- (BOOL)unloadMethod:(NSString *)methodName { + const auto didUnload = _module->unload_method(methodName.UTF8String); + [_inputs removeObjectForKey:methodName]; + [_outputs removeObjectForKey:methodName]; + return didUnload; +} + - (nullable NSSet *)methodNames:(NSError **)error { const auto result = _module->method_names(); if (!result.ok()) { @@ -331,12 +359,21 @@ - (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName - (nullable NSArray *)executeMethod:(NSString *)methodName withInputs:(NSArray *)values error:(NSError **)error { - std::vector inputs; - inputs.reserve(values.count); - for (ExecuTorchValue *value in values) { - inputs.push_back(toEValue(value)); + const char *methodNameString = methodName.UTF8String; + __block auto errorCode = Error::Ok; + [values enumerateObjectsUsingBlock:^(ExecuTorchValue *value, NSUInteger index, BOOL *stop) { + errorCode = _module->set_input(methodNameString, toEValue(value), index); + if (errorCode != Error::Ok) { + *stop = YES; + } + }]; + if (errorCode != Error::Ok) { + if (error) { + *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)errorCode); + } + return nil; } - const auto result = _module->execute(methodName.UTF8String, inputs); + const auto result = _module->execute(methodNameString); if (!result.ok()) { if (error) { *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)result.error()); @@ -423,4 +460,142 @@ - (nullable ExecuTorchMethodMetadata *)methodMetadata:(NSString *)methodName error:error]; } +- (BOOL)setInput:(ExecuTorchValue *)value + error:(NSError **)error NS_SWIFT_NAME(setInput(_:)) { + return [self setInput:value + forMethod:@"forward" + atIndex:0 + error:error]; +} + +- (BOOL)setInput:(ExecuTorchValue *)value + atIndex:(NSInteger)index + error:(NSError **)error { + return [self setInput:value + forMethod:@"forward" + atIndex:index + error:error]; +} + +- (BOOL)setInput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + error:(NSError **)error { + return [self setInput:value + forMethod:methodName + atIndex:0 + error:error]; +} + +- (BOOL)setInput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + atIndex:(NSInteger)index + error:(NSError **)error { + const auto errorCode = _module->set_input(methodName.UTF8String, toEValue(value), index); + if (errorCode != Error::Ok) { + if (error) { + *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)errorCode); + } + return NO; + } + // Cache inputs to keep them alive since ExecuTorchValue owns the actual data. + NSMutableArray *inputs = _inputs[methodName]; + if (!inputs) { + inputs = [NSMutableArray new]; + _inputs[methodName] = inputs; + } + if (index >= inputs.count) { + id placeholder = NSNull.null; + while (inputs.count < index) { + [inputs addObject:placeholder]; + } + [inputs addObject:value]; + } else { + inputs[index] = value; + } + return YES; +} + +- (BOOL)setInputs:(NSArray *)values + error:(NSError **)error { + return [self setInputs:values + forMethod:@"forward" + error:error]; +} + +- (BOOL)setInputs:(NSArray *)values + forMethod:(NSString *)methodName + error:(NSError **)error { + std::vector inputs; + inputs.reserve(values.count); + for (ExecuTorchValue *value in values) { + inputs.push_back(toEValue(value)); + } + const auto errorCode = _module->set_inputs(methodName.UTF8String, inputs); + if (errorCode != Error::Ok) { + if (error) { + *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)errorCode); + } + return NO; + } + // Cache inputs to keep them alive since ExecuTorchValue owns the actual data. + _inputs[methodName] = [values mutableCopy]; + + return YES; +} + +- (BOOL)setOutput:(ExecuTorchValue *)value + error:(NSError **)error { + return [self setOutput:value + forMethod:@"forward" + atIndex:0 + error:error]; +} + +- (BOOL)setOutput:(ExecuTorchValue *)value + atIndex:(NSInteger)index + error:(NSError **)error { + return [self setOutput:value + forMethod:@"forward" + atIndex:index + error:error]; +} + +- (BOOL)setOutput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + error:(NSError **)error { + return [self setOutput:value + forMethod:methodName + atIndex:0 + error:error]; +} + +- (BOOL)setOutput:(ExecuTorchValue *)value + forMethod:(NSString *)methodName + atIndex:(NSInteger)index + error:(NSError **)error { + const auto errorCode = _module->set_output(methodName.UTF8String, toEValue(value), index); + if (errorCode != Error::Ok) { + if (error) { + *error = ExecuTorchErrorWithCode((ExecuTorchErrorCode)errorCode); + } + return NO; + } + // Cache outputs to keep them alive since ExecuTorchValue owns the actual data. + NSMutableArray *outputs = _outputs[methodName]; + if (!outputs) { + outputs = [NSMutableArray new]; + _outputs[methodName] = outputs; + } + if (index >= outputs.count) { + id placeholder = NSNull.null; + while (outputs.count < index) { + [outputs addObject:placeholder]; + } + [outputs addObject:value]; + } else { + outputs[index] = value; + } + return YES; +} + @end diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm index 3cf06207b45..3a2b640b7d7 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm @@ -265,9 +265,15 @@ - (NSString *)description { auto const count = _tensor->numel(); os << "\n count: " << count << ","; os << "\n scalars: ["; + // Create a minimal context for error handling in ET_SWITCH + struct { + [[noreturn]] void fail(torch::executor::Error /* error */) { + ET_CHECK_MSG(false, "Unsupported dtype in description"); + } + } ctx; ET_SWITCH_REALHBBF16_TYPES( static_cast(_tensor->scalar_type()), - nullptr, + ctx, "description", CTYPE, [&] { @@ -488,9 +494,15 @@ - (instancetype)initWithScalars:(NSArray *)scalars "Number of scalars does not match the shape"); std::vector data; data.resize(count * ExecuTorchSizeOfDataType(dataType)); + // Create a minimal context for error handling in ET_SWITCH + struct { + [[noreturn]] void fail(torch::executor::Error /* error */) { + ET_CHECK_MSG(false, "Unsupported dtype in initWithScalars"); + } + } ctx; for (NSUInteger index = 0; index < count; ++index) { ET_SWITCH_REALHBBF16_AND_UINT_TYPES( - static_cast(dataType), nil, "initWithScalars", CTYPE, [&] { + static_cast(dataType), ctx, "initWithScalars", CTYPE, [&] { reinterpret_cast(data.data())[index] = utils::toType(scalars[index]); } ); @@ -801,8 +813,14 @@ + (instancetype)fullTensorWithShape:(NSArray *)shape dataType:(ExecuTorchDataType)dataType shapeDynamism:(ExecuTorchShapeDynamism)shapeDynamism { Scalar fillValue; + // Create a minimal context for error handling in ET_SWITCH + struct { + [[noreturn]] void fail(torch::executor::Error /* error */) { + ET_CHECK_MSG(false, "Unsupported dtype in fullTensor"); + } + } ctx; ET_SWITCH_REALHBBF16_AND_UINT_TYPES( - static_cast(dataType), nil, "fullTensor", CTYPE, [&] { + static_cast(dataType), ctx, "fullTensor", CTYPE, [&] { fillValue = utils::toType(scalar); } ); diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchValue.h b/extension/apple/ExecuTorch/Exported/ExecuTorchValue.h index 31fb1b96cbf..6070d25383c 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchValue.h +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchValue.h @@ -22,12 +22,6 @@ typedef NS_ENUM(uint32_t, ExecuTorchValueTag) { ExecuTorchValueTagDouble, ExecuTorchValueTagInteger, ExecuTorchValueTagBoolean, - ExecuTorchValueTagBooleanList, - ExecuTorchValueTagDoubleList, - ExecuTorchValueTagIntegerList, - ExecuTorchValueTagTensorList, - ExecuTorchValueTagScalarList, - ExecuTorchValueTagOptionalTensorList, } NS_SWIFT_NAME(ValueTag); typedef NSNumber *ExecuTorchScalarValue @@ -107,7 +101,7 @@ __attribute__((objc_subclassing_restricted)) * * @return An float representing the float value. */ - @property(nonatomic, readonly) ExecuTorchFloatValue floatValue NS_SWIFT_NAME(float); +@property(nonatomic, readonly) ExecuTorchFloatValue floatValue NS_SWIFT_NAME(float); /** * Returns YES if the value is of type None. @@ -166,7 +160,7 @@ __attribute__((objc_subclassing_restricted)) * * @return A BOOL indicating whether the value is a float. */ - @property(nonatomic, readonly) BOOL isFloat; +@property(nonatomic, readonly) BOOL isFloat; /** * Creates an instance encapsulating a Tensor. diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchValue.mm b/extension/apple/ExecuTorch/Exported/ExecuTorchValue.mm index 6ba03dc50f9..293cce03819 100644 --- a/extension/apple/ExecuTorch/Exported/ExecuTorchValue.mm +++ b/extension/apple/ExecuTorch/Exported/ExecuTorchValue.mm @@ -43,18 +43,6 @@ static inline ExecuTorchValueTag deduceValueTag(NSNumber *number) { return @"integer"; case ExecuTorchValueTagBoolean: return @"boolean"; - case ExecuTorchValueTagBooleanList: - return @"boolean_list"; - case ExecuTorchValueTagDoubleList: - return @"double_list"; - case ExecuTorchValueTagIntegerList: - return @"integer_list"; - case ExecuTorchValueTagTensorList: - return @"tensor_list"; - case ExecuTorchValueTagScalarList: - return @"scalar_list"; - case ExecuTorchValueTagOptionalTensorList: - return @"optional_tensor_list"; } return @"undefined"; } @@ -233,7 +221,7 @@ - (NSString *)description { [string appendString:@"\n value: "]; if (_value) { NSString *valueDescription = [_value description]; - [string appendString:[_value description]]; + [string appendString:valueDescription]; [string replaceOccurrencesOfString:@"\n" withString:@"\n " options:0 diff --git a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift index a35247f9bce..1cc4a31c4a3 100644 --- a/extension/apple/ExecuTorch/__tests__/ModuleTest.swift +++ b/extension/apple/ExecuTorch/__tests__/ModuleTest.swift @@ -28,6 +28,11 @@ class ModuleTest: XCTestCase { XCTAssertTrue(module.isLoaded()) } + func testInvalidModuleLoad() { + let module = Module(filePath: "invalid/path") + XCTAssertThrowsError(try module.load()) + } + func testLoadMethod() { guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else { XCTFail("Couldn't find the model file") @@ -149,4 +154,43 @@ class ModuleTest: XCTestCase { XCTAssertEqual(methodMetadata.backendNames.count, 0) XCTAssertEqual(methodMetadata.instructionCount, 1) } + + func testSetInputs() { + guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else { + XCTFail("Couldn't find the model file") + return + } + let module = Module(filePath: modelPath) + + XCTAssertNoThrow(try module.setInput(Tensor([2]), at: 1)) + XCTAssertNoThrow(try module.setInput(Tensor([1]))) + XCTAssertEqual(try module.forward(), Tensor([3])) + + XCTAssertNoThrow(try module.setInputs(Tensor([3]), Tensor([4]))) + XCTAssertEqual(try module.forward(), Tensor([7])) + + XCTAssertThrowsError(try module.setInputs(Tensor([1]))) + } + + func testUnloadMethod() { + guard let modelPath = resourceBundle.path(forResource: "add", ofType: "pte") else { + XCTFail("Couldn't find the model file") + return + } + let module = Module(filePath: modelPath) + XCTAssertNoThrow(try module.load("forward")) + XCTAssertTrue(module.isLoaded("forward")) + + XCTAssertNoThrow(try module.setInputs(Tensor([1]), Tensor([2]))) + XCTAssertEqual(try module.forward(), Tensor([3])) + + XCTAssertTrue(module.unload("forward")) + XCTAssertFalse(module.isLoaded("forward")) + XCTAssertFalse(module.unload("forward")) + + XCTAssertThrowsError(try module.forward()) + XCTAssertTrue(module.isLoaded("forward")) + XCTAssertNoThrow(try module.setInputs(Tensor([2]), Tensor([3]))) + XCTAssertEqual(try module.forward(), Tensor([5])) + } } diff --git a/extension/audio/README.md b/extension/audio/README.md new file mode 100644 index 00000000000..ba2b87c73f2 --- /dev/null +++ b/extension/audio/README.md @@ -0,0 +1,11 @@ +# Audio Processing with ExecuTorch + +The file `mel_spectrogram.py` contains the class `WhisperAudioProcessor`, a module which converts a mono waveform audio input (as a 1D tensor) into Mel spectrograms. It applies a Short-Time Fourier Transform (via torch.stft) and a Mel filterbank. It is equivalent to the `WhisperFeatureExtractor` class in HuggingFace Transformers, but is implemented in PyTorch instead of NumPy. `WhisperFeatureExtractor` is used for Whisper, Voxtral, Qwen2 audio and Qwen2.5 omni. For example, the output Mel spectrograms can be fed directly into the Whisper model (encoder+decoder) exported from HF Transformers. + +Since `WhisperAudioProcessor` is written in PyTorch, we can export it with ExecuTorch and run it on device. The defaults for `WhisperAudioProcessor` are 16kHz audio and 80 Mel spectrogram bins and audio chunks of 30 sec. + +Run it as a script + +``` python mel_spectrogram.py ``` + +to export `WhisperFeatureExtractor` (with default constructor arguments) as `whisper_preprocess.pte`, which can run on device (on CPU). diff --git a/extension/audio/TARGETS b/extension/audio/TARGETS new file mode 100644 index 00000000000..53b0cc47625 --- /dev/null +++ b/extension/audio/TARGETS @@ -0,0 +1,28 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_library( + name = "mel_spectrogram_lib", + srcs = ["mel_spectrogram.py"], + deps = [ + "//caffe2:torch", + "//executorch/devtools/backend_debug:delegation_info", + "//executorch/backends/xnnpack/partition:xnnpack_partitioner", + "//executorch/runtime:runtime", + "fbsource//third-party/pypi/datasets:datasets", + "fbsource//third-party/pypi/transformers:transformers", + "fbsource//third-party/pypi/librosa:librosa", + "fbsource//third-party/pypi/soundfile:soundfile" + ] +) + +python_binary( + name = "mel_spectrogram", + main_module = "executorch.extension.audio.mel_spectrogram", + deps = [ + ":mel_spectrogram_lib", + ], +) diff --git a/extension/audio/mel_spectrogram.py b/extension/audio/mel_spectrogram.py new file mode 100644 index 00000000000..b5d8763d67a --- /dev/null +++ b/extension/audio/mel_spectrogram.py @@ -0,0 +1,254 @@ +# 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 logging + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner + +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + to_edge_transform_and_lower, +) + +from torch.export import Dim, export, ExportedProgram + + +class WhisperAudioProcessor(nn.Module): + r""" + Computes Mel spectrograms from mono audio input. + Same as HuggingFace WhisperFeatureExtractor, but implemented in PyTorch + + Args: + feature_size (`int`, defaults to 80): + The feature dimension of the extracted features. + sampling_rate (`int`, defaults to 16000): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + hop_length (`int`, defaults to 160): + Length of the overlaping windows for the STFT used to obtain the Mel Frequency coefficients. + chunk_length (`int`, defaults to 30): + The maximum number of chuncks of `sampling_rate` samples used to trim and pad longer or shorter audio + sequences. + n_fft (`int`, defaults to 400): + Size of the Fourier transform. + padding_value (`float`, *optional*, defaults to 0.0): + Padding value used to pad the audio. Should correspond to silences. + """ + + def __init__( + self, + feature_size: int = 80, + sampling_rate: int = 16000, + hop_length: int = 160, + chunk_length: int = 30, + n_fft: int = 400, + padding_value: float = 0.0, + ) -> None: + super().__init__() + self.feature_size = feature_size + self.sampling_rate = sampling_rate + self.padding_value = padding_value + + self.n_fft = n_fft + self.hop_length = hop_length + self.chunk_length = chunk_length + self.n_samples = chunk_length * sampling_rate + self.nb_max_frames = self.n_samples // hop_length + self.sampling_rate = sampling_rate + self.mel_filters = self.get_mel_filters( + sampling_rate, n_fft, n_mels=feature_size + ) + + def get_mel_filters( + self, sr: int, n_fft: int, n_mels: int = 128, dtype: torch.dtype = torch.float32 + ) -> torch.Tensor: + # Initialize the weights + n_mels = int(n_mels) + weights = torch.zeros((n_mels, int(1 + n_fft // 2)), dtype=dtype) + + # Center freqs of each FFT bin + fftfreqs = torch.fft.rfftfreq(n=n_fft, d=1.0 / sr, dtype=dtype) + + # 'Center freqs' of mel bands - uniformly spaced between limits + min_mel = 0.0 + max_mel = 45.245640471924965 + + mels = torch.linspace(min_mel, max_mel, n_mels + 2, dtype=dtype) + + # Fill in the linear scale + f_min = 0.0 + f_sp = 200.0 / 3 + freqs = f_min + f_sp * mels + + # And now the nonlinear scale + min_log_hz = 1000.0 # beginning of log region (Hz) + min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels) + logstep = ( + torch.log(torch.tensor(6.4, dtype=dtype)) / 27.0 + ) # step size for log region + + # If we have vector data, vectorize + log_t = mels >= min_log_mel + freqs[log_t] = min_log_hz * torch.exp(logstep * (mels[log_t] - min_log_mel)) + + mel_f = freqs + + fdiff = torch.diff(mel_f) + ramps = torch.subtract(mel_f.unsqueeze(1), fftfreqs.unsqueeze(0)) + + for i in range(n_mels): + # lower and upper slopes for all bins + lower = -ramps[i] / fdiff[i] + upper = ramps[i + 2] / fdiff[i + 1] + + # .. then intersect them with each other and zero + weights[i] = torch.maximum( + torch.tensor(0.0, dtype=dtype), torch.minimum(lower, upper) + ) + + # Slaney-style mel is scaled to be approx constant energy per channel + enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels]) # pyre-ignore[58] + weights *= enorm[:, None] # pyre-ignore[16] + + return weights + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + r""" + Args: + waveform (`torch.Tensor`): Mono waveform input, tensor of (dynamic) shape [num_samples], + + Returns: + torch.Tensor: Output of shape [1, feature_size, nb_max_frames * n_chunks] + n_chunks is the number of chunks of `sampling_rate` samples in the input waveform. + [1, 80, 3000] with default options and 1 chunk + """ + n_chunks = (waveform.shape[0] - 1) // self.n_samples + 1 + waveform = F.pad( + waveform, + (0, self.n_samples * n_chunks - waveform.shape[0]), + mode="constant", + value=self.padding_value, + ) + # Ideally we should do: + # window = torch.hann_window(self.n_fft) + # but this is not currently supported when lowering. + # torch.hann_window has slightly better numerics (worst discrepancy is <1e-5 instead of 1e-4) + window = 0.5 * ( + 1 + - torch.cos( + 2 + * torch.pi + * torch.linspace(0, self.n_fft - 1, self.n_fft, dtype=torch.float32) + / self.n_fft + ) + ) + stft = torch.stft( + waveform, + n_fft=self.n_fft, + hop_length=self.hop_length, + window=window, + center=True, + return_complex=True, + ) + magnitudes = torch.abs(stft)[..., :-1] ** 2 # pyre-ignore[58] + + mel_spec = self.mel_filters @ magnitudes + + log_spec = torch.log10(torch.clamp(mel_spec, min=1e-10)) + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + + return log_spec.unsqueeze(0) + + +def export_processor(model=None, output_file="whisper_preprocess.pte"): + if model is None: + model = WhisperAudioProcessor() + audio_tensor = torch.randn(480000) + chunk_tensor = audio_tensor[:93680] + with torch.no_grad(): + dim = Dim("waveform", min=1600, max=audio_tensor.size(0) * 10) # 10 chunks max + ep: ExportedProgram = export( + model, (chunk_tensor,), dynamic_shapes={"waveform": {0: dim}}, strict=True + ) + logging.debug(ep) + + # to edge + edge: EdgeProgramManager = to_edge_transform_and_lower( + ep, + partitioner=[XnnpackPartitioner()], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + ), + ) + logging.debug(edge.exported_program()) + + # to executorch + exec_prog = edge.to_executorch() + with open(output_file, "wb") as file: + exec_prog.write_to_file(file) + + logging.debug("Done") + + +def main(): + parser = argparse.ArgumentParser( + description="Export WhisperAudioProcessor to ExecutorTorch" + ) + parser.add_argument( + "--feature_size", + type=int, + default=80, + help="The feature dimension of the extracted features", + ) + parser.add_argument( + "--sampling_rate", + type=int, + default=16000, + help="The sampling rate at which audio files should be digitalized (Hz)", + ) + parser.add_argument( + "--hop_length", + type=int, + default=160, + help="Length of overlapping windows for STFT", + ) + parser.add_argument( + "--chunk_length", + type=int, + default=30, + help="Maximum number of chunks of sampling_rate samples", + ) + parser.add_argument( + "--n_fft", type=int, default=400, help="Size of the Fourier transform" + ) + parser.add_argument( + "--output_file", + type=str, + default="whisper_preprocess.pte", + help="Output file path for the exported model", + ) + + args = parser.parse_args() + + model = WhisperAudioProcessor( + feature_size=args.feature_size, + sampling_rate=args.sampling_rate, + hop_length=args.hop_length, + chunk_length=args.chunk_length, + n_fft=args.n_fft, + ) + + export_processor(model, args.output_file) + + +if __name__ == "__main__": + main() diff --git a/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm b/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm index fbf1a6c5889..0f509f2809c 100644 --- a/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm +++ b/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm @@ -87,7 +87,7 @@ @implementation LLaMATests for (NSUInteger index = 2; specialTokens.count < 256; ++index) { [specialTokens addObject:[NSString stringWithFormat:@"<|reserved_special_token_%zu|>", index]]; } - auto __block runner = [[ExecuTorchTextLLMRunner alloc] initWithModelPath:modelPath + auto __block runner = [[ExecuTorchLLMTextRunner alloc] initWithModelPath:modelPath tokenizerPath:tokenizerPath specialTokens:specialTokens]; NSError *error; diff --git a/extension/data_loader/CMakeLists.txt b/extension/data_loader/CMakeLists.txt index 104cd23c977..a5e7a0c4a81 100644 --- a/extension/data_loader/CMakeLists.txt +++ b/extension/data_loader/CMakeLists.txt @@ -24,6 +24,11 @@ if(NOT ET_HAVE_SYS_MMAN_H AND NOT WIN32) "extension/data_loader/mmap_data_loader.cpp" ) endif() +if(WIN32) + list(APPEND _extension_data_loader__srcs + "extension/data_loader/mman_windows.cpp" + ) +endif() list(TRANSFORM _extension_data_loader__srcs PREPEND "${EXECUTORCH_ROOT}/") add_library(extension_data_loader ${_extension_data_loader__srcs}) target_link_libraries(extension_data_loader executorch_core) diff --git a/extension/data_loader/mman_windows.cpp b/extension/data_loader/mman_windows.cpp index 2a7f462f99c..89f9f22f467 100644 --- a/extension/data_loader/mman_windows.cpp +++ b/extension/data_loader/mman_windows.cpp @@ -24,7 +24,7 @@ #include #ifndef STATUS_SECTION_TOO_BIG -#define STATUS_SECTION_TOO_BIG ((NTSTATUS)0xC0000040L) +#define STATUS_SECTION_TOO_BIG 0xC0000040L #endif #ifndef FILE_MAP_EXECUTE diff --git a/extension/evalue_util/test/print_evalue_test.cpp b/extension/evalue_util/test/print_evalue_test.cpp index b881e55d8a8..242cb0af224 100644 --- a/extension/evalue_util/test/print_evalue_test.cpp +++ b/extension/evalue_util/test/print_evalue_test.cpp @@ -267,7 +267,7 @@ TEST(PrintEvalueTest, UnelidedBoolLists) { // case; the other scalar types use the same underlying code, so they don't // need to test this again. { - EValue value(ArrayRef(list.data(), 0ul)); + EValue value(ArrayRef(list.data(), static_cast(0ul))); expect_output(value, "(len=0)[]"); } { @@ -419,7 +419,7 @@ TEST(PrintEvalueTest, UnelidedDoubleLists) { std::array list = {-2.2, -1, 0, INFINITY, NAN, 3.3}; { - EValue value(ArrayRef(list.data(), 0ul)); + EValue value(ArrayRef(list.data(), static_cast(0ul))); expect_output(value, "(len=0)[]"); } { diff --git a/extension/export_util/utils.py b/extension/export_util/utils.py index aa3a736af3c..782b2b0ae63 100644 --- a/extension/export_util/utils.py +++ b/extension/export_util/utils.py @@ -14,7 +14,7 @@ import torch from executorch.exir import EdgeProgramManager, ExecutorchProgramManager, to_edge from executorch.exir.tracer import Value -from torch.export import export, export_for_training, ExportedProgram +from torch.export import export, ExportedProgram _EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig( @@ -108,7 +108,7 @@ def export_to_exec_prog( ) -> ExecutorchProgramManager: m = model.eval() # pre-autograd export. eventually this will become torch.export - m = export_for_training(m, example_inputs, strict=True).module() + m = export(m, example_inputs, strict=True).module() core_aten_ep = _to_core_aten( m, diff --git a/extension/flat_tensor/flat_tensor_data_map.cpp b/extension/flat_tensor/flat_tensor_data_map.cpp index 3a69dc8b92c..478ce9d63cf 100644 --- a/extension/flat_tensor/flat_tensor_data_map.cpp +++ b/extension/flat_tensor/flat_tensor_data_map.cpp @@ -56,7 +56,8 @@ Result get_named_data( return Error::NotFound; } for (int i = 0; i < named_data->size(); i++) { - if (std::strncmp( + if (key.size() == named_data->Get(i)->key()->size() && + std::strncmp( named_data->Get(i)->key()->c_str(), key.data(), named_data->Get(i)->key()->size()) == 0) { @@ -73,7 +74,8 @@ Result get_named_data( segments->size()); // Validate the segment. ET_CHECK_OR_RETURN_ERROR( - segments->Get(segment_index)->offset() < segment_end_offset, + (segments->Get(segment_index)->offset() + + segments->Get(segment_index)->size()) <= segment_end_offset, InvalidExternalData, "Invalid segment offset %" PRIu64 " is larger than the segment_base_offset + segment_data_size %" PRIu64 @@ -206,15 +208,21 @@ ET_NODISCARD Result FlatTensorDataMap::get_key( } Result fh = FlatTensorHeader::Parse(header->data(), header->size()); - if (fh.error() == Error::NotFound) { - // No header, throw error. - ET_LOG(Error, "No FlatTensorHeader found."); - return fh.error(); - } else if (fh.error() != Error::Ok) { - // corruption, throw error. - ET_LOG(Error, "Flat tensor header may be corrupt."); - return fh.error(); - } + + ET_CHECK_OR_RETURN_ERROR( + fh.ok(), + InvalidExternalData, + "Failed to parse FlatTensor header with error code %u. File may be corrupt.", + static_cast(fh.error())); + + size_t expected_size = fh->segment_base_offset + fh->segment_data_size; + size_t actual_size = loader->size().get(); + ET_CHECK_OR_RETURN_ERROR( + expected_size == actual_size, + InvalidExternalData, + "File size is too small; file may be corrupted or truncated. Expected %zu from flat_tensor header, received %zu from data loader", + expected_size, + actual_size); // Load flatbuffer data as a segment. Result flat_tensor_data = loader->load( diff --git a/extension/flat_tensor/serialize/CMakeLists.txt b/extension/flat_tensor/serialize/CMakeLists.txt index 39b364797b8..1909bd4de08 100644 --- a/extension/flat_tensor/serialize/CMakeLists.txt +++ b/extension/flat_tensor/serialize/CMakeLists.txt @@ -10,8 +10,12 @@ # ~~~ # The include directory that will contain the generated schema headers. -set(_flat_tensor_schema__include_dir "${CMAKE_BINARY_DIR}/extension/flat_tensor/include") -set(_flat_tensor_schema__output_dir "${_flat_tensor_schema__include_dir}/executorch/extension/flat_tensor/serialize") +set(_flat_tensor_schema__include_dir + "${CMAKE_BINARY_DIR}/extension/flat_tensor/include" +) +set(_flat_tensor_schema__output_dir + "${_flat_tensor_schema__include_dir}/executorch/extension/flat_tensor/serialize" +) # Source root directory for executorch. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/..) @@ -29,9 +33,8 @@ function(generate_flat_tensor_schema _schema_srcs _schema_name) # Generate the headers from the .fbs files. add_custom_command( OUTPUT ${_schema_outputs} - COMMAND - flatc --cpp --cpp-std c++11 --gen-mutable --scoped-enums -o - "${_flat_tensor_schema__output_dir}" ${_schema_srcs} + COMMAND flatc --cpp --cpp-std c++11 --gen-mutable --scoped-enums -o + "${_flat_tensor_schema__output_dir}" ${_schema_srcs} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS flatc ${_schema_srcs} COMMENT "Generating ${_schema_name} headers" @@ -45,7 +48,8 @@ function(generate_flat_tensor_schema _schema_srcs _schema_name) # and some users need an alignment larger than the default, which is typically # 32. target_compile_definitions( - ${_schema_name} INTERFACE FLATBUFFERS_MAX_ALIGNMENT=${EXECUTORCH_FLATBUFFERS_MAX_ALIGNMENT} + ${_schema_name} + INTERFACE FLATBUFFERS_MAX_ALIGNMENT=${EXECUTORCH_FLATBUFFERS_MAX_ALIGNMENT} ) target_include_directories( diff --git a/extension/flat_tensor/serialize/serialize.cpp b/extension/flat_tensor/serialize/serialize.cpp index 9930de6bba6..ff9d568fdf4 100644 --- a/extension/flat_tensor/serialize/serialize.cpp +++ b/extension/flat_tensor/serialize/serialize.cpp @@ -76,6 +76,7 @@ runtime::Error save_ptd( // Write the tensors. size_t total_segment_size = 0; uint32_t i = 0; + size_t tensor_count = tensor_map.size(); for (const auto& [name, tensor] : tensor_map) { auto key = builder.CreateString(name); // Write the tensor layouts. @@ -99,7 +100,11 @@ runtime::Error save_ptd( /*_fbb=*/builder, /*offset=*/total_segment_size, /*size=*/tensor.nbytes())); - total_segment_size += aligned_size(tensor.nbytes(), tensor_alignment); + + // Do not pad the last tensor. + total_segment_size += (i == tensor_count - 1) + ? tensor.nbytes() + : aligned_size(tensor.nbytes(), tensor_alignment); i++; } diff --git a/extension/flat_tensor/serialize/serialize.py b/extension/flat_tensor/serialize/serialize.py index 7f3332c4303..3d8a0c75cc4 100644 --- a/extension/flat_tensor/serialize/serialize.py +++ b/extension/flat_tensor/serialize/serialize.py @@ -1,11 +1,13 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2025 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 importlib.resources as _resources import json import math import os @@ -13,10 +15,10 @@ from dataclasses import dataclass from typing import ClassVar, Dict, List, Literal, Optional -import pkg_resources +import executorch.extension.flat_tensor.serialize as serialize_package + from executorch.exir._serialize._cord import Cord from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass - from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile from executorch.exir._serialize._program import _insert_flatbuffer_header from executorch.exir._serialize.data_serializer import ( @@ -24,9 +26,7 @@ DataPayload, DataSerializer, ) - from executorch.exir._serialize.padding import aligned_size, pad_to, padding_required - from executorch.extension.flat_tensor.serialize.flat_tensor_schema import ( DataSegment, FlatTensor, @@ -49,12 +49,12 @@ def _serialize_to_flatbuffer(flat_tensor: FlatTensor) -> Cord: schema_path = os.path.join(d, "flat_tensor.fbs") with open(schema_path, "wb") as schema_file: schema_file.write( - pkg_resources.resource_string(__name__, "flat_tensor.fbs") + _resources.read_binary(serialize_package, "flat_tensor.fbs") ) scalar_type_path = os.path.join(d, "scalar_type.fbs") with open(scalar_type_path, "wb") as scalar_type_file: scalar_type_file.write( - pkg_resources.resource_string(__name__, "scalar_type.fbs") + _resources.read_binary(serialize_package, "scalar_type.fbs") ) json_path = os.path.join(d, "flat_tensor.json") with open(json_path, "wb") as json_file: @@ -72,13 +72,13 @@ def _deserialize_to_flat_tensor(flatbuffer: bytes) -> FlatTensor: schema_path = os.path.join(d, "flat_tensor.fbs") with open(schema_path, "wb") as schema_file: schema_file.write( - pkg_resources.resource_string(__name__, "flat_tensor.fbs") + _resources.read_binary(serialize_package, "flat_tensor.fbs") ) scalar_type_path = os.path.join(d, "scalar_type.fbs") with open(scalar_type_path, "wb") as scalar_type_file: scalar_type_file.write( - pkg_resources.resource_string(__name__, "scalar_type.fbs") + _resources.read_binary(serialize_package, "scalar_type.fbs") ) bin_path = os.path.join(d, "flat_tensor.bin") diff --git a/extension/flat_tensor/targets.bzl b/extension/flat_tensor/targets.bzl index 3bc36dad9d4..f91e28a2268 100644 --- a/extension/flat_tensor/targets.bzl +++ b/extension/flat_tensor/targets.bzl @@ -1,7 +1,7 @@ -load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime") def define_common_targets(): - for aten_mode in [True, False]: + for aten_mode in get_aten_mode_options(): aten_suffix = "_aten" if aten_mode else "" runtime.cxx_library( name = "flat_tensor_data_map" + aten_suffix, diff --git a/extension/flat_tensor/test/CMakeLists.txt b/extension/flat_tensor/test/CMakeLists.txt index c3296dc61f3..fd3d6792f90 100644 --- a/extension/flat_tensor/test/CMakeLists.txt +++ b/extension/flat_tensor/test/CMakeLists.txt @@ -23,7 +23,7 @@ add_custom_command( "${CMAKE_CURRENT_BINARY_DIR}/ModuleAddMulProgram.ptd" COMMAND ${PYTHON_EXECUTABLE} -m test.models.export_program --modules "ModuleAddMul" - --external-constants --outdir "${CMAKE_CURRENT_BINARY_DIR}" 2> /dev/null + --external-constants --outdir "${CMAKE_CURRENT_BINARY_DIR}" WORKING_DIRECTORY ${EXECUTORCH_ROOT} ) diff --git a/extension/flat_tensor/test/flat_tensor_data_map_test.cpp b/extension/flat_tensor/test/flat_tensor_data_map_test.cpp index 37e1cd2edac..0872a988333 100644 --- a/extension/flat_tensor/test/flat_tensor_data_map_test.cpp +++ b/extension/flat_tensor/test/flat_tensor_data_map_test.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include @@ -17,6 +18,8 @@ #include using namespace ::testing; +using executorch::extension::BufferDataLoader; +using executorch::extension::FileDataLoader; using executorch::extension::FlatTensorDataMap; using executorch::extension::FlatTensorHeader; using executorch::runtime::DataLoader; @@ -24,7 +27,6 @@ using executorch::runtime::Error; using executorch::runtime::FreeableBuffer; using executorch::runtime::Result; using executorch::runtime::TensorLayout; -using torch::executor::util::FileDataLoader; class FlatTensorDataMapTest : public ::testing::Test { protected: @@ -51,7 +53,7 @@ TEST_F(FlatTensorDataMapTest, LoadFlatTensorDataMap) { EXPECT_EQ(data_map.error(), Error::Ok); } -TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetMetadata) { +TEST_F(FlatTensorDataMapTest, GetMetadata) { Result data_map = FlatTensorDataMap::load(data_map_loader_.get()); EXPECT_EQ(data_map.error(), Error::Ok); @@ -93,7 +95,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetMetadata) { EXPECT_EQ(const_c_res.error(), Error::NotFound); } -TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetData) { +TEST_F(FlatTensorDataMapTest, GetData) { Result data_map = FlatTensorDataMap::load(data_map_loader_.get()); EXPECT_EQ(data_map.error(), Error::Ok); @@ -114,7 +116,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_GetData) { EXPECT_EQ(data_c_res.error(), Error::NotFound); } -TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_Keys) { +TEST_F(FlatTensorDataMapTest, GetKeys) { Result data_map = FlatTensorDataMap::load(data_map_loader_.get()); EXPECT_EQ(data_map.error(), Error::Ok); @@ -138,7 +140,7 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_Keys) { EXPECT_EQ(key2_res.error(), Error::InvalidArgument); } -TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_LoadInto) { +TEST_F(FlatTensorDataMapTest, LoadInto) { Result data_map = FlatTensorDataMap::load(data_map_loader_.get()); EXPECT_EQ(data_map.error(), Error::Ok); @@ -160,3 +162,23 @@ TEST_F(FlatTensorDataMapTest, FlatTensorDataMap_LoadInto) { } free(data); } + +TEST_F(FlatTensorDataMapTest, LoadAndCheckSize) { + Result data_map = + FlatTensorDataMap::load(data_map_loader_.get()); + EXPECT_EQ(data_map.error(), Error::Ok); + + // Truncate the file. + size_t trunc_size = data_map_loader_->size().get() - 8; + Result truncated_file = data_map_loader_->load( + 0, + trunc_size, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::External)); + ASSERT_EQ(truncated_file.error(), Error::Ok); + + BufferDataLoader truncated_loader = + BufferDataLoader(truncated_file->data(), trunc_size); + Result truncated_program = + FlatTensorDataMap::load(&truncated_loader); + ASSERT_EQ(truncated_program.error(), Error::InvalidExternalData); +} diff --git a/extension/flat_tensor/test/targets.bzl b/extension/flat_tensor/test/targets.bzl index 4d798cc1a7c..d1272ec5d97 100644 --- a/extension/flat_tensor/test/targets.bzl +++ b/extension/flat_tensor/test/targets.bzl @@ -45,6 +45,7 @@ def define_common_targets(is_fbcode=False): "flat_tensor_data_map_test.cpp", ], deps = [ + "//executorch/extension/data_loader:buffer_data_loader", "//executorch/extension/data_loader:file_data_loader", "//executorch/extension/flat_tensor:flat_tensor_data_map", "//executorch/runtime/core:named_data_map", diff --git a/extension/flat_tensor/test/test_serialize.cpp b/extension/flat_tensor/test/test_serialize.cpp index 57a0253485b..35a1e9ee8dc 100644 --- a/extension/flat_tensor/test/test_serialize.cpp +++ b/extension/flat_tensor/test/test_serialize.cpp @@ -86,8 +86,10 @@ TEST_F(FlatTensorSerializeTest, ValidFlatTensorSerialized) { const uint64_t segment_offset = 48 + 280 + 8; // 8 is padding. EXPECT_EQ(*(uint64_t*)(header_buffer + 24), segment_offset); - // Segment total size, 8 bytes of data (2 floats), 24 bytes of padding. - const uint64_t segment_size = 32; + // Segment total size = 20 + // linear.bias: 4 bytes + 12 bytes of padding. + // linear.weight: 4 bytes + 0 padding (last segment). + const uint64_t segment_size = 20; EXPECT_EQ(*(uint64_t*)(header_buffer + 32), segment_size); // Check Flatbuffer diff --git a/extension/kernel_util/make_boxed_from_unboxed_functor.h b/extension/kernel_util/make_boxed_from_unboxed_functor.h index 8f3d63db449..1710f876b20 100644 --- a/extension/kernel_util/make_boxed_from_unboxed_functor.h +++ b/extension/kernel_util/make_boxed_from_unboxed_functor.h @@ -41,7 +41,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -61,6 +63,49 @@ namespace extension { // internal namespace to avoid conflicts with other extensions. namespace kernel_util_internal { +// Template trait to check if a type is a non-const tensor +template +struct is_nonconst_tensor : std::false_type {}; + +template <> +struct is_nonconst_tensor : std::true_type {}; + +// Template trait to check if a type is a non-const tensor +// Count non-const tensors in a typelist +template +struct count_nonconst_tensors; + +template <> +struct count_nonconst_tensors> { + static constexpr size_t value = 0; +}; + +template +struct count_nonconst_tensors> { + static constexpr size_t value = 0; +}; + +template <> +struct count_nonconst_tensors> { + static constexpr size_t value = 1; +}; + +template +struct count_nonconst_tensors> { + private: + static constexpr size_t tail_tensor_count = + count_nonconst_tensors>::value; + static constexpr size_t tail_args_count = sizeof...(Tail); + static constexpr bool is_head_a_tensor = is_nonconst_tensor::value; + static constexpr bool all_tail_args_are_tensor = + tail_tensor_count == tail_args_count; + + public: + static constexpr size_t value = (is_head_a_tensor && all_tail_args_are_tensor) + ? tail_tensor_count + 1 + : tail_tensor_count; +}; + template struct decay_if_not_tensor final { using type = std::decay_t; @@ -110,16 +155,29 @@ struct evalue_to_arg>> final { } }; -template +template < + class Functor, + size_t nonconst_tensors_to_log, + size_t... evalue_arg_indices, + typename... ArgTypes> void call_functor_with_args_from_stack( executorch::runtime::KernelRuntimeContext& ctx, executorch::runtime::Span stack, std::index_sequence, typelist*) { + executorch::runtime::internal::EventTracerProfileOpScope + event_tracer_op_scope(ctx.internal_event_tracer(), Functor::func_name_); + EXECUTORCH_SCOPE_PROF(Functor::func_name_); (*Functor::func_ptr())( ctx, evalue_to_arg::type>::call( *stack[evalue_arg_indices])...); + constexpr size_t num_inputs = + std::index_sequence::size(); + for (size_t i = num_inputs - nonconst_tensors_to_log; i < num_inputs; ++i) { + executorch::runtime::internal::event_tracer_log_evalue( + ctx.internal_event_tracer(), *stack[i]); + } } } // namespace kernel_util_internal @@ -154,11 +212,16 @@ struct WrapUnboxedIntoFunctor { executorch::runtime::Span stack) { constexpr size_t num_inputs = kernel_util_internal::size::value; - return kernel_util_internal::call_functor_with_args_from_stack( - ctx, - stack, - std::make_index_sequence(), - static_cast(nullptr)); + constexpr size_t num_nonconst_tensors = + kernel_util_internal::count_nonconst_tensors< + ContextRemovedArgsType>::value; + static_assert(num_nonconst_tensors == 1, "Invalid number of inputs"); + return kernel_util_internal:: + call_functor_with_args_from_stack( + ctx, + stack, + std::make_index_sequence(), + static_cast(nullptr)); } }; @@ -181,11 +244,14 @@ static executorch::runtime::Kernel make_boxed_kernel( #define EXECUTORCH_LIBRARY(ns, op_name, func) \ _EXECUTORCH_LIBRARY_IMPL(ns, op_name, func, ET_UID) -#define _EXECUTORCH_LIBRARY_IMPL(ns, op_name, func, uid) \ - static auto ET_CONCATENATE(res_##ns##_, uid) = \ - ::executorch::runtime::register_kernel( \ - ::executorch::extension::make_boxed_kernel( \ - #ns "::" op_name, EXECUTORCH_FN(func))) +#define _EXECUTORCH_LIBRARY_IMPL(ns, op_name, func, uid) \ + static constexpr const char ET_CONCATENATE(name_of_op_, uid)[] = \ + #ns "::" op_name; \ + static auto ET_CONCATENATE(res_##ns##_, uid) = \ + ::executorch::runtime::register_kernel( \ + ::executorch::extension::make_boxed_kernel( \ + #ns "::" op_name, \ + EXECUTORCH_FN(func, ET_CONCATENATE(name_of_op_, uid)))) namespace torch { namespace executor { diff --git a/extension/kernel_util/meta_programming.h b/extension/kernel_util/meta_programming.h index 027568fe687..e3b3649dd4f 100644 --- a/extension/kernel_util/meta_programming.h +++ b/extension/kernel_util/meta_programming.h @@ -32,12 +32,13 @@ template using is_function_type_t = typename is_function_type::type; // A compile-time wrapper around a function pointer -template +template struct CompileTimeFunctionPointer final { static_assert( is_function_type::value, "EXECUTORCH_FN can only wrap function types."); using FuncType = FuncType_; + static constexpr const char* func_name_ = func_name; static constexpr FuncType* func_ptr() { return func_ptr_; @@ -47,15 +48,17 @@ struct CompileTimeFunctionPointer final { // Check if a given type is a compile-time function pointer template struct is_compile_time_function_pointer : std::false_type {}; -template +template struct is_compile_time_function_pointer< - CompileTimeFunctionPointer> : std::true_type {}; + CompileTimeFunctionPointer> + : std::true_type {}; -#define EXECUTORCH_FN_TYPE(func) \ +#define EXECUTORCH_FN_TYPE(func, name) \ ::executorch::extension::kernel_util_internal::CompileTimeFunctionPointer< \ std::remove_pointer_t>, \ - func> -#define EXECUTORCH_FN(func) EXECUTORCH_FN_TYPE(func)() + func, \ + name> +#define EXECUTORCH_FN(func, name) EXECUTORCH_FN_TYPE(func, name)() /** * strip_class: helper to remove the class type from pointers to `operator()`. diff --git a/extension/llm/README.md b/extension/llm/README.md index 0f71088eea1..de8e4e6d619 100644 --- a/extension/llm/README.md +++ b/extension/llm/README.md @@ -10,7 +10,7 @@ Commonly used methods in this class include: - _source_transform_: execute a series of source transform passes. Some transform passes include - weight only quantization, which can be done at source (eager mode) level. - replace some torch operators to a custom operator. For example, _replace_sdpa_with_custom_op_. -- _torch.export_for_training_: get a graph that is ready for pt2 graph-based quantization. +- _torch.export_: get a graph that is ready for pt2 graph-based quantization. - _pt2e_quantize_ with passed in quantizers. - util functions in _quantizer_lib.py_ can help to get different quantizers based on the needs. - _export_to_edge_: export to edge dialect diff --git a/extension/llm/apple/CMakeLists.txt b/extension/llm/apple/CMakeLists.txt index aa7da842004..1755f09b67f 100644 --- a/extension/llm/apple/CMakeLists.txt +++ b/extension/llm/apple/CMakeLists.txt @@ -43,6 +43,7 @@ set_source_files_properties( "-fno-rtti" ) -set_target_properties(extension_llm_apple PROPERTIES - XCODE_ATTRIBUTE_BUILD_LIBRARY_FOR_DISTRIBUTION YES +set_target_properties( + extension_llm_apple PROPERTIES XCODE_ATTRIBUTE_BUILD_LIBRARY_FOR_DISTRIBUTION + YES ) diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLM.h b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLM.h index 42e08d18030..11cdaf63d0b 100644 --- a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLM.h +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLM.h @@ -6,4 +6,6 @@ * LICENSE file in the root directory of this source tree. */ -#import "ExecuTorchTextLLMRunner.h" +#import "ExecuTorchLLMError.h" +#import "ExecuTorchLLMMultimodalRunner.h" +#import "ExecuTorchLLMTextRunner.h" diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.h b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.h new file mode 100644 index 00000000000..13b51a0652b --- /dev/null +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.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. + */ + +#import + +NS_ASSUME_NONNULL_BEGIN + +FOUNDATION_EXPORT NSErrorDomain const ExecuTorchLLMErrorDomain NS_SWIFT_NAME(ErrorDomain); + +NS_ASSUME_NONNULL_END diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/App.swift b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.m similarity index 64% rename from examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/App.swift rename to extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.m index ceddbde1e61..52a9ac8c9da 100644 --- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/App.swift +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMError.m @@ -6,13 +6,6 @@ * LICENSE file in the root directory of this source tree. */ -import SwiftUI +#import "ExecuTorchLLMError.h" -@main -struct App: SwiftUI.App { - var body: some Scene { - WindowGroup { - ContentView() - } - } -} +NSErrorDomain const ExecuTorchLLMErrorDomain = @"org.pytorch.executorch.llm.error"; diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.h b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.h new file mode 100644 index 00000000000..747286b9ec3 --- /dev/null +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.h @@ -0,0 +1,196 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#import + +NS_ASSUME_NONNULL_BEGIN + +/** + Types of multimodal inputs supported by the ExecuTorch LLM APIs. + Must be in sync with the C++ enum in llm/runner/multimodal_input.h +*/ +typedef NS_ENUM(NSInteger, ExecuTorchLLMMultimodalInputType) { + ExecuTorchLLMMultimodalInputTypeText, + ExecuTorchLLMMultimodalInputTypeImage, + ExecuTorchLLMMultimodalInputTypeAudio, + ExecuTorchLLMMultimodalInputTypeUnsupported, +} NS_SWIFT_NAME(MultimodalInputType); + +/** + A container for image inputs used with multimodal generation APIs. +*/ +NS_SWIFT_NAME(Image) +__attribute__((deprecated("This API is experimental."))) +@interface ExecuTorchLLMImage : NSObject + +/** + Initializes an image container with the provided data and dimensions. + + @param data Raw image bytes. + @param width Image width in pixels. + @param height Image height in pixels. + @param channels Number of channels. + @return An initialized ExecuTorchLLMImage instance. +*/ +- (instancetype)initWithData:(NSData *)data + width:(NSInteger)width + height:(NSInteger)height + channels:(NSInteger)channels + NS_DESIGNATED_INITIALIZER; + +@property(nonatomic, readonly) NSData *data; +@property(nonatomic, readonly) NSInteger width; +@property(nonatomic, readonly) NSInteger height; +@property(nonatomic, readonly) NSInteger channels; + ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; + +@end + +/** + A container for pre-processed audio features. +*/ +NS_SWIFT_NAME(Audio) +__attribute__((deprecated("This API is experimental."))) +@interface ExecuTorchLLMAudio : NSObject + +/** + Initializes an audio features container with the provided data and shape. + + @param data Feature buffer. + @param batchSize Batch dimension size. + @param bins Number of frequency bins. + @param frames Number of time frames. + @return An initialized ExecuTorchLLMAudio instance. +*/ +- (instancetype)initWithData:(NSData *)data + batchSize:(NSInteger)batchSize + bins:(NSInteger)bins + frames:(NSInteger)frames + NS_DESIGNATED_INITIALIZER; + +@property(nonatomic, readonly) NSData *data; +@property(nonatomic, readonly) NSInteger batchSize; +@property(nonatomic, readonly) NSInteger bins; +@property(nonatomic, readonly) NSInteger frames; + ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; + +@end + +/** + A tagged container for a single multimodal input item used by + multimodal generation APIs. +*/ +NS_SWIFT_NAME(MultimodalInput) +__attribute__((deprecated("This API is experimental."))) +@interface ExecuTorchLLMMultimodalInput : NSObject + +/** + Creates a text input. + + @param text The UTF-8 text to provide as input. + @return A retained ExecuTorchLLMMultimodalInput instance of type Text. +*/ ++ (instancetype)inputWithText:(NSString *)text + NS_SWIFT_NAME(init(_:)) + NS_RETURNS_RETAINED; + +/** + Creates an image input. + + @param image The image payload to provide as input. + @return A retained ExecuTorchLLMMultimodalInput instance of type Image. +*/ ++ (instancetype)inputWithImage:(ExecuTorchLLMImage *)image + NS_SWIFT_NAME(init(_:)) + NS_RETURNS_RETAINED; + +/** + Creates an audio-features input. + + @param audio The pre-processed audio features to provide as input. + @return A retained ExecuTorchLLMMultimodalInput instance of type Audio. +*/ ++ (instancetype)inputWithAudio:(ExecuTorchLLMAudio *)audio + NS_SWIFT_NAME(init(audio:)) + NS_RETURNS_RETAINED; + +@property(nonatomic, readonly) ExecuTorchLLMMultimodalInputType type; +@property(nonatomic, readonly, nullable) NSString *text; +@property(nonatomic, readonly, nullable) ExecuTorchLLMImage *image; +@property(nonatomic, readonly, nullable) ExecuTorchLLMAudio *audio; + ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; + +@end + +/** + A wrapper class for the C++ llm::MultimodalLLMRunner that provides + Objective-C APIs to load models, manage tokenization, accept mixed + input modalities, generate text sequences, and stop the runner. +*/ +NS_SWIFT_NAME(MultimodalRunner) +__attribute__((deprecated("This API is experimental."))) +@interface ExecuTorchLLMMultimodalRunner : NSObject + +/** + Initializes a multimodal LLM runner with the given model and tokenizer paths. + + @param modelPath File system path to the serialized model. + @param tokenizerPath File system path to the tokenizer data. + @return An initialized ExecuTorchLLMMultimodalRunner instance. +*/ +- (instancetype)initWithModelPath:(NSString *)modelPath + tokenizerPath:(NSString *)tokenizerPath + NS_DESIGNATED_INITIALIZER; + +/** + Checks whether the underlying model has been successfully loaded. + + @return YES if the model is loaded, NO otherwise. +*/ +- (BOOL)isLoaded; + +/** + Loads the model into memory, returning an error if loading fails. + + @param error On failure, populated with an NSError explaining the issue. + @return YES if loading succeeds, NO if an error occurred. +*/ +- (BOOL)loadWithError:(NSError **)error; + +/** + Generates text given a list of multimodal inputs, up to a specified sequence length. + Invokes the provided callback for each generated token. + + @param inputs An ordered array of multimodal inputs. + @param seq_len The maximum number of tokens to generate. + @param callback A block called with each generated token as an NSString. + @param error On failure, populated with an NSError explaining the issue. + @return YES if generation completes successfully, NO if an error occurred. +*/ +- (BOOL)generate:(NSArray *)inputs + sequenceLength:(NSInteger)seq_len +withTokenCallback:(nullable void (^)(NSString *))callback + error:(NSError **)error; + +/** + Stops any ongoing generation and cleans up internal resources. +*/ +- (void)stop; + ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; + +@end + +NS_ASSUME_NONNULL_END diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.mm b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.mm new file mode 100644 index 00000000000..dcc5dc98806 --- /dev/null +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMMultimodalRunner.mm @@ -0,0 +1,219 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#import "ExecuTorchLLMMultimodalRunner.h" + +#import "ExecuTorchLLMError.h" + +#import + +using namespace executorch::extension; +using namespace executorch::runtime; + +@implementation ExecuTorchLLMImage + +- (instancetype)initWithData:(NSData *)data + width:(NSInteger)width + height:(NSInteger)height + channels:(NSInteger)channels { + if (self = [super init]) { + _data = [data copy]; + _width = width; + _height = height; + _channels = channels; + } + return self; +} + +- (id)copyWithZone:(NSZone *)zone { + return self; +} + +@end + +@implementation ExecuTorchLLMAudio + +- (instancetype)initWithData:(NSData *)data + batchSize:(NSInteger)batchSize + bins:(NSInteger)bins + frames:(NSInteger)frames { + if (self = [super init]) { + _data = [data copy]; + _batchSize = batchSize; + _bins = bins; + _frames = frames; + } + return self; +} + +- (id)copyWithZone:(NSZone *)zone { + return self; +} + +@end + +@interface ExecuTorchLLMMultimodalInput () + +- (instancetype)initWithType:(ExecuTorchLLMMultimodalInputType)type + text:(NSString * __nullable)text + image:(ExecuTorchLLMImage * __nullable)image + audio:(ExecuTorchLLMAudio * __nullable)audio + NS_DESIGNATED_INITIALIZER; + +@end + +@implementation ExecuTorchLLMMultimodalInput + ++ (instancetype)inputWithText:(NSString *)text { + return [[self alloc] initWithType:ExecuTorchLLMMultimodalInputTypeText + text:text + image:nil + audio:nil]; +} + ++ (instancetype)inputWithImage:(ExecuTorchLLMImage *)image { + return [[self alloc] initWithType:ExecuTorchLLMMultimodalInputTypeImage + text:nil + image:image + audio:nil]; +} + ++ (instancetype)inputWithAudio:(ExecuTorchLLMAudio *)audio { + return [[self alloc] initWithType:ExecuTorchLLMMultimodalInputTypeAudio + text:nil + image:nil + audio:audio]; +} + +- (instancetype)initWithType:(ExecuTorchLLMMultimodalInputType)type + text:(NSString * __nullable)text + image:(ExecuTorchLLMImage * __nullable)image + audio:(ExecuTorchLLMAudio * __nullable)audio { + if (self = [super init]) { + _type = type; + _text = [text copy]; + _image = image; + _audio = audio; + } + return self; +} + +- (id)copyWithZone:(NSZone *)zone { + return self; +} + +@end + +@implementation ExecuTorchLLMMultimodalRunner { + NSString *_modelPath; + NSString *_tokenizerPath; + std::unique_ptr _runner; +} + +- (instancetype)initWithModelPath:(NSString*)modelPath + tokenizerPath:(NSString*)tokenizerPath { + self = [super init]; + if (self) { + _modelPath = [modelPath copy]; + _tokenizerPath = [tokenizerPath copy]; + } + return self; +} + +- (BOOL)isLoaded { + return _runner && _runner->is_loaded(); +} + +- (BOOL)loadWithError:(NSError**)error { + if (![self isLoaded]) { + _runner = llm::create_multimodal_runner( + _modelPath.UTF8String, + llm::load_tokenizer(_tokenizerPath.UTF8String) + ); + if (!_runner) { + if (error) { + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain + code:-1 + userInfo:@{NSLocalizedDescriptionKey: @"Failed to create runner"}]; + } + return NO; + } + } + auto status = _runner->load(); + if (status != Error::Ok) { + if (error) { + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain + code:(NSInteger)status + userInfo:nil]; + } + return NO; + } + return YES; +} + +- (BOOL)generate:(NSArray *)inputs + sequenceLength:(NSInteger)seq_len +withTokenCallback:(nullable void (^)(NSString *))callback + error:(NSError **)error { + if (![self loadWithError:error]) { + return NO; + } + std::vector nativeInputs; + for (ExecuTorchLLMMultimodalInput *input in inputs) { + switch (input.type) { + case ExecuTorchLLMMultimodalInputTypeText: + nativeInputs.emplace_back(llm::MultimodalInput(input.text.UTF8String)); + break; + case ExecuTorchLLMMultimodalInputTypeImage: { + ExecuTorchLLMImage *image = input.image; + std::vector data((uint8_t *)image.data.bytes, (uint8_t *)image.data.bytes + image.data.length); + nativeInputs.emplace_back(llm::MultimodalInput(llm::Image{ + .data = std::move(data), + .width = (int32_t)image.width, + .height = (int32_t)image.height, + .channels = (int32_t)image.channels + })); + break; + } + default: { + if (error) { + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain + code:-2 + userInfo:@{NSLocalizedDescriptionKey: @"Failed to create input"}]; + } + return NO; + } + } + } + auto status = _runner->generate( + std::move(nativeInputs), + llm::GenerationConfig{.seq_len = static_cast(seq_len)}, + [callback](const std::string& token) { + if (callback) { + callback(@(token.c_str())); + } + } + ); + if (status != Error::Ok) { + if (error) { + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain + code:(NSInteger)status + userInfo:nil]; + } + return NO; + } + return YES; +} + +- (void)stop { + if (_runner) { + _runner->stop(); + } +} + +@end diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.h b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.h similarity index 87% rename from extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.h rename to extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.h index 74518605e90..b2c628fadf6 100644 --- a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.h +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.h @@ -10,16 +10,14 @@ NS_ASSUME_NONNULL_BEGIN -FOUNDATION_EXPORT NSErrorDomain const ExecuTorchTextLLMRunnerErrorDomain; - /** A wrapper class for the C++ llm::TextLLMRunner that provides Objective-C APIs to load models, manage tokenization with custom special tokens, generate text sequences, and stop the runner. */ -NS_SWIFT_NAME(TextLLMRunner) +NS_SWIFT_NAME(TextRunner) __attribute__((deprecated("This API is experimental."))) -@interface ExecuTorchTextLLMRunner : NSObject +@interface ExecuTorchLLMTextRunner : NSObject /** Initializes a text LLM runner with the given model and tokenizer paths, @@ -28,11 +26,12 @@ __attribute__((deprecated("This API is experimental."))) @param modelPath File system path to the serialized model. @param tokenizerPath File system path to the tokenizer data. @param tokens An array of NSString special tokens to use during tokenization. - @return An initialized ExecuTorchTextLLMRunner instance. + @return An initialized ExecuTorchLLMTextRunner instance. */ - (instancetype)initWithModelPath:(NSString *)modelPath tokenizerPath:(NSString *)tokenizerPath - specialTokens:(NSArray *)tokens; + specialTokens:(NSArray *)tokens + NS_DESIGNATED_INITIALIZER; /** Checks whether the underlying model has been successfully loaded. @@ -69,6 +68,9 @@ withTokenCallback:(nullable void (^)(NSString *))callback */ - (void)stop; ++ (instancetype)new NS_UNAVAILABLE; +- (instancetype)init NS_UNAVAILABLE; + @end NS_ASSUME_NONNULL_END diff --git a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.mm b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.mm similarity index 84% rename from extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.mm rename to extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.mm index dd57d25dcde..ac50b000704 100644 --- a/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchTextLLMRunner.mm +++ b/extension/llm/apple/ExecuTorchLLM/Exported/ExecuTorchLLMTextRunner.mm @@ -6,16 +6,16 @@ * LICENSE file in the root directory of this source tree. */ -#import "ExecuTorchTextLLMRunner.h" +#import "ExecuTorchLLMTextRunner.h" + +#import "ExecuTorchLLMError.h" #import using namespace executorch::extension; using namespace executorch::runtime; -NSErrorDomain const ExecuTorchTextLLMRunnerErrorDomain = @"ExecuTorchTextLLMRunnerErrorDomain"; - -@implementation ExecuTorchTextLLMRunner { +@implementation ExecuTorchLLMTextRunner { NSString *_modelPath; NSString *_tokenizerPath; std::unique_ptr> _specialTokens; @@ -49,7 +49,7 @@ - (BOOL)loadWithError:(NSError**)error { ); if (!_runner) { if (error) { - *error = [NSError errorWithDomain:ExecuTorchTextLLMRunnerErrorDomain + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain code:-1 userInfo:@{NSLocalizedDescriptionKey: @"Failed to create runner"}]; } @@ -59,7 +59,7 @@ - (BOOL)loadWithError:(NSError**)error { auto status = _runner->load(); if (status != Error::Ok) { if (error) { - *error = [NSError errorWithDomain:ExecuTorchTextLLMRunnerErrorDomain + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain code:(NSInteger)status userInfo:nil]; } @@ -79,12 +79,14 @@ - (BOOL)generate:(NSString*)prompt prompt.UTF8String, llm::GenerationConfig{.seq_len = static_cast(seq_len)}, [callback](const std::string& token) { - if (callback) callback(@(token.c_str())); + if (callback) { + callback(@(token.c_str())); + } } ); if (status != Error::Ok) { if (error) { - *error = [NSError errorWithDomain:ExecuTorchTextLLMRunnerErrorDomain + *error = [NSError errorWithDomain:ExecuTorchLLMErrorDomain code:(NSInteger)status userInfo:nil]; } diff --git a/extension/llm/apple/ExecuTorchLLM/__tests__/MultimodalRunnerTest.swift b/extension/llm/apple/ExecuTorchLLM/__tests__/MultimodalRunnerTest.swift new file mode 100644 index 00000000000..55bcbb0f407 --- /dev/null +++ b/extension/llm/apple/ExecuTorchLLM/__tests__/MultimodalRunnerTest.swift @@ -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. + */ + +import ExecuTorchLLM +import XCTest + +class MultimodalRunnerTest: XCTestCase { + func test() { + let bundle = Bundle(for: type(of: self)) + guard let modelPath = bundle.path(forResource: "llava", ofType: "pte"), + let tokenizerPath = bundle.path(forResource: "tokenizer", ofType: "bin") else { + XCTFail("Couldn't find model or tokenizer files") + return + } + return + let runner = MultimodalRunner(modelPath: modelPath, tokenizerPath: tokenizerPath) + var text = "" + + do { + try runner.generate([MultimodalInput("hello")], sequenceLength: 2) { token in + text += token + } + } catch { + XCTFail("Failed to generate text with error \(error)") + } + XCTAssertEqual("hello,", text.lowercased()) + } +} diff --git a/extension/llm/apple/ExecuTorchLLM/__tests__/TextLLMRunnerTest.swift b/extension/llm/apple/ExecuTorchLLM/__tests__/TextRunnerTest.swift similarity index 89% rename from extension/llm/apple/ExecuTorchLLM/__tests__/TextLLMRunnerTest.swift rename to extension/llm/apple/ExecuTorchLLM/__tests__/TextRunnerTest.swift index 030da22295d..42dbac8ae30 100644 --- a/extension/llm/apple/ExecuTorchLLM/__tests__/TextLLMRunnerTest.swift +++ b/extension/llm/apple/ExecuTorchLLM/__tests__/TextRunnerTest.swift @@ -35,7 +35,7 @@ struct SpecialTokens { } } -class TextLLMRunnerTest: XCTestCase { +class TextRunnerTest: XCTestCase { func test() { let bundle = Bundle(for: type(of: self)) guard let modelPath = bundle.path(forResource: "llama3_2-1B", ofType: "pte"), @@ -43,7 +43,7 @@ class TextLLMRunnerTest: XCTestCase { XCTFail("Couldn't find model or tokenizer files") return } - let runner = TextLLMRunner(modelPath: modelPath, tokenizerPath: tokenizerPath, specialTokens: SpecialTokens.defaultSpecialTokens()) + let runner = TextRunner(modelPath: modelPath, tokenizerPath: tokenizerPath, specialTokens: SpecialTokens.defaultSpecialTokens()) var text = "" do { diff --git a/extension/llm/custom_ops/CMakeLists.txt b/extension/llm/custom_ops/CMakeLists.txt index c5eba4b7a19..1678dc80296 100644 --- a/extension/llm/custom_ops/CMakeLists.txt +++ b/extension/llm/custom_ops/CMakeLists.txt @@ -25,15 +25,9 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) # -# The `__srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../../executorch_srcs.cmake" -) - -extract_sources(${EXECUTORCH_SRCS_FILE}) - -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() set(_common_include_directories $ diff --git a/extension/llm/custom_ops/custom_ops.py b/extension/llm/custom_ops/custom_ops.py index 947eae6c0d0..3c3243142cf 100644 --- a/extension/llm/custom_ops/custom_ops.py +++ b/extension/llm/custom_ops/custom_ops.py @@ -33,7 +33,7 @@ package_path = Path(__file__).parent.resolve() logging.info(f"Looking for libcustom_ops_aot_lib.so in {package_path}") - libs = list(package_path.glob("**/libcustom_ops_aot_lib.*")) + libs = list(package_path.glob("**/*custom_ops_aot_lib.*")) assert len(libs) == 1, f"Expected 1 library but got {len(libs)}" logging.info(f"Loading custom ops library: {libs[0]}") diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index 4bedd2d5ea9..8b8cc627afc 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -24,7 +24,7 @@ #endif #include -#include +#include namespace torch { namespace executor { diff --git a/extension/llm/custom_ops/op_tile_crop_aot.py b/extension/llm/custom_ops/op_tile_crop_aot.py index 701aabc441c..e03a979ffc5 100644 --- a/extension/llm/custom_ops/op_tile_crop_aot.py +++ b/extension/llm/custom_ops/op_tile_crop_aot.py @@ -13,7 +13,7 @@ tile_crop = torch.ops.preprocess.tile_crop.default assert tile_crop is not None except: - libs = list(Path(__file__).parent.resolve().glob("libcustom_ops_aot_lib.*")) + libs = list(Path(__file__).parent.resolve().glob("*custom_ops_aot_lib.*")) assert len(libs) == 1, f"Expected 1 library but got {len(libs)}" logging.info(f"Loading custom ops library: {libs[0]}") torch.ops.load_library(libs[0]) diff --git a/extension/llm/custom_ops/targets.bzl b/extension/llm/custom_ops/targets.bzl index 545f6516bb7..6c31b24171c 100644 --- a/extension/llm/custom_ops/targets.bzl +++ b/extension/llm/custom_ops/targets.bzl @@ -1,3 +1,4 @@ +load("@fbsource//xplat/executorch/build:build_variables.bzl", "EXTENSION_LLM_CUSTOM_OPS_BUCK_SRCS") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load( "@fbsource//xplat/executorch/kernels/optimized:lib_defs.bzl", @@ -13,7 +14,7 @@ def _get_quantized_sdpa_deps(): if runtime.is_oss: return [] else: - return ["//pytorch/ao/torchao/experimental/kernels/cpu/interface:interface"] + return ["//pytorch/ao/torchao/csrc/cpu/torch_free_kernels/interface:interface"] def _get_quantized_preproc_flags(): if runtime.is_oss: @@ -30,12 +31,7 @@ def define_common_targets(): for mkl_dep in ["", "_mkl_noomp"]: runtime.cxx_library( name = "custom_ops" + mkl_dep, - srcs = [ - "op_fallback.cpp", - "op_fast_hadamard_transform.cpp", - "op_sdpa.cpp", - "op_update_cache.cpp", - ], + srcs = EXTENSION_LLM_CUSTOM_OPS_BUCK_SRCS, exported_headers = [ "op_fallback.h", "op_fast_hadamard_transform.h", diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py index 6db881c5274..01000f3564c 100644 --- a/extension/llm/export/builder.py +++ b/extension/llm/export/builder.py @@ -34,7 +34,7 @@ from executorch.extension.llm.export.export_passes import RemoveRedundantTransposes from pytorch_tokenizers import get_tokenizer -from torch.export import export_for_training, ExportedProgram +from torch.export import export, ExportedProgram from torch.nn.attention import SDPBackend from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e from torchao.quantization.pt2e.quantizer import ComposableQuantizer, Quantizer @@ -96,6 +96,7 @@ def __init__( metadata: Optional[dict] = None, dynamic_shapes: Optional[Any] = None, save_exported_program: bool = False, + generate_etrecord: bool = False, ): # Store necessary constructor arguments. self.model = model @@ -113,9 +114,11 @@ def __init__( self.calibration_data = calibration_data self.tokenizer_path = tokenizer_path self.verbose = verbose - self.metadata = metadata + self.metadata = metadata if metadata is not None else {} + self.metadata["get_max_seq_len"] = max_seq_len self.dynamic_shapes = dynamic_shapes self.save_exported_program = save_exported_program + self.generate_etrecord = generate_etrecord # Note: treat this as the source of truth for the result of # torch.export'ing a model. If the overall ExportedProgram is needed, @@ -130,18 +133,20 @@ def __init__( self.output_dir = "." self._saved_pte_filename = None - def __post_init__(self): - """ - Post init function to update metadata based on dynamic shape - """ - dynamic_shape = self._get_dynamic_shape() - if dynamic_shape is not None: - token_dim = dynamic_shape[0][1] - if self.verbose: - logging.info( - f"Metadata 'get_max_seq_len' is being updated to match torch.export's dynamic shape max: {token_dim.max}" + # Try to resolve dynamic shapes if not specified explicitly. + if not self.dynamic_shapes and self.enable_dynamic_shape: + if not self.use_kv_cache: + # Only one input argument: tokens + # Here we -1 due to export limitation: https://gist.github.com/larryliu0820/419022a57e24d5e64150e325a685eaad + self.dynamic_shapes = ( + {1: torch.export.Dim("token_dim", max=self.max_seq_len - 1)}, + ) + else: + # Two input arguments: tokens and input_pos but input_pos is static shape + self.dynamic_shapes = ( + {1: torch.export.Dim("token_dim", max=self.max_seq_len)}, + {"input_pos": {0: 1}}, ) - self.metadata["get_max_seq_len"] = token_dim.max def set_output_dir(self, output_dir: str) -> "LLMEdgeManager": """ @@ -187,25 +192,6 @@ def source_transform( return self def _get_dynamic_shape(self) -> Any: - if self.dynamic_shapes: - return self.dynamic_shapes - - if self.enable_dynamic_shape: - if not self.use_kv_cache: - # Only one input argument: tokens - # Here we -1 due to export limitation: https://gist.github.com/larryliu0820/419022a57e24d5e64150e325a685eaad - self.dynamic_shapes = ( - {1: torch.export.Dim("token_dim", max=self.max_seq_len - 1)}, - ) - else: - # Two input arguments: tokens and input_pos but input_pos is static shape - self.dynamic_shapes = ( - {1: torch.export.Dim("token_dim", max=self.max_seq_len)}, - {"input_pos": {0: 1}}, - ) - else: - # Two input arguments: tokens and input_pos but both are of static shape - self.dynamic_shapes = None return self.dynamic_shapes def _get_edge_config(self) -> EdgeCompileConfig: @@ -232,7 +218,7 @@ def _export(self, module: Optional[torch.nn.Module] = None) -> ExportedProgram: logging.info(f"inputs: {self.example_inputs}") logging.info(f"kwargs: {self.example_kwarg_inputs}") logging.info(f"dynamic shapes: {dynamic_shape}") - exported_module = export_for_training( + exported_module = export( self.model if not module else module, self.example_inputs, kwargs=self.example_kwarg_inputs, @@ -244,7 +230,7 @@ def _export(self, module: Optional[torch.nn.Module] = None) -> ExportedProgram: def export(self) -> "LLMEdgeManager": """ Exports the model pre-autograd. This is not a full export, since it uses - torch.export_for_training() to keep autograd-safe ops from getting decomposed. + torch.export.export() to keep autograd-safe ops from getting decomposed. The full torch.export() if called later on during to_edge() or to_edge_transform_and_lower(). """ @@ -255,9 +241,7 @@ def export(self) -> "LLMEdgeManager": self.pre_autograd_graph_module = exported_module.module() if self.save_exported_program: export_output = f"{self.modelname}.pt2" - logging.info( - f"Saving torch.export()/export_for_training() result to {export_output}" - ) + logging.info(f"Saving torch.export() result to {export_output}") torch.export.save(exported_module, export_output) return self @@ -481,6 +465,7 @@ def to_edge_transform_and_lower( partitioner=partitioners, compile_config=edge_config, constant_methods=self.metadata, + generate_etrecord=self.generate_etrecord, ) if self.verbose: logging.info(f"Exported graph:\n{self.edge_manager.exported_program()}") diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index ab14a0b4a49..d756d1886ad 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -37,12 +37,16 @@ class ModelType(str, Enum): llama3_2 = "llama3_2" llama3_2_vision = "llama3_2_vision" static_llama = "static_llama" - qwen2_5 = "qwen2_5" + qwen2_5_0_5b = "qwen2_5_0_5b" + qwen2_5_1_5b = "qwen2_5_1_5b" qwen3_0_6b = "qwen3_0_6b" qwen3_1_7b = "qwen3_1_7b" qwen3_4b = "qwen3_4b" phi_4_mini = "phi_4_mini" smollm2 = "smollm2" + lfm2_350m = "lfm2_350m" + lfm2_700m = "lfm2_700m" + lfm2_1_2b = "lfm2_1_2b" class PreqMode(str, Enum): @@ -211,6 +215,9 @@ class ExportConfig: so_library: Shared library to specify custom quantized operators. export_only: Whether to stop right after torch.export() and just save the exported .pt2 graph file. + foundation_weights_file: configure the foundation weights of a model + to be placed in a separate file, external to the PTE. Pass the + intended file name here. """ max_seq_length: int = 128 @@ -219,6 +226,7 @@ class ExportConfig: output_name: Optional[str] = None so_library: Optional[str] = None export_only: bool = False + foundation_weights_file: Optional[str] = None def __post_init__(self): if self.max_context_length < self.max_seq_length: @@ -311,7 +319,12 @@ class QuantizationConfig: """ # Constants. - QMODE_OPTIONS: ClassVar[List[str]] = ["int8", "8da4w", "8da4w-gptq", "vulkan_4w"] + QMODE_OPTIONS: ClassVar[List[str]] = [ + "int8", + "8da4w", + "8da4w-gptq", + "4w", + ] AO_QUANT_PATTERNS: ClassVar[List[str]] = [ r"torchao:8da(\d+)w", r"torchao:fpa(\d+)w", @@ -413,6 +426,7 @@ class VulkanConfig: """ enabled: bool = False + force_fp16: bool = False @dataclass @@ -545,6 +559,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig": # noqa: C901 llm_config.export.so_library = args.so_library if hasattr(args, "export_only"): llm_config.export.export_only = args.export_only + if hasattr(args, "foundation_weights_file"): + llm_config.export.foundation_weights_file = args.foundation_weights_file # QuantizationConfig if hasattr(args, "quantization_mode"): @@ -595,6 +611,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig": # noqa: C901 # Vulkan if hasattr(args, "vulkan"): llm_config.backend.vulkan.enabled = args.vulkan + if hasattr(args, "vulkan_force_fp16"): + llm_config.backend.vulkan.force_fp16 = args.vulkan_force_fp16 # QNN if hasattr(args, "qnn"): diff --git a/extension/llm/export/partitioner_lib.py b/extension/llm/export/partitioner_lib.py index 7b093a7f1a3..5fe220f7dd9 100644 --- a/extension/llm/export/partitioner_lib.py +++ b/extension/llm/export/partitioner_lib.py @@ -32,7 +32,9 @@ def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True): def get_vulkan_partitioner( - dtype_override: Optional[str] = None, enable_dynamic_shape: bool = False + dtype_override: Optional[str] = None, + enable_dynamic_shape: bool = False, + force_fp16: bool = False, ): assert ( dtype_override == "fp32" or dtype_override is None @@ -41,7 +43,9 @@ def get_vulkan_partitioner( VulkanPartitioner, ) - return VulkanPartitioner({"require_dynamic_shapes": enable_dynamic_shape}) + return VulkanPartitioner( + {"require_dynamic_shapes": enable_dynamic_shape, "force_fp16": force_fp16} + ) def get_mps_partitioner(use_kv_cache: bool = False): diff --git a/extension/llm/export/quantizer_lib.py b/extension/llm/export/quantizer_lib.py index d87c722363f..2d87c86d113 100644 --- a/extension/llm/export/quantizer_lib.py +++ b/extension/llm/export/quantizer_lib.py @@ -108,7 +108,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 -DBUCK2= -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 . -j16\n" 'To find the location of the lib: find cmake-out -name "libquantized_ops_aot_lib*"\n' "Then specify the said library via -s _srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}. +# The `__srcs` lists are defined by executorch_load_build_variables. # -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../../executorch_srcs.cmake" -) - -extract_sources(${EXECUTORCH_SRCS_FILE}) - -include(${EXECUTORCH_SRCS_FILE}) +executorch_load_build_variables() # build llm runner library list(TRANSFORM _extension_llm_runner__srcs PREPEND "${EXECUTORCH_ROOT}/") add_library(extension_llm_runner STATIC ${_extension_llm_runner__srcs}) -set(runner_deps executorch_core extension_module extension_tensor tokenizers) +set(runner_deps executorch_core extension_module extension_tensor + tokenizers::tokenizers +) + +# depend on arange_utils +if(NOT TARGET kernels_util_all_deps) + add_subdirectory( + ${EXECUTORCH_ROOT}/kernels/portable/cpu/util + ${CMAKE_CURRENT_BINARY_DIR}/kernels_util + ) +endif() +list(APPEND runner_deps kernels_util_all_deps) target_link_libraries(extension_llm_runner PUBLIC ${runner_deps}) set_target_properties( diff --git a/extension/llm/runner/README.md b/extension/llm/runner/README.md new file mode 100644 index 00000000000..ab8ec8964dd --- /dev/null +++ b/extension/llm/runner/README.md @@ -0,0 +1,527 @@ +# LLM Runner Framework for ExecutorTorch + +This directory contains the LLM Runner framework for ExecutorTorch, providing high-level C++ APIs for running Large Language Models with both text-only and multimodal capabilities. + +## Overview + +The LLM Runner framework provides two main runner classes: + +- **TextLLMRunner**: For text-only language models (e.g., Llama, GPT, etc.) +- **MultimodalRunner**: For multimodal models that can process text, images, and audio (e.g., LLaVA, CLIP-based models) + +Both runners are built on a modular architecture with dependency injection, providing clean separation of concerns and efficient resource management. + +## Architecture Overview + +## MultimodalRunner Architecture + +The MultimodalRunner supports mixed inputs (text, images, audio) and generates text outputs: + +``` +MultimodalRunner Supported Model Architecture: +┌─────────────────────────────────────────────────────────────────────────┐ +│ Multimodal LLM Architecture │ +└─────────────────────────────────────────────────────────────────────────┘ + Input: std::vector + ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ + │ Image │ │ Audio │ │ Text │ + │ [224x │ │ [16kHz │ │ "What" │ + │ 224x3] │ │ audio] │ │ │ + └─────────────────┘ └─────────────────┘ └─────────────────┘ + │ │ │ + ▼ ▼ ▼ + ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ ◄─┐ + │ Encoder │ │ Encoder │ │ Text Tokenizer │ │ + │ (Vision) │ │ (Audio) │ │ & Embedding │ │ + │ │ │ │ │ │ │ + │ pixels → embed │ │ waveform→embed │ │ tokens → embed │ │ + └─────────────────┘ └─────────────────┘ └─────────────────┘ │ + │ │ │ │ + ▼ ▼ ▼ │ + ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │ + │ [D_emb] │ │ [D_emb] │ │ [D_emb] │ │ + │ Embedding │ │ Embedding │ │ Embedding │ │ + └─────────────────┘ └─────────────────┘ └─────────────────┘ │ + │ │ │ │ + └────────────────────┼────────────────────┘ │ + │ │ + ▼ │ + ┌─────────────────────────────┐ │ + │ Text Decoder Block │ │ + │ (Transformer Layers) │ │ + │ │ │ + │ ┌─────────────────────┐ │ │ + │ │ Self-Attention │ │ │ + │ │ + Feed Forward │ │ │ + │ │ (with KV Cache) │ │ │ + │ └─────────────────────┘ │ │ + │ │ │ │ + │ ▼ │ │ + │ Token Generation │ │ + │ (pos_ tracking) │ │ + └─────────────────────────────┘ │ + │───────────────────────────────────────┘ + │ (Autoregressive) + ▼ + ┌─────────────────┐ + │ Generated Text │ + │ "This image │ + │ shows a cat │ + │ sitting..." │ + └─────────────────┘ +``` + +## Key Features + +### TextLLMRunner +- **Text-only processing**: Optimized for pure language models +- **Efficient tokenization**: Support for multiple tokenizer formats +- **KV cache management**: Automatic position tracking for efficient inference +- **Streaming generation**: Token-by-token callbacks for real-time output +- **Configuration-driven**: Comprehensive control via `GenerationConfig` + +### MultimodalRunner +- **Mixed input support**: Process text, images, and audio in any order +- **Type-safe inputs**: `MultimodalInput` class with compile-time type checking +- **Modular encoders**: Separate processing pipelines for different modalities +- **Unified generation**: Single API for complex multimodal workflows +- **Extensible design**: Easy to add support for new modalities + +## Quick Start + +### TextLLMRunner Example + +```cpp +#include +#include + +int main() { + // Load tokenizer and create runner + auto tokenizer = load_tokenizer("tokenizer.bin"); + auto runner = create_text_llm_runner("model.pte", std::move(tokenizer)); + + // Configure generation + GenerationConfig config; + config.max_new_tokens = 100; + config.temperature = 0.7f; + config.echo = true; + + // Set up callbacks + auto token_callback = [](const std::string& token) { + std::cout << token << std::flush; + }; + + // Generate text + auto error = runner->generate( + "Hello, how are you?", // prompt + config, // configuration + token_callback // token callback + ); + + return error == executorch::runtime::Error::Ok ? 0 : 1; +} +``` + +### MultimodalRunner Example + +```cpp +#include +#include +#include + +int main() { + // Load tokenizer and create runner + auto tokenizer = load_tokenizer("tokenizer.bin"); + auto runner = create_multimodal_runner("model.pte", std::move(tokenizer)); + + // Create multimodal inputs + std::vector inputs; + inputs.emplace_back(make_text_input("What do you see in this image?")); + + // Load and add image + Image image = load_image("photo.jpg"); // Your image loading function + inputs.emplace_back(make_image_input(std::move(image))); + + // Configure generation + GenerationConfig config; + config.max_new_tokens = 150; + config.temperature = 0.7f; + config.echo = true; + + // Set up callbacks + auto token_callback = [](const std::string& token) { + std::cout << token << std::flush; + }; + + auto stats_callback = [](const Stats& stats) { + std::cout << "\nGenerated " << stats.num_generated_tokens << " tokens" << std::endl; + }; + + // Generate text + auto error = runner->generate(inputs, config, token_callback, stats_callback); + + return error == executorch::runtime::Error::Ok ? 0 : 1; +} +``` + +## Core Components + +### Component Architecture + +``` + + ┌─────────────────┐ + │ IRunner │ + │ <> │ + │ │ + │ + is_loaded() │ + │ + load() │ + │ + generate() │ + │ + stop() │ + └─────────────────┘ + △ + │ + │ implements + │ + │ + │ + │ + ┌──────┴──────────┐ ┌─────────────────┐ + │ TextLLMRunner │ │MultimodalRunner │ + │ │ │ │ + │ - tokenizer_ │ │ - tokenizer_ │ + ┌─────┼ - module_ │ │ - module_ ┼─────┐ + │ ┌───┼ - stats_ │ │ - stats_ ┼───┐ │ + │ │ ┌─┼ - metadata_ │ │ - metadata_ ┼─┐ │ │ + │ │ │ │ - temperature_ │ │ - pos_ │ │ │ │ + │ │ │ └─────────────────┘ └─────────────────┘ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ │ │ │ + │ │ │ ┌─────────────────┐ │ │ │ + │ │ │ │TextTokenGenerat-│ │ │ │ + │ │ │ │or │ │ │ │ + │ │ │ │ │ │ │ │ + │ │ │ │ - tokenizer_* │ │ │ │ + │ │ │ consists │ - text_decoder_ │ consists │ │ │ + │ │ └──────────────►│ runner_ │◄───────────────┘ │ │ + │ │ │ - eos_ids_ │ │ │ + │ │ │ - use_kv_cache_ │ │ │ + │ │ │ - stats_* │ │ │ + │ │ │ │ │ │ + │ │consists │ + generate() │ consists │ │ + │ │ └────────┬────────┘ │ │ + │ │ ┌──────────────┴───────────────┐ │ │ + │ │ ▼ uses ▼ │ │ + │ │ ┌─────────────────┐ ┌─────────────────┐ │ │ + │ │ │TextDecoderRunner│ │MultimodalDecode-│ │ │ + │ │ │ │ │rRunner │ │ │ + │ │ │ - module_* │ extends │ - module_* │ │ │ + │ └──►│ - should_stop_ │◄─────────┼ - should_stop_ │◄──┘ │ + │ │ │ │ │ │ + │ │ + step() │ │ + step() │ │ + │ │ + logits_to_ │ │ + logits_to_ │ │ + │ │ token() │ │ token() │ │ + │ └─────────────────┘ └─────────────────┘ │ + │ ▲ ▲ │ + │ │ uses │ │ + │consists ├─────────────────────────────┤ │ + │ ┌───────┴─────────┐ │ │ + │ │ TextPrefiller │ │ consists│ + │ │ │ ┌────────┴────────┐ │ + │ │ - text_decoder_ │ │ MultimodalPrefi-│ │ + │ │ runner_ │ │ller │ │ + └────►│ - use_kv_cache_ │ │ - module_* │ │ + │ - enable_ │ │ │◄────┘ + │ parallel_ │ │ + prefill() │ + │ prefill_ │ │ + logits_to_ │ + │ │ │ token() │ + │ + prefill() │ └─────────────────┘ + ├─────────────────┘ +``` + +### 1. Tokenizer +**Purpose**: Converts between text and token IDs + +**Supported Formats**: +- HF JSON (Hugging Face tokenizer format) +- TikToken (OpenAI's tokenizer format) +- SentencePiece (Google's tokenizer format) +- BPE (Byte-pair encoding tokenizer) + +**Key Methods**: +```cpp +virtual Result> encode(const std::string& text, int8_t bos = 1, int8_t eos = 0) = 0; +virtual Result decode(uint64_t prev_token, uint64_t token) = 0; +virtual uint64_t bos_tok() const = 0; +virtual uint64_t eos_tok() const = 0; +``` + +### 2. TextDecoderRunner +**Purpose**: Executes the transformer decoder part of the model + +**Key Responsibilities**: +- Executes transformer decoder layers +- Manages KV cache during inference +- Handles both prefill and decode phases +- Provides low-level model execution interface + +### 3. TextPrefiller +**Purpose**: Handles the prefill phase for text inputs + +**Key Features**: +- Parallel token processing for efficiency +- KV cache management +- Batch processing support +- Integration with tokenizer + +**Configuration**: +```cpp +TextPrefiller( + TextDecoderRunner* text_decoder_runner, + bool use_kv_cache, + bool enable_parallel_prefill, + int64_t max_seq_len +); +``` + +### 4. ImagePrefiller (MultimodalRunner only) +**Purpose**: Processes image inputs through vision encoders + +**Key Features**: +- Vision encoder integration +- Pixel data to embedding conversion +- Multiple image format support +- KV cache integration + +**Image Format**: +```cpp +struct Image { + int32_t width; + int32_t height; + int32_t channels; + std::vector data; // Raw pixel data +}; +``` + +### 5. TextTokenGenerator +**Purpose**: Handles autoregressive token generation + +**Key Features**: +- Temperature-based sampling +- EOS token detection +- Token-by-token callbacks +- Performance statistics tracking + +**Usage**: +```cpp +int64_t num_tokens = text_token_generator->generate( + {start_token}, // Initial tokens + current_pos, // Starting position + max_new_tokens, // Maximum tokens to generate + temperature, // Sampling temperature + token_callback // Callback for each token +); +``` + +### 6. GenerationConfig +**Purpose**: Comprehensive configuration for text generation + +**Key Parameters**: +```cpp +struct GenerationConfig { + int32_t max_new_tokens = -1; // Max tokens to generate (-1 = use available) + int32_t seq_len = 1024; // Total sequence length + float temperature = 0.8f; // Sampling temperature + bool echo = true; // Echo input prompt + int8_t num_bos = 1; // Number of BOS tokens + int8_t num_eos = 1; // Number of EOS tokens + bool warming = false; // Warmup run flag +}; +``` + +### 7. MultimodalInput (MultimodalRunner only) +**Purpose**: Type-safe wrapper for mixed input types + +**Key Features**: +- `std::variant` internally +- Type-safe access methods +- Exception-based and safe access patterns +- Move semantics for efficiency + +**API**: +```cpp +// Type checking +bool is_text() const; +bool is_image() const; + +// Direct access (throws on type mismatch) +const std::string& get_text() const; +const Image& get_image() const; + +// Safe access (returns nullptr on type mismatch) +const std::string* try_get_text() const; +const Image* try_get_image() const; + +// Factory functions +MultimodalInput make_text_input(const std::string& text); +MultimodalInput make_image_input(Image&& image); +``` + +## Helper Functions + +The framework provides utility functions in `llm_runner_helper.h`: + +### load_tokenizer() +```cpp +std::unique_ptr load_tokenizer( + const std::string& tokenizer_path, + std::unique_ptr> special_tokens = nullptr, + std::optional pattern = std::nullopt, + size_t bos_token_index = 0, + size_t eos_token_index = 1 +); +``` + +### create_text_llm_runner() +```cpp +std::unique_ptr create_text_llm_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path = std::nullopt, + float temperature = -1.0f +); +``` + +### create_multimodal_runner() +```cpp +std::unique_ptr create_multimodal_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path = std::nullopt, + float temperature = 0.8f +); +``` + +### get_llm_metadata() +```cpp +std::unordered_map get_llm_metadata( + tokenizers::Tokenizer* tokenizer, + Module* module +); +``` + +## Configuration and Tuning + +### Generation Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `max_new_tokens` | `int32_t` | `-1` | Maximum new tokens to generate (-1 = use available context) | +| `seq_len` | `int32_t` | `1024` | Total sequence length including prompt | +| `temperature` | `float` | `0.8f` | Sampling temperature (0.0 = deterministic, 1.0+ = creative) | +| `echo` | `bool` | `true` | Whether to echo the input prompt | +| `num_bos` | `int8_t` | `1` | Number of beginning-of-sequence tokens | +| `num_eos` | `int8_t` | `1` | Number of end-of-sequence tokens | +| `warming` | `bool` | `false` | Whether this is a warmup run | + +### Performance Tuning + +**Memory Optimization**: +- Use KV cache for efficient autoregressive generation +- Enable parallel prefill for faster prompt processing +- Set appropriate `seq_len` based on available memory + +**Sampling Strategies**: +- Low temperature (0.1-0.3) for factual, deterministic output +- High temperature (0.7-1.0) for creative, diverse output +- Set `max_new_tokens` to prevent runaway generation + +**Monitoring**: +```cpp +auto stats_callback = [](const Stats& stats) { + std::cout << "Model load time: " + << (stats.model_load_end_ms - stats.model_load_start_ms) << "ms" << std::endl; + std::cout << "Inference time: " + << (stats.inference_end_ms - stats.inference_start_ms) << "ms" << std::endl; + std::cout << "Tokens/second: " << stats.tokens_per_second() << std::endl; +}; +``` + +## Supported Models + +### TextLLMRunner +- **Llama family**: Llama 2, Llama 3, Code Llama +- **GPT models**: GPT-2, GPT-3.5, GPT-4 (compatible architectures) +- **Phi models**: Phi-3-mini and variants +- **Custom models**: Any transformer-based text generation model + +### MultimodalRunner + +**Note**: The MultimodalRunner currently supports **EarlyFusion** model architectures only. EarlyFusion is a type of fused model architecture where pretrained encoder(s) are combined with a pretrained decoder (LLM) at the model input and not in internal layers. This is a popular architecture for multimodal models, with a full overview available in [The Evolution of Multimodal Model Architectures](https://arxiv.org/abs/2405.17927). This module works both for decoders in which the encoder tokens are inside the vocab and outside the vocab. + +**Supported EarlyFusion Models**: +- **LLaVA**: Large Language and Vision Assistant +- **CLIP-based models**: Contrastive Language-Image Pre-training +- **Gemma3 4B**: Multimodal variant with vision capabilities +- **Voxtral**: Audio-text multimodal models +- **Custom EarlyFusion models**: Any model with separate encoders that fuse at the input level + +**DeepFusion Models (Not Currently Supported)**: +DeepFusion is another popular model architecture type where a pretrained encoder is combined with a pretrained decoder (LLM) in the internal decoder layers. A common deep fusion architecture is to fuse the encoder input into the decoder with interspersed cross-attention layers. DeepFusion models are currently out of scope because they require significant model definition rewrites to work with torch.export. + +**Examples of DeepFusion models (not supported)**: +- **Llama 3.2 Vision**: Uses cross-attention layers for vision-text fusion +- **Other cross-attention based multimodal models** + +For DeepFusion support, consider using the model's native inference framework or wait for future ExecutorTorch updates that may include DeepFusion architecture support. + +## Building and Integration + +### CMake Integration +```cmake +find_package(executorch REQUIRED) +target_link_libraries(your_target + executorch::extension_llm_runner + executorch::extension_module +) +``` + +### Required Headers +```cpp +// For TextLLMRunner +#include + +// For MultimodalRunner +#include +#include + +// Helper functions +#include + +// Configuration +#include +``` + +## Advanced Usage + +### Custom Sampling +```cpp +// Custom temperature per generation +GenerationConfig config; +config.temperature = 0.1f; // Very deterministic +runner->generate(factual_prompt, config, callback); + +config.temperature = 1.2f; // Very creative +runner->generate(creative_prompt, config, callback); +``` + +### Memory Monitoring +```cpp +#include + +auto stats_callback = [](const Stats& stats) { + double rss_mb = get_rss_bytes() / 1024.0 / 1024.0; + std::cout << "RSS: " << rss_mb << " MiB" << std::endl; +}; +``` diff --git a/extension/llm/runner/audio.h b/extension/llm/runner/audio.h new file mode 100644 index 00000000000..868765950af --- /dev/null +++ b/extension/llm/runner/audio.h @@ -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. + */ + +// A simple audio struct. + +#pragma once +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +/** + * Audio inputs as a raw audio tensor, for use when the audio processing + * into a mel spectrogram is baked into the audio encoder with torch.export. + */ +struct ET_EXPERIMENTAL RawAudio { + std::vector data; + int32_t batch_size; + int32_t n_channels; // For mono, use n_channels = 1. + int32_t n_samples; +}; + +/** + * Pre-processed audio inputs, ready to feed directly into an audio + * encoder. + */ +struct ET_EXPERIMENTAL Audio { + std::vector data; + int32_t batch_size; + int32_t n_bins; + int32_t n_frames; +}; + +} // namespace llm +} // namespace extension +} // namespace executorch + +namespace torch { +namespace executor { +// TODO(T197294990): Remove these deprecated aliases once all users have moved +// to the new `::executorch` namespaces. +using ::executorch::extension::llm::Audio; +} // namespace executor +} // namespace torch diff --git a/extension/llm/runner/constants.h b/extension/llm/runner/constants.h index fc6ddcb451c..4ba88203c50 100644 --- a/extension/llm/runner/constants.h +++ b/extension/llm/runner/constants.h @@ -21,7 +21,8 @@ inline constexpr auto kUseSDPAWithKVCache = "use_sdpa_with_kv_cache"; // Multimodal method name conventions inline constexpr auto kImageEncoderMethod = "image_encoder"; +inline constexpr auto kAudioEncoderMethod = "audio_encoder"; inline constexpr auto kTokenEmbeddingMethod = "token_embedding"; -inline constexpr auto kTextModelMethod = "text_model"; +inline constexpr auto kTextModelMethod = "text_decoder"; } // namespace executorch::extension::llm diff --git a/extension/llm/runner/io_manager/io_manager.h b/extension/llm/runner/io_manager/io_manager.h index ce158c23b6e..fc9a8f0641b 100644 --- a/extension/llm/runner/io_manager/io_manager.h +++ b/extension/llm/runner/io_manager/io_manager.h @@ -8,12 +8,8 @@ #pragma once -#include - +#include #include -#include -#include -#include namespace executorch { namespace extension { @@ -29,6 +25,13 @@ namespace llm { */ class ET_EXPERIMENTAL IOManager { public: + /** + * @brief Construct an IOManager bound to a Module. + * + * @param module The Module used for querying method metadata and execution. + */ + explicit IOManager(ET_MODULE_NAMESPACE::Module& module) : module_(module) {} + /** * @brief Virtual destructor to allow proper cleanup in derived classes. */ @@ -38,20 +41,28 @@ class ET_EXPERIMENTAL IOManager { * @brief Load the IO manager with method metadata for prefill and * decode operations. * - * @param program The program prefill and decode methods are loaded from. * @param prefill_method The prefill method to initialize with. * @param decode_method The decode method to initialize with. */ ET_NODISCARD virtual runtime::Error load( - const executorch::ET_RUNTIME_NAMESPACE::Program& program, - executorch::ET_RUNTIME_NAMESPACE::Method& prefill_method, - executorch::ET_RUNTIME_NAMESPACE::Method& decode_method) { - (void)program; + const std::string& prefill_method, + const std::string& decode_method) { (void)prefill_method; (void)decode_method; return runtime::Error::Ok; } + /** + * @brief Load the IO manager using the default method names. + * + * Uses "forward" for both prefill and decode. + * + * @return Error code. + */ + ET_NODISCARD runtime::Error load() { + return load("forward", "forward"); + } + /** * @brief Reset the IO manager state. * @@ -59,13 +70,24 @@ class ET_EXPERIMENTAL IOManager { * @param decode_method The decode method to reset with. */ ET_NODISCARD virtual runtime::Error reset( - executorch::ET_RUNTIME_NAMESPACE::Method& prefill_method, - executorch::ET_RUNTIME_NAMESPACE::Method& decode_method) { + const std::string& prefill_method, + const std::string& decode_method) { (void)prefill_method; (void)decode_method; return runtime::Error::Ok; } + /** + * @brief Reset the IO manager state using the default method names. + * + * Uses "forward" for both prefill and decode. + * + * @return Error code. + */ + ET_NODISCARD runtime::Error reset() { + return reset("forward", "forward"); + } + /** * @brief Prepare inputs for the prefill phase of LLM inference. * @@ -73,19 +95,22 @@ class ET_EXPERIMENTAL IOManager { * @param start_pos The tensor containing the starting position of the current * input within the context. * @param prefill_method The prefill method to prepare inputs for. - * @return std::vector Vector of prepared inputs + * @return std::vector Vector of prepared inputs * for the prefill method. */ - virtual runtime::Result> - prepare_prefill( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, - executorch::ET_RUNTIME_NAMESPACE::Method& prefill_method) { - if (prefill_method.inputs_size() != 2) { + virtual runtime::Result> prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& prefill_method) { + auto method_meta = module_.method_meta(prefill_method); + if (!method_meta.ok()) { + return method_meta.error(); + } + if (method_meta->num_inputs() != 2) { ET_LOG( Error, "Expected 2 inputs for prefill method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", - prefill_method.inputs_size()); + method_meta->num_inputs()); return runtime::Error::InvalidState; } // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done @@ -93,6 +118,21 @@ class ET_EXPERIMENTAL IOManager { return std::vector{input, start_pos}; } + /** + * @brief Prepare inputs for the prefill phase using the default method name. + * + * Uses "forward" as the prefill method. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position. + * @return Vector of prepared inputs for the prefill method. + */ + runtime::Result> prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos) { + return prepare_prefill(input, start_pos, "forward"); + } + /** * @brief Prepare inputs for the decode phase of LLM inference. * @@ -100,19 +140,22 @@ class ET_EXPERIMENTAL IOManager { * @param start_pos The tensor containing the starting position of the current * input within the context. * @param decode_method The decode method to prepare inputs for. - * @return std::vector Vector of prepared inputs + * @return std::vector Vector of prepared inputs * for the decode method. */ - virtual runtime::Result> - prepare_decode( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, - executorch::ET_RUNTIME_NAMESPACE::Method& decode_method) { - if (decode_method.inputs_size() != 2) { + virtual runtime::Result> prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& decode_method) { + auto method_meta = module_.method_meta(decode_method); + if (!method_meta.ok()) { + return method_meta.error(); + } + if (method_meta->num_inputs() != 2) { ET_LOG( Error, "Expected 2 inputs for decode method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", - decode_method.inputs_size()); + method_meta->num_inputs()); return runtime::Error::InvalidState; } // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done @@ -120,6 +163,21 @@ class ET_EXPERIMENTAL IOManager { return std::vector{input, start_pos}; } + /** + * @brief Prepare inputs for the decode phase using the default method name. + * + * Uses "forward" as the decode method. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position. + * @return Vector of prepared inputs for the decode method. + */ + runtime::Result> prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos) { + return prepare_decode(input, start_pos, "forward"); + } + /** * @brief Process and update internal state with outputs from the prefill * phase. @@ -128,14 +186,27 @@ class ET_EXPERIMENTAL IOManager { * @param model_outputs Vector of outputs from the prefill method execution. */ ET_NODISCARD virtual runtime::Error update_prefill( - executorch::ET_RUNTIME_NAMESPACE::Method& prefill_method, - const std::vector& model_outputs) { - (void)prefill_method; + const std::vector& model_outputs, + const std::string& prefill_method) { (void)model_outputs; + (void)prefill_method; // No post inference work to do. return runtime::Error::Ok; } + /** + * @brief Process outputs from the prefill phase using the default method. + * + * Uses "forward" as the prefill method. + * + * @param model_outputs Vector of outputs from the prefill execution. + * @return Error code. + */ + ET_NODISCARD runtime::Error update_prefill( + const std::vector& model_outputs) { + return update_prefill(model_outputs, "forward"); + } + /** * @brief Process and update internal state with outputs from the decode * phase. @@ -144,13 +215,32 @@ class ET_EXPERIMENTAL IOManager { * @param model_outputs Vector of outputs from the decode method execution. */ ET_NODISCARD virtual runtime::Error update_decode( - const executorch::ET_RUNTIME_NAMESPACE::Method& decode_method, - const std::vector& model_outputs) { - (void)decode_method; + const std::vector& model_outputs, + const std::string& decode_method) { (void)model_outputs; + (void)decode_method; // No post inference work to do. return runtime::Error::Ok; } + + /** + * @brief Process outputs from the decode phase using the default method. + * + * Uses "forward" as the decode method. + * + * @param model_outputs Vector of outputs from the decode execution. + * @return Error code. + */ + ET_NODISCARD runtime::Error update_decode( + const std::vector& model_outputs) { + return update_decode(model_outputs, "forward"); + } + + private: + /** + * @brief Reference to the Module used for method metadata and execution. + */ + ET_MODULE_NAMESPACE::Module& module_; }; } // namespace llm diff --git a/extension/llm/runner/io_manager/targets.bzl b/extension/llm/runner/io_manager/targets.bzl index e538572c51b..5b891b24376 100644 --- a/extension/llm/runner/io_manager/targets.bzl +++ b/extension/llm/runner/io_manager/targets.bzl @@ -1,8 +1,8 @@ -load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "get_aten_mode_options", "runtime") def define_common_targets(): - for aten in (True, False): + for aten in get_aten_mode_options(): aten_suffix = "_aten" if aten else "" # Interface for IOManager. No concrete impl from this dep. @@ -11,10 +11,9 @@ def define_common_targets(): exported_headers = [ "io_manager.h", ], - deps = [ + exported_deps = [ "//executorch/extension/tensor:tensor" + aten_suffix, - "//executorch/runtime/core/exec_aten:lib" + aten_suffix, - "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix, + "//executorch/extension/module:module" + aten_suffix, ], visibility = [ "@EXECUTORCH_CLIENTS", diff --git a/extension/llm/runner/io_manager/test/TARGETS b/extension/llm/runner/io_manager/test/TARGETS index 6db0a7c590b..e214060942a 100644 --- a/extension/llm/runner/io_manager/test/TARGETS +++ b/extension/llm/runner/io_manager/test/TARGETS @@ -10,14 +10,12 @@ define_common_targets() runtime.cxx_test( name = "test_io_manager", - srcs = ["test_io_manager.cpp"], + srcs = [ + "test_io_manager.cpp", + ], deps = [ "//executorch/extension/llm/runner/io_manager:io_manager", - "//executorch/extension/llm/runner/io_manager:io_manager", - "//executorch/extension/module:module", - "//executorch/extension/tensor:tensor", - "//executorch/runtime/executor:program", - "//executorch/kernels/portable:generated_lib", + "//executorch/kernels/portable:generated_lib", ], env = { "KVCACHE_CACHE_POS": "$(location fbcode//executorch/test/models:exported_programs[ModuleKVCacheCachePos.pte])", diff --git a/extension/llm/runner/io_manager/test/test_io_manager.cpp b/extension/llm/runner/io_manager/test/test_io_manager.cpp index bc265e8d083..7c31ff9ea18 100644 --- a/extension/llm/runner/io_manager/test/test_io_manager.cpp +++ b/extension/llm/runner/io_manager/test/test_io_manager.cpp @@ -7,74 +7,45 @@ */ #include -#include -#include -#include -#include + #include using namespace ::testing; -using executorch::extension::Module; -using executorch::extension::llm::IOManager; -using executorch::runtime::Error; -using executorch::runtime::EValue; -using executorch::runtime::Method; -using executorch::runtime::Program; -using executorch::runtime::Result; +using namespace ::executorch::extension; +using namespace ::executorch::runtime; // Test fixture for IOManager tests class IOManagerTest : public Test { protected: void SetUp() override { - executorch::runtime::runtime_init(); - module_ = std::make_unique(std::getenv("KVCACHE_CACHE_POS")); - io_manager_ = std::make_unique(); - auto err = module_->load_method("forward"); - EXPECT_EQ(err, Error::Ok); + io_manager_ = std::make_unique(*module_); + EXPECT_EQ(module_->load_forward(), Error::Ok); } protected: std::unique_ptr module_; - - std::unique_ptr io_manager_; + std::unique_ptr io_manager_; }; // Test that load() returns Error::Ok (no-op) TEST_F(IOManagerTest, LoadReturnsOk) { - auto* program = module_->program().get(); - auto* prefill_method = module_->method("forward").get(); - auto* decode_method = module_->method("forward").get(); - - auto result = io_manager_->load(*program, *prefill_method, *decode_method); - - EXPECT_EQ(result, Error::Ok); + EXPECT_EQ(io_manager_->load(), Error::Ok); } // Test that reset() returns Error::Ok (no-op) TEST_F(IOManagerTest, ResetReturnsOk) { - auto* prefill_method = module_->method("forward").get(); - auto* decode_method = module_->method("forward").get(); - - auto result = io_manager_->reset(*prefill_method, *decode_method); - - EXPECT_EQ(result, Error::Ok); + EXPECT_EQ(io_manager_->reset(), Error::Ok); } // Test that prepare_prefill() returns the input tensors when method has 2 // inputs TEST_F(IOManagerTest, PreparePrefillReturnsInputsWhenValidInputCount) { - auto* prefill_method = module_->method("forward").get(); - - // Create test tensors std::vector input_data = {1.0f, 2.0f, 3.0f, 4.0f}; std::vector start_pos_data = {0}; - auto input_ptr = executorch::extension::make_tensor_ptr({1, 4}, input_data); - auto start_pos_ptr = - executorch::extension::make_tensor_ptr({1}, start_pos_data); - - auto result = - io_manager_->prepare_prefill(input_ptr, start_pos_ptr, *prefill_method); + auto input_ptr = make_tensor_ptr({1, 4}, input_data); + auto start_pos_ptr = make_tensor_ptr({1}, start_pos_data); + auto result = io_manager_->prepare_prefill(input_ptr, start_pos_ptr); EXPECT_EQ(result.error(), Error::Ok); auto outputs = result.get(); @@ -87,17 +58,12 @@ TEST_F(IOManagerTest, PreparePrefillReturnsInputsWhenValidInputCount) { // Test that prepare_decode() returns the input tensors when method has 2 inputs TEST_F(IOManagerTest, PrepareDecodeReturnsInputsWhenValidInputCount) { - auto* decode_method = module_->method("forward").get(); - - // Create test tensors std::vector input_data = {5.0f, 6.0f, 7.0f, 8.0f}; std::vector start_pos_data = {10}; - auto input_ptr = executorch::extension::make_tensor_ptr({1, 4}, input_data); - auto start_pos_ptr = - executorch::extension::make_tensor_ptr({1}, start_pos_data); + auto input_ptr = make_tensor_ptr({1, 4}, input_data); + auto start_pos_ptr = make_tensor_ptr({1}, start_pos_data); - auto result = - io_manager_->prepare_decode(input_ptr, start_pos_ptr, *decode_method); + auto result = io_manager_->prepare_decode(input_ptr, start_pos_ptr); EXPECT_EQ(result.error(), Error::Ok); auto outputs = result.get(); @@ -110,49 +76,31 @@ TEST_F(IOManagerTest, PrepareDecodeReturnsInputsWhenValidInputCount) { // Test that update_prefill() returns Error::Ok (no-op) TEST_F(IOManagerTest, UpdatePrefillReturnsOk) { - auto* prefill_method = module_->method("forward").get(); - - // Create dummy model outputs std::vector model_outputs; std::vector output_data = {0.1f, 0.2f, 0.3f}; - auto output_tensor = - executorch::extension::make_tensor_ptr({1, 3}, output_data); + auto output_tensor = make_tensor_ptr({1, 3}, output_data); model_outputs.emplace_back(*output_tensor); - auto result = io_manager_->update_prefill(*prefill_method, model_outputs); - - EXPECT_EQ(result, Error::Ok); + EXPECT_EQ(io_manager_->update_prefill(model_outputs), Error::Ok); } // Test that update_decode() returns Error::Ok (no-op) TEST_F(IOManagerTest, UpdateDecodeReturnsOk) { - auto* decode_method = module_->method("forward").get(); - - // Create dummy model outputs std::vector model_outputs; std::vector output_data = {0.4f, 0.5f, 0.6f}; - auto output_tensor = - executorch::extension::make_tensor_ptr({1, 3}, output_data); + auto output_tensor = make_tensor_ptr({1, 3}, output_data); model_outputs.emplace_back(*output_tensor); - auto result = io_manager_->update_decode(*decode_method, model_outputs); - - EXPECT_EQ(result, Error::Ok); + EXPECT_EQ(io_manager_->update_decode(model_outputs), Error::Ok); } // Test that prepare_prefill() correctly passes through different tensor shapes TEST_F(IOManagerTest, PreparePrefillPassesThroughDifferentTensorShapes) { - auto* prefill_method = module_->method("forward").get(); - - // Create test tensors with different shapes std::vector input_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; std::vector start_pos_data = {5, 10}; - auto input_ptr = executorch::extension::make_tensor_ptr({2, 3}, input_data); - auto start_pos_ptr = - executorch::extension::make_tensor_ptr({2}, start_pos_data); - - auto result = - io_manager_->prepare_prefill(input_ptr, start_pos_ptr, *prefill_method); + auto input_ptr = make_tensor_ptr({2, 3}, input_data); + auto start_pos_ptr = make_tensor_ptr({2}, start_pos_data); + auto result = io_manager_->prepare_prefill(input_ptr, start_pos_ptr); EXPECT_EQ(result.error(), Error::Ok); auto outputs = result.get(); @@ -165,18 +113,12 @@ TEST_F(IOManagerTest, PreparePrefillPassesThroughDifferentTensorShapes) { // Test that prepare_decode() correctly passes through different tensor shapes TEST_F(IOManagerTest, PrepareDecodePassesThroughDifferentTensorShapes) { - auto* decode_method = module_->method("forward").get(); - - // Create test tensors with different shapes std::vector input_data = { 7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f}; std::vector start_pos_data = {15, 20, 25}; - auto input_ptr = executorch::extension::make_tensor_ptr({2, 4}, input_data); - auto start_pos_ptr = - executorch::extension::make_tensor_ptr({3}, start_pos_data); - - auto result = - io_manager_->prepare_decode(input_ptr, start_pos_ptr, *decode_method); + auto input_ptr = make_tensor_ptr({2, 4}, input_data); + auto start_pos_ptr = make_tensor_ptr({3}, start_pos_data); + auto result = io_manager_->prepare_decode(input_ptr, start_pos_ptr); EXPECT_EQ(result.error(), Error::Ok); auto outputs = result.get(); @@ -189,42 +131,22 @@ TEST_F(IOManagerTest, PrepareDecodePassesThroughDifferentTensorShapes) { // Test that update methods handle empty model outputs TEST_F(IOManagerTest, UpdateMethodsHandleEmptyModelOutputs) { - auto* prefill_method = module_->method("forward").get(); - auto* decode_method = module_->method("forward").get(); - - // Create empty model outputs std::vector empty_outputs; - auto prefill_result = - io_manager_->update_prefill(*prefill_method, empty_outputs); - auto decode_result = - io_manager_->update_decode(*decode_method, empty_outputs); - - EXPECT_EQ(prefill_result, Error::Ok); - EXPECT_EQ(decode_result, Error::Ok); + EXPECT_EQ(io_manager_->update_prefill(empty_outputs), Error::Ok); + EXPECT_EQ(io_manager_->update_decode(empty_outputs), Error::Ok); } // Test that update methods handle multiple model outputs TEST_F(IOManagerTest, UpdateMethodsHandleMultipleModelOutputs) { - auto* prefill_method = module_->method("forward").get(); - auto* decode_method = module_->method("forward").get(); - - // Create multiple model outputs std::vector model_outputs; std::vector output1_data = {0.1f, 0.2f}; std::vector output2_data = {0.3f, 0.4f, 0.5f}; - auto output1_tensor = - executorch::extension::make_tensor_ptr({1, 2}, output1_data); - auto output2_tensor = - executorch::extension::make_tensor_ptr({1, 3}, output2_data); + auto output1_tensor = make_tensor_ptr({1, 2}, output1_data); + auto output2_tensor = make_tensor_ptr({1, 3}, output2_data); model_outputs.emplace_back(*output1_tensor); model_outputs.emplace_back(*output2_tensor); - auto prefill_result = - io_manager_->update_prefill(*prefill_method, model_outputs); - auto decode_result = - io_manager_->update_decode(*decode_method, model_outputs); - - EXPECT_EQ(prefill_result, Error::Ok); - EXPECT_EQ(decode_result, Error::Ok); + EXPECT_EQ(io_manager_->update_prefill(model_outputs), Error::Ok); + EXPECT_EQ(io_manager_->update_decode(model_outputs), Error::Ok); } diff --git a/extension/llm/runner/irunner.h b/extension/llm/runner/irunner.h index 4c2efc91203..5bd5ef9d04e 100644 --- a/extension/llm/runner/irunner.h +++ b/extension/llm/runner/irunner.h @@ -130,7 +130,9 @@ class ET_EXPERIMENTAL IRunner { * given position in KV cache. * * @param prompt The input prompt to generate from - * @param start_pos The starting position in KV cache of the input + * @param start_pos The starting position in KV cache of the input. Note: + * Depending on the actual implementation, a runner may manage the position + * internally, and this may not be respected. * @param config Generation configuration parameters * @param token_callback Callback function called for each generated token * @param stats_callback Callback function for generation statistics @@ -146,6 +148,16 @@ class ET_EXPERIMENTAL IRunner { * Stop the generation process. */ virtual void stop() = 0; + /** + * Force remove prefilled tokens and reset KV cache start position + * + * For some existing runners, overriding this method is not needed because + * start_pos is passed as an argument to generate_from_pos. + * + * This method removes the prefilled tokens from the KV cache and resets the + * start position to 0. + */ + virtual void reset() {}; }; } // namespace llm diff --git a/extension/llm/runner/llm_runner_helper.cpp b/extension/llm/runner/llm_runner_helper.cpp index 555d6eed08c..f12de5f1d87 100644 --- a/extension/llm/runner/llm_runner_helper.cpp +++ b/extension/llm/runner/llm_runner_helper.cpp @@ -5,23 +5,27 @@ * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ - +// @lint-ignore-every CLANGTIDY facebook-hte-Deprecated // Implementation of helper utilities for creating and configuring LLM runners +#include #include +#include +#include +#include #include #include #include #include +#include #include #include #include #include +#include #include -namespace executorch { -namespace extension { -namespace llm { +namespace executorch::extension::llm { using ::executorch::extension::Module; using ::executorch::runtime::Error; @@ -33,6 +37,18 @@ std::unique_ptr load_tokenizer( size_t bos_token_index, size_t eos_token_index) { runtime::runtime_init(); + auto tekken_tokenizer = std::make_unique(); + // Prevent the case where tekken tokenizer accidentally successfully loads a + // HuggingFace tokenizer, which is also .json. + static constexpr std::string_view tekken_name = "tekken.json"; + if (tokenizer_path.size() >= tekken_name.size() && + tokenizer_path.rfind(tekken_name) == + tokenizer_path.size() - tekken_name.size()) { + if (tekken_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded tekken tokenizer"); + return tekken_tokenizer; + } + } auto json_tokenizer = std::make_unique(); if (json_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { ET_LOG(Info, "Loaded json tokenizer"); @@ -71,9 +87,8 @@ std::unique_ptr load_tokenizer( return nullptr; } -std::unordered_map get_llm_metadata( - tokenizers::Tokenizer* tokenizer, - Module* module) { +::executorch::runtime::Result> +get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module) { // Initialize metadata with default values std::unordered_map metadata({ {llm::kEnableDynamicShape, false}, @@ -87,10 +102,20 @@ std::unordered_map get_llm_metadata( auto method_names_result = module->method_names(); if (method_names_result.error() != Error::Ok) { ET_LOG(Error, "Failed reading method names"); - return metadata; + return ::executorch::runtime::Error::InvalidArgument; } const auto& method_names = method_names_result.get(); + // Error out if the max seq len metadata method is not present, since + // it is hard to figure out from just the .pte itself. + if (!method_names.count(llm::kMaxSeqLen)) { + ET_LOG( + Error, + "Required metadata method %s not found in model", + llm::kMaxSeqLen); + return ::executorch::runtime::Error::InvalidArgument; + } + for (auto& pair : metadata) { const auto& method_name = pair.first; auto& value = pair.second; @@ -107,6 +132,18 @@ std::unordered_map get_llm_metadata( } ET_LOG(Info, "Metadata: %s = %" PRId64, method_name.c_str(), value); } + + // If kMaxContextLen method not found but kMaxSeqLen is + // available, set kMaxContextLen to the value of kMaxSeqLen. + if (!method_names.count(llm::kMaxContextLen) && + method_names.count(llm::kMaxSeqLen)) { + metadata[llm::kMaxContextLen] = metadata[llm::kMaxSeqLen]; + ET_LOG( + Info, + "Setting kMaxContextLen to kMaxSeqLen value: %" PRId64, + metadata[llm::kMaxContextLen]); + } + // Set tokenizer-related metadata metadata[llm::kBosId] = tokenizer->bos_tok(); metadata[llm::kVocabSize] = tokenizer->vocab_size(); @@ -163,13 +200,18 @@ std::unique_ptr create_text_llm_runner( // Get metadata from Module ET_LOG(Info, "Reading metadata from model"); - auto metadata = llm::get_llm_metadata(tokenizer.get(), module.get()); + auto metadata_result = llm::get_llm_metadata(tokenizer.get(), module.get()); + if (metadata_result.error() != Error::Ok) { + ET_LOG(Error, "Failed to get metadata from model"); + return nullptr; + } + auto metadata = metadata_result.get(); auto eos_ids = std::make_unique>( llm::get_eos_ids(tokenizer.get(), module.get())); // Create IOManager - std::unique_ptr io_manager = std::make_unique(); + std::unique_ptr io_manager = std::make_unique(*module); // Create text_decoder_runner. Use a shared_ptr so that it can be shared with // TextPrefiller and TextTokenGenerator @@ -205,6 +247,70 @@ std::unique_ptr create_text_llm_runner( temperature); } -} // namespace llm -} // namespace extension -} // namespace executorch +std::unique_ptr create_multimodal_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path) { + // Sanity check tokenizer + if (!tokenizer || !tokenizer->is_loaded()) { + ET_LOG(Error, "Tokenizer is null or not loaded"); + return nullptr; + } + + // Create the Module + std::unique_ptr module; + if (data_path.has_value()) { + module = std::make_unique( + model_path, data_path.value(), Module::LoadMode::File); + } else { + module = std::make_unique(model_path, Module::LoadMode::File); + } + + // Get metadata from Module + ET_LOG(Info, "Reading metadata from model"); + auto metadata_result = get_llm_metadata(tokenizer.get(), module.get()); + if (metadata_result.error() != Error::Ok) { + ET_LOG(Error, "Failed to get metadata from model"); + return nullptr; + } + auto metadata = metadata_result.get(); + + auto eos_ids = std::make_unique>( + get_eos_ids(tokenizer.get(), module.get())); + + // Create IOManager + std::unique_ptr io_manager = std::make_unique(*module); + + // Create text_decoder_runner + auto text_decoder_runner = + std::make_unique(module.get(), io_manager.get()); + + // Create multimodal_prefiller + auto multimodal_prefiller = std::make_unique( + module.get(), + text_decoder_runner.get(), + tokenizer.get(), + io_manager.get()); + + // Create text_token_generator with stats + auto stats = std::make_unique(); + auto text_token_generator = std::make_unique( + tokenizer.get(), + text_decoder_runner.get(), + metadata.at(kUseKVCache), + std::move(eos_ids), + stats.get()); + + // Create and return the MultimodalRunner instance + return std::make_unique( + std::move(metadata), + std::move(tokenizer), + std::move(module), + std::move(text_decoder_runner), + std::move(multimodal_prefiller), + std::move(io_manager), + std::move(text_token_generator), + std::move(stats)); +} + +} // namespace executorch::extension::llm diff --git a/extension/llm/runner/llm_runner_helper.h b/extension/llm/runner/llm_runner_helper.h index 7e91a39abc4..191ea3ab090 100644 --- a/extension/llm/runner/llm_runner_helper.h +++ b/extension/llm/runner/llm_runner_helper.h @@ -19,12 +19,11 @@ #include #include +#include #include #include -namespace executorch { -namespace extension { -namespace llm { +namespace executorch::extension::llm { // Forward declarations class TextLLMRunner; @@ -61,11 +60,13 @@ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( * * @param tokenizer Initialized tokenizer instance * @param module The model module - * @return std::unordered_map Metadata key-value pairs + * @return Result> Metadata key-value + * pairs on success, or Error::InvalidArgument if required metadata (e.g., + * kMaxSeqLen) is missing from the model */ -ET_EXPERIMENTAL std::unordered_map get_llm_metadata( - tokenizers::Tokenizer* tokenizer, - Module* module); +ET_EXPERIMENTAL ::executorch::runtime::Result< + std::unordered_map> +get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module); /** * @brief Gets EOS token IDs from the model and tokenizer @@ -103,6 +104,21 @@ ET_EXPERIMENTAL std::unique_ptr create_text_llm_runner( std::optional data_path = std::nullopt, float temperature = -1.0f); -} // namespace llm -} // namespace extension -} // namespace executorch +/** + * @brief Creates a MultimodalRunner instance with dependency injection + * + * This factory function creates and initializes a MultimodalRunner with all + * necessary components for multimodal text generation. + * + * @param model_path Path to the model file + * @param tokenizer Initialized tokenizer instance + * @param data_path Optional path to additional .ptd required by the model + * @return std::unique_ptr Initialized MultimodalRunner + * instance, or nullptr on failure + */ +ET_EXPERIMENTAL std::unique_ptr create_multimodal_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path = std::nullopt); + +} // namespace executorch::extension::llm diff --git a/extension/llm/runner/multimodal_decoder_runner.h b/extension/llm/runner/multimodal_decoder_runner.h new file mode 100644 index 00000000000..f76b8c64028 --- /dev/null +++ b/extension/llm/runner/multimodal_decoder_runner.h @@ -0,0 +1,106 @@ +/* + * 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 executorch::extension::llm { + +class ET_EXPERIMENTAL MultimodalDecoderRunner + : public executorch::extension::llm::TextDecoderRunner { + public: + explicit MultimodalDecoderRunner(Module* module, IOManager* io_manager) + : TextDecoderRunner(module, io_manager) {} + + /** + * Step the LLM Decoder with the given tokens and start position. + * @param tokens The tokens to the LLM. + * @param start_pos The start position of the tokens. + * @return The logits tensor. + */ + inline executorch::runtime::Result step( + executorch::extension::TensorPtr& tokens, + int64_t start_pos) override { + // run token embedding + auto token_embedding_outputs = + ET_UNWRAP(module_->execute(kTokenEmbeddingMethod, tokens)); + + // Return the logits tensor + return decode(token_embedding_outputs[0], start_pos); + } + + /** + * Decode the embeddings to logits. + * @param embeddings The embeddings tensor. + * @param start_pos The start position of the embeddings. + * @return The logits tensor. + */ + inline executorch::runtime::Result decode( + const runtime::EValue& embeddings, + int64_t start_pos) { + auto start_pos_tensor = ::executorch::extension::from_blob( + &start_pos, {1}, executorch::aten::ScalarType::Long); + // run text model + auto outputs_res = ET_UNWRAP( + module_->execute(kTextModelMethod, {start_pos_tensor, embeddings})); + + ET_CHECK_MSG( + outputs_res.size() == 1, + "More then one output returned from executing LLM."); + ET_CHECK_MSG( + outputs_res[0].isTensor(), + "Non Tensor Output returned from executing LLM"); + + // Return the logits tensor + return outputs_res[0].toTensor(); + } + + /** + * Load the Module for text decode purpose. + * @return The error code. + */ + inline executorch::runtime::Error load() override { + if (is_method_loaded()) { + return executorch::runtime::Error::Ok; + } + ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kTokenEmbeddingMethod)); + ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kTextModelMethod)); + return executorch::runtime::Error::Ok; + } + + /** + * Check if the required methods in the Module is loaded. + * @return True if the Module is loaded, false otherwise. + */ + inline bool is_method_loaded() override { + executorch::runtime::Result> methods_res = + module_->method_names(); + if (methods_res.error() != executorch::runtime::Error::Ok) { + ET_CHECK_MSG(false, "Failed to get method names"); + } + std::unordered_set methods = methods_res.get(); + bool methods_exist = methods.find(kTokenEmbeddingMethod) != methods.end() && + methods.find(kTextModelMethod) != methods.end(); + if (!methods_exist) { + for (const auto& method : methods) { + ET_LOG(Error, "Method: %s", method.c_str()); + } + ET_CHECK_MSG( + methods_exist, + "Missing required methods (%s, %s) in the model", + kTokenEmbeddingMethod, + kTextModelMethod); + } + bool methods_loaded = module_->is_method_loaded(kTokenEmbeddingMethod) && + module_->is_method_loaded(kTextModelMethod); + return methods_loaded; + } +}; + +} // namespace executorch::extension::llm diff --git a/extension/llm/runner/multimodal_input.h b/extension/llm/runner/multimodal_input.h new file mode 100644 index 00000000000..728d8aef08f --- /dev/null +++ b/extension/llm/runner/multimodal_input.h @@ -0,0 +1,327 @@ +/* + * 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. + */ + +// @lint-ignore-every CLANGTIDY facebook-hte-Deprecated +// A generic multimodal input class that can hold either image or text data. + +#pragma once + +#include +#include +#include +#include +#include + +namespace executorch::extension::llm { + +/** + * A generic class to hold either image, text, or audio data for multimodal + * inputs. This allows the generate() API to take a std::vector of these objects + * instead of separate image, text, and audio parameters. + */ +class ET_EXPERIMENTAL MultimodalInput { + public: + /// Type of multimodal input data + enum class Type { + TEXT, ///< Text string input + IMAGE, ///< Processed image input + AUDIO, ///< Processed audio input + RAW_AUDIO, ///< Raw unprocessed audio input (straight from audio file) + UNSUPPORTED ///< Unsupported input type + }; + + // Constructors + explicit MultimodalInput(const std::string& text) : data_(text) {} + explicit MultimodalInput(std::string&& text) : data_(std::move(text)) {} + explicit MultimodalInput(const Image& image) : data_(image) {} + explicit MultimodalInput(Image&& image) : data_(std::move(image)) {} + explicit MultimodalInput(const Audio& audio) : data_(audio) {} + explicit MultimodalInput(Audio&& audio) : data_(std::move(audio)) {} + explicit MultimodalInput(const RawAudio& raw_audio) : data_(raw_audio) {} + explicit MultimodalInput(RawAudio&& raw_audio) + : data_(std::move(raw_audio)) {} + + // Copy constructor and assignment + MultimodalInput(const MultimodalInput& other) = default; + MultimodalInput& operator=(const MultimodalInput& other) = default; + + // Move constructor and assignment + MultimodalInput(MultimodalInput&& other) noexcept = default; + MultimodalInput& operator=(MultimodalInput&& other) noexcept = default; + + // Destructor + ~MultimodalInput() = default; + + /** + * Check if this input contains text data. + * @return true if this input contains text, false otherwise. + */ + bool is_text() const noexcept { + return std::holds_alternative(data_); + } + + /** + * Check if this input contains image data. + * @return true if this input contains an image, false otherwise. + */ + bool is_image() const noexcept { + return std::holds_alternative(data_); + } + + /** + * Check if this input contains audio data. + * @return true if this input contains audio, false otherwise. + */ + bool is_audio() const noexcept { + return std::holds_alternative