From 599dac3aa3dc911c29a9207edd13c16e8a37579d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 18:58:17 -0700 Subject: [PATCH 01/37] Catch missing scatter and index_put kernels when lowering to AOTI backends An AOTI backend (CUDA, Metal) checks at lowering time that every kernel Inductor leaves to ATen exists in the backend's shim library, so a missing one fails the export with a clear list. The check watches two code paths that emit a C shim call, but scatter and index_put fallbacks write their call directly. A program that needs one of those kernels lowers without an error, and the process then dies on the first run with an undefined symbol. For example, Inductor leaves an int64 scatter_add_ to aoti_torch_cuda_scatter_reduce_two_out, which the CUDA shim does not have. Watch those two paths too. They do not pass the kernel name to a method the check can see, so read it from the line they write. This checks the name the generated code really calls. --- backends/aoti/aoti_backend.py | 40 ++++++++++++++++++++++++- backends/cuda/tests/test_cuda_export.py | 18 +++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 14f3b7cbd11..4a2ed0be168 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -7,6 +7,7 @@ import contextlib import hashlib import os +import re import typing from abc import ABC, abstractmethod from enum import Enum @@ -168,7 +169,8 @@ def release_moved_tensors( def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str]): """ Context manager to collect unsupported fallback kernels during compilation. - Monitors both extern kernel calls and runtime lookup. + Monitors extern kernel calls, runtime lookup, and the scatter and + index_put fallbacks, which call their kernels directly. """ supported_kernels = cls.get_supported_fallback_kernels() @@ -178,6 +180,10 @@ def collect_unsupported_fallback_kernels(cls, missing_fallback_kernels: Set[str] original_generate_fallback_kernel_with_runtime_lookup_aot = ( CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot ) + original_generate_scatter_fallback = CppWrapperCpu._generate_scatter_fallback + original_generate_index_put_fallback = ( + CppWrapperCpu._generate_index_put_fallback + ) def generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels( self, kernel: str, *args: Any, **kwargs: Any @@ -200,10 +206,36 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker self, op_overload, *args, **kwargs ) + def collect_kernels_it_writes(original: Any) -> Any: + # These fallbacks write their shim call directly, so read the + # kernel names from the lines they write. + def generate_and_collect_unsupported_kernels( + self, *args: Any, **kwargs: Any + ) -> None: + lines: List[str] = [] + writeline = self.writeline + self.writeline = lambda line: (lines.append(str(line)), writeline(line)) + try: + return original(self, *args, **kwargs) + finally: + self.writeline = writeline + for line in lines: + for kernel in re.findall(r"\b(aoti_torch_\w+)\(", line): + if kernel not in supported_kernels: + missing_fallback_kernels.add(kernel) + + return generate_and_collect_unsupported_kernels + CppWrapperCpu.generate_c_shim_extern_kernel_call = ( generate_c_shim_extern_kernel_call_and_collect_unsupported_kernels ) CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_kernels + CppWrapperCpu._generate_scatter_fallback = collect_kernels_it_writes( + original_generate_scatter_fallback + ) + CppWrapperCpu._generate_index_put_fallback = collect_kernels_it_writes( + original_generate_index_put_fallback + ) try: yield @@ -214,6 +246,12 @@ def generate_fallback_kernel_with_runtime_lookup_aot_and_collect_unsupported_ker CppWrapperCpu.generate_fallback_kernel_with_runtime_lookup_aot = ( original_generate_fallback_kernel_with_runtime_lookup_aot ) + CppWrapperCpu._generate_scatter_fallback = ( + original_generate_scatter_fallback + ) + CppWrapperCpu._generate_index_put_fallback = ( + original_generate_index_put_fallback + ) @classmethod def preprocess( diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8329f42e522..d08c835ca38 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -562,6 +562,24 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) + def test_missing_scatter_and_index_put_kernels_fail_at_export(self): + """Inductor leaves these int64 ops to ATen kernels the CUDA shim lacks.""" + + class Int64Accumulate(torch.nn.Module): + def forward(self, counts: torch.Tensor, idx: torch.Tensor): + ones = torch.ones_like(idx) + return ( + counts.scatter_add(0, idx, ones), + counts.index_put((idx,), ones, accumulate=True), + ) + + inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(Int64Accumulate(), inputs) + message = str(context.exception) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", message) + self.assertIn("aoti_torch_index_put_out", message) + def test_device_info_propagated_to_cuda_delegate_outputs(self): """ Verify that, for a CUDA-delegated graph, every memory-planned tensor's From 7013bcb2e7ad0614c58e2a52857454edef58a7c1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 18:59:37 -0700 Subject: [PATCH 02/37] Keep the OpenVINO blob alive while its compiled model is in use Since OpenVINO 2026.4, ov_core_import_model wraps the blob it is given in an ov::Tensor and the compiled model can keep reading that memory. Earlier versions read it through a stream during the import only. OpenvinoBackend::init frees the blob right after the import. When the program is loaded from a buffer, Free() does nothing, so nothing goes wrong. When it is loaded from a file, the blob is really freed: with a memory-mapped file the first execute crashes with a segmentation fault, even for a two-layer model, and with a file read into memory the model reads freed memory. Free the blob in init only when the OpenVINO runtime is older than 2026.4, as before. With newer versions keep it: ExecuTorch owns it and frees it after destroy(), so it lives exactly as long as the compiled model. The op tests loaded every program from a buffer, so they never saw this. They now load from a file, as most users do. --- backends/openvino/runtime/OpenvinoApi.h | 9 ++++++++ backends/openvino/runtime/OpenvinoBackend.cpp | 22 ++++++++++++++++++- backends/openvino/runtime/OpenvinoBackend.h | 2 ++ .../tests/ops/base_openvino_op_test.py | 20 ++++++++++++----- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/backends/openvino/runtime/OpenvinoApi.h b/backends/openvino/runtime/OpenvinoApi.h index d25e265db07..e4c3e7f737a 100644 --- a/backends/openvino/runtime/OpenvinoApi.h +++ b/backends/openvino/runtime/OpenvinoApi.h @@ -35,6 +35,11 @@ typedef struct { size_t size; } ov_available_devices_t; +typedef struct { + const char* buildNumber; + const char* description; +} ov_version_t; + // Intentionally partial — only OV_STATUS_OK is needed for success checks. // The full enum is defined in openvino/c/ov_common.h. typedef enum { @@ -64,6 +69,8 @@ typedef enum { } ov_element_type_e; // Function pointer types for each OpenVINO C API function we use. +using ov_get_openvino_version_fn = ov_status_e (*)(ov_version_t*); +using ov_version_free_fn = void (*)(ov_version_t*); using ov_core_create_fn = ov_status_e (*)(ov_core_t**); using ov_core_free_fn = void (*)(ov_core_t*); using ov_core_get_available_devices_fn = @@ -102,6 +109,8 @@ struct DlCloser { using DlHandle = std::unique_ptr; struct OpenvinoFunctions { + ov_get_openvino_version_fn get_openvino_version = nullptr; + ov_version_free_fn version_free = nullptr; ov_core_create_fn core_create = nullptr; ov_core_free_fn core_free = nullptr; ov_core_get_available_devices_fn core_get_available_devices = nullptr; diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index 0689e5d4680..fad6ef91cb1 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -6,6 +6,7 @@ */ #include +#include #include #include #include @@ -229,6 +230,8 @@ bool OpenvinoBackend::ensure_loaded() const { return; \ } + LOAD_SYM(get_openvino_version, ov_get_openvino_version) + LOAD_SYM(version_free, ov_version_free) LOAD_SYM(core_create, ov_core_create) LOAD_SYM(core_free, ov_core_free) LOAD_SYM(core_get_available_devices, ov_core_get_available_devices) @@ -255,6 +258,19 @@ bool OpenvinoBackend::ensure_loaded() const { #undef LOAD_SYM + // Before 2026.4, ov_core_import_model reads the blob through a stream, so + // the compiled model holds its own copy. + ov_version_t version = {}; + int major = 0; + int minor = 0; + if (ov_.get_openvino_version(&version) == OV_STATUS_OK) { + if (version.buildNumber && + std::sscanf(version.buildNumber, "%d.%d", &major, &minor) == 2) { + import_copies_blob_ = major < 2026 || (major == 2026 && minor < 4); + } + ov_.version_free(&version); + } + loaded_ = true; ET_LOG( Info, @@ -333,7 +349,11 @@ exr::Result OpenvinoBackend::init( return exr::Error::Internal; } - processed->Free(); + // Since OpenVINO 2026.4 the compiled model can keep reading the blob, so it + // is then kept until ExecuTorch frees it after destroy(). + if (import_copies_blob_) { + processed->Free(); + } ov_infer_request_t* infer_request = nullptr; status = diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index 72b0193dc2d..510e49ee690 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -55,6 +55,8 @@ class OpenvinoBackend final : public ::exr::BackendInterface { mutable OpenvinoFunctions ov_; mutable std::once_flag load_flag_; mutable bool loaded_ = false; + // Whether core_import_model copies the blob, so init can free it. + mutable bool import_copies_blob_ = false; }; } // namespace openvino diff --git a/backends/openvino/tests/ops/base_openvino_op_test.py b/backends/openvino/tests/ops/base_openvino_op_test.py index c429845548a..f6bc928e4d4 100644 --- a/backends/openvino/tests/ops/base_openvino_op_test.py +++ b/backends/openvino/tests/ops/base_openvino_op_test.py @@ -1,3 +1,5 @@ +import os +import tempfile import unittest import executorch @@ -64,12 +66,18 @@ def execute_layer_test( ref_output, ] - # Load model from buffer and execute - runtime = Runtime.get() - program = runtime.load_program(exec_prog.buffer) - method = program.load_method("forward") - assert method is not None - outputs = method.execute(sample_inputs) + # Load from a file: unlike a buffer, it really frees the delegate + # data when the backend frees it. + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp_dir: + pte_path = os.path.join(tmp_dir, "model.pte") + with open(pte_path, "wb") as f: + exec_prog.write_to_file(f) + program = Runtime.get().load_program(pte_path) + method = program.load_method("forward") + assert method is not None + outputs = method.execute(sample_inputs) + # Windows cannot delete a mapped file, so release it first. + del method, program # Compare the outputs with the reference outputs self.assertTrue(len(ref_output) == len(outputs)) From 8a14a650116e483a61a16e7e2d8173a291e05ce9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 20:43:40 -0700 Subject: [PATCH 03/37] Qualcomm AI Engine Direct - Fix DecomposeWrapWithAutocast for real regions torch.export records a `with torch.autocast(...)` block as a wrap_with_autocast call on a submodule. DecomposeWrapWithAutocast inlines that submodule into the graph. It had three bugs: 1. It mapped the submodule's placeholders to the operands by name. When the region reads inputs or parameters, the placeholders are named like arg2_1, arg0_1, arg1_1, not after the operands, so the pass failed with a KeyError such as 'arg2_1'. For example: with torch.autocast("cpu", enabled=False): y = self.linear(x) 2. It only looked at the top graph. A region nested in another region was copied out with its get_attr, whose submodule name only means something inside the outer submodule. In the top module it resolved to the outer body, so the program silently ran the wrong code. 3. It copied the ops of an enabled region as they are. The submodule does not hold the dtype casts that autocast adds when the region runs, so they were lost. A linear on float32 inputs inside `torch.autocast("cpu", dtype=torch.bfloat16)` returned float32 with different values, instead of bfloat16. Map placeholders to operands by position. Trace an enabled region again under its own autocast, which is what torch.export does when it decomposes this operator, so the casts become ops in the graph. Regions nested inside it are traced with their own autocast settings. This retrace also splits composite ops into their parts, for example a linear becomes t and addmm, so the pass now runs before the shared decomposition passes of the export and annotation pipelines. A disabled region changes no dtype, so it is still copied as is, after its nested regions are inlined. --- .../_passes/decompose_wrap_with_autocast.py | 24 ++++++++-- backends/qualcomm/_passes/qnn_pass_manager.py | 5 ++- backends/qualcomm/tests/test_passes.py | 45 +++++++++++++++++++ 3 files changed, 68 insertions(+), 6 deletions(-) diff --git a/backends/qualcomm/_passes/decompose_wrap_with_autocast.py b/backends/qualcomm/_passes/decompose_wrap_with_autocast.py index e1663631f08..380f3e9d91b 100644 --- a/backends/qualcomm/_passes/decompose_wrap_with_autocast.py +++ b/backends/qualcomm/_passes/decompose_wrap_with_autocast.py @@ -10,6 +10,7 @@ import torch from executorch.exir.pass_base import ExportPass, PassResult from executorch.exir.passes import dead_code_elimination_pass +from torch.fx.experimental.proxy_tensor import make_fx from .utils import merge_decomposed_graph @@ -55,13 +56,28 @@ def _replace(self, gm: torch.fx.GraphModule) -> None: if isinstance(node.target, torch._higher_order_ops.wrap.WrapWithAutocast): submod, _ = self._get_submod(gm, node) n_args = node.args + device_type, dtype, enabled, cache_enabled = n_args[:4] + if enabled: + # Retrace the region under its autocast, as torch.export does, + # so the casts autocast applies become ops. Nested regions are + # traced through with their own autocast. + with torch.autocast(device_type, dtype, enabled, cache_enabled): + submod = make_fx(submod, tracing_mode="fake")( + *[arg.meta["val"] for arg in n_args[5:]] + ) + else: + # Inline nested regions first. Their submodule names only make + # sense inside this submodule. + self._replace(submod) input_submod = n_args[4] decomposed_module = submod with graph.inserting_before(node): - # remap is used to map original node values to new node values, - # which ensures that reference to nodes are correctly updated in the new graph - # remap = {"expand_1": node.args[5], "to_4": node.args[6]} - remap = {n_args[i].name: n_args[i] for i in range(5, len(n_args))} + # remap maps each placeholder of the submodule to the operand + # passed in its position. Their names need not match. + placeholders = [ + n for n in submod.graph.nodes if n.op == "placeholder" + ] + remap = {p.name: arg for p, arg in zip(placeholders, n_args[5:])} merge_decomposed_graph( remap=remap, target_node=node, diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 9d6e0f747fd..e8745524512 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -172,6 +172,8 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): """Return annotation pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ RemoveRedundancy, + # Retraced autocast regions may emit ops that the passes below decompose. + DecomposeWrapWithAutocast, RecomposePixelUnshuffle, RecomposeRmsNorm, ReplaceArangeArgs, @@ -194,7 +196,6 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): DecomposeTriu, DecomposeTrunc, DecomposeVar, - DecomposeWrapWithAutocast, DecomposeEinsum, DecomposeEmpty, DecomposeExpM1, @@ -218,6 +219,7 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): def get_export_passes(cls): """Return export pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ + DecomposeWrapWithAutocast, ConvertComplexStft, DecomposeAsStrided, DecomposeBinaryAlpha, @@ -239,7 +241,6 @@ def get_export_passes(cls): # We keep div and floor operations in floating-point to maintain precision. # This pass is needed before to_edge pipeline to avoid mixed type for div operator with RemoveMixedTypeOperators pass. DecomposeFloorDivide, - DecomposeWrapWithAutocast, ConvertSquareToPow, LiftConstantScalarOperands, InsertReshapeForReduceOps, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index d8c5c32c103..9b2b73f69d3 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -7,6 +7,7 @@ ConvertBmmToMatmul, ConvertMhaToSha, DecomposeColIm, + DecomposeWrapWithAutocast, ExpandBroadcastTensorShape, FoldQDQ, InsertIOQDQ, @@ -140,6 +141,50 @@ def forward(self, x): # Must not raise. LiftConstantScalarOperands().call(gm) + def test_decompose_wrap_with_autocast_nested_regions(self): + # A region that reads inputs and parameters names its placeholders + # like arg2_1, arg0_1, arg1_1, not after the operands passed to it. + # Both nested submodules are named submod_1, each in its own module. + # The second linear runs in bfloat16 only because of the outer region. + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(16, 16) + + def forward(self, x): + with torch.autocast("cpu", dtype=torch.bfloat16): + with torch.autocast("cpu", enabled=False): + y = self.linear(x) + z = self.linear(y * 2) + return y + 1, z + + # A disabled region is copied, not retraced, so its nested region + # is inlined on its own. + class DisabledOuterModel(Model): + def forward(self, x): + with torch.autocast("cpu", enabled=False): + y = x * 2 + with torch.autocast("cpu", dtype=torch.bfloat16): + z = self.linear(y) + return y + 1, z + + x = torch.randn(4, 16) + for model in (Model().eval(), DisabledOuterModel().eval()): + with self.subTest(model=type(model).__name__): + gm = torch.export.export(model, (x,)).module() + + DecomposeWrapWithAutocast()(gm) + + self.assertFalse( + any( + isinstance( + n.target, torch._higher_order_ops.wrap.WrapWithAutocast + ) + for n in gm.graph.nodes + ) + ) + torch.testing.assert_close(gm(x), model(x)) + def test_build_op_wrappers_returns_context_binary(self): op_name = "ctx_loader_build" ctx_bin = b"qnn_context_binary" From 87551de2a3ae26c6f4580a1843cf69c0d6435e18 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 21:59:21 -0700 Subject: [PATCH 04/37] Lower detach in AOTI backends Edge turns `x.detach()` into `aten.detach_copy`. Inductor has no lowering or C shim for it, so an AOTI backend (CUDA, Metal) cannot lower any graph that calls detach, for example: return (x * 2).detach() + 1 fails with a "missing fallback kernels" error that lists aten::detach_copy. ReplaceViewCopyWithViewPass already maps copy ops Inductor cannot lower (slice_copy, select_copy, split_copy) back to views. Map detach_copy to alias_copy, which Inductor decomposes, so no fallback kernel is needed. A view such as alias would let a method that returns both y and y.detach() hand the same output buffer back twice, which the CUDA runtime then frees twice. alias_copy gives a fresh buffer, and Edge accepts it in both dim order modes, unlike clone. The pass now also walks the branch graphs of control flow such as torch.cond, so every copy it maps (slice_copy, select_copy, split_copy and detach_copy) is rewritten there too. test_cuda_export lowers these cases. test_output_stride_rearrange runs the y, y.detach() case and compares both outputs with eager. --- .../passes/replace_view_copy_with_view.py | 31 ++++++++++----- backends/cuda/tests/test_cuda_export.py | 38 +++++++++++++++++++ .../tests/test_output_stride_rearrange.py | 16 ++++++++ 3 files changed, 76 insertions(+), 9 deletions(-) diff --git a/backends/aoti/passes/replace_view_copy_with_view.py b/backends/aoti/passes/replace_view_copy_with_view.py index 460d1f09981..c881e650e7a 100644 --- a/backends/aoti/passes/replace_view_copy_with_view.py +++ b/backends/aoti/passes/replace_view_copy_with_view.py @@ -31,6 +31,12 @@ # natively by inductor, so map it here like the other view-copy ops. torch.ops.aten.split_copy.Tensor: torch.ops.aten.split.Tensor, ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor, + # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either. + # Without autograd it is ``alias_copy``, which Inductor can decompose, + # unlike ``detach_copy``. A view would let two outputs share one handle, + # and ``clone`` fails Edge verification in dim order mode. + torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default, + ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default, } @@ -40,19 +46,26 @@ class ReplaceViewCopyWithViewPass(ExportPass): def call(self, graph_module: fx.GraphModule) -> PassResult: graph_changed = False - for node in graph_module.graph.nodes: - if node.op != "call_function" or node.target not in _VIEW_TARGETS: + # Control flow branches, such as those of torch.cond, are submodules. + for module in graph_module.modules(): + if not isinstance(module, fx.GraphModule): continue - if self._has_blocking_user(node, node.users.keys()): - continue + module_changed = False + for node in module.graph.nodes: + if node.op != "call_function" or node.target not in _VIEW_TARGETS: + continue + + if self._has_blocking_user(node, node.users.keys()): + continue - node.target = _VIEW_TARGETS[node.target] - graph_changed = True + node.target = _VIEW_TARGETS[node.target] + module_changed = True - if graph_changed: - graph_module.graph.lint() - graph_module.recompile() + if module_changed: + module.graph.lint() + module.recompile() + graph_changed = True return PassResult(graph_module, graph_changed) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index d08c835ca38..b49fa5ac753 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -328,6 +328,44 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: edge_program_manager = self._export_to_cuda_with_lower(module, inputs) self.assertIsNotNone(edge_program_manager, "Simple add operation export failed") + def test_detach(self): + """Edge turns detach into detach_copy, which has no CUDA kernel of its own.""" + + class DetachModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x * 2).detach() + 1 + + class DetachedOutputModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + y = x * 2 + return y, y.detach() + + class DetachInCondModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = torch.cond(x.sum() > 0, lambda t: t.sin().detach(), torch.cos, (x,)) + # Without this detach, Edge decomposes the branch's detach. + return y.detach() + 1 + + x = torch.randn(4, 8) + for module in (DetachModule(), DetachedOutputModule(), DetachInCondModule()): + with self.subTest(module=type(module).__name__): + # The default config also verifies the program the backend rewrites. + edge_program_manager = to_edge_transform_and_lower( + export(module, (x,), strict=True), + partitioner=[ + CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + ], + ) + self.assertTrue( + any( + "executorch_call_delegate" in str(node.target) + for node in edge_program_manager.exported_program().graph.nodes + ) + ) + edge_program_manager.to_executorch() + def test_conv2d(self): """Test CUDA export for 2D convolution.""" diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 24df64d6de9..aee72aba147 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -256,6 +256,22 @@ def forward(self, a, b): "simple add should match exactly", ) + def test_detached_output(self): + """A method that returns y and y.detach() gets two separate outputs.""" + + class DetachedOutputModule(nn.Module): + def forward(self, x): + y = x * 2 + return y, y.detach() + + x = torch.randn(4, 64, device="cuda") + module = DetachedOutputModule().eval() + outputs = _run_et_aoti(module, (x,)) + expected = module(x) + self.assertEqual(len(outputs), len(expected)) + for output, reference in zip(outputs, expected): + torch.testing.assert_close(output.cpu(), reference.cpu()) + if __name__ == "__main__": unittest.main() From 9eba5f4c2fd1f9cacf19245a95ea756ced1b450e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 23:09:26 -0700 Subject: [PATCH 05/37] Copy a strided CUDA delegate output into its ETensor AOTInductor can return a method output as a view into a larger buffer with strides that do not match a contiguous tensor, for example a slice `y[:, :, :14]` of a [1, 50, 32] result. The CUDA backend copies outputs into the planned ETensor only from contiguous SlimTensors, so execute() failed with "SlimTensor must be contiguous" although lowering accepted the program. The copy already has a slow path that rearranges elements through host memory when the strides differ from the ETensor's. Let it take non-contiguous views too: copy the span the view reaches instead of nbytes(), and check that span against the storage, with checked arithmetic. Wrapping a SlimTensor without a copy still requires a contiguous one. The first time a GPU output takes this path, the backend logs it once, because the path syncs the stream on every call. A CUDA graph replay copies each output's raw bytes. For such a strided output, and also for a contiguous output whose ETensor uses another dim order, the replay was silently wrong. A raw copy is right when the output is contiguous and the ETensor uses the default dim order, or its strides match the ETensor's on every dim larger than 1 (dims of size 1 never move an element). For any other layout, warmup now turns graphs off and the method runs without a graph. If only the capture call sees such a layout, for example because the input shape changed, that call returns an error instead of wrong values, and later calls run without a graph. --- backends/cuda/runtime/cuda_backend.cpp | 50 +++++++++++ .../runtime/test/test_cuda_output_copy.cpp | 19 ++++ backends/cuda/runtime/utils.h | 74 +++++++++++----- .../tests/test_output_stride_rearrange.py | 87 +++++++++++++++++++ 4 files changed, 210 insertions(+), 20 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 921c15ae7af..712ed5ff545 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -87,6 +87,28 @@ constexpr char kUseSharedCudaStream[] = "use_shared_cuda_stream"; constexpr char kEnableCudaGraphForMethod[] = "enable_cuda_graph_for_method"; constexpr int kCudaGraphWarmupSteps = 3; constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods"; + +// A CUDA graph replay copies an output's raw bytes into its ETensor, which is +// right only when every element sits at the same offset in both. Size-1 dims +// never move an element, so their strides may differ. +bool output_is_raw_copyable(const SlimTensor* output, const Tensor* et_output) { + if (!output->is_contiguous() || + output->dim() != static_cast(et_output->dim())) { + return false; + } + // The ETensor still has its previous shape here, so its strides can be stale + // for dynamic shapes; the default dim order is enough on its own. + if (executorch::runtime::is_contiguous_dim_order( + et_output->dim_order().data(), et_output->dim())) { + return true; + } + for (size_t d = 0; d < output->dim(); d++) { + if (output->size(d) != 1 && output->stride(d) != et_output->strides()[d]) { + return false; + } + } + return true; +} } // anonymous namespace bool CudaGraphState::start_warmup() { @@ -1053,6 +1075,23 @@ class ET_EXPERIMENTAL CudaBackend final "cudaStreamEndCapture failed: %s", cudaGetErrorString(gerr)); + // Warmup checks this too, but the capture step's inputs can differ. + // The outputs belong to the captured graph, so leave them to it. + for (size_t i = 0; i < n_outputs; i++) { + if (pre_run_outputs[i] != slim_outputs[i] && + !output_is_raw_copyable( + slim_outputs[i], &(args[i + n_inputs]->toTensor()))) { + std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); + ET_LOG( + Error, + "CUDA graph: output %zu of '%s' is not laid out like its " + "ETensor; this call fails and later calls run without a graph", + i, + handle->method_name.c_str()); + return Error::NotSupported; + } + } + gerr = cudaGraphInstantiate( &handle->cuda_graph_state.graph_exec, handle->cuda_graph_state.graph, @@ -1140,6 +1179,17 @@ class ET_EXPERIMENTAL CudaBackend final } else { // Case (b): AOTI returned its own buffer. D2D copy the result into the // planned GPU ET buffer, then free AOTI's buffer. + // + // A graph replay could not land this output, so run the method eagerly. + if (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && + !output_is_raw_copyable(slim_outputs[i], et_output)) { + handle->cuda_graph_state.phase = CudaGraphPhase::Disabled; + ET_LOG( + Info, + "CUDA graph: disabled for '%s', output %zu is not laid out like its ETensor", + handle->method_name.c_str(), + i); + } ET_CHECK_OK_OR_RETURN_ERROR( copy_slimtensor_to_device_etensor_async( slim_outputs[i], et_output, cuda_stream), diff --git a/backends/cuda/runtime/test/test_cuda_output_copy.cpp b/backends/cuda/runtime/test/test_cuda_output_copy.cpp index 34f36934f73..1952b23a18f 100644 --- a/backends/cuda/runtime/test/test_cuda_output_copy.cpp +++ b/backends/cuda/runtime/test/test_cuda_output_copy.cpp @@ -84,6 +84,15 @@ TEST_F( EXPECT_EQ(copied(), (std::vector{4, 5, 6, 7})); } +// AOTInductor can return a slice of a larger buffer, which is not contiguous. +TEST_F(CudaOutputCopyTest, CopiesAStridedView) { + slim::SlimTensor columns = slim::from_blob( + source_, {2, 4}, slim::c10::ScalarType::Float, slim::DEFAULT_CUDA_DEVICE); + columns.as_strided_({2, 2}, {4, 1}, /*storage_offset=*/1); + ASSERT_EQ(copy(columns), Error::Ok); + EXPECT_EQ(copied(), (std::vector{1, 2, 5, 6})); +} + TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { // from_blob sizes the storage to fit its view, so moving the view past the // end takes as_strided_. The moved view is still inside the allocation. @@ -91,6 +100,16 @@ TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { past_end.as_strided_({2, 2}, {2, 1}, /*storage_offset=*/4); EXPECT_EQ(copy(past_end), Error::InvalidArgument); + // Its nbytes() would fit after the offset, but its last element would not. + slim::SlimTensor columns = slim::from_blob( + source_, {2, 4}, slim::c10::ScalarType::Float, slim::DEFAULT_CUDA_DEVICE); + columns.as_strided_({2, 2}, {4, 1}, /*storage_offset=*/3); + EXPECT_EQ(copy(columns), Error::InvalidArgument); + + // Its extent in bytes does not fit in size_t. + columns.as_strided_({2, 2}, {(int64_t{1} << 62) - 1, 1}, 0); + EXPECT_EQ(copy(columns), Error::InvalidArgument); + EXPECT_EQ( copy(view(static_cast(source_) + 4, /*storage_offset=*/-4)), Error::InvalidArgument); diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index 30a79500c16..6ca1c64d493 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -23,6 +23,26 @@ namespace executorch::backends::cuda { namespace { +// Bytes from a view's first element to one past its last, which equals +// nbytes() for a contiguous view. False if that does not fit in size_t. +inline bool _extent_nbytes( + const executorch::backends::aoti::slim::SlimTensor* slim_tensor, + size_t* extent_nbytes) { + size_t last = 0; + for (size_t i = 0; i < slim_tensor->dim(); ++i) { + size_t term = 0; + if (::c10::mul_overflows( + static_cast(slim_tensor->sizes()[i] - 1), + static_cast(slim_tensor->strides()[i]), + &term) || + ::c10::add_overflows(last, term, &last)) { + return false; + } + } + return !::c10::add_overflows(last, size_t{1}, &last) && + !::c10::mul_overflows(last, slim_tensor->itemsize(), extent_nbytes); +} + inline executorch::runtime::Error _check_tensor_metadata( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, executorch::runtime::etensor::Tensor* etensor) { @@ -34,30 +54,28 @@ inline executorch::runtime::Error _check_tensor_metadata( ET_CHECK_OR_RETURN_ERROR( etensor != nullptr, InvalidArgument, "etensor pointer cannot be nullptr"); - // Check that SlimTensor is contiguous - ET_CHECK_OR_RETURN_ERROR( - slim_tensor->is_contiguous(), - InvalidArgument, - "SlimTensor must be contiguous"); - - // The copy reads nbytes() from data_ptr(), which adds the offset, so the - // view has to start and end inside its storage. + // The copy reads every element the view reaches, starting at data_ptr(), + // which adds the offset, so that span has to stay inside the storage. const int64_t storage_offset = slim_tensor->storage_offset(); ET_CHECK_OR_RETURN_ERROR( storage_offset >= 0, InvalidArgument, "SlimTensor storage_offset must not be negative, got %ld", static_cast(storage_offset)); - const size_t nbytes = slim_tensor->nbytes(); - if (nbytes > 0) { + if (slim_tensor->numel() > 0) { + size_t span_nbytes = 0; + ET_CHECK_OR_RETURN_ERROR( + _extent_nbytes(slim_tensor, &span_nbytes), + InvalidArgument, + "SlimTensor view extent does not fit in size_t"); const size_t storage_nbytes = slim_tensor->storage()->nbytes(); ET_CHECK_OR_RETURN_ERROR( - nbytes <= storage_nbytes && + span_nbytes <= storage_nbytes && static_cast(storage_offset) <= - (storage_nbytes - nbytes) / slim_tensor->itemsize(), + (storage_nbytes - span_nbytes) / slim_tensor->itemsize(), InvalidArgument, - "SlimTensor view of %zu bytes at storage_offset %ld runs past its %zu byte storage", - nbytes, + "SlimTensor view spanning %zu bytes at storage_offset %ld runs past its %zu byte storage", + span_nbytes, static_cast(storage_offset), storage_nbytes); } @@ -119,7 +137,7 @@ inline bool _strides_match( return true; } -// Element-by-element copy from a contiguous CPU buffer (src_strides layout) +// Element-by-element copy from a CPU buffer (src_strides layout) // to an ETensor (dst_strides layout), rearranging data to match. inline void _strided_copy( void* dst, @@ -201,18 +219,22 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( } } else { // Slow path: strides differ (e.g., AOTI delegate output layout differs - // from .pte's dim_order). Copy to a temp CPU buffer, rearrange + // from .pte's dim_order, or AOTI returned a strided view such as a slice). + // Copy the span the view reaches to a temp CPU buffer, rearrange // element-by-element to match the ETensor's expected layout, then move the // result to the destination (CPU stays in place; GPU gets an H2D copy). - std::vector tmp(nbytes); + // _check_tensor_metadata above already checked that this fits. + size_t span_nbytes = 0; + _extent_nbytes(slim_tensor, &span_nbytes); + std::vector tmp(span_nbytes); if (slim_tensor->is_cpu()) { - std::memcpy(tmp.data(), src_data, nbytes); + std::memcpy(tmp.data(), src_data, span_nbytes); } else { if (stream) { ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(stream)); } - ET_CUDA_CHECK_OR_RETURN_ERROR( - cudaMemcpy(tmp.data(), src_data, nbytes, cudaMemcpyDeviceToHost)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy( + tmp.data(), src_data, span_nbytes, cudaMemcpyDeviceToHost)); } const size_t ndim = slim_tensor->dim(); @@ -241,6 +263,14 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( src_strides_vec, dst_strides_vec); } else { + static const bool logged = [] { + ET_LOG( + Info, + "CUDA backend: an output is not laid out like its ETensor, so it " + "is copied through host memory with a stream sync on every call"); + return true; + }(); + (void)logged; // Rearrange into a CPU staging buffer, then copy to the GPU destination. std::vector rearranged(nbytes); _strided_copy( @@ -363,6 +393,10 @@ inline executorch::runtime::Error wrap_slimtensor_to_etensor( const executorch::backends::aoti::slim::SlimTensor* slim_tensor, executorch::runtime::etensor::Tensor* etensor) { ET_CHECK_OK_OR_RETURN_ERROR(_check_tensor_metadata(slim_tensor, etensor)); + ET_CHECK_OR_RETURN_ERROR( + slim_tensor->is_contiguous(), + InvalidArgument, + "SlimTensor must be contiguous"); // Set data pointer to point directly to SlimTensor's data (zero-copy) etensor->unsafeGetTensorImpl()->set_data( diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index aee72aba147..b1bf91638a7 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -273,5 +273,92 @@ def forward(self, x): torch.testing.assert_close(output.cpu(), reference.cpu()) +class _SlicedLinear(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(64, 32) + + def forward(self, x): + return self.proj(x)[:, :, :14] + + +class TestStridedOutputWithCudaGraph(unittest.TestCase): + """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA not available") + if torch.version.hip is not None: + raise unittest.SkipTest("CUDA graphs are not supported on ROCm") + + def _run_with_cuda_graph(self, module, example, inputs, dynamic_shapes=None): + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import to_edge_transform_and_lower + from executorch.runtime import Runtime + + with torch.no_grad(): + ep = torch.export.export(module, example, dynamic_shapes=dynamic_shapes) + spec = CudaBackend.generate_method_name_compile_spec("forward") + program = to_edge_transform_and_lower( + ep, partitioner=[CudaPartitioner([spec])] + ).to_executorch() + runtime = Runtime.get() + option = "enable_cuda_graph_for_method" + results = [] + with tempfile.TemporaryDirectory() as tmpdir: + pte = os.path.join(tmpdir, "model.pte") + with open(pte, "wb") as f: + program.write_to_file(f) + program.write_tensor_data_to_file(tmpdir) + ptd = glob.glob(os.path.join(tmpdir, "*.ptd"))[0] + runtime.backend_registry.set_option("CudaBackend", {option: "forward"}) + try: + method = runtime.load_program(pte, data_path=ptd).load_method("forward") + finally: + # The option is process wide. + runtime.backend_registry.set_option("CudaBackend", {option: ""}) + for x in inputs: + try: + results.append(method.execute([x])[0]) + except RuntimeError as e: + results.append(e) + return results + + def _assert_matches_eager(self, module, x, out): + self.assertIsInstance(out, torch.Tensor) + with torch.no_grad(): + expected = module(x.cuda()).cpu() + torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4) + + def test_sliced_output_matches_eager_past_warmup(self): + torch.manual_seed(0) + module = _SlicedLinear().cuda().eval() + # Three warmup calls, the capture call, then replays. + inputs = [torch.randn(1, 50, 64) for _ in range(6)] + outputs = self._run_with_cuda_graph(module, (inputs[0].cuda(),), inputs) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + self._assert_matches_eager(module, x, out) + + def test_strided_output_first_seen_at_capture(self): + torch.manual_seed(0) + module = _SlicedLinear().cuda().eval() + # With one row the slice is contiguous, so only the capture call sees a + # strided output. + inputs = [torch.randn(1, rows, 64) for rows in (1, 1, 1, 50, 50, 50)] + rows = torch.export.Dim("rows", min=1, max=64) + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({1: rows},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + if i == 3: + self.assertIsInstance(out, RuntimeError) + else: + self._assert_matches_eager(module, x, out) + + if __name__ == "__main__": unittest.main() From 22cc1f018388d01ae66a9ec6ba7c4f5d98b5bc05 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 00:29:25 -0700 Subject: [PATCH 06/37] Do not fold a tensor that an in-place op writes in constant_prop_pass torch.export keeps in-place ops in the program it returns. In t = torch.zeros(6) t[1:3] = x[:2] return t * 2 `t` and its slice look constant to constant_prop_pass, so the pass turned them into constants, while `copy_` still writes into them at run time. Decomposing the program then fails with "Constant ... is mutated in the forward method". A buffer written in place has the same problem, and there the folded reads are silently wrong, because buffers_to_mutate does not list it until the program is decomposed. Find the tensors an op writes in place, through any argument, including out= and lists, as insert_write_back_for_buffers_pass already reads them. Then do not fold them, or any node that shares their storage. The storage is read from the fake tensor in node.meta["val"], so views that composite ops return before decomposition, like torch.atleast_2d, torch.meshgrid and torch.broadcast_tensors, count as well. A program loaded with torch.export.load does not share storage between views, so the tensor that a written view or getitem reads is also kept. Views from those composite ops are only tracked before the program is saved. A write inside a torch.no_grad() or torch.autocast block sits in the submodule of a wrap_with_set_grad_enabled or wrap_with_autocast op. Look inside that submodule and map each written placeholder to the operand passed in its place. Operands that are only read, such as weights, can still be folded. For other higher order ops, every operand counts as written once a submodule writes one of its inputs. The final dead code cleanup keeps these ops, which have no users but write, and only cleans the top-level graph, because Graph.eliminate_dead_code would also drop a region nested inside another one. Programs without in-place ops, such as Edge programs, are unchanged. --- exir/passes/BUCK | 1 + exir/passes/constant_prop_pass.py | 130 +++++++++++++++++++++++++++++- exir/tests/test_passes.py | 63 +++++++++++++++ 3 files changed, 191 insertions(+), 3 deletions(-) diff --git a/exir/passes/BUCK b/exir/passes/BUCK index 6eeed352f34..9deed4abcea 100644 --- a/exir/passes/BUCK +++ b/exir/passes/BUCK @@ -180,6 +180,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/exir/dialects/edge:lib", "//executorch/exir/operator:util", "//executorch/exir/passes:replace_aten_with_edge_pass", + ":insert_write_back_for_buffers_pass", ], ) diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 324ff21db48..44b0121268c 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -7,6 +7,7 @@ # pyre-unsafe import logging +import operator from collections import OrderedDict from typing import cast, Mapping, Optional @@ -15,6 +16,9 @@ from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.operator.util import _QUANT_PRIMITIVES +from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + _schema_paired_args, +) from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge from torch._export.utils import ( get_buffer, @@ -26,6 +30,7 @@ from torch._guards import detect_fake_mode from torch.export import ExportedProgram from torch.export.exported_program import InputKind, InputSpec, TensorArgument +from torch.multiprocessing.reductions import StorageWeakRef from torch.utils import _pytree as pytree # Avoid propagating constants for `exir.ops.edge.aten.full.default`. @@ -110,7 +115,12 @@ def get_constant_placeholder_dict( Returns a dictionary of placeholder node -> constant tensor. """ const_node_to_tensor: OrderedDict[torch.fx.Node, torch.Tensor] = OrderedDict() + # Before run_decompositions, buffers_to_mutate does not list a buffer that + # is written in place, so it looks like a constant here. + written = _nodes_written_in_place(exported_program.graph_module) for node in exported_program.graph.find_nodes(op="placeholder"): + if node in written: + continue if is_param(exported_program, node): const_node_to_tensor[node] = cast( torch.Tensor, get_param(exported_program, node) @@ -126,6 +136,106 @@ def get_constant_placeholder_dict( return const_node_to_tensor +# Higher order ops whose submodule takes, in order, the operands after it. +_WRAPPER_HOPS = ( + torch.ops.higher_order.wrap_with_set_grad_enabled, + torch.ops.higher_order.wrap_with_autocast, +) + + +def _storages(node: torch.fx.Node) -> set[StorageWeakRef]: + # Only strided tensors have a storage to share. Sparse ones raise here. + return { + StorageWeakRef(val.untyped_storage()) + for val in pytree.tree_leaves(node.meta.get("val")) + if isinstance(val, torch.Tensor) and val.layout == torch.strided + } + + +def _hop_written_operands( + graph_module: torch.fx.GraphModule, node: torch.fx.Node +) -> list[torch.fx.Node]: + """ + Returns the operands of a higher order op, such as a no_grad or autocast + region, that an in-place op inside its submodules writes. + """ + written = [] + for i, arg in enumerate(node.args): + if not isinstance(arg, torch.fx.Node) or arg.op != "get_attr": + continue + submodule = getattr(graph_module, arg.target) + if not isinstance(submodule, torch.fx.GraphModule): + continue + submodule_written = _nodes_written_in_place(submodule) + placeholders = submodule.graph.find_nodes(op="placeholder") + if node.target in _WRAPPER_HOPS: + written += [ + operand + for placeholder, operand in zip(placeholders, node.args[i + 1 :]) + if placeholder in submodule_written + and isinstance(operand, torch.fx.Node) + ] + elif any(p in submodule_written for p in placeholders): + # The operand layout is op specific, so any operand may be written. + written += [ + leaf + for leaf in pytree.tree_leaves(node.args) + if isinstance(leaf, torch.fx.Node) + ] + return written + + +def _nodes_written_in_place( + graph_module: torch.fx.GraphModule, +) -> set[torch.fx.Node]: + """ + Returns the nodes whose value an in-place op writes, directly, through a + view or inside a higher order op. Such a value has to stay a runtime + tensor, not a constant. + """ + written = set() + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + schema = getattr(node.target, "_schema", None) + if schema is None: + written.update(_hop_written_operands(graph_module, node)) + continue + for arg, schema_arg in _schema_paired_args(node, schema): + alias_info = schema_arg.alias_info + if alias_info is not None and alias_info.is_write: + written.update( + leaf + for leaf in pytree.tree_leaves(arg) + if isinstance(leaf, torch.fx.Node) + ) + if not written: + return written + # Every view of a written tensor shares its storage, including views that + # composite ops like atleast_2d or meshgrid return before decomposition. + written_storages = set().union(*(_storages(node) for node in written)) + written.update( + node + for node in graph_module.graph.nodes + if not _storages(node).isdisjoint(written_storages) + ) + # A deserialized program does not share storage between a view and its + # base, so a write through a view also marks the tensor the view reads. + pending = list(written) + while pending: + node = pending.pop() + if node.op != "call_function" or not node.args: + continue + base = node.args[0] + is_view = node.target is operator.getitem or getattr( + node.target, "is_view", False + ) + if is_view and isinstance(base, torch.fx.Node) and base not in written: + written.add(base) + pending.append(base) + return written + + def get_propagated_const_tensor_dict( exported_program: ExportedProgram, custom_skip_targets: Optional[set[EdgeOpOverload]], @@ -142,9 +252,12 @@ def get_propagated_const_tensor_dict( # Default set of targets to skip. all_skip_targets = _DEFAULT_SKIP_TARGETS + written_in_place = _nodes_written_in_place(exported_program.graph_module) + for node in exported_program.graph.nodes: if ( node.op != "call_function" + or node in written_in_place or node.target is memory.alloc or node.target in all_skip_targets # Ops with side effects (RNG draws, mutation) have to run at @@ -426,8 +539,19 @@ def constant_prop_pass( _update_output_node_and_specs(exported_program, _skip_dim_order=_skip_dim_order) - # Cleanup the graph. - exported_program.graph.eliminate_dead_code() - exported_program.graph_module.recompile() + # Cleanup the graph. A higher order op that writes its operands has no + # users, yet it has to stay. Graph.eliminate_dead_code also cleans + # submodules without this rule, which would drop a nested region, so only + # the top-level graph is cleaned here. + graph_module = exported_program.graph_module + for node in reversed(graph_module.graph.nodes): + if ( + node.op not in ("placeholder", "output") + and not node.users + and not node.is_impure() + and not _hop_written_operands(graph_module, node) + ): + graph_module.graph.erase_node(node) + graph_module.recompile() return exported_program diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index eb917f55490..81c5298aab1 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2915,6 +2915,69 @@ 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_keeps_tensor_written_in_place(self) -> None: + """ + torch.export keeps in-place ops. A tensor or buffer that one writes, + directly, through a view, as an out= argument, in a list or inside a + no_grad or autocast block, has to stay a runtime tensor, or the folded + value is wrong and decomposing can fail. + """ + + class WrittenInPlace(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + for name in "bcdef": + self.register_buffer(name, torch.zeros(6)) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + t = torch.zeros(2, 6) + rows = t.unbind(0) + rows[0].add_(x) + out = torch.zeros(6) + torch.mul(x, 2, out=out) + listed = torch.zeros(6) + torch._foreach_add_([listed], [x]) + self.b[1:3].copy_(x[:2]) + no_grad = torch.zeros(6) + with torch.no_grad(): + no_grad[1:3] = x[:2] + self.c[1:3].copy_(x[:2]) + autocast = torch.zeros(6) + with torch.autocast("cpu", enabled=False): + autocast[2:4] = x[:2] + nested = torch.zeros(6) + with torch.no_grad(): + with torch.autocast("cpu", enabled=False): + nested[1:3] = x[:2] + torch.atleast_2d(self.d)[0].add_(x) + torch.meshgrid(self.e, x[:1], indexing="ij")[0][:, 0].add_(x) + torch.broadcast_tensors(self.f, x.view(1, 6))[0][0].add_(x) + written = rows[0] + rows[1] + out + listed + no_grad + autocast + nested + buffers = self.b + self.c + self.d + self.e + self.f + return written + buffers * 2, self.b + + x = torch.randn(6) + ep = constant_prop_pass(export(WrittenInPlace(), (x,), strict=False)) + self.assertEqual(len(ep.constants), 0) + torch.testing.assert_close(ep.module()(x), WrittenInPlace()(x)) + + # A saved and loaded program no longer shares storage between views. + class WrittenThroughSlice(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("b", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.b[:2].copy_(x) + return self.b * 2 + + x = torch.randn(2) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "written.pt2") + torch.export.save(export(WrittenThroughSlice(), (x,), strict=False), path) + ep = constant_prop_pass(torch.export.load(path)) + torch.testing.assert_close(ep.module()(x), WrittenThroughSlice()(x)) + def test_constant_prop_pass_detaches_folded_parameters(self) -> None: class M(torch.nn.Module): def __init__(self): From 4d8b5603dd99407683b49f3fbf34a6f57729932c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:09:48 -0700 Subject: [PATCH 07/37] [MLX] Check tensor list inputs for unsupported dtypes The MLX partitioner decides which nodes the delegate can take, and part of that is rejecting nodes whose inputs have a dtype MLX cannot hold, such as float64. The check only looked at arguments that are a single node, so a tensor list argument (the inputs of aten.cat, for example) was never checked. As a result a cat of float64 tensors was put into the delegate. Its float64 inputs became delegate inputs, and lowering then failed with "placeholder has unsupported dtype: torch.float64" instead of leaving the cat on the CPU. Sinusoidal timestep embeddings computed in float64 hit this. Check every node inside each positional argument, so such a cat stays on the CPU like any other float64 op. --- backends/mlx/builder/program_builder.py | 11 ++++++----- backends/mlx/test/test_partitioner.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py index 797cdb1bdf2..c7ed031d0a9 100644 --- a/backends/mlx/builder/program_builder.py +++ b/backends/mlx/builder/program_builder.py @@ -95,12 +95,13 @@ def _check_input_dtypes(node: Node) -> Optional[str]: None if all input dtypes are supported, otherwise an error message string describing which input (arg position or kwarg name) has an unsupported dtype """ - # Check positional args + # Check positional args, including tensor lists such as aten.cat's for i, arg in enumerate(node.args): - if isinstance(arg, Node): - dtype_error = _check_dtype(arg) - if dtype_error is not None: - return f"arg[{i}] ({arg.name}) {dtype_error}" + for a in pytree.tree_leaves(arg): + if isinstance(a, Node): + dtype_error = _check_dtype(a) + if dtype_error is not None: + return f"arg[{i}] ({a.name}) {dtype_error}" # Check kwargs for kwarg_name, kwarg_val in node.kwargs.items(): diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index d8b7b8e485a..065e68a1a4d 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -150,6 +150,20 @@ def test_consumed_constant_still_raises(self): builder.build() +class TestMLXPartitionerUnsupportedDtypes(unittest.TestCase): + def test_float64_tensor_list_stays_off_the_delegate(self): + class Float64Cat(nn.Module): + def forward(self, x, y): + return x + torch.cat([y.sin(), y.cos()]).float() + + _, delegates, err = _run( + Float64Cat().eval(), + (torch.randn(8), torch.randn(4, dtype=torch.float64)), + ) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-6) + + class Sdpa(nn.Module): def __init__(self, is_causal: bool = False): super().__init__() From 494ea5887bf9c8577058f0af31ecbfe1cf7e5c7c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:15:50 -0700 Subject: [PATCH 08/37] [MLX] Share one constant slot between placeholders of the same weight When a weight is used by several partitions, to_backend copies its placeholder once per extra user. Every copy keeps the original target, so one MLX subgraph can get several placeholders for the same weight (a layer applied at more than one step, for example). The builder made a separate constant slot for each of them, but the named slots are keyed by target, so only one of those slots got a name. Loading the program then failed with "load_constants: no name found for constant tensor N". Give placeholders that share a target one constant slot. The weight is then stored and loaded once per subgraph. --- backends/mlx/builder/program_builder.py | 19 ++++++++++++++----- backends/mlx/test/test_partitioner.py | 17 +++++++++++++++++ 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py index c7ed031d0a9..d5ab4124579 100644 --- a/backends/mlx/builder/program_builder.py +++ b/backends/mlx/builder/program_builder.py @@ -462,7 +462,7 @@ def _make_io_slots(self): # noqa: C901 ) output_kind_targets = defaultdict(set) - constant_tensors = [] + constant_tensors = {} user_inputs = [] user_outputs = [] mutable_buffers = [] @@ -491,17 +491,17 @@ def _make_io_slots(self): # noqa: C901 if isinstance(arg, TensorArgument): if kind == InputKind.PARAMETER: # Parameters are treated as constants (not mutated) - constant_tensors.append(name) + constant_tensors[name] = target elif kind == InputKind.BUFFER: if target in output_kind_targets[OutputKind.BUFFER_MUTATION]: mutable_buffers.append(name) else: # Non-mutated buffers (like lifted tensor constants) are constants - constant_tensors.append(name) + constant_tensors[name] = target elif kind == InputKind.USER_INPUT: user_inputs.append(name) elif kind == InputKind.CONSTANT_TENSOR: - constant_tensors.append(name) + constant_tensors[name] = target else: raise NotImplementedError( f"Support for input {arg} is not implemented" @@ -520,12 +520,21 @@ def _make_io_slots(self): # noqa: C901 else: raise NotImplementedError(f"Support for input {arg} is not implemented") + target_to_constant_slot = {} for node in self.ep.graph.nodes: if node.op == "placeholder": if node.users == {}: continue if node.name in constant_tensors: - self.make_or_get_slot(node, id_space=IdSpace.Constant) + # A constant shared with other partitions is duplicated once per + # user, so several placeholders here can name the same target. + target = constant_tensors[node.name] + if target in target_to_constant_slot: + self.set_slot(node, target_to_constant_slot[target]) + else: + target_to_constant_slot[target] = self.make_or_get_slot( + node, id_space=IdSpace.Constant + ) elif node.name in user_inputs: val = node.meta.get("val", None) if isinstance(val, torch.Tensor) and not val.is_contiguous(): diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 065e68a1a4d..f20ab617825 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -337,6 +337,23 @@ def test_zero_head_count_declines_instead_of_raising(self): self.assertEqual(get_mlx_node_counts(path).get("SdpaNode", 0), 0) +class TestMLXPartitionerSharedConstants(unittest.TestCase): + def test_weight_used_twice_in_one_delegate_and_once_in_another(self): + class SharedLinear(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x): + y = self.linear(x) + self.linear(x * 2) + # elu is not lowered, so the last use lands in a second delegate. + return self.linear(torch.nn.functional.elu(y)) + + _, delegates, err = _run(SharedLinear().eval(), (torch.randn(2, 4),)) + self.assertGreater(delegates, 1) + self.assertLess(err, 1e-5) + + class TestMLXPartitionerGroupedKeys(unittest.TestCase): """Repeated KV heads are fused only when the caller enables the shared passes.""" From 852d46c33b4bbf0f0d17b1f2a118c7cd9b8c83c7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:57:02 -0700 Subject: [PATCH 09/37] Skip CoreML ops on bfloat16 tensors coremltools has no bfloat16 type. The support query that CoreMLPartitioner asks coremltools only checks the op name, so an op that reads or produces a bfloat16 tensor was still delegated, and coremltools 9.0 then raised "KeyError: torch.bfloat16" during lowering. The whole export failed instead of leaving that op on the CPU. Models that run part of their forward pass under bfloat16 autocast hit this. Override support for nodes with a bfloat16 input or output, the same way the partitioner already skips other ops coremltools cannot convert. That rejection is per call, but ops_to_not_decompose keeps an op whole as soon as one of its calls is supported. With attention run once in float32 and once in bfloat16, the bfloat16 call was kept whole, rejected by the partitioner, and left with no kernel ("Missing out variants"). Keep an op whole only when every call of it is supported, as the MLX partitioner does. --- .../coreml/partition/coreml_partitioner.py | 38 +++++++++++++++---- .../coreml/test/test_coreml_partitioner.py | 33 ++++++++++++++++ 2 files changed, 64 insertions(+), 7 deletions(-) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index b4e550d8479..9e6572014cf 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -83,6 +83,18 @@ def _is_arg_min_max_over_flattened_input(node: torch.fx.Node) -> bool: return dim is None +def _reads_or_produces_bfloat16(node: torch.fx.Node) -> bool: + """coremltools 9.0 has no bfloat16 type. + + It raises ``KeyError: torch.bfloat16`` on any op that reads or produces a + bfloat16 tensor, instead of reporting the op as unsupported. + """ + val = node.meta.get("val") + vals = list(val) if isinstance(val, (list, tuple)) else [val] + vals += [arg.meta.get("val") for arg in node.all_input_nodes] + return any(isinstance(v, torch.Tensor) and v.dtype == torch.bfloat16 for v in vals) + + def _is_view_op(op: torch._ops.OpOverload) -> bool: schema = op._schema if len(schema.arguments) == 0: @@ -159,6 +171,13 @@ def should_override_support(self, node) -> bool: ) return True + if _reads_or_produces_bfloat16(node): + self.log_once( + f"{node.target} on a bfloat16 tensor is not supported by CoreML. " + "Overriding op support." + ) + return True + # TODO: enable this after bugs in ExecuTorch's partitioner are fixed # # If lower_full_graph=False, do not partition nodes with symbolic args because it can result in symbolic args # # in the placeholders due to partitioning, which CoreML does not support @@ -362,21 +381,26 @@ def ops_to_not_decompose( torch.ops.aten.repeat_interleave.self_int, torch.ops.aten.repeat_interleave.self_Tensor, ] + rejected = set() for node in ep.graph.nodes: if node.op == "call_function" and isinstance( node.target, torch._ops.OpOverload ): try: - if ( - op_support.is_node_supported(None, node) - and node.target not in do_not_decompose_blocklist - and not _is_view_op(node.target) - ): - do_not_decompose.append(node.target) + supported = op_support.is_node_supported(None, node) except Exception as e: # CoreML's op_support.is_node_supported will sometimes throw # for unsupported ops, rather than returning False self.log_once( f"Encountered exception when checking node support, treating node as unsupported: {e}" ) - return do_not_decompose, None + supported = False + if not supported: + rejected.add(node.target) + elif node.target not in do_not_decompose_blocklist and not _is_view_op( + node.target + ): + do_not_decompose.append(node.target) + # Without a per-node filter, a preserved op stays whole at every call, so + # one rejected call would be left with no kernel. + return [op for op in do_not_decompose if op not in rejected], None diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 78dd114eeda..941c3e14b4c 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -433,6 +433,39 @@ def forward(self, x): self.assertIn("executorch_call_delegate", op_names) self.assertNotIn("aten.argmax.default", op_names) + def test_bfloat16_ops_are_skipped(self): + """ + coremltools has no bfloat16 type and raises on ops that use it. The + partitioner must leave those ops to the portable backend and still + delegate the rest, also when the same op has a float32 call that Core + ML takes (as under bfloat16 autocast). + """ + + class Model(torch.nn.Module): + def forward(self, q): + a = torch.nn.functional.scaled_dot_product_attention(q, q, q) + b = q.to(torch.bfloat16) + c = torch.nn.functional.scaled_dot_product_attention(b, b, b) + return a + c.float() + + model = Model().eval() + x = torch.randn(1, 2, 4, 8) + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("executorch_call_delegate", op_names) + self.assertNotIn("aten.scaled_dot_product_attention.default", op_names) + if _TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2)) + def test_tied_embedding_is_quantized_with_its_linear(self): """A weight that is both an embedding table and an output projection lowers. From c59c664bac7fed12c7b477a7aadb31602443ae61 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:05:44 -0700 Subject: [PATCH 10/37] Skip CoreML ops on tensors above rank 5 A Core ML tensor has at most 5 dimensions. The support query that CoreMLPartitioner asks coremltools only checks the op name, so an op on a rank 6 tensor (for example a stack of 5-D tensors) was still delegated, and coremltools 9.0 then raised "Core ML only supports tensors with rank <= 5" during lowering. The whole export failed instead of leaving that op on the CPU. coremltools does rewrite layout chains through a higher rank, such as the view, permute, view of a window partition, as long as each step has static shapes and a single user. So override support for a node when the tensors above rank 5 it is connected to reach anything other than such a chain of view, permute and clone ops. Those ops then run on the portable backend, and window partitions stay delegated as before. --- .../coreml/partition/coreml_partitioner.py | 71 +++++++++++++++++++ .../coreml/test/test_coreml_partitioner.py | 55 ++++++++++++++ 2 files changed, 126 insertions(+) diff --git a/backends/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index 9e6572014cf..4453818f956 100644 --- a/backends/apple/coreml/partition/coreml_partitioner.py +++ b/backends/apple/coreml/partition/coreml_partitioner.py @@ -63,6 +63,65 @@ ) +# Ops that coremltools can carry through a tensor above rank 5, by rewriting a +# reshape -> transpose -> reshape chain into rank 4 steps. +_HIGH_RANK_LAYOUT_TARGETS = frozenset( + [ + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.clone.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + ] +) + + +def _produces_tensor_above_rank_5(node: torch.fx.Node) -> bool: + val = node.meta.get("val") + vals = list(val) if isinstance(val, (list, tuple)) else [val] + return any(isinstance(v, torch.Tensor) and v.dim() > 5 for v in vals) + + +def _touches_unsupported_high_rank_tensor( + node: torch.fx.Node, is_skipped: Callable[[torch.fx.Node], bool] +) -> bool: + """Core ML tensors have at most rank 5. + + coremltools 9.0 raises "Core ML only supports tensors with rank <= 5" for + any tensor above rank 5 that it cannot rewrite away. It only rewrites a + static, unbranched layout chain, so such a tensor is fine only when every + node connected to it through tensors above rank 5 is a delegated view, + permute or clone, and every node producing such a tensor has a static shape + and a single user. + """ + seen, todo = set(), [node] + while todo: + n = todo.pop() + if n in seen: + continue + seen.add(n) + high_rank_inputs = [ + arg for arg in n.all_input_nodes if _produces_tensor_above_rank_5(arg) + ] + high_rank_output = _produces_tensor_above_rank_5(n) + if not high_rank_inputs and not high_rank_output: + continue + if ( + n.op != "call_function" + or n.target not in _HIGH_RANK_LAYOUT_TARGETS + or is_skipped(n) + ): + return True + if high_rank_output and ( + len(n.users) != 1 + or any(isinstance(d, torch.SymInt) for d in n.meta["val"].shape) + ): + return True + todo += high_rank_inputs + if high_rank_output: + todo += list(n.users) + return False + + _ARG_MIN_MAX_TARGETS = ( torch.ops.aten.argmax.default, torch.ops.aten.argmin.default, @@ -164,6 +223,18 @@ def should_override_support(self, node) -> bool: ) return True + if _touches_unsupported_high_rank_tensor( + node, + lambda n: self.should_skip_op_for_delegation( + getattr(n.target, "__name__", "").lower() + ), + ): + self.log_once( + f"{node.target} on a tensor above rank 5 is not supported by " + "CoreML. Overriding op support." + ) + return True + if _is_arg_min_max_over_flattened_input(node): self.log_once( "torch.ops.aten.{argmax, argmin}.default with dim=None is " diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index 941c3e14b4c..0129eb1d33f 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -386,6 +386,61 @@ def forward(self, x): self.assertIn("executorch_call_delegate", op_names) self.assertNotIn("aten.randn.default", op_names) + def test_ops_above_rank_5_are_skipped(self): + """ + Core ML tensors have at most rank 5. coremltools only rewrites an + unbranched view, permute, view chain through a higher rank (a window + partition), and raises on any other op above rank 5. The partitioner + must leave those ops to the portable backend and keep such a layout + chain delegated. + """ + + class Model(torch.nn.Module): + def forward(self, x): + windows = x.view(2, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5) + windows = windows.reshape(-1, 4, 8).relu() + branched = x.view(2, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5) + branched = ( + branched.reshape(-1, 4, 8).sum() + branched.reshape(32, 8).sum() + ) + stacked = torch.stack([x.unsqueeze(0), x.unsqueeze(0) * 2]) + return windows.sum() + branched + stacked.sum(0).relu().sum() + + model = Model().eval() + x = torch.randn(2, 4, 4, 8) + ep = torch.export.export(model, (x,), strict=True) + edge = executorch.exir.to_edge_transform_and_lower( + ep, partitioner=[CoreMLPartitioner()] + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertIn("executorch_call_delegate", op_names) + self.assertIn("aten.sum.dim_IntList", op_names) + self.assertEqual(op_names.count("aten.permute_copy.default"), 1) + if _TEST_RUNTIME: + program = Runtime.get().load_program(edge.to_executorch().buffer) + out = program.load_method("forward").execute((x,))[0] + self.assertTrue(torch.allclose(out, model(x), atol=5e-2, rtol=5e-2)) + + # A skipped permute leaves the chain's views with no rank 6 rewrite. + edge = executorch.exir.to_edge_transform_and_lower( + ep, + partitioner=[ + CoreMLPartitioner( + skip_ops_for_coreml_delegation=["aten.permute_copy.default"] + ) + ], + ) + op_names = [ + n.target.__name__ + for n in edge.exported_program().graph.nodes + if n.op == "call_function" + ] + self.assertEqual(op_names.count("aten.permute_copy.default"), 2) + def test_argmax_argmin_dim_none_is_skipped(self): """ Regression test for https://github.com/pytorch/executorch/issues/11715. From 567d1d8687bda5cb173fc8136d522b65e0fc1479 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:25:20 -0700 Subject: [PATCH 11/37] Vulkan: load programs with a constant of more than 4 dims A Vulkan program failed to load when a delegated op used a constant tensor with more than 4 dims, for example x * c with a 5-D buffer c: Tensor.h:790: (sizes_.size() <= 4) is false! The Vulkan backend stores tensors with more than 4 dims in buffers, and the partitioner delegates such ops on purpose. 5-D activations already work. But add_prepack_standard_node, which uploads a constant at load time, still passed the tensor's sizes, strides and numel as push constants for buffer storage. Those 4-D push constants assert on a 5-D tensor, and the nchw_to_buffer shader no longer reads them: it gets everything from the BufferMetadata UBO, the same way add_staging_to_tensor_node already works. Drop the unused push constants for buffer storage. --- backends/vulkan/runtime/graph/ops/impl/Staging.cpp | 7 +------ backends/vulkan/test/test_vulkan_delegate.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp index 37c8b4d4f3f..bc451668fd9 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Staging.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Staging.cpp @@ -190,12 +190,7 @@ void add_prepack_standard_node( } std::vector pcs; - if (graph.is_buffer_storage(tensor)) { - pcs = { - graph.sizes_pc_of(tensor), - graph.strides_pc_of(tensor), - graph.numel_pc_of(tensor)}; - } else if (is_nchw_to_bitw8_shader(shader)) { + if (is_nchw_to_bitw8_shader(shader)) { pcs = {graph.sizes_pc_of(tensor)}; } diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 01c74110da6..68681688457 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -3049,6 +3049,19 @@ def forward(self, x, y): self.lower_module_and_test_output(high_dim_module, sample_inputs) + def test_vulkan_backend_high_dimensional_constant(self): + class HighDimConstantModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("c", torch.rand(1, 2, 3, 4, 5)) + + def forward(self, x): + return x * self.c + + sample_inputs = (torch.rand(size=(1, 2, 3, 4, 5), dtype=torch.float32),) + + self.lower_module_and_test_output(HighDimConstantModule(), sample_inputs) + def test_vulkan_backend_torchao_wo_quantized_linear(self): in_features = 1024 out_features = 512 From dc5795ddbca035ab8cac450dcc2ebea3cc2b7d93 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:35:07 -0700 Subject: [PATCH 12/37] Arm backend: Keep unused delegate inputs in the TOSA graph The TOSA backend skipped every top-level placeholder without users. That is fine for constants, but a delegate input can also lose all its users during the Arm pass pipeline. Two examples: - ConvertFullLikeToFullPass turns zeros_like(y), ones_like(y) and full_like(y, v) into full, so y is no longer used. - RemoveNoopPass folds torch.empty(...).copy_(x) into x, so the empty tensor, which stays on the CPU, is no longer used. The runtime still passes every delegate argument, and the VGF backend binds them by position. With the unused input missing from the graph, inputs after it were read from the wrong slots, and outputs were written one slot early, into the last input tensor. The program ran without an error but returned wrong values (all zeros in the copy example). Keep unused placeholders that are delegate inputs. The model converter can still remove dead inputs, which the VGF runtime already handles. Ethos-U cannot keep them: Vela drops unused inputs too, and the Ethos-U runtime also binds arguments by position. Instead of silently computing on the wrong tensors, the Ethos-U backend now raises a clear error for such a delegate. The new test case expects that error on Ethos-U. Two Ethos-U unit tests that called preprocess with no program now pass a small exported one, since preprocess reads the program's inputs. --- backends/arm/ethosu/backend.py | 16 ++++++++++++ backends/arm/test/misc/test_ethosu_scratch.py | 3 ++- .../test/misc/test_external_vela_blocks.py | 4 ++- backends/arm/test/ops/test_copy.py | 26 ++++++++++++++++--- backends/arm/tosa/backend.py | 12 ++++++--- 5 files changed, 53 insertions(+), 8 deletions(-) diff --git a/backends/arm/ethosu/backend.py b/backends/arm/ethosu/backend.py index a7e71bc9423..9072a28be28 100644 --- a/backends/arm/ethosu/backend.py +++ b/backends/arm/ethosu/backend.py @@ -114,6 +114,22 @@ def preprocess( # which can be passed on to next compilation step. tosa_preprocess = TOSABackend._preprocess(edge_program, tosa_compile_spec) + unused_inputs = [ + node.name + for node in edge_program.graph_module.graph.nodes + if node.op == "placeholder" + and not node.users + and node.name in edge_program.graph_signature.user_inputs + ] + if unused_inputs: + raise RuntimeError( + "A delegate input is not used after lowering, for example an " + "input that only gives its shape to zeros_like or full_like, or " + "that copy_ overwrites. Vela removes unused inputs, but the " + "Ethos-U runtime binds delegate arguments by position, so this " + f"delegate would compute on the wrong tensors. Unused: {unused_inputs}" + ) + compile_result = EthosUBackend._compile_tosa_flatbuffer( tosa_preprocess.processed_bytes, compile_spec ) diff --git a/backends/arm/test/misc/test_ethosu_scratch.py b/backends/arm/test/misc/test_ethosu_scratch.py index 75224052073..fbf303d653c 100644 --- a/backends/arm/test/misc/test_ethosu_scratch.py +++ b/backends/arm/test/misc/test_ethosu_scratch.py @@ -51,7 +51,8 @@ def vela_main(args): return_value=PreprocessResult(processed_bytes=b"tosa"), ), ): - return EthosUBackend.preprocess(None, compile_spec._to_list()) + edge_program = torch.export.export(torch.nn.ReLU(), (torch.ones(1),)) + return EthosUBackend.preprocess(edge_program, compile_spec._to_list()) return compile_with_scratch diff --git a/backends/arm/test/misc/test_external_vela_blocks.py b/backends/arm/test/misc/test_external_vela_blocks.py index 3f3ff6b2b72..a299274691f 100644 --- a/backends/arm/test/misc/test_external_vela_blocks.py +++ b/backends/arm/test/misc/test_external_vela_blocks.py @@ -10,6 +10,7 @@ from unittest.mock import patch import numpy as np +import torch from executorch.backends.arm import arm_vela from executorch.backends.arm.arm_vela import ( _BLOCK_ALIGNMENT, @@ -193,7 +194,8 @@ def compile_tosa_flatbuffer( side_effect=compile_tosa_flatbuffer, ), ): - result = EthosUBackend.preprocess(None, compile_spec._to_list()) + edge_program = torch.export.export(torch.nn.ReLU(), (torch.ones(1),)) + result = EthosUBackend.preprocess(edge_program, compile_spec._to_list()) assert result.data_store_output is not None external_data = result.data_store_output.external_data diff --git a/backends/arm/test/ops/test_copy.py b/backends/arm/test/ops/test_copy.py index 556e952a2ba..ca02901ae70 100644 --- a/backends/arm/test/ops/test_copy.py +++ b/backends/arm/test/ops/test_copy.py @@ -1,4 +1,4 @@ -# Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -55,6 +55,11 @@ def forward(self, x): return x * 2 + y.copy_(x) +class CopyOverUnusedInput(torch.nn.Module): + def forward(self, x, y): + return y.clone().copy_(x) + x + + test_suite = { "copy_output": lambda: ( CopyOutput, @@ -80,9 +85,24 @@ def forward(self, x): CopyParallelToOtherOp, (torch.rand(1, 2, 3, 4, dtype=torch.float32),), ), + "copy_over_unused_input": lambda: ( + CopyOverUnusedInput, + ( + torch.rand(1, 2, 3, 4, dtype=torch.float32), + torch.rand(1, 2, 3, 4, dtype=torch.float32), + ), + ), } +ethos_u_xfails = { + "copy_over_unused_input": ( + "Vela drops unused inputs, but the Ethos-U runtime binds delegate " + "inputs by position, so lowering rejects them", + RuntimeError, + ), +} + aten_op = "torch.ops.aten.copy_.default" exir_op = "executorch_exir_dialects_edge__ops_aten_copy_default" @@ -114,7 +134,7 @@ def test_copy_tosa_INT(input_data): pipeline.run() -@common.parametrize("input_data", test_suite) +@common.parametrize("input_data", test_suite, xfails=ethos_u_xfails) @common.XfailIfNoCorstone300 def test_copy_u55_INT(input_data): module, input_tensor = input_data() @@ -128,7 +148,7 @@ def test_copy_u55_INT(input_data): pipeline.run() -@common.parametrize("input_data", test_suite) +@common.parametrize("input_data", test_suite, xfails=ethos_u_xfails) @common.XfailIfNoCorstone320 def test_copy_u85_INT(input_data): module, input_tensor = input_data() diff --git a/backends/arm/tosa/backend.py b/backends/arm/tosa/backend.py index 2a72acc20cc..a395607a79a 100644 --- a/backends/arm/tosa/backend.py +++ b/backends/arm/tosa/backend.py @@ -341,9 +341,15 @@ def _annotate_control_flow_region_names( if node.op == "call_function": process_call_function(node, tosa_graph, node_visitors, tosa_spec) elif node.op == "placeholder": - if len(node.users) == 0 and submodule_name is None: - # In top level module, we don't need to handle unused placeholders. - # In submodules, we do need to handle them to preserve call signature. + if ( + len(node.users) == 0 + and submodule_name is None + and node.name not in edge_program.graph_signature.user_inputs + ): + # Unused parameters, buffers and constants are skipped. + # Unused user inputs stay because the runtime binds + # delegate arguments by position. In submodules, all + # placeholders are kept to preserve the call signature. continue process_placeholder( node, From 0afe355aa18ee9e976a6396864efda7730ddacb6 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:27:52 -0700 Subject: [PATCH 13/37] NXP backend: quantize programs from a non-strict torch.export NeutronQuantizer did not work on programs from a non-strict torch.export (the default of torch.export.export), because two places that read the source_fn_stack node metadata assumed it is there, and a non-strict export does not record it: 1. get_aten_node_target_partitions skipped every node without it. The same conv, relu and linear model got 10 quantize and dequantize nodes with strict=True and 2 with strict=False, so almost nothing was delegated. The function only uses the metadata to group partitions, and its caller flattens the groups anyway. 2. MoveActivationBeforeConcat, which runs before annotation, raised KeyError: 'source_fn_stack' on a conv, cat, relu pattern. Group by the node's target when the metadata is missing, and copy the metadata to the new activation nodes only when it is there. The new test quantizes a non-strict conv, cat, relu model, which goes through both. --- .../move_activation_before_concat.py | 11 +++---- backends/nxp/quantizer/utils.py | 17 +++++----- .../nxp/tests/generic_tests/test_quantizer.py | 31 +++++++++++++++++++ 3 files changed, 43 insertions(+), 16 deletions(-) diff --git a/backends/nxp/aten_passes/move_activation_before_concat.py b/backends/nxp/aten_passes/move_activation_before_concat.py index 4ae288fbc1e..07a34995fab 100644 --- a/backends/nxp/aten_passes/move_activation_before_concat.py +++ b/backends/nxp/aten_passes/move_activation_before_concat.py @@ -80,12 +80,11 @@ def _is_concat(node_: Node) -> bool: kwargs=activation.kwargs, ) - new_activation.meta["source_fn_stack"] = [ - ( - new_activation.name, - activation.meta["source_fn_stack"][-1][-1], - ) - ] + # A non-strict torch.export does not record source_fn_stack. + if source_fn_stack := activation.meta.get("source_fn_stack"): + new_activation.meta["source_fn_stack"] = [ + (new_activation.name, source_fn_stack[-1][-1]) + ] new_activation.meta["val"] = input_node.meta["val"] # Replace the uses of the input node with the new activation node. diff --git a/backends/nxp/quantizer/utils.py b/backends/nxp/quantizer/utils.py index 4178c5c17fb..3d7779cbb7e 100644 --- a/backends/nxp/quantizer/utils.py +++ b/backends/nxp/quantizer/utils.py @@ -124,19 +124,16 @@ def get_aten_node_target_partitions( modules: Dict[Type, Dict[str, list[torch.fx.Node]]] = {} for node in graph.nodes: - # The metadata source_fn should contain a tuple of a unique name for the - # source, and the source function if the node is decomposed from a - # function, or the type of module if the node is decomposed from a leaf - # module - # TODO(matthiascremon): look into ways to avoid using source_fn_stack - if (source_fn_st := node.meta.get("source_fn_stack")) is None: - continue - - source_fn = source_fn_st[-1] if node.target not in wanted_original_aten_op: continue - diff_modules = modules.setdefault(source_fn[1], {}) + # The metadata source_fn_stack names the function or leaf module the + # node came from. A non-strict torch.export does not record it, so fall + # back to the node's target. + source_fn_st = node.meta.get("source_fn_stack") + source = source_fn_st[-1][1] if source_fn_st else node.target + + diff_modules = modules.setdefault(source, {}) partition = diff_modules.setdefault(node.name, []) partition.append(node) diff --git a/backends/nxp/tests/generic_tests/test_quantizer.py b/backends/nxp/tests/generic_tests/test_quantizer.py index 886301fb149..1eb7cfd9d16 100644 --- a/backends/nxp/tests/generic_tests/test_quantizer.py +++ b/backends/nxp/tests/generic_tests/test_quantizer.py @@ -147,6 +147,37 @@ def test_quantizer_conv2d(): assert nodes[12].args[0].target == torch.ops.aten.conv2d.default +def test_quantizer_non_strict_export(): + class ConvConcatRelu(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x): + return torch.relu(torch.cat((self.conv(x), self.conv(x)), dim=1)) + + example_input = (torch.ones(1, 4, 8, 8),) + exported_model = torch.export.export( + ConvConcatRelu().eval(), example_input, strict=False + ) + + m = _prepare_for_quantization(exported_model) + m(*example_input) + m = convert_pt2e(m) + + convs = [n for n in m.graph.nodes if n.target == torch.ops.aten.conv2d.default] + assert len(convs) == 2 + for conv in convs: + assert ( + conv.args[0].target + == torch.ops.quantized_decomposed.dequantize_per_tensor.default + ) + assert ( + conv.args[1].target + == torch.ops.quantized_decomposed.dequantize_per_channel.default + ) + + def test_quantizer_linear(): model = models.LinearModule(bias=True) model.eval() From 6806f0eea294e974e1d463ef5feb2f174cc66e63 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 06:43:34 -0700 Subject: [PATCH 14/37] Apply the annotation filter to cat in XNNPACKQuantizer XNNPACKQuantizer passes a filter to every annotator, so that module type, module name and other scoped configs quantize only the nodes in their scope. The cat annotator ignored it and annotated every cat in the graph. With only set_module_type(nn.Linear, ...), a cat after the linear still got observers. Apply the filter in the cat annotator like in the others. Cats outside a scoped config are now left unquantized; a global config quantizes them as before. --- .../quantizer/xnnpack_quantizer_utils.py | 2 ++ .../test/quantizer/test_xnnpack_quantizer.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index a207d77674d..757b9a6c452 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1065,6 +1065,8 @@ def _annotate_cat( for cat_node in gm.graph.nodes: if cat_node.target != torch.ops.aten.cat.default: continue + if filter_fn and not filter_fn(cat_node): + continue if _is_annotated([cat_node]): continue diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index b66166fedb6..8b37c98c60b 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -566,6 +566,28 @@ def forward(self, x): ] self._test_quantizer(m, example_inputs, quantizer, node_occurrence, node_list) + def test_set_module_type_skips_cat_outside_the_module(self): + class M(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = torch.nn.Linear(5, 5) + + def forward(self, x): + return torch.cat([self.linear(x), x], 1) + + quantizer = XNNPACKQuantizer().set_module_type( + torch.nn.Linear, get_symmetric_quantization_config() + ) + m = export(M().eval(), (torch.randn(3, 5),), strict=True).module() + m = quantizer.annotate(m) + annotated = { + node.target + for node in m.graph.nodes + if node.meta.get("quantization_annotation") is not None + } + self.assertIn(torch.ops.aten.linear.default, annotated) + self.assertNotIn(torch.ops.aten.cat.default, annotated) + def test_set_module_type_case_2(self): class M(torch.nn.Module): def __init__(self) -> None: From 1c049b8268921d32630d036dca3d5eac97bcc35e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:28:31 -0700 Subject: [PATCH 15/37] Do not quantize bfloat16 tensors in XNNPACKQuantizer XNNPACKQuantizer annotated ops that read or produce bfloat16 tensors, such as linears in a model that casts its activations to bfloat16. The quantize, dequantize and choose_qparams nodes it produced then ran on bfloat16 tensors. XNNPACK does not take those patterns, so they stayed on the portable quantized kernels, which do not take bfloat16 and abort the process at run time ("Expected input to be Float tensor received: 15"). Skip nodes that read or produce a bfloat16 tensor, both in the filter that every annotator applies (cat included, since the previous commit) and when annotations are propagated to ops that share their input's quantization parameters. Such ops stay unquantized and run correctly. Float32 and float16 models quantize as before; the test checks that a float32 linear feeding a bfloat16 op is still quantized. --- .../xnnpack/quantizer/xnnpack_quantizer.py | 5 +- .../quantizer/xnnpack_quantizer_utils.py | 13 +++++ .../test/quantizer/test_xnnpack_quantizer.py | 49 +++++++++++++++++++ 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer.py b/backends/xnnpack/quantizer/xnnpack_quantizer.py index 3c82a65ad71..e1c37e961c6 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer.py @@ -13,6 +13,7 @@ _convert_scalars_to_attrs, OP_TO_ANNOTATOR, propagate_annotation, + touches_bfloat16, ) from torchao.quantization.pt2e import ( FakeQuantize, @@ -399,9 +400,9 @@ def _annotate_all_patterns( if quantization_config is None: return model - # Create a combined filter function, which returns True only when - # both filter_fn and self.filter_fn return True. def combined_filter_fn(n: Node) -> bool: + if touches_bfloat16(n): + return False combined_filter = [self.filter_fn, filter_fn] return all(f(n) for f in combined_filter if f is not None) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index 757b9a6c452..4e5b6228f55 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1125,10 +1125,23 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool: ] +def touches_bfloat16(node: Node) -> bool: + """Whether a node reads or produces a bfloat16 tensor. The quantized kernels + do not take bfloat16 and XNNPACK leaves such patterns to them, where they + abort at run time, so these nodes are not quantized.""" + return any( + isinstance(n.meta.get("val"), torch.Tensor) + and n.meta["val"].dtype == torch.bfloat16 + for n in (node, *node.all_input_nodes) + ) + + def propagate_annotation(model: torch.fx.GraphModule) -> None: for n in model.graph.nodes: if n.op != "call_function" or not _is_share_obs_or_fq_op(n.target): continue + if touches_bfloat16(n): + continue prev_node = n.args[0] if not isinstance(prev_node, Node): diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 8b37c98c60b..f3322bc83ad 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -734,6 +734,55 @@ def test_dynamic_linear(self): qconfig_mapping, ) + def test_bf16_is_not_quantized(self): + class Cat(torch.nn.Module): + def forward(self, x): + return torch.cat([x, x.float()], 1) + + class LinearMean(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(8, 8) + + def forward(self, x): + return torch.mean(self.linear(x), -1, dtype=torch.bfloat16) + + linear = torch.nn.Linear(8, 8).to(torch.bfloat16) + conv = torch.nn.Conv2d(3, 3, 1).to(torch.bfloat16) + x2d = torch.randn(9, 8, dtype=torch.bfloat16) + x4d = torch.randn(1, 3, 4, 4, dtype=torch.bfloat16) + cases = [ + (linear, x2d), + (torch.nn.Sequential(linear, torch.nn.ReLU()), x2d), + (conv, x4d), + # The cat output is float32 but one of its inputs is bfloat16. + (Cat(), x4d), + # The linear is float32; only the mean it feeds is bfloat16. + (LinearMean(), x2d.float()), + ] + for (m_eager, x), is_dynamic in itertools.product(cases, (False, True)): + quantizer = XNNPACKQuantizer().set_global( + get_symmetric_quantization_config( + is_per_channel=True, is_dynamic=is_dynamic + ) + ) + m = export(m_eager.eval(), (x,), strict=True).module() + m = prepare_pt2e(m, quantizer) + m(x) + m = convert_pt2e(m) + quantized = [ + node.args[0].meta["val"].dtype + for node in m.graph.nodes + if node.target + in ( + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, + ) + ] + self.assertNotIn(torch.bfloat16, quantized) + if isinstance(m_eager, LinearMean): + self.assertTrue(quantized) + def test_dynamic_linear_int4_weight(self): quantizer = XNNPACKQuantizer() quantization_config = get_symmetric_quantization_config( From 8aad730a502316aed75cfbc816d8bf5bf066a4e7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:15:43 -0700 Subject: [PATCH 16/37] Keep the output node last when to_backend fuses a partition to_backend could fail with "Expected output node, got op=call_function" when a runtime check outside a partition reads a value the partition computes, for example torch._check(t.all().item()) where the tensor ops are delegated and the item() and the check are not. After fusing the partition, create_submodule_from_nodes re-sorted the graph with legalize_graph. That sort places a node as soon as its inputs are ready, so the output node could be placed before check nodes that have no users. Later code calls graph.output_node(), which requires the output node to be last. Use stable_topological_sort instead. It keeps the original order wherever the dependencies allow it, so the output node stays last. This changes the order of the remaining top-level nodes of every program that to_backend touches: they now keep their original order instead of the order the old sort produced. --- exir/lowered_backend_module.py | 6 +++--- exir/tests/targets.bzl | 1 + exir/tests/test_delegate.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index 8a381c2b0f0..ba8f0d6fdb8 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -43,11 +43,11 @@ OutputSpec, TensorArgument, ) +from torch.fx.passes.tools_common import stable_topological_sort from torch.fx.passes.utils.fuser_utils import ( erase_nodes, fuse_as_graphmodule, insert_subgm, - legalize_graph, NodeList, topo_sort, ) @@ -931,10 +931,10 @@ def create_submodule_from_nodes( erase_nodes(gm, sorted_nodes) # Topological sort original gm with newly created sub_gm - # TODO : T153794167 Get rid of support for skipping legalize graph in create_submodule_from_nodes + # TODO : T153794167 Get rid of support for skipping the sort in create_submodule_from_nodes # once we transition to using fuse_by_partitions. if not skip_legalize_graph: - legalize_graph(gm) + stable_topological_sort(gm) # Get the call_module node submodule_node = None diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index 0bd79bfa7dc..59ae7b904b2 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -307,6 +307,7 @@ def define_common_targets(is_fbcode = False): "//executorch/exir:lowered_backend_module", "//executorch/exir:schema", "//executorch/exir/dialects:lib", + "//executorch/exir/dialects/edge:lib", ], ) diff --git a/exir/tests/test_delegate.py b/exir/tests/test_delegate.py index d89d3f2bbd8..915c832a591 100644 --- a/exir/tests/test_delegate.py +++ b/exir/tests/test_delegate.py @@ -11,6 +11,7 @@ import torch from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.lowered_backend_module import ( create_submodule_from_nodes, LoweredBackendModule, @@ -246,6 +247,33 @@ def forward(self, x, y): new_res = prog.exported_program().module()(*inputs) self.assertTrue(torch.allclose(new_res, orig_res)) + def test_create_submodule_keeps_output_last(self) -> None: + class Model(torch.nn.Module): + def forward(self, x): + y = x * 2 + torch._check((y > -1e9).all().item()) + return y + + inputs = (torch.randn(4),) + prog = to_edge( + export(Model(), inputs), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + gm = prog.exported_program().graph_module + + # The tensor ops are delegated; item() and the check on it are not. + node_list = [ + node + for node in gm.graph.nodes + if node.op == "call_function" and isinstance(node.target, EdgeOpOverload) + ] + create_submodule_from_nodes(gm, node_list, "tag") + + self.assertEqual(list(gm.graph.nodes)[-1].op, "output") + self.assertTrue( + torch.allclose(prog.exported_program().module()(*inputs), inputs[0] * 2) + ) + def test_create_submodule_list_return(self) -> None: """ Original graph: From ff75c247998de2479a782383c62f6a7b8abbb163 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 18:42:20 -0700 Subject: [PATCH 17/37] Lower torch._check_tensor_all and torch.sym_ite through to_edge A model that uses a tensor-wide check, for example torch._check_tensor_all or transformers' torch_compilable_check, could not go through to_edge. Two things stopped it: 1. torch.export turns the check into aten._is_all_true, which is not a core ATen op and has no decomposition, so the ATen verifier rejected it. 2. The check then reads the result with item() and compares it through torch.sym_ite. The Edge verifier rejected every Python symbolic function (torch.sym_ite, sym_max, sym_min, sym_not, sym_float, sym_int), on data-dependent values and on shapes alike, even though ExecuTorch lowers each of them to its own executorch_prim op, which has a runtime kernel. Like other runtime asserts, the check is removed later by to_executorch unless a backend takes it into a delegate; the problem was that the program never got that far. Decompose aten._is_all_true to aten.all in the default decomposition table, and let the Edge verifier accept the Python symbolic functions that ExecuTorch maps to prim ops, the same way it already accepts those prim ops. The runtime's sym_max and sym_min kernels only take ints, so the verifier still rejects them, also inside control flow branches, when an argument is a float or a bool. --- exir/emit/test/test_emit.py | 5 +--- exir/program/test/test_program.py | 38 ++++++++++++++++++++++++++++++- exir/tracer.py | 2 ++ exir/verification/verifier.py | 27 +++++++++++++++++++++- 4 files changed, 66 insertions(+), 6 deletions(-) diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index e36f075a7b0..8e465561b50 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -2537,10 +2537,7 @@ def forward(self, x): model, (test_inputs[0],), dynamic_shapes=dynamic_shapes ) - edge_program = to_edge( - exported_program, - compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), - ) + edge_program = to_edge(exported_program) et_program = edge_program.to_executorch() program_buffer = et_program.buffer et_module = _load_for_executorch_from_buffer(program_buffer) diff --git a/exir/program/test/test_program.py b/exir/program/test/test_program.py index 4f6ef27b3a7..45b6ed0b1c2 100644 --- a/exir/program/test/test_program.py +++ b/exir/program/test/test_program.py @@ -34,7 +34,7 @@ from executorch.extension.pybindings.portable_lib import ( _load_for_executorch_from_buffer, ) -from torch._export.verifier import Verifier +from torch._export.verifier import SpecViolationError, Verifier from torch.export import Dim, export, ExportedProgram from torch.export._trace import _export from torch.fx.passes.infra.pass_manager import PassManager @@ -337,6 +337,42 @@ def forward(self, x, y): ) edge_manager.to_executorch() + def test_check_tensor_all(self): + class M(torch.nn.Module): + def forward(self, x): + torch._check_tensor_all(x.abs() < 1e9) + return x * 2 + + to_edge(export(M(), (torch.randn(4),))).to_executorch() + + def test_float_sym_max_min_rejected(self): + # The runtime's sym_max and sym_min kernels only take ints. + def model(sym_op): + class M(torch.nn.Module): + def forward(self, x, y): + return x + sym_op(y.item(), 2) + + return M() + + for sym_op in (torch.sym_max, torch.sym_min): + to_edge(export(model(sym_op), (torch.ones(1), torch.tensor(3)))) + with self.assertRaisesRegex(SpecViolationError, "only supports int"): + to_edge(export(model(sym_op), (torch.ones(1), torch.tensor(1.25)))) + + class Cond(torch.nn.Module): + def forward(self, pred, x, y): + return torch.cond( + pred, + lambda x, y: x + torch.sym_max(y.item(), 2), + lambda x, y: x - y, + (x, y), + ) + + with self.assertRaisesRegex(SpecViolationError, "only supports int"): + to_edge( + export(Cond(), (torch.tensor(True), torch.ones(1), torch.tensor(1.25))) + ) + def test_data_dependent(self): with torch.library._scoped_library("mylib", "FRAGMENT") as lib: torch.library.define( diff --git a/exir/tracer.py b/exir/tracer.py index 60a24c507d3..822abc3c021 100644 --- a/exir/tracer.py +++ b/exir/tracer.py @@ -664,6 +664,8 @@ def _default_decomposition_table( # A program already decomposed with PyTorch's table has aten.empty_permuted # where it had aten.empty_like. decomps[torch.ops.aten.empty_permuted.default] = _empty_permuted + # torch._check_tensor_all exports to this op, which has no decomposition in core. + decomps[torch.ops.aten._is_all_true.default] = torch.ops.aten.all.default # pyre-fixme[7]: Expected `Dict[OpOverload, typing.Callable[..., executorch.exir.... never_decompose = [] diff --git a/exir/verification/verifier.py b/exir/verification/verifier.py index e15c907de18..183838acb0e 100644 --- a/exir/verification/verifier.py +++ b/exir/verification/verifier.py @@ -19,7 +19,10 @@ from executorch.exir.error import ExportError, ExportErrorType from executorch.exir.lowered_backend_module import LoweredBackendModule from executorch.exir.passes.dim_order_ops_registry import DimOrderOpsMap -from executorch.exir.passes.executorch_prim_ops_registry import _EXECUTORCH_SYM_OPS +from executorch.exir.passes.executorch_prim_ops_registry import ( + _EXECUTORCH_SYM_OPS, + _PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS, +) from executorch.exir.passes.replace_aten_with_edge_pass import DISALLOW_LIST from executorch.exir.verification.arg_validator import ( EdgeOpArgValidator, @@ -49,6 +52,26 @@ def _check_tensors_are_contiguous(gm: GraphModule) -> None: ) +def _check_sym_max_min_args_are_ints(gm: GraphModule) -> None: + # The executorch_prim sym_max and sym_min kernels only take ints. + for module in gm.modules(): + if not isinstance(module, GraphModule): + continue + for node in module.graph.nodes: + if node.op != "call_function" or node.target not in ( + torch.sym_max, + torch.sym_min, + ): + continue + for arg in node.args: + val = arg.meta.get("val") if isinstance(arg, torch.fx.Node) else arg + if isinstance(val, bool) or not isinstance(val, (int, torch.SymInt)): + raise SpecViolationError( + f"{node.target.__name__} only supports int arguments, but " + f"{node.name} has {val!r}" + ) + + def _check_valid_dim_order_ops(op, use_dim_order) -> None: if use_dim_order: if op in DimOrderOpsMap: @@ -294,6 +317,7 @@ def check_valid_edge_op(self, op): in [operator.getitem] + DISALLOW_LIST + list(_EXECUTORCH_SYM_OPS) + + list(_PYTHON_SYM_OPS_TO_EXECUTORCH_SYM_OPS) + self._core_aten_ops_exception_list ): return @@ -315,6 +339,7 @@ def check_additional(self, gm: GraphModule) -> None: if not self.enable: return _check_tensors_are_contiguous(gm) + _check_sym_max_min_args_are_ints(gm) _check_tensor_args_matching_op_allowed_dtype(gm) def is_valid(self, gm: GraphModule) -> bool: From d23bc28c98c1462d1cd1fd666d8c3ec3a9f1713b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:19:19 -0700 Subject: [PATCH 18/37] Add the aoti_torch_as_strided shim to the CUDA backend A program lowered with the CUDA backend could pass lowering and then kill the process on its first run: undefined symbol: aoti_torch_as_strided Inductor's C++ wrapper calls aoti_torch_as_strided when it allocates a buffer with a larger shape than its logical shape and then gives it the logical shape. ExecuTorch's CUDA shim library did not provide it, so the compiled model could not be loaded. Add aoti_torch_as_strided with the PyTorch shim's signature. It returns a view with the given sizes and strides and the tensor's own storage offset, built on the existing aoti_torch__reinterpret_tensor. Also regenerate the checked-in Windows import library, backends/cuda/runtime/aoti_cuda_shims.lib, as its README describes, so a wrapper lowered for a Windows target can link against the new shim. Its symbol list is the old one plus aoti_torch_as_strided, and its members carry no timestamps. --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 43802 -> 44602 bytes backends/cuda/runtime/shims/memory.cpp | 16 +++++++ backends/cuda/runtime/shims/memory.h | 17 +++++++ .../test_aoti_torch__reinterpret_tensor.cpp | 45 ++++++++++++++++++ 4 files changed, 78 insertions(+) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index c0d61c611008c3897350d7b0950a8e0229607ba3..9f48cd086077a24be2d6cff506f6c665f6ba9b88 100644 GIT binary patch delta 1598 zcmYk*drZ}J6aet=`2|evt;6`SYmE>8C%*Sl4Ra9nMX&Y;_R&5QFtp4b^&X+HLbhpp$?)jbX_j}Iw{BF-z z+SO65E49gEc%mmiTLOg;RYEw&ga~;}3F{U@irAn;FBU?Y{hKo94I!l19%Vt15RyAl zNqAfcDJfJ*2@*nD?5CuSHwlq8F2vGLlng#t%IZ*7?i4~=&6`TOhLUd+LMqaflAnZ- z*1xZ8@(Cf8)+yy_LP(XoM%p&2R5iWIO|xAMO7$`!q*_jr>c^CZ4~38#uPTSyg^-%P z%G10?dX@{4+BjG0=u}Q{@ltn^a++^JT732!A*7xUl-?mBBp(+fy_l|CvnFBou#?*Lm$mXJIo5N(PH4Gy; z0r%QMsl{3Bhd08lXw{tbM`#UDTU8-oU!?W_JH8ot9fx-6T3M|;OMcxRh{5m(gZ@e7 z>~*uOG!mW%^kq~FfFJcorT})$*6GrmG9d4)`G5!J=yY>7=MK-aq2-Za^2B5TCdV4I zI)bBT;%sD@%X!WPI#tGUq+NI7P`p3ho^PVN^PT_WcsZ@uU$$m-QzBe_5zgY(f#>Q;TfaX9SZkk(d48(J3mIcg|a4$00u_^(6E9&r)@I zaEhZwx*ZRs22uN>OrY1(S^$HVyD-AvOw~)$fjXBS0i4KmVJS!Q(mAp_gCn13+c7%B zM#gfkFehv3K89D>am?RFK`YV$JuB4-V5d9Tm4Nou2Ayv#0C3PdtJd?TIa3{b3!FG& zGUMHO8EBmrg4g2i^8*XsH>Kcsd^nC28mji#bu+NH&_S&QD*!K+ z81#1`?;k>Ii%S5T*BkVFBX@Xjvkf;DTX1ZeA7!sg11#C3(^?-t@ojNX$i}6Bk)^tt zf6Z10>H+>Vu{j>_K^YfP&3Vhb4m=zXK<{oz2D(wf{oA$&Q|nd_P{%g@HM%31O1vJx z4OKj=a=!EnyB!!g-9&S@Cj!pel|b!PTv-$@_qyoe4uiYj!&50Q1c=7|J)v}c_d>w; zYB(m(i`!?#Pp4a{V{bCxEA>;|M(&qU7kTUW4eD)h)7SN!`}2XBxGc~~`TKaPyBggz z(!kM=o9txl=U!JHQn!v5j@YjpTm|@8vzz`Z^h`W1qx_g$78||~=xFwi!&Se8`?A6uv8VJ@9iaW3KmHSj- zT7tl}Ru}!+$94Ta;GUd2Bx3JCINf-GC-%il>Mhh33cO{Fpwkz))MJD3S1$5@D~XC< R_JFo18SkVnpy9!C?LX0A2WJ2P delta 1487 zcmXZbeN2^A7y$5dZkHDW?#t!97Yh))5KJfVzAwlXP~=Nc;-f%5)X5Y@G>}h;#b6@H zN0^AnW1$R|ff<6Vwx?PnCjQad%5weD92-Wh`b*MSZApz)&*5;l-`(ElzUMj5^PF>M z9$AJTS_Tr&x#JRCYo9VN#@K6&IeQs*cvBAYK#fG zBAF|U3EP*Y9ltUrBz_|$-8sit(jCT9`lR$~#)QmVDT@*dId&=UDPuzZ&r;#nj0q(- zr9MHA*|;R>8}|iIcnbaPZh^%Kok< zXwaUb;y`jT$EQiuQ|~Sf%FhTR6SKG&^f$Y(=MB0_zclIicQZ8`3Ta;9*w8{pd%s!7rB;fCVanM&4zsQ6 zBhPK|!_Dp494@sPxak0Gn|xLSSGzLU<$<1^G?mH@9UHo6XSH3!husE;JzYA^^-<*U zc@3}i#&Wd2Nj-hh?FY4e3YK5gIo`j}LDTK^!)k z@oQTHxspFRa0ci>49{>OV diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 84d12b77732..86be88e4833 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -387,6 +387,22 @@ AOTITorchError aoti_torch__reinterpret_tensor( return Error::Ok; } +AOTITorchError aoti_torch_as_strided( + SlimTensor* self, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + SlimTensor** ret_new_tensor) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr, InvalidArgument, "aoti_torch_as_strided: self is null"); + return aoti_torch__reinterpret_tensor( + self, + static_cast(self->dim()), + sizes_ptr, + strides_ptr, + self->storage_offset(), + ret_new_tensor); +} + AOTITorchError aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) { (void)non_blocking; // SlimTensor::copy_() is always synchronous for now diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index 942d14bc679..b1b369a58d9 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -180,6 +180,23 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch__reinterpret_tensor( int64_t storage_offset, SlimTensor** ret_new_tensor); +/** + * Creates a view of a tensor with new sizes and strides and the same storage + * offset, like torch.as_strided(self, sizes, strides). The view has as many + * dimensions as self. + * + * @param self Original tensor (must not be null) + * @param sizes_ptr Pointer to array of dimension sizes + * @param strides_ptr Pointer to array of strides for each dimension + * @param ret_new_tensor Output parameter for the tensor view + * @return AOTITorchError error code (Error::Ok on success) + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_as_strided( + SlimTensor* self, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + SlimTensor** ret_new_tensor); + /** * Copies data from source tensor to destination tensor. * diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch__reinterpret_tensor.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch__reinterpret_tensor.cpp index d15b5106e0f..0be919b88e2 100644 --- a/backends/cuda/runtime/shims/tests/test_aoti_torch__reinterpret_tensor.cpp +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch__reinterpret_tensor.cpp @@ -692,3 +692,48 @@ TEST_F(AOTITorchReinterpretTensorSlimTest, ChainedViews_CUDA) { EXPECT_EQ(aoti_torch_delete_tensor_object(view1), Error::Ok); EXPECT_EQ(aoti_torch_delete_tensor_object(view2), Error::Ok); } + +// Inductor allocates some buffers with a larger shape and then calls +// aoti_torch_as_strided to give them their logical shape. +TEST_F(AOTITorchReinterpretTensorSlimTest, AsStridedKeepsStorageOffset_CUDA) { + if (!isCudaAvailable()) { + GTEST_SKIP() << "CUDA not available"; + } + + std::vector sizes = {4, 4}; + Tensor* orig_tensor = createTestTensor( + sizes, + {}, + static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::DeviceType::CUDA), + 0); + ASSERT_NE(orig_tensor, nullptr); + + // Rows 1 to 3 of the 4x4 tensor. + std::vector offset_sizes = {3, 4}; + std::vector strides = {4, 1}; + Tensor* offset_view = nullptr; + ASSERT_EQ( + aoti_torch__reinterpret_tensor( + orig_tensor, 2, offset_sizes.data(), strides.data(), 4, &offset_view), + Error::Ok); + + std::vector new_sizes = {2, 3}; + Tensor* view_tensor = nullptr; + EXPECT_EQ( + aoti_torch_as_strided( + offset_view, new_sizes.data(), strides.data(), &view_tensor), + Error::Ok); + ASSERT_NE(view_tensor, nullptr); + + EXPECT_EQ(view_tensor->dim(), 2); + EXPECT_EQ(view_tensor->size(0), 2); + EXPECT_EQ(view_tensor->size(1), 3); + EXPECT_EQ(view_tensor->stride(0), 4); + EXPECT_EQ(view_tensor->stride(1), 1); + EXPECT_EQ(view_tensor->data_ptr(), offset_view->data_ptr()); + + EXPECT_EQ(aoti_torch_delete_tensor_object(orig_tensor), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(offset_view), Error::Ok); + EXPECT_EQ(aoti_torch_delete_tensor_object(view_tensor), Error::Ok); +} From 7b1507a43d09ede5de136268eebe78ec2a080405 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 11:33:24 -0700 Subject: [PATCH 19/37] Stop GroupBasedPartitioner from making partitions that depend on each other Lowering a dynamically quantized pi0 model to XNNPACK failed with "Invalid partition, found dependency cycles". Before merging two partitions, GroupBasedPartitioner checks that no node outside them is both downstream and upstream of the merged partition. That check looks at single nodes. But every other partition also becomes a single node once it is fused. So a path can leave the merged partition, enter another partition at one node, leave it from a different node that the first one does not reach, and come back. The two partitions then depend on each other, and fusing them fails. Before accepting a merge, also follow paths through the other partitions: whenever the merged partition reaches a node of another partition, add everything that partition reaches as a whole, and reject the merge as soon as that reaches something upstream of it. Only partitions that will be returned count, so partitions that are later filtered out (for example ones without enough compute nodes) do not block a merge. Their compute node count now comes from one helper shared with the filter. To stay fast on large graphs, the check only visits the partitions it actually reaches, and caches what each partition reaches; a merge only drops the cached entries that reach the merged nodes. On synthetic graphs of 4k to 16k nodes the partitioner takes about twice as long as before. The merge loop itself is unchanged and still tries every remaining pair, so a graph where the check rejects most merges now keeps many partitions and pays for all their pairs (about 2 s for 1600 partitions on a 4k-node chain of groups); before, such graphs had few partitions only because merges that cannot be fused were accepted. If caller-provided node groups already depend on each other, what a partition reaches is found by walking around that cycle directly. Two partitions whose nodes all already lie on one cycle through such groups can still be merged, since every cycle through the result already existed. A partition that will be fused counts as one node there; any other partition counts as its separate nodes, so a group that is later filtered out (for example an add with a getitem) can still merge to remove a cycle. Other merges that would leave the merged partition on a cycle are refused. Because a merge that would need a later merge to become valid is now rejected, a few graphs end up with a different number of partitions. The merge loop is still a single greedy pass, so caller node_groups that depend on each other through a plain node can still give partitions that cannot be fused, as before this change. --- .../group_partitioner.py | 236 ++++++++++++++++-- exir/backend/test/test_group_partitioner.py | 222 +++++++++++++++- 2 files changed, 431 insertions(+), 27 deletions(-) diff --git a/exir/backend/canonical_partitioners/group_partitioner.py b/exir/backend/canonical_partitioners/group_partitioner.py index 0ddb7670d41..2a458eae37d 100644 --- a/exir/backend/canonical_partitioners/group_partitioner.py +++ b/exir/backend/canonical_partitioners/group_partitioner.py @@ -35,7 +35,7 @@ class _DependencyViewer: """ def __init__(self, graph_module: GraphModule, nodes): - all_nodes = list(graph_module.graph.nodes) + all_nodes = self.nodes = list(graph_module.graph.nodes) self.positions = {node: i for i, node in enumerate(all_nodes)} self.downstreams = {} self.upstreams = {} @@ -127,12 +127,185 @@ def __init__( self.node_to_group[node] = i self.all_nodes_in_groups.add(node) - def _can_merge_partitions(self, masks1, masks2): - """Check if merging two partitions would create a cycle.""" - nodes, downstreams, upstreams = (m1 | m2 for m1, m2 in zip(masks1, masks2)) - # Fused on its own, the merged partition has a cycle exactly when a node - # outside it is downstream of it and also upstream of it. - return not (downstreams & upstreams & ~nodes) + def _compute_weight(self, node): + """How much a node counts toward keeping its partition from being filtered.""" + if node.op != "call_function": + return 0 + assert callable(node.target) + target_name = _get_qualified_name(node.target) + default_non_compute_ops = {"torch.ops.aten.view", "_operator.getitem"} + non_compute_ops = default_non_compute_ops.union(set(self.non_compute_ops or [])) + weight = 0 if target_name in non_compute_ops else 1 + if ( + self.allowed_single_node_partition_ops + and target_name in self.allowed_single_node_partition_ops + ): + weight += 1 + return weight + + def _is_kept(self, partition, compute_weight): + """Whether a partition will be fused into one node, if it does not grow any + further. A single node needs no fusing: it keeps the same edges either way.""" + if self.allows_single_node_partition: + return len(partition.nodes) > 1 + return compute_weight > 1 + + def _closure( + self, partition_id, partition_masks, assignment, nodes, kept_nodes, closures + ): + """Bitmask of a kept partition and of every node it reaches once each kept + partition is fused into a single node. Cached in closures.""" + stack = [partition_id] + while stack: + current = stack[-1] + if current in closures: + stack.pop() + continue + current_nodes, reach, _ = partition_masks[current] + covered = current_nodes + missing = None + while todo := reach & kept_nodes & ~covered: + position = (todo & -todo).bit_length() - 1 + other = assignment[nodes[position]] + if other not in closures: + missing = other + break + reach |= closures[other] + covered |= closures[other] + if missing is None: + closures[current] = current_nodes | reach + stack.pop() + elif missing in stack: + # Caller node_groups that depend on each other: the partitions + # on the stack reach each other, so walk each one directly. + for p in reversed(stack): + closures[p] = self._walk( + partition_masks[p], + partition_masks, + assignment, + nodes, + kept_nodes, + closures, + ) + stack.clear() + else: + stack.append(missing) + return closures[partition_id] + + def _walk(self, masks, partition_masks, assignment, nodes, kept_nodes, closures): + """Bitmask of the given nodes and of every node they reach once each kept + partition is fused into a single node, following the cached closures and + the edges of the other partitions.""" + covered, reach, _ = masks + reach |= covered + while todo := reach & kept_nodes & ~covered: + position = (todo & -todo).bit_length() - 1 + other = assignment[nodes[position]] + if other in closures: + reach |= closures[other] + covered |= closures[other] + else: + other_nodes, other_reach, _ = partition_masks[other] + reach |= other_nodes | other_reach + covered |= other_nodes + return reach + + def _already_on_one_cycle( + self, + nodes1, + nodes2, + partition_masks, + assignment, + dependency_viewer, + kept_nodes, + closures, + ): + """Whether everything in two partitions already lies on one cycle, as with + caller node_groups that depend on each other. Merging them then puts no + node on a cycle that was not on one before. A kept partition counts as + one node, since it is fused, and any other partition as its own nodes.""" + nodes = dependency_viewer.nodes + # Closure and node mask of each part that is one node in the fused graph. + parts = [] + for n in (nodes1, nodes2): + if n & kept_nodes: + closure = self._closure( + assignment[nodes[(n & -n).bit_length() - 1]], + partition_masks, + assignment, + nodes, + kept_nodes, + closures, + ) + parts.append((closure, n)) + continue + while n: + bit = n & -n + n ^= bit + closure = self._walk( + dependency_viewer.masks_of([nodes[bit.bit_length() - 1]]), + partition_masks, + assignment, + nodes, + kept_nodes, + closures, + ) + parts.append((closure, bit)) + first_closure, first_nodes = parts[0] + return all( + first_closure & part_nodes and closure & first_nodes + for closure, part_nodes in parts[1:] + ) + + def _can_merge_partitions( + self, + masks1, + masks2, + partition_masks, + assignment, + dependency_viewer, + kept_nodes, + closures, + ): + """Whether two partitions can be merged. No node outside the merged + partition may be both downstream and upstream of it. Once the other kept + partitions are fused, it must not be on a cycle through them either, + unless everything in the two partitions already lies on one cycle.""" + nodes = dependency_viewer.nodes + merged, downstreams, upstreams = (m1 | m2 for m1, m2 in zip(masks1, masks2)) + # Fused on its own, the merged partition has a cycle when a node outside + # it is downstream of it and also upstream of it. + if downstreams & upstreams & ~merged: + return False + + # Every other kept partition becomes a single node when it is fused, so + # reaching one of its nodes reaches everything downstream of all of + # them. Without this, two partitions that are each free of cycles can + # still depend on each other. + covered = merged + while reached := downstreams & kept_nodes & ~covered: + position = (reached & -reached).bit_length() - 1 + closure = self._closure( + assignment[nodes[position]], + partition_masks, + assignment, + nodes, + kept_nodes, + closures, + ) + downstreams |= closure + if downstreams & upstreams & ~merged: + return self._already_on_one_cycle( + masks1[0], + masks2[0], + partition_masks, + assignment, + dependency_viewer, + kept_nodes, + closures, + ) + covered |= closure + return True def _process_all_nodes( self, @@ -186,6 +359,18 @@ def _merge_partitions(self, partitions_by_id, assignment): partition_id: dependency_viewer.masks_of(partition.nodes) for partition_id, partition in partitions_by_id.items() } + compute_weights = { + partition_id: sum(self._compute_weight(n) for n in partition.nodes) + for partition_id, partition in partitions_by_id.items() + } + # Nodes of the partitions that will be fused, so a merge check can + # follow paths through them. + kept_nodes = 0 + for partition_id, partition in partitions_by_id.items(): + if self._is_kept(partition, compute_weights[partition_id]): + kept_nodes |= partition_masks[partition_id][0] + # Closures of kept partitions, see _closure. + closures = {} # Set to track removed partitions from initial static list so we can skip them already_merged = set() @@ -200,13 +385,28 @@ def _merge_partitions(self, partitions_by_id, assignment): if p2 in already_merged: continue - # Try to merge partitions if it doesn't create cycles - if self._can_merge_partitions(partition_masks[p1], partition_masks[p2]): + if self._can_merge_partitions( + partition_masks[p1], + partition_masks[p2], + partition_masks, + assignment, + dependency_viewer, + kept_nodes, + closures, + ): self._perform_partition_merge(p1, p2, partitions_by_id, assignment) partition_masks[p1] = tuple( m1 | m2 for m1, m2 in zip(partition_masks[p1], partition_masks.pop(p2)) ) + compute_weights[p1] += compute_weights.pop(p2) + merged_nodes = partition_masks[p1][0] + if self._is_kept(partitions_by_id[p1], compute_weights[p1]): + kept_nodes |= merged_nodes + # Only closures that reach the merged nodes change. + stale = [p for p, c in closures.items() if c & merged_nodes] + for p in stale: + del closures[p] # Mark p2 as merged already_merged.add(p2) @@ -257,26 +457,10 @@ def _filter_single_node_partitions(self, partitions_by_id): if self.allows_single_node_partition: return - default_non_compute_ops = {"torch.ops.aten.view", "_operator.getitem"} - non_compute_ops = default_non_compute_ops.union(set(self.non_compute_ops or [])) partitions_to_remove = [] for id, partition in partitions_by_id.items(): - compute_node_count = 0 - for node in partition.nodes: - if node.op == "call_function": - assert callable(node.target) - target_name = _get_qualified_name(node.target) - - if target_name not in non_compute_ops: - compute_node_count += 1 - - if ( - self.allowed_single_node_partition_ops - and target_name in self.allowed_single_node_partition_ops - ): - compute_node_count += 1 - + compute_node_count = sum(self._compute_weight(n) for n in partition.nodes) if compute_node_count <= 1: partitions_to_remove.append(id) diff --git a/exir/backend/test/test_group_partitioner.py b/exir/backend/test/test_group_partitioner.py index 0f97fe60eb7..feef7687922 100644 --- a/exir/backend/test/test_group_partitioner.py +++ b/exir/backend/test/test_group_partitioner.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import operator import random import unittest from typing import List @@ -15,7 +16,7 @@ from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import create_op_support, OperatorSupportBase -from torch.fx.passes.utils.fuser_utils import validate_partition +from torch.fx.passes.utils.fuser_utils import fuse_by_partitions, validate_partition class TestGroupPartitioner(unittest.TestCase): @@ -1949,3 +1950,222 @@ def supports(nodes): ) for partition in partitioner.propose_partitions(): self.assertTrue(validate_partition(list(partition.nodes))) + + def test_no_cycle_through_another_partition(self): + """ + Test that merging stops when the merged partition and another partition + would depend on each other once both are fused: + + x -> a -> sin -> c (c and d are one node group) + x -> d -> sin -> b (a and b would be merged) + """ + graph = torch.fx.Graph() + x = graph.placeholder("x") + add, sin = torch.ops.aten.add.Tensor, torch.ops.aten.sin.default + a = graph.call_function(add, (x, x)) + c = graph.call_function(add, (graph.call_function(sin, (a,)), x)) + d = graph.call_function(add, (x, x)) + b = graph.call_function(add, (graph.call_function(sin, (d,)), x)) + graph.output((b, c)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + partitions = GroupBasedPartitioner( + graph_module, + create_op_support(lambda _, node: node.target == add), + node_groups=[[c, d]], + allows_single_node_partition=True, + ).propose_partitions() + + self.assertCountEqual([set(p.nodes) for p in partitions], [{a}, {b}, {c, d}]) + # Before the fix this raised "Invalid partition, found dependency cycles". + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + def test_merge_check_through_node_groups_that_depend_on_each_other(self): + """ + Test that a merge check can walk through node groups that depend on + each other until they are merged ({d, g} and {e, f}: d -> f, e -> g): + + x -> a -> b -> c, b -> d, b -> e, c -> e (a and c are one group) + d -> f, e -> f, e -> g, b -> g, f -> h, c -> h + + Before the fix, a, b and c merged, but {d, g} and {e, f, h} were left + depending on each other and fusing raised "Invalid partition, found + dependency cycles". + """ + graph = torch.fx.Graph() + x = graph.placeholder("x") + add = torch.ops.aten.add.Tensor + a = graph.call_function(add, (x, x)) + b = graph.call_function(add, (x, a)) + c = graph.call_function(add, (b, b)) + d = graph.call_function(add, (b, b)) + e = graph.call_function(add, (c, b)) + f = graph.call_function(add, (d, e)) + g = graph.call_function(add, (b, e)) + h = graph.call_function(add, (f, c)) + graph.output((g, h)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + partitions = GroupBasedPartitioner( + graph_module, + create_op_support(lambda _, node: node.target == add), + node_groups=[[a, c], [d, g], [e, f]], + allows_single_node_partition=True, + ).propose_partitions() + + self.assertCountEqual( + [set(p.nodes) for p in partitions], [{a, b, c}, {d, e, f, g}, {h}] + ) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + def test_cycle_check_through_chained_and_filtered_partitions(self): + """ + Test that the path from a to b is followed through two node groups in a + row, so a and b are not merged: + + x -> a -> sin -> c1 (c1 and d1 are one node group) + x -> d1 -> sin -> c2 (c2 and d2 are one node group) + x -> d2 -> sin -> b + + and that a path through a node group that is filtered out (one add and + one non-compute mul) does not stop a and b from merging: + + x -> a -> sin -> c (c and m are one node group) + x -> m -> sin -> b + """ + add, mul, sin = ( + torch.ops.aten.add.Tensor, + torch.ops.aten.mul.Tensor, + torch.ops.aten.sin.default, + ) + supported = create_op_support(lambda _, node: node.target in (add, mul)) + + graph = torch.fx.Graph() + x = graph.placeholder("x") + a = graph.call_function(add, (x, x)) + c1 = graph.call_function(add, (graph.call_function(sin, (a,)), x)) + d1 = graph.call_function(add, (x, x)) + c2 = graph.call_function(add, (graph.call_function(sin, (d1,)), x)) + d2 = graph.call_function(add, (x, x)) + b = graph.call_function(add, (graph.call_function(sin, (d2,)), x)) + graph.output((b, c1, c2)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + partitions = GroupBasedPartitioner( + graph_module, + supported, + node_groups=[[c1, d1], [c2, d2]], + allows_single_node_partition=True, + ).propose_partitions() + self.assertCountEqual( + [set(p.nodes) for p in partitions], [{a}, {b}, {c1, d1}, {c2, d2}] + ) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + graph = torch.fx.Graph() + x = graph.placeholder("x") + a = graph.call_function(add, (x, x)) + c = graph.call_function(add, (graph.call_function(sin, (a,)), x)) + m = graph.call_function(mul, (x, x)) + b = graph.call_function(add, (graph.call_function(sin, (m,)), x)) + graph.output((b, c)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + partitions = GroupBasedPartitioner( + graph_module, + supported, + non_compute_ops=["torch.ops.aten.mul.Tensor"], + node_groups=[[c, m]], + ).propose_partitions() + self.assertCountEqual([set(p.nodes) for p in partitions], [{a, b}]) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + def test_merge_node_groups_that_already_depend_on_each_other(self): + """ + Test that node groups which already depend on each other in a cycle + ({a, d} -> {b, e} -> {c, f} -> {a, d}) are merged until it is gone: + + x -> a -> e, x -> b -> f, x -> c -> d + """ + graph = torch.fx.Graph() + x = graph.placeholder("x") + add = torch.ops.aten.add.Tensor + a = graph.call_function(add, (x, x)) + b = graph.call_function(add, (x, x)) + c = graph.call_function(add, (x, x)) + d = graph.call_function(add, (c, x)) + e = graph.call_function(add, (a, x)) + f = graph.call_function(add, (b, x)) + graph.output((d, e, f)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + partitions = GroupBasedPartitioner( + graph_module, + create_op_support(lambda _, node: node.target == add), + node_groups=[[e, b], [a, d], [c, f]], + ).propose_partitions() + + self.assertCountEqual([set(p.nodes) for p in partitions], [{a, b, c, d, e, f}]) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + def test_filtered_node_group_merged_only_where_it_closes_a_cycle(self): + """ + Test that a node group which is filtered out (one add and one getitem) + is merged when it already lies on a cycle with the other partition: + + x -> a -> b -> c, a -> split -> getitem -> d -> f, b -> f + ({a, c}, {b, f} and {getitem, d} are node groups) + + {a, c} and {b, f} depend on each other and only fuse together with + split, getitem and d. Before the fix, merging {getitem, d} was refused + and fusing raised "Invalid partition, found dependency cycles". + + And that it is not merged otherwise: + + x -> a -> a3, x -> split -> getitem -> a2 -> a3, a -> a4 + ({a, a3} and {a4, getitem} are node groups, split is not supported) + + a2 reaches {a4, getitem} only at a4, through {a, a3}, so merging them + would make a partition that depends on {a, a3} both ways. + """ + add, split = torch.ops.aten.add.Tensor, torch.ops.aten.split.Tensor + + graph = torch.fx.Graph() + x = graph.placeholder("x") + a = graph.call_function(add, (x, x)) + b = graph.call_function(add, (a, x)) + c = graph.call_function(add, (b, x)) + s = graph.call_function(split, (a, 1)) + getitem = graph.call_function(operator.getitem, (s, 0)) + d = graph.call_function(add, (getitem, x)) + f = graph.call_function(add, (d, b)) + graph.output((c, f)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + partitions = GroupBasedPartitioner( + graph_module, + create_op_support( + lambda _, node: node.target in (add, split, operator.getitem) + ), + node_groups=[[a, c], [b, f], [getitem, d]], + ).propose_partitions() + self.assertCountEqual( + [set(p.nodes) for p in partitions], [{a, b, c, s, getitem, d, f}] + ) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) + + graph = torch.fx.Graph() + x = graph.placeholder("x") + a = graph.call_function(add, (x, x)) + getitem = graph.call_function( + operator.getitem, (graph.call_function(split, (x, 1)), 0) + ) + a2 = graph.call_function(add, (getitem, getitem)) + a3 = graph.call_function(add, (a2, a)) + a4 = graph.call_function(add, (a, x)) + graph.output((a3, a4)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + partitions = GroupBasedPartitioner( + graph_module, + create_op_support(lambda _, node: node.target in (add, operator.getitem)), + node_groups=[[a, a3], [a4, getitem]], + ).propose_partitions() + self.assertCountEqual([set(p.nodes) for p in partitions], [{a, a2, a3}]) + fuse_by_partitions(graph_module, [p.nodes for p in partitions]) From fda25256717ddba495bdf6f7d9ca17d055db7b68 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:46:22 -0700 Subject: [PATCH 20/37] Cortex-M: Handle a constant operand in MatmulToBmmPass MatmulToBmmPass read the operand shapes with ProxyValue.to_tensor(). When an operand of the matmul is a constant (a parameter, a buffer or a tensor attribute), ExportPass passes it to call_operator as a plain torch.Tensor, so CortexMQuantizer crashed in prepare_pt2e for every such matmul: AttributeError: 'Tensor' object has no attribute 'to_tensor' For example softmax(x) @ codebook, where codebook is a tensor attribute. Read the shapes through .data, which works for both kinds of operand. A constant operand of a shape the pass rewrites, such as a batched rank-3 codebook used as the right-hand side, now becomes bmm like any other operand and lowers to cortex_m.quantized_batch_matmul, which already supports a constant right-hand side. Shapes the pass does not rewrite, such as rank-3 @ rank-2, no longer crash but stay matmul. --- .../cortex_m/passes/matmul_to_bmm_pass.py | 9 ++++---- .../cortex_m/test/ops/test_batch_matmul.py | 22 +++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/backends/cortex_m/passes/matmul_to_bmm_pass.py b/backends/cortex_m/passes/matmul_to_bmm_pass.py index 0d778a25d2a..a7af74c5b97 100644 --- a/backends/cortex_m/passes/matmul_to_bmm_pass.py +++ b/backends/cortex_m/passes/matmul_to_bmm_pass.py @@ -45,10 +45,11 @@ def call_operator( if op != torch.ops.aten.matmul.default: return super().call_operator(op, args, kwargs, meta) - lhs = cast(ProxyValue, args[0]) - rhs = cast(ProxyValue, args[1]) - lhs_shape = lhs.to_tensor().shape - rhs_shape = rhs.to_tensor().shape + # A constant operand is a plain tensor; .data works for both kinds. + lhs = cast(ProxyValue | torch.Tensor, args[0]) + rhs = cast(ProxyValue | torch.Tensor, args[1]) + lhs_shape = lhs.data.shape + rhs_shape = rhs.data.shape lhs_rank = len(lhs_shape) rhs_rank = len(rhs_shape) diff --git a/backends/cortex_m/test/ops/test_batch_matmul.py b/backends/cortex_m/test/ops/test_batch_matmul.py index 926973cb438..874735372ad 100644 --- a/backends/cortex_m/test/ops/test_batch_matmul.py +++ b/backends/cortex_m/test/ops/test_batch_matmul.py @@ -88,6 +88,24 @@ def forward(self, lhs, rhs): return torch.matmul(lhs, rhs) +class CortexMMatmulConstantRhs(torch.nn.Module): + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_bmm_default": 1, + } + + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 1, + "executorch_exir_dialects_edge__ops_aten_bmm_default": 0, + } + + def __init__(self, rhs_shape): + super().__init__() + self.rhs = torch.randn(rhs_shape) + + def forward(self, lhs): + return torch.matmul(lhs, self.rhs) + + test_cases = { "bmm_small": McuTestCase( CortexMBmm(), @@ -137,6 +155,10 @@ def forward(self, lhs, rhs): CortexMMatmulBroadcast(), (torch.randn(1, 4, 8), torch.randn(2, 8, 4)), ), + "matmul_const_rhs_rank3": McuTestCase( + CortexMMatmulConstantRhs((2, 8, 4)), + (torch.randn(2, 4, 8),), + ), } From 6efe853be9932f20ab25180c54ef9895064e8c20 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:18:55 -0700 Subject: [PATCH 21/37] exir: Give tensor constants static shapes in ExportPass When ExportPass meets a real tensor, for example a get_attr constant, it makes a fake tensor for the node metadata. Only parameters were faked with static_shapes=True. Any other tensor that the fake mode had not seen before got a fresh symbol for every dimension when the fake mode has a ShapeEnv, which is the case for exported programs. Such tensors are common: constant folding (for example the Arm quantizer's ConstantFoldingPass) creates new constant tensors, and the next ExportPass gives them symbolic sizes like (s45, s96) instead of (9, 2). Later passes then see symbolic shapes where everything is static. With the LeRobot diffusion policy, the Arm quantizer then failed in DecomposeVarPass: AssertionError: a.node.constant should not be None A real tensor is constant data with known sizes, so fake all of them with static shapes, as the existing TODO suggested. --- exir/pass_base.py | 13 +++++-------- exir/tests/test_pass_infra.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/exir/pass_base.py b/exir/pass_base.py index 39c649c822d..f4e5f17e974 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -450,14 +450,11 @@ def make_val( try: assert self.fake_tensor_mode is not None - # TODO we should allocate static shapes - # for param/buffer values - if isinstance(x, torch.nn.Parameter): - fake_tensor = self.fake_tensor_mode.from_tensor( - x, static_shapes=True - ) - else: - fake_tensor = self.fake_tensor_mode.from_tensor(x) + # A real tensor here is constant data, so its sizes + # are known and must not become new symbols. + fake_tensor = self.fake_tensor_mode.from_tensor( + x, static_shapes=True + ) except UnsupportedFakeTensorException: # TODO: This is just a workaround to get over the # x.as_subclass error diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 168485c2905..e33cb1b68ba 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -502,6 +502,28 @@ def test_export_pass_preserves_symbolic_input_metadata(self) -> None: self.assertEqual(self._symbolic_input_shape(new_input), original_snapshot) + def test_export_pass_keeps_static_shapes_for_tensor_constants(self) -> None: + class AddConstant(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.constant = torch.randn(9, 2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.constant + + graph_module = export(AddConstant(), (torch.randn(9, 2),), strict=True).module( + check_guards=False + ) + (constant,) = [n for n in graph_module.graph.nodes if n.op == "get_attr"] + # A constant made by an earlier pass, e.g. constant folding, is a + # tensor the fake mode has not seen before. + setattr(graph_module, str(constant.target), torch.randn(9, 2)) + + new_graph_module = ExportPass()(graph_module).graph_module + (constant,) = [n for n in new_graph_module.graph.nodes if n.op == "get_attr"] + + self.assertEqual([type(dim) for dim in constant.meta["val"].shape], [int, int]) + def test_export_pass_tracks_replayed_unbacked_symbols(self) -> None: def data_dependent_slice( tensor: torch.Tensor, length: torch.Tensor From 5120ad6bef2144b065b83203a826587d6464f207 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:26:14 -0700 Subject: [PATCH 22/37] [MLX] Read the base of an out-of-place slice update as an input aten.slice_scatter is functional: it returns a new tensor and leaves its base alone. The MLX builder lowers it to a SliceUpdateNode whose output slot differs from its base slot, and the runtime handles that case without writing the base. Still, exec_slice_update always fetched the base through tensor_ref(), which is the accessor for tensors an op writes, and which refuses constants. So a slice_scatter whose base is a buffer or another constant failed at run time with "tensor_ref: cannot mutate constant tensor". Read the base with const_tensor_ref(), and use tensor_ref() only for the in-place form, which is the one that writes it. --- backends/mlx/runtime/MLXInterpreter.h | 8 +++++--- backends/mlx/test/test_partitioner.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 135baa2dcad..8e046f1609b 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -1211,7 +1211,8 @@ inline void exec_slice_update( // st.set_tensor does a std::move. Keeping for now, but should investigate and // possibly remove in future. const bool in_place = (n.out.idx == n.dst.idx); - array& dst = st.tensor_ref(n.dst); + // Only the in-place form writes dst. Otherwise dst may be a constant. + const array& dst = st.const_tensor_ref(n.dst); const array& upd = st.const_tensor_ref(n.update); const int rank = static_cast(dst.ndim()); @@ -1252,9 +1253,10 @@ inline void exec_slice_update( return; } if (n.step == 1) { - dst = slice_update(dst, upd, to_shape(vstart), to_shape(vstop), s); + st.tensor_ref(n.dst) = + slice_update(dst, upd, to_shape(vstart), to_shape(vstop), s); } else { - dst = slice_update( + st.tensor_ref(n.dst) = slice_update( dst, upd, to_shape(vstart), to_shape(vstop), to_shape(vstrides), s); } diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index f20ab617825..5897fe58bee 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -428,6 +428,21 @@ def test_rank2_unequal_lengths_use_causal_correction(self): self.assertLess(err, 1e-4) +class TestMLXPartitionerConstantBase(unittest.TestCase): + def test_slice_scatter_into_a_buffer(self): + class ScatterIntoBuffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("base", torch.zeros(4, 8)) + + def forward(self, x): + return torch.slice_scatter(self.base, x, 1, 0, 2) * 2 + + counts, _, err = _run(ScatterIntoBuffer().eval(), (torch.randn(4, 2),)) + self.assertEqual(counts.get("SliceUpdateNode", 0), 1) + self.assertLess(err, 1e-6) + + class TestMLXPartitionerMixedSupport(unittest.TestCase): """An operator is preserved from decomposition per operator, not per call.""" From ef39424f481005dfa04c5e5c2e57f65317d2aeb8 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:08:54 -0700 Subject: [PATCH 23/37] Qualcomm AI Engine Direct - Decompose a stack that HTP Pack cannot take The partitioner keeps aten.stack from being decomposed so that it can lower to the QNN Pack op. HTP's Pack takes inputs of rank 1 to 4 (QNN HTP op definition), so its output has 2 to 5 dims. A stack of 0-d or 5-D tensors is rejected at op validation. The node then stays on CPU, but because it was kept from decomposition, to_edge_transform_and_lower fails with "aten.stack ... was not decomposed or delegated". Let filter_fn decompose a stack whose output rank is outside 2 to 5, the same way it already does for integer stacks. The unsqueeze and cat it decomposes into are checked by the partitioner like any other op. --- backends/qualcomm/partition/utils.py | 5 +++++ backends/qualcomm/tests/test_passes.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 5fa5633ac70..1bc339d4cc7 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -148,6 +148,11 @@ def filter_fn(node: torch.fx.Node) -> bool: torch.int64, ]: return False + # HTP Pack takes inputs of rank 1 to 4, so its output has 2 to 5 dims. + if node.target == torch.ops.aten.stack.default and not ( + 2 <= node.meta["val"].dim() <= 5 + ): + return False # QNN Stft only produces the real-valued spectrum and has no normalization. # ConvertComplexStft rewrites the supported complex cases into this form. # Without a window, QNN's rectangular default spans n_fft rather than diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 9b2b73f69d3..aaabaed4a0e 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -828,6 +828,28 @@ def forward(self, x): if "QNN" in str(e) or "qnn" in str(e): self.skipTest(f"QNN SDK not available: {e}") + def test_stack_outside_pack_ranks_is_decomposed(self): + class Stack(torch.nn.Module): + def forward(self, x): + return torch.stack([x, x]) + 1 + + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + # HTP Pack takes inputs of rank 1 to 4. Lowering must not fail with + # "aten.stack ... was not decomposed or delegated". + for shape in [(), (1, 2, 3, 4, 5)]: + with self.subTest(input_rank=len(shape)): + try: + to_edge_transform_and_lower_to_qnn( + Stack().eval(), (torch.randn(shape),), compiler_specs + ) + except RuntimeError as e: + if "Failed to initialize QNN backend" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + def _assert_cpu_conv_keeps_its_weight_dequantize( self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True ): From 13cb7d20ee6e8b9c2f25e45ffb7a05fce9eac08a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:03:13 -0700 Subject: [PATCH 24/37] Qualcomm AI Engine Direct - Keep partitions with no runtime input on CPU When an op stays on CPU, the ops before it can form a partition that is fed only by constants (parameters, buffers, lifted constants), or by nothing at all, for example a lone aten.full or a chain of casts of lifted scalars. HTP refuses such a graph ("No graph inputs present for graph", or "Graph has zero node"), so context binary generation failed and the whole lowering stopped. For HTP, drop a partition when none of its inputs is a runtime value, so its ops run on CPU. A mutable buffer counts as a runtime input, since QNN reads it as a graph input in the partition that updates it. tag_nodes no longer assumes at least one partition is left, which dropping can now cause. Other backends keep their partitions. GPU accepts graphs without inputs, and LPAI places its quantize and dequantize boundaries for CPU fallback before partitioning, so dropping a partition there afterwards would hand a float tensor to a delegate that expects a quantized one. LPAI refuses graphs without inputs too ("Graph has no input tensors"), so such a partition still fails there as before. Fixing that needs the decision to move before LPAI places its boundaries. Two HTP cases still fail as before: a dropped partition that holds an op the partitioner keeps from decomposition, such as a linear on a buffer (that op cannot run on CPU undecomposed, so lowering now fails with "was not decomposed or delegated"), and a partition whose only runtime input is a mutable buffer that a different partition updates. --- .../qualcomm/partition/qnn_partitioner.py | 44 ++++++++++++++++++- backends/qualcomm/tests/test_passes.py | 28 ++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index f29ac21b7f6..fbab3b9bc40 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -14,7 +14,10 @@ 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.builders.utils import ( + is_mutable_buffer_input, + is_parameter, +) from executorch.backends.qualcomm.qnn_preprocess import QnnBackend from executorch.backends.qualcomm.serialization.qc_schema import ( HtpArch, @@ -253,7 +256,14 @@ def generate_partitions( edge_program.graph_module, op_support=self.op_support_checker, ) - return self._drop_orphaned_parameter_dequants(partitions, edge_program) + partitions = self._drop_orphaned_parameter_dequants(partitions, edge_program) + # LPAI places its CPU fallback boundaries before partitioning, and GPU + # accepts graphs without inputs, so only HTP partitions are dropped. + if self.backend == QnnExecuTorchBackendType.kHtpBackend: + partitions = self._drop_partitions_without_runtime_inputs( + partitions, edge_program + ) + return partitions @staticmethod def _drop_orphaned_parameter_dequants( @@ -286,6 +296,32 @@ def _drop_orphaned_parameter_dequants( ) return [partition for partition in partitions if partition.nodes] + @staticmethod + def _drop_partitions_without_runtime_inputs( + partitions: List[Partition], edge_program: torch.export.ExportedProgram + ) -> List[Partition]: + """ + HTP refuses a graph that reads nothing at runtime, e.g. a subgraph fed + only by constants once a CPU op splits it from the rest. Such a + partition stays on CPU. + """ + + def is_runtime_input(node: torch.fx.Node) -> bool: + return not is_parameter(node, edge_program) or is_mutable_buffer_input( + node, edge_program + ) + + return [ + partition + for partition in partitions + if any( + is_runtime_input(arg) + for node in partition.nodes + for arg in node.all_input_nodes + if arg not in partition.nodes + ) + ] + def tag_nodes( self, partitions: List[Partition], edge_program: torch.export.ExportedProgram ) -> None: @@ -298,6 +334,10 @@ def tag_nodes( node.meta["delegation_tag"] = delegation_tag self.partition_tags[delegation_tag] = self.delegation_spec + # Dropping partitions can leave none, and there is no tag to reuse. + if not partitions: + return + # need to take care of consumed constants consumed_constants = ( *edge_program.graph_signature.inputs_to_buffers, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index aaabaed4a0e..8de383f86c7 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -850,6 +850,34 @@ def forward(self, x): self.skipTest(f"QNN SDK not available: {e}") raise + def test_partitioner_keeps_constant_only_subgraph_on_cpu(self): + class ConstantBranch(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("table", torch.randn(4, 16)) + + def forward(self, x): + return x + torch.relu(self.table * 2 - 1) + + sample_input = (torch.randn(4, 16),) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + # With relu on CPU, table * 2 - 1 reads nothing at runtime. + try: + edge = to_edge_transform_and_lower_to_qnn( + ConstantBranch().eval(), + sample_input, + compiler_specs, + skip_node_op_set={"aten.relu.default"}, + ) + except RuntimeError as e: + if "Failed to initialize QNN backend" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + edge.to_executorch() + def _assert_cpu_conv_keeps_its_weight_dequantize( self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True ): From 9efbc22fdc5b94c7696c7aca0c41cc3b20c2f1c8 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:39:53 -0700 Subject: [PATCH 25/37] Qualcomm AI Engine Direct - Remove a 0-d select only when its input can replace it Remove0DTensor drops every select that returns a 0-d tensor and feeds its users the whole input instead. That is only right when the input holds just the selected element and every user broadcasts it, so that a one-element tensor gives the same result as the 0-d element. The pass was written for that case (torch.zeros(1)[0] added to a tensor). Otherwise: - For a lookup such as table[t] on a longer buffer, the users got the whole table: lowering failed later with a broadcast error, or, when the shapes happened to broadcast, the program silently computed the wrong values. - For a one-element input, a user that depends on the rank, such as torch.stack, got a tensor of shape (1,) instead of a 0-d one and produced a different shape. Only drop the select when its input is known to have one element and every user, looking through quantize and dequantize, is an add, sub, mul or div whose output has at least one dimension. Other 0-d selects are left in the graph and run on CPU. That includes a select used by another op such as sqrt, or by a dtype cast (x * int_buffer[0] with a float x): the result is still right, but some ops run on CPU instead of in the delegate. Casts are not looked through, because a dim-order cast keeps the 0-d layout in its own arguments and fails at runtime on CPU when given a one-element tensor. Two kinds of graph then lower only with other changes in place. A stack of such selects needs the partitioner to decompose stacks that HTP Pack cannot take. A select and cast left on CPU can leave an HTP partition fed only by constants, which has to stay on CPU as well. So this change lands after both. The pass's own test in tests/rework selected from a four-element tensor and checked that the select was removed, so it asserted the wrong behavior; it now uses a one-element input. --- backends/qualcomm/_passes/remove_0d_tensor.py | 30 ++++++++++++- backends/qualcomm/tests/rework/src/pattern.py | 2 +- backends/qualcomm/tests/test_passes.py | 42 +++++++++++++++++++ 3 files changed, 71 insertions(+), 3 deletions(-) diff --git a/backends/qualcomm/_passes/remove_0d_tensor.py b/backends/qualcomm/_passes/remove_0d_tensor.py index 4056910f6ef..07298686a17 100644 --- a/backends/qualcomm/_passes/remove_0d_tensor.py +++ b/backends/qualcomm/_passes/remove_0d_tensor.py @@ -5,14 +5,18 @@ # LICENSE file in the root directory of this source tree. import torch +from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult from executorch.exir.passes import dead_code_elimination_pass +from torch.fx.experimental.symbolic_shapes import statically_known_true class Remove0DTensor(ExportPass): """ - QNN does not allow 0D tensor, we remove the node that will output an 0D tensor. + QNN does not allow 0D tensor, so a select whose 0D output is the only element of + its input is replaced by that input, when every user broadcasts it into an + output of rank one or more and so gets the same result. Before adding operations to the list of nodes to be removed, please ensure that it will not change the logic. """ @@ -21,13 +25,35 @@ class Remove0DTensor(ExportPass): exir_ops.edge.aten.select_copy.int, } + broadcast_ops = { + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.div.Tensor, + } + def __init__(self, quantization_capture=False) -> None: super().__init__() + def _only_broadcast(self, node: torch.fx.Node) -> bool: + return all( + ( + self._only_broadcast(user) + if user.target in q_ops or user.target in dq_ops + else user.target in self.broadcast_ops and user.meta["val"].dim() > 0 + ) + for user in node.users + ) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph = graph_module.graph for node in graph.nodes: - if node.target in self.remove_ops and len(node.meta["val"].shape) == 0: + if ( + node.target in self.remove_ops + and len(node.meta["val"].shape) == 0 + and statically_known_true(node.args[0].meta["val"].numel() == 1) + and self._only_broadcast(node) + ): for user_n in list(node.users.keys()): user_n.replace_input_with(node, node.args[0]) graph.erase_node(node) diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..150f9160669 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -5000,7 +5000,7 @@ def test( assertions: Assertions, pass_pipeline: PassPipeline, ): - inputs = (torch.randn(4), torch.randn(3, 4)) + inputs = (torch.randn(1), torch.randn(3, 4)) gm = pass_pipeline.lower_edge_ep( module=Remove0DTensor._ZeroDim(), sample_input=inputs, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 8de383f86c7..ec71db72acd 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -13,6 +13,7 @@ InsertIOQDQ, InsertReshapeForReduceOps, LiftConstantScalarOperands, + Remove0DTensor, RemoveRedundancy, ) from executorch.backends.qualcomm._passes.qnn_pass_manager import ( @@ -432,6 +433,47 @@ def forward(self, x): torch.equal(*out, ref), f"Output mismatch: got {out}, expected {ref}" ) + def test_remove_0d_tensor_keeps_select_it_cannot_replace(self): + class TableLookup(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("table", torch.arange(4.0)) + + def forward(self, x): + return x * self.table[2] + + class StackSelected(torch.nn.Module): + def forward(self, x): + return torch.stack([x[0], x[0]]) + + class CastSelected(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("scale", torch.tensor([3])) + + def forward(self, x): + return x * self.scale[0] + + # A longer input is not the selected element, stacking the one-element + # input instead of the 0-d element changes the shape, and a cast keeps + # the 0-d layout in its arguments. + for model, x in ( + (TableLookup(), torch.randn(4)), + (StackSelected(), torch.randn(1)), + (CastSelected(), torch.randn(4)), + ): + with self.subTest(model=type(model).__name__): + ep = to_edge(torch.export.export(model, (x,), strict=True)) + ep = ep.exported_program() + Remove0DTensor()(ep.graph_module) + self.assertTrue( + any( + n.target == exir_ops.edge.aten.select_copy.int + for n in ep.graph.nodes + ) + ) + self.assertTrue(torch.equal(ep.module()(x), model(x))) + def test_mha_to_sha(self): from executorch.backends.qualcomm.utils.utils import convert_linear_to_conv2d from executorch.examples.models.llama.model_args import ModelArgs From 9871bbfab8a014c73b341095e80fa760a646b720 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:46:59 -0700 Subject: [PATCH 26/37] Qualcomm AI Engine Direct - Recompose RMSNorm only over the last dim with its weight RecomposeRmsNorm takes any constant operand of the last multiply as the RMSNorm weight, uses its shape as normalized_shape, and never looks at how the mean reduces. Two things go wrong: - When a norm is followed by a multiply with a 0-d constant (a scalar scale), it builds rms_norm(x, normalized_shape=[]), which fails the next time the graph is traced: "Expected normalized_shape to be at least 1-dimensional". - When the mean reduces a dim other than the last one, several dims, or drops the reduced dim (keepdim=False), it still builds an rms_norm over the last dim, so the program silently computes different values. The backend's RmsNorm builder normalizes the last dim with a weight of that size, so only recompose that form: the mean must reduce exactly the last dim and keep it (a 1-D input may reduce all dims or drop it), and the constant must have the shape of that dim. Anything else is left as the decomposed ops. --- .../qualcomm/_passes/recompose_rms_norm.py | 19 +++++++++++-- backends/qualcomm/tests/test_passes.py | 28 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/backends/qualcomm/_passes/recompose_rms_norm.py b/backends/qualcomm/_passes/recompose_rms_norm.py index 0572fe6b4ee..f9cdd676345 100644 --- a/backends/qualcomm/_passes/recompose_rms_norm.py +++ b/backends/qualcomm/_passes/recompose_rms_norm.py @@ -127,10 +127,25 @@ def call(self, graph_module: torch.fx.GraphModule): if gamma_node is None: continue + input_node = self._get_input_node(rms_norm_pattern[6]) + # Only the last-dim form the RmsNorm builder supports is recomposed. + # Anything else, e.g. a 0-d scale, is left decomposed. + input_shape = list(input_node.meta["val"].shape) + rank = len(input_shape) + mean_node = rms_norm_pattern[5] + mean_dims = mean_node.args[1] or range(rank) + keepdim = len(mean_node.args) > 2 and mean_node.args[2] + if ( + rank == 0 + or [dim % rank for dim in mean_dims] != [rank - 1] + or not (keepdim or rank == 1) + or list(gamma_node.meta["val"].shape) != input_shape[-1:] + ): + continue + eps = rms_norm_pattern[4].args[1] if isinstance(eps, torch.fx.Node): eps = eps.meta["val"].constant.item() - input_node = self._get_input_node(rms_norm_pattern[6]) with graph.inserting_before(last_mul_node): # args schema @@ -140,7 +155,7 @@ def call(self, graph_module: torch.fx.GraphModule): self.rms_norm_target, ( input_node, - list(gamma_node.meta["val"].shape), + input_shape[-1:], gamma_node, eps, ), diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index ec71db72acd..7d8111034b9 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -13,6 +13,7 @@ InsertIOQDQ, InsertReshapeForReduceOps, LiftConstantScalarOperands, + RecomposeRmsNorm, Remove0DTensor, RemoveRedundancy, ) @@ -474,6 +475,33 @@ def forward(self, x): ) self.assertTrue(torch.equal(ep.module()(x), model(x))) + def test_recompose_rms_norm_only_takes_last_dim_with_weight(self): + class ScaledNorm(torch.nn.Module): + def __init__(self, dim, scale, keepdim=True): + super().__init__() + self.dim = dim + self.keepdim = keepdim + self.register_buffer("scale", scale) + + def forward(self, x): + ms = x.pow(2).mean(self.dim, keepdim=self.keepdim) + return x * torch.rsqrt(ms + 1e-6) * self.scale + + x = torch.randn(32, 32) + # A 0-d scale is not a weight, and a mean over dim 0, over all dims, or + # without keepdim is not a last-dim RMSNorm. + for model in ( + ScaledNorm(-1, torch.tensor(2.0)), + ScaledNorm(0, torch.randn(32)), + ScaledNorm(None, torch.randn(32)), + ScaledNorm(-1, torch.randn(32), keepdim=False), + ): + with self.subTest(dim=model.dim, keepdim=model.keepdim): + ep = to_edge(torch.export.export(model, (x,), strict=True)) + ep = ep.exported_program() + RecomposeRmsNorm()(ep.graph_module) + torch.testing.assert_close(ep.module()(x), model(x)) + def test_mha_to_sha(self): from executorch.backends.qualcomm.utils.utils import convert_linear_to_conv2d from executorch.examples.models.llama.model_args import ModelArgs From b9617ee666974725d9e1a8a5844fb03f4e47df74 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:54:04 -0700 Subject: [PATCH 27/37] Qualcomm AI Engine Direct - Leave bfloat16 nodes on CPU The Qualcomm backend has no bfloat16 tensor type (QNN_TENSOR_TYPE_MAP has no entry for it), but nothing checked data types before a node visitor was asked to define the node's tensors. An unquantized bfloat16 node (for example an explicit x.to(torch.bfloat16) region in a float model) therefore crashed lowering with KeyError: torch.bfloat16 in NodeVisitor.get_data_type instead of staying on CPU. Add has_unsupported_dtype, true when the node's output or one of its inputs is an unquantized tensor of a data type the backend cannot represent, and use it in two places: - The partitioner skips such a node, so it falls back to CPU like any other unsupported node. - filter_fn does not keep such a node from decomposition unless every such tensor is quantized: each such input comes from a dequantize op and the output feeds a quantize op. Otherwise an op the backend normally lowers whole, such as a bfloat16 linear, would stay on CPU undecomposed. That fails with "was not decomposed or delegated", or, for a linear whose weight alone is dequantized (weight-only quantization), needs a linear kernel that the portable CPU kernels do not have. A quantized tensor takes its QNN type from its encoding, so quantized bfloat16 regions are still delegated as before. A bfloat16 graph input or output of a quantized model is still not supported. The test covers a quantized, an fp16 and a weight-only quantized bfloat16 linear. It lowers the quantized model first, because some QNN SDK versions crash on a quantized lowering that follows an fp16 one in the same process. --- .../qualcomm/partition/qnn_partitioner.py | 9 ++- backends/qualcomm/partition/utils.py | 53 +++++++++++++- backends/qualcomm/tests/test_passes.py | 69 +++++++++++++++++++ 3 files changed, 129 insertions(+), 2 deletions(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index fbab3b9bc40..4f3cfd4bd2d 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -60,6 +60,7 @@ filter_fn, generate_qnn_executorch_option, get_skip_decomp_table, + has_unsupported_dtype, warn_if_dilated_conv_may_not_fit_vtcm, ) @@ -125,7 +126,7 @@ def __init__( for soc_model in target_socs ] - def is_node_supported(self, _, node: torch.fx.Node) -> bool: + def is_node_supported(self, _, node: torch.fx.Node) -> bool: # noqa: C901 if node.op != "call_function" or node.target in not_supported_operator: return False @@ -172,6 +173,12 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: logger.info(f"[{self.phase}] {node.target.__name__} | Skipped") return False + if has_unsupported_dtype(node): + logger.info( + f"[{self.phase}] {node.target.__name__} | Skipped, unsupported dtype" + ) + return False + # TODO: workaround for https://github.com/pytorch/executorch/issues/23096. # Remove this once the fp16 path performs the TCM-fit check the quantized # path does. diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 1bc339d4cc7..7880221a1b2 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -8,11 +8,21 @@ from typing import List, Optional import torch +from executorch.backends.qualcomm.builders.node_visitor import ( + dq_ops, + q_ops, + QNN_TENSOR_TYPE_MAP, +) -from executorch.backends.qualcomm.utils.constants import QCOM_QNN_COMPILE_SPEC +from executorch.backends.qualcomm.utils.constants import ( + QCOM_BYPASS_NODE, + QCOM_QNN_COMPILE_SPEC, + QCOM_QUANT_ATTRS, +) from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.dialects._ops import ops as exir_ops +from torch.utils import _pytree as pytree logger = logging.getLogger(__name__) @@ -136,8 +146,49 @@ def warn_if_dilated_conv_may_not_fit_vtcm( return estimate +def has_unsupported_dtype(node: torch.fx.Node) -> bool: + """ + Whether the node reads or writes an unquantized tensor of a data type the + backend has no QNN type for (e.g. bfloat16). Such a node has to stay on + CPU. A quantized tensor takes its QNN type from its encoding instead. + """ + return any( + _produces_unsupported_dtype(n) + for n in (node, *node.all_input_nodes) + if QCOM_QUANT_ATTRS not in n.meta and not n.meta.get(QCOM_BYPASS_NODE) + ) + + +def _produces_unsupported_dtype(node: torch.fx.Node) -> bool: + return any( + isinstance(val, torch.Tensor) and val.dtype not in QNN_TENSOR_TYPE_MAP + for val in pytree.tree_leaves(node.meta.get("val")) + ) + + +def _unsupported_dtypes_are_quantized(node: torch.fx.Node) -> bool: + # filter_fn runs before quantization attributes are annotated, so the + # surrounding quantize and dequantize ops are the only sign of them. + inputs_quantized = all( + arg.target in _DQ_OPS + for arg in node.all_input_nodes + if _produces_unsupported_dtype(arg) + ) + output_quantized = not _produces_unsupported_dtype(node) or any( + user.target in _Q_OPS for user in node.users + ) + return inputs_quantized and output_quantized + + +_Q_OPS = {op._op for op in q_ops} +_DQ_OPS = {op._op for op in dq_ops} + + # Logic to determine whether to skip decompose and has higher priority than get_skip_decomp_table() def filter_fn(node: torch.fx.Node) -> bool: + # The node will stay on CPU, so it must not be kept from decomposition. + if has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node): + return False # QNN does not support int32/int64 IO for the following OPs. potential_i32_i64_io_ops = [ torch.ops.aten.stack.default, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 7d8111034b9..a4ffa89826a 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -62,6 +62,7 @@ from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import OutputKind from torch.library import Library +from torchao.quantization import IntxWeightOnlyConfig, PerAxis, quantize_ from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -948,6 +949,74 @@ def forward(self, x): raise edge.to_executorch() + def test_partitioner_falls_back_on_bfloat16(self): + class Bf16Region(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(16, 16).to(torch.bfloat16) + + def forward(self, x): + return self.linear(x.to(torch.bfloat16)).to(torch.float32) + 1 + + sample_input = (torch.randn(4, 16),) + # Quantized first: some QNN SDKs crash on a quantized lowering that + # follows an fp16 one in the same process. + for quant_dtype, weight_only in ( + (QuantDtype.use_8a8w, False), + (None, False), + (None, True), + ): + with self.subTest(quant_dtype=quant_dtype, weight_only=weight_only): + module = Bf16Region().eval() + if weight_only: + quantize_( + module, + IntxWeightOnlyConfig( + weight_dtype=torch.int8, granularity=PerAxis(0) + ), + ) + if quant_dtype is not None: + quantizer = QnnQuantizer() + quantizer.set_default_quant_config(quant_dtype=quant_dtype) + module = prepare_pt2e( + torch.export.export(module, sample_input).module(), quantizer + ) + module(*sample_input) + module = convert_pt2e(module) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec( + use_fp16=quant_dtype is None + ), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + module, sample_input, compiler_specs + ) + except RuntimeError as e: + if "Failed to initialize QNN backend" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + cpu_nodes = [ + node + for node in edge.exported_program().graph.nodes + if node.op == "call_function" and "delegate" not in str(node.target) + ] + cpu_dtypes = { + node.meta["val"].dtype + for node in cpu_nodes + if isinstance(node.meta.get("val"), torch.Tensor) + } + # Unquantized bfloat16 stays on CPU; quantized bfloat16 has a + # QNN type and is delegated. + self.assertEqual(torch.bfloat16 in cpu_dtypes, quant_dtype is None) + # Portable kernels have no linear, so one left on CPU must be + # decomposed. + self.assertNotIn( + exir_ops.edge.aten.linear.default, + [node.target for node in cpu_nodes], + ) + def _assert_cpu_conv_keeps_its_weight_dequantize( self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True ): From 78cc8bd2654ebc25ae62a513db8664b0e3ba3e20 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 08:34:33 -0700 Subject: [PATCH 28/37] Qualcomm AI Engine Direct - Leave nodes with data-dependent shapes on CPU QNN needs a size for every tensor dimension. For a dynamic dimension the node visitor passes the SymInt's hint, but a data-dependent size (for example the result of boolean mask indexing, x[mask]) has no hint. Nothing checked for that before a node visitor was asked to define such a tensor, and lowering crashed with a TypeError from the TensorWrapper constructor instead of leaving the node on CPU. Add has_data_dependent_shape, true when the node's output or one of its inputs has a dimension without a hint, and use it in two places: - The partitioner skips such a node, so it falls back to CPU like any other unsupported node. - filter_fn does not keep such a node from decomposition. Otherwise an op the backend normally lowers whole, such as a stack of masked tensors, would stay on CPU undecomposed and fail with "was not decomposed or delegated". A linear on masked rows also needs partitions without a runtime input to stay on CPU on HTP: once the linear is decomposed on CPU, its weight permute is fed only by the weight. That is handled by "Keep partitions with no runtime input on CPU", so this change lands after it. --- .../qualcomm/partition/qnn_partitioner.py | 7 +++++ backends/qualcomm/partition/utils.py | 17 ++++++++++++ backends/qualcomm/tests/test_passes.py | 27 +++++++++++++++++++ 3 files changed, 51 insertions(+) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 4f3cfd4bd2d..640d6fbf478 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -60,6 +60,7 @@ filter_fn, generate_qnn_executorch_option, get_skip_decomp_table, + has_data_dependent_shape, has_unsupported_dtype, warn_if_dilated_conv_may_not_fit_vtcm, ) @@ -179,6 +180,12 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: # noqa: C901 ) return False + if has_data_dependent_shape(node): + logger.info( + f"[{self.phase}] {node.target.__name__} | Skipped, data-dependent shape" + ) + return False + # TODO: workaround for https://github.com/pytorch/executorch/issues/23096. # Remove this once the fp16 path performs the TCM-fit check the quantized # path does. diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 7880221a1b2..19fa9b1ad55 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -184,11 +184,28 @@ def _unsupported_dtypes_are_quantized(node: torch.fx.Node) -> bool: _DQ_OPS = {op._op for op in dq_ops} +def has_data_dependent_shape(node: torch.fx.Node) -> bool: + """ + Whether the node reads or writes a tensor with a data-dependent size (e.g. + after boolean mask indexing). QNN needs a size for every dimension, and + such a dimension has no hint to give it, so the node has to stay on CPU. + """ + return any( + isinstance(dim, torch.SymInt) and dim.node.hint is None + for n in (node, *node.all_input_nodes) + for val in pytree.tree_leaves(n.meta.get("val")) + if isinstance(val, torch.Tensor) + for dim in val.shape + ) + + # Logic to determine whether to skip decompose and has higher priority than get_skip_decomp_table() def filter_fn(node: torch.fx.Node) -> bool: # The node will stay on CPU, so it must not be kept from decomposition. if has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node): return False + if has_data_dependent_shape(node): + return False # QNN does not support int32/int64 IO for the following OPs. potential_i32_i64_io_ops = [ torch.ops.aten.stack.default, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index a4ffa89826a..cca5e991338 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -1017,6 +1017,33 @@ def forward(self, x): [node.target for node in cpu_nodes], ) + def test_partitioner_falls_back_on_data_dependent_shape(self): + class MaskedStack(torch.nn.Module): + def forward(self, x, mask): + y = x[mask] + return torch.stack([y, y], dim=-1).sum((0, 2)) + x + + sample_input = (torch.randn(8, 16), torch.randn(8) > 0) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + MaskedStack().eval(), sample_input, compiler_specs + ) + except RuntimeError as e: + if "Failed to initialize QNN backend" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + cpu_ops = { + node.target + for node in edge.exported_program().graph.nodes + if node.op == "call_function" + } + self.assertIn(exir_ops.edge.aten.index.Tensor, cpu_ops) + edge.to_executorch() + def _assert_cpu_conv_keeps_its_weight_dequantize( self, quant_dtype, use_fp16, skip_node_id_set=None, emit=True ): From e8f91f79d2aeba840b7eb4a2828b5fa55504d922 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 10:39:11 -0700 Subject: [PATCH 29/37] Qualcomm AI Engine Direct - Cast a top-of-graph int64 constant before its user Passes that build a scalar constant with create_const_node (for example DecomposeRemainder for i % 3 on an int64 tensor) insert its get_attr at the very top of the graph, above the placeholders. I64toI32 then put the int32 cast right after that get_attr, also above the placeholders. When lift_constant_tensor_pass later turned the get_attr into a placeholder, it placed it with the other placeholders, below the cast, and lowering failed: "Argument _lifted_tensor_constant0 of Node aten__to_copy_default was used before it has been defined". Insert the cast for a get_attr right before its first user instead, which is below the placeholders no matter where the get_attr sits. Placeholders keep the cast right after them, as before. --- backends/qualcomm/_passes/i64_to_i32.py | 10 +++++++++- backends/qualcomm/tests/test_passes.py | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/_passes/i64_to_i32.py b/backends/qualcomm/_passes/i64_to_i32.py index a26956b5569..148564f8cc0 100644 --- a/backends/qualcomm/_passes/i64_to_i32.py +++ b/backends/qualcomm/_passes/i64_to_i32.py @@ -143,7 +143,15 @@ def _cast_constant_to_int32(self, graph_module: torch.fx.GraphModule): # If the scalar operand has an int64 data type, it should be explicitly cast to int32. node_val = n.meta["val"] if self._is_tensor_of_dtype(node_val, torch.int64): - with graph_module.graph.inserting_after(n): + # A get_attr constant can sit above the placeholders. Cast it + # right before its first user, so the cast stays below it once + # it is lifted to a placeholder. + anchor = n + if n.op == "get_attr" and n.users: + anchor = next( + user for user in graph_module.graph.nodes if user in n.users + ).prev + with graph_module.graph.inserting_after(anchor): args = (n,) to_dst_node = graph_module.graph.create_node( "call_function", diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index cca5e991338..58ce89a6e30 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -503,6 +503,26 @@ def forward(self, x): RecomposeRmsNorm()(ep.graph_module) torch.testing.assert_close(ep.module()(x), model(x)) + def test_i64_to_i32_casts_a_top_of_graph_constant_after_placeholders(self): + # DecomposeRemainder puts the int64 constant 3 at the top of the graph. + class IntRemainder(torch.nn.Module): + def forward(self, x, i): + return x * (i % 3).to(x.dtype) + + sample_input = (torch.randn(2, 4), torch.randint(0, 9, (2, 4))) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + try: + to_edge_transform_and_lower_to_qnn( + IntRemainder().eval(), sample_input, compiler_specs + ) + except RuntimeError as e: + if "Failed to initialize QNN backend" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + def test_mha_to_sha(self): from executorch.backends.qualcomm.utils.utils import convert_linear_to_conv2d from executorch.examples.models.llama.model_args import ModelArgs From a98089060a0131b730037c2d1f99201a49708bcb Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:18:28 -0700 Subject: [PATCH 30/37] Delegate ELU to XNNPACK The XNNPACK partitioner never delegated aten.elu, so ELU ran on the portable kernel even though XNNPACK has an ELU node and ExecuTorch already has its partitioner config and serializer. The config was commented out of the partitioner's config list in 2024 while waiting for a PyTorch pin update, and was never turned back on. Turn the config on. XNNPACK's ELU only has alpha, while aten.elu also has scale and input_scale, and SELU decomposes to elu with scale != 1. XNNPACK also stores alpha in the compute dtype: an alpha that overflows it gives NaN, and on x86 an infinite alpha gives NaN and a zero or subnormal alpha gives +0 where eager gives -0. So the config keeps elu nodes on the CPU kernel unless alpha is a positive normal number in the compute dtype and scale and input_scale are 1. The config and the serializer read these values from positional or keyword arguments, and the serializer handles an elu call that leaves alpha at its default. As with other XNNPACK ops, delegated results can still differ from eager for values near zero: subnormal inputs and results are flushed to zero, and the fp16 kernel returns +0 for an input of -0. The config already listed static quantization, so in a statically quantized model an elu between quantized ops is now delegated as a quantized XNNPACK ELU as well (checked to match the reference on a conv, elu, conv model). Enable the existing fp32 and fp16 ELU tests. The fp32 test also checks that SELU, an input_scale of 2 and an infinite alpha stay out of XNNPACK and that an elu without alpha is delegated, and the fp16 test checks that an alpha too large for fp16 stays out, all matching eager. --- backends/xnnpack/operators/op_elu.py | 4 +- backends/xnnpack/partition/config/__init__.py | 4 +- .../partition/config/generic_node_configs.py | 25 ++++++ backends/xnnpack/test/ops/test_elu.py | 80 +++++++++++++++++-- docs/source/backends/xnnpack/op-support.csv | 2 +- 5 files changed, 106 insertions(+), 9 deletions(-) diff --git a/backends/xnnpack/operators/op_elu.py b/backends/xnnpack/operators/op_elu.py index 9bf2f6ef0fd..ede80d0da52 100644 --- a/backends/xnnpack/operators/op_elu.py +++ b/backends/xnnpack/operators/op_elu.py @@ -39,7 +39,9 @@ def define_node( input_id = vals_to_ids[get_input_node(node, 0)] # alpha - alpha = cast(float, node.args[1]) + alpha = cast( + float, node.args[1] if len(node.args) > 1 else node.kwargs.get("alpha", 1.0) + ) # output output_id = vals_to_ids[node] diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index 7775a674510..54e2e7b8b46 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -30,7 +30,7 @@ CosConfig, DeQuantizedPerTensorConfig, DivConfig, - # EluConfig, + EluConfig, ExpConfig, FloorConfig, GeluConfig, @@ -91,7 +91,7 @@ ConvolutionConfig, ClampConfig, DivConfig, - # EluConfig, # Waiting for PyTorch Pin Update + EluConfig, ExpConfig, FloorConfig, GeluConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index d445b5ca8ac..11c8bce12f2 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -369,6 +369,31 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: def get_original_aten(self) -> Optional[torch._ops.OpOverload]: return torch.ops.aten.elu.default + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if not self.check_common_constraints(node, ep): + return False + # XNNPACK's ELU has only alpha, stored in the compute dtype. An alpha + # that overflows it gives NaN, and on x86 an infinite, zero or + # subnormal alpha is also mishandled. SELU decomposes to elu with + # scale != 1. + alpha, scale, input_scale = ( + node.args[i] if len(node.args) > i else node.kwargs.get(name, 1) + for i, name in enumerate(("alpha", "scale", "input_scale"), start=1) + ) + finfo = torch.finfo(node.meta["val"].dtype) + if ( + isinstance(alpha, bool) + or not finfo.tiny <= alpha <= finfo.max + or scale != 1 + or input_scale != 1 + ): + why( + node, + reason="ELU only supports a positive normal alpha, scale = 1 and input_scale = 1", + ) + return False + return True + class SoftmaxConfig(GenericNodePartitionerConfig): target_name = "_softmax.default" diff --git a/backends/xnnpack/test/ops/test_elu.py b/backends/xnnpack/test/ops/test_elu.py index 68a0c687779..60f375cfefa 100644 --- a/backends/xnnpack/test/ops/test_elu.py +++ b/backends/xnnpack/test/ops/test_elu.py @@ -7,7 +7,8 @@ import unittest import torch -from executorch.backends.xnnpack.test.tester import Tester +from executorch.backends.xnnpack.test.tester import RunPasses, Tester +from executorch.exir.dialects._ops import ops as exir_ops class TestElu(unittest.TestCase): @@ -26,6 +27,14 @@ class ELUFunctional(torch.nn.Module): def forward(self, x): return torch.nn.functional.elu(x, alpha=1.2) + class DefaultAlphaELU(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.elu.default(x) + + class InputScaleELU(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.elu.default(x, 1.0, 1.0, 2.0) + def _test_elu(self, inputs): ( Tester(self.ELU(), inputs) @@ -43,15 +52,76 @@ def _test_elu(self, inputs): .run_method_and_compare_outputs() ) - @unittest.skip("PyTorch Pin Update Required") - def _test_fp16_elu(self): + def test_fp16_elu(self): inputs = (torch.randn(1, 3, 3).to(torch.float16),) self._test_elu(inputs) + # An alpha that overflows fp16 would make XNNPACK return NaN. + ( + Tester(torch.nn.ELU(alpha=1e5), inputs) + .export() + .to_edge_transform_and_lower() + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) - @unittest.skip("PyTorch Pin Update Required") - def _test_fp32_elu(self): + def test_fp32_elu(self): inputs = (torch.randn(1, 3, 3),) self._test_elu(inputs) + # XNNPACK's ELU has only alpha, so an elu with scale != 1 (SELU) or + # input_scale != 1 stays on the CPU. It turns an infinite alpha into NaN. + for module in ( + torch.nn.SELU(), + self.InputScaleELU(), + torch.nn.ELU(alpha=float("inf")), + ): + with self.subTest(module=module): + ( + Tester(module, inputs) + .export() + .to_edge_transform_and_lower() + .check_not(["torch.ops.higher_order.executorch_call_delegate"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + # aten.elu called without alpha uses the default alpha of 1. + ( + Tester(self.DefaultAlphaELU(), inputs) + .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_elu_keyword_arguments(self): + def move_elu_arguments_to_kwargs(ep): + for node in ep.graph.nodes: + if node.target == exir_ops.edge.aten.elu.default: + names = ("alpha", "scale", "input_scale") + node.kwargs = dict(zip(names, node.args[1:])) + node.args = node.args[:1] + return ep + + # A graph pass can leave alpha and the scales as keyword arguments. + inputs = (torch.randn(1, 3, 3),) + for module, delegates in ((self.ELUFunctional(), 1), (torch.nn.SELU(), 0)): + ( + Tester(module, inputs) + .export() + .to_edge() + .run_passes(RunPasses(pass_functions=[move_elu_arguments_to_kwargs])) + .partition() + .check_count( + {"torch.ops.higher_order.executorch_call_delegate": delegates} + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) @unittest.skip("Update Quantizer to quantize Elu") def _test_qs8_elu(self): diff --git a/docs/source/backends/xnnpack/op-support.csv b/docs/source/backends/xnnpack/op-support.csv index 5350fed8d12..1b2077ee13b 100644 --- a/docs/source/backends/xnnpack/op-support.csv +++ b/docs/source/backends/xnnpack/op-support.csv @@ -14,7 +14,7 @@ conv2d,"fp16, fp32","PT2E: static or dynamic int8 activations 8-bit weights, symmetric per-tensor or per-channel",constant weights dequantize_per_tensor,"fp16, fp32",, div,"fp16, fp32",, -elu,"fp16, fp32",, +elu,"fp16, fp32",,"alpha positive and normal in the compute dtype, scale = 1, input_scale = 1" exp,"fp16, fp32",, floor,"fp16, fp32",, gelu,"fp16, fp32",, From 3067a86747846f7fd204ccf74d995c0acd4c3a06 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 16:50:46 -0700 Subject: [PATCH 31/37] Stop XNNPACK from dropping the scales of weight-only quantized linears A linear quantized with torchao's IntxWeightOnlyConfig gave wrong outputs after lowering to XNNPACK. With int4 weights in groups of 32 (the llama "4w" mode) the largest error on a small MLP against the eager quantized model was over 4000, and int8 per-channel weights were just as wrong. The quantized weight is a tensor subclass, so its dequantize_affine runs inside the linear call and lands in the linear's source partition. When the linear is decomposed to mm or addmm, the partitioner finds the weight by walking back from the gemm until it leaves the source partition. It walked straight through the dequant to the int8 data, saw a plain parameter, and delegated the gemm as fp32. ConvertToLinearPass does the same walk when it rebuilds the linear inside the delegate, so the int8 values were used as float weights, with the scales dropped. Stop both walks at a dequant node. With float activations the gemm is then seen as quantized, and since the XNNPACK partitioner has no path for float activations with quantized weights, it stays out of the delegate. Lowering such a model now fails loudly at to_executorch (no kernel for dequantize_affine) instead of running with wrong weights. With dynamically quantized activations the gemm is delegated and keeps its weight scales. Dynamically quantized linears from to_edge_transform_and_lower (8da4w and 8da8w: 8-bit dynamic activations with 4-bit or 8-bit weights) are unchanged, because the linear itself is delegated and is not decomposed. With the older to_edge() and to_backend() flow, 8da4w and 8da8w linears that used to be delegated with wrong weights now fail loudly the same way. Add a test with an int4 weight-only linear, with and without bias: with float activations its gemm stays out of the delegate, and with dynamically quantized activations it is delegated and matches eager. --- backends/xnnpack/_passes/convert_to_linear.py | 4 + .../xnnpack/partition/config/gemm_configs.py | 4 + backends/xnnpack/test/ops/test_linear.py | 73 +++++++++++++++++++ 3 files changed, 81 insertions(+) diff --git a/backends/xnnpack/_passes/convert_to_linear.py b/backends/xnnpack/_passes/convert_to_linear.py index 2cef71bf927..8259aae4e98 100644 --- a/backends/xnnpack/_passes/convert_to_linear.py +++ b/backends/xnnpack/_passes/convert_to_linear.py @@ -13,6 +13,7 @@ from executorch.backends.transforms.addmm_mm_to_linear import ( apply_addmm_mm_to_linear_transform, ) +from executorch.backends.xnnpack.utils.quant_utils import is_dequant from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -49,6 +50,9 @@ def find( # TODO(T171263351) - fix source partitioning for lifted graphs if not node or node in args or node.op == "placeholder": return node + # A dequant inside the partition carries the quantization, so keep it. + if is_dequant(node): + return node if kind == "args": other = node.args[index] elif kind == "users": diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index 4df6c4f1141..239a424965b 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -532,9 +532,13 @@ def get_deps_from_src_partition( """ def find_partition_args(input_node): + # A quantized tensor subclass dequantizes inside the linear call. + # Stop at that dequant, or its integer input is taken for a float + # weight. while ( len(input_node.all_input_nodes) != 0 and input_node not in src_partition.input_nodes + and not is_dequant(input_node) ): input_node = input_node.all_input_nodes[0] return input_node diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index e309bc514a7..7e4cb4ab586 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -37,6 +37,7 @@ from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import ( Int8DynamicActivationIntxWeightConfig, + IntxWeightOnlyConfig, quantize_, ) from torchao.utils import unwrap_tensor_subclass @@ -951,6 +952,78 @@ def test_qd8_linear_transposed_scale(self): .run_method_and_compare_outputs() ) + @unittest.skipIf(not torchao_installed, "Requires TorchAO") + def test_linear_int4_weight_only_dequant_in_source_partition(self): + # An int4 weight-only tensor subclass dequantizes inside the linear + # call, so its dequant is inside the linear's source partition and + # must not be skipped when looking for the weight. + # Activations are quantized by hand so that their quantize and + # dequantize stay outside the linear's source partition. + class DynamicActivationLinear(torch.nn.Module): + def __init__(self, linear): + super().__init__() + self.linear = linear + + def forward(self, x): + block, qmin, qmax = [1, 64], -128, 127 + scale, zero_point = torch.ops.torchao.choose_qparams_affine( + x, + "ASYMMETRIC", + block, + torch.int8, + qmin, + qmax, + None, + None, + torch.int8, + ) + q = torch.ops.torchao.quantize_affine( + x, block, scale, zero_point, torch.int8, qmin, qmax + ) + dq = torch.ops.torchao.dequantize_affine( + q, block, scale, zero_point, torch.int8, qmin, qmax + ) + return self.linear(dq) + + for use_bias, gemm in ( + (False, "executorch_exir_dialects_edge__ops_aten_mm_default"), + (True, "executorch_exir_dialects_edge__ops_aten_addmm_default"), + ): + mod = torch.nn.Linear(64, 32, bias=use_bias) + quantize_( + mod, + IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)), + ) + unwrap_tensor_subclass(mod) + inputs = (torch.randn(2, 64),) + # The XNNPACK partitioner does not delegate float activations with + # quantized weights, so the gemm must stay out of the delegate. + ( + Tester(mod, inputs) + .export() + .to_edge_transform_and_lower() + .check_count({gemm: 1}) + ) + # With dynamically quantized activations it is delegated and has + # to keep the weight scales. + ( + Tester(DynamicActivationLinear(mod), inputs) + .export() + .to_edge() + .partition( + Partition( + XnnpackPartitioner( + config_precisions=ConfigPrecisionType.DYNAMIC_QUANT + ) + ) + ) + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not([gemm]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs(atol=1e-2, rtol=1e-2) + ) + def test_qd8_per_channel_linear_parallel(self): in_size = 2 input_size = 4 From 547bf45aaf54017ae331b2e2de4f447bb773a22a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 16:41:43 -0700 Subject: [PATCH 32/37] Add a portable kernel for cummin.out torch.cummin has no ExecuTorch kernel. It is not part of the Core ATen opset and has no decomposition, so to_edge rejects it by default; a program that keeps it anyway (with cummin in _core_aten_ops_exception_list, or with IR validity checks turned off) then fails to load with "kernel 'aten::cummin.out' not found". This happens whenever cummin is left to the CPU, for example in a model lowered with XNNPACK or OpenVINO, or a TensorRT program whose cummin falls outside the engine. Add cummin.out to the portable kernels, following cumsum. For each position it writes the running minimum along dim and the index where that minimum was found. It matches PyTorch: on a tie the index moves to the later element, once a NaN appears it stays the minimum (with the index of the latest NaN), and a zero-dimensional input works with dim 0 or -1. All real types, Bool, Half and BFloat16 are supported. Test: op_cummin_test.cpp checks values and indices along both dims for every real type, Bool, NaN propagation, a zero-dimensional input, and that an invalid dim fails. --- kernels/aten/functions.yaml | 2 + kernels/portable/cpu/op_cummin.cpp | 94 ++++++++++++++++ kernels/portable/functions.yaml | 5 + kernels/test/CMakeLists.txt | 1 + kernels/test/op_cummin_test.cpp | 102 ++++++++++++++++++ kernels/test/targets.bzl | 1 + .../executorch/build/build_variables.bzl | 1 + .../kernels/portable/op_registration_util.bzl | 6 ++ 8 files changed, 212 insertions(+) create mode 100644 kernels/portable/cpu/op_cummin.cpp create mode 100644 kernels/test/op_cummin_test.cpp diff --git a/kernels/aten/functions.yaml b/kernels/aten/functions.yaml index 36e79f0dc43..38bc6440c80 100644 --- a/kernels/aten/functions.yaml +++ b/kernels/aten/functions.yaml @@ -133,6 +133,8 @@ - op: cosh.out +- op: cummin.out + - op: cumsum.out - op: dequantize.self_out diff --git a/kernels/portable/cpu/op_cummin.cpp b/kernels/portable/cpu/op_cummin.cpp new file mode 100644 index 00000000000..4e7bb44f47a --- /dev/null +++ b/kernels/portable/cpu/op_cummin.cpp @@ -0,0 +1,94 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using Tensor = executorch::aten::Tensor; +using ScalarType = executorch::aten::ScalarType; + +/** + * Computes the running minimum of self along dim, and the index of that + * minimum. Like PyTorch, a tie moves the index to the later element and a NaN + * becomes the minimum from where it appears on. + */ +std::tuple cummin_out( + KernelRuntimeContext& ctx, + const Tensor& self, + int64_t dim, + Tensor& values, + Tensor& indices) { + std::tuple ret_val(values, indices); + + ET_KERNEL_CHECK( + ctx, dim_is_valid(dim, nonzero_dim(self)), InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dtype(self, values), InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, indices.scalar_type() == ScalarType::Long, InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, + tensor_is_default_dim_order(self) && + tensor_is_default_dim_order(values) && + tensor_is_default_dim_order(indices), + InvalidArgument, + ret_val); + ET_KERNEL_CHECK( + ctx, + resize_tensor(values, self.sizes()) == Error::Ok && + resize_tensor(indices, self.sizes()) == Error::Ok, + InvalidArgument, + ret_val); + + if (self.numel() == 0) { + return ret_val; + } + + dim = dim < 0 ? dim + nonzero_dim(self) : dim; + const size_t dim_size = nonempty_size(self, dim); + const size_t leading_dims = getLeadingDims(self, dim); + const size_t trailing_dims = self.dim() == 0 ? 1 : getTrailingDims(self, dim); + + ET_SWITCH_REALHBBF16_TYPES( + self.scalar_type(), ctx, "cummin.out", CTYPE, [&]() { + const CTYPE* in_data = self.const_data_ptr(); + CTYPE* values_data = values.mutable_data_ptr(); + int64_t* indices_data = indices.mutable_data_ptr(); + + for (size_t i = 0; i < leading_dims; i++) { + const size_t start = i * dim_size * trailing_dims; + for (size_t j = 0; j < dim_size; j++) { + const size_t cur = start + j * trailing_dims; + const size_t prev = cur - trailing_dims; + for (size_t k = 0; k < trailing_dims; k++) { + const CTYPE v = in_data[cur + k]; + if (j == 0 || utils::isnan_override(v) || + (!utils::isnan_override(values_data[prev + k]) && + v <= values_data[prev + k])) { + values_data[cur + k] = v; + indices_data[cur + k] = static_cast(j); + } else { + values_data[cur + k] = values_data[prev + k]; + indices_data[cur + k] = indices_data[prev + k]; + } + } + } + } + }); + + return ret_val; +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index 97f60f75c02..e67976fac40 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -323,6 +323,11 @@ - arg_meta: null kernel_name: torch::executor::cosh_out +- op: cummin.out + kernels: + - arg_meta: null + kernel_name: torch::executor::cummin_out + - op: cumsum.out kernels: - arg_meta: null diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index 86f3aa07319..029c9252998 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -197,6 +197,7 @@ set(all_test_sources "op_copy_test.cpp" "op_cos_test.cpp" "op_cosh_test.cpp" + "op_cummin_test.cpp" "op_cumsum_test.cpp" "op_detach_copy_test.cpp" "op_diagonal_copy_test.cpp" diff --git a/kernels/test/op_cummin_test.cpp b/kernels/test/op_cummin_test.cpp new file mode 100644 index 00000000000..7ba11762ccb --- /dev/null +++ b/kernels/test/op_cummin_test.cpp @@ -0,0 +1,102 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include // Declares the operator +#include +#include +#include +#include + +using namespace executorch::aten; +using torch::executor::testing::TensorFactory; + +class OpCumminOutTest : public OperatorTest { + protected: + std::tuple op_cummin_out( + const Tensor& self, + int64_t dim, + Tensor& values, + Tensor& indices) { + return torch::executor::aten::cummin_outf( + context_, self, dim, values, indices); + } + + template + void test_dtype() { + TensorFactory tf; + TensorFactory tf_long; + // clang-format off + Tensor in = tf.make({2, 4}, { + 3, 1, 2, 1, + 0, 5, 0, 4}); + // clang-format on + Tensor values = tf.zeros({2, 4}); + Tensor indices = tf_long.zeros({2, 4}); + + op_cummin_out(in, /*dim=*/1, values, indices); + // A tie moves the index to the later element, as in PyTorch. + EXPECT_TENSOR_EQ(values, tf.make({2, 4}, {3, 1, 1, 1, 0, 0, 0, 0})); + EXPECT_TENSOR_EQ(indices, tf_long.make({2, 4}, {0, 1, 1, 3, 0, 0, 2, 2})); + + op_cummin_out(in, /*dim=*/-2, values, indices); + EXPECT_TENSOR_EQ(values, tf.make({2, 4}, {3, 1, 2, 1, 0, 1, 0, 1})); + EXPECT_TENSOR_EQ(indices, tf_long.make({2, 4}, {0, 0, 0, 0, 1, 0, 1, 0})); + } +}; + +TEST_F(OpCumminOutTest, AllRealDtypesSupported) { +#define TEST_ENTRY(ctype, dtype) test_dtype(); + ET_FORALL_REALHBF16_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpCumminOutTest, NanBecomesTheMinimum) { + TensorFactory tf; + TensorFactory tf_long; + const float nan = NAN; + Tensor in = tf.make({5}, {2, nan, 1, nan, 0}); + Tensor values = tf.zeros({5}); + Tensor indices = tf_long.zeros({5}); + + op_cummin_out(in, /*dim=*/0, values, indices); + EXPECT_TENSOR_CLOSE(values, tf.make({5}, {2, nan, nan, nan, nan})); + EXPECT_TENSOR_EQ(indices, tf_long.make({5}, {0, 1, 1, 3, 3})); +} + +TEST_F(OpCumminOutTest, Bool) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.make({4}, {true, false, true, false}); + Tensor values = tf.zeros({4}); + Tensor indices = tf_long.zeros({4}); + + op_cummin_out(in, /*dim=*/0, values, indices); + EXPECT_TENSOR_EQ(values, tf.make({4}, {true, false, false, false})); + EXPECT_TENSOR_EQ(indices, tf_long.make({4}, {0, 1, 1, 3})); +} + +TEST_F(OpCumminOutTest, ZeroDim) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.make({}, {2}); + Tensor values = tf.zeros({}); + Tensor indices = tf_long.ones({}); + + op_cummin_out(in, /*dim=*/-1, values, indices); + EXPECT_TENSOR_EQ(values, in); + EXPECT_TENSOR_EQ(indices, tf_long.zeros({})); +} + +TEST_F(OpCumminOutTest, InvalidDimDies) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.ones({2, 3}); + Tensor values = tf.zeros({2, 3}); + Tensor indices = tf_long.zeros({2, 3}); + ET_EXPECT_KERNEL_FAILURE(context_, op_cummin_out(in, 2, values, indices)); +} diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 7555c14c4a8..8289a1ddd5f 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -248,6 +248,7 @@ def define_common_targets(): _common_op_test("op_copy_test", ["aten", "portable"]) _common_op_test("op_cos_test", ["aten", "portable"]) _common_op_test("op_cosh_test", ["aten", "portable"]) + _common_op_test("op_cummin_test", ["aten", "portable"]) _common_op_test("op_cumsum_test", ["aten", "portable"]) _common_op_test("op_detach_copy_test", ["aten", "portable"]) _common_op_test("op_diagonal_copy_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index 82d0c40064e..b0e7f7e7065 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -112,6 +112,7 @@ PORTABLE_KERNELS_SRCS = [ "kernels/portable/cpu/op_copy.cpp", "kernels/portable/cpu/op_cos.cpp", "kernels/portable/cpu/op_cosh.cpp", + "kernels/portable/cpu/op_cummin.cpp", "kernels/portable/cpu/op_cumsum.cpp", "kernels/portable/cpu/op_detach_copy.cpp", "kernels/portable/cpu/op_diagonal_copy.cpp", diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index f3d5739a862..733eed5a780 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -502,6 +502,12 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/pattern:pattern", ], ), + op_target( + name = "op_cummin", + deps = [ + "//executorch/kernels/portable/cpu/util:math_util", + ], + ), op_target( name = "op_cumsum", deps = [ From 9966268b8fe3f88383ced6a21ee9c3c2fd99c51b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 15:54:52 -0700 Subject: [PATCH 33/37] OpenVINO backend: convert partitions with their input types The OpenVINO backend converted each partition with openvino_compile, a helper written for torch.compile. That helper converts the graph first and sets the input types only afterwards. Two things break because of that: 1. Indexing with a boolean mask, like x[mask] or x[mask] = v, is lowered wrong. OpenVINO's converter looks at the index type to choose the boolean-mask lowering, but during conversion the type is still unknown, so the mask is treated as integer indices. A read gives wrong results, and a write fails with "Check 'args_et.is_dynamic() || is_supported_et' failed" once the boolean type is set. 2. A bfloat16 partition input fails with "KeyError: torch.bfloat16", because the helper's type table has no bfloat16 entry. Now the backend builds OpenVINO's PyTorch FX decoder with the partition input types and converts it with the same PyTorch frontend calls the helper used, then compiles the result with openvino.compile_model on the device from the compile specs. The converter sees the real types from the start, and its own type table already covers bfloat16. The helper's other options (model caching and an extra config dictionary) cannot be reached from ExecuTorch compile specs, so nothing else changes. A partition input can also be a SymInt, such as a data-dependent size. It has no dtype, so it is given to the converter as an int64 scalar, the form OpenVINO's FX translators expect. The runtime passes such an input as an int64 tensor of shape [1], so a PrePostProcessor step declares it that way and squeezes it back to a scalar. The runtime takes no other non-tensor input, so a float or bool one now fails at preprocess with a clear error instead of when the program runs. The new tests write through a boolean mask with float32 and bfloat16 data, and slice and view a tensor by a data-dependent size. They fail before this change and pass after it. --- backends/openvino/preprocess.py | 42 ++++++++++++++++--- backends/openvino/tests/ops/test_index_put.py | 28 +++++++++++++ backends/openvino/tests/ops/test_view.py | 11 +++++ 3 files changed, 75 insertions(+), 6 deletions(-) create mode 100644 backends/openvino/tests/ops/test_index_put.py diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index de2c35d2414..7a3f490b258 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -8,6 +8,7 @@ from typing import final, List +import torch from executorch.backends.openvino._passes import DecomposeFloorDividePass from executorch.exir.backend.backend_details import ( @@ -17,9 +18,12 @@ ) from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.passes.memory_format_ops_pass import DimOrderOpsRevertPass -from openvino.frontend.pytorch.torchdynamo.compile import ( # type: ignore[import-untyped] - openvino_compile, +from openvino import compile_model, opset13 # type: ignore[import-untyped] +from openvino.frontend import FrontEndManager # type: ignore[import-untyped] +from openvino.frontend.pytorch.fx_decoder import ( # type: ignore[import-untyped] + TorchFXPythonDecoder, ) +from openvino.preprocess import PrePostProcessor # type: ignore[import-untyped] @final @@ -45,18 +49,44 @@ def preprocess( edge_program._graph_module = result.graph_module input_names = edge_program.graph_signature.user_inputs - args = [] + input_types = [] + input_shapes = [] + int_inputs = [] for node in edge_program.graph.nodes: if node.target in input_names: - args.append(node.meta["val"]) + val = node.meta["val"] + if isinstance(val, torch.Tensor): + input_types.append(val.dtype) + input_shapes.append( + [d if isinstance(d, int) else -1 for d in val.shape] + ) + elif isinstance(val, (int, torch.SymInt)) and not isinstance(val, bool): + int_inputs.append(len(input_types)) + input_types.append(torch.int64) + input_shapes.append([]) + else: + raise RuntimeError( + f"OpenVINO backend: partition input {node.target} is a " + f"{type(val).__name__}; only tensors and ints are supported" + ) compile_options = {} for spec in module_compile_spec: compile_options[spec.key] = spec.value.decode() - compiled = openvino_compile( - edge_program.module(), *args, options=compile_options + # The converter chooses some lowerings (e.g. boolean-mask indexing) from + # the input types, so they must be known before conversion starts. + decoder = TorchFXPythonDecoder( + edge_program.module(), input_shapes=input_shapes, input_types=input_types ) + frontend = FrontEndManager().load_by_framework("pytorch") + ppp = PrePostProcessor(frontend.convert(frontend.load(decoder))) + # The converter takes an int input, like a SymInt size, as a scalar, but + # the runtime passes it as an int64 tensor of shape [1]. + for i in int_inputs: + ppp.input(i).tensor().set_shape([1]) + ppp.input(i).preprocess().custom(lambda x: opset13.squeeze(x, 0).output(0)) + compiled = compile_model(ppp.build(), compile_options.get("device", "CPU")) model_bytes = compiled.export_model() return PreprocessResult(processed_bytes=model_bytes.getvalue()) diff --git a/backends/openvino/tests/ops/test_index_put.py b/backends/openvino/tests/ops/test_index_put.py new file mode 100644 index 00000000000..238ffbc3206 --- /dev/null +++ b/backends/openvino/tests/ops/test_index_put.py @@ -0,0 +1,28 @@ +import torch +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) + + +class TestIndexPutOperator(BaseOpenvinoOpTest): + + def create_model(self): + class IndexPut(torch.nn.Module): + def forward(self, x, mask, values): + x = x.clone() + x[mask] = values[mask] + return x + + return IndexPut() + + def test_index_put_bool_mask(self): + mask = torch.rand(2, 8) > 0.5 + for dtype in [torch.float32, torch.bfloat16]: + with self.subTest(dtype=dtype): + module = self.create_model() + sample_input = ( + torch.zeros(2, 8, dtype=dtype), + mask, + torch.randn(2, 8).to(dtype), + ) + self.execute_layer_test(module, sample_input) diff --git a/backends/openvino/tests/ops/test_view.py b/backends/openvino/tests/ops/test_view.py index 8aef13fffa0..5ea2f8cb1dc 100644 --- a/backends/openvino/tests/ops/test_view.py +++ b/backends/openvino/tests/ops/test_view.py @@ -33,3 +33,14 @@ def test_view(self): sample_input = (torch.randn(params["input_shape"]),) self.execute_layer_test(module, sample_input) + + def test_view_data_dependent_size(self): + class View(torch.nn.Module): + def forward(self, x, n): + n = n.item() + torch._check(n >= 1) + torch._check(n <= x.shape[0]) + return x[:n].view(n, -1).sum(0) + + # The size reaches the OpenVINO partition as a SymInt input. + self.execute_layer_test(View(), (torch.randn(4, 8), torch.tensor(3))) From 82081dcf98989bc75b2a1758900a3a967c37ef3c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 15:52:09 -0700 Subject: [PATCH 34/37] CUDA backend: lower transposed convolution without ATen Exporting a model with ConvTranspose1d or ConvTranspose2d to the CUDA backend fails at the end of the AOTInductor compile with "missing fallback kernels: aoti_torch_cuda_convolution". This backend lets Inductor use only Triton kernels, and on CUDA Inductor has no Triton template for transposed convolution (it has none for 1-D, and the 2-D one is disabled on CUDA), so it falls back to ATen's convolution, which the ExecuTorch CUDA runtime does not have. The diffusion policy in LeRobot (ConditionalUnet1D upsampling) hits this. Add aten.conv_transpose1d and aten.conv_transpose2d to the backend's decomposition table, next to the existing conv1d entry. For finite weights, a transposed convolution equals a stride 1 convolution with the flipped kernel (input and output channels swapped per group) over the input with stride - 1 zeros inserted between elements and padded on each side by dilation * (kernel - 1) minus padding (plus output_padding at the end). Inductor generates Triton code for the zero insertion, the padding and the resulting convolution. The inserted zeros cost extra work (about stride times the multiply-adds per spatial dimension), which is the price of having a Triton path at all. On ROCm, Inductor still uses its Triton template for 2-D transposed convolution, so the 2-D entry is only added for CUDA. A transposed convolution with dilation > 1 becomes a dilated convolution, which has no Triton template, so on CUDA it still reports the missing kernel, as a regular dilated convolution does today. Test: test_conv_transpose lowers a ConvTranspose1d to the CUDA backend and checks the decompositions against conv_transpose1d and conv_transpose2d with stride, padding, output_padding, groups and dilation set, an unbatched input, and padding that crops more than the input. --- backends/cuda/cuda_backend.py | 74 ++++++++++++++++++++++++- backends/cuda/tests/test_cuda_export.py | 38 +++++++++++++ 2 files changed, 111 insertions(+), 1 deletion(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 787d2502b63..052c480f1c4 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -482,6 +482,72 @@ def _moved_program_tensors() -> List[torch.Tensor]: return [t for t in live if t is not None] +def _conv_transpose_as_conv( + input: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, + stride=1, + padding=0, + output_padding=0, + groups: int = 1, + dilation=1, +) -> torch.Tensor: + """ + Rewrite a 1-D or 2-D transposed convolution as a stride 1 convolution with + the flipped kernel, over the input with stride - 1 zeros inserted between + elements and padded on each side. Inductor has no Triton template for + transposed convolution on CUDA, so it would call ATen's convolution, which + this backend does not ship. + """ + n = weight.dim() - 2 + if input.dim() == n + 1: + return _conv_transpose_as_conv( + input.unsqueeze(0), + weight, + bias, + stride, + padding, + output_padding, + groups, + dilation, + ).squeeze(0) + options = [] + for v in (stride, padding, output_padding, dilation): + v = [v] if isinstance(v, int) else list(v) + options.append(v * n if len(v) == 1 else v) + stride, padding, output_padding, dilation = options + x = input + pad = [] + crops = [] + for i, dim in enumerate(range(2, 2 + n)): + if stride[i] > 1: + length = x.shape[dim] + zeros = x.new_zeros(*x.shape[: dim + 1], stride[i] - 1, *x.shape[dim + 1 :]) + x = torch.cat([x.unsqueeze(dim + 1), zeros], dim + 1).flatten(dim, dim + 1) + x = x.narrow(dim, 0, (length - 1) * stride[i] + 1) + before = dilation[i] * (weight.shape[dim] - 1) - padding[i] + after = before + output_padding[i] + # constant_pad_nd takes the last dimension first. Negative padding is + # applied as a crop after padding, since cropping first can need more + # elements than the input has. + pad = [max(before, 0), max(after, 0)] + pad + if before < 0 or after < 0: + crops.append((dim, max(-before, 0), x.shape[dim] + before + after)) + if any(pad): + x = torch.ops.aten.constant_pad_nd.default(x, pad) + for dim, start, length in crops: + x = x.narrow(dim, start, length) + # (in, out / groups, *kernel) -> (out, in / groups, *kernel) + weight = ( + weight.unflatten(0, (groups, -1)) + .transpose(1, 2) + .flatten(0, 1) + .flip(list(range(2, 2 + n))) + ) + conv = conv1d_to_conv2d if n == 1 else torch.ops.aten.conv2d.default + return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups) + + @final @experimental( "This API and all of cuda backend related functionality are experimental." @@ -777,9 +843,15 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: @classmethod def get_decomposition_table(cls) -> Dict[Any, Any]: - return { + table = { torch.ops.aten.conv1d.default: conv1d_to_conv2d, + torch.ops.aten.conv_transpose1d.default: _conv_transpose_as_conv, } + # Inductor disables its transposed conv2d Triton template only on CUDA + # (https://github.com/pytorch/pytorch/issues/187081). + if torch.version.hip is None: + table[torch.ops.aten.conv_transpose2d.input] = _conv_transpose_as_conv + return table @classmethod def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index b49fa5ac753..4bf431c4c49 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -558,6 +558,44 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: edge_program_manager = self._export_to_cuda_with_lower(module, inputs) self.assertIsNotNone(edge_program_manager, "Conv1d operation export failed") + def test_conv_transpose(self): + """Test CUDA export for transposed convolution.""" + module = torch.nn.ConvTranspose1d(4, 6, kernel_size=4, stride=2, padding=1) + module.eval() + inputs = (torch.randn(1, 4, 10),) + + edge_program_manager = self._export_to_cuda_with_lower(module, inputs) + self.assertIsNotNone( + edge_program_manager, "ConvTranspose1d operation export failed" + ) + + # Unbatched input, a crop larger than the input, and uneven 2D options. + cases = [ + ( + torch.ops.aten.conv_transpose1d.default, + torch.nn.functional.conv_transpose1d, + (torch.randn(4, 1), torch.randn(4, 3, 1), torch.randn(6)), + ([10], [4], [9], 2, [1]), + ), + ( + torch.ops.aten.conv_transpose2d.input, + torch.nn.functional.conv_transpose2d, + (torch.randn(2, 4, 5, 6), torch.randn(4, 3, 3, 2), torch.randn(6)), + ([3, 2], [2, 0], [1, 1], 2, [2, 1]), + ), + ] + table = CudaBackend.get_decomposition_table() + for op, reference, tensors, options in cases: + # ROCm leaves transposed conv2d to Inductor's own template. + if ( + torch.version.hip is not None + and op == torch.ops.aten.conv_transpose2d.input + ): + continue + torch.testing.assert_close( + table[op](*tensors, *options), reference(*tensors, *options) + ) + def test_sdpa_single_kernel(self): """ Test CUDA export for model containing single SDPA kernel. From 455b977eb6a7f195ac27f9ea4da37106bf11994f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 07:42:02 -0700 Subject: [PATCH 35/37] Fix constant data of duplicated constant placeholders When a constant is used by more than one partition, duplicate_constant_node gives every extra user of the constant its own placeholder, and every copy keeps the target of the original. So several placeholders can point at the same state_dict or constants entry, and two of them can end up in the same partition. exir already handles this when it removes delegated constants (it skips targets that are still in use). delete_constant_placeholder deleted the entry on the first call. That left the other placeholders pointing at a missing tensor, and deleting a second placeholder with the same target failed with a KeyError. The Arm FuseEqualPlaceholdersPass hits this when it fuses such copies, for example for a weight reused around an op that stays on the CPU: KeyError: 'lin.bias' Now the entry is removed only when no remaining placeholder refers to its target, the same rule exir uses when it removes delegated constants. duplicate_constant_node also stored a deep copy of the tensor under each copy's node name (for example p_weight_copy_0). No input spec refers to that key, so the data was never read, and it stayed in the program after every placeholder was deleted. Stop storing it. --- ...test_create_delete_constant_placeholder.py | 43 +++++++++++++++++++ backends/transforms/utils.py | 15 ++++--- .../duplicate_constant_node_pass.py | 21 ++------- 3 files changed, 56 insertions(+), 23 deletions(-) diff --git a/backends/transforms/test/test_create_delete_constant_placeholder.py b/backends/transforms/test/test_create_delete_constant_placeholder.py index 9444e6a4b63..b85a14b036f 100644 --- a/backends/transforms/test/test_create_delete_constant_placeholder.py +++ b/backends/transforms/test/test_create_delete_constant_placeholder.py @@ -9,6 +9,9 @@ delete_constant_placeholder, ) from executorch.exir import to_edge +from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import ( + duplicate_constant_node, +) from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export from torch.export.graph_signature import InputKind @@ -190,6 +193,46 @@ def test_create_same_name_returns_existing_node(): assert len(exported_program.state_dict) == 1 +def test_delete_placeholders_sharing_target(): + """ + duplicate_constant_node gives every copy the target of the original, so the + tensor must stay until the last placeholder referring to it is deleted. + """ + + class SharedWeight(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + + def forward(self, x): + return x + self.weight, x - self.weight + + exported_program = to_edge( + export(SharedWeight(), (torch.zeros(1),), strict=True) + ).exported_program() + (weight_name,) = exported_program.graph_signature.inputs_to_parameters + duplicate_constant_node(exported_program, weight_name) + + weight_nodes = [ + node + for node in exported_program.graph.nodes + if node.name in exported_program.graph_signature.inputs_to_parameters + ] + assert len(weight_nodes) == 2 + input_node = next( + node + for node in exported_program.graph.nodes + if node.name in exported_program.graph_signature.user_inputs + ) + for node in weight_nodes: + node.replace_all_uses_with(input_node) + + delete_constant_placeholder(exported_program, weight_nodes[0]) + assert "weight" in exported_program.state_dict + delete_constant_placeholder(exported_program, weight_nodes[1]) + assert len(exported_program.state_dict) == 0 + + def test_create_delete_parameter(): _test_create_delete(InputKind.PARAMETER) diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index d3c50c692b1..92794364cb0 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -243,24 +243,24 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod f"Cannot delete input node {node.name} since it has users in the graph." ) - # Remove tensor from state_dict/ constants + # Find the state_dict/ constants entry holding the tensor if node.name in exp_program.graph_signature.inputs_to_parameters: target = exp_program.graph_signature.inputs_to_parameters[node.name] - del exp_program.state_dict[target] + tensors = exp_program.state_dict elif node.name in exp_program.graph_signature.inputs_to_buffers: target = exp_program.graph_signature.inputs_to_buffers[node.name] if target in exp_program.graph_signature.non_persistent_buffers: - del exp_program.constants[target] + tensors = exp_program.constants else: - del exp_program.state_dict[target] + tensors = exp_program.state_dict elif node.name in exp_program.graph_signature.inputs_to_lifted_tensor_constants: target = exp_program.graph_signature.inputs_to_lifted_tensor_constants[ node.name ] - del exp_program.constants[target] + tensors = exp_program.constants else: raise RuntimeError( f"Cannot delete input node {node.name} since it is not a parameter, a buffer, nor a lifted tensor constant." @@ -277,6 +277,11 @@ def delete_constant_placeholder(exp_program: ExportedProgram, node: torch.fx.Nod ) exp_program._graph_signature = new_graph_signature + # Copies made by duplicate_constant_node keep the original target, so keep + # the tensor while another placeholder still refers to it. + if all(spec.target != target for spec in input_specs): + del tensors[target] + # Remove node from graph node.graph.erase_node(node) diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py index 68e4dba42e4..07c3eb5e57a 100644 --- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py +++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py @@ -13,7 +13,6 @@ from torch.export import ExportedProgram from torch.export.exported_program import InputSpec, TensorArgument -from torch.export.graph_signature import InputKind def _get_attribute_or_constants( @@ -44,7 +43,7 @@ def duplicate_constant_node( Args: exported_program: the exported program to be modified. If constants nodes are copied, they will be added as new - placeholder and the state_dict will be updated + placeholders that share the original target candidate_node: the name of the constant node to be duplicated Returns: @@ -117,22 +116,8 @@ def duplicate_constant_node( } users[ith].args = new_args users[ith].kwargs = new_kwargs - if old_input_spec.kind == InputKind.CONSTANT_TENSOR: - exported_program.constants[copy_constant_or_attribute_fqn] = ( - copy.deepcopy(constant_or_attribute) - ) - elif ( - old_input_spec.kind == InputKind.BUFFER - and old_input_spec.persistent is False - ): - # non persistent buffer will be in the .constants - exported_program.constants[copy_constant_or_attribute_fqn] = ( - copy.deepcopy(constant_or_attribute) - ) - else: - exported_program.state_dict[copy_constant_or_attribute_fqn] = ( - copy.deepcopy(constant_or_attribute) - ) + # The copy keeps the original target, so it reads the + # original tensor and needs no data of its own. new_input_specs.append( InputSpec( kind=old_input_spec.kind, From 75cdf6690268bcfa9c80d7114e32be69a999eb0c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 13:59:04 -0700 Subject: [PATCH 36/37] Lower torch._check_tensor_all in the CUDA backend A model with a tensor-wide check, for example torch._check_tensor_all or transformers' torch_compilable_check, could not be lowered to the CUDA backend: Method forward missing fallback kernels: aten::_is_all_true torch.export turns the check into aten._is_all_true. The AOTI partitioner keeps every op for Inductor to handle, but Inductor has no lowering for this one, and ExecuTorch's CUDA shim has no fallback kernel for it. Add it to the CUDA backend's decomposition table as aten.all, which Inductor lowers, next to the existing conv1d entry. --- backends/cuda/cuda_backend.py | 3 +++ backends/cuda/tests/test_cuda_export.py | 17 ++++++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 052c480f1c4..0bd1599cd5a 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -846,6 +846,9 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: table = { torch.ops.aten.conv1d.default: conv1d_to_conv2d, torch.ops.aten.conv_transpose1d.default: _conv_transpose_as_conv, + # torch._check_tensor_all exports to this op; Inductor has no lowering + # for it and the shim no fallback kernel. + torch.ops.aten._is_all_true.default: torch.ops.aten.all.default, } # Inductor disables its transposed conv2d Triton template only on CUDA # (https://github.com/pytorch/pytorch/issues/187081). diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 4bf431c4c49..d962db09117 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -266,6 +266,7 @@ def _export_to_cuda_with_lower( module: torch.nn.Module, inputs: Tuple[torch.Tensor, ...], compile_specs: list[CompileSpec] | None = None, + strict: bool = True, ) -> None: """Helper method to export a module to CUDA backend using to_edge_transform_and_lower. @@ -274,9 +275,10 @@ def _export_to_cuda_with_lower( inputs: The example inputs for the module compile_specs: Optional list of compile specs. If not provided, defaults to only the method name compile spec for "forward" + strict: Whether to export in strict mode """ # Export the model - exported_program = export(module, inputs, strict=True) + exported_program = export(module, inputs, strict=strict) # Create partitioner with compile specs if compile_specs is None: @@ -366,6 +368,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) edge_program_manager.to_executorch() + def test_check_tensor_all(self): + """Test CUDA export of a tensor-wide runtime check.""" + + class CheckModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + torch._check_tensor_all(x > -1e9) + return x * 2 + + # A strict export does not keep the check. + self._export_to_cuda_with_lower( + CheckModule().eval(), (torch.randn(4, 8),), strict=False + ) + def test_conv2d(self): """Test CUDA export for 2D convolution.""" From 204bbc81fc018b5015a3d7cc17325d31753e0fc2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 21:47:49 -0700 Subject: [PATCH 37/37] Refuse another program's CUDA weights file instead of running with it Every CUDA program writes its weights to a file named aoti_cuda_blob.ptd, and keys each weight by its parameter name. When two programs with the same parameter names are exported into one folder, the second export overwrites the first program's weights file. The first program still loads, finds every key it asks for in the other program's file, and silently runs with the other program's weights. For example, a Linear layer exported twice with different seeds: the first program then returned the second one's outputs. Add each weight's SHA-256, which the exporter already computes while it writes the weight, to the weight's key. A program then finds only keys written for its own weights, and loading it with another program's file fails with "CUDA FQN storage ... is missing from named data". The bytes are not hashed again at load, so this catches files that belong to another program, not files corrupted in place. Identical weights still get the same key, so methods that share a weight still store it once. Within one export, a later method that stores different bytes under the same name is still rejected. The file name stays aoti_cuda_blob.ptd, because runners and scripts look for it. Two programs still need their own folders or file names, but a mix-up now fails at load instead of giving wrong results. This also means a .ptd with different weights, for example from a fine-tuned model, needs the .pte exported with it. test_missing_weights_blob.py exports two programs with different weights into one folder and checks that the first one refuses the second one's file. Without this change it loads. Tests that read the stored keys now take them from the delegate's manifest, which is what the runtime reads. --- backends/cuda/cuda_weight_collector.py | 20 +++++++++++++- backends/cuda/tests/test_cuda_partitioner.py | 27 +++++++++++++------ backends/cuda/tests/test_cuda_weight_reuse.py | 4 +-- .../cuda/tests/test_missing_weights_blob.py | 23 +++++++++++++--- 4 files changed, 60 insertions(+), 14 deletions(-) diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 3921e3db0d4..19e9f8c8bb6 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -389,6 +389,8 @@ def __init__(self) -> None: # with the first one on the device and not written again. The weights # must not change until finish() releases them. self._stored_storages: Dict[str, torch.Tensor] = {} + # The key each storage was written under, which ends in its SHA-256. + self._content_keys: Dict[str, str] = {} @contextlib.contextmanager def capture( @@ -599,12 +601,27 @@ def add_preprocess_result( ), ) if data is not None: + # Every program writes its weights to the same file name, so the + # key carries the content: another program's file then fails to + # load instead of running this program with its weights. + content_key = f"{entry.storage_key}:{data.sha256().hex()}" + if ( + self._content_keys.setdefault(entry.storage_key, content_key) + != content_key + ): + data.close() + raise ValueError( + f"CUDA FQN weight {entry.fqn!r} is stored with different data" + ) + entry = replace(entry, storage_key=content_key) self._add_weight(entry, data, external_tag) elif entry.storage_key in artifact.stored: - if entry.storage_key not in self._store.key_to_buffer_idx: + content_key = self._content_keys.get(entry.storage_key) + if content_key not in self._store.key_to_buffer_idx: raise RuntimeError( f"CUDA FQN weight {entry.fqn!r} was reused but never stored" ) + entry = replace(entry, storage_key=content_key) elif not _is_offgraph_kv_fqn(entry.fqn): raise RuntimeError(f"Missing CUDA FQN weight storage for {entry.fqn!r}") serialized_entries.append(entry) @@ -620,3 +637,4 @@ def finish(self) -> None: for result in self._results: result.data_store_output = shared_output self._stored_storages.clear() + self._content_keys.clear() diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index af49ed25e77..996ed2a20b8 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -25,6 +25,7 @@ CUDA_AOTI_METADATA_MAGIC, CudaAotiVariant, CudaWeightCollector, + decode_cuda_aoti_metadata, encode_cuda_aoti_metadata, ) from executorch.backends.cuda.passes.lower_offgraph_kv import ( @@ -431,7 +432,11 @@ def test_same_fqn_with_different_metadata_shares_storage(self) -> None: self.assertNotEqual( first_result.processed_bytes, second_result.processed_bytes ) - weight_key = first_artifact.entries[0].storage_key + weight_key = ( + decode_cuda_aoti_metadata(first_result.processed_bytes) + .entries[0] + .storage_key + ) self.assertIn( weight_key, first_result.data_store_output.external_data["aoti_cuda_blob"], @@ -481,12 +486,14 @@ def test_aoti_local_fqns_are_scoped_by_library(self) -> None: external_data = first_result.data_store_output.external_data[ "aoti_cuda_blob" ] - first_key = "cuda_fqn_weight:cpu:first_so:_tensor_constant0" - second_key = "cuda_fqn_weight:cpu:second_so:_tensor_constant0" - self.assertIn(first_key, external_data) - self.assertIn(second_key, external_data) - self.assertIn(first_key.encode(), first_result.processed_bytes) - self.assertIn(second_key.encode(), second_result.processed_bytes) + for result, library in ((first_result, "first"), (second_result, "second")): + (entry,) = decode_cuda_aoti_metadata(result.processed_bytes).entries + self.assertTrue( + entry.storage_key.startswith( + f"cuda_fqn_weight:cpu:{library}_so:_tensor_constant0:" + ) + ) + self.assertIn(entry.storage_key, external_data) for artifact in (first_artifact, second_artifact): for storage in artifact.storages.values(): @@ -511,7 +518,11 @@ def test_methods_share_one_collected_named_data_store(self) -> None: self.assertIs( first_result.data_store_output, second_result.data_store_output ) - weight_key = first_artifact.entries[0].storage_key + weight_key = ( + decode_cuda_aoti_metadata(first_result.processed_bytes) + .entries[0] + .storage_key + ) self.assertIn( weight_key, first_result.data_store_output.external_data["aoti_cuda_blob"], diff --git a/backends/cuda/tests/test_cuda_weight_reuse.py b/backends/cuda/tests/test_cuda_weight_reuse.py index 9a142ec2629..4413ecfd0f5 100644 --- a/backends/cuda/tests/test_cuda_weight_reuse.py +++ b/backends/cuda/tests/test_cuda_weight_reuse.py @@ -80,10 +80,10 @@ def counting(tensor, directory): self.assertEqual(set(first.storages), {e.storage_key for e in first.entries}) self.assertEqual(second.storages, {}) self.assertEqual(second.stored, {e.storage_key for e in second.entries}) - for entry in first.entries: - self.assertIn(entry.storage_key, collector._store.key_to_buffer_idx) first_manifest = decode_cuda_aoti_metadata(first_result.processed_bytes) second_manifest = decode_cuda_aoti_metadata(second_result.processed_bytes) + for entry in first_manifest.entries: + self.assertIn(entry.storage_key, collector._store.key_to_buffer_idx) self.assertEqual(second_manifest.entries, first_manifest.entries) def test_the_stored_bytes_match_writing_every_method(self) -> None: diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index 410184ee238..fb2556cb3f8 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -51,8 +51,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: "loads through the Python runtime, which the ROCm build does not include", ) class TestMissingWeightsBlob(unittest.TestCase): - def _lower(self, outdir: str) -> str: - torch.manual_seed(0) + def _lower(self, outdir: str, seed: int = 0, name: str = "model.pte") -> str: + torch.manual_seed(seed) module = ModuleWithConstants().eval().cuda() inputs = (torch.randn(4, 64, device="cuda"),) exported = export(module, inputs) @@ -61,7 +61,7 @@ def _lower(self, outdir: str) -> str: exported, partitioner=[CudaPartitioner([spec])] ) program = lowered.to_executorch() - path = os.path.join(outdir, "model.pte") + path = os.path.join(outdir, name) with open(path, "wb") as f: program.write_to_file(f) program.write_tensor_data_to_file(outdir) @@ -160,6 +160,23 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: ): program.load_method("forward") + def test_load_fails_with_another_programs_blob(self) -> None: + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + first = self._lower(outdir, seed=0, name="first.pte") + # Same module with other weights, written to the same blob file name. + second = self._lower(outdir, seed=1, name="second.pte") + (blob,) = [f for f in os.listdir(outdir) if f.endswith(".ptd")] + blob = os.path.join(outdir, blob) + # The blob belongs to the second program, which loads with it. + Runtime.get().load_program(second, data_path=blob).load_method("forward") + program = Runtime.get().load_program(first, data_path=blob) + with self.assertRaisesRegex( + RuntimeError, r"Failed to load method forward, error: 0x:?20" + ): + program.load_method("forward") + def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob.