diff --git a/.ci/scripts/test-native.sh b/.ci/scripts/test-native.sh new file mode 100644 index 00000000000..2f08fc6dabf --- /dev/null +++ b/.ci/scripts/test-native.sh @@ -0,0 +1,67 @@ +#!/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 -euxo pipefail + +component="${1:?Usage: test-native.sh [cmake options]}" +shift +build_dir="cmake-out/native-${component}" +build_jobs="${CMAKE_BUILD_PARALLEL_LEVEL:-$(nproc)}" +cmake_args=() +targets=() + +case "${component}" in + runtime) + targets+=(native_runtime_test) + test_regex='^native_runtime_test$' + ;; + ggml) + cmake_args+=( + -DEXECUTORCH_BUILD_NATIVE_GGML=ON + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON + -DGGML_NATIVE=OFF + ) + targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test) + test_regex='^native_(ggml_engine|llm_runner_utils)_test$' + ;; + vulkan) + cmake_args+=( + -DEXECUTORCH_BUILD_NATIVE_VULKAN=ON + -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON + -DEXECUTORCH_VULKAN_SHADER_COMPILE_NTHREADS="${build_jobs}" + ) + targets+=(native_vulkan_engine run_classification run_hf_llm_vulkan native_vulkan_passes_test native_runner_utils_test) + test_regex='^native_(vulkan_passes|runner_utils)_test$' + ;; + *) + echo "Unknown native component: ${component}" >&2 + exit 1 + ;; +esac + +cmake -S . -B "${build_dir}" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DPYTHON_EXECUTABLE="$(command -v python)" \ + -DEXECUTORCH_BUILD_TESTS=ON \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \ + -DEXECUTORCH_BUILD_NATIVE=ON \ + "${cmake_args[@]}" "$@" +cmake --build "${build_dir}" \ + --parallel "${build_jobs}" \ + --target "${targets[@]}" +ctest --test-dir "${build_dir}/backends/native/test" \ + -R "${test_regex}" --output-on-failure --no-tests=error --timeout 120 + +if [[ "${component}" == "ggml" ]]; then + PYTHONPATH="${PWD}/src:${PYTHONPATH:-}" \ + FLATC_EXECUTABLE="${FLATC_EXECUTABLE:-${PWD}/${build_dir}/third-party/flatc_ep/bin/flatc}" \ + python .ci/scripts/test_native_llm_runner.py \ + --runner "${build_dir}/backends/native/examples/hf_llm/runtime/run_hf_llm_ggml" +fi diff --git a/.ci/scripts/test_native_llm_runner.py b/.ci/scripts/test_native_llm_runner.py new file mode 100644 index 00000000000..9fe9eadf638 --- /dev/null +++ b/.ci/scripts/test_native_llm_runner.py @@ -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. + +"""Compare GGML LLM runner prefill and decoding with a tiny eager HF Llama.""" + +import argparse +import base64 +import re +import subprocess +import tempfile +from pathlib import Path + +import torch +from executorch.backends.native.examples.hf_llm.export_hf_llm import LastTokenLM +from executorch.exir.native import to_native +from transformers import LlamaConfig, LlamaForCausalLM + + +@torch.no_grad() +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--runner", required=True, type=Path) + args = parser.parse_args() + + torch.manual_seed(0) + torch.set_num_threads(4) + max_seq_length = 48 + max_new_tokens = 8 + config = LlamaConfig( + vocab_size=512, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=max_seq_length, + bos_token_id=256, + eos_token_id=257, + attn_implementation="sdpa", + ) + model = LlamaForCausalLM(config).eval() + model.generation_config.use_cache = True + model.generation_config.cache_implementation = "static" + wrapper = LastTokenLM(model, max_seq_length).eval() + seq = torch.export.Dim("seq", min=1, max=max_seq_length - 1) + exported = torch.export.export( + wrapper, + (torch.zeros((1, 6), dtype=torch.long), torch.arange(6)), + dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, + strict=True, + ) + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + package = root / "llama.ptn" + to_native(exported).save(str(package)) + tokenizer = root / "tokenizer.model" + # The byte vocabulary puts Llama's BOS at 256, followed by 255 specials. + tokenizer.write_text( + "".join( + f"{base64.b64encode(bytes([i])).decode()} {i}\n" for i in range(256) + ) + ) + prompt = "abc" + unreachable_eos_id = config.vocab_size + tokens = [config.bos_token_id, *prompt.encode()] + expected = [] + # Recompute the full prefix without a cache for an independent reference. + for _ in range(max_new_tokens): + logits = model(torch.tensor([tokens]), use_cache=False).logits + token = int(logits[0, -1].argmax()) + expected.append(token) + tokens.append(token) + + result = subprocess.run( + [ + str(args.runner.resolve()), + "--device=cpu", + "--cpu_threads=4", + f"--package={package}", + f"--tokenizer={tokenizer}", + f"--prompt={prompt}", + f"--eos_ids={unreachable_eos_id}", + f"--max_new_tokens={max_new_tokens}", + f"--max_seq_length={max_seq_length}", + "--print_token_ids", + ], + capture_output=True, + text=True, + errors="replace", + timeout=120, + ) + print(result.stdout, end="") + print(result.stderr, end="") + result.check_returncode() + match = re.search(r"^token_ids: ([0-9,]+)$", result.stdout, re.MULTILINE) + if match is None: + raise AssertionError("Runner did not report generated token IDs") + actual = [int(token) for token in match[1].split(",")] + if actual != expected: + raise AssertionError(f"Generated tokens {actual} != eager {expected}") + print(f"GGML HF LLM runner matched all {max_new_tokens} eager token IDs") + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 41dd61ebcda..07653ae8bd4 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -43,6 +43,10 @@ on: - examples/models/muse-glimmer/** - extension/pybindings/** - runtime/__init__.py + - exir/passes/convert_constant_dim_order_pass.py + - exir/passes/insert_write_back_for_buffers_pass.py + - exir/passes/propagate_device_config.py + - exir/passes/propagate_device_pass.py workflow_dispatch: inputs: models: @@ -469,8 +473,9 @@ jobs: unittest-cuda: name: unittest-cuda (${{ matrix.shard }}) needs: [changed-files, run-decision] - # This job runs the pybindings device test, so it also has to fire on the - # two pybindings paths the workflow filter lists. + # This job runs the pybindings device test and the device placement test, + # so it also has to fire on the pybindings and device pass paths the + # workflow filter lists. if: | contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || @@ -481,6 +486,10 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || contains(needs.changed-files.outputs.changed-files, 'extension/pybindings') || contains(needs.changed-files.outputs.changed-files, 'runtime/__init__.py') || + contains(needs.changed-files.outputs.changed-files, 'exir/passes/convert_constant_dim_order_pass.py') || + contains(needs.changed-files.outputs.changed-files, 'exir/passes/insert_write_back_for_buffers_pass.py') || + contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_config.py') || + contains(needs.changed-files.outputs.changed-files, 'exir/passes/propagate_device_pass.py') || needs.run-decision.outputs.is-full-run == 'true' # kernel: Triton kernels and their autotuning. system: lowering passes, op # dispatch, export/runtime, and model tests. Together they do not fit the @@ -567,6 +576,21 @@ jobs: echo "::error::a pybindings device test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing" exit 1 fi + + # Same for the device placement test, which skips without a GPU. + set +e + python -m pytest backends/cuda/tests/test_device_placement.py -v -rs -o "addopts=" \ + 2>&1 | tee /tmp/device_placement_test.log + device_placement_status=${PIPESTATUS[0]} + set -e + if [ "${device_placement_status}" -ne 0 ]; then + echo "::error::a device placement test failed" + exit "${device_placement_status}" + fi + if grep -qE "[0-9]+ skipped" /tmp/device_placement_test.log; then + echo "::error::a device placement test skipped in the CUDA job, which means the CUDA backend or a visible GPU is missing" + exit 1 + fi fi # Build ExecuTorch with CUDA support @@ -589,7 +613,7 @@ jobs: python -m pytest "${KERNEL_TESTS[@]}" -v -o "addopts=" --durations=20 else IGNORES=() - for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}"; do + for t in "${KERNEL_TESTS[@]}" "${QUANTIZED_GEMM_TESTS[@]}" backends/cuda/tests/test_device_placement.py; do IGNORES+=("--ignore=${t}") done python -m pytest backends/cuda/tests backends/cuda/passes/tests -v -o "addopts=" --durations=20 \ diff --git a/.github/workflows/test-native.yml b/.github/workflows/test-native.yml new file mode 100644 index 00000000000..c5e0f0b7c76 --- /dev/null +++ b/.github/workflows/test-native.yml @@ -0,0 +1,113 @@ +name: Test Native Engines + +on: + pull_request: + paths: + - .github/workflows/test-native.yml + - .github/workflows/_docker-image.yml + - .ci/docker/** + - .ci/scripts/test-native.sh + - .ci/scripts/test_native_llm_runner.py + - .gitmodules + - CMakeLists.txt + - install_requirements.py + - install_requirements.sh + - requirements-dev.txt + - requirements-examples.txt + - backends/native/** + - backends/vulkan/CMakeLists.txt + - backends/vulkan/cmake/** + - backends/vulkan/runtime/api/** + - backends/vulkan/runtime/vk_api/** + - backends/vulkan/runtime/graph/** + - backends/vulkan/runtime/utils/** + - backends/vulkan/runtime/gen_vulkan_spv.py + - backends/vulkan/third-party/Vulkan-Headers + - backends/vulkan/third-party/VulkanMemoryAllocator + - backends/vulkan/third-party/volk + - backends/vulkan_shared/CMakeLists.txt + - backends/vulkan_shared/runtime/** + - backends/xnnpack/third-party/FXdiv + - backends/xnnpack/third-party/cpuinfo + - backends/xnnpack/third-party/pthreadpool + - examples/models/llama/tokenizer/** + - extension/llm/tokenizers + - extension/llm/export/** + - exir/** + - runtime/** + - third-party/** + - tools/cmake/** + - build/** + - codegen/** + - schema/** + push: + branches: + - main + - release/* + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }} + cancel-in-progress: true + +jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + + test-native: + name: native-${{ matrix.component }} + needs: docker-image + runs-on: mt-l-x86iavx512-8-64 + container: + image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} + timeout-minutes: 60 + permissions: + contents: read + strategy: + fail-fast: false + matrix: + component: [runtime, ggml, vulkan] + defaults: + run: + shell: bash + steps: + - name: Fix workspace permissions + run: sudo chown -R "$(id -u):$(id -g)" "${GITHUB_WORKSPACE}" + + - name: Checkout ExecuTorch + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + submodules: recursive + persist-credentials: false + + - name: Build and test native ${{ matrix.component }} + env: + NATIVE_COMPONENT: ${{ matrix.component }} + run: | + set -exo pipefail + eval "$(/opt/conda/bin/conda shell.bash hook)" + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + + unset SCCACHE_BUCKET SCCACHE_S3_KEY_PREFIX + export SCCACHE_DIR=/tmp/sccache + + packages=(libzip) + if [[ "${NATIVE_COMPONENT}" == "vulkan" ]]; then + packages+=(shaderc) + fi + conda create -y -p /tmp/native-deps -c conda-forge "${packages[@]}" + export CMAKE_PREFIX_PATH="/tmp/native-deps:${CMAKE_PREFIX_PATH:-}" + export PATH="/tmp/native-deps/bin:${PATH}" + + if [[ "${NATIVE_COMPONENT}" == "ggml" ]]; then + ./install_requirements.sh --use-pt-pinned-commit + # The LLM runner test imports executorch from src/ without installing the + # wheel, so install its export-time dependencies directly. + pip install flatbuffers ruamel.yaml tabulate \ + "$(grep -E '^transformers\b' requirements-examples.txt)" + fi + + CMAKE_BUILD_PARALLEL_LEVEL=8 bash .ci/scripts/test-native.sh "${NATIVE_COMPONENT}" diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ae92a457b5..3768d04a083 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -980,7 +980,10 @@ if(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP) endif() if(EXECUTORCH_BUILD_EXTENSION_LLM) - if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) + if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER + OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER + ) + # Both LLM runners need tokenizers with regex lookahead support. set(SUPPORT_REGEX_LOOKAHEAD ON) # llama/runner/CMakeLists.txt builds a shared library libllama_runner.so # that transitively depends on tokenizers. Need to build tokenizers with @@ -991,7 +994,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) set(CMAKE_POSITION_INDEPENDENT_CODE ON) endif() add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/tokenizers) - if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) + if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER + OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER + ) set(CMAKE_POSITION_INDEPENDENT_CODE ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG} ) diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index 93afebb6fa2..18e9006b51a 100644 --- a/backends/cpu/BUCK +++ b/backends/cpu/BUCK @@ -22,6 +22,8 @@ fbcode_target( deps = [ "//caffe2:torch", "//executorch/backends/native:lib", + "//executorch/backends/transforms:fuse_batch_norm_with_conv", + "//executorch/backends/transforms:fuse_batch_norm_with_linear", "//executorch/exir:lib", "//executorch/exir/_serialize:lib", "//executorch/exir/backend:backend_details", @@ -32,6 +34,43 @@ fbcode_target( ], ) +fbcode_target( + _kind = runtime.python_library, + name = "export_torchvision_lib", + srcs = [ + "workflows/__init__.py", + "workflows/export.py", + ], + deps = [ + ":lib", + "//caffe2:torch", + "//executorch/backends/native:lib", + "//executorch/backends/xnnpack/_passes:xnnpack_passes", + "//executorch/backends/xnnpack/recipes:xnnpack_recipes", + "//executorch/export:lib", + "//pytorch/ao:torchao", + "//pytorch/vision:torchvision", # @manual + ], +) + +fbcode_target( + _kind = runtime.python_binary, + name = "export_torchvision", + par_style = "xar", + main_module = "executorch.backends.cpu.workflows.export", + deps = [":export_torchvision_lib"], +) + +fbcode_target( + _kind = runtime.python_test, + name = "test_export_torchvision", + srcs = ["test/test_export_torchvision.py"], + deps = [ + ":export_torchvision_lib", + "//caffe2:torch", + ], +) + fbcode_target( _kind = runtime.python_test, name = "test_aot", @@ -57,3 +96,16 @@ fbcode_target( "//executorch/extension/pybindings:portable_lib", # @manual ], ) + +fbcode_target( + _kind = runtime.python_test, + name = "test_runtime", + srcs = ["test/test_runtime.py"], + preload_deps = [":cpu_backend"], + deps = [ + ":lib", + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", + ], +) diff --git a/backends/cpu/README.md b/backends/cpu/README.md index 0c105e8cdf0..51c58edebeb 100644 --- a/backends/cpu/README.md +++ b/backends/cpu/README.md @@ -6,8 +6,11 @@ The initial workload is pretrained ConvNeXt-Tiny in FP32 on a Linux server. The registered `CPURecipeType.FP32` (`cpu_fp32`) Recipe retains a semantic Native graph inside one `CpuBackend` delegate in a regular ExecuTorch `.pte`. -This first delegate version supports pure operations on statically shaped, -contiguous FP32 tensors. +The partitioner supports pure operations on statically shaped, contiguous FP32, +uint8, int8, int16, int32, int64 and boolean tensors. Integer and boolean operations +use the existing ET kernel routes, subject to each kernel's operator and dtype +support. Runtime-computed scalar arguments and symbolic tensor shapes remain +unsupported. Ordinary integer tensors do not imply quantized execution support. The `cpu_delegate_version` compile spec stores `CPU_DELEGATE_VERSION` as a four-byte little-endian unsigned integer. The loader accepts explicitly @@ -44,7 +47,7 @@ allocation. Sequential steps share the maximum declared scratch capacity. XNNPAC shares compatible convolution-weight conversions within a plan and borrows weights whose physical order already matches. -For a loaded method, shapes, FP32 dtype, contiguous layout, threadpool identity, +For a loaded method, shapes, tensor dtypes, contiguous layout, threadpool identity, thread count and ISA capabilities remain fixed. Input/output pointers may change and are validated on every invocation. Changing an execution property requires reloading the method; failed validation executes no provider. The trace includes diff --git a/backends/cpu/partitioner.py b/backends/cpu/partitioner.py index 3aaf77a4324..808150838ca 100644 --- a/backends/cpu/partitioner.py +++ b/backends/cpu/partitioner.py @@ -23,11 +23,19 @@ # Serialized as a little-endian uint32. Bump for incompatible artifact contracts. CPU_DELEGATE_VERSION = 1 +CPU_SUPPORTED_DTYPES = ( + torch.float32, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.bool, +) class CPUSemanticOperators(OperatorSupportBase): - """Pure tensor operations representable by the initial static FP32 - delegate version. + """Pure static tensor operations on FP32, integer and boolean values. Runtime providers determine implementation coverage. An unresolved operation is reported by name when the CPU runtime prepares the graph. @@ -64,7 +72,7 @@ def is_node_supported( if any(related_node.meta.get("val") is None for related_node in related_nodes): return False return all( - value.dtype == torch.float32 + value.dtype in CPU_SUPPORTED_DTYPES and all(isinstance(size, int) for size in value.shape) and value.is_contiguous() for related_node in related_nodes @@ -74,7 +82,7 @@ def is_node_supported( class CPUPartitioner(Partitioner): - """Retain the static FP32 semantic graph for runtime provider selection.""" + """Retain the static semantic graph for runtime provider selection.""" def ops_to_not_decompose(self, ep: torch.export.ExportedProgram): support = CPUSemanticOperators() diff --git a/backends/cpu/preprocess.py b/backends/cpu/preprocess.py index b0a7f9c9141..42eb88e298b 100644 --- a/backends/cpu/preprocess.py +++ b/backends/cpu/preprocess.py @@ -7,8 +7,17 @@ from copy import deepcopy import torch -from executorch.backends.cpu.partitioner import CPU_DELEGATE_VERSION +from executorch.backends.cpu.partitioner import ( + CPU_DELEGATE_VERSION, + CPU_SUPPORTED_DTYPES, +) from executorch.backends.native.serialization import serialize_graph +from executorch.backends.transforms.fuse_batch_norm_with_conv import ( + FuseBatchNormWithConvPass, +) +from executorch.backends.transforms.fuse_batch_norm_with_linear import ( + FuseBatchNormWithLinearPass, +) from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import ( BackendDetails, @@ -16,8 +25,20 @@ ExportedProgram, PreprocessResult, ) +from executorch.exir.program._program import ( + _transform, + _update_exported_program_graph_module, + lift_constant_tensor_pass, +) +from torch.export.graph_signature import InputKind from torch.utils._pytree import tree_leaves, tree_map +_CONSTANT_INPUT_KINDS = ( + InputKind.PARAMETER, + InputKind.BUFFER, + InputKind.CONSTANT_TENSOR, +) + def _validate_boundary_layouts(graph_module: torch.fx.GraphModule) -> None: for node in graph_module.graph.nodes: @@ -34,8 +55,41 @@ def _validate_boundary_layouts(graph_module: torch.fx.GraphModule) -> None: ) +def _remove_unused_constants(edge_program: ExportedProgram) -> None: + signature = edge_program.graph_signature + mutated = set(signature.buffers_to_mutate.values()) + placeholders = { + node.name: node for node in edge_program.graph.nodes if node.op == "placeholder" + } + kept_specs = [] + for spec in signature.input_specs: + node = placeholders.get(spec.arg.name) + if ( + node is not None + and spec.kind in _CONSTANT_INPUT_KINDS + and spec.target not in mutated + and not node.users + ): + edge_program.graph.erase_node(node) + edge_program.state_dict.pop(spec.target, None) + edge_program.constants.pop(spec.target, None) + else: + kept_specs.append(spec) + signature.input_specs = kept_specs + edge_program.graph_module.recompile() + + +# Ordered passes applied to every program in preprocess, mirroring +# XNNPACKPassManager.passes. Each entry takes the current program, so later +# passes observe earlier passes' rewrites; extend by appending new pass types. +CPU_PREPROCESS_PASSES = [ + FuseBatchNormWithConvPass, + FuseBatchNormWithLinearPass, +] + + class CpuBackend(BackendDetails): - """Serialize provider-neutral Native semantics with contiguous FP32 storage. + """Serialize provider-neutral Native semantics with contiguous tensor storage. The delegate version fixes 64-byte alignment and 64 initialized readable tail bytes. Named constants include the tail in their recorded buffer capacity. @@ -56,7 +110,17 @@ def preprocess( f"expected exactly {expected_specs!r}, got {specs!r}" ) _validate_boundary_layouts(edge_program.graph_module) - graph_module = deepcopy(edge_program.graph_module) + edge_program = deepcopy(edge_program) + edge_program = _update_exported_program_graph_module( + edge_program, edge_program.graph_module + ) + for pass_type in CPU_PREPROCESS_PASSES: + edge_program = _transform(edge_program, pass_type(edge_program)) + # Passes can introduce tensor attributes that must become program inputs. + edge_program = lift_constant_tensor_pass(edge_program) + # Fusion leaves the original weights and batch norm stats unused. + _remove_unused_constants(edge_program) + graph_module = edge_program.graph_module for node in graph_module.graph.nodes: if "val" in node.meta: node.meta["val"] = tree_map( @@ -73,9 +137,9 @@ def preprocess( ) data = NamedDataStore() for name, tensor in constants.items(): - if tensor.dtype != torch.float32 or tensor.device.type != "cpu": + if tensor.dtype not in CPU_SUPPORTED_DTYPES or tensor.device.type != "cpu": raise ValueError( - f"CpuBackend requires FP32 CPU constant {name}, " + f"CpuBackend requires supported CPU constant {name}, " f"got dtype {tensor.dtype} on {tensor.device}" ) if not tensor.is_contiguous(): diff --git a/backends/cpu/runtime/CPUBackend.cpp b/backends/cpu/runtime/CPUBackend.cpp index a3c39b5e879..3ff1396c71d 100644 --- a/backends/cpu/runtime/CPUBackend.cpp +++ b/backends/cpu/runtime/CPUBackend.cpp @@ -67,7 +67,7 @@ class CPUExecution { ET_CHECK_OR_RETURN_ERROR( spec.kind == ptn::OutputKind::UserOutput, NotSupported, - "CpuBackend static FP32 delegate does not support mutation"); + "CpuBackend static delegate does not support mutation"); } buffers_.resize(graph.values.size()); bytes_.reserve(graph.values.size()); @@ -171,7 +171,8 @@ class CPUExecution { const auto& tensor = args[index]->toTensor(); const auto& expected = method_->graph.value(id).tensor_meta(); ET_CHECK_OR_RETURN_ERROR( - tensor.scalar_type() == aten::ScalarType::Float && + tensor.scalar_type() == + static_cast(expected.dtype) && tensor.dim() == expected.sizes.size() && std::equal( expected.sizes.begin(), @@ -182,15 +183,20 @@ class CPUExecution { InvalidArgument, "CPU static shape/dtype/layout/storage mismatch: %s", method_->graph.value(id).name.c_str()); - int64_t stride = 1; - for (size_t reverse = expected.sizes.size(); reverse > 0; --reverse) { - const size_t dim = reverse - 1; - ET_CHECK_OR_RETURN_ERROR( - tensor.dim_order()[dim] == dim && tensor.strides()[dim] == stride, - InvalidArgument, - "CPU binding must be contiguous: %s", - method_->graph.value(id).name.c_str()); - stride *= expected.sizes[dim]; + // Empty bindings carry no data, so their strides and dim order are + // meaningless (PyTorch does not canonicalize them); skip the contiguity + // check rather than rejecting valid empty tensors. + if (tensor.nbytes() != 0) { + int64_t stride = 1; + for (size_t reverse = expected.sizes.size(); reverse > 0; --reverse) { + const size_t dim = reverse - 1; + ET_CHECK_OR_RETURN_ERROR( + tensor.dim_order()[dim] == dim && tensor.strides()[dim] == stride, + InvalidArgument, + "CPU binding must be contiguous: %s", + method_->graph.value(id).name.c_str()); + stride *= expected.sizes[dim]; + } } } // EValue exposes the logical span, but no backing capacity or readable diff --git a/backends/cpu/runtime/KernelProvider.cpp b/backends/cpu/runtime/KernelProvider.cpp index eba874f4f31..5e032c4c47b 100644 --- a/backends/cpu/runtime/KernelProvider.cpp +++ b/backends/cpu/runtime/KernelProvider.cpp @@ -5,6 +5,7 @@ // LICENSE file in the root directory of this source tree. #include +#include #include #include @@ -90,10 +91,16 @@ Result tensor_bytes(const ptn::Value& value) { value.name.c_str()); const auto& meta = value.tensor_meta(); ET_CHECK_OR_RETURN_ERROR( - meta.dtype == ptn::ScalarType::Float && meta.is_contiguous() && - meta.sizes.size() <= 16, + (meta.dtype == ptn::ScalarType::Float || + meta.dtype == ptn::ScalarType::Byte || + meta.dtype == ptn::ScalarType::Char || + meta.dtype == ptn::ScalarType::Short || + meta.dtype == ptn::ScalarType::Int || + meta.dtype == ptn::ScalarType::Long || + meta.dtype == ptn::ScalarType::Bool) && + meta.is_contiguous() && meta.sizes.size() <= 16, NotSupported, - "CPU static delegate requires contiguous FP32 rank <= 16: %s", + "CPU static delegate requires supported dtype and contiguous rank <= 16: %s", value.name.c_str()); size_t elements = 1; for (auto dimension = meta.sizes.rbegin(); dimension != meta.sizes.rend(); @@ -107,7 +114,13 @@ Result tensor_bytes(const ptn::Value& value) { value.name.c_str()); elements *= size; } - return elements * sizeof(float); + const size_t element_size = ptn::element_size(meta.dtype); + ET_CHECK_OR_RETURN_ERROR( + elements <= SIZE_MAX / element_size, + NotSupported, + "CPU tensor byte size overflows: %s", + value.name.c_str()); + return elements * element_size; } Result KernelImplementation::requirements( diff --git a/backends/cpu/runtime/KernelProvider.h b/backends/cpu/runtime/KernelProvider.h index a2571d25d52..0db278cd7e2 100644 --- a/backends/cpu/runtime/KernelProvider.h +++ b/backends/cpu/runtime/KernelProvider.h @@ -131,7 +131,7 @@ class KernelImplementation { return false; } /// Pure metadata query, before compile can read constants or allocate - /// storage. The default declares contiguous FP32 storage for every region + /// storage. The default declares contiguous tensor storage for every region /// boundary. virtual runtime::Result requirements( const KernelRegion& region, diff --git a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp index 5110d7eb108..8dc54961c0a 100644 --- a/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp +++ b/backends/cpu/runtime/providers/reference/ReferenceProvider.cpp @@ -29,6 +29,9 @@ struct KernelSpec { bool tensor_list_output = false; // Schema positions declared float or float?, rather than Scalar. std::string_view float_args{}; + // Tensor-list inputs are declared Tensor?[]. The serializer emits a plain + // tensor list when every entry is present, but the kernel unboxes optionals. + bool optional_tensor_lists = false; }; aten::ScalarType to_aten(ptn::ScalarType dtype) { @@ -103,6 +106,7 @@ constexpr auto kKernels = std::to_array({ nullptr, false, "56"}, + {"torch.ops.aten._neg_view_copy.default", "aten::neg.out", 1, 1, nullptr}, {"torch.ops.aten._pdist_forward.default", "aten::_pdist_forward.out", 2, @@ -296,8 +300,22 @@ constexpr auto kKernels = std::to_array({ {"torch.ops.aten.gt.Scalar", "aten::gt.Scalar_out", 2, 1, nullptr}, {"torch.ops.aten.gt.Tensor", "aten::gt.Tensor_out", 2, 1, nullptr}, {"torch.ops.aten.hardtanh.default", "aten::hardtanh.out", 3, 1, nullptr}, - {"torch.ops.aten.index.Tensor", "aten::index.Tensor_out", 2, 1, nullptr}, - {"torch.ops.aten.index_put.default", "aten::index_put.out", 4, 1, nullptr}, + {"torch.ops.aten.index.Tensor", + "aten::index.Tensor_out", + 2, + 1, + nullptr, + false, + {}, + true}, + {"torch.ops.aten.index_put.default", + "aten::index_put.out", + 4, + 1, + nullptr, + false, + {}, + true}, {"torch.ops.aten.index_select.default", "aten::index_select.out", 3, @@ -1004,6 +1022,20 @@ class ReferenceExecutable final : public Executable { reason.data()); return Error::NotSupported; } + if (node.target == "torch.ops.aten._neg_view_copy.default" && + tensor_bytes(context_.graph.value(node.inputs[0].arg.as_tensor().id)) + .get() == 0) { + const auto& input = + context_.graph.value(node.inputs[0].arg.as_tensor().id) + .tensor_meta(); + const auto& output = + context_.graph.value(node.outputs[0].value_id).tensor_meta(); + ET_CHECK_OR_RETURN_ERROR( + input.dtype == output.dtype && input.sizes == output.sizes, + InvalidProgram, + "Empty negative view copy requires matching input/output metadata"); + continue; + } auto instruction = std::make_unique(); instruction->node = &node; instruction->kernel = route->kernel; @@ -1037,11 +1069,16 @@ class ReferenceExecutable final : public Executable { if (input.arg.kind() == ptn::ArgKind::TensorList || input.arg.kind() == ptn::ArgKind::OptionalTensorList) { auto argument = std::make_unique(); - const bool optional = + const bool serialized_optional = input.arg.kind() == ptn::ArgKind::OptionalTensorList; - const auto& ids = optional ? input.arg.as_optional_tensor_list().ids - : input.arg.as_tensor_list().ids; - const auto error = bind_tensor_list(ids, optional, false, *argument); + const auto& ids = serialized_optional + ? input.arg.as_optional_tensor_list().ids + : input.arg.as_tensor_list().ids; + const auto error = bind_tensor_list( + ids, + serialized_optional || spec->optional_tensor_lists, + false, + *argument); if (error != Error::Ok) { return error; } diff --git a/backends/cpu/test/test_aot.py b/backends/cpu/test/test_aot.py index 61442987b50..42fd78b9546 100644 --- a/backends/cpu/test/test_aot.py +++ b/backends/cpu/test/test_aot.py @@ -14,12 +14,16 @@ CPUPartitioner, CPUSemanticOperators, ) -from executorch.backends.cpu.preprocess import CpuBackend +from executorch.backends.cpu.preprocess import _remove_unused_constants, CpuBackend from executorch.backends.cpu.recipes import CPURecipeType from executorch.backends.native.serialization import deserialize_graph, serialize_graph -from executorch.exir import to_edge_transform_and_lower +from executorch.backends.transforms.fuse_batch_norm_with_linear import ( + FuseBatchNormWithLinearPass, +) +from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.program._program import _transform, lift_constant_tensor_pass from executorch.exir.schema import DelegateCall, Tensor from executorch.export import export, ExportRecipe @@ -213,6 +217,64 @@ def test_constants_have_recorded_readable_tail(self): self.assertEqual(len(buffer), tensor.numel() * tensor.element_size() + 64) self.assertEqual(buffer[-64:], bytes(64)) + def test_preprocess_fuses_batch_norm_into_conv_and_linear(self): + class ConvBnLinearBn(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.bn1 = torch.nn.BatchNorm2d(4) + self.linear = torch.nn.Linear(4, 4) + self.bn2 = torch.nn.BatchNorm1d(4) + + def forward(self, x): + y = self.bn1(self.conv(x)) + return self.bn2(self.linear(y.mean(dim=(2, 3)))) + + model = ConvBnLinearBn().eval() + edge = to_edge( + torch.export.export(model, (torch.randn(1, 4, 8, 8),)), + compile_config=EdgeCompileConfig( + preserve_ops=[torch.ops.aten.linear.default] + ), + ).exported_program() + self.assertTrue( + any("batch_norm" in str(node.target) for node in edge.graph.nodes) + ) + result = CpuBackend.preprocess( + edge, + [ + CompileSpec( + "cpu_delegate_version", + CPU_DELEGATE_VERSION.to_bytes(4, "little"), + ) + ], + ) + graph = deserialize_graph(bytes(result.processed_bytes)) + targets = [node.target for node in graph.nodes] + self.assertNotIn( + "torch.ops.aten._native_batch_norm_legit_no_training.default", targets + ) + self.assertIn("torch.ops.aten.convolution.default", targets) + self.assertIn("torch.ops.aten.linear.default", targets) + # Only the fused conv and linear weights and biases remain. + self.assertEqual(len(result.data_store_output.pte_data), 4) + + def test_remove_unused_constants_prunes_backing_store(self): + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.BatchNorm1d(4)) + ep = to_edge( + torch.export.export(model.eval(), (torch.randn(2, 4),)), + compile_config=EdgeCompileConfig( + preserve_ops=[torch.ops.aten.linear.default] + ), + ).exported_program() + ep = _transform(ep, FuseBatchNormWithLinearPass(ep)) + ep = lift_constant_tensor_pass(ep) + self.assertIn("0.weight", ep.state_dict) + _remove_unused_constants(ep) + targets = {spec.target for spec in ep.graph_signature.input_specs} + self.assertNotIn("0.weight", ep.state_dict) + self.assertLessEqual(set(ep.state_dict) | set(ep.constants), targets) + def test_preprocess_rejects_compile_spec_mismatch(self): ep = torch.export.export(torch.nn.Linear(2, 3), (torch.ones(1, 2),)) for specs in ( @@ -274,7 +336,7 @@ def test_preprocess_rejects_unsupported_constants(self): ) with self.assertRaisesRegex( ValueError, - f"CpuBackend requires FP32 CPU constant weight, " + f"CpuBackend requires supported CPU constant weight, " f"got dtype {dtype} on {device}", ): CpuBackend.preprocess( @@ -404,7 +466,7 @@ def test_getitem_validates_output_metadata(self): node.meta["val"] = producer.meta["val"][0] support = CPUSemanticOperators() self.assertTrue(support.is_node_supported({}, node)) - for val in (None, 1, (), torch.ones(2, dtype=torch.int32)): + for val in (None, 1, (), torch.ones(2, dtype=torch.float64)): with self.subTest(val=val): node.meta["val"] = val self.assertFalse(support.is_node_supported({}, node)) @@ -414,7 +476,7 @@ class Add(torch.nn.Module): def forward(self, x): return x + x - ep = torch.export.export(Add(), (torch.ones(2, dtype=torch.int32),)) + ep = torch.export.export(Add(), (torch.ones(2, dtype=torch.float64),)) edge = to_edge_transform_and_lower(ep, partitioner=[CPUPartitioner()]) self.assertEqual( len(edge.to_executorch().executorch_program.execution_plan[0].delegates), 0 diff --git a/backends/cpu/test/test_export_torchvision.py b/backends/cpu/test/test_export_torchvision.py new file mode 100644 index 00000000000..05b3daa9aea --- /dev/null +++ b/backends/cpu/test/test_export_torchvision.py @@ -0,0 +1,130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import argparse +import json +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import torch +from executorch.backends.cpu.workflows.export import ( + capture, + image_size_for, + main, + synthetic_image, +) + + +class TorchvisionExportTest(unittest.TestCase): + def test_pretrained_vit_h_14_input_size(self) -> None: + self.assertEqual(518, image_size_for("vit_h_14", None)) + + def test_random_weights_keep_model_defaults(self) -> None: + self.assertEqual(224, image_size_for("vit_h_14", None, random_weights=True)) + self.assertEqual(299, image_size_for("inception_v3", None, random_weights=True)) + + def test_explicit_image_size_overrides_weights(self) -> None: + self.assertEqual(384, image_size_for("vit_h_14", 384)) + self.assertEqual(384, image_size_for("vit_h_14", 384, random_weights=True)) + + def test_random_weight_vit_capture_respects_image_size(self) -> None: + with tempfile.TemporaryDirectory() as directory: + args = argparse.Namespace( + model="vit_b_16", + image_size=384, + random_weights=True, + pt2=None, + output=Path(directory), + ) + + exported, image = capture(args) + + self.assertEqual((1, 3, 384, 384), tuple(image.shape)) + with torch.no_grad(): + self.assertEqual((1, 1000), tuple(exported.module()(image).shape)) + + def test_pretrained_vit_rejects_incompatible_image_size(self) -> None: + with tempfile.TemporaryDirectory() as directory: + args = argparse.Namespace( + model="vit_b_16", + image_size=384, + random_weights=False, + pt2=None, + output=Path(directory), + ) + + with self.assertRaisesRegex(ValueError, "image_size.*224.*384"): + capture(args) + + def test_reused_capture_preserves_input_shape(self) -> None: + model = torch.nn.Sequential(torch.nn.Flatten(), torch.nn.Linear(3 * 6 * 10, 4)) + image = synthetic_image((1, 3, 6, 10)) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source.pt2" + torch.export.save(torch.export.export(model.eval(), (image,)), source) + output = root / "output" + output.mkdir() + args = argparse.Namespace(model="vit_h_14", pt2=source, output=output) + + exported, actual_image = capture(args) + + torch.testing.assert_close(image, actual_image) + torch.testing.assert_close(model(image), exported.module()(actual_image)) + self.assertEqual(image.numel() * 4, (output / "input.bin").stat().st_size) + self.assertEqual(4 * 4, (output / "eager.bin").stat().st_size) + self.assertEqual( + source.read_bytes(), (output / "vit_h_14.pt2").read_bytes() + ) + + def test_reused_capture_rejects_image_size(self) -> None: + argv = [ + "export", + "--pt2", + "capture.pt2", + "--image-size", + "8", + "--output", + "out", + "--source-revision", + "test", + ] + with mock.patch("sys.argv", argv), self.assertRaises(SystemExit): + main() + + def test_manifest_lists_nodes_from_every_cpu_delegate(self) -> None: + model = torch.nn.Sequential( + torch.nn.Conv2d(3, 4, 3), + torch.nn.MaxPool2d(2), + torch.nn.Conv2d(4, 4, 3), + ).eval() + captured = torch.export.export(model, (synthetic_image((1, 3, 8, 8)),)) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + torch.export.save(captured, path / "capture.pt2") + argv = [ + "export", + "--pt2", + str(path / "capture.pt2"), + "--output", + directory, + "--source-revision", + "test", + ] + with mock.patch("sys.argv", argv): + main() + manifest = json.loads((path / "export.json").read_text()) + + self.assertEqual("capture", manifest["model"]) + self.assertEqual({"CpuBackend": 2}, manifest["delegates"]) + self.assertEqual( + ["aten_convolution_default", "aten_convolution_default_1"], + manifest["semantic_nodes"], + ) diff --git a/backends/cpu/test/test_runtime.py b/backends/cpu/test/test_runtime.py new file mode 100644 index 00000000000..4e821d76eec --- /dev/null +++ b/backends/cpu/test/test_runtime.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. + +from __future__ import annotations + +import unittest +from contextlib import contextmanager + +import torch +from executorch.backends.cpu.partitioner import CPUPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) + + +class CPURuntimeTest(unittest.TestCase): + @contextmanager + def _lower(self, model, inputs): + edge = to_edge_transform_and_lower( + torch.export.export(model, inputs), partitioner=[CPUPartitioner()] + ) + program = edge.to_executorch() + plan = program.executorch_program.execution_plan[0] + self.assertEqual([delegate.id for delegate in plan.delegates], ["CpuBackend"]) + self.assertEqual(plan.operators, []) + # The runtime borrows the serialized program buffer. + yield _load_for_executorch_from_buffer(program.buffer) + + def test_integer_inputs_and_outputs_preserve_precision(self): + class Add(torch.nn.Module): + def forward(self, x, y): + return x + y + + for dtype, values in ( + (torch.uint8, [0, 128, 250]), + (torch.int8, [-100, 0, 100]), + (torch.int16, [-30000, 0, 30000]), + (torch.int32, [-(2**24 + 1), 0, 2**24 + 1]), + (torch.int64, [-(2**40 + 3), 0, 2**40 + 3]), + ): + with self.subTest(dtype=dtype): + model = Add() + x = torch.tensor(values, dtype=dtype) + y = torch.ones_like(x) + with self._lower(model, (x, y)) as runtime: + for inputs in ((x, y), (y, x)): + actual = runtime.forward(inputs)[0] + torch.testing.assert_close( + actual, model(*inputs), rtol=0, atol=0 + ) + + def test_integer_indices_and_mixed_outputs(self): + class Select(torch.nn.Module): + def forward(self, x, indices): + selected = torch.index_select(x, 1, indices) + positive = selected > 0 + return ( + torch.where(positive, selected, -selected), + selected.argmax(dim=1), + positive, + ) + + model = Select() + x = torch.tensor([[-3.0, 5.0, 2.0], [7.0, -2.0, -4.0]]) + for dtype in (torch.int32, torch.int64): + with self.subTest(dtype=dtype): + indices = torch.tensor([2, 0], dtype=dtype) + with self._lower(model, (x, indices)) as runtime: + for current in (indices, indices.flip(0)): + actual = runtime.forward((x, current)) + for output, expected in zip( + actual, model(x, current), strict=True + ): + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_boolean_inputs_and_integer_constants(self): + class MaskedAdd(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.tensor([2**40 + 3, -7])) + + def forward(self, x, mask): + return torch.where(mask, x + self.offset, x), torch.logical_not(mask) + + model = MaskedAdd() + x = torch.tensor([2**40 + 1, -(2**40 + 5)]) + mask = torch.tensor([True, False]) + with self._lower(model, (x, mask)) as runtime: + for current in (mask, ~mask): + actual = runtime.forward((x, current)) + for output, expected in zip(actual, model(x, current), strict=True): + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_empty_tensors_execute(self): + class Add(torch.nn.Module): + def forward(self, x): + return x + 1 + + model = Add() + x = torch.empty(2, 0, 3) + with self._lower(model, (x,)) as runtime: + actual = runtime.forward((x,))[0] + torch.testing.assert_close(actual, model(x)) + + def test_neg_view_copy_routes_to_neg(self): + class NegView(torch.nn.Module): + def forward(self, x): + return torch.ops.aten._neg_view_copy.default(x) + + model = NegView() + x = torch.randn(2, 3) + with self._lower(model, (x,)) as runtime: + actual = runtime.forward((x,))[0] + torch.testing.assert_close(actual, model(x)) + + def test_neg_view_copy_empty_bool(self): + class NegView(torch.nn.Module): + def forward(self, x): + return torch.ops.aten._neg_view_copy.default(x) + + model = NegView() + x = torch.empty(0, dtype=torch.bool) + expected = model(x) + with self._lower(model, (x,)) as runtime: + actual = runtime.forward((x,))[0] + torch.testing.assert_close(actual, expected) + + def test_optional_integer_index_lists(self): + class Index(torch.nn.Module): + def __init__(self, accumulate): + super().__init__() + self.accumulate = accumulate + + def forward(self, x, indices, values): + updated = torch.index_put(x, (indices,), values, self.accumulate) + return updated, updated[indices], updated[:, indices] + + x = torch.arange(9, dtype=torch.float32).reshape(3, 3) + indices = torch.tensor([0, 2]) + values = torch.tensor([[10.0, 20.0, 30.0], [-1.0, -2.0, -3.0]]) + inputs = (x, indices, values) + for accumulate in (False, True): + with self.subTest(accumulate=accumulate): + model = Index(accumulate) + with self._lower(model, inputs) as runtime: + actual = runtime.forward(inputs) + for output, expected in zip(actual, model(*inputs), strict=True): + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_batch_norm_fused_into_conv_and_linear(self): + class ConvBnLinearBn(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.bn1 = torch.nn.BatchNorm2d(4) + self.linear = torch.nn.Linear(4, 8) + self.bn2 = torch.nn.BatchNorm1d(8) + + def forward(self, x): + y = self.bn1(self.conv(x)) + return self.bn2(self.linear(y.mean(dim=(2, 3)))) + + torch.manual_seed(0) + model = ConvBnLinearBn().eval() + for bn in (model.bn1, model.bn2): + bn.running_mean.uniform_(-1, 1) + bn.running_var.uniform_(0.5, 2) + bn.weight.data.uniform_(0.5, 2) + bn.bias.data.uniform_(-1, 1) + inputs = (torch.randn(2, 4, 8, 8),) + with self._lower(model, inputs) as runtime: + torch.testing.assert_close( + runtime.forward(inputs)[0], model(*inputs), rtol=1e-5, atol=1e-5 + ) diff --git a/backends/cpu/workflows/README.md b/backends/cpu/workflows/README.md new file mode 100644 index 00000000000..e16b24679c1 --- /dev/null +++ b/backends/cpu/workflows/README.md @@ -0,0 +1,27 @@ +# CPU torchvision exports + +List classifiers or export a PTE with the registered CPU or XNNPACK FP32 recipe: + +```sh +python -m executorch.backends.cpu.workflows.export --list-models +python -m executorch.backends.cpu.workflows.export \ + --model convnext_tiny --backend cpu --output /tmp/convnext-cpu \ + --source-revision REV +``` + +Inside fbsource, run `buck2 run fbcode//executorch/backends/cpu:export_torchvision --` +with the same arguments. + +The default model is `convnext_tiny` with pretrained weights. Use +`--random-weights` for deterministic random weights, or `--pt2` to reuse a capture. +Input size defaults to the pretrained weights' crop size, or the model default +for random weights; `--image-size` overrides it. For random-weight ViTs, this +also configures the model's input size. Pretrained ViTs reject sizes that do not +match their weights. Reused captures keep their recorded input shape, and the +model name defaults to the capture's file stem. +Exports include `model.pte`, the captured `.pt2`, `input.bin`, `eager.bin`, +`golden.bin`, and an `export.json` manifest with artifact hashes and preprocessing +provenance. `input.bin` is the raw float32 input; `eager.bin` and `golden.bin` +hold the eager and edge-program logits. A runner can feed `input.bin` to +`model.pte` and compare its output against `eager.bin` to validate the PTE before +measuring latency. diff --git a/backends/cpu/workflows/__init__.py b/backends/cpu/workflows/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/cpu/workflows/__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/backends/cpu/workflows/export.py b/backends/cpu/workflows/export.py new file mode 100644 index 00000000000..705572b34fb --- /dev/null +++ b/backends/cpu/workflows/export.py @@ -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. + +from __future__ import annotations + +import argparse +import hashlib +import json +import logging +import math +import shutil +from collections import Counter +from pathlib import Path +from typing import Any + +import torch +import torchvision +from executorch.backends.cpu.partitioner import CPU_DELEGATE_VERSION +from executorch.backends.cpu.recipes import CPURecipeType +from executorch.backends.native.serialization import deserialize_graph +from executorch.backends.native.serialization.schema import OpKind +from executorch.backends.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack.recipes import XNNPackRecipeType +from executorch.export import export, ExportRecipe, ExportSession +from torchvision.models import get_model, get_model_weights, list_models + + +logger: logging.Logger = logging.getLogger(__name__) +RECIPE_TYPES = { + "cpu": CPURecipeType.FP32, + "xnnpack": XNNPackRecipeType.FP32, +} +DEFAULT_MODEL = "convnext_tiny" +DEFAULT_IMAGE_SIZE = 224 +IMAGE_SIZE_OVERRIDES = {"inception_v3": 299} + + +def classification_models() -> list[str]: + return sorted(list_models(torchvision.models)) + + +def image_size_for( + model: str, override: int | None, *, random_weights: bool = False +) -> int: + if override is not None: + return override + if not random_weights: + return get_model_weights(model).DEFAULT.transforms().crop_size[0] + return IMAGE_SIZE_OVERRIDES.get(model, DEFAULT_IMAGE_SIZE) + + +def synthetic_image(shape: tuple[int, ...]) -> torch.Tensor: + return (torch.arange(math.prod(shape)) % 251).float().div(250).reshape(shape) + + +def capture( + args: argparse.Namespace, +) -> tuple[torch.export.ExportedProgram, torch.Tensor]: + torch.manual_seed(0) + torch.set_num_threads(1) + destination = args.output / f"{args.model}.pt2" + if args.pt2: + exported = torch.export.load(args.pt2) + input_name = exported.graph_signature.user_inputs[0] + input_node = next( + node for node in exported.graph.nodes if node.name == input_name + ) + image = synthetic_image(tuple(input_node.meta["val"].shape)) + if args.pt2.resolve() != destination.resolve(): + shutil.copyfile(args.pt2, destination) + else: + size = image_size_for( + args.model, args.image_size, random_weights=args.random_weights + ) + weights = None if args.random_weights else get_model_weights(args.model).DEFAULT + model_kwargs = {"image_size": size} if args.model.startswith("vit_") else {} + model = get_model(args.model, weights=weights, **model_kwargs).eval() + image = synthetic_image((1, 3, size, size)) + exported = torch.export.export(model, (image,)) + torch.export.save(exported, destination) + with torch.no_grad(): + exported.module()(image).contiguous().numpy().tofile(args.output / "eager.bin") + image.numpy().tofile(args.output / "input.bin") + return exported, image + + +def write_manifest( + args: argparse.Namespace, + spec: dict[str, Any], + exported: torch.export.ExportedProgram, + recipe: ExportRecipe, + session: ExportSession, + image: torch.Tensor, +) -> None: + is_cpu = spec["backend"] == "cpu" + passes = [] if is_cpu else XNNPACKPassManager(exported).passes + files = {} + artifact = "model.pte" + for name in (f"{args.model}.pt2", "input.bin", "eager.bin", "golden.bin", artifact): + path = args.output / name + with path.open("rb") as handle: + digest = hashlib.sha256() + for chunk in iter(lambda: handle.read(8 << 20), b""): + digest.update(chunk) + files[name] = {"sha256": digest.hexdigest(), "bytes": path.stat().st_size} + program = session.get_executorch_program() + semantic_nodes = None + if is_cpu: + semantic_nodes = [ + node.name + for delegate in program.execution_plan[0].delegates + if delegate.id == "CpuBackend" + for node in deserialize_graph( + bytes(program.backend_delegate_data[delegate.processed.index].data) + ).nodes + if node.op_kind == OpKind.CALL_FUNCTION + ] + manifest = { + "version": 1, + "model": args.model, + "runtime": "et", + "artifact": artifact, + "preprocess": "CpuBackend" if is_cpu else "XnnpackBackend", + "source_revision": args.source_revision, + "versions": { + "torch": torch.__version__, + "torchvision": torchvision.__version__, + }, + "recipe": spec, + "quantization": None, + "cpu_delegate_version": CPU_DELEGATE_VERSION if is_cpu else None, + "recipe_passes": {"aten": [], "edge": [], "edge_manager": []}, + "backend_default_passes": [f"{p.__module__}.{p.__qualname__}" for p in passes], + "preserved_ops": sorted( + { + str(op) + for partitioner in recipe.lowering_recipe.partitioners + for op in partitioner.ops_to_not_decompose(exported)[0] + } + ), + "stages": [stage.name for stage in session.get_stage_artifacts()], + "delegates": dict(Counter(d.id for d in program.execution_plan[0].delegates)), + "semantic_nodes": semantic_nodes, + "input": { + "shape": list(image.shape), + "dtype": "float32", + "pattern": "(index % 251) / 250", + }, + "files": files, + } + for name, value in (("recipe.json", spec), ("export.json", manifest)): + (args.output / name).write_text(json.dumps(value, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path) + parser.add_argument("--backend", choices=tuple(RECIPE_TYPES), default="cpu") + parser.add_argument("--pt2", type=Path) + parser.add_argument("--model") + parser.add_argument("--image-size", type=int, default=None) + parser.add_argument("--source-revision") + parser.add_argument("--random-weights", action="store_true") + parser.add_argument( + "--list-models", + action="store_true", + help="Print supported torchvision classification models and exit", + ) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO) + if args.list_models: + for name in classification_models(): + print(name) + return + if args.output is None or args.source_revision is None: + parser.error("--output and --source-revision are required") + if args.pt2: + if args.random_weights or args.image_size is not None: + parser.error("--random-weights and --image-size cannot be used with --pt2") + args.model = args.model or args.pt2.stem + else: + args.model = args.model or DEFAULT_MODEL + if args.model not in classification_models(): + parser.error(f"Unknown torchvision model {args.model!r}; see --list-models") + if args.image_size is not None and args.image_size <= 0: + parser.error("--image-size must be positive") + recipe_type = RECIPE_TYPES[args.backend] + spec = { + "version": 1, + "backend": args.backend, + "recipe": recipe_type.value, + "kwargs": {}, + } + args.output.mkdir(parents=True, exist_ok=True) + exported, image = capture(args) + recipe = ExportRecipe.get_recipe(recipe_type) + session = export(exported, export_recipe=recipe) + edge_module = session.get_edge_program_manager().exported_program().module() + with torch.no_grad(): + edge_module(image).contiguous().numpy().tofile(args.output / "golden.bin") + artifact = args.output / "model.pte" + artifact.write_bytes(session.get_pte_buffer()) + write_manifest(args, spec, exported, recipe, session, image) + logger.info("Exported %s to %s", spec["recipe"], artifact) + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/tests/test_device_placement.py b/backends/cuda/tests/test_device_placement.py new file mode 100644 index 00000000000..b54242802a1 --- /dev/null +++ b/backends/cuda/tests/test_device_placement.py @@ -0,0 +1,147 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import tempfile +import unittest +from pathlib import Path + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.passes.propagate_device_pass import ( + PropagateDeviceConfig, + TARGET_DEVICE_COMPILE_SPEC_KEY, +) + + +class EachAddOnCuda(Partitioner): + """Lowers each add to its own CUDA delegate. Parameters stay outside, so + they cross the delegate boundary as in a partial lowering.""" + + def partition(self, exported_program) -> PartitionResult: + spec = DelegationSpec( + CudaBackend.__name__, + [ + CudaBackend.generate_method_name_compile_spec("forward"), + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, b"cuda:0"), + ], + ) + tags = {} + adds = [ + node + for node in exported_program.graph_module.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.edge.aten.add.Tensor + ] + for i, node in enumerate(adds): + node.meta["delegation_tag"] = f"add{i}" + tags[f"add{i}"] = spec + return PartitionResult( + tagged_exported_program=exported_program, partition_tags=tags + ) + + +def load(model, inputs, directory, config=None): + """Lower model with EachAddOnCuda and load its forward method from a file + in directory.""" + from executorch.runtime import Runtime + + edge = to_edge_transform_and_lower( + torch.export.export(model, inputs, strict=True), + partitioner=[EachAddOnCuda()], + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + path = Path(directory) / "model.pte" + path.write_bytes(edge.to_executorch(config or ExecutorchBackendConfig()).buffer) + return Runtime.get().load_program(path).load_method("forward") + + +def run(model, inputs) -> list[torch.Tensor]: + """Lower model with EachAddOnCuda, run it in the runtime, return its outputs.""" + with tempfile.TemporaryDirectory() as directory: + method = load(model, inputs, directory) + return [output.cpu() for output in method.execute(list(inputs))] + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class CudaDevicePlacementTest(unittest.TestCase): + def test_transposed_parameter_read_by_a_delegate(self): + """The program stores a transposed parameter contiguous and copies + those bytes to the device. to_edge makes the parameter contiguous + before the delegate is compiled, so the delegate expects that layout.""" + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + weight = torch.arange(6.0).reshape(3, 2).t() + self.weight = torch.nn.Parameter(weight, requires_grad=False) + + def forward(self, x): + return torch.add(x, self.weight) + + model = Model() + x = torch.arange(6.0).reshape(2, 3) * 10 + (output,) = run(model, (x,)) + torch.testing.assert_close(output, model(x)) + + def test_constant_read_by_two_delegates(self): + """Two delegates read a constant stored in the program from outside + them, so the program copies it to the device for them.""" + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("table", torch.arange(6.0).reshape(2, 3)) + + def forward(self, x): + return torch.add(x, self.table), torch.add(x * 2, self.table) + + model = Model() + x = torch.arange(6.0).reshape(2, 3) * 10 + for output, expected in zip(run(model, (x,)), model(x), strict=True): + torch.testing.assert_close(output, expected) + + def test_input_on_device_also_read_on_host(self): + """With skip_h2d_for_method_inputs, the caller passes the input on the + device. A delegate reads it there, and a host op reads a copy of it in + host memory.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y), torch.sin(x) + + config = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + model = Model() + inputs = (torch.randn(2, 3), torch.randn(2, 3)) + with tempfile.TemporaryDirectory() as directory: + method = load(model, inputs, directory, config) + outputs = method.execute([t.cuda() for t in inputs]) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) diff --git a/backends/native/BUCK b/backends/native/BUCK index 18026b82d45..e0afa1058fc 100644 --- a/backends/native/BUCK +++ b/backends/native/BUCK @@ -36,6 +36,7 @@ fbcode_target( "fbsource//third-party/pypi/flatbuffers:flatbuffers", "//caffe2:torch", "//executorch/backends/transforms:collapse_view_copy", + "//executorch/backends/transforms:fuse_batch_norm_with_conv", "//executorch/backends/transforms:fuse_gqa_with_sdpa", "//executorch/backends/transforms:fuse_rms_norm", "//executorch/backends/transforms:normalize_sdpa_input_rank", diff --git a/backends/native/CMakeLists.txt b/backends/native/CMakeLists.txt index e111c8f5453..2b4b8ec1a7c 100644 --- a/backends/native/CMakeLists.txt +++ b/backends/native/CMakeLists.txt @@ -36,7 +36,7 @@ install( PATTERN "ggml" EXCLUDE ) -if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) +if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) add_subdirectory(examples) endif() diff --git a/backends/native/README.md b/backends/native/README.md index a784537989f..6fafb89f11d 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -33,6 +33,11 @@ quantize/dequantize ops. TorchAO int4 group-wise weights of `linear` and stored packed with their scheme in tensor metadata. Each method must lower wholly to one delegate; `to_native` raises otherwise. +Immutable constant keys are namespaced as `/`, even for a single +method. A key mutated by any method becomes `state/` in every method, so +methods share its runtime state. Method names may not be `state` or start with +`state/`. `to_native` raises if namespaced keys collide. + A `.ptn` is an uncompressed zip with fixed entry names: | Entry | Contents | @@ -110,10 +115,22 @@ target_link_libraries(my_app PRIVATE native_api native_vulkan_engine) With `EXECUTORCH_BUILD_TESTS=ON`, `native_runtime_test` covers the shared runtime and API. `native_vulkan_passes_test` covers the Vulkan graph rewrites without -requiring a GPU. When the vision runner below is enabled, -`native_runner_utils_test` covers its shared runner utilities. Run these with +requiring a GPU. When either runner below is enabled, +`native_runner_utils_test` covers their shared runner utilities. Run these with `ctest --test-dir cmake-out -R native_`. +The `Test Native Engines` CI workflow has separate Linux jobs for the shared +runtime, GGML, and Vulkan, with each test suite run once. To reproduce with +dependencies installed, run `bash .ci/scripts/test-native.sh runtime`, +`bash .ci/scripts/test-native.sh ggml`, or +`bash .ci/scripts/test-native.sh vulkan` from the OSS checkout. The GGML job +also runs `native_llm_runner_utils_test` and executes `run_hf_llm_ggml` on a tiny +seeded Hugging Face Llama with a generated byte tokenizer. Its eight generated +token IDs must match eager PyTorch, covering prefill and KV-cache decoding without +downloading a model. The GGML job requires the Python export dependencies +installed by the `Test Native Engines` workflow. The Vulkan job builds +`run_classification` and `run_hf_llm_vulkan`; its tests do not require a GPU. + ### GGML engine Native GGML builds the GGML subtree of the llama.cpp submodule at @@ -170,3 +187,31 @@ Export the `.ptn`, input, reference logits, and labels with [`export.py`](../../examples/native/vision/export.py). The runner prints the top-5 logits and, when `--expected` is given, exits nonzero if the logits differ from the reference beyond `--rtol`/`--atol` (both default to `1e-3`). + +## Native LLM runner + +`EXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON` builds `run_hf_llm_vulkan` with Native +Vulkan and `run_hf_llm_ggml` with Native GGML. It requires gflags and +`tokenizers::tokenizers` (the vendored tokenizer is built automatically in an +OSS checkout; otherwise provide an installed tokenizers package with regex +lookahead support). It accepts Llama 3 TikToken tokenizer files. + +To build the text runner for GGML CPU, initialize the llama.cpp submodule as above: + +```sh +cmake -S . -B cmake-out -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_NATIVE=ON -DEXECUTORCH_BUILD_NATIVE_GGML=ON \ + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON +cmake --build cmake-out --target run_hf_llm_ggml +cmake --install cmake-out --prefix install +install/bin/run_hf_llm_ggml --device cpu --cpu_threads 4 \ + --package llama.ptn --tokenizer tokenizer.model \ + --prompt "The capital of France is" --max_new_tokens 32 +``` + +The GGML runner accepts `--device cpu|gpu`, `--device_index`, `--cpu_threads`, +and `--max_cached_graphs` (default 1, 0 rebuilds each call). For GPU execution, +enable the desired `GGML_*` backend when configuring ExecuTorch, as above. + +With `EXECUTORCH_BUILD_TESTS=ON`, `native_llm_runner_utils_test` exercises the +LLM generation utilities. diff --git a/backends/native/examples/CMakeLists.txt b/backends/native/examples/CMakeLists.txt index 81b683a7d05..2f6235040eb 100644 --- a/backends/native/examples/CMakeLists.txt +++ b/backends/native/examples/CMakeLists.txt @@ -4,5 +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. +if(NOT TARGET gflags AND NOT TARGET gflags::gflags) + find_package(gflags REQUIRED) +endif() +if(TARGET gflags::gflags) + set(_native_runner_gflags gflags::gflags) +else() + set(_native_runner_gflags gflags) +endif() + add_subdirectory(runner_utils) -add_subdirectory(vision) +if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) + add_subdirectory(vision) +endif() +if(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) + add_subdirectory(hf_llm/runtime) +endif() diff --git a/backends/native/examples/hf_llm/runtime/CMakeLists.txt b/backends/native/examples/hf_llm/runtime/CMakeLists.txt new file mode 100644 index 00000000000..5e3db657663 --- /dev/null +++ b/backends/native/examples/hf_llm/runtime/CMakeLists.txt @@ -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. + +if(NOT TARGET tokenizers::tokenizers) + if(EXISTS "${EXECUTORCH_ROOT}/extension/llm/tokenizers/CMakeLists.txt") + set(SUPPORT_REGEX_LOOKAHEAD ON) + add_subdirectory( + ${EXECUTORCH_ROOT}/extension/llm/tokenizers + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers + ) + else() + find_package(tokenizers CONFIG REQUIRED) + endif() +endif() + +add_library(native_llm_runner_utils STATIC LlmRunnerUtils.cpp) +target_compile_features(native_llm_runner_utils PUBLIC cxx_std_20) +target_include_directories( + native_llm_runner_utils PUBLIC $ +) +add_library( + native_llm_runner STATIC + LlmRunner.cpp + ${EXECUTORCH_ROOT}/examples/models/llama/tokenizer/llama_tiktoken.cpp +) +target_link_libraries( + native_llm_runner + PUBLIC native_engine + PRIVATE ${_native_runner_gflags} native_llm_runner_utils native_package + native_runner_utils native_runtime tokenizers::tokenizers +) +foreach(target native_llm_runner native_llm_runner_utils) + target_compile_options( + ${target} PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) +endforeach() + +set(_llm_runners) +if(EXECUTORCH_BUILD_NATIVE_VULKAN) + add_executable(run_hf_llm_vulkan run_hf_llm_vulkan.cpp) + target_link_libraries(run_hf_llm_vulkan PRIVATE native_vulkan_engine) + list(APPEND _llm_runners run_hf_llm_vulkan) +endif() +if(EXECUTORCH_BUILD_NATIVE_GGML) + add_executable(run_hf_llm_ggml run_hf_llm_ggml.cpp) + target_link_libraries(run_hf_llm_ggml PRIVATE native_ggml_engine) + list(APPEND _llm_runners run_hf_llm_ggml) +endif() +foreach(target IN LISTS _llm_runners) + target_link_libraries( + ${target} PRIVATE native_llm_runner ${_native_runner_gflags} + ) + target_compile_options( + ${target} PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) + install(TARGETS ${target} RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) +endforeach() diff --git a/backends/native/examples/vision/CMakeLists.txt b/backends/native/examples/vision/CMakeLists.txt index af7ea765224..76bfdb282f3 100644 --- a/backends/native/examples/vision/CMakeLists.txt +++ b/backends/native/examples/vision/CMakeLists.txt @@ -4,19 +4,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -if(NOT TARGET gflags AND NOT TARGET gflags::gflags) - find_package(gflags REQUIRED) -endif() -if(TARGET gflags::gflags) - set(_native_vision_gflags gflags::gflags) -else() - set(_native_vision_gflags gflags) -endif() - add_executable(run_classification run_classification.cpp) target_link_libraries( run_classification - PRIVATE ${_native_vision_gflags} native_engine native_package + PRIVATE ${_native_runner_gflags} native_engine native_package native_runner_utils native_runtime native_vulkan_engine ) target_compile_options( diff --git a/backends/native/passes/__init__.py b/backends/native/passes/__init__.py index e043b9fae89..7b1508747d1 100644 --- a/backends/native/passes/__init__.py +++ b/backends/native/passes/__init__.py @@ -27,7 +27,7 @@ default pass list. """ -from typing import List, Union +from typing import Collection, List, Union from executorch.backends.native.passes.fuse_rope import FuseRoPEPass from executorch.backends.native.passes.reinplace import ( @@ -43,6 +43,9 @@ ) from executorch.backends.transforms.collapse_view_copy import CollapseViewCopyPass +from executorch.backends.transforms.fuse_batch_norm_with_conv import ( + FuseBatchNormWithConvProgramPass, +) from executorch.backends.transforms.fuse_gqa_with_sdpa import FuseGQAWithSDPAPass from executorch.backends.transforms.fuse_rms_norm import FuseRMSNormPass from executorch.backends.transforms.normalize_sdpa_input_rank import ( @@ -56,6 +59,7 @@ "backend_inplace_aten_variants", "BACKEND_INPLACE_OPS", "CollapseViewCopyPass", + "FuseBatchNormWithConvProgramPass", "FuseGQAWithSDPAPass", "FuseRMSNormPass", "FuseRoPEPass", @@ -67,14 +71,19 @@ ] -def get_default_passes() -> List[Union[ExportPass, ExportedProgramPassBase]]: +def get_default_passes( + *, mutable_buffers: Collection[str] = () +) -> List[Union[ExportPass, ExportedProgramPassBase]]: """Passes enabled by default for the native backend. + Frozen evaluation BatchNorm folds into convolution first. mutable_buffers + additionally protects buffers mutated by other methods in the same package. GQA fusion runs before SDPA rank normalization, then view_copy collapsing and CSE settle the graph. Reinplace runs last, rewriting functional ops into their in-place edge forms. """ return [ + FuseBatchNormWithConvProgramPass(mutable_buffers), FuseRoPEPass(), FuseGQAWithSDPAPass(), NormalizeSDPAInputRankPass(), diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py index 827d19e6370..c841c3560e4 100644 --- a/backends/native/serialization/graph_serialize.py +++ b/backends/native/serialization/graph_serialize.py @@ -31,6 +31,7 @@ from enum import IntEnum from typing import ( Any, + Callable, cast, get_args, get_origin, @@ -1487,8 +1488,8 @@ def _iter_graph_args(value: object) -> Iterator[GraphArg]: yield from _iter_graph_args(item) -def collect_data_keys(program: Program) -> set[str]: - """Return every out-of-line data key the program references. +def map_data_keys(method: Method, fn: Callable[[str], str]) -> None: + """Replace every out-of-line data key in ``method`` with ``fn(key)``, in place. Covers constant references (NamedTensorRef.data_key) plus the quant keys (AffineGroup scale and zero-point) carried on any TensorMeta.quant, whether @@ -1496,30 +1497,39 @@ def collect_data_keys(program: Program) -> set[str]: recursing into HOP subgraphs wherever a schema field carries a ``GraphArg``. PackedQuant carries no external keys. """ - keys: set[str] = set() - def add_meta(meta: Optional[TensorMeta]) -> None: + def map_meta(meta: Optional[TensorMeta]) -> None: if meta is None or meta.quant is None: return scheme = meta.quant.scheme if isinstance(scheme, AffineGroup): - keys.add(scheme.scale_data_key) + scheme.scale_data_key = fn(scheme.scale_data_key) if scheme.zero_point_data_key: - keys.add(scheme.zero_point_data_key) + scheme.zero_point_data_key = fn(scheme.zero_point_data_key) def visit_graph(graph: Graph) -> None: for value in graph.tensor_values or []: - add_meta(value.meta) + map_meta(value.meta) for node in graph.nodes: for graph_arg in _iter_graph_args(node): visit_graph(graph_arg.graph) - for method in program.methods: - for constant in method.constants or []: - keys.add(constant.data_key) - add_meta(constant.meta) - visit_graph(method.graph) + for constant in method.constants or []: + constant.data_key = fn(constant.data_key) + map_meta(constant.meta) + visit_graph(method.graph) + +def collect_data_keys(program: Program) -> set[str]: + """Return every out-of-line data key the program references.""" + keys: set[str] = set() + + def add(key: str) -> str: + keys.add(key) + return key + + for method in program.methods: + map_data_keys(method, add) return keys diff --git a/backends/native/test/BUCK b/backends/native/test/BUCK index a0d85273697..c3f3eeb546d 100644 --- a/backends/native/test/BUCK +++ b/backends/native/test/BUCK @@ -3,6 +3,17 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_test, + name = "test_fuse_conv_batch_norm", + srcs = ["test_fuse_conv_batch_norm.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/native:lib", + "//executorch/exir:lib", + ], +) + fbcode_target( _kind = runtime.python_test, name = "test_serialize", diff --git a/backends/native/test/CMakeLists.txt b/backends/native/test/CMakeLists.txt index 68b9b119f94..fb321e8736e 100644 --- a/backends/native/test/CMakeLists.txt +++ b/backends/native/test/CMakeLists.txt @@ -79,3 +79,22 @@ if(TARGET native_runner_utils) $<$>:-fexceptions> ) endif() + +if(TARGET native_llm_runner_utils) + add_executable( + native_llm_runner_utils_test + ../examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp + ) + target_link_libraries( + native_llm_runner_utils_test PRIVATE native_llm_runner_utils + GTest::gtest_main + ) + add_test(NAME native_llm_runner_utils_test + COMMAND native_llm_runner_utils_test + ) + target_compile_options( + native_llm_runner_utils_test + PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) +endif() diff --git a/backends/native/test/runtime/ggml/test_ggml_ops.py b/backends/native/test/runtime/ggml/test_ggml_ops.py index 20fe7013937..54d31653a77 100644 --- a/backends/native/test/runtime/ggml/test_ggml_ops.py +++ b/backends/native/test/runtime/ggml/test_ggml_ops.py @@ -173,10 +173,9 @@ def forward(self, x): return self.relu6(self.bn(self.conv(x))) -class UnfoldableBN(torch.nn.Module): - """Batch norms no convolution can absorb: one reads a convolution output - that is also read elsewhere, one the sole reader of a convolution whose - weight is shared, one a ReLU.""" +class SharedConvBN(torch.nn.Module): + """Batch norms reading a shared convolution output, a convolution with + shared weights, and a ReLU.""" def __init__(self): super().__init__() @@ -305,6 +304,7 @@ def _run_ggml( resize: bool = False, max_cached_graphs: int = 1, graph_ops: collections.Counter | None = None, + method: str = "forward", ): """Runs each step's inputs in order on one executable and returns each step's outputs as (bytes, shape) pairs. `graph_ops`, if given, collects @@ -325,6 +325,7 @@ def _run_ggml( command = [ f"--device={os.environ.get('GGML_PTN_DEVICE', 'cpu')}", f"--package={package}", + f"--method={method}", f"--inputs={';'.join(input_steps)}", f"--outputs={';'.join(output_steps)}", f"--max_cached_graphs={max_cached_graphs}", @@ -732,15 +733,16 @@ def test_pointwise_conv_bn_relu6(self): ) self.assertEqual(ops["SQRT"] + ops["MUL"], 0) - def test_unfoldable_batch_norms(self): + def test_batch_norms_with_shared_convolution(self): ops = collections.Counter() self._check( - UnfoldableBN(), (torch.randn(1, 4, 5, 6),), graph_ops=ops, **_MATMUL_TOL + SharedConvBN(), (torch.randn(1, 4, 5, 6),), graph_ops=ops, **_MATMUL_TOL ) - # Scale and shift come from load; each graph has three batch norms - # (one scale multiply each) and two convolutions. + # AOT folds the shared-weight branch; two batch norm scales and two + # convolutions remain in each graph. self.assertEqual(ops["SQRT"], 0) - self.assertEqual(2 * ops["MUL"], 3 * ops["MUL_MAT"]) + self.assertGreater(ops["MUL_MAT"], 0) + self.assertEqual(ops["MUL"], ops["MUL_MAT"]) def test_batch_norm_of_computed_statistics(self): # Statistics computed in the graph cannot fold at load. @@ -767,6 +769,41 @@ def test_batch_norm_no_affine(self): norm.running_var.uniform_(0.5, 2.0) self._check(norm, (torch.randn(1, 4, 5, 6),), atol=1e-5) + def test_batch_norm_statistics_written_by_another_method(self): + class Update(torch.nn.Module): + def __init__(self, bn): + super().__init__() + self.bn = bn + + def forward(self, delta): + self.bn.running_mean.add_(delta) + return self.bn.running_mean + 1 + + model = ConvBN(groups=1, bias=True).eval() + methods = { + "predict": (model, (torch.randn(1, 8, 9, 9),)), + "update": (Update(model.bn), (torch.randn(8),)), + } + with tempfile.TemporaryDirectory() as tmp: + outdir = Path(tmp) + package = outdir / "model.ptn" + to_native( + { + name: torch.export.export(module, inputs) + for name, (module, inputs) in methods.items() + } + ).save(str(package)) + for name, (module, inputs) in methods.items(): + with self.subTest(method=name), torch.no_grad(): + expected = copy.deepcopy(module)(*inputs) + ((raw, shape),) = _run_ggml( + package, [inputs], outdir, 1, method=name + )[0] + actual = torch.frombuffer(bytearray(raw), dtype=torch.float32) + torch.testing.assert_close( + actual.reshape(shape), expected, atol=1e-4, rtol=1e-5 + ) + def test_max_pool2d(self): self._check( torch.nn.MaxPool2d(3, stride=2, padding=1), (torch.randn(1, 4, 11, 9),) diff --git a/backends/native/test/test_fuse_conv_batch_norm.py b/backends/native/test/test_fuse_conv_batch_norm.py new file mode 100644 index 00000000000..bba0f60972b --- /dev/null +++ b/backends/native/test/test_fuse_conv_batch_norm.py @@ -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. + +from __future__ import annotations + +import copy +import unittest + +import torch +from executorch.backends.native import get_default_compile_config +from executorch.backends.native.passes import get_default_passes +from executorch.exir import to_edge + + +class ConvBatchNorm(torch.nn.Module): + def __init__(self, *, groups=1, bias=True, affine=True): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1, groups=groups, bias=bias) + self.bn = torch.nn.BatchNorm2d(4, affine=affine) + with torch.no_grad(): + self.bn.running_mean.copy_(torch.tensor([-0.5, 0.2, 0.7, 1.0])) + self.bn.running_var.copy_(torch.tensor([0.3, 0.8, 1.2, 2.0])) + if affine: + self.bn.weight.copy_(torch.tensor([0.4, 0.8, 1.1, 1.5])) + self.bn.bias.copy_(torch.tensor([-0.7, 0.1, 0.3, 0.9])) + + def forward(self, x): + return self.bn(self.conv(x)) + + +class FuseConvBatchNormPassTest(unittest.TestCase): + def _edge(self, model, inputs): + return to_edge( + torch.export.export(model, inputs), + compile_config=get_default_compile_config(), + ) + + def _batch_norms(self, program): + return [ + node for node in program.graph.nodes if "batch_norm" in str(node.target) + ] + + def test_default_pipeline_folds_eval_batch_norm(self): + for config in ( + {}, + {"groups": 2, "bias": False}, + {"groups": 4, "affine": False}, + ): + with self.subTest(config=config): + model = ConvBatchNorm(**config).eval() + inputs = (torch.randn(1, 4, 8, 8),) + state = copy.deepcopy(model.state_dict()) + expected = model(*inputs) + edge = self._edge(model, inputs) + self.assertEqual(len(self._batch_norms(edge.exported_program())), 1) + program = edge.transform(get_default_passes()).exported_program() + self.assertFalse(self._batch_norms(program)) + torch.testing.assert_close(program.module()(*inputs), expected) + torch.testing.assert_close(model.state_dict(), state) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 8db23a465ef..f29ac21b7f6 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -14,6 +14,7 @@ PER_CHANNEL_GROUP_ENCODING, ) from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader +from executorch.backends.qualcomm.builders.utils import is_parameter from executorch.backends.qualcomm.qnn_preprocess import QnnBackend from executorch.backends.qualcomm.serialization.qc_schema import ( HtpArch, @@ -41,6 +42,7 @@ PartitionResult, ) from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer +from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import ExportedProgram from torch.fx.passes.infra.partitioner import Partition from torch.fx.passes.operator_support import OperatorSupportBase @@ -61,6 +63,13 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) +# dequantize ops with a CPU out variant; others cannot fall back +CPU_DQ_OPS = { + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, +} + class QnnOperatorSupport(OperatorSupportBase): def __init__( @@ -240,10 +249,42 @@ def generate_partitions( self.skip_node_id_set, self.skip_node_op_set, ) - return generate_partitions_from_list_of_nodes( + partitions = generate_partitions_from_list_of_nodes( edge_program.graph_module, op_support=self.op_support_checker, ) + return self._drop_orphaned_parameter_dequants(partitions, edge_program) + + @staticmethod + def _drop_orphaned_parameter_dequants( + partitions: List[Partition], edge_program: torch.export.ExportedProgram + ) -> List[Partition]: + """ + A parameter dequantize whose consumer is not delegated stays on CPU. A QNN + Dequantize alone rejects per-channel encodings. + """ + for partition in partitions: + stack = [ + node + for node in partition.nodes + if node.target in CPU_DQ_OPS + and node.meta.get(QCOM_BYPASS_NODE, False) + and is_parameter(node.args[0], edge_program) + ] + while stack: + node = stack.pop() + if node not in partition.nodes or any( + user in partition.nodes for user in node.users + ): + continue + partition.remove_node(node) + stack.extend( + arg + for arg in node.args + if isinstance(arg, torch.fx.Node) + and arg.meta.get(QCOM_BYPASS_NODE, False) + ) + return [partition for partition in partitions if partition.nodes] def tag_nodes( self, partitions: List[Partition], edge_program: torch.export.ExportedProgram diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..d8c5c32c103 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -783,6 +783,83 @@ def forward(self, x): if "QNN" in str(e) or "qnn" in str(e): self.skipTest(f"QNN SDK not available: {e}") + def _assert_cpu_conv_keeps_its_weight_dequantize( + self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True + ): + """A quantized conv kept off QNN must take its per-channel weight dequantize + with it. FoldQDQ force-passes that dequantize, so it used to be lowered alone + as a QNN Dequantize, which rejects per-channel encodings and fails context + binary generation (#19908).""" + + class ConvDepthwiseConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.pw1 = torch.nn.Conv2d(8, 16, 1) + self.dw = torch.nn.Conv2d(16, 16, 3, padding=1, groups=16) + self.pw2 = torch.nn.Conv2d(16, 8, 1) + + def forward(self, x): + return self.pw2(torch.relu(self.dw(torch.relu(self.pw1(x))))) + + sample_input = (torch.randn(1, 8, 16, 16),) + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + quant_dtype=quant_dtype, is_conv_per_channel=True + ) + prepared = prepare_pt2e( + torch.export.export(ConvDepthwiseConv().eval(), sample_input).module(), + quantizer, + ) + prepared(*sample_input) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=use_fp16), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + convert_pt2e(prepared), + sample_input, + compiler_specs, + skip_node_id_set=skip_node_id_set, + ) + except RuntimeError as e: + if "QNN" in str(e) or "qnn" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + + # inspected before to_executorch(), which swaps in out-variant ops + def cpu_nodes(target): + return [ + node + for node in edge.exported_program().graph.nodes + if node.op == "call_function" and node.target == target + ] + + convs = cpu_nodes(exir_ops.edge.aten.convolution.default) + self.assertEqual([conv.name for conv in convs], ["aten_convolution_default_1"]) + dequants = cpu_nodes( + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default + ) + self.assertTrue(dequants) + for dequant in dequants: + self.assertEqual(list(dequant.users), convs) + if emit: + edge.to_executorch() + + def test_skipped_conv_keeps_its_weight_dequantize_on_cpu(self): + self._assert_cpu_conv_keeps_its_weight_dequantize( + QuantDtype.use_16a8w, + use_fp16=False, + skip_node_id_set={"aten_convolution_default_1"}, + ) + + def test_rejected_conv_keeps_its_weight_dequantize_on_cpu(self): + # the HTP rejects fp16-activation/int8-weight depthwise conv. Not emitted: + # InsertCastForFpActQuantizedWeight casts the CPU conv's bias meta to fp16. + self._assert_cpu_conv_keeps_its_weight_dequantize( + QuantDtype.use_fp16a8w, use_fp16=True, emit=False + ) + def test_lowering_accepts_exported_program_without_reexport(self): """to_edge_transform_and_lower_to_qnn must lower a caller-supplied ExportedProgram as captured, never re-exporting it. diff --git a/backends/transforms/fuse_batch_norm_with_conv.py b/backends/transforms/fuse_batch_norm_with_conv.py index a936a0b77e4..f858d1c52b8 100644 --- a/backends/transforms/fuse_batch_norm_with_conv.py +++ b/backends/transforms/fuse_batch_norm_with_conv.py @@ -4,145 +4,191 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations + import operator -from typing import Optional, Tuple +from collections.abc import Collection import torch - -from executorch.backends.transforms.utils import get_param_tensor, is_param_node -from executorch.exir import ExportedProgram +from executorch.backends.transforms.utils import ( + create_constant_placeholder, + get_param_tensor, + is_param_node, +) from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass, PassResult - +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import ( + ExportedProgramPassBase, + ExportedProgramPassResult, + ExportPass, + PassResult, +) +from torch._subclasses.fake_tensor import unset_fake_temporarily +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind +from torch.fx import GraphModule, Node from torch.nn.utils.fusion import fuse_conv_bn_weights class FuseBatchNormWithConvPass(ExportPass): - """ - Batch Norm can be implemented using 1x1 Depthwise Convolution. However doing so will increase - memory usage since we serialize new weights to represent the convolution. In most cases, - Batch norm is used after convolution. The 1x1 depthwise convolution can then be fused - with the previous convolution + """Fold frozen evaluation BatchNorm into a single-consumer convolution. + + GraphModule callers retain get_attr parameters. Use + FuseBatchNormWithConvProgramPass for signature-aware lifted constants. """ - def __init__(self, exported_program: ExportedProgram): + def __init__( + self, + exported_program: ExportedProgram, + *, + mutable_buffers: Collection[str] = (), + ) -> None: super().__init__() self.exported_program = exported_program + self.mutable_buffers = frozenset(mutable_buffers) + + def call(self, graph_module: GraphModule) -> PassResult: + modified = self._fuse_graph(graph_module, lift_constants=False) + if modified: + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, modified) - def call(self, graph_module: torch.fx.GraphModule): + def _fuse_graph(self, graph_module: GraphModule, *, lift_constants: bool) -> bool: graph = graph_module.graph - counter = 0 - for conv in graph.nodes: - # We want to discover a chain of conv -> batch_norm. - # Only proceed if the current node is a conv node, and has a single - # user/successor. + for node in graph.nodes: + target = node.target + if isinstance(target, EdgeOpOverload): + target = target._op + if isinstance(target, torch._ops.OpOverload) and target._schema.is_mutable: + return False + modified = False + for conv in list(graph.nodes): if ( conv.target != exir_ops.edge.aten.convolution.default or len(conv.users) != 1 ): continue - - # The single user of conv op must be batch_norm. If not, bail. - bn = list(conv.users.keys())[0] - if ( - bn.target != exir_ops.edge.aten.native_batch_norm.default - and bn.target - != exir_ops.edge.aten._native_batch_norm_legit_no_training.default - ): - continue - + bn = next(iter(conv.users)) if not self.can_fuse(conv, bn, self.exported_program): continue - - # Get the parameters from conv op - assert len(conv.args) == 9 - - conv_weight = get_param_tensor(self.exported_program, conv.args[1]) - assert conv_weight is not None - - conv_bias = get_param_tensor(self.exported_program, conv.args[2]) - - # Get the parameters from the batchnorm op - assert ( - bn.target == exir_ops.edge.aten.native_batch_norm.default - and len(bn.args) == 8 - ) or ( - bn.target - == exir_ops.edge.aten._native_batch_norm_legit_no_training.default - and len(bn.args) == 7 - ) - bn_weight = get_param_tensor(self.exported_program, bn.args[1]) - bn_bias = get_param_tensor(self.exported_program, bn.args[2]) - - running_mean = get_param_tensor(self.exported_program, bn.args[3]) - assert running_mean is not None - - running_var = get_param_tensor(self.exported_program, bn.args[4]) - assert running_var is not None - - # args[7] for native_batch_norm, but args[6] for - # _native_batch_norm_legit_no_training (which doesn't have training - # as an arg) - eps = bn.args[-1] - - # Compute the updated weight and bias after fusing conv op - # with batchnorm op. - fused_weight, fused_bias = self._fuse_conv_bn_weights( - conv_weight, - conv_bias, - running_mean, - running_var, - eps, - bn_weight, - bn_bias, - transposed=bool(conv.args[6]), - groups=conv.args[8], + args = (conv.args[1], conv.args[2], *bn.args[1:5]) + if any( + isinstance(arg, Node) + and self.exported_program.graph_signature.inputs_to_buffers.get( + arg.name + ) + in self.mutable_buffers + for arg in args + ): + continue + weight, bias, scale, shift, mean, var = ( + get_param_tensor(self.exported_program, arg) for arg in args ) - - # Modify the graph by updating the weight and bias of conv op - # with the fused weight and bias params, and replacing all the users - # of getitem(batchnorm) with the conv op. - with graph.inserting_before(conv): - fused_weight_name = f"_fused_with_bn_weight_{counter}" - graph_module.register_parameter(fused_weight_name, fused_weight) - fused_weight_node = graph.get_attr(fused_weight_name) - fused_bias_name = f"_fused_with_bn_bias_{counter}" - graph_module.register_parameter(fused_bias_name, fused_bias) - fused_bias_node = graph.get_attr(fused_bias_name) - - # Update the weight and bias of conv op - conv_args = list(conv.args) + ([None] if len(conv.args) == 2 else []) - conv_args[1] = fused_weight_node - conv_args[2] = fused_bias_node - conv.args = tuple(conv_args) - # Remove any use of batchnorm from the graph - for user in bn.users.copy(): - assert user.target == operator.getitem + with unset_fake_temporarily(), torch.no_grad(): + fused_weight, fused_bias = self._fuse_conv_bn_weights( + weight, + bias, + mean, + var, + bn.args[-1], + scale, + shift, + transposed=bool(conv.args[6]), + groups=conv.args[8], + ) + weight_node = self._fused_constant(conv, fused_weight, lift_constants) + bias_node = self._fused_constant(conv, fused_bias, lift_constants) + conv.args = (conv.args[0], weight_node, bias_node, *conv.args[3:]) + for user in list(bn.users): user.replace_all_uses_with(conv) graph.erase_node(user) - graph.erase_node(bn) + modified = True + if modified: + graph.lint() + graph_module.recompile() + return modified + + def _fused_constant(self, conv: Node, data: torch.Tensor, lift: bool) -> Node: + data = data.detach() + stem = "_fused_bn" + name = stem + program = self.exported_program + graph = conv.graph + occupied = set(program.state_dict) | set(program.constants) + occupied.update(graph._graph_namespace._used_names) + index = 0 + while name in occupied: + index += 1 + name = f"{stem}_{index}" + if lift: + with graph.inserting_before(next(iter(graph.nodes))): + return create_constant_placeholder( + program, graph, name, InputKind.CONSTANT_TENSOR, data + ) + graph.owning_module.register_parameter( + name, torch.nn.Parameter(data, requires_grad=False) + ) + with graph.inserting_before(conv): + return graph.get_attr(name) - counter += 1 - - graph_module.recompile() - # To Regenerate meta data and shape information, retrace module - graph_module = super().call(graph_module).graph_module - - return PassResult(graph_module, True) + @staticmethod + def can_fuse(conv: Node, bn: Node, program: ExportedProgram) -> bool: + if ( + conv.target != exir_ops.edge.aten.convolution.default + or len(conv.args) != 9 + or conv.kwargs + or len(conv.users) != 1 + ): + return False + if bn.target == exir_ops.edge.aten.native_batch_norm.default: + if len(bn.args) != 8 or bn.args[5] is not False: + return False + elif ( + bn.target == exir_ops.edge.aten._native_batch_norm_legit_no_training.default + ): + if len(bn.args) != 7: + return False + else: + return False + if ( + bn.kwargs + or bn.args[0] is not conv + or not bn.users + or any( + user.target is not operator.getitem or user.args != (bn, 0) + for user in bn.users + ) + or conv.args[1] is None + or bn.args[3] is None + or bn.args[4] is None + ): + return False + signature = program.graph_signature + for arg in (conv.args[1], conv.args[2], *bn.args[1:5]): + if arg is None: + continue + if not isinstance(arg, Node) or not is_param_node(program, arg): + return False + if ( + signature.inputs_to_buffers.get(arg.name) + in signature.buffers_to_mutate.values() + ): + return False + return True @staticmethod def _fuse_conv_bn_weights( conv_weight: torch.Tensor, - conv_bias: Optional[torch.Tensor], + conv_bias: torch.Tensor | None, running_mean: torch.Tensor, running_var: torch.Tensor, eps: float, - bn_weight: Optional[torch.Tensor], - bn_bias: Optional[torch.Tensor], + bn_weight: torch.Tensor | None, + bn_bias: torch.Tensor | None, transposed: bool, groups: int, - ) -> Tuple[torch.nn.Parameter, torch.nn.Parameter]: + ) -> tuple[torch.nn.Parameter, torch.nn.Parameter]: """ A transposed conv weight is [in, out/groups, *kernel], so its output channels sit on dim 1, split into one block per group along dim 0. @@ -187,34 +233,17 @@ def _fuse_conv_bn_weights( ) return fused_weight, fused_bias - @staticmethod - def can_fuse( - conv: torch.fx.Node, bn: torch.fx.Node, program: ExportedProgram - ) -> bool: - """ - Determine whether a batch norm node can be fused with a preceding conv node. - """ - - # All the users of batchnorm node must be getitem ops. batchnorm - # returns a 3-element tuple. Each user must only access the first - # element of the tuple. - if [ - (user.target == operator.getitem and user.args[1] == 0) for user in bn.users - ].count(False): - return False - - conv_weights = conv.args[1] - bn_weights = bn.args[1] - # Check that the weights for conv and batchnorm are both params - if not isinstance(conv_weights, torch.fx.Node) or not isinstance( - bn_weights, torch.fx.Node - ): - return False +class FuseBatchNormWithConvProgramPass(ExportedProgramPassBase): + """ExportedProgram adapter for the shared Conv–BatchNorm transformation.""" - if [is_param_node(program, node) for node in {conv_weights, bn_weights}].count( - False - ): - return False + def __init__(self, mutable_buffers: Collection[str] = ()) -> None: + super().__init__() + self.mutable_buffers = frozenset(mutable_buffers) - return True + def call(self, program: ExportedProgram) -> ExportedProgramPassResult: + fusion = FuseBatchNormWithConvPass( + program, mutable_buffers=self.mutable_buffers + ) + modified = fusion._fuse_graph(program.graph_module, lift_constants=True) + return ExportedProgramPassResult(program, modified) diff --git a/backends/transforms/fuse_batch_norm_with_linear.py b/backends/transforms/fuse_batch_norm_with_linear.py new file mode 100644 index 00000000000..66c61d78977 --- /dev/null +++ b/backends/transforms/fuse_batch_norm_with_linear.py @@ -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. + +import operator + +import torch + +from executorch.backends.transforms.utils import get_param_tensor, is_param_node +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.nn.utils.fusion import fuse_linear_bn_weights + + +class FuseBatchNormWithLinearPass(ExportPass): + """ + Fold an inference batch norm into the preceding linear's weight and bias. + """ + + def __init__(self, exported_program: ExportedProgram): + super().__init__() + self.exported_program = exported_program + + def call(self, graph_module: torch.fx.GraphModule): + graph = graph_module.graph + counter = 0 + for linear in graph.nodes: + # We want to discover a chain of linear -> batch_norm. + # Only proceed if the current node is a linear node, and has a + # single user/successor. + if ( + linear.target != exir_ops.edge.aten.linear.default + or len(linear.users) != 1 + ): + continue + + # The single user of linear op must be batch_norm. If not, bail. + bn = list(linear.users.keys())[0] + if ( + bn.target != exir_ops.edge.aten.native_batch_norm.default + and bn.target + != exir_ops.edge.aten._native_batch_norm_legit_no_training.default + ): + continue + + if not self.can_fuse(linear, bn, self.exported_program): + continue + + # Linear has 2 or 3 args: (input, weight, bias). + assert len(linear.args) in (2, 3) + + linear_weight = get_param_tensor(self.exported_program, linear.args[1]) + assert linear_weight is not None + + linear_bias = ( + get_param_tensor(self.exported_program, linear.args[2]) + if len(linear.args) == 3 + else None + ) + + # Get the parameters from the batchnorm op + assert ( + bn.target == exir_ops.edge.aten.native_batch_norm.default + and len(bn.args) == 8 + ) or ( + bn.target + == exir_ops.edge.aten._native_batch_norm_legit_no_training.default + and len(bn.args) == 7 + ) + bn_weight = get_param_tensor(self.exported_program, bn.args[1]) + bn_bias = get_param_tensor(self.exported_program, bn.args[2]) + + running_mean = get_param_tensor(self.exported_program, bn.args[3]) + assert running_mean is not None + + running_var = get_param_tensor(self.exported_program, bn.args[4]) + assert running_var is not None + + if bn_bias is None: + bn_bias = torch.zeros_like(running_mean) + + # args[7] for native_batch_norm, but args[6] for + # _native_batch_norm_legit_no_training (which doesn't have training + # as an arg) + eps = bn.args[-1] + + # Compute the updated weight and bias after fusing linear op + # with batchnorm op. + fused_weight, fused_bias = fuse_linear_bn_weights( + linear_weight, + linear_bias, + running_mean, + running_var, + eps, + bn_weight, + bn_bias, + ) + + # Modify the graph by updating the weight and bias of linear op + # with the fused weight and bias params, and replacing all the users + # of getitem(batchnorm) with the linear op. + with graph.inserting_before(linear): + fused_weight_name = f"_fused_with_bn_weight_{counter}" + graph_module.register_parameter(fused_weight_name, fused_weight) + fused_weight_node = graph.get_attr(fused_weight_name) + fused_bias_name = f"_fused_with_bn_bias_{counter}" + graph_module.register_parameter(fused_bias_name, fused_bias) + fused_bias_node = graph.get_attr(fused_bias_name) + + # Update the weight and bias of linear op + linear_args = list(linear.args) + ([None] if len(linear.args) == 2 else []) + linear_args[1] = fused_weight_node + linear_args[2] = fused_bias_node + linear.args = tuple(linear_args) + # Remove any use of batchnorm from the graph + for user in bn.users.copy(): + assert user.target == operator.getitem + user.replace_all_uses_with(linear) + graph.erase_node(user) + + graph.erase_node(bn) + + counter += 1 + + if counter == 0: + return PassResult(graph_module, False) + + graph_module.recompile() + # To Regenerate meta data and shape information, retrace module + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, True) + + @staticmethod + def can_fuse( + linear: torch.fx.Node, bn: torch.fx.Node, program: ExportedProgram + ) -> bool: + """ + Determine whether a batch norm node can be fused with a preceding linear node. + """ + + # All the users of batchnorm node must be getitem ops. batchnorm + # returns a 3-element tuple. Each user must only access the first + # element of the tuple. + if [ + (user.target == operator.getitem and user.args[1] == 0) for user in bn.users + ].count(False): + return False + + if bn.target == exir_ops.edge.aten.native_batch_norm.default and bn.args[5]: + return False + + # Linear and batch norm share their channel axis only for rank-2 outputs. + linear_output = linear.meta.get("val") + if not isinstance(linear_output, torch.Tensor) or linear_output.dim() != 2: + return False + + params = (linear.args[1], bn.args[1], bn.args[3], bn.args[4]) + params += tuple( + param for param in linear.args[2:3] + bn.args[2:3] if param is not None + ) + return all( + isinstance(param, torch.fx.Node) and is_param_node(program, param) + for param in params + ) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 6b153c0fc6c..8314bbb5695 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -129,7 +129,21 @@ def define_common_targets(): ":utils", "//caffe2:torch", "//executorch/exir:pass_base", - "//executorch/exir:sym_util", + "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", + ], + ) + + runtime.python_library( + name = "fuse_batch_norm_with_linear", + srcs = ["fuse_batch_norm_with_linear.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + ":utils", + "//caffe2:torch", + "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", ], ) @@ -511,6 +525,19 @@ def define_common_targets(): ], ) + runtime.python_test( + name = "test_fuse_batch_norm_with_linear", + srcs = [ + "test/test_fuse_batch_norm_with_linear.py", + ], + deps = [ + "//caffe2:torch", + ":fuse_batch_norm_with_linear", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_quantize_fused_convbn_bias_pass", srcs = [ diff --git a/backends/transforms/test/test_fuse_batch_norm_with_conv.py b/backends/transforms/test/test_fuse_batch_norm_with_conv.py index ee9a439966e..ef5012ee675 100644 --- a/backends/transforms/test/test_fuse_batch_norm_with_conv.py +++ b/backends/transforms/test/test_fuse_batch_norm_with_conv.py @@ -9,6 +9,7 @@ import torch from executorch.backends.transforms.fuse_batch_norm_with_conv import ( FuseBatchNormWithConvPass, + FuseBatchNormWithConvProgramPass, ) from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops @@ -46,18 +47,35 @@ def _check(self, conv, bn, input_shape): x = torch.randn(input_shape) expected = module(x) - ep = to_edge( - export(module, (x,), strict=True), - compile_config=EdgeCompileConfig(_check_ir_validity=False), - ).exported_program() - FuseBatchNormWithConvPass(ep)(ep.graph_module) - - self.assertFalse(any(n.target in _BN_OPS for n in ep.graph_module.graph.nodes)) - torch.testing.assert_close(ep.module()(x), expected, rtol=1e-5, atol=1e-5) + for program_pass in (False, True): + with self.subTest(program_pass=program_pass): + ep = to_edge( + export(module, (x,), strict=True), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + if program_pass: + FuseBatchNormWithConvProgramPass()(ep) + ep.validate() + else: + FuseBatchNormWithConvPass(ep)(ep.graph_module) + self.assertFalse(any(n.target in _BN_OPS for n in ep.graph.nodes)) + torch.testing.assert_close( + ep.module()(x), expected, rtol=1e-5, atol=1e-5 + ) def test_conv2d(self): self._check(torch.nn.Conv2d(4, 8, 3), torch.nn.BatchNorm2d(8), (1, 4, 8, 8)) + def test_conv1d(self): + self._check(torch.nn.Conv1d(4, 8, 3), torch.nn.BatchNorm1d(8), (1, 4, 16)) + + def test_conv3d(self): + self._check( + torch.nn.Conv3d(4, 8, 3, bias=False), + torch.nn.BatchNorm3d(8), + (1, 4, 6, 6, 6), + ) + def test_grouped_conv2d(self): self._check( torch.nn.Conv2d(4, 8, 3, groups=2, bias=False), @@ -91,3 +109,192 @@ def test_grouped_conv_transpose2d(self): torch.nn.BatchNorm2d(6), (1, 4, 8, 8), ) + + +class ConvBatchNorm(torch.nn.Module): + def __init__(self, *, groups=1, bias=True, affine=True): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1, groups=groups, bias=bias) + self.bn = torch.nn.BatchNorm2d(4, affine=affine) + with torch.no_grad(): + self.bn.running_mean.copy_(torch.tensor([-0.5, 0.2, 0.7, 1.0])) + self.bn.running_var.copy_(torch.tensor([0.3, 0.8, 1.2, 2.0])) + if affine: + self.bn.weight.copy_(torch.tensor([0.4, 0.8, 1.1, 1.5])) + self.bn.bias.copy_(torch.tensor([-0.7, 0.1, 0.3, 0.9])) + + def forward(self, x): + return self.bn(self.conv(x)) + + +class TestBatchNormFusionSafety(unittest.TestCase): + def _edge(self, model, inputs): + return to_edge( + export(model, inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + + def _batch_norms(self, program): + return [ + node for node in program.graph.nodes if "batch_norm" in str(node.target) + ] + + def _apply(self, program, mutable_buffers=()): + return FuseBatchNormWithConvProgramPass(mutable_buffers)(program).modified + + def test_shared_convolution_output_is_preserved(self): + class SharedOutput(ConvBatchNorm): + def forward(self, x): + y = self.conv(x) + return y, self.bn(y) + + model = SharedOutput().eval() + inputs = (torch.randn(1, 4, 8, 8),) + program = self._edge(model, inputs).exported_program() + self.assertFalse(self._apply(program)) + torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + + def test_shared_weights_are_not_modified(self): + class SharedWeight(ConvBatchNorm): + def forward(self, x): + return self.bn(self.conv(x)), self.conv(x + 1) + + model = SharedWeight().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = model(*inputs) + program = self._edge(model, inputs).exported_program() + self.assertTrue(self._apply(program)) + self.assertFalse(self._batch_norms(program)) + torch.testing.assert_close(program.module()(*inputs), expected) + + def test_channels_last_weights_preserve_strided_output(self): + class StridedOutput(ConvBatchNorm): + def forward(self, x): + return super().forward(x).as_strided((4, 64), (1, 4)) + + model = StridedOutput().eval().to(memory_format=torch.channels_last) + inputs = (torch.randn(1, 4, 8, 8),) + expected = model(*inputs) + program = self._edge(model, inputs).exported_program() + torch.testing.assert_close(program.module()(*inputs), expected) + self.assertTrue(self._apply(program)) + self.assertFalse(self._batch_norms(program)) + torch.testing.assert_close(program.module()(*inputs), expected) + + def test_training_batch_norm_is_not_folded(self): + program = self._edge( + ConvBatchNorm().train(), (torch.randn(2, 4, 8, 8),) + ).exported_program() + self.assertTrue(self._batch_norms(program)) + self.assertFalse(self._apply(program)) + + def test_runtime_statistics_are_not_folded(self): + class RuntimeStats(ConvBatchNorm): + def forward(self, x, mean, var): + return torch.nn.functional.batch_norm( + self.conv(x), + mean, + var, + self.bn.weight, + self.bn.bias, + training=False, + ) + + inputs = (torch.randn(1, 4, 8, 8), torch.zeros(4), torch.ones(4)) + program = self._edge(RuntimeStats().eval(), inputs).exported_program() + self.assertTrue(self._batch_norms(program)) + self.assertFalse(self._apply(program)) + + def test_mutated_statistics_are_not_folded(self): + class MutableStats(ConvBatchNorm): + def forward(self, x): + self.bn.running_mean.add_(1) + return super().forward(x) + + program = self._edge( + MutableStats().eval(), (torch.randn(1, 4, 8, 8),) + ).exported_program() + self.assertTrue(self._batch_norms(program)) + self.assertFalse(self._apply(program)) + + def test_statistics_mutated_by_another_method_are_not_folded(self): + program = self._edge( + ConvBatchNorm().eval(), (torch.randn(1, 4, 8, 8),) + ).exported_program() + self.assertFalse(self._apply(program, {"bn.running_mean"})) + self.assertTrue(self._batch_norms(program)) + + def test_auxiliary_batch_norm_outputs_are_preserved(self): + class AuxiliaryOutputs(ConvBatchNorm): + def forward(self, x): + return torch.ops.aten.native_batch_norm.default( + self.conv(x), + self.bn.weight, + self.bn.bias, + self.bn.running_mean, + self.bn.running_var, + False, + 0.1, + 1e-5, + ) + + inputs = (torch.randn(1, 4, 8, 8),) + model = AuxiliaryOutputs().eval() + program = self._edge(model, inputs).exported_program() + self.assertFalse(self._apply(program)) + torch.testing.assert_close(program.module()(*inputs), model(*inputs)) + + def test_consecutive_batch_norms_fold_into_distinct_constants(self): + class Stacked(ConvBatchNorm): + def __init__(self): + super().__init__() + self.tail = ConvBatchNorm(groups=2, bias=False) + + def forward(self, x): + return self.tail(torch.relu(super().forward(x))) + + model = Stacked().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = model(*inputs) + program = self._edge(model, inputs).exported_program() + self.assertEqual(len(self._batch_norms(program)), 2) + self.assertTrue(self._apply(program)) + self.assertFalse(self._batch_norms(program)) + convolutions = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.aten.convolution.default + ] + self.assertEqual(len(convolutions), 2) + self.assertEqual( + len({node.args[i] for node in convolutions for i in (1, 2)}), 4 + ) + torch.testing.assert_close(program.module()(*inputs), expected) + + def test_optional_bias_and_affine(self): + model = ConvBatchNorm(groups=4, bias=False, affine=False).eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = model(*inputs) + program = self._edge(model, inputs).exported_program() + self.assertTrue(self._apply(program)) + self.assertFalse(self._batch_norms(program)) + torch.testing.assert_close(program.module()(*inputs), expected) + self.assertFalse(self._apply(program)) + + +class TestGraphModuleBatchNormFusionSafety(TestBatchNormFusionSafety): + def _apply(self, program, mutable_buffers=()): + fusion = FuseBatchNormWithConvPass(program, mutable_buffers=mutable_buffers) + return fusion(program.graph_module).modified + + def test_subclass_can_fuse_is_honored(self): + class NeverFuse(FuseBatchNormWithConvPass): + @staticmethod + def can_fuse(conv, bn, program): + return False + + program = self._edge( + ConvBatchNorm().eval(), (torch.randn(1, 4, 8, 8),) + ).exported_program() + self.assertFalse(NeverFuse(program)(program.graph_module).modified) + self.assertTrue(self._batch_norms(program)) diff --git a/backends/transforms/test/test_fuse_batch_norm_with_linear.py b/backends/transforms/test/test_fuse_batch_norm_with_linear.py new file mode 100644 index 00000000000..99ed4fbfe1f --- /dev/null +++ b/backends/transforms/test/test_fuse_batch_norm_with_linear.py @@ -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. + +import copy +import unittest + +import torch +from executorch.backends.transforms.fuse_batch_norm_with_linear import ( + FuseBatchNormWithLinearPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export + +_BN_OPS = { + exir_ops.edge.aten.native_batch_norm.default, + exir_ops.edge.aten._native_batch_norm_legit_no_training.default, +} + + +class _LinearBn(torch.nn.Module): + def __init__(self, linear: torch.nn.Module, bn: torch.nn.Module): + super().__init__() + self.linear = linear + self.bn = bn + + def forward(self, x): + return self.bn(self.linear(x)) + + +class _LinearBnWithRuntimeParameter(_LinearBn): + def __init__(self, parameter: str): + super().__init__(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8)) + self.parameter = parameter + + def forward(self, x, parameter): + x = torch.nn.functional.linear( + x, + self.linear.weight, + parameter if self.parameter == "linear_bias" else self.linear.bias, + ) + return torch.nn.functional.batch_norm( + x, + parameter if self.parameter == "running_mean" else self.bn.running_mean, + parameter if self.parameter == "running_var" else self.bn.running_var, + self.bn.weight, + parameter if self.parameter == "bn_bias" else self.bn.bias, + training=False, + eps=self.bn.eps, + ) + + +def _seeded(module: torch.nn.Module) -> torch.nn.Module: + torch.manual_seed(0) + bn = module.bn + bn.running_mean.uniform_(-1, 1) + bn.running_var.uniform_(0.5, 2) + bn.weight.data.uniform_(0.5, 2) + if bn.bias is not None: + bn.bias.data.uniform_(-1, 1) + return module.eval() + + +class TestFuseBatchNormWithLinear(unittest.TestCase): + def _check(self, module, inputs, *, fuse=True): + module = _seeded(module) + expected = module(*inputs) + + edge = to_edge( + export(module, inputs, strict=True), + compile_config=EdgeCompileConfig( + preserve_ops=[torch.ops.aten.linear.default], + ), + ) + self.assertTrue( + any(n.target in _BN_OPS for n in edge.exported_program().graph.nodes) + ) + result = FuseBatchNormWithLinearPass(edge.exported_program())( + copy.deepcopy(edge.exported_program().graph_module) + ) + self.assertEqual(result.modified, fuse) + edge = edge.transform([FuseBatchNormWithLinearPass(edge.exported_program())]) + ep = edge.exported_program() + + self.assertEqual(any(n.target in _BN_OPS for n in ep.graph.nodes), not fuse) + torch.testing.assert_close(ep.module()(*inputs), expected, rtol=1e-5, atol=1e-5) + + def test_linear(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8)), + (torch.randn(4, 16),), + ) + + def test_linear_no_bias(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8, bias=False), torch.nn.BatchNorm1d(8)), + (torch.randn(4, 16),), + ) + + def test_linear_wide(self): + self._check( + _LinearBn(torch.nn.Linear(64, 32), torch.nn.BatchNorm1d(32)), + (torch.randn(2, 64),), + ) + + def test_batch_norm_no_bias(self): + self._check( + _LinearBn(torch.nn.Linear(16, 8), torch.nn.BatchNorm1d(8, bias=False)), + (torch.randn(4, 16),), + ) + + def test_non_rank_two_is_not_fused(self): + for bn_type, spatial_shape in ( + (torch.nn.BatchNorm1d, ()), + (torch.nn.BatchNorm2d, (3,)), + (torch.nn.BatchNorm3d, (3, 2)), + ): + for channels in (4, 8): + with self.subTest(batch_norm=bn_type.__name__, channels=channels): + self._check( + _LinearBn(torch.nn.Linear(16, 8), bn_type(channels)), + (torch.randn(2, channels, *spatial_shape, 16),), + fuse=False, + ) + + def test_runtime_parameters_are_not_fused(self): + for parameter in ("linear_bias", "bn_bias", "running_mean", "running_var"): + with self.subTest(parameter=parameter): + self._check( + _LinearBnWithRuntimeParameter(parameter), + (torch.randn(4, 16), torch.rand(8) + 0.5), + fuse=False, + ) diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index 3bef9899981..2eb714839c6 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/backends/vulkan/custom_ops_lib.py @@ -1068,14 +1068,32 @@ def embedding_q4gsw_impl( unpacked = torch.stack([low, high], dim=-1).reshape(weight.shape[0], -1) else: unpacked = torch.stack([high, low], dim=-1).reshape(weight.shape[0], -1) - # Dequantize using per-group scales - num_groups = weight_scales.shape[1] if weight_scales.dim() > 1 else 1 + # The runtime shader indexes scales as [vocab_size, groups_per_row]. + vocab_size = weight.shape[0] + if weight_scales.dim() == 1: + if weight_scales.shape[0] != vocab_size: + raise ValueError( + "embedding_q4gsw weight_scales must contain one scale per row" + ) + # Per-row scales form a single group, so group_size spans the full row. + num_groups = 1 + scales = weight_scales.reshape(vocab_size, 1, 1) + elif weight_scales.dim() == 2: + if weight_scales.shape[0] != vocab_size: + raise ValueError( + "embedding_q4gsw weight_scales first dimension must match weight rows" + ) + num_groups = weight_scales.shape[1] + scales = weight_scales.unsqueeze(-1) + else: + raise ValueError("embedding_q4gsw weight_scales must be 1-D or 2-D") + + if group_size <= 0 or num_groups * group_size != unpacked.shape[1]: + raise ValueError( + "embedding_q4gsw group_size and weight_scales must tile embedding rows" + ) + unpacked_groups = unpacked.reshape(weight.shape[0], num_groups, group_size) - scales = ( - weight_scales.unsqueeze(-1) - if weight_scales.dim() > 1 - else weight_scales.reshape(1, 1, 1) - ) dequantized = unpacked_groups.float() * scales.float() dequantized = dequantized.reshape(weight.shape[0], -1) return torch.nn.functional.embedding(indices, dequantized) diff --git a/backends/vulkan/patterns/sdpa.py b/backends/vulkan/patterns/sdpa.py index f67799f9b76..0ca3a225bf1 100644 --- a/backends/vulkan/patterns/sdpa.py +++ b/backends/vulkan/patterns/sdpa.py @@ -27,8 +27,22 @@ def is_custom_sdpa_node(node: Any) -> bool: return utils.node_has_target(node, "llama::custom_sdpa") -def is_sdpa_with_kv_cache_node(node: Any) -> bool: - return utils.node_has_target(node, "llama::sdpa_with_kv_cache") +def find_cache_update( + cache_node: torch.fx.Node, + start_pos_node: torch.fx.Node, + nodes_before_attention: set[torch.fx.Node], +) -> Optional[torch.fx.Node]: + for user in cache_node.users: + if ( + is_update_cache_node(user) + and user.args[1] is cache_node + and user.args[2] is start_pos_node + and not user.users + and user in nodes_before_attention + ): + return user + + return None class CausalSDPAMatch(PatternMatch): @@ -54,32 +68,44 @@ def __init__(self, custom_sdpa_node: torch.fx.Node) -> None: else: self.scale_node = None - # try to find update key cache node - self.update_key_cache_node = None - for user in self.key_cache_node.users: - if is_update_cache_node(user): - self.update_key_cache_node = user + nodes_before_attention = set() + for node in custom_sdpa_node.graph.nodes: + if node is self.anchor_node: break + nodes_before_attention.add(node) + + self.update_key_cache_node = find_cache_update( + self.key_cache_node, self.start_pos_node, nodes_before_attention + ) self.key_projection_node = None if self.update_key_cache_node is not None: self.key_projection_node = self.update_key_cache_node.args[0] - # find update value cache node - self.update_value_cache_node = None - for user in self.value_cache_node.users: - if is_update_cache_node(user): - self.update_value_cache_node = user - break + self.update_value_cache_node = find_cache_update( + self.value_cache_node, self.start_pos_node, nodes_before_attention + ) self.value_projection_node = None if self.update_value_cache_node is not None: self.value_projection_node = self.update_value_cache_node.args[0] - # We have additional optional arguments but we don't need to capture them - # since the new op doesn't use them - - self.match_found = True + key_cache_users = {self.anchor_node, self.update_key_cache_node} + value_cache_users = {self.anchor_node, self.update_value_cache_node} + self.match_found = ( + self.update_key_cache_node is not None + and self.key_projection_node is not None + and self.update_value_cache_node is not None + and self.value_projection_node is not None + and self.key_cache_node is not self.value_cache_node + and self.update_key_cache_node is not self.update_value_cache_node + and self.key_projection_node is not self.key_cache_node + and self.key_projection_node is not self.value_cache_node + and self.value_projection_node is not self.key_cache_node + and self.value_projection_node is not self.value_cache_node + and set(self.key_cache_node.users) == key_cache_users + and set(self.value_cache_node.users) == value_cache_users + ) @register_pattern_detector("causal_sdpa") @@ -101,22 +127,9 @@ def find_causal_sdpa_patterns( ## -def find_singleton_start_pos_node(graph_module: torch.fx.GraphModule): - for node in graph_module.graph.nodes: - if is_update_cache_node(node): - return node.args[2] - - if is_sdpa_with_kv_cache_node(node): - return node.args[5] - - raise Exception( - "Could not find an instance of llama::update_cache or sdpa_with_kv_cache" - ) - - @register_pattern_replacement("causal_sdpa") def replace_custom_sdpa_with_causal_sdpa( - ep: ExportedProgram, + ep: Optional[ExportedProgram], graph_module: torch.fx.GraphModule, match: CausalSDPAMatch, ): @@ -125,8 +138,6 @@ def replace_custom_sdpa_with_causal_sdpa( assert match.update_value_cache_node is not None assert match.value_projection_node is not None - singleton_start_pos_node = find_singleton_start_pos_node(graph_module) - with graph_module.graph.inserting_before(match.anchor_node): new_node = graph_module.graph.create_node( "call_function", @@ -137,7 +148,7 @@ def replace_custom_sdpa_with_causal_sdpa( match.value_projection_node, match.key_cache_node, match.value_cache_node, - singleton_start_pos_node, + match.start_pos_node, 1, match.attn_mask_node, match.dropout_p_node, diff --git a/backends/vulkan/runtime/api/Context.cpp b/backends/vulkan/runtime/api/Context.cpp index 05c7e5a635c..492d258d989 100644 --- a/backends/vulkan/runtime/api/Context.cpp +++ b/backends/vulkan/runtime/api/Context.cpp @@ -18,7 +18,7 @@ #endif #ifndef VULKAN_QUERY_POOL_SIZE -#define VULKAN_QUERY_POOL_SIZE 4096u +#define VULKAN_QUERY_POOL_SIZE 65536u #endif namespace vkcompute { diff --git a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp index bc315727bd2..c83a8a318c0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/SDPA.cpp @@ -695,7 +695,6 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { VK_CHECK_COND( graph.val_is_none(dropout_p) || graph.extract_scalar(dropout_p) == 0); - VK_CHECK_COND(graph.val_is_none(scale)); // is_causal is assumed to be true in the current implementation. VK_CHECK_COND( graph.val_is_none(is_causal) || graph.extract_scalar(is_causal)); @@ -762,7 +761,9 @@ void sdpa_impl(ComputeGraph& graph, const std::vector& args) { utils::kWidthPacked); 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)); + const float scale_val = graph.val_is_none(scale) + ? 1.0f / std::sqrt(static_cast(head_dim_size)) + : graph.extract_scalar(scale); add_sdpa_compute_attn_weights_node( graph, @@ -816,10 +817,18 @@ void sdpa_with_kv_cache_impl( (void)sequence_len; utils::StorageType cache_storage = graph.storage_type_of(q_projected); - const ValueRef k_cache = - graph.add_tensor_like(k_cache_data, cache_storage, utils::kWidthPacked); - const ValueRef v_cache = - graph.add_tensor_like(v_cache_data, cache_storage, utils::kWidthPacked); + const ValueRef k_cache = prepack_standard( + graph, + k_cache_data, + cache_storage, + utils::kWidthPacked, + /*passthrough=*/true); + const ValueRef v_cache = prepack_standard( + graph, + v_cache_data, + cache_storage, + utils::kWidthPacked, + /*passthrough=*/true); update_cache_impl(graph, {k_projected, k_cache, input_pos_symint, -1}); update_cache_impl(graph, {v_projected, v_cache, input_pos_symint, -1}); @@ -855,7 +864,6 @@ void compute_attn_weight_with_kv_cache_impl( const ValueRef is_causal = args[arg_idx++]; (void)is_causal; const ValueRef scale = args[arg_idx++]; - (void)scale; // Output tensors const ValueRef out = args[arg_idx++]; @@ -863,16 +871,26 @@ void compute_attn_weight_with_kv_cache_impl( (void)sequence_len; const utils::StorageType cache_storage = graph.storage_type_of(q_projected); - const ValueRef k_cache = - graph.add_tensor_like(k_cache_data, cache_storage, utils::kWidthPacked); - const ValueRef v_cache = - graph.add_tensor_like(v_cache_data, cache_storage, utils::kWidthPacked); + const ValueRef k_cache = prepack_standard( + graph, + k_cache_data, + cache_storage, + utils::kWidthPacked, + /*passthrough=*/true); + const ValueRef v_cache = prepack_standard( + graph, + v_cache_data, + cache_storage, + utils::kWidthPacked, + /*passthrough=*/true); update_cache_impl(graph, {k_projected, k_cache, input_pos_symint, -1}); update_cache_impl(graph, {v_projected, v_cache, input_pos_symint, -1}); 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)); + const float scale_val = graph.val_is_none(scale) + ? 1.0f / std::sqrt(static_cast(head_dim_size)) + : graph.extract_scalar(scale); add_sdpa_compute_attn_weights_node( graph, diff --git a/backends/vulkan/runtime/vk_api/QueryPool.h b/backends/vulkan/runtime/vk_api/QueryPool.h index 8357a06521e..a7a0a9c9bcf 100644 --- a/backends/vulkan/runtime/vk_api/QueryPool.h +++ b/backends/vulkan/runtime/vk_api/QueryPool.h @@ -20,7 +20,7 @@ #include #ifndef VULKAN_QUERY_POOL_SIZE -#define VULKAN_QUERY_POOL_SIZE 4096u +#define VULKAN_QUERY_POOL_SIZE 65536u #endif namespace vkcompute { diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index e395e49ae78..bd425874169 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -162,6 +162,10 @@ def maybe_add_constant_tensor(self, node: Node) -> int: # Serialize tensor data to bytes tensor = tensor.contiguous() + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if tensor.nbytes != tensor.untyped_storage().nbytes(): + tensor = tensor.clone() size = tensor.untyped_storage().nbytes() if size > 0: diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index b20073493af..0871cb455a3 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -283,6 +283,8 @@ def serialize_constant_tensors( ): vk_graph.constants.append(VkBytes(current_offset, 0)) elif isinstance(tensor, torch.Tensor): + # The bytes are read from the start of the storage in memory order. + tensor = tensor.clone(memory_format=torch.contiguous_format) array_type = ctypes.c_char * tensor.untyped_storage().nbytes() array = ctypes.cast( tensor.untyped_storage().data_ptr(), diff --git a/backends/vulkan/test/op_tests/utils/aten_types.py b/backends/vulkan/test/op_tests/utils/aten_types.py index 52da866ae01..380091fc756 100644 --- a/backends/vulkan/test/op_tests/utils/aten_types.py +++ b/backends/vulkan/test/op_tests/utils/aten_types.py @@ -25,7 +25,7 @@ OPT_LAYOUT = "::std::optional" OPT_MEMORY_FORMAT = "::std::optional" OPT_SCALAR_TYPE = "::std::optional" -STRING = "std::string_view" +STRING = "::std::string_view" TWO_TENSOR_TUPLE = "::std::tuple" THREE_TENSOR_TUPLE = "::std::tuple" TENSOR_VECTOR = "::std::vector" diff --git a/backends/vulkan/test/op_tests/utils/gen_computegraph.py b/backends/vulkan/test/op_tests/utils/gen_computegraph.py index 657cbfdd926..7fde3f7dbe6 100644 --- a/backends/vulkan/test/op_tests/utils/gen_computegraph.py +++ b/backends/vulkan/test/op_tests/utils/gen_computegraph.py @@ -121,9 +121,8 @@ def __init__( ctype = cpp.argumenttype_type( arg.type, mutable=arg.is_write, binds=arg.name ) - cpp_type = ctype.cpp_type(strip_ref=True).replace( - "c10::string_view", STRING - ) + cpp_type = ctype.cpp_type(strip_ref=True) + cpp_type = cpp_type.replace("c10::string_view", STRING) self.args.append( ATenArg(name=arg.name, cpp_type=cpp_type, default=arg.default) diff --git a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.cpp b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.cpp index 84454d8723d..a1e02529193 100644 --- a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.cpp @@ -59,10 +59,14 @@ TestCase create_test_case(const EmbeddingConfig& config) { weight.set_constant(true); test_case.add_input_spec(weight); - // Weight scales: [vocab_size, groups_per_row] + // Weight scales: [vocab_size] for per-row, otherwise + // [vocab_size, groups_per_row]. int64_t groups_per_row = config.embed_dim / config.group_size; + std::vector scales_shape = config.per_row_scales_1d + ? std::vector{config.vocab_size} + : std::vector{config.vocab_size, groups_per_row}; ValueSpec weight_scales( - {config.vocab_size, groups_per_row}, + scales_shape, config.scales_dtype, utils::kBuffer, utils::kWidthPacked, @@ -124,7 +128,8 @@ void embedding_4bit_reference(TestCase& tc) { int64_t packed_dim = weight_spec.sizes[1]; int64_t embed_dim = packed_dim * 2; - int64_t groups_per_row = scales_spec.sizes[1]; + int64_t groups_per_row = + scales_spec.sizes.size() == 1 ? 1 : scales_spec.sizes[1]; int64_t num_indices = 1; for (auto s : indices_spec.sizes) { diff --git a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.h b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.h index 0a2103350a8..a161653ce41 100644 --- a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.h +++ b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/embedding_q4gsw.h @@ -27,6 +27,7 @@ struct EmbeddingConfig { vkapi::ScalarType scales_dtype = vkapi::kHalf; utils::StorageType storage_type = utils::kBuffer; bool is_linear_weight = false; + bool per_row_scales_1d = false; }; TestCase create_test_case(const EmbeddingConfig& config); diff --git a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/test_sets/embedding_q4gsw_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/test_sets/embedding_q4gsw_sets.cpp index 614eeb721a7..764a34102cf 100644 --- a/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/test_sets/embedding_q4gsw_sets.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/embedding_q4gsw/test_sets/embedding_q4gsw_sets.cpp @@ -266,7 +266,32 @@ REGISTER_TEST_CASE_SET("embedding_q4gsw", "correctness") { .scales_dtype = vkapi::kFloat, .storage_type = utils::kTexture3D}, }; - return {generate_test_cases(configs), embedding_4bit_reference}; + auto test_cases = generate_test_cases(configs); + for (const auto storage_type : {utils::kBuffer, utils::kTexture3D}) { + test_cases.push_back(create_test_case( + {.vocab_size = 16, + .embed_dim = 1536, + .group_size = 1536, + .indices_shape = {1, 4}, + .test_case_name = "gemma4_token_embedding_per_row", + .dtype = vkapi::kFloat, + .scales_dtype = vkapi::kFloat, + .storage_type = storage_type, + .is_linear_weight = true, + .per_row_scales_1d = true})); + test_cases.push_back(create_test_case( + {.vocab_size = 16, + .embed_dim = 8960, + .group_size = 8960, + .indices_shape = {1, 4}, + .test_case_name = "gemma4_ple_embedding_per_row", + .dtype = vkapi::kFloat, + .scales_dtype = vkapi::kFloat, + .storage_type = storage_type, + .is_linear_weight = true, + .per_row_scales_1d = true})); + } + return {test_cases, embedding_4bit_reference}; } REGISTER_TEST_CASE_SET("embedding_q4gsw", "performance") { diff --git a/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.cpp b/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.cpp index 24e39a3e110..cceb98c8916 100644 --- a/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.cpp @@ -34,10 +34,15 @@ TestCase create_sdpa_test_case( std::to_string(config.n_heads) + " Hkv" + std::to_string(config.n_kv_heads) + " S" + std::to_string(config.seq_len) + " C" + - std::to_string(config.context_len); + std::to_string(config.context_len) + " P" + + std::to_string(config.input_pos >= 0 + ? config.input_pos + : config.context_len - config.seq_len); - const std::string suffix = - "[" + config.model + " " + config.regime + " " + impl + "]"; + const std::string suffix = "[" + config.model + " " + config.regime + " " + + impl + " scale=" + + (config.scale.has_value() ? std::to_string(*config.scale) : "default") + + "]"; test_case.set_name(make_test_label( prefix, dtype_str, dtype_str, shape, storage_str, suffix)); @@ -64,7 +69,29 @@ TestCase create_sdpa_test_case( storage_type, utils::kWidthPacked, DataGenType::RANDOM); - + ValueSpec k_projected( + {1, config.seq_len, config.n_kv_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + ValueSpec v_projected( + {1, config.seq_len, config.n_kv_heads, config.head_dim}, + dtype, + storage_type, + utils::kWidthPacked, + DataGenType::RANDOM); + const int64_t input_pos = config.input_pos >= 0 + ? config.input_pos + : config.context_len - config.seq_len; + ValueSpec input_pos_spec(utils::safe_downcast(input_pos)); + + ValueSpec scale; + if (config.scale.has_value()) { + scale = ValueSpec(*config.scale); + } else { + scale.set_none(true); + } ValueSpec impl_selector = ValueSpec::make_string(impl); // out: [1, S, n_heads, head_dim] @@ -78,6 +105,10 @@ TestCase create_sdpa_test_case( test_case.add_input_spec(q); test_case.add_input_spec(k_cache); test_case.add_input_spec(v_cache); + test_case.add_input_spec(k_projected); + test_case.add_input_spec(v_projected); + test_case.add_input_spec(input_pos_spec); + test_case.add_input_spec(scale); test_case.add_input_spec(impl_selector); test_case.add_output_spec(output); @@ -121,7 +152,7 @@ static std::vector as_float_data(const ValueSpec& spec) { } // Reference: causal SDPA over the KV cache. -// q:[1,S,H,D], k/v cache:[1,C,Hkv,D], input_pos = C - S. +// q:[1,S,H,D], k/v cache:[1,C,Hkv,D], input_pos + S <= C. // For query row s (absolute position input_pos + s), attends to cache // positions [0, input_pos + s]. GQA: head h maps to kv head h / (H/Hkv). void sdpa_reference_impl(TestCase& test_case) { @@ -138,17 +169,38 @@ void sdpa_reference_impl(TestCase& test_case) { const int64_t C = k_sizes[1]; const int64_t Hkv = k_sizes[2]; - if (C > kRefContextLenLimit) { + if (C > 512 || (C > kRefContextLenLimit && S > 18)) { throw std::invalid_argument("sdpa reference: perf shape, skipping"); } - const int64_t input_pos = C - S; + const int64_t input_pos = test_case.inputs()[5].get_int_value(); + const int64_t context_len = input_pos + S; + if (input_pos < 0 || context_len > C) { + throw std::invalid_argument("sdpa reference: invalid input_pos"); + } const int64_t heads_per_kv = H / Hkv; - const float scale = 1.0f / std::sqrt(static_cast(D)); + const ValueSpec& scale_spec = test_case.inputs()[6]; + const float scale = scale_spec.is_none() + ? 1.0f / std::sqrt(static_cast(D)) + : scale_spec.get_float_value(); const auto q_data = as_float_data(q); - const auto k_data = as_float_data(k); - const auto v_data = as_float_data(v); + auto k_data = as_float_data(k); + auto v_data = as_float_data(v); + if (test_case.inputs()[7].get_string_value() == "fused") { + const auto k_projected_data = as_float_data(test_case.inputs()[3]); + const auto v_projected_data = as_float_data(test_case.inputs()[4]); + for (int64_t s = 0; s < S; ++s) { + for (int64_t hk = 0; hk < Hkv; ++hk) { + for (int64_t d = 0; d < D; ++d) { + const int64_t projected_idx = (s * Hkv + hk) * D + d; + const int64_t cache_idx = ((input_pos + s) * Hkv + hk) * D + d; + k_data[cache_idx] = k_projected_data[projected_idx]; + v_data[cache_idx] = v_projected_data[projected_idx]; + } + } + } + } ValueSpec& output = test_case.outputs()[0]; auto& ref = output.get_ref_float_data(); @@ -162,7 +214,7 @@ void sdpa_reference_impl(TestCase& test_case) { return (c * Hkv + hk) * D + d; }; - std::vector scores(C); + std::vector scores(context_len); for (int64_t s = 0; s < S; ++s) { const int64_t attend_len = input_pos + s + 1; // causal for (int64_t h = 0; h < H; ++h) { diff --git a/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.h b/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.h index a04e2784704..b1fb126afc9 100644 --- a/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.h +++ b/backends/vulkan/test/ops/benchmark/cases/sdpa/sdpa.h @@ -9,6 +9,7 @@ #include #include +#include #include namespace executorch { @@ -32,6 +33,8 @@ struct SDPAConfig { int64_t context_len; // total KV length (kv_len) std::string model; // label only std::string regime; // "decode" / "prefill", label only + std::optional scale = std::nullopt; + int64_t input_pos = -1; }; TestCase create_sdpa_test_case( diff --git a/backends/vulkan/test/ops/benchmark/cases/sdpa/test_sets/sdpa_sets.cpp b/backends/vulkan/test/ops/benchmark/cases/sdpa/test_sets/sdpa_sets.cpp index 95b1b341713..de1b23948d6 100644 --- a/backends/vulkan/test/ops/benchmark/cases/sdpa/test_sets/sdpa_sets.cpp +++ b/backends/vulkan/test/ops/benchmark/cases/sdpa/test_sets/sdpa_sets.cpp @@ -113,6 +113,113 @@ REGISTER_TEST_CASE_SET("sdpa", "correctness") { // unless padded — see sdpa_impl). .accu_cases = { + {.configs = + { + {256, + 8, + 1, + 1, + 32, + "Gemma-4-E2B", + "sliding_decode", + 1.0f}, + {512, + 8, + 1, + 1, + 16, + "Gemma-4-E2B", + "full_decode", + 1.0f}, + {64, + 8, + 2, + 1, + 32, + "default_scale", + "decode", + std::nullopt}, + }, + .storage_types = {utils::kTexture3D}, + .impls = {"production"}}, + {.configs = + { + {256, + 8, + 1, + 1, + 32, + "Gemma-4-E2B", + "sliding_decode", + 1.0f}, + {512, + 8, + 1, + 1, + 16, + "Gemma-4-E2B", + "full_decode", + 1.0f}, + }, + .storage_types = {utils::kTexture3D}, + .impls = {"gqa_tile2", "gqa_base", "non_gqa"}}, + {.configs = + { + {256, + 8, + 1, + 4, + 4, + "Gemma-4-E2B", + "sliding_prefill", + 1.0f}, + {512, + 8, + 1, + 4, + 4, + "Gemma-4-E2B", + "full_prefill", + 1.0f}, + {256, + 8, + 1, + 18, + 512, + "Gemma-4-E2B", + "sliding_prefill_real", + 1.0f, + 0}, + {512, + 8, + 1, + 18, + 512, + "Gemma-4-E2B", + "full_prefill_real", + 1.0f, + 0}, + {256, + 8, + 1, + 1, + 512, + "Gemma-4-E2B", + "sliding_decode_prefix", + 1.0f, + 17}, + {512, + 8, + 1, + 1, + 512, + "Gemma-4-E2B", + "full_decode_prefix", + 1.0f, + 17}, + }, + .storage_types = {utils::kTexture3D}, + .impls = {"fused"}}, // Cover D=64 (D4=16) and D=128 (D4=32) with the // vendor-default GQA and the per-query-head shaders, across // texture + buffer. diff --git a/backends/vulkan/test/ops/benchmark/test_ops/impl/TestSDPA.cpp b/backends/vulkan/test/ops/benchmark/test_ops/impl/TestSDPA.cpp index 4a59210cf9c..9ab1630331b 100644 --- a/backends/vulkan/test/ops/benchmark/test_ops/impl/TestSDPA.cpp +++ b/backends/vulkan/test/ops/benchmark/test_ops/impl/TestSDPA.cpp @@ -15,6 +15,11 @@ namespace vkcompute { +void sdpa_impl(ComputeGraph& graph, const std::vector& args); +void sdpa_with_kv_cache_impl( + ComputeGraph& graph, + const std::vector& args); + // Bare-minimum mirror of the LLM decode/prefill path in sdpa_impl() (see // SDPA.cpp), stripped of the input-validation VK_CHECK_CONDs. Building the // three SDPA nodes directly lets the test forward a shader_override knob to the @@ -25,6 +30,7 @@ static void test_sdpa_impl( const ValueRef k_cache, const ValueRef v_cache, const ValueRef input_pos_symint, + const ValueRef scale, const ValueRef out, const ValueRef shader_override) { const int64_t num_q_heads = graph.size_at(-2, q_projected); @@ -74,7 +80,9 @@ static void test_sdpa_impl( utils::kWidthPacked); 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)); + const float scale_val = graph.val_is_none(scale) + ? 1.0f / std::sqrt(static_cast(head_dim_size)) + : graph.extract_scalar(scale); add_sdpa_compute_attn_weights_node( graph, @@ -110,11 +118,8 @@ static void test_sdpa_impl( // Test wrapper for the LLM KV-cache SDPA op (llama.custom_sdpa.default). // // The production op reads the dynamic context length from an `input_pos` -// symint, but input_pos is not a free parameter: the op enforces -// context_len = S + input_pos, and context_len is exactly the KV-cache's dim -// -3. Since q is [B=1, S, H, D] and the caches are [B=1, context_len, H_kv, D], -// input_pos is fully determined by the tensor shapes. Deriving it here from -// those shapes keeps the single source of truth (the cache size) authoritative. +// symint. The KV-cache's dim -3 is its capacity, so input_pos + S must not +// exceed that dimension. // // Decode vs prefill is selected automatically inside the nodes via // is_single_token() (S == 1 -> coop/GEMV shaders, S > 1 -> tiled shaders). For @@ -126,12 +131,17 @@ static void test_sdpa_impl( // "non_gqa" -> force the per-query-head coop shader // impl_selector has no effect on prefill (tiled). // -// Args: q, k_cache, v_cache, impl_selector, out +// Args: q, k_cache, v_cache, k_projected, v_projected, input_pos, scale, +// impl_selector, out void test_sdpa(ComputeGraph& graph, const std::vector& args) { int arg_idx = 0; const ValueRef q = args.at(arg_idx++); const ValueRef k_cache = args.at(arg_idx++); const ValueRef v_cache = args.at(arg_idx++); + const ValueRef k_projected = args.at(arg_idx++); + const ValueRef v_projected = args.at(arg_idx++); + const ValueRef input_pos_ref = args.at(arg_idx++); + const ValueRef scale = args.at(arg_idx++); const ValueRef impl_selector_str = args.at(arg_idx++); const ValueRef out = args.at(arg_idx++); @@ -139,18 +149,52 @@ void test_sdpa(ComputeGraph& graph, const std::vector& args) { VK_CHECK_COND( impl_selector == "default" || impl_selector == "gqa" || impl_selector == "gqa_tile2" || impl_selector == "gqa_base" || - impl_selector == "non_gqa", + impl_selector == "non_gqa" || impl_selector == "production" || + impl_selector == "fused", "test_sdpa: impl_selector must be one of {default, gqa, gqa_tile2, " - "gqa_base, non_gqa}"); + "gqa_base, non_gqa, production, fused}"); const int64_t seq_len = graph.size_at(-3, q); - const int64_t context_len = graph.size_at(-3, k_cache); - VK_CHECK_COND(context_len >= seq_len); - const int32_t input_pos_val = - utils::safe_downcast(context_len - seq_len); + const int64_t cache_capacity = graph.size_at(-3, k_cache); + const int32_t input_pos_val = graph.extract_scalar(input_pos_ref); + VK_CHECK_COND(input_pos_val >= 0); + VK_CHECK_COND(input_pos_val + seq_len <= cache_capacity); const ValueRef input_pos_symint = graph.add_symint(input_pos_val); + if (impl_selector == "fused") { + sdpa_with_kv_cache_impl( + graph, + {q, + k_projected, + v_projected, + k_cache, + v_cache, + input_pos_symint, + graph.add_scalar(seq_len), + graph.add_none(), + graph.add_scalar(0.0), + graph.add_scalar(true), + scale, + out}); + return; + } + + if (impl_selector == "production") { + sdpa_impl( + graph, + {q, + k_cache, + v_cache, + input_pos_symint, + graph.add_none(), + graph.add_scalar(0.0), + graph.add_scalar(true), + scale, + out}); + return; + } + // shader_override (see SDPA.h): kDummyValueRef auto; otherwise a // kShaderOverride* scalar forcing the AV shader family / variant. ValueRef shader_override = kDummyValueRef; @@ -165,7 +209,14 @@ void test_sdpa(ComputeGraph& graph, const std::vector& args) { } test_sdpa_impl( - graph, q, k_cache, v_cache, input_pos_symint, out, shader_override); + graph, + q, + k_cache, + v_cache, + input_pos_symint, + scale, + out, + shader_override); } REGISTER_OPERATORS { diff --git a/backends/vulkan/test/targets.bzl b/backends/vulkan/test/targets.bzl index aa28733f0de..c9b3e9b836c 100644 --- a/backends/vulkan/test/targets.bzl +++ b/backends/vulkan/test/targets.bzl @@ -65,6 +65,10 @@ def define_common_targets(is_fbcode = False): srcs = [ "test_vulkan_passes.py", ], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], deps = [ "//caffe2:torch", "//executorch/backends/vulkan/_passes:vulkan_passes", @@ -73,10 +77,29 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/vulkan:vulkan_preprocess", "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer", "//executorch/exir:lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/extension/pybindings:portable_lib", # @manual + "//executorch/kernels/portable:custom_ops_generated_lib", "//pytorch/ao:torchao", # @manual ] ) + python_unittest( + name = "test_sdpa_patterns", + srcs = ["test_sdpa_patterns.py"], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/patterns:vulkan_patterns", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/extension/pybindings:portable_lib", # @manual + "//executorch/kernels/portable:custom_ops_generated_lib", + ], + ) + python_unittest( name = "test_vulkan_delegate_header", srcs = [ diff --git a/backends/vulkan/test/test_sdpa_patterns.py b/backends/vulkan/test/test_sdpa_patterns.py new file mode 100644 index 00000000000..73d4fe4a3c7 --- /dev/null +++ b/backends/vulkan/test/test_sdpa_patterns.py @@ -0,0 +1,284 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch + +from executorch.backends.vulkan.patterns.sdpa import ( + CausalSDPAMatch, + is_update_cache_node, + replace_custom_sdpa_with_causal_sdpa, +) +from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + +class CausalSDPAPatternTest(unittest.TestCase): + def test_rejects_shared_kv_without_cache_updates(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_preserves_explicit_scale(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertTrue(match.match_found) + self.assertEqual(match.scale_node, 1.0) + + def test_rejects_cache_shared_with_later_sdpa(self) -> None: + graph = torch.fx.Graph() + q_donor = graph.placeholder("q_donor") + q_shared = graph.placeholder("q_shared") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + donor_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q_donor, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + shared_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q_shared, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output((donor_sdpa, shared_sdpa)) + + match = CausalSDPAMatch(donor_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_live_cache_update_result(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + updated_k = graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output((custom_sdpa, updated_k)) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_cache_used_as_update_source(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + other_k_cache = graph.placeholder("other_k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k_cache, other_k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_mismatched_cache_update_position(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + update_pos = graph.placeholder("update_pos") + attention_pos = graph.placeholder("attention_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, update_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, update_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, attention_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_cache_updates_after_attention(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_aliased_key_and_value_caches(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + projected = graph.placeholder("projected") + cache = graph.placeholder("cache") + start_pos = graph.placeholder("start_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(projected, cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, cache, cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_rejects_cache_updated_from_itself(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k_cache, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, 1.0), + ) + graph.output(custom_sdpa) + + match = CausalSDPAMatch(custom_sdpa) + + self.assertFalse(match.match_found) + + def test_replacement_preserves_arguments_and_erases_updates(self) -> None: + graph = torch.fx.Graph() + q = graph.placeholder("q") + k = graph.placeholder("k") + v = graph.placeholder("v") + k_cache = graph.placeholder("k_cache") + v_cache = graph.placeholder("v_cache") + start_pos = graph.placeholder("start_pos") + scale = graph.placeholder("scale") + graph.call_function( + torch.ops.llama.update_cache.default, + args=(k, k_cache, start_pos), + ) + graph.call_function( + torch.ops.llama.update_cache.default, + args=(v, v_cache, start_pos), + ) + custom_sdpa = graph.call_function( + torch.ops.llama.custom_sdpa.default, + args=(q, k_cache, v_cache, start_pos, None, 0.0, True, scale), + ) + custom_sdpa.meta["val"] = None + graph.output(custom_sdpa) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + match = CausalSDPAMatch(custom_sdpa) + + self.assertTrue(match.match_found) + replace_custom_sdpa_with_causal_sdpa(None, graph_module, match) + graph_module.graph.lint() + + fused_nodes = [ + node + for node in graph_module.graph.nodes + if node.target is torch.ops.llama.sdpa_with_kv_cache.default + ] + self.assertEqual(len(fused_nodes), 1) + self.assertIs(fused_nodes[0].args[5], start_pos) + self.assertIs(fused_nodes[0].args[10], scale) + self.assertFalse( + any(is_update_cache_node(node) for node in graph_module.graph.nodes) + ) diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 8b95b1df2a1..2eb739a760e 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -6,7 +6,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 ctypes import random import unittest from types import SimpleNamespace @@ -267,6 +266,8 @@ def test_serialize_vulkan_binary(self) -> None: shaders=[], ) const_tensors = self._generate_random_const_tensors(5) + const_tensors.append(torch.randn(3, 4, 4)[1]) + const_tensors.append(torch.randn(3, 4, 4)[1].t()) serialized_binary = serialize_vulkan_graph(vk_graph, const_tensors, []) @@ -305,14 +306,7 @@ def test_serialize_vulkan_binary(self) -> None: constant_data_bytes = constant_data_payload[offset : offset + length] - array_type = ctypes.c_char * tensor.untyped_storage().nbytes() - array = ctypes.cast( - tensor.untyped_storage().data_ptr(), - ctypes.POINTER(array_type), - ).contents - - tensor_bytes = bytes(array) - self.assertEqual(constant_data_bytes, tensor_bytes) + self.assertEqual(constant_data_bytes, tensor.numpy().tobytes()) def test_serialize_deserialize_vkgraph(self) -> None: in_vk_graph = VkGraph( diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py index aebb40f8ade..1332da0428d 100644 --- a/backends/vulkan/test/test_vulkan_graph_builder.py +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -108,6 +108,27 @@ def forward(self, first, second): self.assertEqual(len(builder.input_ids), 2) +class TestVkGraphBuilderConstants(unittest.TestCase): + def test_view_constant_serializes_only_its_own_bytes(self) -> None: + class AddViewBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("b", torch.randn(3, 4)[1]) + + def forward(self, x): + return x + self.b + + module = AddViewBuffer() + edge = to_edge(torch.export.export(module, (torch.randn(4),), strict=True)) + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + (data,) = builder.named_data_store.get_named_data_store_output().buffers + self.assertEqual(bytes(data), module.b.numpy().tobytes()) + + class TestVkGraphBuilderScalarLists(unittest.TestCase): """A scalar list serializes to exactly one Value. bool is a subclass of int, so a bool list also satisfies the all-int check, and an empty list diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index eed4137ac5e..5e79ceec20d 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -304,6 +304,125 @@ def conv_input_ranks(graph_module: torch.fx.GraphModule) -> List[int]: class TestVulkanPasses(unittest.TestCase): + def test_embedding_q4gsw_reference_supports_per_row_scales(self): + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + vocab_size = 3 + embed_dim = 32 + unpacked = ( + torch.arange(vocab_size * embed_dim, dtype=torch.int8) + .reshape(vocab_size, embed_dim) + .remainder(16) + .sub(8) + ) + packed = ((unpacked[:, 1::2] + 8) << 4) | (unpacked[:, 0::2] + 8) + packed = packed.to(torch.uint8) + scales = torch.tensor([0.25, 0.5, 0.75], dtype=torch.float32) + indices = torch.tensor([[2, 0], [1, 2]], dtype=torch.long) + + output = torch.ops.et_vk.embedding_q4gsw.default( + packed, scales, embed_dim, indices, True + ) + expected_weight = unpacked.float() * scales.unsqueeze(1) + expected = torch.nn.functional.embedding(indices, expected_weight) + + self.assertEqual(tuple(output.shape), (2, 2, embed_dim)) + self.assertTrue(torch.equal(output, expected)) + + def test_embedding_q4gsw_reference_supports_groupwise_scales(self): + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + vocab_size = 3 + embed_dim = 32 + group_size = 8 + unpacked = ( + torch.arange(vocab_size * embed_dim, dtype=torch.int8) + .reshape(vocab_size, embed_dim) + .remainder(16) + .sub(8) + ) + packed = ((unpacked[:, 1::2] + 8) << 4) | (unpacked[:, 0::2] + 8) + packed = packed.to(torch.uint8) + scales = torch.arange(1, 13, dtype=torch.float32).reshape(vocab_size, 4) + indices = torch.tensor([2, 0, 1], dtype=torch.long) + + output = torch.ops.et_vk.embedding_q4gsw.default( + packed, scales, group_size, indices, True + ) + expected_weight = ( + unpacked.reshape(vocab_size, 4, group_size).float() * scales.unsqueeze(-1) + ).reshape(vocab_size, embed_dim) + expected = torch.nn.functional.embedding(indices, expected_weight) + + self.assertEqual(tuple(output.shape), (3, embed_dim)) + self.assertTrue(torch.equal(output, expected)) + + def test_embedding_q4gsw_reference_rejects_invalid_scale_shape(self): + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + packed = torch.zeros((3, 16), dtype=torch.uint8) + indices = torch.tensor([0], dtype=torch.long) + + with self.assertRaisesRegex(ValueError, "weight_scales"): + torch.ops.et_vk.embedding_q4gsw.default( + packed, + torch.ones(4, dtype=torch.float32), + 8, + indices, + True, + ) + + def test_fuse_quantized_embedding_with_per_row_scales(self): + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + vocab_size = 64 + embed_dim = 128 + + class PerRowQuantizedEmbedding(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "weight", + torch.zeros((vocab_size, embed_dim // 2), dtype=torch.uint8), + ) + self.register_buffer( + "scales", torch.ones(vocab_size, dtype=torch.float32) + ) + + def forward(self, indices): + return torch.ops.quantized_decomposed.embedding_4bit.dtype( + self.weight, + self.scales, + None, + -8, + 7, + indices, + dtype=torch.float32, + ) + + model = PerRowQuantizedEmbedding() + sample_inputs = (torch.tensor([[0, 1, 2]], dtype=torch.long),) + program = torch.export.export(model, sample_inputs, strict=True) + edge_program = to_edge( + program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + ep = edge_program._edge_programs["forward"] + fuse_pass = FusePatternsPass() + fuse_pass._exported_program = ep + + result = fuse_pass.call(ep.graph_module) + + self.assertTrue(result.modified) + self.assertEqual(op_node_count(ep.graph_module, "embedding_q4gsw.default"), 1) + fused_node = next( + node + for node in ep.graph.nodes + if get_target_canonical_name(node) == "embedding_q4gsw.default" + ) + self.assertEqual(fused_node.args[2], embed_dim) + self.assertEqual(tuple(fused_node.args[1].meta["val"].shape), (vocab_size,)) + def test_conv1d_as_conv2d_rewrites_eligible_conv(self): # Output equality cannot show this: an unrewritten conv1d computes the # same numbers. What distinguishes the two is the rank the convolution diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index 1963e4494fd..238bb992784 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -288,6 +288,10 @@ def get_quant_params( if quant_params.per_channel_group: scale = scale.to(torch.bfloat16) + # The bytes are read from the start of the storage in memory order. + # Scales are small, so always copy them into storage of their own. + scale = scale.clone(memory_format=torch.contiguous_format) + num_bytes = scale.untyped_storage().nbytes() scale_array = ctypes.cast( scale.untyped_storage().data_ptr(), @@ -619,6 +623,11 @@ def get_serialized_buffer_index( if quant_params is not None and quant_params.is_qc4w: const_val = self.convert_to_qc4w(const_val) + # The bytes are read from the start of the storage, so a view of a + # larger tensor needs its own copy. + if const_val.nbytes != const_val.untyped_storage().nbytes(): + const_val = const_val.clone() + size = const_val.untyped_storage().nbytes() array_type = ctypes.c_char * size array = ctypes.cast( diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index 154fa2fa559..e309bc514a7 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -781,6 +781,25 @@ def test_fp32_addmm(self): uses_bias=True, ) + def test_fp32_linear_view_weight(self): + class ViewWeightLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.randn(3, 8, 8)[1]) + + def forward(self, x): + return torch.nn.functional.linear(x, self.weight) + + ( + Tester(ViewWeightLinear(), (torch.randn(2, 8),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_fp32_linear_fused_relu(self): for use_bias in (True, False): for num_batch_dims in range(1, 3): @@ -898,6 +917,40 @@ def test_linear_qd8_per_token_groupwise_unsupported_groupsize(self): lin_mod, inputs, group_size=30, use_bias=False, atol=1e-2 ) + @unittest.skipIf(not torchao_installed, "Requires TorchAO") + def test_qd8_linear_transposed_scale(self): + mod = torch.nn.Linear(64, 4, bias=False) + with torch.no_grad(): + # Rows with very different scales make a misordered scale obvious. + mod.weight.mul_(torch.arange(1.0, 5.0).unsqueeze(1)) + quantize_( + mod, + Int8DynamicActivationIntxWeightConfig( + weight_dtype=torch.int4, weight_granularity=PerGroup(32) + ), + ) + unwrap_tensor_subclass(mod) + # XNNPACK stores groupwise scales as bf16, so round them for an exact comparison. + scale = mod.weight.scale.to(torch.bfloat16).float() + mod.weight.scale = scale.t().contiguous().t() + ( + Tester(mod, (torch.randn(2, 64),)) + .export() + .to_edge_transform_and_lower( + ToEdgeTransformAndLower( + [ + XnnpackPartitioner( + config_precisions=ConfigPrecisionType.DYNAMIC_QUANT + ) + ] + ) + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qd8_per_channel_linear_parallel(self): in_size = 2 input_size = 4 diff --git a/docs/source/backends-overview.md b/docs/source/backends-overview.md index 1d626a04ba4..5e3d575ca4b 100644 --- a/docs/source/backends-overview.md +++ b/docs/source/backends-overview.md @@ -22,7 +22,7 @@ Backends are the bridge between your exported model and the hardware it runs on. |--------------------------------------------------------------|---------------|---------------|---------------------------------| | [XNNPACK](backends/xnnpack/xnnpack-overview.md) | All | CPU | General-purpose, fallback | | [CUDA](backends/cuda/cuda-overview.md) | Linux/Windows | GPU | NVIDIA GPU acceleration | -| [TensorRT](backends/tensorrt/tensorrt-overview.md) | Linux | GPU | NVIDIA GPU acceleration via TensorRT | +| [TensorRT](backends/tensorrt/tensorrt-overview.md) | Linux/Windows | GPU | NVIDIA GPU acceleration via TensorRT | | [Core ML](backends/coreml/coreml-overview.md) | iOS, macOS | NPU/GPU/CPU | Apple devices, high performance | | [MLX](/backends/mlx/mlx-overview.md) | iOS (experimental), macOS | GPU | Apple Silicon GPU (MLX) | | [Vulkan](backends/vulkan/vulkan-overview.md) | Android, Linux, Windows | GPU | Android devices (mature); Desktops (experimental) | diff --git a/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md b/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md index 5fe030d1de3..1e8ee993398 100644 --- a/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md +++ b/docs/source/backends/arm-ethos-u/tutorials/ethos-u-getting-started.md @@ -25,6 +25,8 @@ If you are already familiar with this delegate, you may want to jump directly to This tutorial serves as an introduction to using ExecuTorch to deploy PyTorch models on Arm® Ethos™-U targets. It is based on `ethos_u_minimal_example.ipynb`, provided in Arm’s examples folder. +This should not be here. + ## Prerequisites ### Hardware diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index ef57aca7ae5..32effc7e520 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -7,24 +7,29 @@ The CUDA backend is the ExecuTorch solution for running models on NVIDIA GPUs. I - **Optimized GPU Execution**: Uses AOTInductor to generate highly optimized CUDA kernels for model operators - **Triton Kernel Support**: Leverages Triton for GEMM (General Matrix Multiply), convolution, and SDPA (Scaled Dot-Product Attention) kernels. - **Quantization Support**: INT4 weight quantization with tile-packed format for improved performance and reduced memory footprint -- **Cross-Platform**: Supports both Linux and Windows platforms +- **Cross-Platform**: Runs on Linux, and on Windows from C++. Programs for Windows are exported on Linux. - **Multiple Model Support**: Works with various models including LLMs, vision-language models, and audio models ## Target Requirements Below are the requirements for running a CUDA-delegated ExecuTorch model: -- **Hardware**: NVIDIA GPU with CUDA compute capability -- **CUDA Toolkit**: CUDA 11.x or later (CUDA 12.x recommended) -- **Operating System**: Linux or Windows -- **Drivers**: PyTorch-Compatible NVIDIA GPU drivers installed +- **Hardware**: NVIDIA GPU with compute capability 8.0 (Ampere) or newer. This is the floor the + prebuilt packages are built for. +- **CUDA**: 13 for the prebuilt packages. On Linux they bring the CUDA runtime with them. On + Windows it comes from a CUDA 13 toolkit install. A build from source also supports CUDA 12.6. +- **Operating System**: Linux, or Windows from C++ for programs exported on Linux (see + [CUDA on Windows](../../using-executorch-cpp.md#cuda-on-windows)) +- **Drivers**: An NVIDIA driver that supports the CUDA version you use: CUDA 13 for the prebuilt + packages. ## Development Requirements To develop and export models using the CUDA backend: -- **Python**: Python 3.8+ +- **Python**: Python 3.11 to 3.15 - **PyTorch**: PyTorch with CUDA support +- **CUDA Toolkit**: export compiles kernels with `nvcc`, the CUDA compiler, so it needs a CUDA toolkit - **ExecuTorch**: Install ExecuTorch with CUDA backend support ## Using the CUDA Backend @@ -83,7 +88,7 @@ For a complete working example, see the [CUDA export script](https://github.com/ ## Runtime Integration -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. +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. When building from source, pass `-DEXECUTORCH_BUILD_CUDA=ON` when configuring the CMake build to compile the CUDA backend. @@ -118,7 +123,7 @@ exec_program = et_program.to_executorch() ```python # CPU tensors. The runtime copies them to the GPU and copies results back. -outputs = method([torch.randn(4, 64)]) +outputs = method.execute([torch.randn(4, 64)]) ``` This is the simplest option and the right default. The cost is a copy in each direction on @@ -140,24 +145,22 @@ exec_program = et_program.to_executorch( skip_d2h_for_method_outputs=True, ), enable_non_cpu_memory_planning=True, - # Required alongside the skips. Memory planning reserves a buffer for graph inputs by - # default, and the runtime fills a planned input by copying the caller's memory into that - # buffer, which reintroduces the copy you just asked to skip. On device memory that copy is - # a host memcpy into a device pointer, which is undefined. - # Outputs stay planned, so the runtime allocates them on the device and hands back a - # pointer to its own buffer. See the note below for when to unplan them instead. - memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + # Keeps the method from reserving its own input and output memory. + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), ) ) ``` ```python # Now the tensors must already be on the GPU. Passing CPU tensors is a caller error. -outputs = method([torch.randn(4, 64, device="cuda")]) +outputs = method.execute([torch.randn(4, 64, device="cuda")]) ``` Outputs also stay on the device, which is what you want when the next stage of your pipeline is -also on the GPU. +also on the GPU. From Python, each output comes back in GPU memory that the next call does not +overwrite, unless you pass `clone_outputs=False`. ### Things worth knowing @@ -170,24 +173,17 @@ it to `False` while asking for either skip raises a `ValueError`, because copy i during device-aware memory planning. **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the -program reserves its own input buffer and the runtime copies the caller's memory into it, so the -copy comes back at run time. On device memory that copy is a host memcpy into a device pointer, -which is undefined, so the program crashes rather than returning a wrong answer. - -**Outputs have two valid shapes, and which one you want depends on the language.** - -Leave outputs planned, which is the default, when the caller cannot supply device memory itself. -The planned buffer is allocated on the device because `enable_non_cpu_memory_planning` is on, and -`skip_d2h_for_method_outputs` means nothing copies it back, so the result stays on the device and -the caller receives a pointer to the runtime's own buffer. That buffer belongs to the method and is -reused by the next execution, so read or copy the result before running again. This is the shape -every Python caller wants, and it is what the exporters in this repository use, for example -`examples/models/gemma4_31b/export.py`. - -Set `alloc_graph_output=False` only when the caller supplies the output memory itself, which today -means a C++ caller handing the runtime a device tensor through `Module::set_output`. The Python -bindings allocate host memory for an unplanned output, so a device resident method exported this -way cannot run from Python. +program reserves its own input buffer on the device, and the runtime refuses to copy the caller's +tensor into it. The call fails with error `0x10` (`NotSupported`), and the log says +`Planned-buffer copy needs host memory on both sides`. + +**Leave outputs unplanned too**, via `alloc_graph_output=False`. Then each result lands in memory +you own, with no extra copy. A C++ caller must give each unplanned output a buffer once with +`Module::set_output`, before the first run: GPU memory when the output copy is skipped, host +memory otherwise. Pass each output's index to `set_output`; the C++ example below sets output 0. If +you keep outputs planned instead, they live in the method's own memory and the next run overwrites +them. Python's `execute()` then copies each output on every call, unless you pass +`clone_outputs=False`. **Per-method selection is on the outer config, not on the flags.** Pass a dict of `PropagateDeviceConfig` keyed by method name: @@ -208,6 +204,29 @@ Do not pass a dict to `skip_h2d_for_method_inputs` itself. The pass reads that f truthiness, so any non-empty dict enables the skip for every method regardless of the values inside it. +**Some copies stay even with the skip flags.** In the cases below, `to_executorch` keeps a copy +that runs on every call and logs a warning. To remove the copy, have the partitioner put more of +the graph in the delegate. + +- An input stays on the device of the first delegate that reads it. Other readers, such as an + operator outside the delegate, a delegate on another device, or the method output without + `skip_d2h_for_method_outputs`, read a copy in host memory. A delegate on another device then + gets a second copy, from the host. +- A constant, such as a weight, that a delegate reads from outside is copied to each device that + reads it. +- A returned output that an operator outside the delegate writes in place, such as a mutable + custom operator left on the host, is copied to the host and returned there. The method + metadata gives the device of each output. + +Add `error_on_per_call_copies=True` to `PropagateDeviceConfig` to make `to_executorch` fail +instead of warning. It fails only on a copy that one of the skip flags you set was meant to remove. + +**`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if +the method changes an input, even when the delegate makes the change. This includes a mutable +buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every +operator that uses it in the delegate. The CUDA partitioner already moves the buffer in, with +`tag_mutated_buffer`, once those operators are there. + **The choice is baked into the `.pte`.** A program exported with copies expects host tensors; one exported without them expects device tensors. Passing the wrong kind is a caller error, not something the runtime corrects, so the two are not interchangeable at run time. @@ -221,6 +240,11 @@ device right after the scalar type. Nothing is copied: ```cpp auto input = from_blob(device_data, {rows, columns}, ScalarType::Float, DeviceType::CUDA); +auto output = from_blob(device_result, {rows, out_columns}, ScalarType::Float, DeviceType::CUDA); +// Once, before the first run. Every run then writes its result to device_result. +if (module.set_output(output) != Error::Ok) { + return 1; +} auto result = module.forward(input); ``` @@ -235,6 +259,7 @@ across, which costs one transfer: ```cpp auto host_input = make_tensor_ptr({rows, columns}, std::move(data)); auto input = clone_tensor_ptr_to(host_input, DeviceType::CUDA); +// Set the output with module.set_output first, as above. auto result = module.forward(input); ``` diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index d85b427d688..411cc262961 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -1,87 +1,97 @@ # TensorRT Backend The TensorRT backend runs ExecuTorch models on NVIDIA GPUs with -[TensorRT](https://developer.nvidia.com/tensorrt), NVIDIA's inference library. At export -time the model graph is compiled into TensorRT engines. The engines are stored inside the -`.pte` file, and at run time the delegate hands each one back to TensorRT. +[TensorRT](https://developer.nvidia.com/tensorrt), NVIDIA's inference library. At export time, +TensorRT compiles the parts of the model it supports into engines (programs optimized for your +GPU), and by default the engines are stored in the `.pte` file. At run time, the TensorRT delegate +hands each engine back to TensorRT. The delegate is built and shipped by the [Torch-TensorRT](https://github.com/pytorch/TensorRT) -project, so it comes from its wheel and not from the ExecuTorch wheel. Everything after -installation is the normal ExecuTorch API: the same `.pte` file, the same `Module` class, the -same runtime. +project, not by ExecuTorch. You export with Torch-TensorRT, and you run with the normal ExecuTorch +runtime and its usual APIs. The +[Torch-TensorRT ExecuTorch guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html) +covers every option in detail. + +This page describes Torch-TensorRT 2.15 and newer. ## Features -- **TensorRT engines inside a `.pte`.** One file holds the program and its engines, so the - runtime loads a model the usual way. -- **Works with the CUDA backend in one program.** Operators TensorRT cannot convert can be - compiled by the ExecuTorch [CUDA backend](../cuda/cuda-overview.md) instead of falling back - to the CPU. Both delegates live in the same `.pte` and run in the same method. -- **GPU resident inputs and outputs.** The copies around the method boundary can be turned - off, so a pipeline that already holds its data on the GPU does not pay for them. -- **CUDA graph replay.** An engine can be recorded once and replayed with a single launch, - which helps models that launch many short kernels. This needs a Torch-TensorRT build newer - than the 2026-10-06 nightly, see [Feature availability](#feature-availability). -- **Shared activation scratch.** Many engines in one program can share one scratch buffer - instead of each keeping its own. -- **Shared engines.** Loading the same program twice keeps one copy of the engine weights in - GPU memory. This also needs a build newer than the 2026-10-06 nightly. -- **Caller-owned CUDA stream.** The application can give every delegate the same stream, - including a CUDA green context stream, so the model stays inside one part of the GPU. +- **TensorRT engines inside a `.pte`.** By default, one file holds the program and its engines. +- **Works with the CUDA backend in one program.** Operators that TensorRT cannot compile can go to + the ExecuTorch [CUDA backend](../cuda/cuda-overview.md) instead of the CPU, so the whole model + stays on the GPU. +- **Inputs and outputs that stay on the GPU.** The copies at the edges of the program can be + turned off, for a pipeline that already has its data on the GPU. +- **Your own CUDA stream.** The stream can also belong to a CUDA green context, which keeps the + model inside a fixed part of the GPU. TensorRT warns that this setup can crash, as + [Performance Options](#performance-options) explains. +- **CUDA graph replay.** The GPU work of an engine is recorded once and then replayed with a single + launch, which saves CPU time on every call. +- **Shared activation scratch.** Many engines can share one GPU buffer for their temporary results + (activations), which saves memory. +- **Shared engines.** Loading the same program twice keeps one copy of the engine weights in GPU + memory. +- **Weight streaming.** An engine can keep part of its weights in CPU memory. +- **Zero-copy KV cache.** The engine of a large language model (LLM) can update its key-value (KV) + cache in place. +- **Engines in a separate data file.** The engines can be kept in a `.ptd` file next to the `.pte`. +- **Optimization profiles.** One engine can carry tunings for different input sizes, and a C++ + application picks one per call. ## Target Requirements -- **Hardware**: NVIDIA GPU. -- **Operating system**: Linux, on x86_64 and on aarch64. -- **CUDA**: 13.x. The delegate links CUDA 13, so a CUDA 12 build is not supported. -- **Drivers**: an NVIDIA driver that matches the CUDA version. +- **Hardware**: an NVIDIA GPU that CUDA 13 supports. The prebuilt ExecuTorch CUDA packages need + compute capability 8.0 (Ampere) or newer. +- **Operating system**: Linux on x86-64 with glibc (the GNU C library) 2.28 or newer, Linux on + aarch64 with glibc 2.35 or newer, or Windows on x86-64. +- **CUDA**: 13. The prebuilt delegate links the CUDA 13 runtime. +- **Drivers**: an NVIDIA driver that supports CUDA 13. + +On Linux, TensorRT and the CUDA runtime come as dependencies of the Python packages, so you do not +need a system install of either. On Windows, TensorRT comes the same way, but the delegate loads +the CUDA runtime from a CUDA 13 toolkit install. Install the toolkit, and make sure `CUDA_PATH` or +`CUDA_HOME` points at it. -TensorRT itself and the CUDA runtime arrive as dependencies of the wheels, so a system -install of either is not required. +By default, TensorRT builds each engine for the operating system, the GPU and the TensorRT version +it is exported with. So export your programs on the kind of machine that will run them, with the +same packages. + +On Windows, ExecuTorch's CUDA backend cannot export a program, and ExecuTorch's Python runtime there +does not include the CUDA backend. So +[sending unsupported operators to the CUDA backend](#sending-unsupported-operators-to-the-cuda-backend) +is a Linux setup. A program that uses only TensorRT still runs from Python on Windows, because there +the delegate registers the CUDA memory allocator itself. ## Development Requirements -One command installs everything needed to export a model. The `executorch` extra pulls in a -CUDA build of ExecuTorch, the delegate wheel, Torch-TensorRT, and PyTorch: +One command installs everything, in a fresh virtual environment: ```bash -pip install --pre "torch-tensorrt[executorch]" \ - --index-url https://download.pytorch.org/whl/nightly/cu132 \ - --extra-index-url https://pypi.org/simple \ - --extra-index-url https://pypi.nvidia.com +pip install "torch_tensorrt[executorch]" ``` -Swap `cu132` for the channel that matches your CUDA, for example `cu134` for CUDA 13.4. Keep -PyTorch, ExecuTorch, and Torch-TensorRT on the same channel. - -All three indexes matter. Without `--pre` pip takes the latest stable Torch-TensorRT from the -public index, which pairs with an older ExecuTorch release than this page describes. Without -NVIDIA's index the TensorRT package resolves to a source build that takes about twenty -minutes and then fails. - -A CUDA build of ExecuTorch is required, not only to build against. The delegate needs a -library that only the CUDA wheels carry, so a CPU-only build installs and then fails on -import. - -If you also want the CUDA backend to pick up the operators TensorRT rejects, install a CUDA -toolkit as well. That backend compiles those operators at export time with `nvcc`. - -### Feature availability +That gives you Torch-TensorRT, PyTorch, a CUDA build of ExecuTorch, and the TensorRT delegate, +which is the `torch-tensorrt-executorch-runtime` package. On Windows, add the CUDA package index, +because PyPI has only the CPU builds of PyTorch and ExecuTorch for Windows: -Two features on this page are newer than the 2026-10-06 Linux nightly of Torch-TensorRT and -its delegate package, which was the newest one when this page was written: +```bash +pip install "torch_tensorrt[executorch]" --extra-index-url https://download.pytorch.org/whl/cu132 +``` -- [CUDA graph replay](#cuda-graph-replay). With the 2026-10-06 build, `torch_tensorrt.save` - rejects `use_cuda_graphs` with a `TypeError`. -- [Shared engines](#shared-engines). The 2026-10-06 delegate keeps every engine private, so - loading a program twice still holds two copies of its weights. +The +[installation section](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation) +of the Torch-TensorRT guide has the other options. -Both are on the Torch-TensorRT main branch and in its 2.15 release branch. To use them, -install a wheel newer than 2026-10-06 once one is published, or build Torch-TensorRT from -source. Everything else on this page works with the 2026-10-06 build. +The CUDA build of ExecuTorch is required even to run a program: the delegate needs ExecuTorch's +`extension_cuda` library, which only the CUDA builds ship. To send the operators that TensorRT +rejects to the CUDA backend, you also need a CUDA toolkit, because that backend compiles them at +export time with `nvcc`, the CUDA compiler. PyTorch finds `nvcc` on your `PATH` or under +`CUDA_HOME`. ## Using the TensorRT Backend +The Python examples on this page run in order, and later ones reuse names from earlier ones. + ### Exporting a model Compile the model with Torch-TensorRT, then save it in ExecuTorch format: @@ -93,68 +103,83 @@ import torch_tensorrt class MyModel(torch.nn.Module): def forward(self, x): - return x + 1 + return torch.relu(x) + 1 with torch.no_grad(): model = MyModel().eval().cuda() - example_input = (torch.randn(2, 3, 4, 4).cuda(),) + example_inputs = (torch.randn(2, 3, 4, 4, device="cuda"),) - exported_program = torch.export.export(model, example_input) + exported_program = torch.export.export(model, example_inputs) trt_gm = torch_tensorrt.dynamo.compile( exported_program, - arg_inputs=example_input, + arg_inputs=example_inputs, min_block_size=1, ) - torch_tensorrt.save( trt_gm, "model.pte", output_format="executorch", - arg_inputs=example_input, - retrace=False, ) ``` -`retrace=False` is required. Re-exporting the compiled graph drops the engine, and the error -only shows up much later when the program is loaded. +The Torch-TensorRT page on +[saving models](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/saving_models.html#executorch-save) +describes the other options of `save`. It also shows how to put several methods in one file, and +how to let the engine update a KV cache in place. -### Letting the CUDA backend take what TensorRT rejects +### Sending unsupported operators to the CUDA backend -TensorRT has no converter for every ATen operator. By default the leftover operators run on -the CPU, which costs a copy in each direction. Pass a `CudaPartitioner` and they are compiled -by the ExecuTorch CUDA backend instead, so the whole model stays on the GPU: +TensorRT cannot convert every PyTorch operator. By default the ones it rejects run on ExecuTorch's +CPU kernels, with a copy in each direction on every call. If an operator such as `erfinv` has no +CPU kernel, the method fails to load. Pass a `CudaPartitioner`, and the ExecuTorch CUDA backend +compiles those operators instead: ```python from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner -torch_tensorrt.save( - trt_gm, - "coalesced.pte", - output_format="executorch", - arg_inputs=example_input, - retrace=False, - partitioners=[ - CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) - ], -) + +class CoalescedModel(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + +with torch.no_grad(): + coalesced_model = CoalescedModel().eval().cuda() + coalesced_inputs = (torch.randn(64, 64, device="cuda"),) + + coalesced_gm = torch_tensorrt.dynamo.compile( + torch.export.export(coalesced_model, coalesced_inputs), + arg_inputs=coalesced_inputs, + min_block_size=1, + ) + torch_tensorrt.save( + coalesced_gm, + "coalesced.pte", + output_format="executorch", + partitioners=[ + CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) + ], + ) ``` -The TensorRT partitioner always runs first, and the `CudaPartitioner` picks up the rest. For -a model such as `cos(erfinv(tanh(x)))`, where TensorRT cannot take `erfinv`, the result is one -`.pte` whose delegate list reads `['TensorRTBackend', 'CudaBackend', 'TensorRTBackend']`. +TensorRT takes what it can first, and the `CudaPartitioner` takes the rest. TensorRT cannot take +`erfinv`, so this program's delegates are `['TensorRTBackend', 'CudaBackend', 'TensorRTBackend']`. +The Torch-TensorRT docs call such a program *coalesced*. -The CUDA backend writes its external weights to a `.ptd` file beside the `.pte` under a fixed -name. Save each coalesced model into its own directory. A second export into the same -directory overwrites the first model's weights, and the first model then still loads and -returns a wrong answer with no error. +The export writes the weights of the CUDA part to a data file, `aoti_cuda_blob.ptd`, next to the +`.pte`. Pass that file when you load the program: if the CUDA part has weights, loading fails +without it. The file has the same name for every model, so a second export into the same directory +overwrites the first one's weights, and the first model can then return wrong answers with no +error. Save each model into its own directory. ### Keeping inputs and outputs on the GPU -By default ExecuTorch inserts a host to device copy before the first delegate and a device to -host copy after the last one, so a method is safe to call with CPU tensors. If your data is -already on the GPU, turn those copies off: +By default, the program copies its inputs to the GPU before the first delegate, and its outputs +back after the last one, so you can call it with CPU tensors. If your data is already on the GPU, +turn those copies off. In the flag names, h2d means host (CPU) to device (GPU), and d2h the other +way: ```python from executorch.exir import ExecutorchBackendConfig @@ -165,45 +190,43 @@ torch_tensorrt.save( trt_gm, "device_resident.pte", output_format="executorch", - arg_inputs=example_input, - retrace=False, backend_config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True, ), - enable_non_cpu_memory_planning=True, - memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), ), ) ``` -Three points are easy to miss: - -- **Both skip flags need `enable_non_cpu_memory_planning=True`.** Copy insertion happens - during device aware memory planning. Setting it to `False` raises a `ValueError`. -- **Inputs must be unplanned**, with `MemoryPlanningPass(alloc_graph_input=False)`. Without - it the program reserves its own input buffer and the runtime copies the caller's memory - into it, which puts the copy straight back. -- **Leave the outputs planned when Python runs the program.** The program's own device arena - then holds the output, and Python returns a copy of it as a CUDA tensor. With - `alloc_graph_output=False` the caller provides the output memory instead. A C++ caller does - that with `Module::set_output`. The Python bindings do it correctly starting with ExecuTorch - 1.6.0.dev20261008, but the 2026-10-06 delegate wheel pins ExecuTorch 1.6.0.dev20260925, so - with that wheel keep the outputs planned. - -The choice is baked into the `.pte`. A program exported this way wants CUDA tensors, so pass -CUDA tensors. A host tensor works only when TensorRT is the first to read that input. On most -GPUs it then stages a copy on every call, which is the cost the export existed to remove. When -the CUDA backend reads it first, as it can in a coalesced program, it refuses the host memory -and the call fails. +- The two skip flags work only with `enable_non_cpu_memory_planning`, which is on by default. + With either flag set, turning it off makes the export raise a `ValueError`. +- `alloc_graph_input=False` is required. Without it the program keeps its own input buffer, and + every call fails with error `0x10` (`NotSupported`). The log then says + `Planned-buffer copy needs host memory on both sides`. +- `alloc_graph_output=False` keeps the outputs out of the program's memory. From Python, the + runtime allocates GPU memory for each output, and the next call does not overwrite an output you + still hold, unless you pass `clone_outputs=False`. From C++, you must give the method your own + GPU buffer once with `Module::set_output`, before the first run, as the + [CUDA backend page](../cuda/cuda-overview.md#things-worth-knowing) shows. If you keep outputs + planned instead, the next call overwrites them, so Python's `execute()` copies them on every call + unless you pass `clone_outputs=False`. + +Pass CUDA tensors to a program exported this way. The Python runtime refuses a CPU tensor. From +C++, what happens to a CPU tensor depends on which delegate reads that input first: + +- TensorRT: the call works, but copies the tensor on every call. +- The CUDA backend: the call fails. ## Runtime Integration ### Python -Import the delegate package once, anywhere before a program is loaded. The import is what -registers the backend. Nothing else about your code changes: +Import the delegate package once, before you load a program. The import registers the delegate. +ExecuTorch's own CUDA backend needs no import, and nothing else in your code changes: ```python from pathlib import Path @@ -217,57 +240,36 @@ forward = program.load_method("forward") outputs = forward.execute((torch.ones(2, 3, 4, 4),)) ``` -If the delegate cannot be loaded, that import raises straight away instead of leaving the -failure to show up later as a program that will not load. - -A coalesced program needs no extra registration. Both backends are registered, and the -program records which parts go where. Its weights are a different matter. If the part the -CUDA backend compiled has weights, they are in the `aoti_cuda_blob.ptd` file saved beside the -`.pte`, and the runtime does not look for that file on its own. Pass it when you load: +If the delegate cannot load, the import raises a `DelegateCompatibilityError` right away. If the +export wrote a data file, such as `aoti_cuda_blob.ptd` from the CUDA backend, pass it when you load: ```python program = Runtime.get().load_program( Path("coalesced.pte"), data_path=Path("aoti_cuda_blob.ptd") ) +outputs = program.load_method("forward").execute((torch.ones(64, 64),)) ``` -Without it, loading a program whose CUDA part has weights fails. - -### C++ against the installed wheels - -The wheels ship a prebuilt delegate and a CMake package, so a C++ application can link them -without building anything from source: - -```cmake -find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) -find_package(executorch_backend_tensorrt REQUIRED) +A program that keeps its inputs and outputs on the GPU takes CUDA tensors and returns CUDA tensors: -target_link_libraries(my_app PRIVATE - executorch::runtime - executorch::backend_cuda - executorch::backend_tensorrt - executorch::kernels_optimized +```python +program = Runtime.get().load_program(Path("device_resident.pte")) +outputs = program.load_method("forward").execute( + (torch.ones(2, 3, 4, 4, device="cuda"),) ) ``` -`kernels_optimized` supplies the `et_copy` operators that move data across the method -boundary. `backend_cuda` is needed by any coalesced program, and also by a TensorRT only -program that still has those boundary copies, because it registers the device allocator they -use. +### C++ -The two packages live in two distributions, so point CMake at both. ExecuTorch is a namespace -package, so its path has to come from its distribution metadata: +The Python packages ship the prebuilt delegate and its CMake package, so a C++ application can link +it without building anything from source. The steps below are for Linux. On Windows, link the +same components and copy the DLLs beside your program, as +[CUDA on Windows](../../using-executorch-cpp.md#cuda-on-windows) shows, but export TensorRT programs +on Windows itself. -```bash -cmake -DCMAKE_PREFIX_PATH="$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import importlib.metadata as m, torch_tensorrt_executorch_runtime as r, pathlib; print(str(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake") + ";" + str(pathlib.Path(r.__file__).parent))')" ... -``` - -CMake 3.28 or newer is needed for the form above, because the `backend_cuda` component -rejects older versions. CMake is not part of the wheels, so install it yourself. - -There is no header to include. The delegate registers itself with the backend registry from a -static initializer inside the shared library, and everything after that is the ordinary -ExecuTorch C++ API: +**1. Write the program.** Running a program needs no header from the delegate. The delegate +registers itself when its library loads, and the rest is the ordinary ExecuTorch C++ API. Save +this as `main.cpp`: ```cpp #include @@ -280,8 +282,7 @@ using namespace executorch::extension; int main() { Module module("model.pte"); - std::vector data(2 * 3 * 4 * 4, 1.0f); - auto input = make_tensor_ptr({2, 3, 4, 4}, std::move(data)); + auto input = make_tensor_ptr({2, 3, 4, 4}, std::vector(2 * 3 * 4 * 4, 1.0f)); const auto outputs = module.forward(input); if (!outputs.ok()) { @@ -294,198 +295,158 @@ int main() { } ``` -For a coalesced program whose CUDA part has weights, pass the data file to the constructor -as well, for the same reason as in Python: - -```cpp -Module module("coalesced.pte", "aoti_cuda_blob.ptd"); -``` - -Linking the target also records the wheel's own library directory in your binary, which is -right for an application built against an installed wheel and wrong for anything you plan to -redistribute. Turn it off with `EXECUTORCH_BACKEND_TENSORRT_EMBED_RUNPATH` and -`CMAKE_SKIP_BUILD_RPATH`, both set before `find_package`. - -### Building from source - -The delegate source ships inside `libtorchtrt.tar.gz` as -`torch_tensorrt/src/torch_tensorrt/executorch/`. Turn on ExecuTorch's CUDA backend and the -extensions a `Module` app uses, add the delegate next to ExecuTorch, and link the target it -provides: +**2. Write `CMakeLists.txt`.** ```cmake -set(EXECUTORCH_BUILD_CUDA ON) -set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) -set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) -set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) -set(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP ON) -set(EXECUTORCH_BUILD_EXTENSION_MODULE ON) - -add_subdirectory("executorch") -add_subdirectory("torch_tensorrt/src/torch_tensorrt/executorch") - -target_link_libraries(my_runner PRIVATE - executorch - executorch::backends - executorch::extensions - executorch::kernels - executorch::backend_tensorrt -) -``` - -Turn on the CUDA backend even for a program that only uses TensorRT. It builds the shared -`extension_cuda` library the delegate links, and it registers the device allocator the -boundary copies use. Without it the delegate stops at configure time. It also needs PyTorch -installed in the Python environment that CMake finds. The other options are the extensions -the CUDA backend and the `Module` class depend on, and CMake stops with an error if one of -them is missing. - -Building the delegate needs CUDA Toolkit 12.5 or newer. `libextension_cuda` stays a shared -library on purpose, so that every CUDA-capable delegate in the process reads the same caller -stream. - -## Performance Options - -### Choosing the CUDA stream - -With no stream chosen, both delegates run on `cudaStreamPerThread`, the default stream of the -calling thread. A coalesced program run from one thread is therefore ordered correctly as it -is, with nothing to add. - -To run every delegate on a stream of your own, for example a green context stream, scope a -guard over the whole execution: - -```cpp -#include - -using namespace executorch::extension; - -cuda::CallerStreamGuard guard(stream); -module.forward(input); -``` - -The guard lives in the `extension_cuda` library, so add it to the link recipe above: - -```cmake -find_package(executorch REQUIRED COMPONENTS backend_cuda extension_cuda kernels_optimized) - -target_link_libraries(my_app PRIVATE executorch::extension_cuda) -``` - -One guard reaches every CUDA-capable delegate, because they all resolve the same shared -`libextension_cuda`. The stream must be on the same device as the engine. The CUDA backend -refuses a caller stream for a method that uses its own CUDA graphs. - -### Green contexts - -Because both delegates honor the caller's stream, that stream can be a CUDA green context -stream. A green context holds a fixed number of streaming multiprocessors, so the whole model -is confined to that part of the GPU and the rest stays free for other work. Create the stream -with `cuGreenCtxStreamCreate` and scope the same guard over it. - -This has been exercised on a 108 SM card with a green context holding 8 of them, running a -program split across the TensorRT delegate and the CUDA backend. - -### CUDA graph replay - -This feature needs a Torch-TensorRT build newer than the 2026-10-06 nightly, see -[Feature availability](#feature-availability). - -Replay is off by default. When it is on, the delegate records an engine's kernel launches -once as a CUDA graph and then replays the whole engine with one launch. It helps engines with -fixed shapes that launch many short kernels, where the CPU launch work is a large part of -each call. +cmake_minimum_required(VERSION 3.28) +project(my_app CXX) -Turn it on at export time: +find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) +find_package(executorch_backend_tensorrt REQUIRED) -```python -torch_tensorrt.save( - trt_gm, - "model.pte", - output_format="executorch", - arg_inputs=example_input, - retrace=False, - use_cuda_graphs=True, +add_executable(my_app main.cpp) +target_link_libraries(my_app PRIVATE + executorch::runtime + executorch::backend_cuda + executorch::backend_tensorrt + executorch::kernels_optimized ) ``` -It is not free. Each engine keeps one stable device buffer per input and output, and every -call copies each input in and each output out. For an engine with few kernels, or with large -inputs and outputs, those copies can cost more than the launches they save. Shapes that -change on every call never replay and still pay the copies. - -Some engines never replay, even with replay on. An engine on the shared activation scratch -below always takes the ordinary path, so turning on both options gives no replay. Engines -with aliased outputs, GPUs without stream-ordered memory, and drivers older than CUDA 12.5 -also keep the ordinary path, as does a call on a green context stream. The delegate logs the -reason once. - -Recording is also not safe next to every kind of work. While a recording is running, creating -or destroying a TensorRT engine elsewhere in the process is unsafe, and so is a whole device -sync such as `cudaDeviceSynchronize` or `torch.cuda.synchronize()`. Leave replay off when -other threads may do any of that. A C++ host can refuse a program's saved request for one -load: +- `kernels_optimized` provides the CPU operators: the copies between the CPU and the GPU at the + edges of the program, and any operator that stays on the CPU. +- `backend_cuda` registers the allocator for the program's GPU memory, so a program that uses only + TensorRT needs it too. +- You need CMake 3.28 or newer. On older versions `find_package(executorch ...)` fails, because the + package provides the `executorch::` targets only from 3.28. CMake is not part of the packages, so + install it yourself. -```cpp -#include - -using namespace executorch::extension; -using namespace executorch::runtime; +**3. Build and run.** The two CMake packages are inside two Python packages, so point CMake at +both: -BackendOptions<1> options; -options.set_option("use_cuda_graphs", false); -LoadBackendOptionsMap by_backend; -by_backend.set_options("TensorRTBackend", options.view()); -module.load(by_backend); +```bash +ET_CMAKE=$(python -c 'import importlib.metadata as m, pathlib; print(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake")') +DELEGATE=$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import pathlib, torch_tensorrt_executorch_runtime as r; print(pathlib.Path(r.__file__).parent)') +cmake -S . -B build -DCMAKE_PREFIX_PATH="${ET_CMAKE};${DELEGATE}" +cmake --build build +./build/my_app ``` -### Shared activation scratch +`TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1` in the second command lets Python import the +delegate package to read its path, without loading the delegate. Run `./build/my_app` from the +folder that holds `model.pte`. -A TensorRT execution context allocates its own activation scratch and holds it for as long as -the context lives. A model lowered to many small engines pays that cost many times over, and -can run out of GPU memory on the engine count alone. This option backs all of a device's -contexts from one buffer instead, grown to the largest size any call has asked for: +If the export wrote a data file, pass it to the constructor: ```cpp -#include - -using namespace executorch::runtime; - -BackendOptions<1> options; -options.set_option("use_shared_activation_scratch", true); -set_option("TensorRTBackend", options.view()); +Module module("coalesced.pte", "aoti_cuda_blob.ptd"); ``` -Set it before loading the methods whose contexts should use the pool. The memory saved is the -sum of the separate scratch sizes less the largest of them. - -The cost is parallelism. Contexts that share one buffer do not run at the same time on the -device. The backend holds a per device lock from the claim on the buffer through the enqueue, -so two calls on one device are serialized at submission. The pool also never shrinks, and a -device that has run a pooled engine must not be reset with `cudaDeviceReset()`. - -### Shared engines +The Torch-TensorRT guide explains +[how to ship such an application](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#run-from-c) +to other machines, and links to the instructions for building the delegate from source. -This feature needs a Torch-TensorRT build newer than the 2026-10-06 nightly, see -[Feature availability](#feature-availability). - -When the same program is loaded twice, for example once per robot arm, handles loaded from -the same engine bytes, for the same device and the same weight streaming request, share one -engine. The engine weights are in GPU memory once. Each handle still gets its own execution -context and buffers. +## Performance Options -This is on by default. Measured on an 8 GB Jetson Orin Nano, loading one robot policy a second -time grew GPU memory by 806 MiB without sharing and by 122 MiB with it. Those numbers are the -cost of the second load, not the total. Finding a match means hashing the engine bytes on -every load, which costs roughly 0.15 s per GB on a Jetson AGX Thor. Pass the -`use_shared_engines` load option as `false` to keep one module's engines private. +Each option below links to its section of the Torch-TensorRT guide. The guide's section +[Setting a runtime option](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#setting-a-runtime-option) +shows how to set the options that apply at run time, from C++ or from Python. + +- [**Choosing the CUDA stream.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#choosing-the-cuda-stream) + By default, both delegates run on `cudaStreamPerThread`, the default stream of the calling + thread, so a program split between them runs in the right order on one thread. To use your own + stream, put one guard around the call. It reaches every delegate: + + ```cpp + #include + + using namespace executorch::extension; + + cuda::CallerStreamGuard guard(stream); + module.forward(input); + ``` + + The guard is in the `extension_cuda` library. Your own CUDA calls, such as creating the stream, + also need the CUDA runtime library from a CUDA 13 toolkit. In `CMakeLists.txt`, replace the + `find_package(executorch ...)` line from step 2 with the first line below, add the second line + after it, and add the third line after `add_executable`: + + ```cmake + find_package(executorch REQUIRED COMPONENTS backend_cuda extension_cuda kernels_optimized) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(my_app PRIVATE executorch::extension_cuda CUDA::cudart) + ``` + + - The stream must be on the same device as the engine. + - The CUDA backend refuses a caller stream for a method that uses its own CUDA graph option. + - Synchronize the stream before you read GPU outputs, because the CUDA backend can return while + its work is still queued. + - The stream can belong to a CUDA green context. A green context owns a fixed number of the GPU's + streaming multiprocessors (SMs), so the model stays inside that part of the GPU. Creating its + stream with `cuGreenCtxStreamCreate` also needs `CUDA::cuda_driver`. TensorRT then warns that + the engine was built without green context constraints, which it says may cause crashes or + undefined behavior. Torch-TensorRT has no option to build engines with them. +- [**CUDA graph replay.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#cuda-graph-replay) + Turn it on when you save. It helps engines with fixed shapes that launch many short kernels: + + ```python + torch_tensorrt.save( + trt_gm, + "cuda_graphs.pte", + output_format="executorch", + use_cuda_graphs=True, + ) + ``` + + - It costs a copy of every input and output, and an engine whose input shapes change on every + call never replays. Calls on a green context stream, or on a stream of another CUDA context, + do not replay either, and the guide lists a few rarer cases. + - A program saved with replay on can still have it turned off at load, with the + `use_cuda_graphs` option. + - While a graph is being recorded, no other thread may: + - create or destroy a TensorRT engine or execution context (the object that runs an engine) + outside this delegate, + - call `cudaDeviceSynchronize()` or `torch.cuda.synchronize()`, + - or load a program that uses the CUDA backend, which can make such a sync. + + Recording happens again after every shape change, so leave replay off when other threads may + do these things. +- [**Shared activation scratch.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#shared-activation-scratch) + The `use_shared_activation_scratch` option, set for the whole process before the program loads, + lets all engines on a device share one buffer for their temporary results. That saves memory for + a model split into many engines, but two engines on the shared scratch never run at the same + time. The scratch never shrinks, and a device that used it must not be reset with + `cudaDeviceReset()`. +- [**Shared engines.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#shared-engines) + On by default: loading the same program twice, on the same device and with the same weight + streaming budget, keeps one copy of each engine in GPU memory. Engines are matched by a hash of + their bytes, which a crafted program could match on purpose, so every program in the process must + be trusted. To keep one program's engines private, pass the `use_shared_engines` option as + `false` when you load it. +- [**Weight streaming.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#weight-streaming) + An engine compiled with `enable_weight_streaming=True` can keep part of its weights in CPU memory. + Set how many bytes stay on the GPU with `weight_streaming_budget_per_engine` in + `torch_tensorrt.save`, or with the `weight_streaming_budget` load option. The load option takes + the number of bytes as a decimal string, or as an integer below 2 GiB. +- [**Engines in a separate data file.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#keep-the-engines-in-a-separate-file) + Pass `external_engine_data="engines"` to `torch_tensorrt.save` to write the engines to + `engines.ptd` instead of the `.pte`, and pass that file when you load the program. + The Python loader takes one data file, so for a program that also uses the CUDA backend, pass + `external_engine_data="aoti_cuda_blob"` to share its file. +- [**Optimization profiles.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#optimization-profiles) + One engine can be tuned for several ranges of input shapes, for example one for decoding one + token at a time and one for a long prompt. A C++ application picks the profile for each call + with `torch_tensorrt::executorch_backend::OptimizationProfileGuard`, from the + `torch_tensorrt/executorch/TensorRTBackend.h` header. Without the guard, every engine runs + profile 0, and Python cannot choose a profile. The Python packages do not carry that header, and + the guide explains where to get it. ## Examples - [Export a static shape model](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_static_shape.py) -- [Export a coalesced TensorRT and CUDA model](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_coalesced.py) +- [Export a model whose input shapes can change](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_dynamic_shape.py) +- [Export a model split between TensorRT and the CUDA backend](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_coalesced.py) - [Export a model that keeps its data on the GPU](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_device_resident.py) -- [C++ reference runner, including the green context option](https://github.com/pytorch/TensorRT/tree/main/examples/executorch_reference_runner) - -The Torch-TensorRT documentation covers the export options in more detail, including dynamic -shapes, programs with more than one method, and a zero copy KV cache for decoder models. +- [Export a decoder step whose KV cache stays in memory you own](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_kv_cache_decode.py) +- [Export one engine with two optimization profiles](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_multi_profile.py) +- [C++ and Python reference runners](https://github.com/pytorch/TensorRT/tree/main/examples/executorch_reference_runner) diff --git a/docs/source/compiler-delegate-and-partitioner.md b/docs/source/compiler-delegate-and-partitioner.md index 5dec025fb4c..2d8ef530971 100644 --- a/docs/source/compiler-delegate-and-partitioner.md +++ b/docs/source/compiler-delegate-and-partitioner.md @@ -132,6 +132,36 @@ number in the program, lets the planner reuse the bytes for tensors whose lifetimes do not overlap the delegate call, and makes the requirement visible to the memory planning inspection tools. +### Sharing one handle across calls + +A method can call the same lowered module more than once. For example, you lower +a step once with `to_backend`, call it in a loop from a wrapper module, and +export the wrapper again. By default each call gets its own handle (the object +that `init()` returns), so `init()` runs once per call. A backend can ask for +one handle for all of these calls: + +```python +return PreprocessResult( + processed_bytes=blob, + share_handle_across_calls=True, +) +``` + +Then `init()` runs once, and every call runs `execute()` on that handle with its +own arguments. Set it only if `execute()` does all of the following, in every +runtime version of the backend and with any runtime options: + +- It reads its arguments again on every call. +- It keeps nothing on the handle that is tied to one call or that changes later + results, for example saved argument pointers or a cache that each call + updates. +- It writes outputs only into the output tensors it is given. + +Sharing applies to calls of one lowered module. Separate partitions get separate +lowered modules, so they keep separate handles even when their blobs are +identical. A backend that declares scratch always gets one handle per call, +because the runtime keeps scratch on the delegate and each call has its own. + ## Backend Interfaces: Runtime Initialization and Execution During the runtime, the compiled blob from the `preprocess` function will be diff --git a/docs/source/getting-started.md b/docs/source/getting-started.md index 3ffde9e9f9a..b03a1cc9759 100644 --- a/docs/source/getting-started.md +++ b/docs/source/getting-started.md @@ -37,8 +37,10 @@ pip install executorch torch \ | NVIDIA GPU, CUDA 13.2 | `cu132` | | NVIDIA GPU, CUDA 13.4 | `cu134` | -The CUDA packages are built for Linux, on x86_64 and ARM64. Use the `cpu` -variant on macOS and on Windows. If your CUDA version is not in the table, +The CUDA packages are built for Linux, on x86_64 and ARM64, and for Windows on +x86_64. On Windows they run CUDA programs from C++, and programs for the CUDA +backend are exported on Linux (see [CUDA on Windows](using-executorch-cpp.md#cuda-on-windows)). +Use the `cpu` variant on macOS. If your CUDA version is not in the table, choose the closest lower one with the same major version, because CUDA works across minor versions but not across major ones. There is no package for CUDA 12, so on CUDA 12 use the `cpu` variant or build from source. On CUDA 13.0 or 13.1 no @@ -54,8 +56,7 @@ To get a change that has landed on the `main` branch but is not in a release yet, use a nightly build. These are rebuilt every day. Put `nightly/` in front of the variant name, for example `nightly/cu132`, and add `--pre` to the command, otherwise pip skips development versions. Nightly builds cover the same -variants. CUDA 13.4 is the newest, and until the next release it is in nightly -builds only. +variants. The second index is needed because a bare `--index-url` replaces PyPI instead of adding to it, and some dependencies are published only on PyPI. diff --git a/docs/source/runtime-python-api-reference.rst b/docs/source/runtime-python-api-reference.rst index 9345f7af53c..6473e5536bf 100644 --- a/docs/source/runtime-python-api-reference.rst +++ b/docs/source/runtime-python-api-reference.rst @@ -18,7 +18,7 @@ For detailed information on how APIs evolve and the deprecation process, please :members: method_names, load_method .. autoclass:: Method - :members: execute, metadata + :members: execute, set_output, metadata .. autoclass:: ExecuTorchResult :members: shape, strides, dtype, device, nbytes, to_host, from_host diff --git a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py index ae10b22f879..25e36237fec 100644 --- a/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py +++ b/examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py @@ -1292,7 +1292,9 @@ def __init__( def _tag_ios(self, node, fixed_point_type): quant_io_type = None - # tag sharding io + # tag sharding io: for quantized models, where adjacent partitions are directly + # connected (partition1 -> llama_fallback -> partition2) without redundant explicit Q/DQ nodes in between + # (partition1 -> Q -> llama_fallback -> DQ -> partition2). if exir_ops.edge.llama.fallback.default in [ u.target for u in list(node.users.keys()) ] + [node.target]: @@ -1337,7 +1339,16 @@ def compile(self, request: Request): self.dep_table[TagQuantIO] = [SplitGraph] if not request_data.skip_quantize: - fixed_point_type = {"io_type": torch.float32} + fixed_point_type = {} + + # Tag sharding I/O as uint16 for A16 quantization. + # TODO: Add uint8 dtype handling for A8 quantization when a use case arises. + if self.quant_recipe.get_act_bit_width() == 16: + fixed_point_type["io_type"] = torch.uint16 + else: + raise RuntimeError( + f"unknown io bit width: {self.quant_recipe.get_act_bit_width()}" + ) # setup quantized IO self.passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index a3a77e25373..95cebb75ed9 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -125,6 +125,7 @@ def to_backend( compile_specs=compile_specs, named_data_store_output=preprocess_result.data_store_output, scratch_specs=preprocess_result.scratch_specs, + share_handle_across_calls=preprocess_result.share_handle_across_calls, ) lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map @@ -620,6 +621,7 @@ def lower_all_submodules_to_backend( compile_specs=compile_spec, named_data_store_output=preprocess_result.data_store_output, scratch_specs=preprocess_result.scratch_specs, + share_handle_across_calls=preprocess_result.share_handle_across_calls, ) lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map, diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index 838d6fd6eea..ca7d3b6e66a 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -74,6 +74,16 @@ class PreprocessResult: # in this order. They are not arguments of the delegate call. scratch_specs: Optional[List[DelegateScratchSpec]] = None + # Whether every call of this lowered module in a method may run on one + # delegate handle, so the runtime initializes the backend once instead of + # once per call. Set it only if execute(), in every runtime version of the + # backend, reads its arguments again on every call, keeps nothing on the + # handle that is tied to one call or that changes later results, and writes + # outputs only into the output tensors it is given. Ignored when + # scratch_specs is not empty, because the runtime keeps scratch on the + # delegate and each call has its own. + share_handle_across_calls: bool = False + """ How to create a backend (for example, BackendWithCompilerDemo): diff --git a/exir/backend/test/BUCK b/exir/backend/test/BUCK index 4b0574edbc6..99913f65989 100644 --- a/exir/backend/test/BUCK +++ b/exir/backend/test/BUCK @@ -294,9 +294,12 @@ fbcode_target(_kind = runtime.python_test, "fbsource//third-party/pypi/hypothesis:hypothesis", ":backend_with_compiler_demo", ":demo_backend", + ":op_partitioner_demo", "//caffe2:torch", "//executorch/exir:lib", "//executorch/exir:schema", + "//executorch/exir:tensor", + "//executorch/exir/_serialize:lib", "//executorch/exir/backend:backend_api", "//executorch/exir/backend:compile_spec_schema", "//executorch/exir/tests:models", diff --git a/exir/backend/test/backend_with_compiler_demo.py b/exir/backend/test/backend_with_compiler_demo.py index aa60fa2154c..86e5e8d78f4 100644 --- a/exir/backend/test/backend_with_compiler_demo.py +++ b/exir/backend/test/backend_with_compiler_demo.py @@ -139,4 +139,6 @@ def preprocess( ), debug_handle_map=debug_handle_map, _delegate_info_meta="test", + # The runtime keeps no state and reads its arguments on every call. + share_handle_across_calls=True, ) diff --git a/exir/backend/test/test_delegate_scratch.py b/exir/backend/test/test_delegate_scratch.py index 68740982b29..755ef839240 100644 --- a/exir/backend/test/test_delegate_scratch.py +++ b/exir/backend/test/test_delegate_scratch.py @@ -97,6 +97,21 @@ def preprocess( assert TWO_BUFFER_SIZES[1] % ALIGNMENT != 0 +@final +class SharedHandleScratchTestBackend(BackendDetails): + """Asks for one handle across calls and for scratch, which cannot be shared.""" + + @staticmethod + def preprocess( + edge_program: ExportedProgram, compile_specs: List[CompileSpec] + ) -> PreprocessResult: + return PreprocessResult( + processed_bytes=b"shared-handle-scratch-backend-blob", + scratch_specs=[DelegateScratchSpec(nbytes=TWO_BUFFER_SIZES[0])], + share_handle_across_calls=True, + ) + + @final class TwoBufferScratchTestBackend(BackendDetails): """Declares two buffers, to pin down ordering and independent sizing.""" @@ -584,14 +599,14 @@ def test_scratch_is_inserted_when_view_copies_are_kept(self): (delegate_call,) = _delegate_calls(plan) self.assertEqual(_scratch_sizes(delegate_call), [12 * BYTES_PER_ELEMENT]) - def _lowered_single_delegate(self): + def _lowered_single_delegate(self, backend_id="DelegateScratchTestBackend"): exported = torch.export.export( SingleDelegateModule().eval(), (torch.randn(3, 4), torch.randn(3, 4)) ) edge = to_edge_transform_and_lower( exported, compile_config=EdgeCompileConfig(_check_ir_validity=False) ) - return to_backend("DelegateScratchTestBackend", edge.exported_program(), []) + return to_backend(backend_id, edge.exported_program(), []) def test_lowered_module_emits_its_own_program_with_scratch(self): # program() and buffer() run their own small pipeline rather than @@ -602,6 +617,33 @@ def test_lowered_module_emits_its_own_program_with_scratch(self): self.assertEqual(_scratch_sizes(delegate_call), [12 * BYTES_PER_ELEMENT]) self.assertTrue(lowered.buffer()) + def test_a_delegate_with_scratch_is_not_shared_across_calls(self): + # The runtime resolves scratch onto delegates_[delegate_index], so each + # call site that declares scratch needs its own entry even when the + # backend allows one handle across calls. + lowered = self._lowered_single_delegate("SharedHandleScratchTestBackend") + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return self.lowered(self.lowered(x, y), y) + + inputs = (torch.randn(3, 4), torch.randn(3, 4)) + plan = ( + exir.to_edge(torch.export.export(TwoCalls(), inputs)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + + first, second = _delegate_calls(plan) + self.assertEqual(len(plan.delegates), 2) + self.assertNotEqual(first.delegate_index, second.delegate_index) + self.assertEqual(_scratch_sizes(first), [TWO_BUFFER_SIZES[0]]) + self.assertEqual(_scratch_sizes(second), [TWO_BUFFER_SIZES[0]]) + def test_emitting_without_the_pass_is_an_error_not_a_silent_drop(self): # Any emit path that skips the pass must fail. The count comes from the # graph, not the declaration, so otherwise the program ships without diff --git a/exir/backend/test/test_lowered_backend_module.py b/exir/backend/test/test_lowered_backend_module.py index 57575f8a797..b2f2e1cc233 100644 --- a/exir/backend/test/test_lowered_backend_module.py +++ b/exir/backend/test/test_lowered_backend_module.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 copy +import io import operator import unittest from unittest.mock import patch @@ -12,14 +14,17 @@ import torch from executorch import exir -from executorch.exir import to_edge +from executorch.exir import EdgeProgramManager, to_edge, to_edge_transform_and_lower +from executorch.exir._serialize import _PTEFile, _serialize_pte_binary from executorch.exir.backend.backend_api import to_backend from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.test.backend_with_compiler_demo import ( BackendWithCompilerDemo, ) from executorch.exir.backend.test.demo_backend import DemoBackend -from executorch.exir.schema import DelegateCall, Program +from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo +from executorch.exir.schema import DelegateCall, DelegateScratch, Program +from executorch.exir.tensor import make_allocation_info from executorch.extension.pybindings.portable_lib import ( # @manual _load_for_executorch_from_buffer, @@ -107,6 +112,206 @@ def forward(self, x): torch.allclose(model_outputs[0], expected_res, atol=1e-03, rtol=1e-03) ) + def _repeated_sin_calls(self, backend_id, compile_specs, num_calls): + """Lowers sin once and returns it with a module that calls it num_calls times.""" + + class SinModule(torch.nn.Module): + def forward(self, x): + return torch.sin(x) + + model_inputs = (torch.ones(1),) + lowered_sin_module = to_backend( + backend_id, + to_edge(export(SinModule(), model_inputs, strict=True)).exported_program(), + compile_specs, + ) + + class RepeatedCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.sin = lowered_sin_module + + def forward(self, x): + for _ in range(num_calls): + x = self.sin(x) + return x + + edge = to_edge(export(RepeatedCalls(), model_inputs, strict=True)) + return lowered_sin_module, edge, model_inputs + + def _delegate_indices(self, plan): + return [ + instruction.instr_args.delegate_index + for instruction in plan.chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + + def test_calls_of_one_lowered_module_share_one_delegate(self): + lowered_sin_module, edge, _ = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=3, + ) + self.assertTrue(copy.deepcopy(lowered_sin_module).share_handle_across_calls) + + plan = edge.to_executorch().executorch_program.execution_plan[0] + self.assertEqual(len(plan.delegates), 1) + self.assertEqual(self._delegate_indices(plan), [0, 0, 0]) + args = [ + tuple(instruction.instr_args.args) + for instruction in plan.chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + self.assertEqual(len(set(args)), 3) + + def test_shared_delegate_matches_separate_calls_at_runtime(self): + lowered_sin_module, edge, model_inputs = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=3, + ) + program = edge.to_executorch() + self.assertEqual(len(program.executorch_program.execution_plan[0].delegates), 1) + + expected = model_inputs[0] + single_call = _load_for_executorch_from_buffer(lowered_sin_module.buffer()) + for _ in range(3): + expected = single_call.forward([expected])[0] + outputs = _load_for_executorch_from_buffer(program.buffer).forward( + list(model_inputs) + ) + self.assertTrue(torch.equal(outputs[0], expected)) + + def test_sharing_survives_a_serde_round_trip(self): + _, edge, _ = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=2, + ) + buffer = io.BytesIO() + exir.save(edge.exported_program(), buffer) + buffer.seek(0) + + plan = ( + EdgeProgramManager(exir.load(buffer)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + self.assertEqual(self._delegate_indices(plan), [0, 0]) + + def test_runtime_rejects_a_shared_delegate_with_scratch(self): + _, edge, model_inputs = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=2, + ) + + def delegate_calls(program): + return [ + instruction.instr_args + for instruction in program.execution_plan[0].chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + + program = edge.to_executorch().executorch_program + scratch = [DelegateScratch(allocation=make_allocation_info(1, 0), size=4)] + + def run(program): + return _load_for_executorch_from_buffer( + bytes(_serialize_pte_binary(pte_file=_PTEFile(program=program))) + ).forward(list(model_inputs)) + + # The same scratch on two separate delegates is valid, so the errors + # below come from the sharing alone. + separate = copy.deepcopy(program) + delegates = separate.execution_plan[0].delegates + delegates.append(copy.deepcopy(delegates[0])) + delegate_calls(separate)[1].delegate_index = 1 + for call in delegate_calls(separate): + call.scratch = scratch + run(separate) + + # The emitter never writes these, but another producer could. + for calls_with_scratch in ([0, 1], [0], [1]): + shared = copy.deepcopy(program) + for index in calls_with_scratch: + delegate_calls(shared)[index].scratch = scratch + with self.assertRaisesRegex(RuntimeError, "error: 0x23"): + run(shared) + + def test_partitioned_module_shares_one_delegate_across_calls(self): + class AddModule(torch.nn.Module): + def forward(self, x): + return x + x + + model_inputs = (torch.ones(2, 2),) + add = ( + to_edge_transform_and_lower( + export(AddModule(), model_inputs, strict=True), + partitioner=[AddMulPartitionerDemo()], + ) + .exported_program() + .module() + ) + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.add = add + + def forward(self, x): + return self.add(self.add(x)) + + plan = ( + to_edge(export(TwoCalls(), model_inputs, strict=False)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + self.assertEqual(self._delegate_indices(plan), [0, 0]) + + def test_each_call_uses_the_delegate_of_its_own_module(self): + class SinModule(torch.nn.Module): + def forward(self, x): + return torch.sin(x) + + model_inputs = (torch.ones(1),) + edge_sin = to_edge(export(SinModule(), model_inputs, strict=True)) + compile_specs = [CompileSpec("max_value", bytes([1]))] + first, second = ( + to_backend( + BackendWithCompilerDemo.__name__, + edge_sin.exported_program(), + compile_specs, + ) + for _ in range(2) + ) + + class TwoModules(torch.nn.Module): + def __init__(self): + super().__init__() + self.first = first + self.second = second + + def forward(self, x): + return self.first(self.second(self.first(x))) + + program = ( + to_edge(export(TwoModules(), model_inputs, strict=True)) + .to_executorch() + .executorch_program + ) + # Identical blobs are stored once, but separate modules keep separate + # delegates, and a repeated call goes back to its own module's delegate. + self.assertEqual(len(program.backend_delegate_data), 1) + self.assertEqual(self._delegate_indices(program.execution_plan[0]), [0, 1, 0]) + + def test_calls_keep_their_own_delegate_unless_the_backend_allows_sharing(self): + _, edge, _ = self._repeated_sin_calls(DemoBackend.__name__, [], num_calls=2) + + plan = edge.to_executorch().executorch_program.execution_plan[0] + self.assertEqual(len(plan.delegates), 2) + self.assertEqual(self._delegate_indices(plan), [0, 1]) + def test_emit_lowered_backend_module(self): module_list = [ models.Emformer(), diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index d8db9145232..efdfaab00cd 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -153,6 +153,10 @@ class _EmitterState: spec2id_dict: Dict[TensorSpec, int] = field(default_factory=dict) # Only literal arguments with no schema aliasing enter this per-plan pool. constant_cache: Dict[Union[int, Tuple[int, ...]], int] = field(default_factory=dict) + # pyre-ignore[11]: Annotation `LoweredBackendModule` is not defined as a type. + delegate_cache: Dict["LoweredBackendModule", int] = field( # noqa: F821 + default_factory=dict + ) def spec2id(self, spec: TensorSpec) -> int: """Map a TensorSpec to value index in the values array.""" @@ -1499,14 +1503,26 @@ def _emit_delegate( ) ) - backend_delegate = BackendDelegate( - id=lowered_module.backend_id, - processed=BackendDelegateDataReference( - location=DataLocation.INLINE, index=delegate_index - ), - compile_specs=lowered_module.compile_specs, - ) - self.emitter_state.delegates.append(backend_delegate) + backend_delegate_index = self.emitter_state.delegate_cache.get(lowered_module) + if backend_delegate_index is None: + backend_delegate = BackendDelegate( + id=lowered_module.backend_id, + processed=BackendDelegateDataReference( + location=DataLocation.INLINE, index=delegate_index + ), + compile_specs=lowered_module.compile_specs, + ) + self.emitter_state.delegates.append(backend_delegate) + backend_delegate_index = len(self.emitter_state.delegates) - 1 + # Scratch is planned per call site but the runtime keeps it on the + # delegate, so a delegate with scratch cannot be shared. + if ( + lowered_module.share_handle_across_calls + and not lowered_module.scratch_specs + ): + self.emitter_state.delegate_cache[lowered_module] = ( + backend_delegate_index + ) # TODO(angelayi) Will need to emit the kwargs too, in the correct order according to the # function's spec and with default arguments. This requires us to store the function's spec # in to_backend() @@ -1521,7 +1537,7 @@ def _emit_delegate( self.chain.instructions.append( Instruction( DelegateCall( - delegate_index=len(self.emitter_state.delegates) - 1, + delegate_index=backend_delegate_index, args=delegate_args, scratch=self._emit_delegate_scratch(scratch_specs) or None, ) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index dd3b2d8a6ca..8a381c2b0f0 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -71,6 +71,7 @@ class LoweredBackendModule(torch.nn.Module): _scratch_specs: List[ DelegateScratchSpec ] # Scratch buffers the backend needs at execute() time + _share_handle_across_calls: bool # Whether all calls may share one handle meta: Optional[Dict[str, Any]] # Metadata for the lowered module def __init__( @@ -81,6 +82,7 @@ def __init__( compile_specs: List[CompileSpec], named_data_store_output: Optional[NamedDataStoreOutput] = None, scratch_specs: Optional[List[DelegateScratchSpec]] = None, + share_handle_across_calls: bool = False, ) -> None: super().__init__() self._original_exported_program = edge_program @@ -89,6 +91,7 @@ def __init__( self._compile_specs = compile_specs self._named_data_store_output = named_data_store_output self._scratch_specs = scratch_specs or [] + self._share_handle_across_calls = share_handle_across_calls self.meta = None # pyre-ignore @@ -118,6 +121,7 @@ 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, scratch_specs=copy.deepcopy(self._scratch_specs, memo), + share_handle_across_calls=self._share_handle_across_calls, ) res.meta = copy.copy(getattr(self, "meta", {})) return res @@ -164,6 +168,13 @@ def scratch_specs(self) -> List[DelegateScratchSpec]: """ return self._scratch_specs + @property + def share_handle_across_calls(self) -> bool: + """ + Returns whether every call of this module in a method may share one delegate handle + """ + return self._share_handle_across_calls + # TODO(chenlai): consolidate the seriailization config with serialize_to_flatbuffer api def buffer( self, diff --git a/exir/native/native.py b/exir/native/native.py index ad85f110734..4df444e5340 100644 --- a/exir/native/native.py +++ b/exir/native/native.py @@ -38,13 +38,15 @@ from executorch.backends.native.preprocess import NativeDelegateInfo from executorch.backends.native.serialization.graph_serialize import ( + _compile_to_bytes, _same_tensor, collect_data_keys, deserialize_program, + map_data_keys, merge_programs, validate_program, ) -from executorch.backends.native.serialization.schema import Program +from executorch.backends.native.serialization.schema import OutputKind, Program from executorch.exir._serialize._ptn import write_ptn from torch.utils import _pytree as pytree @@ -52,6 +54,8 @@ from torch.export import ExportedProgram from torch.fx import GraphModule +_STATE_NAMESPACE = "state" + class NativeProgramManager: """One or more methods lowered for the native runtime, ready to save. @@ -255,13 +259,7 @@ def _delegate_info(lowered: object, method_name: str) -> NativeDelegateInfo: def _merge_constants( per_method: dict[str, dict[str, torch.Tensor]], ) -> dict[str, torch.Tensor]: - """Union per-method constants by fqn. - - serialize_program treats an fqn as one buffer across every method, so hold to - that here rather than letting a later method silently clobber an earlier one. - _same_tensor short-circuits on identity and shared storage, which is the usual - case for a weight referenced by several methods. - """ + """Merge constants by their final data keys, rejecting differing values.""" merged: dict[str, torch.Tensor] = {} source: dict[str, str] = {} for method_name, constants in per_method.items(): @@ -325,11 +323,11 @@ def _validate_merged_program( def _normalize_programs( programs: "ExportedProgram | dict[str, ExportedProgram]", ) -> dict[str, ExportedProgram]: - """Validate and normalize the public input before starting expensive lowering.""" + """Validate the public input, then functionalize each method for lowering.""" from torch.export import ExportedProgram as TorchExportedProgram if isinstance(programs, TorchExportedProgram): - return {"forward": programs} + programs = {"forward": programs} if not isinstance(programs, dict): raise TypeError( "to_native: programs must be an ExportedProgram or a dict mapping " @@ -347,12 +345,18 @@ def _normalize_programs( ) if not method_name: raise ValueError("to_native: method names must be non-empty strings.") + if method_name.split("/")[0] == _STATE_NAMESPACE: + raise ValueError( + f"to_native: method name {method_name!r} is reserved; " + f"{_STATE_NAMESPACE!r} namespaces mutable state keys." + ) if not isinstance(program, TorchExportedProgram): raise TypeError( f"to_native: method {method_name!r} must map to an " f"ExportedProgram, got {type(program).__name__}." ) - method_programs[method_name] = program + # Functionalize so each method's buffer writes appear in its signature. + method_programs[method_name] = program.run_decompositions({}) return method_programs @@ -364,6 +368,11 @@ def to_native( ``programs`` is a single ``ExportedProgram`` (lowered as the sole ``forward`` method) or a dict mapping method name to ``ExportedProgram``. + Immutable constant keys are namespaced as ``/``, including for a + single method. Keys mutated by any method become ``state/`` in every + method, so state remains shared across methods. Equal immutable tensors + deduplicate on save. + Nothing about lowering is configurable yet. The passes, the partitioner, and the edge compile config are all owned here, because correctness of the package depends on them. Exposing the compile config is worth reconsidering. @@ -378,9 +387,10 @@ def to_native( Raises: TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name dictionary containing only ``ExportedProgram`` values. - ValueError: If the method dictionary is empty, a method name is empty, a - method does not fully delegate to the native backend, a constant - differs between methods, or a data key the graph references has no + ValueError: If the method dictionary is empty, a method name is empty or + in the reserved ``state`` namespace, a method does not fully delegate + to the native backend, a mutable constant differs between methods, + namespaced keys conflict, or a data key the graph references has no backing tensor. """ method_programs = _normalize_programs(programs) @@ -394,9 +404,18 @@ def to_native( get_lowered_submodules, ) + # Methods in a package share mutable buffers, so a buffer one method writes + # is not constant in the others. Passes that fold buffer values into weights + # must skip any buffer written by any method, not just their own. + mutable_buffers = { + buffer + for program in method_programs.values() + for buffer in program.graph_signature.buffers_to_mutate.values() + } + edge = to_edge_transform_and_lower( method_programs, - transform_passes=get_default_passes(), + transform_passes=get_default_passes(mutable_buffers=mutable_buffers), partitioner=[ NativePartitioner(external_constants_tag=None, _serialize_as_ptn=True) ], @@ -437,13 +456,44 @@ def to_native( method_blobs[method_name] = module.processed_bytes method_constants[method_name] = _delegate_info(module, method_name).constants - ptg_blob = merge_programs(method_blobs) - constants = _merge_constants(method_constants) + program = deserialize_program(merge_programs(method_blobs)) + state_keys = _mutable_data_keys(program) | { + buffer.fqn + for method in program.methods + for buffer in method.mutable_buffers or [] + } + state = {key: f"{_STATE_NAMESPACE}/{key}" for key in state_keys} + data_keys = set(state.values()) + for method in program.methods: + constants_for_method = method_constants[method.name] + namespaced = { + key: f"{method.name}/{key}" + for key in constants_for_method + if key not in state + } + conflicts = set(namespaced.values()) & data_keys + if conflicts: + raise ValueError( + "to_native: namespaced constant keys already exist: " + f"{sorted(conflicts)}" + ) + data_keys.update(namespaced.values()) + mapping = state | namespaced + map_data_keys(method, lambda key, mapping=mapping: mapping.get(key, key)) + for buffer in method.mutable_buffers or []: + buffer.fqn = state[buffer.fqn] + for spec in method.output_specs or []: + if spec.kind == OutputKind.BUFFER_MUTATION and spec.target is not None: + spec.target = state.get(spec.target, spec.target) + method_constants[method.name] = { + mapping.get(key, key): tensor + for key, tensor in constants_for_method.items() + } - program = deserialize_program(ptg_blob) + constants = _merge_constants(method_constants) _validate_merged_program(program, constants) - return NativeProgramManager(ptg=ptg_blob, constants=constants) + return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants) __all__ = ["NativeProgramManager", "to_native"] diff --git a/exir/native/test/test_native.py b/exir/native/test/test_native.py index a45277bbc91..c2ee8591b7d 100644 --- a/exir/native/test/test_native.py +++ b/exir/native/test/test_native.py @@ -24,10 +24,16 @@ Graph, Method, MutableBufferSpec, + OutputKind, Program, ) from executorch.exir import to_edge_transform_and_lower -from executorch.exir._serialize._ptn import ALIASES_ENTRY, PTG_ENTRY, SAFETENSORS_ENTRY +from executorch.exir._serialize._ptn import ( + ALIASES_ENTRY, + PTG_ENTRY, + read_ptn, + SAFETENSORS_ENTRY, +) from executorch.exir.lowered_backend_module import ( executorch_call_delegate, get_lowered_submodules, @@ -60,7 +66,7 @@ def _test_equivalence(self, programs): flow keeps one single-method delegate blob per method and would ship the same constants in a .ptd. The graphs are compared modulo that wrapping, and the full key sets (safetensors owners + aliases) must match the - delegates' named-data keys. + delegates' named-data keys after namespacing immutable constants. """ method_programs = ( programs if isinstance(programs, dict) else {"forward": programs} @@ -87,6 +93,12 @@ def _test_equivalence(self, programs): ) native_program = deserialize_program(ptg_blob) + mutated_keys = { + c.data_key + for method in native_program.methods + for c in method.constants or [] + if c.mutated + } native_graphs = {m.name: m.graph for m in native_program.methods} native_keys = st_keys | set(aliases.keys()) @@ -113,9 +125,17 @@ def _test_equivalence(self, programs): # Constants may be embedded in the .pte (pte_data) or shipped to a # .ptd (external_data) depending on external_constants_tag; collect # both so the key set matches to_native's .ptn/.safetensors. - et_keys |= set(nds.pte_data.keys()) + method_keys = set(nds.pte_data.keys()) for key_to_entry in nds.external_data.values(): - et_keys |= set(key_to_entry.keys()) + method_keys |= set(key_to_entry.keys()) + et_keys |= { + ( + f"state/{key}" + if f"state/{key}" in mutated_keys + else f"{name}/{key}" + ) + for key in method_keys + } # Emission is outside the comparison, but keep proving the ExecuTorch path # still completes end to end through the unchanged NamedDataStore branch. @@ -172,10 +192,17 @@ def test_multi_method(self): "prefill": torch.export.export(model, (torch.randn(2, 8),)), "decode": torch.export.export(model, (torch.randn(1, 8),)), } - # Shared weights resolve to the same fqns across methods, so the key sets - # still line up between the two paths. - _, st_keys, _ = self._test_equivalence(programs) - self.assertEqual(len(st_keys), 2) # one weight + bias, shared by fqn + native_keys, st_keys, aliases = self._test_equivalence(programs) + self.assertEqual( + native_keys, + {"prefill/weight", "prefill/bias", "decode/weight", "decode/bias"}, + ) + self.assertEqual(len(st_keys), 2) + for key in ("weight", "bias"): + self.assertEqual( + aliases.get(f"prefill/{key}", f"prefill/{key}"), + aliases.get(f"decode/{key}", f"decode/{key}"), + ) def test_conv_relu(self): model = nn.Sequential( @@ -186,6 +213,18 @@ def test_conv_relu(self): _, st_keys, _ = self._test_equivalence(ep) self.assertEqual(len(st_keys), 2) # conv weight + bias + def test_conv_batch_norm(self): + for inplace in (False, True): + with self.subTest(inplace=inplace): + model = nn.Sequential( + nn.Conv2d(3, 4, 3, bias=False), + nn.BatchNorm2d(4), + nn.ReLU(inplace=inplace), + ).eval() + ep = torch.export.export(model, (torch.randn(1, 3, 8, 8),)) + _, st_keys, _ = self._test_equivalence(ep) + self.assertEqual(len(st_keys), 2) + def test_tied_weights(self): class Tied(nn.Module): def __init__(self): @@ -239,7 +278,7 @@ def forward(self, x): ep = torch.export.export(ConstAdd(), (torch.randn(2, 4),)) _, st_keys, _ = self._test_equivalence(ep) - self.assertTrue(any("lifted" in k for k in st_keys)) + self.assertEqual(st_keys, {"forward/_lifted_tensor_constant0"}) class ToNativeInputValidationTest(unittest.TestCase): @@ -259,6 +298,12 @@ def test_empty_method_name_is_rejected(self): with self.assertRaisesRegex(ValueError, "non-empty"): to_native({"": object()}) + def test_state_namespace_method_name_is_rejected(self): + for name in ("state", "state/a"): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "reserved"): + to_native({name: object()}) + def test_wrong_method_value_is_rejected(self): with self.assertRaisesRegex(TypeError, "must map to an ExportedProgram"): to_native({"forward": object()}) @@ -393,7 +438,219 @@ def test_dropped_public_input_is_rejected(self): self._check(self._graph_module(drop_input=True)) +class ConstantNamespaceTest(unittest.TestCase): + def _scalar_shift(self, offset: float): + class Shift(nn.Module): + def forward(self, x): + return x + offset + + return torch.export.export(Shift(), (torch.randn(2, 4),)) + + def _read(self, native) -> tuple[Program, dict[str, torch.Tensor]]: + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "m.ptn") + native.save(path) + ptg, constants = read_ptn(path) + return deserialize_program(ptg), constants + + def _lifted_keys(self, program: Program, method_name: str) -> list[str]: + (method,) = [m for m in program.methods if m.name == method_name] + return [c.data_key for c in method.constants or [] if "lifted" in c.data_key] + + def test_distinct_lifted_constants_are_namespaced(self): + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(2.0)} + native = to_native(programs) + self.assertEqual(native.methods, {"a", "b"}) + program, constants = self._read(native) + (a_key,) = self._lifted_keys(program, "a") + (b_key,) = self._lifted_keys(program, "b") + self.assertEqual(a_key, "a/_lifted_tensor_constant0") + self.assertEqual(b_key, "b/_lifted_tensor_constant0") + torch.testing.assert_close( + constants[a_key], torch.full_like(constants[a_key], 1.0), rtol=0, atol=0 + ) + torch.testing.assert_close( + constants[b_key], torch.full_like(constants[b_key], 2.0), rtol=0, atol=0 + ) + + def test_identical_lifted_constants_dedup_on_save(self): + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(1.0)} + native = to_native(programs) + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "m.ptn") + native.save(path) + ptg, constants = read_ptn(path) + with zipfile.ZipFile(path) as pkg: + owners = _safetensors_keys(pkg.read(SAFETENSORS_ENTRY)) + aliases = json.loads(pkg.read(ALIASES_ENTRY)) + + program = deserialize_program(ptg) + (a_key,) = self._lifted_keys(program, "a") + (b_key,) = self._lifted_keys(program, "b") + self.assertEqual(a_key, "a/_lifted_tensor_constant0") + self.assertEqual(b_key, "b/_lifted_tensor_constant0") + self.assertEqual(owners, {a_key}) + self.assertEqual(aliases, {b_key: a_key}) + self.assertIs(constants[a_key], constants[b_key]) + + def test_distinct_parameters_are_namespaced(self): + programs = { + "a": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), + "b": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), + } + program, constants = self._read(to_native(programs)) + self.assertEqual(set(constants), {"a/weight", "a/bias", "b/weight", "b/bias"}) + for method in program.methods: + self.assertEqual( + {c.data_key for c in method.constants}, + {f"{method.name}/weight", f"{method.name}/bias"}, + ) + for key, tensor in programs[method.name].state_dict.items(): + torch.testing.assert_close(constants[f"{method.name}/{key}"], tensor) + + def test_existing_immutable_key_is_also_namespaced(self): + class ConstWithBuffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.ones(4)) + self.register_buffer("forward/weight", torch.zeros(4)) + + def forward(self, x): + return x + self.weight + self.get_buffer("forward/weight") + + ep = torch.export.export(ConstWithBuffer(), (torch.randn(4),)) + program, constants = self._read(to_native(ep)) + self.assertEqual(set(constants), {"forward/weight", "forward/forward/weight"}) + self.assertEqual( + {c.data_key for c in program.methods[0].constants}, set(constants) + ) + torch.testing.assert_close(constants["forward/weight"], torch.ones(4)) + torch.testing.assert_close(constants["forward/forward/weight"], torch.zeros(4)) + + def test_mutable_key_is_namespaced_as_state(self): + class ConstWithMutableBuffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.ones(4)) + self.register_buffer("forward/weight", torch.zeros(4)) + + def forward(self, x): + buffer = self.get_buffer("forward/weight") + buffer.add_(x) + return buffer + self.weight + + ep = torch.export.export(ConstWithMutableBuffer(), (torch.randn(4),)) + _, constants = self._read(to_native(ep)) + self.assertEqual(set(constants), {"forward/weight", "state/forward/weight"}) + torch.testing.assert_close(constants["forward/weight"], torch.ones(4)) + torch.testing.assert_close(constants["state/forward/weight"], torch.zeros(4)) + + def test_zero_initialized_mutable_buffer_is_namespaced_as_state(self): + program, constants = self._read( + to_native( + self._buffer_program(torch.zeros(4), mutate=True, persistent=False) + ) + ) + self.assertEqual(constants, {}) + (method,) = program.methods + self.assertEqual([b.fqn for b in method.mutable_buffers], ["state/shared"]) + self.assertEqual(self._mutation_targets(method), ["state/shared"]) + + def _mutation_targets(self, method: Method) -> list[str]: + return [ + spec.target + for spec in method.output_specs or [] + if spec.kind == OutputKind.BUFFER_MUTATION + ] + + def _buffer_program(self, buffer, mutate, key="shared", persistent=True): + class Buffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer(key, buffer, persistent=persistent) + + def forward(self, x): + value = self.get_buffer(key) + if mutate: + value.add_(x) + return x + value + + return torch.export.export(Buffer(), (torch.ones(4),)) + + def test_namespaced_keys_from_different_methods_must_not_collide(self): + programs = { + "a": self._buffer_program(torch.ones(4), mutate=False, key="b/shared"), + "a/b": self._buffer_program(torch.ones(4), mutate=False), + } + with self.assertRaisesRegex(ValueError, "namespaced constant keys"): + to_native(programs) + + def test_mutable_buffer_stays_shared_with_read_only_method(self): + shared = torch.zeros(4) + programs = { + "read": self._buffer_program(shared, mutate=False), + "update": self._buffer_program(shared, mutate=True), + } + program, constants = self._read(to_native(programs)) + self.assertEqual(set(constants), {"state/shared"}) + torch.testing.assert_close(constants["state/shared"], shared) + for method in program.methods: + (constant,) = method.constants + self.assertEqual(constant.data_key, "state/shared") + self.assertEqual(constant.mutated, method.name == "update") + self.assertEqual( + self._mutation_targets(method), + ["state/shared"] if method.name == "update" else [], + ) + + def test_conflicting_mutable_buffer_values_are_rejected(self): + programs = { + "read": self._buffer_program(torch.ones(4), mutate=False), + "update": self._buffer_program(torch.zeros(4), mutate=True), + } + with self.assertRaisesRegex( + ValueError, "constant 'state/shared' differs between methods" + ): + to_native(programs) + + class MutableBufferDedupTest(unittest.TestCase): + def test_batch_norm_statistics_mutated_by_another_method(self): + class Classifier(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 4, 3) + self.bn = nn.BatchNorm2d(4) + + def forward(self, x): + return self.bn(self.conv(x)) + + class Update(nn.Module): + def __init__(self, bn): + super().__init__() + self.bn = bn + + def forward(self, x): + self.bn.running_mean.add_(x) + return self.bn.running_mean + 1 + + model = Classifier().eval() + native = to_native( + { + "predict": torch.export.export(model, (torch.randn(1, 3, 8, 8),)), + "update": torch.export.export(Update(model.bn), (torch.ones(4),)), + } + ) + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "model.ptn") + native.save(path) + with zipfile.ZipFile(path) as package: + program = deserialize_program(package.read(PTG_ENTRY)) + predict = next(method for method in program.methods if method.name == "predict") + self.assertTrue( + any("batch_norm" in (node.target or "") for node in predict.graph.nodes) + ) + def test_identical_mutable_buffers_stay_separate(self): """Equal bytes must not make two independently mutable buffers one. diff --git a/exir/passes/BUCK b/exir/passes/BUCK index 6c896ea9f5d..7a25cbaa134 100644 --- a/exir/passes/BUCK +++ b/exir/passes/BUCK @@ -523,7 +523,9 @@ fbcode_target(_kind = runtime.python_library, "propagate_device_pass.py", ], deps = [ + ":convert_constant_dim_order_pass", ":device_copy_ops_registry", + ":insert_write_back_for_buffers_pass", ":propagate_device_config", "//caffe2:torch", "//executorch/exir:delegate", diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index ea600e1c222..324ff21db48 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -182,8 +182,23 @@ def get_propagated_const_tensor_dict( and 0 in prop_constant_tensor.stride() ): prop_constant_tensor = prop_constant_tensor.contiguous() + + # A folded constant must be a plain tensor that does not require + # grad. no_grad is not enough when the op returns a view of a + # parameter or the parameter itself. + if isinstance(prop_constant_tensor, torch.Tensor): + prop_constant_tensor = prop_constant_tensor.detach() const_node_to_tensor[node] = prop_constant_tensor + # Only tensors can be lifted into placeholders. A non-tensor value read by + # a node outside the fold stays a runtime op, together with the non-tensor + # values it is computed from. + for node in reversed(list(const_node_to_tensor)): + if not isinstance(const_node_to_tensor[node], torch.Tensor) and any( + user not in const_node_to_tensor for user in node.users + ): + del const_node_to_tensor[node] + return const_node_to_tensor diff --git a/exir/passes/convert_constant_dim_order_pass.py b/exir/passes/convert_constant_dim_order_pass.py index dd98fc9ece6..36037961f62 100644 --- a/exir/passes/convert_constant_dim_order_pass.py +++ b/exir/passes/convert_constant_dim_order_pass.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 logging + import torch from torch.export import ExportedProgram from torch.export.graph_signature import InputKind +logger: logging.Logger = logging.getLogger(__name__) + def _should_transform(tensor: torch.Tensor) -> bool: """ @@ -42,10 +46,20 @@ def _update_placeholder_meta( ), None, ) - if input_spec is None: + if placeholder_node is None: raise RuntimeError(f"Missing placeholder for {input_spec.arg.name}.") + strides = placeholder_node.meta["val"].stride() placeholder_node.meta["val"] = placeholder_node.meta["val"].contiguous() + logger.warning( + "%s '%s' has strides %s, which are neither contiguous nor channels-last. " + "to_edge stores it contiguous, so partitioners and delegates see the " + "contiguous layout, not the strides in the model. Make it contiguous " + "where it is created to avoid this change.", + kind.name.replace("_", " ").capitalize(), + target, + tuple(strides), + ) def convert_constant_dim_order_pass( @@ -54,9 +68,9 @@ def convert_constant_dim_order_pass( """ Normalize lifted tensors to default or channels-last dim order. - Tensor constants and buffers with other dim orders or striding are converted - to contiguous tensors. This pass acts in-place on the unlifted exported - program. + Tensor constants, buffers and parameters with other dim orders or striding + are converted to contiguous tensors. This pass acts in-place on the unlifted + exported program. Args: exported_program: The ExportedProgram to transform. @@ -84,6 +98,13 @@ def convert_constant_dim_order_pass( ) _update_placeholder_meta(exported_program, key, kind) + for key, param in exported_program.named_parameters(): + if _should_transform(param): + exported_program.state_dict[key] = torch.nn.Parameter( + param.contiguous(), requires_grad=param.requires_grad + ) + _update_placeholder_meta(exported_program, key, InputKind.PARAMETER) + # Convert buffers. for key, buffer in exported_program.named_buffers(): if ( diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py index 6071209ffeb..3d6dfeeefda 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -27,9 +27,9 @@ class PropagateDeviceConfig: """Controls whether the runtime copies method inputs and outputs across the device boundary. - Skipping a copy also means the runtime must not reserve its own buffer for the tensor, or it - would fill that buffer from the caller's memory and the copy would come back. Memory planning - allocates graph inputs and outputs by default, so a program using either skip below needs:: + Memory planning reserves buffers for graph inputs and outputs by default. A program that skips + the input copy must use ``alloc_graph_input=False``. Unplanning the outputs as well, as below, + copies the least:: ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( @@ -43,33 +43,63 @@ class PropagateDeviceConfig: ) Without ``alloc_graph_input=False`` the runtime reserves device memory for an input the caller - already owns, then copies into it with a host memcpy, which is undefined for device memory. + already owns and refuses to copy into it, so every call fails with ``NotSupported``. Planned + outputs work, but they live in the method's memory, which the next run overwrites, and the + Python ``execute()`` copies them unless it gets ``clone_outputs=False``. """ # When True, method-level input tensors that feed directly into a device # delegate are NOT wrapped with _h2d_copy. The user must provide tensors - # already on the target device. Useful for pipelines where inputs are - # pre-staged on GPU. + # already on the target device: the device of the first delegate that + # reads the input, in graph order. Any other reader of such an input, such + # as a host op, a delegate on another device, or the method output without + # skip_d2h_for_method_outputs, reads a copy of it in host memory, made on + # every call, and the pass logs a warning. A delegate on another device + # then gets a second copy, from the host. An input that an op writes in + # place on the host makes to_executorch raise. Useful for pipelines where + # inputs are pre-staged on GPU. Constants are not passed + # by the caller, so they are still copied: their data stays on the host, + # in the program or in its external data file. That copy runs on every + # call, once for each device that reads the constant, and stays in device + # memory until the last delegate there has read it. # # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer - # for the input and copies the caller's memory into it, which is what this exists to avoid. - # A dict can be used to set per-method values, keyed by method name. + # for the input and refuses the caller's tensor. + # For per-method values, pass a dict of PropagateDeviceConfig, keyed by method name, as the + # propagate_device_config of ExecutorchBackendConfig. skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False - # When True, device delegate outputs that are directly method outputs - # are NOT wrapped with _d2h_copy. The method outputs stay on device. - # Useful for cross-method GPU pipelines where the next method consumes - # GPU tensors directly. + # When True, device delegate outputs that the method returns stay on + # device, with one exception: when an op outside the delegate, such as a + # mutable custom op left on the host, writes the host copy of such an + # output in place, the method returns that host copy, since the device + # tensor never sees the write. The copy to the host runs on every call, and the pass logs a + # warning about it. + # Useful for cross-method GPU pipelines where the next method consumes GPU + # tensors directly. # - # Pair with MemoryPlanningPass(alloc_graph_output=False) for the same reason as the input flag. - # A dict can be used to set per-method values, keyed by method name. + # Planned outputs work, but they live in the method's memory, which the next run overwrites, + # and the Python runtime copies them. MemoryPlanningPass(alloc_graph_output=False) avoids + # that copy, and then a C++ caller must set each output with Module::set_output. + # For per-method values, pass a dict of PropagateDeviceConfig, keyed by method name, as the + # propagate_device_config of ExecutorchBackendConfig. skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False + # When True, the copies that the skip flags above cannot remove make + # to_executorch raise instead of logging a warning: with + # skip_h2d_for_method_inputs, a constant copied to the device and a method + # input copied to the host for a reader off its device, and with + # skip_d2h_for_method_outputs, a returned output copied to the host for a + # host op that writes it. Each runs on every call. Use it to make sure a + # program that skips its boundary copies has none left. + error_on_per_call_copies: bool = False + def __hash__(self) -> int: return hash( ( str(self.skip_h2d_for_method_inputs), str(self.skip_d2h_for_method_outputs), + self.error_on_per_call_copies, ) ) @@ -79,4 +109,5 @@ def __eq__(self, other: object) -> bool: return ( self.skip_h2d_for_method_inputs == other.skip_h2d_for_method_inputs and self.skip_d2h_for_method_outputs == other.skip_d2h_for_method_outputs + and self.error_on_per_call_copies == other.error_on_per_call_copies ) diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index aac2de75593..4e3ed5035bc 100644 --- a/exir/passes/propagate_device_pass.py +++ b/exir/passes/propagate_device_pass.py @@ -18,6 +18,13 @@ import torch.utils._pytree as pytree from executorch.exir.delegate import executorch_call_delegate from executorch.exir.lowered_backend_module import LoweredBackendModule +from executorch.exir.passes.convert_constant_dim_order_pass import _should_transform +from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + _alias_sets, + _contains_node, + _schema_paired_args, + _schema_writes_input, +) # Re-exported for backward compatibility; the dataclass lives in a lightweight # module so that ExecutorchBackendConfig can reference it without importing the @@ -35,6 +42,16 @@ # with this key and a value encoding the device string (e.g., b"cuda:0"). TARGET_DEVICE_COMPILE_SPEC_KEY = "target_device" +# Ops that read only the sizes, strides or offset of a tensor. Those stay in +# host memory when the data is on a device, so these ops need no copy of it. +_METADATA_READERS = ( + torch.ops.aten.sym_size.default, + torch.ops.aten.sym_size.int, + torch.ops.aten.sym_stride.int, + torch.ops.aten.sym_numel.default, + torch.ops.aten.sym_storage_offset.default, +) + def _parse_device_spec_value(value: bytes) -> tuple[schema.DeviceType, int]: """ @@ -135,36 +152,166 @@ def _tag_specs_with_device( return False -def _clone_spec_with_device( +def _copy_output_spec( spec: TensorSpec, device_type: schema.DeviceType, device_index: int = 0, ) -> TensorSpec: - """Create a copy of a TensorSpec with a different device.""" + """Create the spec of an _h2d_copy or _d2h_copy output from the spec it copies. + + The output lives on ``device_type``. The copy writes it at run time, so it is + never a constant, even when the copied tensor is one. + """ new_spec = copy.copy(spec) new_spec.init_mem_planning_fields() + new_spec.const = False _set_device_on_spec(new_spec, device_type, device_index) return new_spec +def _schema_aliases_input( + node: torch.fx.Node, schema: torch.FunctionSchema, input_node: torch.fx.Node +) -> bool: + """Whether schema marks the result of node as an alias of input_node, as for + a view or an in-place op.""" + returned = set().union(*(_alias_sets(ret.alias_info) for ret in schema.returns)) + return any( + _alias_sets(schema_arg.alias_info) & returned + for arg, schema_arg in _schema_paired_args(node, schema) + if _contains_node(arg, input_node) + ) + + +def _in_place_writers(node: torch.fx.Node) -> list[torch.fx.Node]: + """The ops that write *node* in place, directly or through a view of it, + per the op's schema. + + A mutable op left on the host does this once to_executorch has removed + auto_functionalized, so a copy of *node* can miss the write. If the op + wrote a view of *node* in the model, that removal makes the view again, + for example with as_strided. The result of an in-place op is a view of + its input too. A view_copy is not: it still makes its own tensor here, + and ReplaceViewCopyWithViewPass turns it into a view later only where no + write can tell the difference. Custom op schemas are trusted, as + functionalization trusts them. The write-back pass trusts only aten + schemas, but here that would keep a copy for every custom op on the host, + such as a KV cache update. Users without a schema, such as delegate calls + and getitem, do not write their inputs. + """ + writers: dict[torch.fx.Node, None] = {} + aliases = {node} + pending = [node] + while pending: + alias = pending.pop() + for user in alias.users: + schema = getattr(user.target, "_schema", None) + if schema is None: + continue + if _schema_writes_input(user, schema, alias): + writers[user] = None + if user not in aliases and _schema_aliases_input(user, schema, alias): + aliases.add(user) + pending.append(user) + return list(writers) + + +def _written_between( + node: torch.fx.Node, + after: Optional[torch.fx.Node], + before: torch.fx.Node, +) -> bool: + """Whether an op writes *node* in place after *after*, or anywhere when it + is None, and before *before*, in graph order, which is the order the + method runs in. A write outside that range cannot change what a reader at + *before* of a copy made at *after* sees.""" + writers = _in_place_writers(node) + if not writers: + return False + position = {n: i for i, n in enumerate(before.graph.nodes)} + start = position[after] if after is not None else -1 + return any(start < position[writer] < position[before] for writer in writers) + + +def _reusable_copy( + tensor: torch.fx.Node, + delegate: torch.fx.Node, + device_type: schema.DeviceType, + device_index: int, +) -> Optional[torch.fx.Node]: + """The _h2d_copy of *tensor* to this device that *delegate* can read too. + + Delegates are visited in graph order, so a copy made for an earlier + delegate on this device already comes before *delegate*, and the last such + copy is the newest. It is stale when an op writes *tensor* in place between + that copy and *delegate*, so then there is none. + """ + copy = None + for user in tensor.users: + if ( + user.target == torch.ops.et_copy._h2d_copy.default + and user.meta["spec"].device == device_type + and user.meta["spec"].device_index == device_index + ): + copy = user + if copy is not None and _written_between(tensor, copy, delegate): + return None + return copy + + +def _replace_spec( + readers: list[torch.fx.Node], old_spec: TensorSpec, new_spec: TensorSpec +) -> None: + """Replace *old_spec* with *new_spec* in *readers*, which now read the + tensor with *new_spec*, and in the readers of their results. + + SpecPropPass gives the output node the spec of each tensor it returns, and + an in-place result the spec of the tensor it writes. Memory planning relies + on both, so these specs have to follow the tensor that is read now. + """ + pending = list(readers) + while pending: + node = pending.pop() + spec = node.meta.get("spec") + if any(value is old_spec for value in pytree.tree_leaves(spec)): + node.meta["spec"] = pytree.tree_map( + lambda value: new_spec if value is old_spec else value, spec + ) + pending.extend(node.users) + + class PropagateDevicePass(PassBase): """ After to_backend, walk the graph and insert H2D/D2H copy ops at delegate boundaries based on partitioner-assigned device info. - When a delegate has a target_device CompileSpec (e.g., "cuda:0"): - - For each delegate input: insert et_copy._h2d_copy before the delegate call. - The original input node stays CPU; the h2d_copy output is tagged as device. - - For each delegate output: insert et_copy._d2h_copy after each getitem. - The getitem stays device; the d2h_copy output is tagged as CPU. + When a delegate has a target_device CompileSpec (e.g., "cuda:0") other + than the CPU: + - For each delegate input: insert et_copy._h2d_copy before the delegate call, + or reuse the one an earlier delegate on the same device reads when no op + writes the input in place in between. A reused copy stays in device + memory until its last reader. + The original input node keeps its device; the h2d_copy output is tagged as + device. + - For each delegate output: insert et_copy._d2h_copy after each getitem, + or after the call itself when it returns a single tensor. + The delegate output stays device; the d2h_copy output is tagged as CPU. - Getitem nodes that extract from a delegate call inherit the device. Skip-copy optimizations: - skip_h2d_for_method_inputs: If the input is a graph-level placeholder - feeding directly to a delegate, don't insert H2D — tag the placeholder - as device instead (user provides GPU tensor at runtime). - - skip_d2h_for_method_outputs: If the getitem feeds directly to graph - output, don't insert D2H — the output stays on device. + feeding a delegate, don't insert H2D; tag the placeholder as device + instead (user provides GPU tensor at runtime). Every other reader of the + placeholder reads one D2H copy of it. A placeholder that an op writes in + place on the host is refused. Constants are placeholders too, but their + data is stored on the host, so they still get an H2D copy. + - skip_d2h_for_method_outputs: The graph output reads a delegate output + on device, not its D2H copy; the copy stays only for other readers. When + one of them writes the copy in place, the graph output reads the written + host copy instead. + - error_on_per_call_copies: Raise instead of logging a warning when a + copy that a skip flag asked to avoid has to stay: a constant copied to a + device, a method input copied for a reader off its device, or a + returned output written on the host. """ def __init__( @@ -172,11 +319,13 @@ def __init__( skip_h2d_for_method_inputs: bool = False, skip_d2h_for_method_outputs: bool = False, enable_non_cpu_memory_planning: bool = False, + error_on_per_call_copies: bool = False, ) -> None: super().__init__() self.skip_h2d_for_method_inputs = skip_h2d_for_method_inputs self.skip_d2h_for_method_outputs = skip_d2h_for_method_outputs self.enable_non_cpu_memory_planning = enable_non_cpu_memory_planning + self.error_on_per_call_copies = error_on_per_call_copies if ( skip_h2d_for_method_inputs or skip_d2h_for_method_outputs @@ -193,9 +342,25 @@ def _is_placeholder(self, node: torch.fx.Node) -> bool: """Check if a node is a graph-level input (placeholder).""" return node.op == "placeholder" - def _feeds_directly_to_output(self, node: torch.fx.Node) -> bool: - """Check if all users of a node are output nodes.""" - return all(user.op == "output" for user in node.users) + def _delegate_device( + self, graph_module: torch.fx.GraphModule, node: torch.fx.Node + ) -> Optional[tuple[schema.DeviceType, int]]: + """The target device of a delegate call node, or None when there is none.""" + if node.op != "call_function" or node.target != executorch_call_delegate: + return None + lowered_module = _get_lowered_module(graph_module, node) + if lowered_module is None: + return None + return _get_target_device_from_compile_specs(lowered_module) + + def _report_per_call_copy(self, message: str, skip_requested: bool) -> None: + """Warn about a copy that runs on every call. When a skip flag asked + for no such copy and error_on_per_call_copies is set, raise instead.""" + if skip_requested and self.error_on_per_call_copies: + raise RuntimeError( + f"{message} Raised because error_on_per_call_copies=True." + ) + logger.warning("%s", message) def _insert_h2d_copies( self, @@ -204,7 +369,13 @@ def _insert_h2d_copies( target_device_type: schema.DeviceType, device_index: int, ) -> bool: - """Insert H2D copy nodes for each tensor input to a delegate call.""" + """Make each tensor input of a delegate call read device memory. + + A method input that is not a constant moves to the device under + skip_h2d_for_method_inputs. + Any other input reads an H2D copy, reused from an earlier delegate on + the same device when possible. + """ changed = False new_args = list(node.args) for i, arg in enumerate(node.args[1:], start=1): @@ -214,41 +385,149 @@ def _insert_h2d_copies( if not isinstance(arg_spec, TensorSpec): continue - if self.skip_h2d_for_method_inputs and self._is_placeholder(arg): - # TODO(gasoonjia): support skip_h2d_for_method_inputs for - # multiple-user placeholder inputs. - if len(arg.users) != 1: - raise RuntimeError( - f"skip_h2d_for_method_inputs=True requires placeholder " - f"'{arg.name}' to have exactly one user, but it has " - f"{len(arg.users)} users. The placeholder is shared by " - f"multiple consumers, so its TensorSpec cannot be safely " - f"mutated in-place to the delegate's device. Either disable " - f"skip_h2d_for_method_inputs, or ensure the placeholder is " - f"used exclusively by this delegate." - ) - _set_device_on_spec(arg_spec, target_device_type, device_index) + if ( + self.skip_h2d_for_method_inputs + and self._is_placeholder(arg) + and not arg_spec.const + ): + self._place_input_on_device( + graph_module, arg, target_device_type, device_index + ) changed = True continue - with graph_module.graph.inserting_before(node): - h2d_node = graph_module.graph.call_function( - torch.ops.et_copy._h2d_copy.default, - (arg,), - ) - h2d_spec = _clone_spec_with_device( - arg_spec, target_device_type, device_index + val = arg.meta.get("val") + if ( + arg_spec.const + and isinstance(val, torch.Tensor) + and _should_transform(val) + ): + raise RuntimeError( + f"Constant '{arg.name}' feeds a delegate on " + f"{target_device_type.name.lower()}:{device_index}, but it is " + f"neither contiguous nor channels-last. The program stores " + f"such a constant contiguous, so its device copy would not " + f"match the stored data. Make the constant contiguous where " + f"it is created, for example in the pass that adds it." ) - h2d_node.meta["spec"] = h2d_spec - h2d_node.meta["val"] = arg.meta.get("val") - if "tensor_meta" in arg.meta: - h2d_node.meta["tensor_meta"] = arg.meta["tensor_meta"] - new_args[i] = h2d_node - changed = True + + h2d_node = _reusable_copy(arg, node, target_device_type, device_index) + if h2d_node is None: + if arg_spec.const: + self._report_per_call_copy( + f"Constant '{arg.name}' ({arg_spec.nbytes()} bytes) feeds a " + f"delegate on {target_device_type.name.lower()}:" + f"{device_index} from outside it, so it is copied to the " + f"device on every call. Have the partitioner take it into " + f"the delegate to avoid that copy.", + self.skip_h2d_for_method_inputs, + ) + with graph_module.graph.inserting_before(node): + h2d_node = graph_module.graph.call_function( + torch.ops.et_copy._h2d_copy.default, + (arg,), + ) + h2d_node.meta["spec"] = _copy_output_spec( + arg_spec, target_device_type, device_index + ) + h2d_node.meta["val"] = arg.meta.get("val") + if "tensor_meta" in arg.meta: + h2d_node.meta["tensor_meta"] = arg.meta["tensor_meta"] + new_args[i] = h2d_node + changed = True node.args = tuple(new_args) return changed + def _place_input_on_device( + self, + graph_module: torch.fx.GraphModule, + placeholder: torch.fx.Node, + device_type: schema.DeviceType, + device_index: int, + ) -> None: + """Move a method input to the device for skip_h2d_for_method_inputs. + + Delegates on that device read it there, and so do ops that read only + its shape. Every other reader, such as a host op, a delegate on the + host, a delegate on another device or, without + skip_d2h_for_method_outputs, the method output, reads one D2H copy of it + instead, made on every call. An + input that an op writes in place on the host is refused: the write would + go to the copy, and neither the caller nor the delegates would see it. + Delegates are visited in graph order, so the first one that reads the + input decides its device, and later ones on that device find it placed. + """ + spec = placeholder.meta["spec"] + if (spec.device, spec.device_index) == (device_type, device_index): + return + device = f"{device_type.name.lower()}:{device_index}" + writers = _in_place_writers(placeholder) + if writers: + write_back = next( + ( + writer + for writer in writers + if writer.target == torch.ops.aten.copy_.default + and writer.args[0] is placeholder + ), + None, + ) + refusal = ( + f"skip_h2d_for_method_inputs=True cannot place placeholder " + f"'{placeholder.name}' on {device}, because " + f"'{(write_back or writers[0]).name}' writes it" + ) + if write_back is not None: + raise RuntimeError( + f"{refusal} back on the host (the method mutates it). Either " + f"disable skip_h2d_for_method_inputs, or, for a mutable buffer, " + f"have the partitioner put the buffer and every op that reads " + f"or writes it in the delegate; tag_mutated_buffer tags the " + f"buffer once those ops are there." + ) + raise RuntimeError( + f"{refusal} in place on the host. Either disable " + f"skip_h2d_for_method_inputs, or have the partitioner put that op " + f"in a delegate on {device}." + ) + + _set_device_on_spec(spec, device_type, device_index) + # With skip_d2h_for_method_outputs, the method returns the input where + # the caller put it. + other_readers = [ + user + for user in placeholder.users + if self._delegate_device(graph_module, user) != (device_type, device_index) + and user.target not in _METADATA_READERS + and not (user.op == "output" and self.skip_d2h_for_method_outputs) + ] + if not other_readers: + return + position = {n: i for i, n in enumerate(graph_module.graph.nodes)} + first_reader = min(other_readers, key=position.__getitem__) + self._report_per_call_copy( + f"Method input '{placeholder.name}' ({spec.nbytes()} bytes) stays on " + f"{device} for the delegates there, but '{first_reader.name}' also " + f"reads it, so it is copied to the host on every call, and from there " + f"to any other device that reads it. To avoid that, have the " + f"partitioner put every op that reads it in a delegate on {device}, " + f"and set skip_d2h_for_method_outputs if the method returns it.", + True, + ) + with graph_module.graph.inserting_before(first_reader): + d2h_node = graph_module.graph.call_function( + torch.ops.et_copy._d2h_copy.default, (placeholder,) + ) + d2h_spec = _copy_output_spec(spec, schema.DeviceType.CPU, 0) + d2h_node.meta["spec"] = d2h_spec + d2h_node.meta["val"] = placeholder.meta.get("val") + if "tensor_meta" in placeholder.meta: + d2h_node.meta["tensor_meta"] = placeholder.meta["tensor_meta"] + for reader in other_readers: + reader.replace_input_with(placeholder, d2h_node) + _replace_spec(other_readers, spec, d2h_spec) + def _insert_d2h_for_getitem( self, graph_module: torch.fx.GraphModule, @@ -280,16 +559,22 @@ def _insert_d2h_for_getitem( return False _set_device_on_spec(spec, source_spec.device, source_spec.device_index) + self._insert_d2h_after(graph_module, node, spec) + return True - if self.skip_d2h_for_method_outputs and self._feeds_directly_to_output(node): - return True - + def _insert_d2h_after( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + spec: TensorSpec, + ) -> None: + """Insert a D2H copy after *node*, a device tensor, for all of its users.""" with graph_module.graph.inserting_after(node): d2h_node = graph_module.graph.call_function( torch.ops.et_copy._d2h_copy.default, (node,), ) - d2h_spec = _clone_spec_with_device(spec, schema.DeviceType.CPU, 0) + d2h_spec = _copy_output_spec(spec, schema.DeviceType.CPU, 0) d2h_node.meta["spec"] = d2h_spec d2h_node.meta["val"] = node.meta.get("val") if "tensor_meta" in node.meta: @@ -299,15 +584,54 @@ def _insert_d2h_for_getitem( d2h_node, delete_user_cb=lambda user, _d2h=d2h_node: user != _d2h, ) - # Memory planning finds the graph outputs through the output node's - # spec, so it has to name the copy now rather than the device tensor. - for user in d2h_node.users: - if user.op == "output" and "spec" in user.meta: - user.meta["spec"] = pytree.tree_map( - lambda value: d2h_spec if value is spec else value, - user.meta["spec"], - ) - return True + _replace_spec(list(d2h_node.users), spec, d2h_spec) + + def _keep_method_outputs_on_device( + self, + graph_module: torch.fx.GraphModule, + ) -> None: + """Make the graph output read the device tensor instead of its D2H copy. + + The copy stays for its other users, such as a host op. When one of them + writes the copy in place, the device tensor never sees that write, so the + output keeps reading the written host copy, and a warning says so. + Copying it back to the device would add a second copy on every call. + """ + d2h = torch.ops.et_copy._d2h_copy.default + for node in list(graph_module.graph.nodes): + if ( + node.op != "call_function" + or node.target != d2h + or not any(user.op == "output" for user in node.users) + ): + continue + returned = node.args[0] + writers = _in_place_writers(node) + if writers: + device_spec = returned.meta["spec"] + index = list(graph_module.graph.output_node().args[0]).index(node) + self._report_per_call_copy( + f"Method output {index} ({device_spec.nbytes()} bytes) comes " + f"from a delegate on {device_spec.device.name.lower()}:" + f"{device_spec.device_index} and '{writers[0].target}' writes " + f"it in place on the host, so it is copied to the host on every " + f"call and returned there, even with " + f"skip_d2h_for_method_outputs. A caller that passes its own " + f"output buffers must pass a host buffer for it. Have the " + f"partitioner take that op into the delegate to avoid that copy.", + True, + ) + continue + _replace_spec( + [graph_module.graph.output_node()], + node.meta["spec"], + returned.meta["spec"], + ) + node.replace_all_uses_with( + returned, delete_user_cb=lambda user: user.op == "output" + ) + if not node.users: + graph_module.graph.erase_node(node) def _elide_redundant_same_device_copies( self, @@ -327,8 +651,10 @@ def _elide_redundant_same_device_copies( rewires the consumer straight to that device tensor and drops both copies. The device comparison is the correctness guard: a genuine cross-device move (e.g. ``cuda:0 -> host -> cuda:1``) has differing devices and is left - intact. A D2H feeding other (host) consumers keeps those users and is - removed only once it becomes dead. + intact. So is a round-trip whose host copy an op writes in place before + the copy back, because the device tensor never sees that write. A D2H + feeding other (host) consumers keeps those users and is removed only once + it becomes dead. """ h2d = torch.ops.et_copy._h2d_copy.default d2h = torch.ops.et_copy._d2h_copy.default @@ -356,6 +682,7 @@ def _elide_redundant_same_device_copies( if ( x_spec.device != h2d_spec.device or x_spec.device_index != h2d_spec.device_index + or _written_between(src, None, node) ): continue @@ -371,15 +698,20 @@ def _elide_redundant_same_device_copies( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Three-pass approach: # Pass 1 – For each delegate with a target_device CompileSpec, insert - # H2D copy nodes before delegate inputs and tag the delegate + # H2D copy nodes before delegate inputs, or reuse the copy an + # earlier delegate on the same device reads, and tag the delegate # output specs with the target device. Delegates without a - # target_device are left untouched (no copies, specs stay CPU). + # target_device, or with the CPU as their target_device, are + # left untouched (no copies, specs stay CPU). # Pass 2 – For each getitem that extracts from a device-tagged delegate # (tracked in device_delegates), propagate the device onto the # getitem spec and insert a D2H copy after it so downstream - # non-delegated ops receive CPU tensors. + # non-delegated ops receive CPU tensors. A delegate call that + # returns a single tensor gets the D2H copy itself. # Pass 3 – Elide the redundant device->host->device round-trips left # between consecutive same-device delegates by passes 1 and 2. + # With skip_d2h_for_method_outputs, then make the graph output + # read device tensors instead of D2H copies. changed = False device_delegates: set[torch.fx.Node] = set() @@ -395,7 +727,8 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 ) result = _get_target_device_from_compile_specs(lowered_module) - if result is None: + # A delegate on the host reads and writes host memory directly. + if result is None or result[0] == schema.DeviceType.CPU: continue target_device_type, device_index = result @@ -428,10 +761,20 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 lowered_module.backend_id, ) - # Second pass: propagate device through getitem nodes and insert D2H - # only for delegates that have a target_device. + # Second pass: propagate device through getitem nodes and insert D2H, + # after each getitem or single-tensor call, only for delegates that have + # a target_device. for node in list(graph_module.graph.nodes): - if node.op == "call_function" and node.target == operator.getitem: + # A lowered module that returns one tensor and is called from another + # module has no getitem after its call, so copy its output directly. + if ( + node in device_delegates + and self.enable_non_cpu_memory_planning + and isinstance(node.meta.get("spec"), TensorSpec) + ): + self._insert_d2h_after(graph_module, node, node.meta["spec"]) + changed = True + elif node.op == "call_function" and node.target == operator.getitem: source = node.args[0] if isinstance(source, torch.fx.Node) and source in device_delegates: if self.enable_non_cpu_memory_planning: @@ -457,11 +800,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Third pass: elide the redundant device->host->device round-trips left # between consecutive same-device delegates by the per-boundary H2D/D2H - # insertion above. Only meaningful when passes 1/2 modified the graph - # (there was a device delegate); when nothing changed there are no copies - # to elide, so skip the scan entirely. + # insertion above, then, with skip_d2h_for_method_outputs, make the graph + # output read device tensors. Only meaningful when passes 1/2 modified the + # graph (there was a device delegate); when nothing changed there are no + # copies to elide, so skip the scan entirely. if changed: self._elide_redundant_same_device_copies(graph_module) + if self.skip_d2h_for_method_outputs: + self._keep_method_outputs_on_device(graph_module) graph_module.recompile() return PassResult(graph_module, changed) diff --git a/exir/program/_program.py b/exir/program/_program.py index ed55b345362..b1d896281c0 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -793,6 +793,7 @@ def edge_to_executorch_passes( skip_h2d_for_method_inputs=device_cfg.skip_h2d_for_method_inputs, skip_d2h_for_method_outputs=device_cfg.skip_d2h_for_method_outputs, enable_non_cpu_memory_planning=config.enable_non_cpu_memory_planning, + error_on_per_call_copies=device_cfg.error_on_per_call_copies, ), EdgeToBackendOpsPass(), RemoveGraphAssertsPass(), @@ -823,8 +824,10 @@ def _generate_edge_program( program = remove_unused_parameters_pass(program) if not config._skip_dim_order: - # Normalize constants and buffers before MemoryFormatOpsPass rewrites - # preserve-format copies, which cannot represent arbitrary dim orders. + # Normalize constants, buffers and parameters before partitioning, so a + # delegate sees the layout the program stores, and before + # MemoryFormatOpsPass rewrites preserve-format copies, which cannot + # represent arbitrary dim orders. program = convert_constant_dim_order_pass.convert_constant_dim_order_pass( program ) @@ -870,7 +873,10 @@ def _generate_edge_program( # Lift the tensor constants created in ScalarToTensorPass edge_program = lift_constant_tensor_pass(edge_program) - # Normalize tensor constants introduced by the preceding passes. + # Normalize tensor constants introduced by the preceding passes. With + # _skip_dim_order this is the only call, so it fixes the layout of + # parameters, buffers and constants, but ops computed from them keep the + # old strides. edge_program = convert_constant_dim_order_pass.convert_constant_dim_order_pass( edge_program ) diff --git a/exir/serde/schema.py b/exir/serde/schema.py index f2347af7367..ff4cf26efb2 100644 --- a/exir/serde/schema.py +++ b/exir/serde/schema.py @@ -17,7 +17,7 @@ from executorch.exir.serde.union import _Union # NOTE: Please update this value if any modifications are made to the schema -SCHEMA_VERSION = (5, 5) +SCHEMA_VERSION = (5, 6) TREESPEC_VERSION = 1 @@ -415,3 +415,4 @@ class LoweredBackendModule: named_data_store: Optional[bytes] = None # The delegate scratch buffers, in declaration order. scratch_specs: Optional[List[DelegateScratchSpec]] = None + share_handle_across_calls: bool = False diff --git a/exir/serde/serialize.py b/exir/serde/serialize.py index d9cd9a700a7..677ac072ee0 100644 --- a/exir/serde/serialize.py +++ b/exir/serde/serialize.py @@ -314,6 +314,7 @@ def serialize_bytes(b: bytes) -> str: SerdeDelegateScratchSpec(nbytes=spec.nbytes, label=spec.label) for spec in lowered_module.scratch_specs ], + share_handle_across_calls=lowered_module.share_handle_across_calls, ) json_lowered_module = json.dumps( @@ -636,6 +637,7 @@ def deserialize_lowered_module( DelegateScratchSpec(nbytes=spec.nbytes, label=spec.label) for spec in (serialized_lowered_module.scratch_specs or []) ], + serialized_lowered_module.share_handle_across_calls, ) self.module.register_module(serialized_lowered_module_arg.name, lowered_module) return self.graph.get_attr(serialized_lowered_module_arg.name) diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 1afaece0f3a..2fa00618679 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2902,6 +2902,66 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: module = new_ep.module() self.assertFalse(torch.equal(module(x), module(x))) + def test_constant_prop_pass_detaches_folded_parameters(self) -> None: + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(4, 4)) + + def forward(self, x): + # `to` with the same dtype returns the parameter itself. + return x + self.w.t() + self.w.to(torch.float32) + + m = M() + x = torch.randn(4, 4) + new_ep = constant_prop_pass(export(m, (x,), strict=True)) + + targets = [node.target for node in new_ep.graph.nodes] + self.assertNotIn(torch.ops.aten.t.default, targets) + self.assertNotIn(torch.ops.aten.to.dtype, targets) + self.assertTrue(torch.allclose(new_ep.module()(x), m(x))) + copy.deepcopy(to_edge(new_ep).exported_program()) + + def test_constant_prop_pass_keeps_scalar_read_at_runtime(self) -> None: + """ + A non-tensor value cannot become a constant placeholder. When a node + outside the fold or the program output reads one, it stays in the + graph, while the tensors computed from it still fold. A non-tensor + value read only by folded nodes folds away. + """ + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("n", torch.tensor(1)) + self.register_buffer("s", torch.tensor(2.0)) + self.register_buffer("k", torch.tensor(3)) + self.w = torch.nn.Parameter(torch.randn(4, 3)) + + def forward(self, x): + n = (self.n * 2).item() + 1 + torch._check(n >= 1) + torch._check(n <= x.shape[0]) + # The first reader of `n` is folded, a later one is not. + w = (self.w * self.s.item())[:n] + return x[:n] * w, self.k.item() + + m = M() + x = torch.randn(4, 3) + new_ep = constant_prop_pass(export(m, (x,), strict=True)) + + targets = [node.target for node in new_ep.graph.nodes] + self.assertEqual(targets.count(torch.ops.aten.item.default), 2) + self.assertIn(torch.ops.aten._assert_scalar.default, targets) + self.assertEqual(targets.count(torch.ops.aten.mul.Tensor), 1) + self.assertEqual(targets.count(torch.ops.aten.slice.Tensor), 1) + self.assertNotIn("s", new_ep.state_dict) + out, k = new_ep.module()(x) + expected_out, expected_k = m(x) + self.assertTrue(torch.allclose(out, expected_out)) + self.assertEqual(k, expected_k) + to_edge(new_ep).to_executorch() + def test_constant_prop_pass_zero_stride_tensors(self) -> None: """ Test that constant propagation correctly handles tensors with zero strides @@ -3115,6 +3175,115 @@ def test_convert_constant_dim_order_to_contiguous(self): f"Constant should be contiguous after pass, got strides {modified_const.stride()}", ) + def test_convert_constant_dim_order_parameter(self): + """ + Verify that the convert_constant_dim_order_pass converts a parameter + outside the allowed dim orders to contiguous, keeps it a parameter, and + leaves a channels-last parameter as it is. + """ + + class AddParameter(torch.nn.Module): + def __init__(self, weight: torch.Tensor): + super().__init__() + self.weight = torch.nn.Parameter(weight) + + def forward(self, x): + return x + self.weight + + weights = { + "transposed": (torch.randn(5, 4).t(), True), + "channels_last": ( + torch.randn(2, 3, 4, 5).to(memory_format=torch.channels_last), + False, + ), + } + for description, (weight, converted) in weights.items(): + with self.subTest(description=description): + inputs = (torch.randn(weight.shape),) + ep = torch.export.export(AddParameter(weight), inputs) + logs = ( + self.assertLogs(convert_constant_dim_order_pass.logger, "WARNING") + if converted + else self.assertNoLogs(convert_constant_dim_order_pass.logger) + ) + with logs as captured: + modified_ep = ( + convert_constant_dim_order_pass.convert_constant_dim_order_pass( + copy.deepcopy(ep) + ) + ) + if converted: + self.assertIn( + "Parameter 'weight' has strides (1, 4)", captured.output[0] + ) + + torch.testing.assert_close( + ep.module()(*inputs), modified_ep.module()(*inputs) + ) + param = modified_ep.state_dict["weight"] + self.assertIsInstance(param, torch.nn.Parameter) + self.assertEqual(param.is_contiguous(), converted) + placeholder = next( + n for n in modified_ep.graph.nodes if n.name == "p_weight" + ) + self.assertEqual(placeholder.meta["val"].stride(), param.stride()) + + def test_to_edge_makes_transposed_parameter_contiguous(self): + """ + Verify that to_edge stores a transposed parameter contiguous with and + without _skip_dim_order, so a partitioner sees the stored layout. + """ + + class AddParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(5, 4).t()) + + def forward(self, x): + return self.weight + x + + model = AddParameter() + inputs = (torch.randn(4, 5),) + for skip_dim_order in (False, True): + with self.subTest(skip_dim_order=skip_dim_order): + with self.assertLogs( + convert_constant_dim_order_pass.logger, "WARNING" + ) as captured: + ep = to_edge( + torch.export.export(model, inputs), + compile_config=EdgeCompileConfig( + _skip_dim_order=skip_dim_order + ), + ).exported_program() + # Only one run of the pass converts: the first, or with + # _skip_dim_order the only one. + self.assertEqual(len(captured.output), 1) + self.assertIn("Parameter 'weight'", captured.output[0]) + self.assertTrue(ep.state_dict["weight"].is_contiguous()) + placeholder = next(n for n in ep.graph.nodes if n.name == "p_weight") + self.assertTrue(placeholder.meta["val"].is_contiguous()) + torch.testing.assert_close(ep.module()(*inputs), model(*inputs)) + + def test_convert_constant_dim_order_missing_placeholder(self): + """ + Verify that the convert_constant_dim_order_pass names the placeholder it + cannot find for a lifted tensor it converts. + """ + + class AddParameter(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(5, 4).t()) + + def forward(self, x): + return x + self.weight + + ep = torch.export.export(AddParameter(), (torch.randn(4, 5),)) + spec = next(s for s in ep.graph_signature.input_specs if s.target == "weight") + spec.arg = TensorArgument(name="missing") + with self.assertRaisesRegex(RuntimeError, "Missing placeholder for missing"): + convert_constant_dim_order_pass.convert_constant_dim_order_pass(ep) + class TestMemoryFormatOpsPassPreserveFormat(unittest.TestCase): """ diff --git a/exir/tests/test_propagate_device_pass.py b/exir/tests/test_propagate_device_pass.py index 41902ddced4..612e8a06290 100644 --- a/exir/tests/test_propagate_device_pass.py +++ b/exir/tests/test_propagate_device_pass.py @@ -13,6 +13,7 @@ import executorch.exir.passes._device_copy_ops_registry # noqa: F401 import torch from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir.backend.backend_api import to_backend from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.partitioner import ( DelegationSpec, @@ -29,7 +30,12 @@ from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult +from executorch.exir.passes import ( + MemoryPlanningPass, + propagate_device_pass, + ToOutVarPass, +) from executorch.exir.passes.propagate_device_pass import ( _get_target_device_from_compile_specs, _parse_device_spec_value, @@ -102,6 +108,23 @@ def _collect_device_copy_nodes(gm: torch.fx.GraphModule) -> DeviceCopyNodes: ) +def _lower_add_to_device(inputs: tuple, device: str) -> torch.nn.Module: + """Lower a module that returns ``x + y`` to a delegate on ``device``.""" + + class Add(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y) + + return to_backend( + BackendWithCompilerDemo.__name__, + to_edge(export(Add(), inputs)).exported_program(), + [ + CompileSpec("max_value", bytes([4])), + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")), + ], + ) + + class PerAddDeviceDelegatePartitioner(Partitioner): """Delegates each ``aten.add.Tensor`` as its OWN delegate, each with a per-op ``target_device``. @@ -111,10 +134,11 @@ class PerAddDeviceDelegatePartitioner(Partitioner): directly-connected adds become two directly-connected delegates -- the shape that leaves a redundant ``device -> host -> device`` round-trip between them. ``devices[i]`` is the target device for the i-th add in graph order (the last - entry is reused if there are more adds than devices). + entry is reused if there are more adds than devices). ``None`` gives that add + a delegate without a target device. """ - def __init__(self, devices: List[str]) -> None: + def __init__(self, devices: List[Optional[str]]) -> None: super().__init__() self._devices = list(devices) @@ -129,12 +153,13 @@ def partition(self, exported_program) -> PartitionResult: device = self._devices[min(i, len(self._devices) - 1)] tag = f"tag{i}" node.meta["delegation_tag"] = tag + compile_specs = [CompileSpec("max_value", bytes([4]))] + if device is not None: + compile_specs.append( + CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")) + ) partition_tags[tag] = DelegationSpec( - BackendWithCompilerDemo.__name__, - [ - CompileSpec("max_value", bytes([4])), - CompileSpec(TARGET_DEVICE_COMPILE_SPEC_KEY, device.encode("utf-8")), - ], + BackendWithCompilerDemo.__name__, compile_specs ) return PartitionResult( tagged_exported_program=exported_program, @@ -142,6 +167,80 @@ def partition(self, exported_program) -> PartitionResult: ) +# A mutable op that stays on the host, like a KV cache update. Once +# to_executorch removes auto_functionalized, it adds other into x in place. +_test_lib = torch.library.Library("propagate_device_test", "DEF") +_test_lib.define("accumulate(Tensor(a!) x, Tensor other) -> Tensor") +_test_lib.define( + "accumulate.out(Tensor(a!) x, Tensor other, *, Tensor(b!) out) -> Tensor(b!)" +) + + +@torch.library.impl(_test_lib, "accumulate", "CompositeExplicitAutograd") +def _accumulate(x: torch.Tensor, other: torch.Tensor) -> torch.Tensor: + x.add_(other) + return x.clone() + + +@torch.library.impl(_test_lib, "accumulate.out", "CompositeExplicitAutograd") +def _accumulate_out( + x: torch.Tensor, other: torch.Tensor, out: torch.Tensor +) -> torch.Tensor: + x.add_(other) + return out.copy_(x) + + +def _lower_with_host_writer( + model: torch.nn.Module, + inputs: tuple, + devices: List[str], + et_config: ExecutorchBackendConfig, +) -> torch.fx.GraphModule: + """Delegate each add of a model that also calls accumulate on the host.""" + # torch._export.config turns off auto_functionalized v2 for export in some + # PyTorch builds (see is_fbcode() there). Under v1, a custom op that writes + # through a view writes a view_copy and the write is lost, a limitation + # that predates these passes. Pin v2 so these tests check the same graph in + # every build. + with torch._export.config.patch(enable_auto_functionalized_v2_for_export=True): + edge = to_edge( + export(model, inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + et = edge.to_backend(PerAddDeviceDelegatePartitioner(devices)).to_executorch( + et_config + ) + return et.exported_program().graph_module + + +def _host_writer(gm: torch.fx.GraphModule) -> torch.fx.Node: + (writer,) = [ + n + for n in gm.graph.nodes + if n.target == torch.ops.propagate_device_test.accumulate.out + ] + return writer + + +def _host_written(gm: torch.fx.GraphModule) -> torch.fx.Node: + """The tensor that accumulate writes, looking through the view it may write.""" + written = _host_writer(gm).args[0] + if written.target == torch.ops.aten.as_strided.default: + return written.args[0] + return written + + +def _host_writer_config(write_view: bool, **kwargs) -> ExecutorchBackendConfig: + return ExecutorchBackendConfig( + emit_stacktrace=False, + enable_non_cpu_memory_planning=True, + # A write through a view leaves the aten.as_strided that + # to_executorch makes for it, which has no out variant. + to_out_var_pass=ToOutVarPass(ignore_to_out_var_failure=write_view), + **kwargs, + ) + + class TestPropagateDevicePass(unittest.TestCase): @staticmethod def _collect_tensor_specs(node: torch.fx.Node) -> List[TensorSpec]: @@ -450,6 +549,29 @@ def forward(self, a, b): self.assertEqual(len(device_copy_nodes.h2d_nodes), 0) self.assertEqual(len(device_copy_nodes.d2h_nodes), 0) + def test_no_copy_nodes_for_cpu_target_device(self): + """A delegate whose target device is the CPU reads and writes host + memory, so it gets no copies, like a delegate without a target device.""" + + class Model(torch.nn.Module): + def forward(self, a, b): + return torch.add(a, b) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, enable_non_cpu_memory_planning=True + ) + for pipeline, gm in _lower_model_to_executorch( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + PerAddDeviceDelegatePartitioner(["cpu"]), + et_config, + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.delegate_nodes), 1) + self.assertEqual(nodes.h2d_nodes, []) + self.assertEqual(nodes.d2h_nodes, []) + def test_redundant_same_device_copies_elided(self): """Two directly-connected delegates on the *same* device must not leave a redundant device->host->device round-trip between them. The @@ -493,6 +615,458 @@ def forward(self, a, b, c): f"(1 boundary D2H), got {len(nodes.d2h_nodes)}", ) + def test_same_device_copies_kept_when_host_op_writes_copy(self): + """A host op that writes the host copy of a delegate output in place, + directly or through a view, changes what the next delegate on the same + device reads, so that round-trip stays. The input it only reads still + gets one copy for both delegates.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + a = x + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(a, y) + return a + y + + for write_view in (False, True): + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(write_view), + ) + with self.subTest(write_view=write_view): + written = _host_written(gm) + self.assertEqual(written.target, torch.ops.et_copy._d2h_copy.out) + nodes = _collect_device_copy_nodes(gm) + self.assertIn( + written, + [ + arg.args[0] + for arg in nodes.delegate_nodes[1].args[1:] + if arg.target == torch.ops.et_copy._h2d_copy.out + ], + ) + y_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "y"] + self.assertEqual(len(y_copies), 1) + + def test_host_writes_after_the_device_readers_keep_no_copies(self): + """Host ops that write tensors in place only after the delegates on one + device have read them cannot change what those delegates read. So the + round-trip between the delegates is still removed, and the tensor both + delegates read gets one copy.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + t = x * 2 + a = t + y + b = a + t + # Both writes read b, so they run after the second delegate. + torch.ops.propagate_device_test.accumulate(a, b) + torch.ops.propagate_device_test.accumulate(t, b) + return a, b, t + + gm = _lower_with_host_writer( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(False), + ) + nodes = _collect_device_copy_nodes(gm) + first, second = nodes.delegate_nodes + order = list(gm.graph.nodes) + writers = [ + n + for n in order + if n.target == torch.ops.propagate_device_test.accumulate.out + ] + self.assertEqual(len(writers), 2) + for writer in writers: + self.assertLess(order.index(second), order.index(writer)) + # Only t and y are copied, and the second delegate reads the first + # delegate's output directly. + self.assertEqual(len(nodes.h2d_nodes), 2) + (shared,) = set(first.args[1:]) & set(second.args[1:]) + self.assertIn(shared, nodes.h2d_nodes) + + def test_single_tensor_delegate_output_read_on_host(self): + """A host op that reads the output of a lowered module returning one + tensor reads a host copy of it.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return torch.sin(self.lowered(x, y)) + + et = to_edge(export(Model(), inputs)).to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, enable_non_cpu_memory_planning=True + ) + ) + gm = et.exported_program().graph_module + nodes = _collect_device_copy_nodes(gm) + (d2h,) = nodes.d2h_nodes + self.assertIs(d2h.args[0], nodes.delegate_nodes[0]) + (sin,) = d2h.users + self.assertEqual(sin.target, torch.ops.aten.sin.out) + self._assert_nodes_buffer_device( + [d2h], et._emitter_output.program, DeviceType.CPU, "host op", "D2H copy" + ) + + def test_single_tensor_delegate_output_copied_to_host(self): + """A lowered module that returns one tensor, called from another module, + has no getitem after its call. Its device output still gets a _d2h_copy, + and the next call on the same device reads it without a round-trip. With + skip_h2d_for_method_inputs, y, read by both calls, is placed once.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return self.lowered(self.lowered(x, y), y) + + configs = ( + PropagateDeviceConfig(), + PropagateDeviceConfig(skip_h2d_for_method_inputs=True), + PropagateDeviceConfig(skip_d2h_for_method_outputs=True), + ) + for device_config in configs: + skip_h2d = device_config.skip_h2d_for_method_inputs + skip_d2h = device_config.skip_d2h_for_method_outputs + et = to_edge(export(Model(), inputs)).to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=device_config, + enable_non_cpu_memory_planning=True, + ) + ) + program = et._emitter_output.program + gm = et.exported_program().graph_module + case = f"skip_h2d={skip_h2d}, skip_d2h={skip_d2h}" + with self.subTest(case=case): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.delegate_nodes), 2) + # Only method inputs are copied to the device, x and y once + # each, unless the caller places them. + self.assertEqual(len(nodes.h2d_nodes), 0 if skip_h2d else 2) + self.assertTrue( + all(n.args[0].op == "placeholder" for n in nodes.h2d_nodes) + ) + cuda_ph, _ = self._collect_placeholders_by_device(gm) + self.assertEqual(len(cuda_ph), 2 if skip_h2d else 0) + # Only the method output can be copied back to the host. + self.assertEqual(len(nodes.d2h_nodes), 0 if skip_d2h else 1) + (output,) = gm.graph.output_node().args[0] + self.assertEqual( + output.target, + ( + executorch_call_delegate + if skip_d2h + else torch.ops.et_copy._d2h_copy.out + ), + ) + self._assert_nodes_buffer_device( + nodes.delegate_nodes, program, DeviceType.CUDA, case, "Delegate" + ) + self._assert_nodes_buffer_device( + nodes.d2h_nodes, program, DeviceType.CPU, case, "D2H copy" + ) + + def test_host_inplace_ops_on_delegate_output_write_host_copy(self): + """Host ops that write a delegate output in place, here relu_ and then + mul_ from reinplace_pass, write its host copy. Their results share the + copy's spec, as memory planning requires, and the method returns the + last one from the host.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Call(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return torch.relu(self.lowered(x, y)) * 2 + + class Add(torch.nn.Module): + def forward(self, x, y): + return torch.relu(x + y) * 2 + + cases = { + "single tensor": (Call(), None), + "getitem": (Add(), PerAddDeviceDelegatePartitioner(["cuda:0"])), + } + reinplaced = frozenset( + {exir_ops.edge.aten.relu.default, exir_ops.edge.aten.mul.Tensor} + ) + for case, (model, partitioner) in cases.items(): + for skip_d2h in (False, True): + edge = to_edge(export(model, inputs)) + if partitioner is not None: + edge = edge.to_backend(partitioner) + et = edge.to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + reinplace_extra_ops=reinplaced, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=skip_d2h + ), + enable_non_cpu_memory_planning=True, + ) + ) + gm = et.exported_program().graph_module + with self.subTest(case=case, skip_d2h=skip_d2h): + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + (relu,) = d2h.users + (mul,) = relu.users + self.assertEqual(relu.target, exir_ops.edge.aten.relu_.default) + self.assertEqual(mul.target, exir_ops.edge.aten.mul_.Tensor) + self.assertIs(relu.meta["spec"], d2h.meta["spec"]) + self.assertIs(mul.meta["spec"], d2h.meta["spec"]) + self.assertEqual(gm.graph.output_node().args[0], (mul,)) + self._assert_nodes_buffer_device( + [mul], + et._emitter_output.program, + DeviceType.CPU, + case, + "Output", + ) + + def test_skip_d2h_keeps_output_also_read_on_device(self): + """With skip_d2h_for_method_outputs, a returned delegate output that a + delegate on the same device also reads stays on the device, like the + other outputs.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + + class Calls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + a = self.lowered(x, y) + return a, self.lowered(a, y) + + class Adds(torch.nn.Module): + def forward(self, x, y): + a = torch.add(x, y) + return a, torch.add(a, y) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + cases = { + "single tensor": (Calls(), None), + "getitem": (Adds(), PerAddDeviceDelegatePartitioner(["cuda:0"])), + } + for case, (model, partitioner) in cases.items(): + edge = to_edge(export(model, inputs)) + if partitioner is not None: + edge = edge.to_backend(partitioner) + et = edge.to_executorch(deepcopy(et_config)) + gm = et.exported_program().graph_module + with self.subTest(case=case): + self.assertEqual(len(_collect_device_copy_nodes(gm).d2h_nodes), 0) + outputs = gm.graph.output_node().args[0] + self._assert_nodes_device(outputs, DeviceType.CUDA, case, "Output") + self._assert_nodes_buffer_device( + outputs, et._emitter_output.program, DeviceType.CUDA, case, "Output" + ) + + def test_skip_d2h_keeps_output_also_read_off_device(self): + """With skip_d2h_for_method_outputs, a returned delegate output that a + host op or a delegate on another device also reads stays on the device. + Only that other reader gets the host copy.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = _lower_add_to_device(inputs, "cuda:0") + lowered_cuda_1 = _lower_add_to_device(inputs, "cuda:1") + + class HostReader(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + a = self.lowered(x, y) + return a, torch.sin(a) + + class OtherDeviceReader(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + self.lowered_cuda_1 = lowered_cuda_1 + + def forward(self, x, y): + a = self.lowered(x, y) + return a, self.lowered_cuda_1(a, y) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + for case, model in { + "host op": HostReader(), + "delegate on cuda:1": OtherDeviceReader(), + }.items(): + et = to_edge(export(model, inputs)).to_executorch(deepcopy(et_config)) + gm = et.exported_program().graph_module + with self.subTest(case=case): + output = gm.graph.output_node().args[0][0] + self.assertEqual(output.target, executorch_call_delegate) + self._assert_nodes_device([output], DeviceType.CUDA, case, "Output", 0) + self._assert_nodes_buffer_device( + [output], + et._emitter_output.program, + DeviceType.CUDA, + case, + "Output", + ) + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + self.assertIs(d2h.args[0], output) + self.assertNotIn(gm.graph.output_node(), d2h.users) + + def test_skip_d2h_leaves_returned_output_unplanned(self): + """With skip_d2h_for_method_outputs and unplanned graph outputs, the + caller provides the returned delegate output, while the host copy that + a host op reads is still planned.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y + return a, torch.sin(a) + + et = ( + to_edge(export(Model(), (torch.randn(2, 2), torch.randn(2, 2)))) + .to_backend(PerAddDeviceDelegatePartitioner(["cuda:0"])) + .to_executorch( + ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False), + ) + ) + ) + gm = et.exported_program().graph_module + output = gm.graph.output_node().args[0][0] + (d2h,) = _collect_device_copy_nodes(gm).d2h_nodes + self.assertIs(d2h.args[0], output) + self.assertIsNone(output.meta["spec"].mem_id) + self.assertIsNotNone(d2h.meta["spec"].mem_id) + + def test_skip_d2h_returns_output_written_on_host_from_host(self): + """With skip_d2h_for_method_outputs, a returned delegate output whose + host copy a host op writes in place, directly or through a view, is + returned as that written host copy, because the device tensor never + sees the write. No copy goes back to the device. A warning names the + op that writes it, since the copy to the host runs on every call.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + a = x + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(a.view(-1), y.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(a, y) + return a + + for write_view in (False, True): + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0"], + _host_writer_config( + write_view, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + ), + ) + with self.subTest(write_view=write_view): + (warning,) = logs.output + self.assertIn("Method output 0 (16 bytes)", warning) + self.assertIn( + "'propagate_device_test.accumulate.default' writes it", warning + ) + (output,) = gm.graph.output_node().args[0] + self.assertIs(output, _host_written(gm)) + self.assertEqual(output.target, torch.ops.et_copy._d2h_copy.out) + (output_spec,) = gm.graph.output_node().meta["spec"] + self.assertEqual(output_spec.device, DeviceType.CPU) + self.assertFalse( + any( + copy.args[0] is output + for copy in _collect_device_copy_nodes(gm).h2d_nodes + ) + ) + + def test_skip_d2h_copies_a_written_output_only_for_a_later_delegate(self): + """With skip_d2h_for_method_outputs, a returned output that a host op + writes before a later delegate on the same device reads it is copied + back to the device for that delegate only. The method returns the + written host copy.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y + torch.ops.propagate_device_test.accumulate(a, y) + return a, a + y + + gm = _lower_with_host_writer( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config( + False, + propagate_device_config=PropagateDeviceConfig( + skip_d2h_for_method_outputs=True + ), + ), + ) + (copy,) = [ + n + for n in _collect_device_copy_nodes(gm).h2d_nodes + if n.args[0] is _host_written(gm) + ] + _, second = _collect_device_copy_nodes(gm).delegate_nodes + self.assertIn(copy, second.args[1:]) + self.assertIs(gm.graph.output_node().args[0][0], _host_written(gm)) + def test_cross_device_copies_preserved(self): """A genuine cross-device hand-off (cuda:0 -> cuda:1) between two directly-connected delegates must be preserved -- the round-trip is a @@ -563,14 +1137,14 @@ def forward(self, a, b): 4, f"[{pipeline}] expected four delegates", ) - # a -> {D0, D1} and b -> {D0, D2} give 4 graph-input H2D copies; - # the single graph output gives 1 D2H. All inter-delegate - # round-trips (including delegate 0's output feeding D1 and D2) - # are elided. + # a -> {D0, D1} and b -> {D0, D2} give one graph-input H2D copy + # each, shared by the delegates that read it; the single graph + # output gives 1 D2H. All inter-delegate round-trips (including + # delegate 0's output feeding D1 and D2) are elided. self.assertEqual( len(nodes.h2d_nodes), - 4, - f"[{pipeline}] expected 4 boundary H2D copies, got " + 2, + f"[{pipeline}] expected 2 boundary H2D copies, got " f"{len(nodes.h2d_nodes)}", ) self.assertEqual( @@ -722,6 +1296,47 @@ def forward(self, a, b): f"first delegate's output directly", ) + def test_tensors_shared_by_delegates_on_different_devices(self): + """Delegates on cuda:0 and cuda:1 that read the same method input and + each other's outputs. The input gets one copy for each device, read by + every delegate there, and each output crosses to the other device + through the host.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y # delegate 0 (cuda:0) + b = x + a # delegate 1 (cuda:1) + return a, b, x + b # delegate 2 (cuda:0) + + d2h_out = torch.ops.et_copy._d2h_copy.out + for pipeline, gm in _lower_model_to_executorch( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1", "cuda:0"]), + ExecutorchBackendConfig(enable_non_cpu_memory_planning=True), + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + first, second, third = nodes.delegate_nodes + x_copies = [n for n in nodes.h2d_nodes if n.args[0].name == "x"] + first_x, second_x, third_x = ( + next(arg for arg in delegate.args[1:] if arg in x_copies) + for delegate in (first, second, third) + ) + self.assertEqual(len(x_copies), 2) + self.assertIs(first_x, third_x) + self.assertEqual(first_x.meta["spec"].device_index, 0) + self.assertEqual(second_x.meta["spec"].device_index, 1) + crossings = sorted( + ( + h2d.args[0].args[0].meta["spec"].device_index, + h2d.meta["spec"].device_index, + ) + for h2d in nodes.h2d_nodes + if h2d.args[0].target == d2h_out + ) + self.assertEqual(crossings, [(0, 1), (1, 0)]) + # ---- Integration tests: device consistency after to_executorch ---- def test_device_consistency_cuda_1(self): @@ -1265,6 +1880,546 @@ def forward(self, a, b): "CPU placeholder", ) + def test_constant_feeding_device_delegate_is_copied(self): + """A constant feeding a device delegate gets an _h2d_copy into planned device + memory, with or without skip_h2d_for_method_inputs: its data lives in the + program on the host, so it is not a method input the caller can place.""" + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + + def forward(self, a): + return torch.add(a, self.offset) + + model = Model() + inputs = (torch.randn(2, 2),) + for skip in (False, True): + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=skip + ), + enable_non_cpu_memory_planning=True, + ) + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + results = self._get_executorch_program( + model, inputs, DeviceAwarePartitioner("cuda:0"), et_config + ) + self.assertIn("Constant 'b_offset' (16 bytes)", logs.output[0]) + for pipeline, program, gm in results: + with self.subTest(pipeline=pipeline, skip=skip): + nodes = _collect_device_copy_nodes(gm) + # Input a is copied only when the caller cannot place it. + self.assertEqual(len(nodes.h2d_nodes), 1 if skip else 2) + cuda_ph, _ = self._collect_placeholders_by_device(gm) + self.assertEqual(len(cuda_ph), int(skip)) + self.assertTrue(all(not n.meta["spec"].const for n in cuda_ph)) + copies = [n.meta["spec"] for n in nodes.h2d_nodes] + self.assertTrue(all(not spec.const for spec in copies)) + self._assert_nodes_buffer_device( + nodes.h2d_nodes, program, DeviceType.CUDA, pipeline, "H2D copy" + ) + + def test_error_on_per_call_copies(self): + """error_on_per_call_copies turns the warning about a copy that a skip + flag cannot remove into an error: a constant copied to the device and + an input copied for a host reader under skip_h2d_for_method_inputs, and + a returned output written on the host under skip_d2h_for_method_outputs. + Without the skip flag, a constant copy only logs the warning.""" + + class ReadsConstant(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + + def forward(self, a): + return torch.add(a, self.offset), self.offset * 2 + + class WritesOutputOnHost(torch.nn.Module): + def forward(self, x, y): + a = x + y + torch.ops.propagate_device_test.accumulate(a, y) + return a + + class ReadsInputOnHost(torch.nn.Module): + def forward(self, x, y): + return x + y, torch.sin(x) + + def lower(model, num_inputs, **flags): + return _lower_with_host_writer( + model, + tuple(torch.randn(2, 2) for _ in range(num_inputs)), + ["cuda:0"], + _host_writer_config( + False, + propagate_device_config=PropagateDeviceConfig( + error_on_per_call_copies=True, **flags + ), + ), + ) + + cases = { + "constant": (ReadsConstant(), 1, "skip_h2d_for_method_inputs"), + "input read on host": (ReadsInputOnHost(), 2, "skip_h2d_for_method_inputs"), + "output written on host": ( + WritesOutputOnHost(), + 2, + "skip_d2h_for_method_outputs", + ), + } + for case, (model, num_inputs, flag) in cases.items(): + with self.subTest(case=case): + with self.assertRaisesRegex( + RuntimeError, + "on every call.*Raised because error_on_per_call_copies=True", + ): + lower(model, num_inputs, **{flag: True}) + with self.assertLogs(propagate_device_pass.logger, "WARNING"): + lower(ReadsConstant(), 1) + + def test_constant_copy_keeps_stored_layout(self): + """The copy of a constant moves its bytes as they are, so it must use the + layout the constant is stored in: contiguous for a transposed parameter, + which to_edge makes contiguous, and channels-last for a channels-last + one.""" + + class Model(torch.nn.Module): + def __init__(self, weight): + super().__init__() + self.weight = torch.nn.Parameter(weight) + + def forward(self, a): + return torch.add(a, self.weight) + + weights = { + "transposed": torch.randn(3, 2).t(), + "channels_last": torch.randn(1, 2, 3, 3).to( + memory_format=torch.channels_last + ), + } + for name, weight in weights.items(): + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + for pipeline, _program, gm in self._get_executorch_program( + Model(weight), + (torch.randn(weight.shape),), + DeviceAwarePartitioner("cuda:0"), + et_config, + ): + with self.subTest(pipeline=pipeline, weight=name): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.h2d_nodes), 1) + h2d = nodes.h2d_nodes[0] + # The emitter rewrites the constant's spec to the stored layout. + stored = h2d.args[0].meta["spec"] + self.assertEqual(h2d.meta["spec"].dim_order, stored.dim_order) + self.assertEqual( + tuple(h2d.meta["spec"].stride), tuple(stored.stride) + ) + + def test_constant_with_unstored_layout_is_refused(self): + """A constant that is neither contiguous nor channels-last is stored + contiguous, so its copy would not match the stored data. A pass that runs + after to_edge normalized the constants can create one, and lowering then + fails instead of producing wrong results.""" + + class StoreWeightTransposed(ExportedProgramPassBase): + def call(self, ep): + ep.state_dict["w"] = torch.nn.Parameter( + ep.state_dict["w"].t().contiguous().t() + ) + for node in ep.graph.nodes: + if node.name == "p_w": + node.meta["val"] = node.meta["val"].t().contiguous().t() + return ExportedProgramPassResult(ep, True) + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(2, 3)) + + def forward(self, a): + return torch.add(a, self.w) + + for skip in (False, True): + edge = to_edge_transform_and_lower( + export(Model(), (torch.randn(2, 3),)), + transform_passes=[StoreWeightTransposed()], + partitioner=[DeviceAwarePartitioner("cuda:0")], + ) + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=skip + ), + enable_non_cpu_memory_planning=True, + ) + with self.subTest(skip=skip): + with self.assertRaisesRegex( + RuntimeError, + "Constant 'p_w' feeds a delegate on cuda:0, but it is neither " + "contiguous nor channels-last", + ): + edge.to_executorch(et_config) + + def test_constant_read_by_several_delegates_is_copied_once_per_device(self): + """A constant read by several delegates gets one _h2d_copy for each + device, not one for each delegate.""" + + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + lowered = { + device: _lower_add_to_device(inputs, device) + for device in ("cuda:0", "cuda:1") + } + + class Model(torch.nn.Module): + def __init__(self, devices): + super().__init__() + self.register_buffer("offset", torch.randn(2, 2)) + self.first, self.second, self.third = (lowered[d] for d in devices) + + def forward(self, x): + return ( + self.first(x, self.offset), + self.second(x, self.offset), + self.third(x, self.offset), + ) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, enable_non_cpu_memory_planning=True + ) + cases = { + "one device": (("cuda:0", "cuda:0", "cuda:0"), [0]), + "two devices": (("cuda:0", "cuda:1", "cuda:0"), [0, 1]), + } + for case, (devices, copy_indices) in cases.items(): + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + et = to_edge(export(Model(devices), inputs[:1])).to_executorch( + deepcopy(et_config) + ) + gm = et.exported_program().graph_module + with self.subTest(case=case): + self.assertEqual(len(logs.output), len(copy_indices)) + nodes = _collect_device_copy_nodes(gm) + offset_copies = [ + n for n in nodes.h2d_nodes if n.args[0].name == "b_offset" + ] + self.assertEqual( + sorted(n.meta["spec"].device_index for n in offset_copies), + copy_indices, + ) + for delegate, device in zip(nodes.delegate_nodes, devices): + offset = delegate.args[2] + self.assertIn(offset, offset_copies) + self.assertEqual(f"cuda:{offset.meta['spec'].device_index}", device) + + def test_tensor_written_on_host_is_copied_for_each_delegate(self): + """A tensor that a host op writes in place, directly or through a view, + between two delegates on one device gets a new copy for the second + delegate, so it reads the new value.""" + + class Model(torch.nn.Module): + def __init__(self, write_view): + super().__init__() + self.write_view = write_view + + def forward(self, x, y): + t = x * 2 + a = t + y + if self.write_view: + torch.ops.propagate_device_test.accumulate(t.view(-1), a.view(-1)) + else: + torch.ops.propagate_device_test.accumulate(t, a) + return a, t + y + + for write_view in (False, True): + gm = _lower_with_host_writer( + Model(write_view), + (torch.randn(2, 2), torch.randn(2, 2)), + ["cuda:0", "cuda:0"], + _host_writer_config(write_view), + ) + with self.subTest(write_view=write_view): + writer = _host_writer(gm) + first, second = _collect_device_copy_nodes(gm).delegate_nodes + (second_copy,) = [ + arg + for arg in second.args[1:] + if arg.target == torch.ops.et_copy._h2d_copy.out + and arg.args[0] is _host_written(gm) + ] + nodes = list(gm.graph.nodes) + self.assertLess(nodes.index(first), nodes.index(writer)) + self.assertLess(nodes.index(writer), nodes.index(second_copy)) + + def test_skip_h2d_for_input_shared_by_device_delegates(self): + """With skip_h2d_for_method_inputs=True, a method input read by two + delegates on the same device is placed on that device once, with no + _h2d_copy.""" + + class Model(torch.nn.Module): + def forward(self, a, b): + c = torch.sin(torch.add(a, b)) + return torch.add(a, c) + + model = Model() + inputs = (torch.randn(2, 2), torch.randn(2, 2)) + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + for pipeline, program, gm in self._get_executorch_program( + model, inputs, DeviceAwarePartitioner("cuda:0"), et_config + ): + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(len(nodes.delegate_nodes), 2) + # Only sin's output, computed on the host, is copied. + self.assertEqual(len(nodes.h2d_nodes), 1) + cuda_ph, _ = self._collect_placeholders_by_device(gm) + self.assertEqual(sorted(n.name for n in cuda_ph), ["a", "b"]) + self._assert_nodes_buffer_device( + cuda_ph, program, DeviceType.CUDA, pipeline, "Placeholder" + ) + + def test_skip_h2d_copies_input_for_readers_off_device(self): + """With skip_h2d_for_method_inputs=True, a method input that a device + delegate reads stays on that device. A host op, a delegate on the host, + or a delegate on another device that also reads it reads one _d2h_copy + of it, also when more than one delegate on the device reads it, and a + warning says so.""" + + class HostOp(torch.nn.Module): + def forward(self, a, b): + return torch.add(a, b), torch.sin(a) + + class TwoDelegates(torch.nn.Module): + def forward(self, a, b, c): + return torch.add(a, b), torch.add(a, c) + + class TwoDelegatesAndHostOp(torch.nn.Module): + def forward(self, a, b, c): + return torch.add(a, b), torch.add(a, c), torch.sin(a) + + class CopiedIntoBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a, b): + self.state.copy_(a) + return torch.add(a, b) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + cases = { + "host op": (HostOp(), 2, DeviceAwarePartitioner("cuda:0")), + "two delegates on the device and a host op": ( + TwoDelegatesAndHostOp(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:0"]), + ), + "delegate without a device": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", None]), + ), + "delegate on the cpu": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cpu"]), + ), + "other device index": ( + TwoDelegates(), + 3, + PerAddDeviceDelegatePartitioner(["cuda:0", "cuda:1"]), + ), + # The write-back of the buffer only reads the input. + "buffer write-back": ( + CopiedIntoBuffer(), + 2, + DeviceAwarePartitioner("cuda:0"), + ), + } + for case, (model, num_inputs, partitioner) in cases.items(): + inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs)) + with self.assertLogs(propagate_device_pass.logger, "WARNING") as logs: + results = self._get_executorch_program( + model, inputs, partitioner, et_config + ) + self.assertIn("Method input 'a' (16 bytes) stays on cuda:0", logs.output[0]) + for pipeline, program, gm in results: + with self.subTest(case=case, pipeline=pipeline): + a = next(n for n in gm.graph.nodes if n.name == "a") + self.assertEqual(a.meta["spec"].device, DeviceType.CUDA) + self.assertEqual(a.meta["spec"].device_index, 0) + nodes = _collect_device_copy_nodes(gm) + self.assertIs(nodes.delegate_nodes[0].args[1], a) + # The copy is the only other reader of the input, and is + # not copied again. + (copy,) = [n for n in a.users if n not in nodes.delegate_nodes] + self.assertIn(copy, nodes.d2h_nodes) + self.assertFalse(set(copy.users) & set(nodes.d2h_nodes)) + self._assert_nodes_buffer_device( + [copy], program, DeviceType.CPU, pipeline, "D2H copy" + ) + + def test_skip_h2d_shape_reader_needs_no_copy(self): + """With skip_h2d_for_method_inputs=True, an op that reads only the shape + of an input on the device reads the input there, since its sizes stay + in host memory, so the input gets no copy and no warning.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y), torch.zeros(x.shape[0]) + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + batch = torch.export.Dim("batch", min=2, max=8) + ep = export( + Model(), + (torch.randn(4, 2), torch.randn(4, 2)), + dynamic_shapes={"x": {0: batch}, "y": {0: batch}}, + ) + edge = to_edge( + ep, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).to_backend(DeviceAwarePartitioner("cuda:0")) + with self.assertNoLogs(propagate_device_pass.logger, "WARNING"): + gm = edge.to_executorch(et_config).exported_program().graph_module + (size,) = [n for n in gm.graph.nodes if n.target == torch.ops.aten.sym_size.int] + read = size.args[0] + self.assertEqual(read.op, "placeholder") + self.assertEqual(read.meta["spec"].device, DeviceType.CUDA) + self.assertFalse( + set(read.users) & set(_collect_device_copy_nodes(gm).d2h_nodes) + ) + + def test_skip_both_returns_input_on_device_without_copy(self): + """With both skip flags, a method that returns an input a device + delegate also reads returns it on the device, with no copy and no + warning.""" + + class Model(torch.nn.Module): + def forward(self, x, y): + return torch.add(x, y), x + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True + ), + enable_non_cpu_memory_planning=True, + ) + with self.assertNoLogs(propagate_device_pass.logger, "WARNING"): + results = self._get_executorch_program( + Model(), + (torch.randn(2, 2), torch.randn(2, 2)), + DeviceAwarePartitioner("cuda:0"), + et_config, + ) + for pipeline, _program, gm in results: + with self.subTest(pipeline=pipeline): + nodes = _collect_device_copy_nodes(gm) + self.assertEqual(nodes.h2d_nodes + nodes.d2h_nodes, []) + x = next(n for n in gm.graph.nodes if n.name == "x") + self.assertIn(x, gm.graph.output_node().args[0]) + self.assertEqual(x.meta["spec"].device, DeviceType.CUDA) + + def test_skip_h2d_refuses_placeholder_written_on_host(self): + """A mutated input or mutable buffer is written back on the host, so the + refusal says so, also when a host op reads it too. An input that a host + op writes in place is refused as well, since the write would only reach + its host copy.""" + + class MutableBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a): + out = torch.add(a, self.state) + self.state.add_(a) + return out + + class MutableBufferReadOnHost(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(2, 2)) + + def forward(self, a): + out = torch.add(a, self.state) + read = torch.sin(self.state) + self.state.add_(a) + return out, read + + class MutatedInput(torch.nn.Module): + def forward(self, a, x): + out = torch.add(a, x) + x.add_(a) + return out + + class InputWrittenByHostOp(torch.nn.Module): + def forward(self, a, x): + out = torch.add(a, x) + torch.ops.propagate_device_test.accumulate(x, a) + return out + + et_config = ExecutorchBackendConfig( + emit_stacktrace=False, + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ), + enable_non_cpu_memory_planning=True, + ) + write_back = r"'copy_' writes it back on the host \(the method mutates it\)" + cases = { + "mutable buffer": (MutableBuffer(), 1, "b_state", write_back), + "mutable buffer read on host": ( + MutableBufferReadOnHost(), + 1, + "b_state", + write_back, + ), + "mutated input": (MutatedInput(), 2, "x", write_back), + "input written by a host op": ( + InputWrittenByHostOp(), + 2, + "x", + "'[a-z_]*accumulate[a-z_]*' writes it in place on the host", + ), + } + for case, (model, num_inputs, name, reason) in cases.items(): + inputs = tuple(torch.randn(2, 2) for _ in range(num_inputs)) + edge = to_edge( + export(model, inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).to_backend(DeviceAwarePartitioner("cuda:0")) + with self.subTest(case=case): + with self.assertRaisesRegex( + RuntimeError, f"placeholder '{name}' on cuda:0, because {reason}" + ): + edge.to_executorch(deepcopy(et_config)) + def test_tensorspec_repr_includes_device(self): spec = TensorSpec(dtype=torch.float32, shape=torch.Size([2, 3])) repr_str = repr(spec) diff --git a/extension/llm/custom_ops/op_moe_aot.cpp b/extension/llm/custom_ops/op_moe_aot.cpp index b22c52e0107..02c635387d2 100644 --- a/extension/llm/custom_ops/op_moe_aot.cpp +++ b/extension/llm/custom_ops/op_moe_aot.cpp @@ -11,6 +11,7 @@ #include #include +#include namespace torch { namespace executor { @@ -62,7 +63,7 @@ at::Tensor quantized_moe_ffn_aten( const int64_t dim, const int64_t group_size, const int64_t weight_nbit, - c10::string_view score_func, + std::string_view score_func, const double route_scale) { auto output = at::empty({x.size(0), dim}, x.options().dtype(at::kFloat)); WRAP_TO_ATEN(quantized_moe_ffn_out_no_context, 13) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 25fb2abe576..39ae5082bef 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -3587,6 +3587,118 @@ struct PyMethod final { dlpack_tensor_ptrs_ = std::move(dlpack_tensor_ptrs); torch_inputs_ = std::move(torch_inputs); torch_tensor_ptrs_ = std::move(input_tensors); +#endif + find_inputs_aliased_by_outputs(); + } + + // Writes output `index` straight into `tensor`'s storage from the next + // execute() on, instead of into storage this method allocates. The binding + // persists across calls and the tensor is kept alive until it is replaced. + // Storage the output used before is kept alive, so tensors already returned + // for it with clone_outputs=False stay valid. + void set_output(py::object tensor, size_t index) { + ObjectLock lock(mutex_); +#ifdef USE_ATEN_LIB + // ATen swaps the data pointer of the output's storage in place, so an + // output returned earlier would read the bound tensor, or freed memory once + // it is rebound. + throw std::runtime_error( + "set_output is not supported by the ATen bindings"); +#else + if (!is_torch_tensor(tensor)) { + throw py::type_error("Output storage must be a torch tensor"); + } + const auto num_outputs = storage_->method->outputs_size(); + if (index >= num_outputs) { + throw std::out_of_range( + "output index " + std::to_string(index) + " is out of range for " + + std::to_string(num_outputs) + " outputs"); + } + const auto& output = storage_->method->get_output(index); + if (!output.isTensor()) { + throw py::value_error( + "output " + std::to_string(index) + " is not a tensor"); + } + if (const auto input = aliased_input(index)) { + throw py::value_error(aliased_input_message(index, *input)); + } + for (size_t i = 0; i < num_outputs; ++i) { + if (i != index && &storage_->method->get_output(i) == &output) { + throw py::value_error( + "output " + std::to_string(index) + + " is the same value as output " + std::to_string(i) + + " and cannot be bound"); + } + } + const auto meta = storage_->method->method_meta().output_tensor_meta(index); + THROW_IF_ERROR( + meta.error(), "Failed to get output tensor meta for output %zu", index); + if (meta->is_memory_planned()) { + throw py::value_error( + "output " + std::to_string(index) + + " is memory planned and cannot be bound"); + } + if (!tensor.attr("is_contiguous")().cast()) { + throw py::value_error( + "output " + std::to_string(index) + + " must be bound to a contiguous tensor"); + } + const TorchTensorView view(tensor, index); + const auto device = output.toTensor().device(); + if (view.scalar_type() != meta->scalar_type() || view.device() != device) { + const auto name = [](executorch::aten::Device value) { + return value.is_cpu() ? std::string("cpu") + : "cuda:" + std::to_string(value.index()); + }; + throw py::value_error( + "output " + std::to_string(index) + " is " + + runtime::toString(meta->scalar_type()) + " on " + name(device) + + ", but the bound tensor is " + runtime::toString(view.scalar_type()) + + " on " + name(view.device())); + } + const auto dim_order = meta->dim_order(); + for (size_t d = 0; d < dim_order.size(); ++d) { + if (dim_order[d] != d) { + throw py::value_error( + "output " + std::to_string(index) + + " is not contiguous (e.g. channels-last) and cannot be bound"); + } + } + if (view.nbytes() < meta->nbytes()) { + throw py::value_error( + "output " + std::to_string(index) + " needs " + + std::to_string(meta->nbytes()) + " bytes, but the tensor has " + + std::to_string(view.nbytes())); + } + // Two outputs written into one buffer, or an output written over an input + // a kernel is still reading, would silently corrupt each other. + const auto overlaps = [&](const void* data, size_t nbytes) { + const auto start = reinterpret_cast(view.data()); + const auto other = reinterpret_cast(data); + return nbytes != 0 && view.nbytes() != 0 && start < other + nbytes && + other < start + view.nbytes(); + }; + for (size_t i = 0; i < bound_outputs_.size(); ++i) { + if (i == index || !bound_outputs_[i]) { + continue; + } + const TorchTensorView bound(bound_outputs_[i], i); + if (overlaps(bound.data(), bound.nbytes())) { + throw py::value_error( + "output " + std::to_string(index) + + " would share memory with output " + std::to_string(i) + + ", which is already bound"); + } + } + for (const auto& input : torch_inputs_) { + if (overlaps(input->data(), input->nbytes())) { + throw py::value_error( + "output " + std::to_string(index) + + " would share memory with input " + std::to_string(input->index())); + } + } + bound_outputs_.resize(num_outputs); + bound_outputs_[index] = tensor; #endif } @@ -3634,7 +3746,26 @@ struct PyMethod final { const auto num_outputs = storage_->method->outputs_size(); allocate_output_storages(); std::vector> output_storage_spans(num_outputs); - for (int i = 0; i < storage_->output_storages.size(); ++i) { + caller_outputs_.assign(num_outputs, py::object()); + for (size_t i = 0; i < storage_->output_storages.size(); ++i) { +#ifndef USE_ATEN_LIB + // set_output() refuses to bind anything in ATen builds. + if (is_bound_output(i)) { + auto bound = bound_outputs_[i]; + set_output(bound, i); + const TorchTensorView view(bound, i); + output_storage_spans[i] = + Span(static_cast(view.data()), view.nbytes()); + caller_outputs_[i] = bound; + continue; + } +#endif + // An output that is an unplanned input (e.g. the write-back of a mutated + // input) already points at the tensor fed for it. Giving it storage + // would repoint the input too. + if (aliased_input(i)) { + continue; + } const auto& device_buffer = storage_->device_output_storages[i]; output_storage_spans[i] = device_buffer ? device_buffer->as_span() @@ -3688,8 +3819,39 @@ struct PyMethod final { storage_->method->method_meta().name(), static_cast(get_outputs_status)); + // An output bound with set_output() already lives in the caller's tensor, + // so it is returned as that tensor and taken out of `result` before it + // could be cloned. A dynamic bound output may have been resized smaller + // than its (contiguous) tensor, so view it at its written size. An output + // that is an input is returned like any other: cloned by default, and as + // an alias of the caller's tensor with clone_outputs=False. + std::vector caller_backed(result.size()); + for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { + if (!caller_outputs_[i] || !result[i].isTensor()) { + continue; + } + const auto& caller = caller_outputs_[i]; + const auto& written = result[i].toTensor(); + std::vector sizes( + written.sizes().begin(), written.sizes().end()); + if (caller.attr("shape").cast>() == sizes) { + caller_backed[i] = caller; + } else { + caller_backed[i] = caller.attr("view")(-1) + .attr("narrow")(0, 0, written.numel()) + .attr("view")(py::cast(sizes)); + } + result[i] = EValue(); + } + // Retrieve outputs - return get_outputs_as_py_list(result, clone_outputs); + py::list outputs = get_outputs_as_py_list(result, clone_outputs); + for (size_t i = 0; i < caller_backed.size(); ++i) { + if (caller_backed[i]) { + outputs[i] = caller_backed[i]; + } + } + return outputs; } py::list call(const py::object& inputs, bool clone_outputs = true) { @@ -3778,6 +3940,70 @@ struct PyMethod final { std::vector> buffer_inputs_; std::vector buffer_tensor_ptrs_; std::vector borrowed_inputs_; + // Caller tensors bound as output storage with set_output(), by output index. + std::vector bound_outputs_; + // Per output of the last execute(), the caller tensor it lives in, if any. + std::vector caller_outputs_; + // Per output, the unplanned input it is, if any. Empty until the first + // set_inputs(). + std::vector> inputs_aliased_by_outputs_; + + bool is_bound_output(size_t index) const { + return index < bound_outputs_.size() && bool(bound_outputs_[index]); + } + + std::optional aliased_input(size_t index) const { + return index < inputs_aliased_by_outputs_.size() + ? inputs_aliased_by_outputs_[index] + : std::nullopt; + } + + // An output that carries an input mutation is the same value as that input. + // get_input() marks the input as set, so this only runs once set_inputs() + // has set them all, leaving execute()'s check for unset inputs intact. + void find_inputs_aliased_by_outputs() { + const auto num_outputs = storage_->method->outputs_size(); + if (inputs_aliased_by_outputs_.size() == num_outputs) { + return; + } + inputs_aliased_by_outputs_.assign(num_outputs, std::nullopt); + const auto meta = storage_->method->method_meta(); + for (size_t i = 0; i < num_outputs; ++i) { + const auto output_meta = meta.output_tensor_meta(i); + if (!output_meta.ok() || output_meta->is_memory_planned()) { + continue; + } + for (size_t j = 0; j < storage_->method->inputs_size(); ++j) { +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + const bool same = + &storage_->method->get_output(i) == &storage_->method->get_input(j); +#pragma GCC diagnostic pop + if (same) { + inputs_aliased_by_outputs_[i] = j; + break; + } + } + } + // Only possible for a binding made before the first set_inputs(), when + // aliases were not known yet. + std::string error; + for (size_t i = 0; i < num_outputs; ++i) { + if (const auto input = aliased_input(i); input && is_bound_output(i)) { + bound_outputs_[i] = py::object(); + error += + aliased_input_message(i, *input) + ", so its binding was removed. "; + } + } + if (!error.empty()) { + throw py::value_error(error); + } + } + + static std::string aliased_input_message(size_t output, size_t input) { + return "output " + std::to_string(output) + " is input " + + std::to_string(input) + ", which the method already writes in place"; + } #ifndef USE_ATEN_LIB std::vector> dlpack_inputs_; std::vector dlpack_tensor_ptrs_; @@ -3997,17 +4223,10 @@ struct PyMethod final { output_tensor_meta.error(), "Failed to get output tensor meta for output %zu", i); - if (output_tensor_meta.get().is_memory_planned()) { - // Skip allocating storage for planned memory outputs. - output_storages.emplace_back(); - continue; - } - const auto& output = storage_->method->get_output(i); - bool is_input = false; - for (size_t j = 0; j < storage_->method->inputs_size(); ++j) { - is_input |= &output == &storage_->method->get_input(j); - } - if (is_input) { + if (output_tensor_meta.get().is_memory_planned() || is_bound_output(i) || + aliased_input(i).has_value()) { + // Skip allocating storage for planned memory outputs, and for outputs + // that live in a caller's tensor. output_storages.emplace_back(); continue; } @@ -4776,6 +4995,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { py::class_(m, "ExecuTorchMethod", py::module_local()) .def("set_inputs", &PyMethod::set_inputs, py::arg("inputs"), call_guard) .def("execute", &PyMethod::execute, call_guard) + .def( + "set_output", + &PyMethod::set_output, + py::arg("tensor"), + py::arg("index"), + call_guard) .def( "get_outputs", &PyMethod::get_outputs, diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 3810cc45b8f..0cc6bd784e6 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -10,6 +10,8 @@ from __future__ import annotations from enum import Enum from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +import torch + _T = TypeVar("_T") def experimental(message: str) -> Callable[[_T], _T]: ... @@ -135,6 +137,19 @@ class ExecuTorchProgram: class ExecuTorchMethod: """ExecuTorchMethod is a Python wrapper around a loaded C++ method. + Outputs are cloned when ``clone_outputs`` is True, except for an output + bound with ``set_output``, which is returned as the bound tensor and + overwritten by later calls. + + An input that is not memory planned (e.g. exported with + ``alloc_graph_input=False``) is used in place, so a mutation of it is + written into the tensor passed in for it. When such an input is also an + output (e.g. the write-back of a mutated input, or an input returned + unchanged), that output is a clone like any other by default, and shares + the storage of the tensor passed in with ``clone_outputs=False``. A memory + planned input is copied into the method, so the tensor passed in for it is + not updated. + .. warning:: This API is experimental and subject to change without notice. @@ -143,6 +158,7 @@ class ExecuTorchMethod: # pyre-ignore[2]: "Any" in parameter type annotations. def set_inputs(self, inputs: Any) -> None: ... def execute(self) -> None: ... + def set_output(self, tensor: torch.Tensor, index: int) -> None: ... # pyre-ignore[3]: "Any" in return type annotations. def get_outputs(self, clone_outputs: bool = True) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. diff --git a/pytest-windows.ini b/pytest-windows.ini index 39f238cdc00..26a0ee6dabc 100644 --- a/pytest-windows.ini +++ b/pytest-windows.ini @@ -103,6 +103,8 @@ addopts = --deselect=exir/backend/test/test_compatibility.py::TestCompatibility::test_compatibility_in_runtime --deselect=exir/backend/test/test_compatibility.py::TestCompatibility::test_compatibility_in_runtime_edge_program_manager --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_lowered_backend_module_end_to_end + --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_shared_delegate_matches_separate_calls_at_runtime + --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_runtime_rejects_a_shared_delegate_with_scratch --deselect=exir/backend/test/test_to_backend_multi_method.py::TestToBackendMultiMethod::test_multi_method_end_to_end --deselect=extension/llm/custom_ops/test_sdpa_with_kv_cache.py::SDPATestForSpeculativeDecode::test_sdpa_with_cache_seq_len_130 --deselect=devtools/inspector/tests/inspector_test.py::TestInspector::test_etrecord_populates_correct_edge_dialect_aot_intermediate_outputs diff --git a/runtime/__init__.py b/runtime/__init__.py index a4ebf9a5e72..0aa0d89655c 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -111,7 +111,20 @@ import functools from pathlib import Path from types import ModuleType -from typing import Any, BinaryIO, Dict, List, Optional, Sequence, Set, Union +from typing import ( + Any, + BinaryIO, + Dict, + List, + Optional, + Sequence, + Set, + TYPE_CHECKING, + Union, +) + +if TYPE_CHECKING: + import torch try: from executorch.extension.pybindings.portable_lib import ( # type: ignore[import-not-found] # noqa: F401 @@ -141,6 +154,15 @@ def execute( ) -> Sequence[Any]: """Executes the method with the given inputs. + Outputs are copies, except for an output bound with + :meth:`set_output`, which is returned as the bound tensor and + overwritten by later calls. An input that is not memory planned (e.g. + exported with ``alloc_graph_input=False``) is used in place, so a + mutation of it is written into the tensor passed in for it, while an + output that is that input is still returned as a copy. A memory planned + input is copied into the method, so the tensor passed in for it is not + updated. + Args: inputs: A sequence of input values, typically torch.Tensor objects. Without PyTorch, pass NumPy arrays, other buffers such as @@ -155,6 +177,31 @@ def execute( """ return self._method(inputs, clone_outputs=clone_outputs) + def set_output(self, tensor: "torch.Tensor", index: int) -> None: + """Binds output ``index`` to ``tensor``'s storage for every later call. + + The method then writes that output straight into ``tensor`` and returns + it uncloned (viewed at the written shape when that differs) instead of + a copy of its own storage. Every later call overwrites ``tensor``, so + clone a returned output to keep it. + + The output must not be memory planned (export with + ``alloc_graph_output=False``), must not be returned at another output + index too, and must not be an input, which already lives in the tensor + passed in for it. That last case is only known once inputs have been + set: a binding made before the first call is removed by that call, + which raises. ``tensor`` must have the output's dtype and device, be + contiguous, hold at least as many bytes as the output's largest shape, + and not share memory with another bound output or with an input. The + output itself must be contiguous, not e.g. channels-last. A rejected + binding raises ``ValueError``. The ATen bindings don't support it. + + Args: + tensor: The tensor that stores the output from now on. + index: The index of the output to bind. + """ + self._method.set_output(tensor, index) + @property def metadata(self) -> MethodMeta: """Gets the metadata for the method. diff --git a/runtime/backend/interface.h b/runtime/backend/interface.h index 921d9ed324d..d0e54c3218d 100644 --- a/runtime/backend/interface.h +++ b/runtime/backend/interface.h @@ -92,7 +92,15 @@ class BackendInterface { * * @param[in] handle An opaque handle returned by `init()`. Usually a backend * executable unit. This executable unit should be ready to execute the - * delegate blobs. + * delegate blobs. One handle can serve several calls in a method, each + * with its own `args`. The emitter does this when the backend set + * `PreprocessResult.share_handle_across_calls` and declared no scratch, + * and programs from ExecuTorch 0.7 and earlier can do it for any + * backend. Such programs rely on `execute()` reading `args` again on + * every call, keeping nothing on the handle that is tied to one call or + * that changes later results, and writing outputs only into the output + * tensors in `args`. Setting the flag promises this for every runtime + * version of the backend that may load the program. * @param[in] args The method’s inputs and outputs. * @retval Error::Ok if successful. */ diff --git a/runtime/core/exec_aten/exec_aten.h b/runtime/core/exec_aten/exec_aten.h index ac4fb9a126e..f1d66075fdc 100644 --- a/runtime/core/exec_aten/exec_aten.h +++ b/runtime/core/exec_aten/exec_aten.h @@ -35,7 +35,6 @@ #include // @manual #include // @manual #include // @manual -#include // @manual #include #else // use executor #include // @manual diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index c56a55f35bd..9b8b9e34008 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -116,6 +116,7 @@ class BackendDelegate final { out->backend_ = backend; out->handle_ = nullptr; out->scratch_ = Span>(); + out->has_call_site_ = false; // Pass a pointer to this buffer to the backend. It's safe for the backend // to point its handle to this object, since it will outlive the backend. new (&out->segment_) FreeableBuffer(std::move(processed_data.get())); @@ -154,6 +155,15 @@ class BackendDelegate final { scratch_ = scratch; } + /// Whether an instruction already calls this delegate. + bool has_call_site() const { + return has_call_site_; + } + + void add_call_site() { + has_call_site_ = true; + } + Error Execute( BackendExecutionContext& backend_execution_context, Span args) const { @@ -240,6 +250,7 @@ class BackendDelegate final { const BackendInterface* backend_; DelegateHandle* handle_; Span> scratch_; + bool has_call_site_; }; /** @@ -1143,19 +1154,27 @@ Error Method::init( } chain_instruction_arg_lists[instr_idx] = res.get(); + const auto delegate_index = delegate_call->delegate_index(); + ET_CHECK_OR_RETURN_ERROR( + delegate_index >= 0 && + static_cast(delegate_index) < n_delegate_, + InvalidProgram, + "Delegate index %" PRId32 " >= num delegates %" ET_PRIsize_t, + delegate_index, + n_delegate_); + BackendDelegate& delegate = delegates_[delegate_index]; const auto scratch = delegate_call->scratch(); - if (scratch != nullptr && scratch->size() > 0) { - // The emitter appends one BackendDelegate per delegate call, so - // the index identifies the call site and the buffers can hang - // off the delegate itself. - const auto delegate_index = delegate_call->delegate_index(); - ET_CHECK_OR_RETURN_ERROR( - delegate_index >= 0 && - static_cast(delegate_index) < n_delegate_, - InvalidProgram, - "Delegate index %" PRId32 " >= num delegates %" ET_PRIsize_t, - delegate_index, - n_delegate_); + const bool has_scratch = scratch != nullptr && scratch->size() > 0; + // Scratch is planned for one call site and kept on the delegate, so + // a delegate with scratch must not serve any other call site. + ET_CHECK_OR_RETURN_ERROR( + !delegate.has_call_site() || + (!has_scratch && delegate.scratch().empty()), + InvalidProgram, + "Delegate %" PRId32 " has scratch and more than one call site", + delegate_index); + delegate.add_call_site(); + if (has_scratch) { Result>> buffers = gen_delegate_scratch( method_allocator, memory_manager_->planned_memory(), scratch); if (!buffers.ok()) { @@ -1167,7 +1186,7 @@ Error Method::init( instr_idx); return buffers.error(); } - delegates_[delegate_index].set_scratch(buffers.get()); + delegate.set_scratch(buffers.get()); } } break; case executorch_flatbuffer::InstructionArguments::JumpFalseCall: { diff --git a/runtime/test/targets.bzl b/runtime/test/targets.bzl index 74309a59e01..0ef92d4be23 100644 --- a/runtime/test/targets.bzl +++ b/runtime/test/targets.bzl @@ -11,6 +11,8 @@ def define_common_targets(is_fbcode = False): "//executorch/extension/pybindings/test:make_test", "//executorch/runtime:runtime", "//executorch/devtools/etdump:serialize", + "//executorch/exir:lib", + "//executorch/exir/passes:lib", ], ) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 1b1dcf3cd48..1fadea62692 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -12,7 +12,8 @@ from pathlib import Path import torch - +from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass from executorch.extension.pybindings.test.make_test import ( create_program, ModuleAdd, @@ -21,6 +22,70 @@ from executorch.runtime import Runtime, Verification +class ModuleCacheUpdate(torch.nn.Module): + """Writes one row of a cache passed in as an input, like a static KV cache.""" + + def forward(self, x, cache, pos): + cache.index_copy_(0, pos, x) + return cache.sum(0, keepdim=True) + + +class ModuleRecurrent(torch.nn.Module): + def forward(self, x, state): + return state + x + + +class ModuleInputAndIncrement(torch.nn.Module): + def forward(self, x): + return x, x + 1 + + +class ModuleDouble(torch.nn.Module): + def forward(self, x): + return x * 2 + + +class ModuleDoubleChannelsLast(torch.nn.Module): + def forward(self, x): + return (x * 2).contiguous(memory_format=torch.channels_last) + + +class ModuleTwoOutputs(torch.nn.Module): + def forward(self, x): + return x * 2, x + 1 + + +class ModuleDoubleTwice(torch.nn.Module): + def forward(self, x): + y = x * 2 + return y, y + + +def _cache_inputs(): + return torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) + + +def _load_unplanned_io_method( + module, inputs, dynamic_shapes=None, plan_inputs=False, plan_outputs=False +): + """Loads `forward` exported with its inputs and outputs not memory planned, + unless asked. Returns the program buffer too: the loaded program reads it + in place, so it must outlive the method.""" + config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=plan_inputs, alloc_graph_output=plan_outputs + ) + ) + exported = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + buffer = to_edge(exported).to_executorch(config).buffer + program = Runtime.get().load_program(buffer, verification=Verification.Minimal) + return program.load_method("forward"), buffer + + +def _load_cache_update_method(): + return _load_unplanned_io_method(ModuleCacheUpdate(), _cache_inputs()) + + class RuntimeTest(unittest.TestCase): def test_runtime_import_does_not_load_the_exporter(self): # A fresh interpreter, since this process has already imported exir. @@ -52,6 +117,243 @@ def test_smoke(self): outputs = method.execute(inputs) self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + def test_mutated_input_is_updated_in_place(self): + # An unplanned input the method mutates is written where it lives. Its + # write-back output is a copy by default, and shares the caller's + # storage with clone_outputs=False. + x, cache, _ = _cache_inputs() + method, _buffer = _load_cache_update_method() + expected = torch.zeros(3, 4) + expected[1] = 1.0 + outputs = method.execute((x, cache, torch.tensor([1]))) + self.assertIsNot(outputs[0], cache) + self.assertNotEqual(outputs[0].data_ptr(), cache.data_ptr()) + self.assertTrue(torch.equal(outputs[0], expected)) + self.assertTrue(torch.equal(cache, expected)) + outputs = method._method((x * 2, cache, torch.tensor([2])), clone_outputs=False) + self.assertEqual(outputs[0].data_ptr(), cache.data_ptr()) + expected[2] = 2.0 + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.equal(outputs[1], expected.sum(0, keepdim=True))) + + def test_returned_input_is_cloned_by_default(self): + # An unplanned input returned as an output is an independent copy with + # the default clone_outputs=True, as before set_output existed. + method, _buffer = _load_unplanned_io_method( + ModuleInputAndIncrement(), (torch.tensor([1.0]),) + ) + x = torch.tensor([1.0]) + y, z = method.execute((x,)) + self.assertIsNot(y, x) + self.assertNotEqual(y.data_ptr(), x.data_ptr()) + x.fill_(5.0) + self.assertTrue(torch.equal(y, torch.tensor([1.0]))) + self.assertTrue(torch.equal(z, torch.tensor([2.0]))) + y, z = method._method((x,), clone_outputs=False) + self.assertEqual(y.data_ptr(), x.data_ptr()) + self.assertTrue(torch.equal(z, torch.tensor([6.0]))) + + def test_mutated_planned_input_is_not_updated(self): + x, cache, pos = _cache_inputs() + # The update is planned onto the input's own spec, which is also an + # output, so both must be planned for the cache to be planned. + method, _buffer = _load_unplanned_io_method( + ModuleCacheUpdate(), _cache_inputs(), plan_inputs=True, plan_outputs=True + ) + self.assertTrue(method.metadata.input_tensor_meta(1).is_memory_planned()) + outputs = method.execute((x, cache, pos)) + self.assertIsNot(outputs[0], cache) + self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) + self.assertTrue(torch.equal(outputs[0][1], torch.ones(4))) + + def test_output_fed_back_as_input_is_not_that_input(self): + # An uncloned output fed back in shares the input's memory, but is not + # an input the method mutates. + x, state = torch.ones(2), torch.zeros(2) + method, _buffer = _load_unplanned_io_method(ModuleRecurrent(), (x, state)) + for step in (1, 2): + fed = state + state = method._method((x, fed), clone_outputs=False)[0] + self.assertIsNot(state, fed) + self.assertTrue(torch.equal(state, torch.full((2,), float(step)))) + fresh = torch.zeros(2) + outputs = method._method((x, fresh), clone_outputs=False) + self.assertIsNot(outputs[0], fresh) + self.assertTrue(torch.equal(fresh, torch.zeros(2))) + self.assertTrue(torch.equal(outputs[0], torch.ones(2))) + + def test_set_output_binds_caller_storage(self): + x, cache, _ = _cache_inputs() + method, _buffer = _load_cache_update_method() + # Outputs are [the mutated cache, the sum]: bind the sum to a caller + # tensor, which the method then writes and returns uncloned. + summed = torch.empty(1, 4) + method.set_output(summed, 1) + for row in (1, 2): + outputs = method.execute((x, cache, torch.tensor([row]))) + self.assertIs(outputs[1], summed) + self.assertTrue(torch.equal(summed, cache.sum(0, keepdim=True))) + + def test_set_output_rebinds(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + first, second = torch.zeros(1, 4), torch.zeros(1, 4) + method.set_output(first, 1) + method.set_output(second, 1) + outputs = method.execute((x, cache, pos)) + self.assertIs(outputs[1], second) + self.assertTrue(torch.equal(second, torch.ones(1, 4))) + self.assertTrue(torch.equal(first, torch.zeros(1, 4))) + + def test_failed_set_inputs_then_full_set_aliases_new_input(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method._method.set_inputs((x, cache, pos)) + # Fails on the position, after the new cache was already set: the + # method refuses to run until a full set is installed, and the mutated + # input it then writes is the new one. + new_cache = torch.zeros(3, 4) + with self.assertRaises(RuntimeError): + method._method.set_inputs((x, new_cache, 1)) + with self.assertRaises(RuntimeError): + method._method.execute() + outputs = method._method((x, new_cache, pos), clone_outputs=False) + self.assertEqual(outputs[0].data_ptr(), new_cache.data_ptr()) + self.assertTrue(torch.equal(new_cache[1], torch.ones(4))) + self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) + + def test_failed_input_prevalidation_preserves_previous_inputs(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method._method.set_inputs((x, cache, pos)) + new_cache = torch.zeros_like(cache) + with self.assertRaises(ValueError): + method._method.set_inputs((x, new_cache, torch.tensor([1.0]))) + method._method.execute() + self.assertTrue(torch.equal(cache[1], x[0])) + self.assertTrue(torch.equal(new_cache, torch.zeros_like(cache))) + + def test_aliased_output_rejects_replaced_input_storage(self): + x = torch.ones(2, 2) + method, _buffer = _load_unplanned_io_method(ModuleInputAndIncrement(), (x,)) + method._method.set_inputs((x,)) + method._method.execute() + x.data = torch.zeros_like(x) + with self.assertRaisesRegex(ValueError, "storage changed"): + method._method.get_outputs(clone_outputs=False) + + def test_aliased_output_preserves_written_metadata(self): + for parameter in (False, True): + with self.subTest(parameter=parameter): + x = torch.ones(2, 3) + if parameter: + x = torch.nn.Parameter(x, requires_grad=False) + method, _buffer = _load_unplanned_io_method( + ModuleInputAndIncrement(), (x,) + ) + method._method.set_inputs((x,)) + method._method.execute() + x.transpose_(0, 1) + output = method._method.get_outputs(clone_outputs=False)[0] + self.assertEqual(output.shape, (2, 3)) + self.assertEqual(output.stride(), (3, 1)) + self.assertEqual(output.data_ptr(), x.data_ptr()) + + def test_set_output_keeps_earlier_outputs_alive(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + earlier = method._method((x, cache, pos), clone_outputs=False)[1] + expected = earlier.clone() + method.set_output(torch.empty(1, 4), 1) + method.execute((x * 2, cache, torch.tensor([2]))) + self.assertTrue(torch.equal(earlier, expected)) + + def test_set_output_narrows_dynamic_output(self): + dynamic_shapes = ({0: torch.export.Dim("n", max=8)},) + method, _buffer = _load_unplanned_io_method( + ModuleDouble(), (torch.ones(4, 3),), dynamic_shapes + ) + bound = torch.empty(8, 3) + method.set_output(bound, 0) + x = torch.arange(6.0).reshape(2, 3) + output = method.execute((x,))[0] + self.assertEqual(output.shape, (2, 3)) + self.assertEqual(output.data_ptr(), bound.data_ptr()) + self.assertTrue(torch.equal(output, x * 2)) + + def test_set_output_rejects_bad_bindings(self): + method, _buffer = _load_cache_update_method() + bad_bindings = { + "out of range": (IndexError, torch.empty(1, 4), 2), + "non contiguous": (ValueError, torch.empty(1, 8)[:, ::2], 1), + "too small": (ValueError, torch.empty(1, 2), 1), + "wrong dtype": (ValueError, torch.empty(1, 4, dtype=torch.float64), 1), + } + for name, (error, tensor, index) in bad_bindings.items(): + with self.subTest(name), self.assertRaises(error): + method.set_output(tensor, index) + + def test_set_output_rejects_outputs_it_cannot_own(self): + x = torch.ones(2) + method, _buffer = _load_unplanned_io_method( + ModuleDouble(), (x,), plan_outputs=True + ) + with self.subTest("memory planned"), self.assertRaises(ValueError): + method.set_output(torch.empty(2), 0) + method, _buffer = _load_unplanned_io_method(ModuleDoubleTwice(), (x,)) + with self.subTest("returned twice"), self.assertRaises(ValueError): + method.set_output(torch.empty(2), 0) + image = torch.ones(1, 2, 3, 4) + method, _buffer = _load_unplanned_io_method( + ModuleDoubleChannelsLast(), (image,) + ) + with self.subTest("channels last"), self.assertRaises(ValueError): + method.set_output(torch.empty(1, 2, 3, 4), 0) + + def test_set_output_rejects_mutated_input(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method.execute((x, cache, pos)) + with self.assertRaises(ValueError): + method.set_output(torch.zeros(3, 4), 0) + # Bound before any inputs are set, the first call removes the binding + # and raises, and later calls work. + method, _buffer = _load_cache_update_method() + method.set_output(torch.zeros(3, 4), 0) + fresh = torch.zeros(3, 4) + with self.assertRaises(ValueError): + method.execute((x, fresh, pos)) + outputs = method._method((x, fresh, pos), clone_outputs=False) + self.assertEqual(outputs[0].data_ptr(), fresh.data_ptr()) + self.assertTrue(torch.equal(fresh[1], torch.ones(4))) + + def test_set_output_rejects_shared_memory(self): + x = torch.ones(2) + method, _buffer = _load_unplanned_io_method(ModuleTwoOutputs(), (x,)) + storage = torch.empty(4) + method.set_output(storage[:2], 0) + with self.subTest("another bound output"), self.assertRaises(ValueError): + method.set_output(storage[1:3], 1) + method.set_output(storage[2:], 1) + # Inputs are only known once set, so an output bound over an input is + # refused by the call that sets it. + method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) + method.set_output(x, 0) + with self.subTest("an input"), self.assertRaises(ValueError): + method.execute((x,)) + + def test_execute_without_inputs_raises_after_set_output(self): + # Aliases are found only after inputs are set, since reading an input + # marks it as set; binding first must not hide the missing inputs. + x = torch.ones(2) + for bind in (False, True): + with self.subTest(bind=bind): + method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) + if bind: + method.set_output(torch.empty(2), 0) + with self.assertRaisesRegex(RuntimeError, "no inputs yet"): + method._method.execute() + def test_execute_can_borrow_outputs(self): ep, inputs = create_program(ModuleAdd()) method = Runtime.get().load_program(ep.buffer).load_method("forward") diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index 039c5ce3712..acd332f0bda 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -195,9 +195,14 @@ cmake_dependent_option( "Build the native vision classification runner (run_classification)" OFF "EXECUTORCH_BUILD_NATIVE_VULKAN" OFF ) +cmake_dependent_option( + EXECUTORCH_BUILD_NATIVE_LLM_RUNNER "Build the native LLM runners" OFF + "EXECUTORCH_BUILD_NATIVE_VULKAN OR EXECUTORCH_BUILD_NATIVE_GGML" OFF +) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VULKAN) announce_configured_options(EXECUTORCH_BUILD_NATIVE_GGML) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) +announce_configured_options(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) define_overridable_option( EXECUTORCH_BUILD_WEBGPU "Build the WebGPU backend" BOOL OFF )