From 6207bdfd38b4d98837bae054821a9a3a03543d21 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 08:54:41 -0700 Subject: [PATCH 01/29] Fix crashes, wrong outputs and lowering failures in several backends and in exir (#24222) Fixes found while exporting models to ExecuTorch backends. ## Fixes - **exir:** `constant_prop_pass` keeps tensors that in-place ops write, also through views or in `no_grad`/`autocast` blocks. `to_backend` keeps the output node last. `to_edge` lowers `torch._check_tensor_all` and `torch.sym_ite`. `GroupBasedPartitioner` no longer makes partitions that depend on each other, unless caller node groups already do. Tensor constants get static shapes; deduplicated constants keep their data. - **AOTI (AOTInductor) and CUDA:** the missing-kernel check catches scatter and index_put fallbacks. `detach` lowers, also inside `torch.cond`. Strided outputs are copied. Transposed convolutions (1-D, 2-D, no dilation) and `torch._check_tensor_all` lower. - **XNNPACK:** the quantizer skips bfloat16 tensors, and `cat` respects scoped configs. ELU is delegated when XNNPACK can run it. - **Qualcomm:** autocast regions are inlined correctly. Unquantized bfloat16 nodes, data-dependent shapes and HTP partitions without runtime inputs stay on the CPU. Fixes for stacks, 0-d selects, RMSNorm and int64 constants. - **OpenVINO:** model data stays alive while OpenVINO 2026.4 reads it. Partitions are converted with their input types. - **Core ML:** bfloat16 ops, and ops above rank 5 outside static layout chains, stay on the CPU. - **MLX, Arm, Cortex-M, Vulkan, NXP, portable:** dtype checks for MLX tensor lists, one constant slot per weight, slice updates on a constant base; TOSA keeps unused delegate inputs for VGF; `MatmulToBmmPass` takes a constant operand; Vulkan constants above 4 dims; non-strict programs in the NXP quantizer; `cummin.out`. ## Behavior changes - A CUDA `.pte` loads only the weights it was exported with. A `.ptd` with other weights, for example after fine-tuning, needs a newly exported `.pte`. - CUDA graphs turn off when outputs do not match the ETensor layout. If only the capture call sees such a layout, that call returns an error. - `GroupBasedPartitioner` is about 2x slower on large graphs, and some graphs get a different number of partitions. `to_backend` keeps the original order of top-level nodes. - XNNPACK: weight-only quantized linears, and 8da4w or 8da8w linears (8-bit dynamic activations, 4-bit or 8-bit weights) lowered with `to_backend`, fail at export instead of giving wrong results. Delegated ELU flushes subnormal values to zero. Cats outside a scoped quantizer config are no longer quantized. - Ethos-U raises an error for a delegate with unused inputs instead of returning wrong values. We would like to cherry-pick this PR to release/1.6 after it merges. It needs #23915 there first. ## Testing Each fix has a test that fails without it. Run on an A100 (CUDA, Vulkan), an x86 CPU (OpenVINO, QNN SDK), macOS (Core ML, MLX), the TOSA reference model and VGF emulation (Arm), and the eIQ Neutron SDK (NXP). Not run: Metal, ROCm, Windows CUDA. (cherry picked from commit 054de1faa9a86aa86659a7c13c11361cdc129520) --- backends/aoti/aoti_backend.py | 40 +- .../passes/replace_view_copy_with_view.py | 31 +- .../coreml/partition/coreml_partitioner.py | 109 +++++- .../coreml/test/test_coreml_partitioner.py | 88 +++++ 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 +- .../cortex_m/passes/matmul_to_bmm_pass.py | 9 +- .../cortex_m/test/ops/test_batch_matmul.py | 22 ++ backends/cuda/cuda_backend.py | 77 +++- backends/cuda/cuda_weight_collector.py | 20 +- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 43802 -> 44602 bytes backends/cuda/runtime/cuda_backend.cpp | 50 +++ backends/cuda/runtime/shims/memory.cpp | 16 + backends/cuda/runtime/shims/memory.h | 17 + .../test_aoti_torch__reinterpret_tensor.cpp | 45 +++ .../runtime/test/test_cuda_output_copy.cpp | 19 + backends/cuda/runtime/utils.h | 74 +++- backends/cuda/tests/test_cuda_export.py | 111 +++++- 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 +- .../tests/test_output_stride_rearrange.py | 103 +++++ backends/mlx/builder/program_builder.py | 30 +- backends/mlx/runtime/MLXInterpreter.h | 8 +- backends/mlx/test/test_partitioner.py | 46 +++ .../move_activation_before_concat.py | 11 +- backends/nxp/quantizer/utils.py | 17 +- .../nxp/tests/generic_tests/test_quantizer.py | 31 ++ backends/openvino/preprocess.py | 42 +- 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 +- backends/openvino/tests/ops/test_index_put.py | 28 ++ backends/openvino/tests/ops/test_view.py | 11 + .../_passes/decompose_wrap_with_autocast.py | 24 +- backends/qualcomm/_passes/i64_to_i32.py | 10 +- backends/qualcomm/_passes/qnn_pass_manager.py | 5 +- .../qualcomm/_passes/recompose_rms_norm.py | 19 +- backends/qualcomm/_passes/remove_0d_tensor.py | 30 +- .../qualcomm/partition/qnn_partitioner.py | 89 ++++- backends/qualcomm/partition/utils.py | 75 +++- backends/qualcomm/tests/rework/src/pattern.py | 2 +- backends/qualcomm/tests/test_passes.py | 358 ++++++++++++++++++ ...test_create_delete_constant_placeholder.py | 43 +++ backends/transforms/utils.py | 15 +- .../vulkan/runtime/graph/ops/impl/Staging.cpp | 7 +- backends/vulkan/test/test_vulkan_delegate.py | 13 + backends/xnnpack/_passes/convert_to_linear.py | 4 + backends/xnnpack/operators/op_elu.py | 4 +- backends/xnnpack/partition/config/__init__.py | 4 +- .../xnnpack/partition/config/gemm_configs.py | 4 + .../partition/config/generic_node_configs.py | 25 ++ .../xnnpack/quantizer/xnnpack_quantizer.py | 5 +- .../quantizer/xnnpack_quantizer_utils.py | 15 + backends/xnnpack/test/ops/test_elu.py | 80 +++- backends/xnnpack/test/ops/test_linear.py | 73 ++++ .../test/quantizer/test_xnnpack_quantizer.py | 71 ++++ docs/source/backends/xnnpack/op-support.csv | 2 +- .../duplicate_constant_node_pass.py | 21 +- .../group_partitioner.py | 236 ++++++++++-- exir/backend/test/test_group_partitioner.py | 222 ++++++++++- exir/emit/test/test_emit.py | 5 +- exir/lowered_backend_module.py | 6 +- exir/pass_base.py | 13 +- exir/passes/BUCK | 1 + exir/passes/constant_prop_pass.py | 130 ++++++- exir/program/test/test_program.py | 38 +- exir/tests/targets.bzl | 1 + exir/tests/test_delegate.py | 28 ++ exir/tests/test_pass_infra.py | 178 +++++++++ exir/tests/test_passes.py | 63 +++ exir/tracer.py | 2 + exir/verification/verifier.py | 27 +- 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 + 85 files changed, 3175 insertions(+), 208 deletions(-) create mode 100644 backends/openvino/tests/ops/test_index_put.py create mode 100644 kernels/portable/cpu/op_cummin.cpp create mode 100644 kernels/test/op_cummin_test.cpp 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/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/apple/coreml/partition/coreml_partitioner.py b/backends/apple/coreml/partition/coreml_partitioner.py index b4e550d8479..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, @@ -83,6 +142,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: @@ -152,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 " @@ -159,6 +242,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 +452,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..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. @@ -433,6 +488,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. 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, 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),), + ), } diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 787d2502b63..0bd1599cd5a 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,18 @@ 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, + # 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). + 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/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/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/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/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); +} 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_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8329f42e522..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: @@ -328,6 +330,57 @@ 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_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.""" @@ -520,6 +573,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. @@ -562,6 +653,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 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. diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 24df64d6de9..b1bf91638a7 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -256,6 +256,109 @@ 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()) + + +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() diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py index 797cdb1bdf2..d5ab4124579 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(): @@ -461,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 = [] @@ -490,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" @@ -519,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/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 d8b7b8e485a..5897fe58bee 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__() @@ -323,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.""" @@ -397,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.""" 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() 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/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)) 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))) 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/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/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 854c619061c..b60a80b5c7f 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/_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/_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/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 8db23a465ef..770bf4284d3 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -14,6 +14,10 @@ PER_CHANNEL_GROUP_ENCODING, ) from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader +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, @@ -55,6 +59,8 @@ 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, ) @@ -113,7 +119,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 @@ -160,6 +166,18 @@ 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 + + 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. @@ -244,6 +262,71 @@ def generate_partitions( edge_program.graph_module, op_support=self.op_support_checker, ) + 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( + partitions: List[Partition], edge_program: torch.export.ExportedProgram + ) -> List[Partition]: + """ + A parameter dequantize whose consumer is not delegated stays on CPU. A QNN + Dequantize alone rejects per-channel encodings. + """ + for partition in partitions: + stack = [ + node + for node in partition.nodes + if node.target in CPU_DQ_OPS + and node.meta.get(QCOM_BYPASS_NODE, False) + and is_parameter(node.args[0], edge_program) + ] + while stack: + node = stack.pop() + if node not in partition.nodes or any( + user in partition.nodes for user in node.users + ): + continue + partition.remove_node(node) + stack.extend( + arg + for arg in node.args + if isinstance(arg, torch.fx.Node) + and arg.meta.get(QCOM_BYPASS_NODE, False) + ) + return [partition for partition in partitions if partition.nodes] + + @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 @@ -257,6 +340,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/partition/utils.py b/backends/qualcomm/partition/utils.py index 5fa5633ac70..19fa9b1ad55 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,66 @@ 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} + + +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, @@ -148,6 +216,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/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 606fed01d72..58ce89a6e30 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -7,11 +7,14 @@ ConvertBmmToMatmul, ConvertMhaToSha, DecomposeColIm, + DecomposeWrapWithAutocast, ExpandBroadcastTensorShape, FoldQDQ, InsertIOQDQ, InsertReshapeForReduceOps, LiftConstantScalarOperands, + RecomposeRmsNorm, + Remove0DTensor, RemoveRedundancy, ) from executorch.backends.qualcomm._passes.qnn_pass_manager import ( @@ -59,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 @@ -140,6 +144,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" @@ -387,6 +435,94 @@ 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_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_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 @@ -783,6 +919,228 @@ 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 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 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 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 + ): + """A quantized conv kept off QNN must take its per-channel weight dequantize + with it. FoldQDQ force-passes that dequantize, so it used to be lowered alone + as a QNN Dequantize, which rejects per-channel encodings and fails context + binary generation (#19908).""" + + class ConvDepthwiseConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.pw1 = torch.nn.Conv2d(8, 16, 1) + self.dw = torch.nn.Conv2d(16, 16, 3, padding=1, groups=16) + self.pw2 = torch.nn.Conv2d(16, 8, 1) + + def forward(self, x): + return self.pw2(torch.relu(self.dw(torch.relu(self.pw1(x))))) + + sample_input = (torch.randn(1, 8, 16, 16),) + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + quant_dtype=quant_dtype, is_conv_per_channel=True + ) + prepared = prepare_pt2e( + torch.export.export(ConvDepthwiseConv().eval(), sample_input).module(), + quantizer, + ) + prepared(*sample_input) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=use_fp16), + ) + try: + edge = to_edge_transform_and_lower_to_qnn( + convert_pt2e(prepared), + sample_input, + compiler_specs, + skip_node_id_set=skip_node_id_set, + ) + except RuntimeError as e: + if "QNN" in str(e) or "qnn" in str(e): + self.skipTest(f"QNN SDK not available: {e}") + raise + + # inspected before to_executorch(), which swaps in out-variant ops + def cpu_nodes(target): + return [ + node + for node in edge.exported_program().graph.nodes + if node.op == "call_function" and node.target == target + ] + + convs = cpu_nodes(exir_ops.edge.aten.convolution.default) + self.assertEqual([conv.name for conv in convs], ["aten_convolution_default_1"]) + dequants = cpu_nodes( + exir_ops.edge.quantized_decomposed.dequantize_per_channel.default + ) + self.assertTrue(dequants) + for dequant in dequants: + self.assertEqual(list(dequant.users), convs) + if emit: + edge.to_executorch() + + def test_skipped_conv_keeps_its_weight_dequantize_on_cpu(self): + self._assert_cpu_conv_keeps_its_weight_dequantize( + QuantDtype.use_16a8w, + use_fp16=False, + skip_node_id_set={"aten_convolution_default_1"}, + ) + + def test_rejected_conv_keeps_its_weight_dequantize_on_cpu(self): + # the HTP rejects fp16-activation/int8-weight depthwise conv. Not emitted: + # InsertCastForFpActQuantizedWeight casts the CPU conv's bias meta to fp16. + self._assert_cpu_conv_keeps_its_weight_dequantize( + QuantDtype.use_fp16a8w, use_fp16=True, emit=False + ) + def test_lowering_accepts_exported_program_without_reexport(self): """to_edge_transform_and_lower_to_qnn must lower a caller-supplied ExportedProgram as captured, never re-exporting it. diff --git a/backends/transforms/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/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 c7b66fe0ae4..b1ef1f27db6 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -2782,6 +2782,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 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/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/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/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/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 a207d77674d..4e5b6228f55 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 @@ -1123,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/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/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 diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index b66166fedb6..f3322bc83ad 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: @@ -712,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( 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",, 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, 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]) diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index 78480d3b55c..6dc401def7c 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -2517,10 +2517,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/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/pass_base.py b/exir/pass_base.py index e7c2caa66a9..6f83996b550 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -449,14 +449,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/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/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/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: diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 505fd50af8d..8d00c17cc05 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -495,6 +495,184 @@ 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 + ) -> torch.Tensor: + end = length.item() + torch._check(end >= 0) + return tensor[:end] + + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_inputs = ( + fake_mode.from_tensor(torch.randn(10)), + fake_mode.from_tensor(torch.tensor(5)), + ) + graph_module = make_fx( + data_dependent_slice, + tracing_mode="fake", + _allow_non_fake_inputs=True, + )(*fake_inputs) + + self.assertFalse(shape_env.pending_fresh_unbacked_symbols) + outer_pending = [shape_env.create_unbacked_symint().node.expr for _ in range(3)] + replayed = ExportPass()(graph_module).graph_module + self.assertEqual( + shape_env.pending_fresh_unbacked_symbols, + outer_pending, + ) + + item_node = next( + node + for node in replayed.graph.nodes + if node.target == torch.ops.aten._local_scalar_dense.default + ) + slice_node = next( + node + for node in replayed.graph.nodes + if node.target == torch.ops.aten.slice.Tensor + ) + self.assertTrue(item_node.meta.get("unbacked_bindings")) + self.assertTrue(slice_node.meta.get("unbacked_bindings")) + + def test_export_pass_rejects_unreturned_replay_symbol(self) -> None: + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_input = fake_mode.from_tensor(torch.randn(4)) + graph_module = make_fx( + lambda tensor: tensor + 1, + tracing_mode="fake", + _allow_non_fake_inputs=True, + )(fake_input) + + class InvalidUnbackedSymbolPass(ExportPass): + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + assert self.fake_tensor_mode is not None + assert self.fake_tensor_mode.shape_env is not None + self.fake_tensor_mode.shape_env.create_unbacked_symint() + return super().call_operator(op, args, kwargs, meta) + + with self.assertRaises(PendingUnbackedSymbolNotFound): + InvalidUnbackedSymbolPass()(graph_module) + + def test_export_pass_restored_outer_symbol_remains_unbound(self) -> None: + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_input = fake_mode.from_tensor(torch.randn(4)) + graph_module = make_fx( + lambda tensor: tensor + 1, + tracing_mode="fake", + _allow_non_fake_inputs=True, + )(fake_input) + + outer_symbol = shape_env.create_unbacked_symint() + ExportPass()(graph_module) + self.assertEqual( + shape_env.pending_fresh_unbacked_symbols, + [outer_symbol.node.expr], + ) + with self.assertRaises(PendingUnbackedSymbolNotFound): + compute_unbacked_bindings(shape_env, fake_input) + + def test_export_pass_preserves_multiple_nested_symbol_scopes(self) -> None: + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_input = fake_mode.from_tensor(torch.randn(4)) + graph_module = make_fx( + lambda tensor: tensor + 1, + tracing_mode="fake", + _allow_non_fake_inputs=True, + )(fake_input) + + observed_scopes: list[tuple[int, list[Any], list[Any], Any]] = [] + + class NestedReplayPass(ExportPass): + def __init__(self, depth: int) -> None: + super().__init__() + self.depth = depth + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + if self.depth: + assert self.fake_tensor_mode is not None + nested_shape_env = self.fake_tensor_mode.shape_env + assert nested_shape_env is not None + owned_symbol = nested_shape_env.create_unbacked_symint() + if self.depth % 2 == 0: + nested_shape_env.ignorable_fresh_unbacked_symbols.append( + owned_symbol.node.expr + ) + + NestedReplayPass(self.depth - 1)(graph_module) + observed_scopes.append( + ( + self.depth, + list(nested_shape_env.pending_fresh_unbacked_symbols), + list(nested_shape_env.ignorable_fresh_unbacked_symbols), + owned_symbol.node.expr, + ) + ) + compute_unbacked_bindings(nested_shape_env, owned_symbol) + return super().call_operator(op, args, kwargs, meta) + + caller_symbols = [shape_env.create_unbacked_symint() for _ in range(2)] + shape_env.ignorable_fresh_unbacked_symbols.append(caller_symbols[1].node.expr) + NestedReplayPass(3)(graph_module) + + self.assertEqual([scope[0] for scope in observed_scopes], [1, 2, 3]) + for depth, pending, ignorable, owned_symbol in observed_scopes: + self.assertEqual(pending, [owned_symbol]) + self.assertEqual(ignorable, [owned_symbol] if depth % 2 == 0 else []) + self.assertEqual( + shape_env.pending_fresh_unbacked_symbols, + [symbol.node.expr for symbol in caller_symbols], + ) + self.assertEqual( + shape_env.ignorable_fresh_unbacked_symbols, + [caller_symbols[1].node.expr], + ) + compute_unbacked_bindings(shape_env, caller_symbols[0]) + self.assertFalse(shape_env.pending_fresh_unbacked_symbols) + self.assertFalse(shape_env.ignorable_fresh_unbacked_symbols) + def test_export_pass_matches_symbolic_inputs_by_position(self) -> None: class RenamePlaceholderPass(ExportPass): def placeholder( diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 2fa00618679..d5ba561083a 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2902,6 +2902,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): 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: diff --git a/kernels/aten/functions.yaml b/kernels/aten/functions.yaml index 0e152cbea62..0ae87f6e733 100644 --- a/kernels/aten/functions.yaml +++ b/kernels/aten/functions.yaml @@ -131,6 +131,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 e3a5dba0f20..298142cb9e6 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -318,6 +318,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 0bce3855e5d..913df45439e 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -196,6 +196,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 5078dc996ea..b443c7f8619 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -247,6 +247,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 bed219b04af..94f535094e5 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -111,6 +111,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 0a3ae4409c0..fa4ab1e4498 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -496,6 +496,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 d1174e42328dd31aaef0ce569f7089d593c0aead Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 13:30:50 -0700 Subject: [PATCH 02/29] Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend (#24294) ## What is wrong today The OpenVINO backend has these problems: 1. A 0-d (scalar) input or output fails on every run, because OpenVINO's `ov_shape_create` refuses rank 0. 2. A float32 program does not compute in float32 on many CPUs. OpenVINO picks bfloat16 (two to three significant digits) on x86 CPUs with bfloat16 hardware and float16 on ARM CPUs, and no compile spec can change it. 3. An output whose size depends on an int input keeps its largest size. `x[:n].view(n, -1)` with x of shape (4, 8) and n = 3 returns (4, 8). 4. `cumsum` and `prod` of a bool or integer tensor are not computed in int64, as PyTorch does. A cumsum over a bool mask wraps after 255, so pi0 and pi0.5 get wrong position ids and wrong actions. ## What this change does - **0-d tensors:** a rank-0 tensor keeps its empty shape, which OpenVINO accepts. - **Precision:** when the device is `CPU`, a program computes in float32 by default, both when it is lowered and when it is loaded. A program can ask for a faster precision: ```python OpenvinoPartitioner( [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] ) ``` The value is `f32`, `f16` or `bf16`, checked before conversion starts. Other devices keep OpenVINO's default. On x86 CPUs OpenVINO applies the precision when the runtime loads the program, so the runtime passes it on too. - **Dynamic output sizes:** an output whose size can change goes to a buffer that OpenVINO owns, reused by every call. It starts empty and OpenVINO grows it to the largest size a call writes, so a large upper bound costs nothing until a call needs it. After inference the backend resizes the ExecuTorch output to OpenVINO's shape and copies the data in, so one delegate can also serve several calls, as programs from ExecuTorch 0.7 and earlier do. Static outputs still use the ExecuTorch buffer directly. - **Integer reductions:** the input of `cumsum` and `prod` is cast to int64 before conversion. An integer cumsum that reads only constants is computed during lowering, because OpenVINO's CPU plugin gives it wrong results from the second inference on. This includes a cumsum over `ones_like` or `zeros_like` of an input with a static shape, which OpenVINO folds into a constant too. Each one is folded on its own, and one that cannot be folded is named in a warning. Trade-offs: - float32 is slower. On a Xeon with bfloat16 hardware, a 4-layer MLP takes 2.4 ms in float32 and 1.5 ms in bfloat16. On x86 this also applies to programs exported before this change, because OpenVINO applies the precision when a program loads. To get the old speed, export again with `inference_precision` set to `bf16`. - OpenVINO's CPU plugin computes int64 in 32 bits, so an integer result past the int32 range wraps without an error. That was already true for int64 programs, and now also applies to `cumsum` and `prod` over smaller integer types. ## Tests Each new test fails without its fix: 0-d inputs and outputs; a 256 x 256 addmm checked to 1e-4, and the same addmm lowered with `bf16` checked not to reach 1e-4, so the runtime must apply an explicit precision (both skipped where OpenVINO already computes in float32); an unknown precision refused before conversion; dynamic output sizes, on one program and on one delegate shared by two calls; bool and int32 cumsum and prod; pi0-style masks with int64 and int32 cumsums over two calls, built from a constant and from `ones_like`; and folding when one cumsum cannot be evaluated. Two guards pass either way: integer `sum`, which OpenVINO already computes in int64, and a static output after a dynamic view. The full OpenVINO op suite passes on Linux x86_64 (a CPU with bfloat16 units) with OpenVINO 2026.4.1, through both pytest and the CI test runner. With OpenVINO 2025.1.0 it also passes, except that the prod test is skipped because that release does not delegate prod. (cherry picked from commit 7b14148ec5e5bae0c054dc26dc7be3a17ea8cbfb) --- backends/openvino/README.md | 45 ++++ backends/openvino/_passes/__init__.py | 8 +- .../_passes/cast_reduction_input_pass.py | 45 ++++ .../_passes/fold_constant_cumsum_pass.py | 156 +++++++++++++ backends/openvino/preprocess.py | 55 ++++- backends/openvino/runtime/OpenvinoApi.h | 17 ++ backends/openvino/runtime/OpenvinoBackend.cpp | 203 ++++++++++++++-- backends/openvino/runtime/OpenvinoBackend.h | 14 +- .../tests/ops/base_openvino_op_test.py | 64 ++--- backends/openvino/tests/ops/test_add.py | 9 + backends/openvino/tests/ops/test_addmm.py | 56 +++++ .../openvino/tests/ops/test_reductions.py | 220 ++++++++++++++++++ backends/openvino/tests/ops/test_view.py | 81 +++++++ docs/source/build-run-openvino.md | 45 ++++ 14 files changed, 960 insertions(+), 58 deletions(-) create mode 100644 backends/openvino/_passes/cast_reduction_input_pass.py create mode 100644 backends/openvino/_passes/fold_constant_cumsum_pass.py create mode 100644 backends/openvino/tests/ops/test_reductions.py diff --git a/backends/openvino/README.md b/backends/openvino/README.md index 331aebf87ff..40a245f12a0 100644 --- a/backends/openvino/README.md +++ b/backends/openvino/README.md @@ -48,6 +48,51 @@ print(_get_registered_backend_names()) # Should include 'OpenvinoBackend' ``` +## Precision + +When the device is `CPU`, a program computes in float32 by default. Without this default, +OpenVINO would compute in bfloat16 on x86 CPUs with bfloat16 units, which keeps only two to +three significant digits, and in float16 on ARM CPUs, which keeps about three. +Both are faster. To trade accuracy for speed, ask for the faster precision when +lowering: + +```python +from executorch.backends.openvino.partitioner import OpenvinoPartitioner +from executorch.exir.backend.compile_spec_schema import CompileSpec + +OpenvinoPartitioner( + [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] +) +``` + +On an ARM CPU, ask for `f16` instead: there float32 is about 1.8x slower than +float16, as measured on an Orin Nano. + +The value is OpenVINO's `INFERENCE_PRECISION_HINT`: `f32`, `f16` or `bf16`. An +empty value means the default, and any other value fails at lowering. A CPU +without float16 support computes an `f16` program in float32. OpenVINO has no +bfloat16 on ARM CPUs, so there `bf16` keeps the ARM default, float16 where the +CPU supports it. Which x86 CPUs compute in bfloat16 is up to OpenVINO; see its +[Precision Control](https://docs.openvino.ai/2026/openvino-workflow/running-inference/optimize-inference/precision-control.html) +page. Other devices, such as GPUs, keep OpenVINO's default unless the option is +set. + +The default also applies to a program exported in float16: on a CPU it computes +in float32 unless you lower it with `inference_precision`. + +On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. +So a program lowered without `inference_precision` computes in float32, even +one exported with an ExecuTorch release that did not have this default. That +is slower and can use more memory. Export it again with `inference_precision` +set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +precision when the program is lowered, so a program exported with such an +earlier release keeps computing in float16 until it is exported again. + +OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result +outside the int32 range wraps around without an error. This applies to every +int64 tensor, including the int64 result of `cumsum` and `prod` over smaller +integer types. + ## Directory Structure ``` diff --git a/backends/openvino/_passes/__init__.py b/backends/openvino/_passes/__init__.py index 0db94ca1ec8..bc2853a73e6 100644 --- a/backends/openvino/_passes/__init__.py +++ b/backends/openvino/_passes/__init__.py @@ -4,6 +4,12 @@ # except in compliance with the License. See the license file found in the # LICENSE file in the root directory of this source tree. +from .cast_reduction_input_pass import CastReductionInputPass from .decompose_floor_divide_pass import DecomposeFloorDividePass +from .fold_constant_cumsum_pass import FoldConstantCumsumPass -__all__ = ["DecomposeFloorDividePass"] +__all__ = [ + "CastReductionInputPass", + "DecomposeFloorDividePass", + "FoldConstantCumsumPass", +] diff --git a/backends/openvino/_passes/cast_reduction_input_pass.py b/backends/openvino/_passes/cast_reduction_input_pass.py new file mode 100644 index 00000000000..36e3f3795c0 --- /dev/null +++ b/backends/openvino/_passes/cast_reduction_input_pass.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +REDUCTION_OPS = { + exir_ops.edge.aten.cumsum.default, + exir_ops.edge.aten.prod.default, + exir_ops.edge.aten.prod.dim_int, +} + + +class CastReductionInputPass(PassBase): + """Cast the input of cumsum and prod to the dtype of their result. + + PyTorch reduces bool and integer tensors in int64, but OpenVINO keeps a + narrower type for these two, though not for sum. A cumsum over a bool mask + is computed in 8 bits and wraps after 255, and a cumsum or prod over int32 + returns int32. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + for node in graph_module.graph.nodes: + if node.op != "call_function" or node.target not in REDUCTION_OPS: + continue + source = node.args[0] + dtype = node.meta["val"].dtype + if source.meta["val"].dtype == dtype: + continue + with graph_module.graph.inserting_before(node): + cast = graph_module.graph.call_function( + exir_ops.edge.aten._to_copy.default, (source,), {"dtype": dtype} + ) + cast.meta["val"] = source.meta["val"].to(dtype) + node.replace_input_with(source, cast) + modified = True + if modified: + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/openvino/_passes/fold_constant_cumsum_pass.py b/backends/openvino/_passes/fold_constant_cumsum_pass.py new file mode 100644 index 00000000000..fabc030b06a --- /dev/null +++ b/backends/openvino/_passes/fold_constant_cumsum_pass.py @@ -0,0 +1,156 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import operator + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.constant_prop_pass import _nodes_written_in_place +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.utils import _pytree as pytree + + +def _reads_only_shape(node: torch.fx.Node) -> bool: + """full_like with a constant fill value reads only the shape and dtype of + its input. With a static shape its result is a constant, and OpenVINO + folds it into one too. + """ + if node.target != exir_ops.edge.aten.full_like.default: + return False + if any( + isinstance(a, torch.fx.Node) for a in (*node.args[1:], *node.kwargs.values()) + ): + return False + val = node.meta.get("val") + return isinstance(val, torch.Tensor) and all(isinstance(d, int) for d in val.shape) + + +def _inputs(node: torch.fx.Node) -> list: + return [] if _reads_only_shape(node) else node.all_input_nodes + + +def _on_cpu(value): + if isinstance(value, torch.Tensor): + return value.cpu() + if isinstance(value, torch.device): + return torch.device("cpu") + return value + + +def _constant_nodes(graph_module: torch.fx.GraphModule) -> set: + """Nodes that depend only on constants. Graph order puts inputs first.""" + constants = set() + written = _nodes_written_in_place(graph_module) + for node in graph_module.graph.nodes: + # Only tensors: a higher-order op such as cond reads its subgraphs, + # which may hold a random op, through get_attr nodes too. + if node.op == "get_attr": + value = operator.attrgetter(node.target)(graph_module) + if isinstance(value, torch.Tensor) and node not in written: + constants.add(node) + elif node.op == "call_function" and ( + _reads_only_shape(node) + or ( + torch.Tag.nondeterministic_seeded + not in getattr(node.target, "tags", ()) + and all(n in constants for n in node.all_input_nodes) + ) + ): + constants.add(node) + return constants + + +def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict: + """Values of targets, computing only what they read, on the CPU.""" + needed = set() + stack = list(targets) + while stack: + node = stack.pop() + if node not in needed: + needed.add(node) + stack.extend(_inputs(node)) + order = [n for n in graph_module.graph.nodes if n in needed] + last_use = {i: n for n in order for i in _inputs(n)} + values = {} + for node in order: + if node.op == "get_attr": + values[node] = operator.attrgetter(node.target)(graph_module) + continue + if _reads_only_shape(node): + like = node.args[0].meta["val"] + args = (torch.empty(like.shape, dtype=like.dtype), *node.args[1:]) + kwargs = dict(node.kwargs) + else: + args, kwargs = torch.fx.node.map_arg( + (node.args, node.kwargs), lambda n: values[n] + ) + # The values are the same on any device, and lowering for OpenVINO + # must not need the GPU a constant was made for. + args, kwargs = pytree.tree_map(_on_cpu, (args, kwargs)) + if kwargs.get("pin_memory"): + kwargs = {**kwargs, "pin_memory": False} + values[node] = node.target(*args, **kwargs) + for i in _inputs(node): + if last_use[i] is node and i not in targets: + del values[i] + return values + + +class FoldConstantCumsumPass(PassBase): + """Replace an integer cumsum that reads only constants with its result. + + With an integer cumsum over a constant, OpenVINO's CPU plugin (2025.4 to + 2026.4) gives wrong results from the second inference on: masks built from + it keep only what earlier calls allowed. pi0 builds its attention masks + this way. The OpenVINO nightly no longer does this, so the pass can go once + the oldest supported OpenVINO release has the fix. Only these cumsums are + folded; constants that feed no such cumsum, such as compressed weights, are + left as they are. Runs on an unlifted module, where constants are get_attr + nodes; buffers the graph writes to are not constants, and full_like of an + input with a static shape is one. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + constants = _constant_nodes(graph_module) + cumsums = [ + n + for n in graph.nodes + if n in constants + and n.target == exir_ops.edge.aten.cumsum.default + and not n.meta["val"].dtype.is_floating_point + and not n.meta["val"].dtype.is_complex + ] + modified = False + # One at a time, so a cumsum that cannot be evaluated does not keep + # the others from being folded. + for node in cumsums: + try: + value = _evaluate(graph_module, [node])[node] + except RuntimeError as e: + # Some ops have modes the CPU lacks. + logging.warning( + "OpenVINO: could not fold %s, so on a CPU some OpenVINO " + "releases give wrong results for it from the second " + "inference on: %s", + node.name, + e, + ) + continue + name = f"_folded_cumsum_{node.name}" + while hasattr(graph_module, name): + name += "_" + graph_module.register_buffer(name, value) + with graph.inserting_before(node): + folded = graph.get_attr(name) + folded.meta = node.meta.copy() + node.replace_all_uses_with(folded) + modified = True + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index de2c35d2414..c945fb57067 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -8,7 +8,12 @@ from typing import final, List -from executorch.backends.openvino._passes import DecomposeFloorDividePass +import torch +from executorch.backends.openvino._passes import ( + CastReductionInputPass, + DecomposeFloorDividePass, + FoldConstantCumsumPass, +) from executorch.exir.backend.backend_details import ( BackendDetails, @@ -21,6 +26,10 @@ openvino_compile, ) +# The values of the inference_precision compile spec. The runtime +# (OpenvinoBackend.cpp) accepts the same ones. +_PRECISIONS = ("f32", "f16", "bf16") + @final class OpenvinoBackend(BackendDetails): @@ -39,7 +48,25 @@ def preprocess( Returns: PreprocessResult: The result of preprocessing, including the compiled model bytes. """ - for pass_cls in [DimOrderOpsRevertPass, DecomposeFloorDividePass]: + compile_options = {} + for spec in module_compile_spec: + compile_options[spec.key] = spec.value.decode() + device = compile_options.get("device", "CPU") + requested_precision = compile_options.get("inference_precision") + # Some OpenVINO releases accept an unknown value. Checking before the + # conversion fails fast and keeps it out of the program, where the + # runtime would only fail at load. + if requested_precision and requested_precision not in _PRECISIONS: + raise ValueError( + f"OpenVINO backend: unsupported inference_precision " + f"'{requested_precision}', expected one of {', '.join(_PRECISIONS)}" + ) + + for pass_cls in [ + DimOrderOpsRevertPass, + DecomposeFloorDividePass, + CastReductionInputPass, + ]: result = pass_cls()(edge_program.graph_module) if result and result.graph_module: edge_program._graph_module = result.graph_module @@ -50,13 +77,25 @@ def preprocess( if node.target in input_names: args.append(node.meta["val"]) - 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. + module = edge_program.module() + FoldConstantCumsumPass()(module) + decoder = TorchFXPythonDecoder( + 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)) + # On ARM CPUs OpenVINO fixes the precision here, not when the runtime + # imports the model, so this must follow the same rule as the runtime. + precision = requested_precision or ("f32" if device.startswith("CPU") else None) + config = {"INFERENCE_PRECISION_HINT": precision} if precision else {} + compiled = compile_model(ppp.build(), device, config) model_bytes = compiled.export_model() return PreprocessResult(processed_bytes=model_bytes.getvalue()) diff --git a/backends/openvino/runtime/OpenvinoApi.h b/backends/openvino/runtime/OpenvinoApi.h index d25e265db07..301039ca1b6 100644 --- a/backends/openvino/runtime/OpenvinoApi.h +++ b/backends/openvino/runtime/OpenvinoApi.h @@ -69,6 +69,9 @@ using ov_core_free_fn = void (*)(ov_core_t*); using ov_core_get_available_devices_fn = ov_status_e (*)(const ov_core_t*, ov_available_devices_t*); using ov_available_devices_free_fn = void (*)(ov_available_devices_t*); +// Takes one property key and value, both strings, after the device name. +using ov_core_set_property_fn = + ov_status_e (*)(const ov_core_t*, const char*, ...); using ov_core_import_model_fn = ov_status_e (*)( const ov_core_t*, const char*, @@ -87,10 +90,17 @@ using ov_infer_request_set_input_tensor_by_index_fn = using ov_infer_request_set_output_tensor_by_index_fn = ov_status_e (*)(ov_infer_request_t*, size_t, const ov_tensor_t*); using ov_infer_request_infer_fn = ov_status_e (*)(ov_infer_request_t*); +using ov_infer_request_get_output_tensor_by_index_fn = + ov_status_e (*)(const ov_infer_request_t*, size_t, ov_tensor_t**); using ov_infer_request_free_fn = void (*)(ov_infer_request_t*); using ov_tensor_create_from_host_ptr_fn = ov_status_e (*)(ov_element_type_e, ov_shape_t, void*, ov_tensor_t**); +using ov_tensor_create_fn = + ov_status_e (*)(ov_element_type_e, ov_shape_t, ov_tensor_t**); using ov_tensor_free_fn = void (*)(ov_tensor_t*); +using ov_tensor_get_shape_fn = ov_status_e (*)(const ov_tensor_t*, ov_shape_t*); +using ov_tensor_get_byte_size_fn = ov_status_e (*)(const ov_tensor_t*, size_t*); +using ov_tensor_data_fn = ov_status_e (*)(const ov_tensor_t*, void**); using ov_shape_create_fn = ov_status_e (*)(int64_t, const int64_t*, ov_shape_t*); using ov_shape_free_fn = ov_status_e (*)(ov_shape_t*); @@ -106,6 +116,7 @@ struct OpenvinoFunctions { ov_core_free_fn core_free = nullptr; ov_core_get_available_devices_fn core_get_available_devices = nullptr; ov_available_devices_free_fn available_devices_free = nullptr; + ov_core_set_property_fn core_set_property = nullptr; ov_core_import_model_fn core_import_model = nullptr; ov_compiled_model_create_infer_request_fn compiled_model_create_infer_request = nullptr; @@ -117,9 +128,15 @@ struct OpenvinoFunctions { ov_infer_request_set_output_tensor_by_index_fn infer_request_set_output_tensor_by_index = nullptr; ov_infer_request_infer_fn infer_request_infer = nullptr; + ov_infer_request_get_output_tensor_by_index_fn + infer_request_get_output_tensor_by_index = nullptr; ov_infer_request_free_fn infer_request_free = nullptr; ov_tensor_create_from_host_ptr_fn tensor_create_from_host_ptr = nullptr; + ov_tensor_create_fn tensor_create = nullptr; ov_tensor_free_fn tensor_free = nullptr; + ov_tensor_get_shape_fn tensor_get_shape = nullptr; + ov_tensor_get_byte_size_fn tensor_get_byte_size = nullptr; + ov_tensor_data_fn tensor_data = nullptr; ov_shape_create_fn shape_create = nullptr; ov_shape_free_fn shape_free = nullptr; }; diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index 0689e5d4680..1fa66ba54f1 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include "OpenvinoBackend.h" @@ -233,6 +234,7 @@ bool OpenvinoBackend::ensure_loaded() const { LOAD_SYM(core_free, ov_core_free) LOAD_SYM(core_get_available_devices, ov_core_get_available_devices) LOAD_SYM(available_devices_free, ov_available_devices_free) + LOAD_SYM(core_set_property, ov_core_set_property) LOAD_SYM(core_import_model, ov_core_import_model) LOAD_SYM( compiled_model_create_infer_request, @@ -247,9 +249,16 @@ bool OpenvinoBackend::ensure_loaded() const { infer_request_set_output_tensor_by_index, ov_infer_request_set_output_tensor_by_index) LOAD_SYM(infer_request_infer, ov_infer_request_infer) + LOAD_SYM( + infer_request_get_output_tensor_by_index, + ov_infer_request_get_output_tensor_by_index) LOAD_SYM(infer_request_free, ov_infer_request_free) LOAD_SYM(tensor_create_from_host_ptr, ov_tensor_create_from_host_ptr) + LOAD_SYM(tensor_create, ov_tensor_create) LOAD_SYM(tensor_free, ov_tensor_free) + LOAD_SYM(tensor_get_shape, ov_tensor_get_shape) + LOAD_SYM(tensor_get_byte_size, ov_tensor_get_byte_size) + LOAD_SYM(tensor_data, ov_tensor_data) LOAD_SYM(shape_create, ov_shape_create) LOAD_SYM(shape_free, ov_shape_free) @@ -305,17 +314,55 @@ exr::Result OpenvinoBackend::init( size_t data_size = processed->size(); std::string device = "CPU"; + std::string precision; for (auto& compile_spec : compile_specs) { - if (std::strcmp(compile_spec.key, "device") == 0) { - const char* buf = static_cast(compile_spec.value.buffer); - size_t len = compile_spec.value.nbytes; - // Strip trailing null bytes that may be included in nbytes. - while (len > 0 && buf[len - 1] == '\0') { - --len; - } - if (len > 0) { - device.assign(buf, len); - } + const bool is_device = std::strcmp(compile_spec.key, "device") == 0; + if (!is_device && + std::strcmp(compile_spec.key, "inference_precision") != 0) { + continue; + } + const char* buf = static_cast(compile_spec.value.buffer); + size_t len = compile_spec.value.nbytes; + // Strip trailing null bytes that may be included in nbytes. + while (len > 0 && buf[len - 1] == '\0') { + --len; + } + if (len > 0) { + (is_device ? device : precision).assign(buf, len); + } + } + + // The same rule as preprocess.py, because on x86 CPUs OpenVINO applies the + // precision here, on import. + if (precision.empty() && device.rfind("CPU", 0) == 0) { + precision = "f32"; + } + // The same values as in preprocess.py. Some OpenVINO releases accept an + // unknown value, and others only reject it on import, with a generic + // error. + if (!precision.empty() && precision != "f32" && precision != "f16" && + precision != "bf16") { + ET_LOG( + Error, + "OpenVINO: unsupported inference_precision '%s', expected one of " + "f32, f16, bf16", + precision.c_str()); + ov_.core_free(core); + return exr::Error::InvalidArgument; + } + if (!precision.empty()) { + status = ov_.core_set_property( + core, device.c_str(), "INFERENCE_PRECISION_HINT", precision.c_str()); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, + "OpenVINO: failed to set inference precision '%s' for device '%s' " + "(status=%d)", + precision.c_str(), + device.c_str(), + status); + ov_.core_free(core); + return exr::Error::InvalidArgument; } } @@ -359,9 +406,26 @@ exr::Result OpenvinoBackend::init( ov_.compiled_model_free(compiled_model); return exr::Error::MemoryAllocationFailed; } + size_t num_outputs = 0; + status = ov_.compiled_model_outputs_size(compiled_model, &num_outputs); + ov_tensor_t** owned_outputs = nullptr; + if (status == OV_STATUS_OK && num_outputs > 0) { + owned_outputs = allocator->allocateList(num_outputs); + } + if (status != OV_STATUS_OK || (num_outputs > 0 && !owned_outputs)) { + ET_LOG(Error, "OpenVINO: failed to set up outputs (status=%d)", status); + ov_.infer_request_free(infer_request); + ov_.compiled_model_free(compiled_model); + return exr::Error::Internal; + } + for (size_t i = 0; i < num_outputs; i++) { + owned_outputs[i] = nullptr; + } new (handle) ExecutionHandle; handle->compiled_model = compiled_model; handle->infer_request = infer_request; + handle->owned_outputs = owned_outputs; + handle->num_outputs = num_outputs; return handle; } @@ -449,16 +513,35 @@ exr::Error OpenvinoBackend::execute( "OpenVINO: expected tensor for output %zu", i); auto output_tensor = args[num_inputs + i]->toTensor(); - auto result = create_ov_tensor(output_tensor); - if (!result.ok()) { - return result.error(); + ov_tensor_t*& owned = execution_handle->owned_outputs[i]; + ov_tensor_t* tensor = nullptr; + if (output_tensor.shape_dynamism() == exa::TensorShapeDynamism::STATIC) { + auto result = create_ov_tensor(output_tensor); + if (!result.ok()) { + return result.error(); + } + tensor = result.get(); + } else { + // OpenVINO cannot write more than a borrowed buffer's shape holds, and + // an earlier call may have shrunk that shape. So an output whose size + // can change goes to a buffer OpenVINO owns and grows, and is copied + // out after inference. + if (!owned) { + auto result = create_ov_tensor(output_tensor, /*own_memory=*/true); + if (!result.ok()) { + return result.error(); + } + owned = result.get(); + } + tensor = owned; } - ov_tensor_t* tensor = result.get(); status = ov_.infer_request_set_output_tensor_by_index( execution_handle->infer_request, i, tensor); // Safe to free: see shared_ptr ownership comment on input tensor above. - ov_.tensor_free(tensor); + if (tensor != owned) { + ov_.tensor_free(tensor); + } if (status != OV_STATUS_OK) { return exr::Error::Internal; } @@ -470,6 +553,27 @@ exr::Error OpenvinoBackend::execute( return exr::Error::Internal; } + for (size_t i = 0; i < num_outputs; i++) { + auto output_tensor = args[num_inputs + i]->toTensor(); + if (output_tensor.shape_dynamism() == exa::TensorShapeDynamism::STATIC) { + continue; + } + ov_tensor_t* ov_output = nullptr; + status = ov_.infer_request_get_output_tensor_by_index( + execution_handle->infer_request, i, &ov_output); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, "OpenVINO: failed to get output %zu (status=%d)", i, status); + return exr::Error::Internal; + } + exr::Error error = copy_output(ov_output, output_tensor); + ov_.tensor_free(ov_output); + if (error != exr::Error::Ok) { + ET_LOG(Error, "OpenVINO: failed to copy output %zu", i); + return error; + } + } + return exr::Error::Ok; } @@ -484,6 +588,13 @@ void OpenvinoBackend::destroy(exr::DelegateHandle* handle) const { ExecutionHandle* execution_handle = static_cast(handle); + for (size_t i = 0; i < execution_handle->num_outputs; i++) { + if (execution_handle->owned_outputs[i] && ov_.tensor_free) { + ov_.tensor_free(execution_handle->owned_outputs[i]); + execution_handle->owned_outputs[i] = nullptr; + } + } + if (execution_handle->infer_request && ov_.infer_request_free) { ov_.infer_request_free(execution_handle->infer_request); execution_handle->infer_request = nullptr; @@ -495,8 +606,45 @@ void OpenvinoBackend::destroy(exr::DelegateHandle* handle) const { } } +exr::Error OpenvinoBackend::copy_output( + const ov_tensor_t* ov_output, + const exa::Tensor& output) const { + ov_shape_t shape = {}; + ov_status_e status = ov_.tensor_get_shape(ov_output, &shape); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, "OpenVINO: failed to get an output shape (status=%d)", status); + return exr::Error::Internal; + } + exr::Error error = exr::resize_tensor( + output, + exa::ArrayRef(shape.dims, static_cast(shape.rank))); + ov_.shape_free(&shape); + if (error != exr::Error::Ok) { + return error; + } + size_t nbytes = 0; + void* data = nullptr; + if (ov_.tensor_get_byte_size(ov_output, &nbytes) != OV_STATUS_OK || + ov_.tensor_data(ov_output, &data) != OV_STATUS_OK) { + ET_LOG(Error, "OpenVINO: failed to read an output"); + return exr::Error::Internal; + } + ET_CHECK_OR_RETURN_ERROR( + nbytes == output.nbytes(), + Internal, + "OpenVINO: output has %zu bytes, expected %zu", + nbytes, + output.nbytes()); + if (nbytes > 0) { + std::memcpy(output.mutable_data_ptr(), data, nbytes); + } + return exr::Error::Ok; +} + exr::Result OpenvinoBackend::create_ov_tensor( - const exa::Tensor& tensor) const { + const exa::Tensor& tensor, + bool own_memory) const { ov_element_type_e ov_type = convert_to_openvino_type(tensor.scalar_type()); if (ov_type == OV_ELEMENT_UNDEFINED) { return exr::Error::NotSupported; @@ -517,22 +665,33 @@ exr::Result OpenvinoBackend::create_ov_tensor( dims = dims_heap.get(); } for (int64_t d = 0; d < rank; d++) { - dims[d] = sizes[d]; + dims[d] = own_memory ? 0 : sizes[d]; } - // shape is zero-initialized; shape_free is only needed after a successful - // shape_create (the zero state is safe to skip). + // shape_create refuses rank 0, so a 0-d tensor keeps the zero-initialized + // shape, which describes a scalar. ov_shape_t shape = {}; - ov_status_e status = ov_.shape_create(rank, dims, &shape); + ov_status_e status = OV_STATUS_OK; + if (rank > 0) { + status = ov_.shape_create(rank, dims, &shape); + } dims_heap.reset(); // Release heap dims (no-op if stack was used). if (status != OV_STATUS_OK) { + ET_LOG( + Error, + "OpenVINO: failed to create a shape of rank %" PRId64 " (status=%d)", + rank, + status); return exr::Error::Internal; } ov_tensor_t* ov_tensor = nullptr; - status = ov_.tensor_create_from_host_ptr( - ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); + status = own_memory + ? ov_.tensor_create(ov_type, shape, &ov_tensor) + : ov_.tensor_create_from_host_ptr( + ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); ov_.shape_free(&shape); if (status != OV_STATUS_OK || !ov_tensor) { + ET_LOG(Error, "OpenVINO: failed to create a tensor (status=%d)", status); return exr::Error::Internal; } return ov_tensor; diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index 72b0193dc2d..255aadfe3ee 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -28,6 +28,10 @@ namespace openvino { struct ExecutionHandle { ov_compiled_model_t* compiled_model = nullptr; ov_infer_request_t* infer_request = nullptr; + // A buffer OpenVINO owns for each output whose size can change, made on + // first use and reused by every call; null for the others. + ov_tensor_t** owned_outputs = nullptr; + size_t num_outputs = 0; }; class OpenvinoBackend final : public ::exr::BackendInterface { @@ -49,7 +53,15 @@ class OpenvinoBackend final : public ::exr::BackendInterface { private: bool ensure_loaded() const; ov_element_type_e convert_to_openvino_type(exa::ScalarType scalar_type) const; - exr::Result create_ov_tensor(const exa::Tensor& tensor) const; + // With own_memory, the OpenVINO tensor gets its own buffer instead of + // borrowing the tensor's. That buffer starts empty, with the tensor's rank, + // and OpenVINO grows it to the size it writes. + exr::Result create_ov_tensor( + const exa::Tensor& tensor, + bool own_memory = false) const; + exr::Error copy_output( + const ov_tensor_t* ov_output, + const exa::Tensor& output) const; mutable DlHandle lib_handle_; mutable OpenvinoFunctions ov_; diff --git a/backends/openvino/tests/ops/base_openvino_op_test.py b/backends/openvino/tests/ops/base_openvino_op_test.py index c429845548a..22fa78d1530 100644 --- a/backends/openvino/tests/ops/base_openvino_op_test.py +++ b/backends/openvino/tests/ops/base_openvino_op_test.py @@ -23,6 +23,8 @@ def execute_layer_test( sample_inputs: tuple[torch.Tensor], expected_partitions: int = 1, assert_output_equal: bool = True, + compile_specs: tuple[CompileSpec, ...] = (), + more_inputs: tuple[tuple[torch.Tensor, ...], ...] = (), ): module = module.eval() @@ -30,7 +32,7 @@ def execute_layer_test( aten_dialect: ExportedProgram = export(module, sample_inputs) # Convert to edge dialect and lower the module to the backend with a custom partitioner - compile_spec = [CompileSpec("device", self.device.encode())] + compile_spec = [CompileSpec("device", self.device.encode()), *compile_specs] lowered_module: EdgeProgramManager = to_edge_transform_and_lower( aten_dialect, partitioner=[ @@ -57,30 +59,40 @@ def execute_layer_test( # Execute the model and compare the outputs with the reference outputs if assert_output_equal: - # Execute the module in eager mode to calculate the reference outputs - ref_output = module(*sample_inputs) - if isinstance(ref_output, torch.Tensor): - ref_output = [ - ref_output, - ] + # Later inputs run on the same loaded method, as an app would. + all_inputs = (sample_inputs, *more_inputs) - # 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 + all_outputs = [method.execute(inputs) for inputs in all_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)) - for i in range(len(ref_output)): - self.assertTrue( - torch.allclose( - outputs[i], - ref_output[i], - atol=self.atol, - rtol=self.rtol, - equal_nan=True, - ), - msg=f"ref_output:\n{ref_output[i]}\n\ntest_output:\n{outputs[i]}", - ) + for inputs, outputs in zip(all_inputs, all_outputs): + # Execute the module in eager mode to calculate the reference outputs + ref_output = module(*inputs) + if isinstance(ref_output, torch.Tensor): + ref_output = [ + ref_output, + ] + + # Compare the outputs with the reference outputs + self.assertTrue(len(ref_output) == len(outputs)) + for i in range(len(ref_output)): + self.assertTrue( + torch.allclose( + outputs[i], + ref_output[i], + atol=self.atol, + rtol=self.rtol, + equal_nan=True, + ), + msg=f"ref_output:\n{ref_output[i]}\n\ntest_output:\n{outputs[i]}", + ) diff --git a/backends/openvino/tests/ops/test_add.py b/backends/openvino/tests/ops/test_add.py index 5b68d0ff149..a976f7cc7fd 100644 --- a/backends/openvino/tests/ops/test_add.py +++ b/backends/openvino/tests/ops/test_add.py @@ -20,3 +20,12 @@ def test_add(self): module = self.create_model() sample_input = (torch.randn(2, 5, 1, 3), torch.randn(2, 5, 1, 3)) self.execute_layer_test(module, sample_input) + + def test_add_0d_inputs(self): + module = self.create_model() + for sample_input in ( + (torch.randn(2, 3), torch.tensor(2.5)), + (torch.tensor(1.5), torch.tensor(2.5)), + ): + with self.subTest(shapes=[list(x.shape) for x in sample_input]): + self.execute_layer_test(module, sample_input) diff --git a/backends/openvino/tests/ops/test_addmm.py b/backends/openvino/tests/ops/test_addmm.py index 51c1314db0d..4a690e08b98 100644 --- a/backends/openvino/tests/ops/test_addmm.py +++ b/backends/openvino/tests/ops/test_addmm.py @@ -1,7 +1,15 @@ +# mypy: disable-error-code=import-not-found + + +from unittest import mock + import torch +from executorch.backends.openvino import preprocess from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( BaseOpenvinoOpTest, ) +from executorch.exir.backend.backend_details import CompileSpec +from openvino import Core, Type # type: ignore[import-untyped] class TestAddMMOperator(BaseOpenvinoOpTest): @@ -26,3 +34,51 @@ def test_addmm(self): input_z = torch.randn(4, 4, dtype=torch.float32) sample_input = (input_x, input_y, input_z) self.execute_layer_test(module, sample_input) + + def test_addmm_computes_in_float32(self): + if not self.device.startswith("CPU"): + self.skipTest("only CPUs compute in float32 by default") + if Core().get_property("CPU", "INFERENCE_PRECISION_HINT") == Type.f32: + self.skipTest("OpenVINO already computes in float32 on this CPU") + # OpenVINO would otherwise compute this in bfloat16 on x86 CPUs with + # bfloat16 units, or in float16 on ARM CPUs, far outside 1e-4. + self.atol = self.rtol = 1e-4 + sample_input = tuple(torch.randn(256, 256) for _ in range(3)) + self.execute_layer_test(self.create_model(), sample_input) + # On x86 only the runtime passes the precision on, so this fails if it + # drops an explicit one. + with self.assertRaisesRegex(AssertionError, "ref_output"): + self.execute_layer_test( + self.create_model(), + sample_input, + compile_specs=(CompileSpec("inference_precision", b"bf16"),), + ) + + def test_addmm_passes_the_precision_to_openvino(self): + if not self.device.startswith("CPU"): + self.skipTest("only CPUs compute in float32 by default") + # Checks the lowering half too: on x86 the runtime half alone decides. + self.atol = self.rtol = 0.05 + sample_input = tuple(torch.randn(4, 4) for _ in range(3)) + bf16 = (CompileSpec("inference_precision", b"bf16"),) + for compile_specs, hint in (((), "f32"), (bf16, "bf16")): + with self.subTest(hint=hint), mock.patch.object( + preprocess, "compile_model", wraps=preprocess.compile_model + ) as compile_model: + self.execute_layer_test( + self.create_model(), sample_input, compile_specs=compile_specs + ) + self.assertEqual( + compile_model.call_args.args[2], {"INFERENCE_PRECISION_HINT": hint} + ) + + def test_addmm_refuses_an_unknown_inference_precision_before_converting(self): + sample_input = tuple(torch.randn(4, 4) for _ in range(3)) + with mock.patch.object( + preprocess, "FrontEndManager", side_effect=AssertionError("converted") + ), self.assertRaisesRegex(ValueError, "unsupported inference_precision 'f7'"): + self.execute_layer_test( + self.create_model(), + sample_input, + compile_specs=(CompileSpec("inference_precision", b"f7"),), + ) diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py new file mode 100644 index 00000000000..66378bafb02 --- /dev/null +++ b/backends/openvino/tests/ops/test_reductions.py @@ -0,0 +1,220 @@ +# mypy: disable-error-code=import-not-found + +import itertools +import unittest +from unittest import mock + +import torch +from executorch.backends.openvino._passes import ( + fold_constant_cumsum_pass, + FoldConstantCumsumPass, +) +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from openvino.frontend.pytorch.torchdynamo.op_support import ( # type: ignore[import-untyped] + OperatorSupport, +) + + +class TestIntegerReductions(BaseOpenvinoOpTest): + """PyTorch reduces bool and small integer tensors to int64.""" + + def test_cumsum(self): + class CumSum(torch.nn.Module): + def forward(self, x): + return torch.cumsum(x, 1) + + for x in ( + torch.rand(2, 7) > 0.5, + torch.randint(0, 5, (2, 7), dtype=torch.int32), + ): + with self.subTest(dtype=x.dtype): + self.execute_layer_test(CumSum(), (x,)) + + def test_sum(self): + class Sum(torch.nn.Module): + def forward(self, x): + return x.sum(1) + + # More than 255 true values in a row, which 8 bits cannot count. + mask = torch.arange(800).reshape(2, 400) % 4 != 0 + for x in (mask, torch.randint(0, 5, (2, 7), dtype=torch.int32)): + with self.subTest(dtype=x.dtype): + self.execute_layer_test(Sum(), (x,)) + + def test_prod(self): + class Prod(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + return x.prod() if self.dim is None else x.prod(self.dim) + + support = OperatorSupport([])._support_dict + x = torch.randint(1, 4, (2, 5), dtype=torch.int32) + for dim, op in ((1, "prod.dim_int"), (None, "prod.default")): + with self.subTest(op=op): + if f"torch.ops.aten.{op}" not in support: + self.skipTest(f"this OpenVINO release does not delegate {op}") + self.execute_layer_test(Prod(dim), (x,)) + + def test_attention_masks_with_more_valid_tokens(self): + class AttentionMasks(torch.nn.Module): + """Masks and positions from padding, as pi0 builds them.""" + + def __init__(self, dtype, from_input): + super().__init__() + self.dtype = dtype + self.from_input = from_input + + def forward(self, pad_masks): + if self.from_input: + # Constant too, but only once OpenVINO folds the shape. + att_masks = torch.ones_like(pad_masks[0]) + else: + att_masks = torch.tensor([0] * 765 + [1] * 51, dtype=torch.bool) + cumsum = torch.cumsum(att_masks[None, :], dim=1, dtype=self.dtype) + att_2d = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d & pad_2d, torch.cumsum(pad_masks, dim=1) - 1 + + def pad_masks(valid): + return (torch.arange(816) < valid)[None, :] + + for dtype, from_input in itertools.product( + (torch.int64, torch.int32), (False, True) + ): + with self.subTest(dtype=dtype, from_input=from_input): + # One loaded program, run again with more valid tokens. + self.execute_layer_test( + AttentionMasks(dtype, from_input), + (pad_masks(300),), + more_inputs=((pad_masks(700),),), + ) + + +class TestFoldConstantCumsumPass(unittest.TestCase): + def fold(self, module, inputs): + """Edge program module of module, with constant cumsums folded.""" + edge = to_edge(torch.export.export(module, inputs)) + folded = edge.exported_program().module() + FoldConstantCumsumPass()(folded) + return folded + + def test_buffer_written_by_the_graph_is_not_folded(self): + class State(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("s", torch.arange(3)) + + def forward(self, x): + self.s.add_(1) + return torch.cumsum(self.s, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(State(), (x,)) + eager = State() + for _ in range(2): + torch.testing.assert_close(folded(x), eager(x)) + + def test_folded_name_does_not_replace_a_buffer(self): + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "_folded_cumsum_aten_cumsum_default", torch.ones(3) + ) + + def forward(self, x): + counts = torch.cumsum(torch.ones(3, dtype=torch.int64), 0) + return counts + x + self._folded_cumsum_aten_cumsum_default + + x = torch.zeros(3) + folded = self.fold(Model(), (x,)) + # The folded cumsum would have taken this name. + self.assertTrue(hasattr(folded, "_folded_cumsum_aten_cumsum_default_")) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_a_cumsum_that_cannot_be_evaluated_leaves_the_others_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + ones = torch.ones(3, dtype=torch.int64) + return torch.cumsum(ones, 0) + torch.cumsum(ones * 2, 0) + x + + evaluate = fold_constant_cumsum_pass._evaluate + + def fail_on_the_doubled_one(graph_module, targets): + if any("mul" in str(n.args[0].target) for n in targets): + raise RuntimeError("not on the CPU") + return evaluate(graph_module, targets) + + x = torch.zeros(3, dtype=torch.int64) + with mock.patch.object( + fold_constant_cumsum_pass, "_evaluate", side_effect=fail_on_the_doubled_one + ), self.assertLogs(level="WARNING") as logs: + folded = self.fold(Model(), (x,)) + targets = [n.target for n in folded.graph.nodes] + self.assertEqual(targets.count(exir_ops.edge.aten.cumsum.default), 1) + self.assertIn("wrong results", logs.output[0]) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_float_cumsum_is_not_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.cumsum(torch.ones(3), 0) + x + + folded = self.fold(Model(), (torch.zeros(3),)) + self.assertIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + + def test_random_draws_are_not_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + c = torch.ones(3) + r = torch.cond( + torch.tensor(True), + lambda c: (torch.rand(3) * 100 + c).long(), + lambda c: c.long(), + (c,), + ) + direct = (torch.rand(3) * 100).long() + return torch.cumsum(r, 0) + torch.cumsum(direct, 0) + x + + folded = self.fold(Model(), (torch.zeros(3, dtype=torch.int64),)) + targets = [n.target for n in folded.graph.nodes] + self.assertEqual(targets.count(exir_ops.edge.aten.cumsum.default), 2) + + def test_cumsum_over_a_split_constant_is_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + part = torch.ones(6, dtype=torch.int64).split(3)[1] + return torch.cumsum(part, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(Model(), (x,)) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_constants_are_folded_without_a_gpu(self): + class Model(torch.nn.Module): + def forward(self, x): + ones = torch.ones(3, dtype=torch.int64, device="cuda") + pinned = torch.ones(3, dtype=torch.int64, pin_memory=True) + return torch.cumsum(ones, 0).cpu() + torch.cumsum(pinned, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(Model(), (x,)) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), torch.tensor([2, 4, 6])) diff --git a/backends/openvino/tests/ops/test_view.py b/backends/openvino/tests/ops/test_view.py index 8aef13fffa0..792f39a9056 100644 --- a/backends/openvino/tests/ops/test_view.py +++ b/backends/openvino/tests/ops/test_view.py @@ -1,7 +1,16 @@ +import os +import tempfile + import torch +from executorch.backends.openvino.preprocess import OpenvinoBackend from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( BaseOpenvinoOpTest, ) +from executorch.exir import to_edge +from executorch.exir.backend.backend_api import to_backend +from executorch.exir.backend.backend_details import CompileSpec +from executorch.runtime import Runtime +from torch.export import Dim, export op_params = [ {"input_shape": [2, 3, 2], "target_shape": [2, 6]}, @@ -33,3 +42,75 @@ 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 __init__(self, reduce): + super().__init__() + self.reduce = reduce + + def forward(self, x, n): + n = n.item() + torch._check(n >= 1) + torch._check(n <= x.shape[0]) + y = x[:n].view(n, -1) + return y.sum(0) if self.reduce else y + + # The size reaches the OpenVINO partition as a SymInt input. The second + # call needs a larger output than the first one returned, unless the + # sum makes the output static. + x = torch.randn(4, 8) + for reduce in (False, True): + with self.subTest(reduce=reduce): + self.execute_layer_test( + View(reduce), + (x, torch.tensor(3)), + more_inputs=((x, torch.tensor(4)),), + ) + + def test_one_delegate_called_with_different_output_sizes(self): + class Double(torch.nn.Module): + def forward(self, x): + return x * 2 + + x, y = torch.randn(2, 8), torch.randn(4, 8) + double = to_backend( + OpenvinoBackend.__name__, + to_edge( + export(Double(), (y,), dynamic_shapes=({0: Dim("n", min=1, max=4)},)) + ).exported_program(), + [CompileSpec("device", self.device.encode())], + ) + # Programs from ExecuTorch 0.7 and earlier run every call of a lowered + # module on one delegate handle, each call with its own outputs. + double._share_handle_across_calls = True + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = double + + def forward(self, x, y): + return self.lowered(x), self.lowered(y) + + dynamic_shapes = ( + {0: Dim("rows_x", min=1, max=2)}, + {0: Dim("rows_y", min=1, max=4)}, + ) + program = to_edge( + export(TwoCalls(), (x, y), dynamic_shapes=dynamic_shapes) + ).to_executorch() + self.assertEqual(len(program.executorch_program.execution_plan[0].delegates), 1) + + 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: + program.write_to_file(f) + loaded = Runtime.get().load_program(pte_path) + method = loaded.load_method("forward") + assert method is not None + for inputs in ((x, y), (y[:1], y[:3])): + outputs = method.execute(inputs) + for output, value in zip(outputs, inputs): + torch.testing.assert_close(output, value * 2) + del method, loaded diff --git a/docs/source/build-run-openvino.md b/docs/source/build-run-openvino.md index 69a73672ac1..c5f38a2d0d7 100644 --- a/docs/source/build-run-openvino.md +++ b/docs/source/build-run-openvino.md @@ -102,6 +102,51 @@ Follow the steps below to setup your build environment: ./openvino_build.sh ``` +## Precision + +When the device is `CPU`, a program computes in float32 by default. Without this default, +OpenVINO would compute in bfloat16 on x86 CPUs with bfloat16 units, which keeps only two to +three significant digits, and in float16 on ARM CPUs, which keeps about three. +Both are faster. To trade accuracy for speed, ask for the faster precision when +lowering: + +```python +from executorch.backends.openvino.partitioner import OpenvinoPartitioner +from executorch.exir.backend.compile_spec_schema import CompileSpec + +OpenvinoPartitioner( + [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] +) +``` + +On an ARM CPU, ask for `f16` instead: there float32 is about 1.8x slower than +float16, as measured on an Orin Nano. + +The value is OpenVINO's `INFERENCE_PRECISION_HINT`: `f32`, `f16` or `bf16`. An +empty value means the default, and any other value fails at lowering. A CPU +without float16 support computes an `f16` program in float32. OpenVINO has no +bfloat16 on ARM CPUs, so there `bf16` keeps the ARM default, float16 where the +CPU supports it. Which x86 CPUs compute in bfloat16 is up to OpenVINO; see its +[Precision Control](https://docs.openvino.ai/2026/openvino-workflow/running-inference/optimize-inference/precision-control.html) +page. Other devices, such as GPUs, keep OpenVINO's default unless the option is +set. + +The default also applies to a program exported in float16: on a CPU it computes +in float32 unless you lower it with `inference_precision`. + +On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. +So a program lowered without `inference_precision` computes in float32, even +one exported with an ExecuTorch release that did not have this default. That +is slower and can use more memory. Export it again with `inference_precision` +set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +precision when the program is lowered, so a program exported with such an +earlier release keeps computing in float16 until it is exported again. + +OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result +outside the int32 range wraps around without an error. This applies to every +int64 tensor, including the int64 result of `cumsum` and `prod` over smaller +integer types. + ## Build Instructions for Examples ### AOT step: From 2e117c697617e469fe9e62043afd72ac9d39bc22 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 13:30:50 -0700 Subject: [PATCH 03/29] Fix 0-d tensors, dynamic output sizes and integer reductions, and default to float32 on CPU, in the OpenVINO backend (#24294) ## What is wrong today The OpenVINO backend has these problems: 1. A 0-d (scalar) input or output fails on every run, because OpenVINO's `ov_shape_create` refuses rank 0. 2. A float32 program does not compute in float32 on many CPUs. OpenVINO picks bfloat16 (two to three significant digits) on x86 CPUs with bfloat16 hardware and float16 on ARM CPUs, and no compile spec can change it. 3. An output whose size depends on an int input keeps its largest size. `x[:n].view(n, -1)` with x of shape (4, 8) and n = 3 returns (4, 8). 4. `cumsum` and `prod` of a bool or integer tensor are not computed in int64, as PyTorch does. A cumsum over a bool mask wraps after 255, so pi0 and pi0.5 get wrong position ids and wrong actions. ## What this change does - **0-d tensors:** a rank-0 tensor keeps its empty shape, which OpenVINO accepts. - **Precision:** when the device is `CPU`, a program computes in float32 by default, both when it is lowered and when it is loaded. A program can ask for a faster precision: ```python OpenvinoPartitioner( [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] ) ``` The value is `f32`, `f16` or `bf16`, checked before conversion starts. Other devices keep OpenVINO's default. On x86 CPUs OpenVINO applies the precision when the runtime loads the program, so the runtime passes it on too. - **Dynamic output sizes:** an output whose size can change goes to a buffer that OpenVINO owns, reused by every call. It starts empty and OpenVINO grows it to the largest size a call writes, so a large upper bound costs nothing until a call needs it. After inference the backend resizes the ExecuTorch output to OpenVINO's shape and copies the data in, so one delegate can also serve several calls, as programs from ExecuTorch 0.7 and earlier do. Static outputs still use the ExecuTorch buffer directly. - **Integer reductions:** the input of `cumsum` and `prod` is cast to int64 before conversion. An integer cumsum that reads only constants is computed during lowering, because OpenVINO's CPU plugin gives it wrong results from the second inference on. This includes a cumsum over `ones_like` or `zeros_like` of an input with a static shape, which OpenVINO folds into a constant too. Each one is folded on its own, and one that cannot be folded is named in a warning. Trade-offs: - float32 is slower. On a Xeon with bfloat16 hardware, a 4-layer MLP takes 2.4 ms in float32 and 1.5 ms in bfloat16. On x86 this also applies to programs exported before this change, because OpenVINO applies the precision when a program loads. To get the old speed, export again with `inference_precision` set to `bf16`. - OpenVINO's CPU plugin computes int64 in 32 bits, so an integer result past the int32 range wraps without an error. That was already true for int64 programs, and now also applies to `cumsum` and `prod` over smaller integer types. ## Tests Each new test fails without its fix: 0-d inputs and outputs; a 256 x 256 addmm checked to 1e-4, and the same addmm lowered with `bf16` checked not to reach 1e-4, so the runtime must apply an explicit precision (both skipped where OpenVINO already computes in float32); an unknown precision refused before conversion; dynamic output sizes, on one program and on one delegate shared by two calls; bool and int32 cumsum and prod; pi0-style masks with int64 and int32 cumsums over two calls, built from a constant and from `ones_like`; and folding when one cumsum cannot be evaluated. Two guards pass either way: integer `sum`, which OpenVINO already computes in int64, and a static output after a dynamic view. The full OpenVINO op suite passes on Linux x86_64 (a CPU with bfloat16 units) with OpenVINO 2026.4.1, through both pytest and the CI test runner. With OpenVINO 2025.1.0 it also passes, except that the prod test is skipped because that release does not delegate prod. (cherry picked from commit 7b14148ec5e5bae0c054dc26dc7be3a17ea8cbfb) --- backends/openvino/README.md | 45 ++++ backends/openvino/_passes/__init__.py | 8 +- .../_passes/cast_reduction_input_pass.py | 45 ++++ .../_passes/fold_constant_cumsum_pass.py | 156 +++++++++++++ backends/openvino/preprocess.py | 44 +++- backends/openvino/runtime/OpenvinoApi.h | 17 ++ backends/openvino/runtime/OpenvinoBackend.cpp | 203 ++++++++++++++-- backends/openvino/runtime/OpenvinoBackend.h | 14 +- .../tests/ops/base_openvino_op_test.py | 48 ++-- backends/openvino/tests/ops/test_add.py | 9 + backends/openvino/tests/ops/test_addmm.py | 56 +++++ .../openvino/tests/ops/test_reductions.py | 220 ++++++++++++++++++ backends/openvino/tests/ops/test_view.py | 76 +++++- docs/source/build-run-openvino.md | 45 ++++ 14 files changed, 930 insertions(+), 56 deletions(-) create mode 100644 backends/openvino/_passes/cast_reduction_input_pass.py create mode 100644 backends/openvino/_passes/fold_constant_cumsum_pass.py create mode 100644 backends/openvino/tests/ops/test_reductions.py diff --git a/backends/openvino/README.md b/backends/openvino/README.md index 331aebf87ff..40a245f12a0 100644 --- a/backends/openvino/README.md +++ b/backends/openvino/README.md @@ -48,6 +48,51 @@ print(_get_registered_backend_names()) # Should include 'OpenvinoBackend' ``` +## Precision + +When the device is `CPU`, a program computes in float32 by default. Without this default, +OpenVINO would compute in bfloat16 on x86 CPUs with bfloat16 units, which keeps only two to +three significant digits, and in float16 on ARM CPUs, which keeps about three. +Both are faster. To trade accuracy for speed, ask for the faster precision when +lowering: + +```python +from executorch.backends.openvino.partitioner import OpenvinoPartitioner +from executorch.exir.backend.compile_spec_schema import CompileSpec + +OpenvinoPartitioner( + [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] +) +``` + +On an ARM CPU, ask for `f16` instead: there float32 is about 1.8x slower than +float16, as measured on an Orin Nano. + +The value is OpenVINO's `INFERENCE_PRECISION_HINT`: `f32`, `f16` or `bf16`. An +empty value means the default, and any other value fails at lowering. A CPU +without float16 support computes an `f16` program in float32. OpenVINO has no +bfloat16 on ARM CPUs, so there `bf16` keeps the ARM default, float16 where the +CPU supports it. Which x86 CPUs compute in bfloat16 is up to OpenVINO; see its +[Precision Control](https://docs.openvino.ai/2026/openvino-workflow/running-inference/optimize-inference/precision-control.html) +page. Other devices, such as GPUs, keep OpenVINO's default unless the option is +set. + +The default also applies to a program exported in float16: on a CPU it computes +in float32 unless you lower it with `inference_precision`. + +On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. +So a program lowered without `inference_precision` computes in float32, even +one exported with an ExecuTorch release that did not have this default. That +is slower and can use more memory. Export it again with `inference_precision` +set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +precision when the program is lowered, so a program exported with such an +earlier release keeps computing in float16 until it is exported again. + +OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result +outside the int32 range wraps around without an error. This applies to every +int64 tensor, including the int64 result of `cumsum` and `prod` over smaller +integer types. + ## Directory Structure ``` diff --git a/backends/openvino/_passes/__init__.py b/backends/openvino/_passes/__init__.py index 0db94ca1ec8..bc2853a73e6 100644 --- a/backends/openvino/_passes/__init__.py +++ b/backends/openvino/_passes/__init__.py @@ -4,6 +4,12 @@ # except in compliance with the License. See the license file found in the # LICENSE file in the root directory of this source tree. +from .cast_reduction_input_pass import CastReductionInputPass from .decompose_floor_divide_pass import DecomposeFloorDividePass +from .fold_constant_cumsum_pass import FoldConstantCumsumPass -__all__ = ["DecomposeFloorDividePass"] +__all__ = [ + "CastReductionInputPass", + "DecomposeFloorDividePass", + "FoldConstantCumsumPass", +] diff --git a/backends/openvino/_passes/cast_reduction_input_pass.py b/backends/openvino/_passes/cast_reduction_input_pass.py new file mode 100644 index 00000000000..36e3f3795c0 --- /dev/null +++ b/backends/openvino/_passes/cast_reduction_input_pass.py @@ -0,0 +1,45 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +REDUCTION_OPS = { + exir_ops.edge.aten.cumsum.default, + exir_ops.edge.aten.prod.default, + exir_ops.edge.aten.prod.dim_int, +} + + +class CastReductionInputPass(PassBase): + """Cast the input of cumsum and prod to the dtype of their result. + + PyTorch reduces bool and integer tensors in int64, but OpenVINO keeps a + narrower type for these two, though not for sum. A cumsum over a bool mask + is computed in 8 bits and wraps after 255, and a cumsum or prod over int32 + returns int32. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + for node in graph_module.graph.nodes: + if node.op != "call_function" or node.target not in REDUCTION_OPS: + continue + source = node.args[0] + dtype = node.meta["val"].dtype + if source.meta["val"].dtype == dtype: + continue + with graph_module.graph.inserting_before(node): + cast = graph_module.graph.call_function( + exir_ops.edge.aten._to_copy.default, (source,), {"dtype": dtype} + ) + cast.meta["val"] = source.meta["val"].to(dtype) + node.replace_input_with(source, cast) + modified = True + if modified: + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/openvino/_passes/fold_constant_cumsum_pass.py b/backends/openvino/_passes/fold_constant_cumsum_pass.py new file mode 100644 index 00000000000..fabc030b06a --- /dev/null +++ b/backends/openvino/_passes/fold_constant_cumsum_pass.py @@ -0,0 +1,156 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import operator + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.constant_prop_pass import _nodes_written_in_place +from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.utils import _pytree as pytree + + +def _reads_only_shape(node: torch.fx.Node) -> bool: + """full_like with a constant fill value reads only the shape and dtype of + its input. With a static shape its result is a constant, and OpenVINO + folds it into one too. + """ + if node.target != exir_ops.edge.aten.full_like.default: + return False + if any( + isinstance(a, torch.fx.Node) for a in (*node.args[1:], *node.kwargs.values()) + ): + return False + val = node.meta.get("val") + return isinstance(val, torch.Tensor) and all(isinstance(d, int) for d in val.shape) + + +def _inputs(node: torch.fx.Node) -> list: + return [] if _reads_only_shape(node) else node.all_input_nodes + + +def _on_cpu(value): + if isinstance(value, torch.Tensor): + return value.cpu() + if isinstance(value, torch.device): + return torch.device("cpu") + return value + + +def _constant_nodes(graph_module: torch.fx.GraphModule) -> set: + """Nodes that depend only on constants. Graph order puts inputs first.""" + constants = set() + written = _nodes_written_in_place(graph_module) + for node in graph_module.graph.nodes: + # Only tensors: a higher-order op such as cond reads its subgraphs, + # which may hold a random op, through get_attr nodes too. + if node.op == "get_attr": + value = operator.attrgetter(node.target)(graph_module) + if isinstance(value, torch.Tensor) and node not in written: + constants.add(node) + elif node.op == "call_function" and ( + _reads_only_shape(node) + or ( + torch.Tag.nondeterministic_seeded + not in getattr(node.target, "tags", ()) + and all(n in constants for n in node.all_input_nodes) + ) + ): + constants.add(node) + return constants + + +def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict: + """Values of targets, computing only what they read, on the CPU.""" + needed = set() + stack = list(targets) + while stack: + node = stack.pop() + if node not in needed: + needed.add(node) + stack.extend(_inputs(node)) + order = [n for n in graph_module.graph.nodes if n in needed] + last_use = {i: n for n in order for i in _inputs(n)} + values = {} + for node in order: + if node.op == "get_attr": + values[node] = operator.attrgetter(node.target)(graph_module) + continue + if _reads_only_shape(node): + like = node.args[0].meta["val"] + args = (torch.empty(like.shape, dtype=like.dtype), *node.args[1:]) + kwargs = dict(node.kwargs) + else: + args, kwargs = torch.fx.node.map_arg( + (node.args, node.kwargs), lambda n: values[n] + ) + # The values are the same on any device, and lowering for OpenVINO + # must not need the GPU a constant was made for. + args, kwargs = pytree.tree_map(_on_cpu, (args, kwargs)) + if kwargs.get("pin_memory"): + kwargs = {**kwargs, "pin_memory": False} + values[node] = node.target(*args, **kwargs) + for i in _inputs(node): + if last_use[i] is node and i not in targets: + del values[i] + return values + + +class FoldConstantCumsumPass(PassBase): + """Replace an integer cumsum that reads only constants with its result. + + With an integer cumsum over a constant, OpenVINO's CPU plugin (2025.4 to + 2026.4) gives wrong results from the second inference on: masks built from + it keep only what earlier calls allowed. pi0 builds its attention masks + this way. The OpenVINO nightly no longer does this, so the pass can go once + the oldest supported OpenVINO release has the fix. Only these cumsums are + folded; constants that feed no such cumsum, such as compressed weights, are + left as they are. Runs on an unlifted module, where constants are get_attr + nodes; buffers the graph writes to are not constants, and full_like of an + input with a static shape is one. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + constants = _constant_nodes(graph_module) + cumsums = [ + n + for n in graph.nodes + if n in constants + and n.target == exir_ops.edge.aten.cumsum.default + and not n.meta["val"].dtype.is_floating_point + and not n.meta["val"].dtype.is_complex + ] + modified = False + # One at a time, so a cumsum that cannot be evaluated does not keep + # the others from being folded. + for node in cumsums: + try: + value = _evaluate(graph_module, [node])[node] + except RuntimeError as e: + # Some ops have modes the CPU lacks. + logging.warning( + "OpenVINO: could not fold %s, so on a CPU some OpenVINO " + "releases give wrong results for it from the second " + "inference on: %s", + node.name, + e, + ) + continue + name = f"_folded_cumsum_{node.name}" + while hasattr(graph_module, name): + name += "_" + graph_module.register_buffer(name, value) + with graph.inserting_before(node): + folded = graph.get_attr(name) + folded.meta = node.meta.copy() + node.replace_all_uses_with(folded) + modified = True + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 7a3f490b258..83537f26f38 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -9,7 +9,11 @@ from typing import final, List import torch -from executorch.backends.openvino._passes import DecomposeFloorDividePass +from executorch.backends.openvino._passes import ( + CastReductionInputPass, + DecomposeFloorDividePass, + FoldConstantCumsumPass, +) from executorch.exir.backend.backend_details import ( BackendDetails, @@ -25,6 +29,10 @@ ) from openvino.preprocess import PrePostProcessor # type: ignore[import-untyped] +# The values of the inference_precision compile spec. The runtime +# (OpenvinoBackend.cpp) accepts the same ones. +_PRECISIONS = ("f32", "f16", "bf16") + @final class OpenvinoBackend(BackendDetails): @@ -43,7 +51,25 @@ def preprocess( Returns: PreprocessResult: The result of preprocessing, including the compiled model bytes. """ - for pass_cls in [DimOrderOpsRevertPass, DecomposeFloorDividePass]: + compile_options = {} + for spec in module_compile_spec: + compile_options[spec.key] = spec.value.decode() + device = compile_options.get("device", "CPU") + requested_precision = compile_options.get("inference_precision") + # Some OpenVINO releases accept an unknown value. Checking before the + # conversion fails fast and keeps it out of the program, where the + # runtime would only fail at load. + if requested_precision and requested_precision not in _PRECISIONS: + raise ValueError( + f"OpenVINO backend: unsupported inference_precision " + f"'{requested_precision}', expected one of {', '.join(_PRECISIONS)}" + ) + + for pass_cls in [ + DimOrderOpsRevertPass, + DecomposeFloorDividePass, + CastReductionInputPass, + ]: result = pass_cls()(edge_program.graph_module) if result and result.graph_module: edge_program._graph_module = result.graph_module @@ -70,14 +96,12 @@ def preprocess( 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() - # The converter chooses some lowerings (e.g. boolean-mask indexing) from # the input types, so they must be known before conversion starts. + module = edge_program.module() + FoldConstantCumsumPass()(module) decoder = TorchFXPythonDecoder( - edge_program.module(), input_shapes=input_shapes, input_types=input_types + module, input_shapes=input_shapes, input_types=input_types ) frontend = FrontEndManager().load_by_framework("pytorch") ppp = PrePostProcessor(frontend.convert(frontend.load(decoder))) @@ -86,7 +110,11 @@ def preprocess( 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")) + # On ARM CPUs OpenVINO fixes the precision here, not when the runtime + # imports the model, so this must follow the same rule as the runtime. + precision = requested_precision or ("f32" if device.startswith("CPU") else None) + config = {"INFERENCE_PRECISION_HINT": precision} if precision else {} + compiled = compile_model(ppp.build(), device, config) model_bytes = compiled.export_model() return PreprocessResult(processed_bytes=model_bytes.getvalue()) diff --git a/backends/openvino/runtime/OpenvinoApi.h b/backends/openvino/runtime/OpenvinoApi.h index e4c3e7f737a..b4aa25f9d32 100644 --- a/backends/openvino/runtime/OpenvinoApi.h +++ b/backends/openvino/runtime/OpenvinoApi.h @@ -76,6 +76,9 @@ using ov_core_free_fn = void (*)(ov_core_t*); using ov_core_get_available_devices_fn = ov_status_e (*)(const ov_core_t*, ov_available_devices_t*); using ov_available_devices_free_fn = void (*)(ov_available_devices_t*); +// Takes one property key and value, both strings, after the device name. +using ov_core_set_property_fn = + ov_status_e (*)(const ov_core_t*, const char*, ...); using ov_core_import_model_fn = ov_status_e (*)( const ov_core_t*, const char*, @@ -94,10 +97,17 @@ using ov_infer_request_set_input_tensor_by_index_fn = using ov_infer_request_set_output_tensor_by_index_fn = ov_status_e (*)(ov_infer_request_t*, size_t, const ov_tensor_t*); using ov_infer_request_infer_fn = ov_status_e (*)(ov_infer_request_t*); +using ov_infer_request_get_output_tensor_by_index_fn = + ov_status_e (*)(const ov_infer_request_t*, size_t, ov_tensor_t**); using ov_infer_request_free_fn = void (*)(ov_infer_request_t*); using ov_tensor_create_from_host_ptr_fn = ov_status_e (*)(ov_element_type_e, ov_shape_t, void*, ov_tensor_t**); +using ov_tensor_create_fn = + ov_status_e (*)(ov_element_type_e, ov_shape_t, ov_tensor_t**); using ov_tensor_free_fn = void (*)(ov_tensor_t*); +using ov_tensor_get_shape_fn = ov_status_e (*)(const ov_tensor_t*, ov_shape_t*); +using ov_tensor_get_byte_size_fn = ov_status_e (*)(const ov_tensor_t*, size_t*); +using ov_tensor_data_fn = ov_status_e (*)(const ov_tensor_t*, void**); using ov_shape_create_fn = ov_status_e (*)(int64_t, const int64_t*, ov_shape_t*); using ov_shape_free_fn = ov_status_e (*)(ov_shape_t*); @@ -115,6 +125,7 @@ struct OpenvinoFunctions { ov_core_free_fn core_free = nullptr; ov_core_get_available_devices_fn core_get_available_devices = nullptr; ov_available_devices_free_fn available_devices_free = nullptr; + ov_core_set_property_fn core_set_property = nullptr; ov_core_import_model_fn core_import_model = nullptr; ov_compiled_model_create_infer_request_fn compiled_model_create_infer_request = nullptr; @@ -126,9 +137,15 @@ struct OpenvinoFunctions { ov_infer_request_set_output_tensor_by_index_fn infer_request_set_output_tensor_by_index = nullptr; ov_infer_request_infer_fn infer_request_infer = nullptr; + ov_infer_request_get_output_tensor_by_index_fn + infer_request_get_output_tensor_by_index = nullptr; ov_infer_request_free_fn infer_request_free = nullptr; ov_tensor_create_from_host_ptr_fn tensor_create_from_host_ptr = nullptr; + ov_tensor_create_fn tensor_create = nullptr; ov_tensor_free_fn tensor_free = nullptr; + ov_tensor_get_shape_fn tensor_get_shape = nullptr; + ov_tensor_get_byte_size_fn tensor_get_byte_size = nullptr; + ov_tensor_data_fn tensor_data = nullptr; ov_shape_create_fn shape_create = nullptr; ov_shape_free_fn shape_free = nullptr; }; diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index fad6ef91cb1..da9fa43c325 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include "OpenvinoBackend.h" @@ -236,6 +237,7 @@ bool OpenvinoBackend::ensure_loaded() const { LOAD_SYM(core_free, ov_core_free) LOAD_SYM(core_get_available_devices, ov_core_get_available_devices) LOAD_SYM(available_devices_free, ov_available_devices_free) + LOAD_SYM(core_set_property, ov_core_set_property) LOAD_SYM(core_import_model, ov_core_import_model) LOAD_SYM( compiled_model_create_infer_request, @@ -250,9 +252,16 @@ bool OpenvinoBackend::ensure_loaded() const { infer_request_set_output_tensor_by_index, ov_infer_request_set_output_tensor_by_index) LOAD_SYM(infer_request_infer, ov_infer_request_infer) + LOAD_SYM( + infer_request_get_output_tensor_by_index, + ov_infer_request_get_output_tensor_by_index) LOAD_SYM(infer_request_free, ov_infer_request_free) LOAD_SYM(tensor_create_from_host_ptr, ov_tensor_create_from_host_ptr) + LOAD_SYM(tensor_create, ov_tensor_create) LOAD_SYM(tensor_free, ov_tensor_free) + LOAD_SYM(tensor_get_shape, ov_tensor_get_shape) + LOAD_SYM(tensor_get_byte_size, ov_tensor_get_byte_size) + LOAD_SYM(tensor_data, ov_tensor_data) LOAD_SYM(shape_create, ov_shape_create) LOAD_SYM(shape_free, ov_shape_free) @@ -321,17 +330,55 @@ exr::Result OpenvinoBackend::init( size_t data_size = processed->size(); std::string device = "CPU"; + std::string precision; for (auto& compile_spec : compile_specs) { - if (std::strcmp(compile_spec.key, "device") == 0) { - const char* buf = static_cast(compile_spec.value.buffer); - size_t len = compile_spec.value.nbytes; - // Strip trailing null bytes that may be included in nbytes. - while (len > 0 && buf[len - 1] == '\0') { - --len; - } - if (len > 0) { - device.assign(buf, len); - } + const bool is_device = std::strcmp(compile_spec.key, "device") == 0; + if (!is_device && + std::strcmp(compile_spec.key, "inference_precision") != 0) { + continue; + } + const char* buf = static_cast(compile_spec.value.buffer); + size_t len = compile_spec.value.nbytes; + // Strip trailing null bytes that may be included in nbytes. + while (len > 0 && buf[len - 1] == '\0') { + --len; + } + if (len > 0) { + (is_device ? device : precision).assign(buf, len); + } + } + + // The same rule as preprocess.py, because on x86 CPUs OpenVINO applies the + // precision here, on import. + if (precision.empty() && device.rfind("CPU", 0) == 0) { + precision = "f32"; + } + // The same values as in preprocess.py. Some OpenVINO releases accept an + // unknown value, and others only reject it on import, with a generic + // error. + if (!precision.empty() && precision != "f32" && precision != "f16" && + precision != "bf16") { + ET_LOG( + Error, + "OpenVINO: unsupported inference_precision '%s', expected one of " + "f32, f16, bf16", + precision.c_str()); + ov_.core_free(core); + return exr::Error::InvalidArgument; + } + if (!precision.empty()) { + status = ov_.core_set_property( + core, device.c_str(), "INFERENCE_PRECISION_HINT", precision.c_str()); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, + "OpenVINO: failed to set inference precision '%s' for device '%s' " + "(status=%d)", + precision.c_str(), + device.c_str(), + status); + ov_.core_free(core); + return exr::Error::InvalidArgument; } } @@ -379,9 +426,26 @@ exr::Result OpenvinoBackend::init( ov_.compiled_model_free(compiled_model); return exr::Error::MemoryAllocationFailed; } + size_t num_outputs = 0; + status = ov_.compiled_model_outputs_size(compiled_model, &num_outputs); + ov_tensor_t** owned_outputs = nullptr; + if (status == OV_STATUS_OK && num_outputs > 0) { + owned_outputs = allocator->allocateList(num_outputs); + } + if (status != OV_STATUS_OK || (num_outputs > 0 && !owned_outputs)) { + ET_LOG(Error, "OpenVINO: failed to set up outputs (status=%d)", status); + ov_.infer_request_free(infer_request); + ov_.compiled_model_free(compiled_model); + return exr::Error::Internal; + } + for (size_t i = 0; i < num_outputs; i++) { + owned_outputs[i] = nullptr; + } new (handle) ExecutionHandle; handle->compiled_model = compiled_model; handle->infer_request = infer_request; + handle->owned_outputs = owned_outputs; + handle->num_outputs = num_outputs; return handle; } @@ -469,16 +533,35 @@ exr::Error OpenvinoBackend::execute( "OpenVINO: expected tensor for output %zu", i); auto output_tensor = args[num_inputs + i]->toTensor(); - auto result = create_ov_tensor(output_tensor); - if (!result.ok()) { - return result.error(); + ov_tensor_t*& owned = execution_handle->owned_outputs[i]; + ov_tensor_t* tensor = nullptr; + if (output_tensor.shape_dynamism() == exa::TensorShapeDynamism::STATIC) { + auto result = create_ov_tensor(output_tensor); + if (!result.ok()) { + return result.error(); + } + tensor = result.get(); + } else { + // OpenVINO cannot write more than a borrowed buffer's shape holds, and + // an earlier call may have shrunk that shape. So an output whose size + // can change goes to a buffer OpenVINO owns and grows, and is copied + // out after inference. + if (!owned) { + auto result = create_ov_tensor(output_tensor, /*own_memory=*/true); + if (!result.ok()) { + return result.error(); + } + owned = result.get(); + } + tensor = owned; } - ov_tensor_t* tensor = result.get(); status = ov_.infer_request_set_output_tensor_by_index( execution_handle->infer_request, i, tensor); // Safe to free: see shared_ptr ownership comment on input tensor above. - ov_.tensor_free(tensor); + if (tensor != owned) { + ov_.tensor_free(tensor); + } if (status != OV_STATUS_OK) { return exr::Error::Internal; } @@ -490,6 +573,27 @@ exr::Error OpenvinoBackend::execute( return exr::Error::Internal; } + for (size_t i = 0; i < num_outputs; i++) { + auto output_tensor = args[num_inputs + i]->toTensor(); + if (output_tensor.shape_dynamism() == exa::TensorShapeDynamism::STATIC) { + continue; + } + ov_tensor_t* ov_output = nullptr; + status = ov_.infer_request_get_output_tensor_by_index( + execution_handle->infer_request, i, &ov_output); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, "OpenVINO: failed to get output %zu (status=%d)", i, status); + return exr::Error::Internal; + } + exr::Error error = copy_output(ov_output, output_tensor); + ov_.tensor_free(ov_output); + if (error != exr::Error::Ok) { + ET_LOG(Error, "OpenVINO: failed to copy output %zu", i); + return error; + } + } + return exr::Error::Ok; } @@ -504,6 +608,13 @@ void OpenvinoBackend::destroy(exr::DelegateHandle* handle) const { ExecutionHandle* execution_handle = static_cast(handle); + for (size_t i = 0; i < execution_handle->num_outputs; i++) { + if (execution_handle->owned_outputs[i] && ov_.tensor_free) { + ov_.tensor_free(execution_handle->owned_outputs[i]); + execution_handle->owned_outputs[i] = nullptr; + } + } + if (execution_handle->infer_request && ov_.infer_request_free) { ov_.infer_request_free(execution_handle->infer_request); execution_handle->infer_request = nullptr; @@ -515,8 +626,45 @@ void OpenvinoBackend::destroy(exr::DelegateHandle* handle) const { } } +exr::Error OpenvinoBackend::copy_output( + const ov_tensor_t* ov_output, + const exa::Tensor& output) const { + ov_shape_t shape = {}; + ov_status_e status = ov_.tensor_get_shape(ov_output, &shape); + if (status != OV_STATUS_OK) { + ET_LOG( + Error, "OpenVINO: failed to get an output shape (status=%d)", status); + return exr::Error::Internal; + } + exr::Error error = exr::resize_tensor( + output, + exa::ArrayRef(shape.dims, static_cast(shape.rank))); + ov_.shape_free(&shape); + if (error != exr::Error::Ok) { + return error; + } + size_t nbytes = 0; + void* data = nullptr; + if (ov_.tensor_get_byte_size(ov_output, &nbytes) != OV_STATUS_OK || + ov_.tensor_data(ov_output, &data) != OV_STATUS_OK) { + ET_LOG(Error, "OpenVINO: failed to read an output"); + return exr::Error::Internal; + } + ET_CHECK_OR_RETURN_ERROR( + nbytes == output.nbytes(), + Internal, + "OpenVINO: output has %zu bytes, expected %zu", + nbytes, + output.nbytes()); + if (nbytes > 0) { + std::memcpy(output.mutable_data_ptr(), data, nbytes); + } + return exr::Error::Ok; +} + exr::Result OpenvinoBackend::create_ov_tensor( - const exa::Tensor& tensor) const { + const exa::Tensor& tensor, + bool own_memory) const { ov_element_type_e ov_type = convert_to_openvino_type(tensor.scalar_type()); if (ov_type == OV_ELEMENT_UNDEFINED) { return exr::Error::NotSupported; @@ -537,22 +685,33 @@ exr::Result OpenvinoBackend::create_ov_tensor( dims = dims_heap.get(); } for (int64_t d = 0; d < rank; d++) { - dims[d] = sizes[d]; + dims[d] = own_memory ? 0 : sizes[d]; } - // shape is zero-initialized; shape_free is only needed after a successful - // shape_create (the zero state is safe to skip). + // shape_create refuses rank 0, so a 0-d tensor keeps the zero-initialized + // shape, which describes a scalar. ov_shape_t shape = {}; - ov_status_e status = ov_.shape_create(rank, dims, &shape); + ov_status_e status = OV_STATUS_OK; + if (rank > 0) { + status = ov_.shape_create(rank, dims, &shape); + } dims_heap.reset(); // Release heap dims (no-op if stack was used). if (status != OV_STATUS_OK) { + ET_LOG( + Error, + "OpenVINO: failed to create a shape of rank %" PRId64 " (status=%d)", + rank, + status); return exr::Error::Internal; } ov_tensor_t* ov_tensor = nullptr; - status = ov_.tensor_create_from_host_ptr( - ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); + status = own_memory + ? ov_.tensor_create(ov_type, shape, &ov_tensor) + : ov_.tensor_create_from_host_ptr( + ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); ov_.shape_free(&shape); if (status != OV_STATUS_OK || !ov_tensor) { + ET_LOG(Error, "OpenVINO: failed to create a tensor (status=%d)", status); return exr::Error::Internal; } return ov_tensor; diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index 510e49ee690..d75fa725ab0 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -28,6 +28,10 @@ namespace openvino { struct ExecutionHandle { ov_compiled_model_t* compiled_model = nullptr; ov_infer_request_t* infer_request = nullptr; + // A buffer OpenVINO owns for each output whose size can change, made on + // first use and reused by every call; null for the others. + ov_tensor_t** owned_outputs = nullptr; + size_t num_outputs = 0; }; class OpenvinoBackend final : public ::exr::BackendInterface { @@ -49,7 +53,15 @@ class OpenvinoBackend final : public ::exr::BackendInterface { private: bool ensure_loaded() const; ov_element_type_e convert_to_openvino_type(exa::ScalarType scalar_type) const; - exr::Result create_ov_tensor(const exa::Tensor& tensor) const; + // With own_memory, the OpenVINO tensor gets its own buffer instead of + // borrowing the tensor's. That buffer starts empty, with the tensor's rank, + // and OpenVINO grows it to the size it writes. + exr::Result create_ov_tensor( + const exa::Tensor& tensor, + bool own_memory = false) const; + exr::Error copy_output( + const ov_tensor_t* ov_output, + const exa::Tensor& output) const; mutable DlHandle lib_handle_; mutable OpenvinoFunctions ov_; diff --git a/backends/openvino/tests/ops/base_openvino_op_test.py b/backends/openvino/tests/ops/base_openvino_op_test.py index f6bc928e4d4..4e20b41cbf6 100644 --- a/backends/openvino/tests/ops/base_openvino_op_test.py +++ b/backends/openvino/tests/ops/base_openvino_op_test.py @@ -25,6 +25,8 @@ def execute_layer_test( sample_inputs: tuple[torch.Tensor], expected_partitions: int = 1, assert_output_equal: bool = True, + compile_specs: tuple[CompileSpec, ...] = (), + more_inputs: tuple[tuple[torch.Tensor, ...], ...] = (), ): module = module.eval() @@ -32,7 +34,7 @@ def execute_layer_test( aten_dialect: ExportedProgram = export(module, sample_inputs) # Convert to edge dialect and lower the module to the backend with a custom partitioner - compile_spec = [CompileSpec("device", self.device.encode())] + compile_spec = [CompileSpec("device", self.device.encode()), *compile_specs] lowered_module: EdgeProgramManager = to_edge_transform_and_lower( aten_dialect, partitioner=[ @@ -59,12 +61,8 @@ def execute_layer_test( # Execute the model and compare the outputs with the reference outputs if assert_output_equal: - # Execute the module in eager mode to calculate the reference outputs - ref_output = module(*sample_inputs) - if isinstance(ref_output, torch.Tensor): - ref_output = [ - ref_output, - ] + # Later inputs run on the same loaded method, as an app would. + all_inputs = (sample_inputs, *more_inputs) # Load from a file: unlike a buffer, it really frees the delegate # data when the backend frees it. @@ -75,20 +73,28 @@ def execute_layer_test( program = Runtime.get().load_program(pte_path) method = program.load_method("forward") assert method is not None - outputs = method.execute(sample_inputs) + all_outputs = [method.execute(inputs) for inputs in all_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)) - for i in range(len(ref_output)): - self.assertTrue( - torch.allclose( - outputs[i], - ref_output[i], - atol=self.atol, - rtol=self.rtol, - equal_nan=True, - ), - msg=f"ref_output:\n{ref_output[i]}\n\ntest_output:\n{outputs[i]}", - ) + for inputs, outputs in zip(all_inputs, all_outputs): + # Execute the module in eager mode to calculate the reference outputs + ref_output = module(*inputs) + if isinstance(ref_output, torch.Tensor): + ref_output = [ + ref_output, + ] + + # Compare the outputs with the reference outputs + self.assertTrue(len(ref_output) == len(outputs)) + for i in range(len(ref_output)): + self.assertTrue( + torch.allclose( + outputs[i], + ref_output[i], + atol=self.atol, + rtol=self.rtol, + equal_nan=True, + ), + msg=f"ref_output:\n{ref_output[i]}\n\ntest_output:\n{outputs[i]}", + ) diff --git a/backends/openvino/tests/ops/test_add.py b/backends/openvino/tests/ops/test_add.py index 5b68d0ff149..a976f7cc7fd 100644 --- a/backends/openvino/tests/ops/test_add.py +++ b/backends/openvino/tests/ops/test_add.py @@ -20,3 +20,12 @@ def test_add(self): module = self.create_model() sample_input = (torch.randn(2, 5, 1, 3), torch.randn(2, 5, 1, 3)) self.execute_layer_test(module, sample_input) + + def test_add_0d_inputs(self): + module = self.create_model() + for sample_input in ( + (torch.randn(2, 3), torch.tensor(2.5)), + (torch.tensor(1.5), torch.tensor(2.5)), + ): + with self.subTest(shapes=[list(x.shape) for x in sample_input]): + self.execute_layer_test(module, sample_input) diff --git a/backends/openvino/tests/ops/test_addmm.py b/backends/openvino/tests/ops/test_addmm.py index 51c1314db0d..4a690e08b98 100644 --- a/backends/openvino/tests/ops/test_addmm.py +++ b/backends/openvino/tests/ops/test_addmm.py @@ -1,7 +1,15 @@ +# mypy: disable-error-code=import-not-found + + +from unittest import mock + import torch +from executorch.backends.openvino import preprocess from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( BaseOpenvinoOpTest, ) +from executorch.exir.backend.backend_details import CompileSpec +from openvino import Core, Type # type: ignore[import-untyped] class TestAddMMOperator(BaseOpenvinoOpTest): @@ -26,3 +34,51 @@ def test_addmm(self): input_z = torch.randn(4, 4, dtype=torch.float32) sample_input = (input_x, input_y, input_z) self.execute_layer_test(module, sample_input) + + def test_addmm_computes_in_float32(self): + if not self.device.startswith("CPU"): + self.skipTest("only CPUs compute in float32 by default") + if Core().get_property("CPU", "INFERENCE_PRECISION_HINT") == Type.f32: + self.skipTest("OpenVINO already computes in float32 on this CPU") + # OpenVINO would otherwise compute this in bfloat16 on x86 CPUs with + # bfloat16 units, or in float16 on ARM CPUs, far outside 1e-4. + self.atol = self.rtol = 1e-4 + sample_input = tuple(torch.randn(256, 256) for _ in range(3)) + self.execute_layer_test(self.create_model(), sample_input) + # On x86 only the runtime passes the precision on, so this fails if it + # drops an explicit one. + with self.assertRaisesRegex(AssertionError, "ref_output"): + self.execute_layer_test( + self.create_model(), + sample_input, + compile_specs=(CompileSpec("inference_precision", b"bf16"),), + ) + + def test_addmm_passes_the_precision_to_openvino(self): + if not self.device.startswith("CPU"): + self.skipTest("only CPUs compute in float32 by default") + # Checks the lowering half too: on x86 the runtime half alone decides. + self.atol = self.rtol = 0.05 + sample_input = tuple(torch.randn(4, 4) for _ in range(3)) + bf16 = (CompileSpec("inference_precision", b"bf16"),) + for compile_specs, hint in (((), "f32"), (bf16, "bf16")): + with self.subTest(hint=hint), mock.patch.object( + preprocess, "compile_model", wraps=preprocess.compile_model + ) as compile_model: + self.execute_layer_test( + self.create_model(), sample_input, compile_specs=compile_specs + ) + self.assertEqual( + compile_model.call_args.args[2], {"INFERENCE_PRECISION_HINT": hint} + ) + + def test_addmm_refuses_an_unknown_inference_precision_before_converting(self): + sample_input = tuple(torch.randn(4, 4) for _ in range(3)) + with mock.patch.object( + preprocess, "FrontEndManager", side_effect=AssertionError("converted") + ), self.assertRaisesRegex(ValueError, "unsupported inference_precision 'f7'"): + self.execute_layer_test( + self.create_model(), + sample_input, + compile_specs=(CompileSpec("inference_precision", b"f7"),), + ) diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py new file mode 100644 index 00000000000..66378bafb02 --- /dev/null +++ b/backends/openvino/tests/ops/test_reductions.py @@ -0,0 +1,220 @@ +# mypy: disable-error-code=import-not-found + +import itertools +import unittest +from unittest import mock + +import torch +from executorch.backends.openvino._passes import ( + fold_constant_cumsum_pass, + FoldConstantCumsumPass, +) +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from openvino.frontend.pytorch.torchdynamo.op_support import ( # type: ignore[import-untyped] + OperatorSupport, +) + + +class TestIntegerReductions(BaseOpenvinoOpTest): + """PyTorch reduces bool and small integer tensors to int64.""" + + def test_cumsum(self): + class CumSum(torch.nn.Module): + def forward(self, x): + return torch.cumsum(x, 1) + + for x in ( + torch.rand(2, 7) > 0.5, + torch.randint(0, 5, (2, 7), dtype=torch.int32), + ): + with self.subTest(dtype=x.dtype): + self.execute_layer_test(CumSum(), (x,)) + + def test_sum(self): + class Sum(torch.nn.Module): + def forward(self, x): + return x.sum(1) + + # More than 255 true values in a row, which 8 bits cannot count. + mask = torch.arange(800).reshape(2, 400) % 4 != 0 + for x in (mask, torch.randint(0, 5, (2, 7), dtype=torch.int32)): + with self.subTest(dtype=x.dtype): + self.execute_layer_test(Sum(), (x,)) + + def test_prod(self): + class Prod(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + return x.prod() if self.dim is None else x.prod(self.dim) + + support = OperatorSupport([])._support_dict + x = torch.randint(1, 4, (2, 5), dtype=torch.int32) + for dim, op in ((1, "prod.dim_int"), (None, "prod.default")): + with self.subTest(op=op): + if f"torch.ops.aten.{op}" not in support: + self.skipTest(f"this OpenVINO release does not delegate {op}") + self.execute_layer_test(Prod(dim), (x,)) + + def test_attention_masks_with_more_valid_tokens(self): + class AttentionMasks(torch.nn.Module): + """Masks and positions from padding, as pi0 builds them.""" + + def __init__(self, dtype, from_input): + super().__init__() + self.dtype = dtype + self.from_input = from_input + + def forward(self, pad_masks): + if self.from_input: + # Constant too, but only once OpenVINO folds the shape. + att_masks = torch.ones_like(pad_masks[0]) + else: + att_masks = torch.tensor([0] * 765 + [1] * 51, dtype=torch.bool) + cumsum = torch.cumsum(att_masks[None, :], dim=1, dtype=self.dtype) + att_2d = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d & pad_2d, torch.cumsum(pad_masks, dim=1) - 1 + + def pad_masks(valid): + return (torch.arange(816) < valid)[None, :] + + for dtype, from_input in itertools.product( + (torch.int64, torch.int32), (False, True) + ): + with self.subTest(dtype=dtype, from_input=from_input): + # One loaded program, run again with more valid tokens. + self.execute_layer_test( + AttentionMasks(dtype, from_input), + (pad_masks(300),), + more_inputs=((pad_masks(700),),), + ) + + +class TestFoldConstantCumsumPass(unittest.TestCase): + def fold(self, module, inputs): + """Edge program module of module, with constant cumsums folded.""" + edge = to_edge(torch.export.export(module, inputs)) + folded = edge.exported_program().module() + FoldConstantCumsumPass()(folded) + return folded + + def test_buffer_written_by_the_graph_is_not_folded(self): + class State(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("s", torch.arange(3)) + + def forward(self, x): + self.s.add_(1) + return torch.cumsum(self.s, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(State(), (x,)) + eager = State() + for _ in range(2): + torch.testing.assert_close(folded(x), eager(x)) + + def test_folded_name_does_not_replace_a_buffer(self): + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "_folded_cumsum_aten_cumsum_default", torch.ones(3) + ) + + def forward(self, x): + counts = torch.cumsum(torch.ones(3, dtype=torch.int64), 0) + return counts + x + self._folded_cumsum_aten_cumsum_default + + x = torch.zeros(3) + folded = self.fold(Model(), (x,)) + # The folded cumsum would have taken this name. + self.assertTrue(hasattr(folded, "_folded_cumsum_aten_cumsum_default_")) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_a_cumsum_that_cannot_be_evaluated_leaves_the_others_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + ones = torch.ones(3, dtype=torch.int64) + return torch.cumsum(ones, 0) + torch.cumsum(ones * 2, 0) + x + + evaluate = fold_constant_cumsum_pass._evaluate + + def fail_on_the_doubled_one(graph_module, targets): + if any("mul" in str(n.args[0].target) for n in targets): + raise RuntimeError("not on the CPU") + return evaluate(graph_module, targets) + + x = torch.zeros(3, dtype=torch.int64) + with mock.patch.object( + fold_constant_cumsum_pass, "_evaluate", side_effect=fail_on_the_doubled_one + ), self.assertLogs(level="WARNING") as logs: + folded = self.fold(Model(), (x,)) + targets = [n.target for n in folded.graph.nodes] + self.assertEqual(targets.count(exir_ops.edge.aten.cumsum.default), 1) + self.assertIn("wrong results", logs.output[0]) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_float_cumsum_is_not_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.cumsum(torch.ones(3), 0) + x + + folded = self.fold(Model(), (torch.zeros(3),)) + self.assertIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + + def test_random_draws_are_not_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + c = torch.ones(3) + r = torch.cond( + torch.tensor(True), + lambda c: (torch.rand(3) * 100 + c).long(), + lambda c: c.long(), + (c,), + ) + direct = (torch.rand(3) * 100).long() + return torch.cumsum(r, 0) + torch.cumsum(direct, 0) + x + + folded = self.fold(Model(), (torch.zeros(3, dtype=torch.int64),)) + targets = [n.target for n in folded.graph.nodes] + self.assertEqual(targets.count(exir_ops.edge.aten.cumsum.default), 2) + + def test_cumsum_over_a_split_constant_is_folded(self): + class Model(torch.nn.Module): + def forward(self, x): + part = torch.ones(6, dtype=torch.int64).split(3)[1] + return torch.cumsum(part, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(Model(), (x,)) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), Model()(x)) + + def test_constants_are_folded_without_a_gpu(self): + class Model(torch.nn.Module): + def forward(self, x): + ones = torch.ones(3, dtype=torch.int64, device="cuda") + pinned = torch.ones(3, dtype=torch.int64, pin_memory=True) + return torch.cumsum(ones, 0).cpu() + torch.cumsum(pinned, 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = self.fold(Model(), (x,)) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), torch.tensor([2, 4, 6])) diff --git a/backends/openvino/tests/ops/test_view.py b/backends/openvino/tests/ops/test_view.py index 5ea2f8cb1dc..792f39a9056 100644 --- a/backends/openvino/tests/ops/test_view.py +++ b/backends/openvino/tests/ops/test_view.py @@ -1,7 +1,16 @@ +import os +import tempfile + import torch +from executorch.backends.openvino.preprocess import OpenvinoBackend from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( BaseOpenvinoOpTest, ) +from executorch.exir import to_edge +from executorch.exir.backend.backend_api import to_backend +from executorch.exir.backend.backend_details import CompileSpec +from executorch.runtime import Runtime +from torch.export import Dim, export op_params = [ {"input_shape": [2, 3, 2], "target_shape": [2, 6]}, @@ -36,11 +45,72 @@ def test_view(self): def test_view_data_dependent_size(self): class View(torch.nn.Module): + def __init__(self, reduce): + super().__init__() + self.reduce = reduce + 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) + y = x[:n].view(n, -1) + return y.sum(0) if self.reduce else y + + # The size reaches the OpenVINO partition as a SymInt input. The second + # call needs a larger output than the first one returned, unless the + # sum makes the output static. + x = torch.randn(4, 8) + for reduce in (False, True): + with self.subTest(reduce=reduce): + self.execute_layer_test( + View(reduce), + (x, torch.tensor(3)), + more_inputs=((x, torch.tensor(4)),), + ) + + def test_one_delegate_called_with_different_output_sizes(self): + class Double(torch.nn.Module): + def forward(self, x): + return x * 2 + + x, y = torch.randn(2, 8), torch.randn(4, 8) + double = to_backend( + OpenvinoBackend.__name__, + to_edge( + export(Double(), (y,), dynamic_shapes=({0: Dim("n", min=1, max=4)},)) + ).exported_program(), + [CompileSpec("device", self.device.encode())], + ) + # Programs from ExecuTorch 0.7 and earlier run every call of a lowered + # module on one delegate handle, each call with its own outputs. + double._share_handle_across_calls = True + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = double + + def forward(self, x, y): + return self.lowered(x), self.lowered(y) + + dynamic_shapes = ( + {0: Dim("rows_x", min=1, max=2)}, + {0: Dim("rows_y", min=1, max=4)}, + ) + program = to_edge( + export(TwoCalls(), (x, y), dynamic_shapes=dynamic_shapes) + ).to_executorch() + self.assertEqual(len(program.executorch_program.execution_plan[0].delegates), 1) - # The size reaches the OpenVINO partition as a SymInt input. - self.execute_layer_test(View(), (torch.randn(4, 8), torch.tensor(3))) + 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: + program.write_to_file(f) + loaded = Runtime.get().load_program(pte_path) + method = loaded.load_method("forward") + assert method is not None + for inputs in ((x, y), (y[:1], y[:3])): + outputs = method.execute(inputs) + for output, value in zip(outputs, inputs): + torch.testing.assert_close(output, value * 2) + del method, loaded diff --git a/docs/source/build-run-openvino.md b/docs/source/build-run-openvino.md index 69a73672ac1..c5f38a2d0d7 100644 --- a/docs/source/build-run-openvino.md +++ b/docs/source/build-run-openvino.md @@ -102,6 +102,51 @@ Follow the steps below to setup your build environment: ./openvino_build.sh ``` +## Precision + +When the device is `CPU`, a program computes in float32 by default. Without this default, +OpenVINO would compute in bfloat16 on x86 CPUs with bfloat16 units, which keeps only two to +three significant digits, and in float16 on ARM CPUs, which keeps about three. +Both are faster. To trade accuracy for speed, ask for the faster precision when +lowering: + +```python +from executorch.backends.openvino.partitioner import OpenvinoPartitioner +from executorch.exir.backend.compile_spec_schema import CompileSpec + +OpenvinoPartitioner( + [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")] +) +``` + +On an ARM CPU, ask for `f16` instead: there float32 is about 1.8x slower than +float16, as measured on an Orin Nano. + +The value is OpenVINO's `INFERENCE_PRECISION_HINT`: `f32`, `f16` or `bf16`. An +empty value means the default, and any other value fails at lowering. A CPU +without float16 support computes an `f16` program in float32. OpenVINO has no +bfloat16 on ARM CPUs, so there `bf16` keeps the ARM default, float16 where the +CPU supports it. Which x86 CPUs compute in bfloat16 is up to OpenVINO; see its +[Precision Control](https://docs.openvino.ai/2026/openvino-workflow/running-inference/optimize-inference/precision-control.html) +page. Other devices, such as GPUs, keep OpenVINO's default unless the option is +set. + +The default also applies to a program exported in float16: on a CPU it computes +in float32 unless you lower it with `inference_precision`. + +On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. +So a program lowered without `inference_precision` computes in float32, even +one exported with an ExecuTorch release that did not have this default. That +is slower and can use more memory. Export it again with `inference_precision` +set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +precision when the program is lowered, so a program exported with such an +earlier release keeps computing in float16 until it is exported again. + +OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result +outside the int32 range wraps around without an error. This applies to every +int64 tensor, including the int64 result of `cumsum` and `prod` over smaller +integer types. + ## Build Instructions for Examples ### AOT step: From 333bc8f14ceaa87a1d2df8a1919039d0ecd22a37 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 13:08:48 -0700 Subject: [PATCH 04/29] Fail the OpenVINO test runner when tests fail The OpenVINO CI job runs backends/openvino/tests/test_runner.py. When a test failed, the runner printed "completed with failures" but still exited with code 0, so the CI job stayed green and the failure went unnoticed. Now the runner exits with code 1 when any test fails, so CI turns red. This is the test runner fix from #24305. Co-authored-by: suryasidd --- backends/openvino/tests/test_runner.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backends/openvino/tests/test_runner.py b/backends/openvino/tests/test_runner.py index 7d8c6b968c2..0c839e0c347 100644 --- a/backends/openvino/tests/test_runner.py +++ b/backends/openvino/tests/test_runner.py @@ -1,4 +1,5 @@ import argparse +import sys import unittest @@ -71,3 +72,4 @@ def parse_arguments(): print("OpenVINO backend tests completed successfully") else: print("OpenVINO backend tests completed with failures") + sys.exit(1) From 22140a8d94804768510d581c94e783b0168267a4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 08:46:02 -0700 Subject: [PATCH 05/29] Say that a CUDA program fails to load with another model's weights file The TensorRT guide said that when a second export overwrites aoti_cuda_blob.ptd, the first model can return wrong answers with no error. CUDA weight keys now carry a hash of their contents, so the first model's weights are missing from the new file and loading fails instead. Docs only. --- docs/source/backends/tensorrt/tensorrt-overview.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index 411cc262961..3bbb6f2249a 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -171,8 +171,8 @@ The Torch-TensorRT docs call such a program *coalesced*. The export writes the weights of the CUDA part to a data file, `aoti_cuda_blob.ptd`, next to the `.pte`. Pass that file when you load the program: if the CUDA part has weights, loading fails without it. The file has the same name for every model, so a second export into the same directory -overwrites the first one's weights, and the first model can then return wrong answers with no -error. Save each model into its own directory. +overwrites the first one's weights, and the first model then fails to load. Save each model into +its own directory. ### Keeping inputs and outputs on the GPU From 39e69af3171eb3eea56eb6b48f1e5af4389083c1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 09:02:27 -0700 Subject: [PATCH 06/29] Qualcomm AI Engine Direct - Decompose ops that read only constants for HTP On HTP, the partitioner keeps a partition that reads nothing at runtime on the CPU, because HTP refuses such a graph. That happens after decomposition, so an op that QNN asks to keep whole, such as linear or hardswish, could land on the CPU still undecomposed, and to_edge_transform_and_lower failed with "was not decomposed or delegated". The filter returned with the ops to keep whole now also leaves out any op that reads only constants, so it is decomposed like other CPU ops. An op that reads a runtime input is still kept whole. The constant-only partition test now has a linear in its constant branch, and fails without this change. --- .../qualcomm/partition/qnn_partitioner.py | 36 ++++++++++++++----- backends/qualcomm/tests/test_passes.py | 7 ++-- 2 files changed, 33 insertions(+), 10 deletions(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 640d6fbf478..b399e8b52a7 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -68,6 +68,15 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) + +def _is_runtime_input( + node: torch.fx.Node, edge_program: torch.export.ExportedProgram +) -> bool: + return not is_parameter(node, edge_program) or is_mutable_buffer_input( + node, edge_program + ) + + # dequantize ops with a CPU out variant; others cannot fall back CPU_DQ_OPS = { exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, @@ -319,17 +328,11 @@ def _drop_partitions_without_runtime_inputs( 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) + _is_runtime_input(arg, edge_program) for node in partition.nodes for arg in node.all_input_nodes if arg not in partition.nodes @@ -399,4 +402,21 @@ def ops_to_not_decompose( The filter function (`filter_fn`) can be used to further refine which nodes are not decomposed. (advanced use case) """ do_not_decompose = get_skip_decomp_table() - return (do_not_decompose, filter_fn) + if self.backend != QnnExecuTorchBackendType.kHtpBackend: + return (do_not_decompose, filter_fn) + # An op that reads only constants can end up in a partition that + # _drop_partitions_without_runtime_inputs keeps on CPU, so it must be + # decomposed like any other CPU op. + constant_only = set() + for node in ep.graph.nodes: + if node.op == "placeholder": + if not _is_runtime_input(node, ep): + constant_only.add(node) + elif node.op == "call_function" and all( + arg in constant_only for arg in node.all_input_nodes + ): + constant_only.add(node) + return ( + do_not_decompose, + lambda node: node not in constant_only and filter_fn(node), + ) diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 58ce89a6e30..43ef8c8e8a0 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -946,16 +946,19 @@ class ConstantBranch(torch.nn.Module): def __init__(self): super().__init__() self.register_buffer("table", torch.randn(4, 16)) + self.register_buffer("weight", torch.randn(16, 16)) def forward(self, x): - return x + torch.relu(self.table * 2 - 1) + table = torch.nn.functional.linear(self.table, self.weight) + return x + torch.relu(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. + # With relu on CPU, the branch before it reads nothing at runtime. It + # also holds linear, which QNN otherwise keeps from decomposing. try: edge = to_edge_transform_and_lower_to_qnn( ConstantBranch().eval(), From a9c1ac49dd8e55b01b974cdcc22f87a68c3bd1e3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 10:57:37 -0700 Subject: [PATCH 07/29] Qualcomm AI Engine Direct - Say that HTP decomposes every constant-only op The comment said ops that read only constants are decomposed like other CPU ops. Partitions are not known when the filter is built, so the rule also decomposes such ops whose partition stays delegated. The comment and the ops_to_not_decompose docstring now say so. Comment only. --- backends/qualcomm/partition/qnn_partitioner.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index b399e8b52a7..c40e0e60627 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -399,14 +399,17 @@ def ops_to_not_decompose( """ Determines which op should not be decomposed during partitioning. The list of operators is obtained from `get_skip_decomp_table()`. - The filter function (`filter_fn`) can be used to further refine which nodes are not decomposed. (advanced use case) + The filter refines which nodes are not decomposed: `filter_fn`, and on + HTP it also leaves out ops that read only constants. """ do_not_decompose = get_skip_decomp_table() if self.backend != QnnExecuTorchBackendType.kHtpBackend: return (do_not_decompose, filter_fn) # An op that reads only constants can end up in a partition that - # _drop_partitions_without_runtime_inputs keeps on CPU, so it must be - # decomposed like any other CPU op. + # _drop_partitions_without_runtime_inputs keeps on CPU after + # decomposition, so every such op is decomposed. Partitions are not + # known yet, so this also decomposes ones whose partition stays + # delegated; folding constant-only subgraphs first would avoid that. constant_only = set() for node in ep.graph.nodes: if node.op == "placeholder": From f1f41d1334723edad0e6123618f544bb98615f08 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:39:37 -0700 Subject: [PATCH 08/29] Qualcomm AI Engine Direct - Keep quantize and dequantize ops whole on HTP The HTP rule that decomposes ops reading only constants also covered the quantize and dequantize ops in the skip table, such as dequantize_per_channel_group on a frozen weight. QNN builds the weight encoding from those ops. With today's torch they are not decomposed anyway, so nothing changed in practice, but the filter should not ask for it. They are now kept whole, and ops after them still count as reading only constants. A new test checks the filter itself and needs no QNN SDK, so it runs in CI: a constant-only linear is decomposed, a linear on an input and a constant-only group-wise dequantize are kept whole. --- .../qualcomm/partition/qnn_partitioner.py | 13 ++++- backends/qualcomm/tests/test_passes.py | 55 ++++++++++++++++++- 2 files changed, 65 insertions(+), 3 deletions(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index c40e0e60627..3bff323a736 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -400,7 +400,8 @@ def ops_to_not_decompose( Determines which op should not be decomposed during partitioning. The list of operators is obtained from `get_skip_decomp_table()`. The filter refines which nodes are not decomposed: `filter_fn`, and on - HTP it also leaves out ops that read only constants. + HTP it also leaves out ops that read only constants, except quantize + and dequantize ops. """ do_not_decompose = get_skip_decomp_table() if self.backend != QnnExecuTorchBackendType.kHtpBackend: @@ -419,7 +420,15 @@ def ops_to_not_decompose( arg in constant_only for arg in node.all_input_nodes ): constant_only.add(node) + # Quantize and dequantize ops stay whole: QNN builds weight encodings + # from them. + quant_ops = { + op + for op in do_not_decompose + if op.namespace in ("quantized_decomposed", "torchao") + } return ( do_not_decompose, - lambda node: node not in constant_only and filter_fn(node), + lambda node: (node not in constant_only or node.target in quant_ops) + and filter_fn(node), ) diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 43ef8c8e8a0..e009c7b1466 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -31,7 +31,10 @@ ) from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader from executorch.backends.qualcomm.builders.utils import is_parameter -from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport +from executorch.backends.qualcomm.partition.qnn_partitioner import ( + QnnOperatorSupport, + QnnPartitioner, +) from executorch.backends.qualcomm.qnn_preprocess import QnnBackend from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype from executorch.backends.qualcomm.serialization.qc_schema import ( @@ -972,6 +975,56 @@ def forward(self, x): raise edge.to_executorch() + def test_htp_keeps_whole_only_ops_that_read_a_runtime_input(self): + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("table", torch.randn(4, 16)) + self.register_buffer("weight", torch.randn(16, 16)) + self.register_buffer( + "qweight", torch.randint(-8, 8, (16, 32), dtype=torch.int8) + ) + self.register_buffer("scales", torch.rand(16, 2) + 0.1) + self.register_buffer("zero_points", torch.zeros(16, 2).long()) + + def forward(self, x, y): + weight = torch.ops.quantized_decomposed.dequantize_per_channel_group( + self.qweight, + self.scales, + self.zero_points, + -8, + 7, + torch.int8, + 16, + torch.float32, + ) + table = torch.nn.functional.linear(self.table, self.weight) + return torch.nn.functional.linear(x, self.weight) + table, ( + torch.nn.functional.linear(y, weight) + ) + + ep = torch.export.export(Model(), (torch.randn(4, 16), torch.randn(4, 32))) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep) + linears = { + node.args[0].name: node + for node in ep.graph.nodes + if node.target == torch.ops.aten.linear.default + } + placeholders = {b: p for p, b in ep.graph_signature.inputs_to_buffers.items()} + self.assertFalse(keep_whole(linears[placeholders["table"]])) + self.assertTrue(keep_whole(linears["x"])) + (dequantize,) = ( + node + for node in ep.graph.nodes + if node.target + == torch.ops.quantized_decomposed.dequantize_per_channel_group.default + ) + self.assertTrue(keep_whole(dequantize)) + def test_partitioner_falls_back_on_bfloat16(self): class Bf16Region(torch.nn.Module): def __init__(self): From d59b4020c8bddda0d4f33e09edb16b3a9788f49c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:37:57 -0700 Subject: [PATCH 09/29] Qualcomm AI Engine Direct - Keep reporting kept-whole ops left on CPU After lowering, exir asks the partitioner for its filter again and reports any kept-whole op still on the CPU. On HTP the filter now let constant-only ops through, so that report went quiet. For example RecomposeRmsNorm rebuilt a constant-only rms_norm that then stayed on the CPU, and the program failed later in to_executorch with "Missing out variants" instead of the clear error from before. On the lowered (edge) program the filter is plain filter_fn again, so these cases report as they did before the HTP change. The filter test checks this too. --- backends/qualcomm/partition/qnn_partitioner.py | 4 +++- backends/qualcomm/tests/test_passes.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 3bff323a736..f824578c487 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -404,7 +404,9 @@ def ops_to_not_decompose( and dequantize ops. """ do_not_decompose = get_skip_decomp_table() - if self.backend != QnnExecuTorchBackendType.kHtpBackend: + # On the lowered program this filter tells exir which kept-whole ops + # may stay on CPU, and none may, so it stays plain filter_fn there. + if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE": return (do_not_decompose, filter_fn) # An op that reads only constants can end up in a partition that # _drop_partitions_without_runtime_inputs keeps on CPU after diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index e009c7b1466..1d97ca1cba3 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -1024,6 +1024,24 @@ def forward(self, x, y): == torch.ops.quantized_decomposed.dequantize_per_channel_group.default ) self.assertTrue(keep_whole(dequantize)) + # After lowering, exir uses the filter to report a kept-whole op left + # on CPU, so there it must not let a constant-only one through. + edge = to_edge( + ep, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + preserve_ops=[torch.ops.aten.linear.default], + ), + ).exported_program() + _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(edge) + placeholders = {b: p for p, b in edge.graph_signature.inputs_to_buffers.items()} + (constant_linear,) = ( + node + for node in edge.graph.nodes + if node.target == exir_ops.edge.aten.linear.default + and node.args[0].name == placeholders["table"] + ) + self.assertTrue(keep_whole(constant_linear)) def test_partitioner_falls_back_on_bfloat16(self): class Bf16Region(torch.nn.Module): From 8f342434bb843c945f65237da3a70bc9871c8851 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:34:59 -0700 Subject: [PATCH 10/29] Keep the first OpenVINO partition input when only a folded cumsum read it A position cumsum such as cumsum(ones_like(input_ids)) - 1 reads only the shape of input_ids. Once FoldConstantCumsumPass folds it, nothing in the partition reads input_ids. When that input comes first, OpenVINO's PyTorch converter takes it for forward's self and removes it. The runtime still passes every partition input, so every call failed with "expected 2 args (inputs=1 + outputs=1), got 3". An unused input in any other position is kept by the converter. Not folding is no fix, because masks built from such a cumsum then go wrong on OpenVINO's CPU plugin from the second call on. Instead, when the converted model has exactly one parameter fewer than the partition has inputs, preprocess puts the first one back, so the compiled model takes the same inputs the delegate passes. The test leaves both the first input and a later one unused. --- backends/openvino/preprocess.py | 17 ++++++++++++++++- backends/openvino/tests/ops/test_reductions.py | 15 +++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 83537f26f38..86038ccd65e 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -27,6 +27,9 @@ from openvino.frontend.pytorch.fx_decoder import ( # type: ignore[import-untyped] TorchFXPythonDecoder, ) +from openvino.frontend.pytorch.utils import ( # type: ignore[import-untyped] + pt_to_ov_type_map, +) from openvino.preprocess import PrePostProcessor # type: ignore[import-untyped] # The values of the inference_precision compile spec. The runtime @@ -104,7 +107,19 @@ def preprocess( module, input_shapes=input_shapes, input_types=input_types ) frontend = FrontEndManager().load_by_framework("pytorch") - ppp = PrePostProcessor(frontend.convert(frontend.load(decoder))) + model = frontend.convert(frontend.load(decoder)) + # The converter takes an unused first input for forward's self and + # removes it, but the runtime passes every partition input. The cumsum + # fold can leave the first input unused. + parameters = list(model.get_parameters()) + if len(parameters) == len(input_types) - 1: + for parameter in parameters: + model.remove_parameter(parameter) + first = opset13.parameter( + input_shapes[0], pt_to_ov_type_map[str(input_types[0])] + ) + model.add_parameters([first, *parameters]) + ppp = PrePostProcessor(model) # 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: diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index 66378bafb02..d91833111ae 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -96,6 +96,21 @@ def pad_masks(valid): more_inputs=((pad_masks(700),),), ) + def test_cumsum_over_ones_like_of_inputs_read_nowhere_else(self): + class Positions(torch.nn.Module): + def forward(self, input_ids, offset, mask): + positions = torch.cumsum(torch.ones_like(input_ids), 1) - 1 + return positions + offset * torch.cumsum(torch.ones_like(mask), 1) + + # Once the cumsums are folded, nothing in the partition reads input_ids + # or mask: the first input and one after it. + input_ids = torch.zeros(1, 8, dtype=torch.int64) + mask = torch.zeros(1, 8, dtype=torch.int32) + self.execute_layer_test( + Positions(), + (input_ids, torch.arange(8)[None], mask), + ) + class TestFoldConstantCumsumPass(unittest.TestCase): def fold(self, module, inputs): From 8625490c92e36dd2bf3abffeccf18a84688479dc Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:39:42 -0700 Subject: [PATCH 11/29] Give OpenVINO-owned output buffers the shape OpenVINO compiled ExecuTorch can plan an output with a dynamic size while OpenVINO compiles a fixed shape for it, for example x[torch.ones(4, dtype=bool)] once OpenVINO folds the mask. The buffer the runtime makes for such an output started empty, and OpenVINO refuses an output tensor whose shape differs from a fixed one, so every call failed with 0x1. Before the owned buffers, these programs ran. The buffer now starts with the shape of the infer request's own output tensor: empty along each dynamic dimension and the full shape where OpenVINO fixed it. --- backends/openvino/runtime/OpenvinoBackend.cpp | 54 +++++++++++++++---- backends/openvino/runtime/OpenvinoBackend.h | 14 ++--- backends/openvino/tests/ops/test_view.py | 9 ++++ 3 files changed, 60 insertions(+), 17 deletions(-) diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index da9fa43c325..efa31857106 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -547,11 +547,11 @@ exr::Error OpenvinoBackend::execute( // can change goes to a buffer OpenVINO owns and grows, and is copied // out after inference. if (!owned) { - auto result = create_ov_tensor(output_tensor, /*own_memory=*/true); - if (!result.ok()) { - return result.error(); + exr::Error error = create_owned_output( + execution_handle->infer_request, i, output_tensor, &owned); + if (error != exr::Error::Ok) { + return error; } - owned = result.get(); } tensor = owned; } @@ -662,9 +662,43 @@ exr::Error OpenvinoBackend::copy_output( return exr::Error::Ok; } +exr::Error OpenvinoBackend::create_owned_output( + ov_infer_request_t* infer_request, + size_t index, + const exa::Tensor& output, + ov_tensor_t** owned) const { + ov_element_type_e ov_type = convert_to_openvino_type(output.scalar_type()); + if (ov_type == OV_ELEMENT_UNDEFINED) { + return exr::Error::NotSupported; + } + // The request's own output has the compiled shape, empty along each dynamic + // dimension. OpenVINO refuses any other shape for a dimension it fixed, + // even when ExecuTorch planned that output with a dynamic size. + ov_tensor_t* current = nullptr; + ov_shape_t shape = {}; + ov_status_e status = ov_.infer_request_get_output_tensor_by_index( + infer_request, index, ¤t); + if (status == OV_STATUS_OK) { + status = ov_.tensor_get_shape(current, &shape); + ov_.tensor_free(current); + } + if (status == OV_STATUS_OK) { + status = ov_.tensor_create(ov_type, shape, owned); + ov_.shape_free(&shape); + } + if (status != OV_STATUS_OK || !*owned) { + ET_LOG( + Error, + "OpenVINO: failed to create a buffer for output %zu (status=%d)", + index, + status); + return exr::Error::Internal; + } + return exr::Error::Ok; +} + exr::Result OpenvinoBackend::create_ov_tensor( - const exa::Tensor& tensor, - bool own_memory) const { + const exa::Tensor& tensor) const { ov_element_type_e ov_type = convert_to_openvino_type(tensor.scalar_type()); if (ov_type == OV_ELEMENT_UNDEFINED) { return exr::Error::NotSupported; @@ -685,7 +719,7 @@ exr::Result OpenvinoBackend::create_ov_tensor( dims = dims_heap.get(); } for (int64_t d = 0; d < rank; d++) { - dims[d] = own_memory ? 0 : sizes[d]; + dims[d] = sizes[d]; } // shape_create refuses rank 0, so a 0-d tensor keeps the zero-initialized // shape, which describes a scalar. @@ -705,10 +739,8 @@ exr::Result OpenvinoBackend::create_ov_tensor( } ov_tensor_t* ov_tensor = nullptr; - status = own_memory - ? ov_.tensor_create(ov_type, shape, &ov_tensor) - : ov_.tensor_create_from_host_ptr( - ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); + status = ov_.tensor_create_from_host_ptr( + ov_type, shape, tensor.mutable_data_ptr(), &ov_tensor); ov_.shape_free(&shape); if (status != OV_STATUS_OK || !ov_tensor) { ET_LOG(Error, "OpenVINO: failed to create a tensor (status=%d)", status); diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index d75fa725ab0..2ab05fdebc6 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -53,12 +53,14 @@ class OpenvinoBackend final : public ::exr::BackendInterface { private: bool ensure_loaded() const; ov_element_type_e convert_to_openvino_type(exa::ScalarType scalar_type) const; - // With own_memory, the OpenVINO tensor gets its own buffer instead of - // borrowing the tensor's. That buffer starts empty, with the tensor's rank, - // and OpenVINO grows it to the size it writes. - exr::Result create_ov_tensor( - const exa::Tensor& tensor, - bool own_memory = false) const; + exr::Result create_ov_tensor(const exa::Tensor& tensor) const; + // Creates a buffer OpenVINO owns and grows for an output whose size can + // change. It starts with the shape OpenVINO compiled for that output. + exr::Error create_owned_output( + ov_infer_request_t* infer_request, + size_t index, + const exa::Tensor& output, + ov_tensor_t** owned) const; exr::Error copy_output( const ov_tensor_t* ov_output, const exa::Tensor& output) const; diff --git a/backends/openvino/tests/ops/test_view.py b/backends/openvino/tests/ops/test_view.py index 792f39a9056..a5a71c607b2 100644 --- a/backends/openvino/tests/ops/test_view.py +++ b/backends/openvino/tests/ops/test_view.py @@ -68,6 +68,15 @@ def forward(self, x, n): more_inputs=((x, torch.tensor(4)),), ) + def test_dynamic_size_output_that_openvino_makes_static(self): + class Select(torch.nn.Module): + def forward(self, x): + return x[torch.ones(4, dtype=torch.bool)] + + # ExecuTorch plans the output with a dynamic size, but OpenVINO folds + # the mask and compiles a fixed output shape. + self.execute_layer_test(Select(), (torch.randn(4, 8),)) + def test_one_delegate_called_with_different_output_sizes(self): class Double(torch.nn.Module): def forward(self, x): From e89b46ae51ac9d9f31f9defe39d7603529f20d39 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:44:17 -0700 Subject: [PATCH 12/29] Evaluate OpenVINO full_like folds on the CPU under any default device FoldConstantCumsumPass evaluates a full_like from an empty tensor of its input's shape. That tensor went to the default device, so lowering for OpenVINO under torch.device("cuda") or set_default_device("cuda") failed on a CPU-only PyTorch with "Torch not compiled with CUDA enabled", and on a CUDA build it used the GPU for nothing. The pass evaluates everything else on the CPU already; this tensor now goes there too. --- .../_passes/fold_constant_cumsum_pass.py | 5 ++++- backends/openvino/tests/ops/test_reductions.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/backends/openvino/_passes/fold_constant_cumsum_pass.py b/backends/openvino/_passes/fold_constant_cumsum_pass.py index fabc030b06a..f346096c832 100644 --- a/backends/openvino/_passes/fold_constant_cumsum_pass.py +++ b/backends/openvino/_passes/fold_constant_cumsum_pass.py @@ -82,7 +82,10 @@ def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict: continue if _reads_only_shape(node): like = node.args[0].meta["val"] - args = (torch.empty(like.shape, dtype=like.dtype), *node.args[1:]) + args = ( + torch.empty(like.shape, dtype=like.dtype, device="cpu"), + *node.args[1:], + ) kwargs = dict(node.kwargs) else: args, kwargs = torch.fx.node.map_arg( diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index d91833111ae..6f99859dcc8 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -233,3 +233,19 @@ def forward(self, x): [n.target for n in folded.graph.nodes], ) torch.testing.assert_close(folded(x), torch.tensor([2, 4, 6])) + + def test_full_like_is_folded_under_another_default_device(self): + class Model(torch.nn.Module): + def forward(self, x): + return torch.cumsum(torch.ones_like(x), 0) + x + + x = torch.zeros(3, dtype=torch.int64) + folded = to_edge(torch.export.export(Model(), (x,))).exported_program().module() + # Meta, unlike CUDA, fails on every build: it never reaches the CPU. + with torch.device("meta"): + FoldConstantCumsumPass()(folded) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), torch.tensor([1, 2, 3])) From b2e71b0681a6697add01adfc9b056fe923584018 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:44:46 -0700 Subject: [PATCH 13/29] Let the last OpenVINO inference_precision entry win at load too With inference_precision given twice, for example bf16 and then an empty value, lowering keeps the last entry and so compiles with the float32 default, but the runtime skipped empty entries and applied bf16. On x86, where the runtime decides, the program then computed in bfloat16. The runtime now takes the last entry as well, including an empty one, which means the default. --- backends/openvino/runtime/OpenvinoBackend.cpp | 4 +--- backends/openvino/tests/ops/test_addmm.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index efa31857106..e0a0e95c7ae 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -343,9 +343,7 @@ exr::Result OpenvinoBackend::init( while (len > 0 && buf[len - 1] == '\0') { --len; } - if (len > 0) { - (is_device ? device : precision).assign(buf, len); - } + (is_device ? device : precision).assign(buf, len); } // The same rule as preprocess.py, because on x86 CPUs OpenVINO applies the diff --git a/backends/openvino/tests/ops/test_addmm.py b/backends/openvino/tests/ops/test_addmm.py index 4a690e08b98..d4052638afc 100644 --- a/backends/openvino/tests/ops/test_addmm.py +++ b/backends/openvino/tests/ops/test_addmm.py @@ -47,12 +47,17 @@ def test_addmm_computes_in_float32(self): self.execute_layer_test(self.create_model(), sample_input) # On x86 only the runtime passes the precision on, so this fails if it # drops an explicit one. + bf16 = CompileSpec("inference_precision", b"bf16") with self.assertRaisesRegex(AssertionError, "ref_output"): self.execute_layer_test( - self.create_model(), - sample_input, - compile_specs=(CompileSpec("inference_precision", b"bf16"),), + self.create_model(), sample_input, compile_specs=(bf16,) ) + # The last value wins, as in lowering, and an empty one is the default. + self.execute_layer_test( + self.create_model(), + sample_input, + compile_specs=(bf16, CompileSpec("inference_precision", b"")), + ) def test_addmm_passes_the_precision_to_openvino(self): if not self.device.startswith("CPU"): From 05b5d8f5b2a5c0528c4b72f40f8aba2adc4ebaa0 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 18:37:45 -0700 Subject: [PATCH 14/29] Return OpenVINO-owned output buffers as a Result, like create_ov_tensor create_owned_output reported errors through its return value and its buffer through an out-parameter, and took a whole tensor only to read its dtype. It now returns exr::Result and takes the dtype, like create_ov_tensor beside it, so its call site reads like the one for create_ov_tensor again. No behavior change. --- backends/openvino/runtime/OpenvinoBackend.cpp | 23 ++++++++++--------- backends/openvino/runtime/OpenvinoBackend.h | 5 ++-- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index e0a0e95c7ae..8004a418211 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -545,11 +545,12 @@ exr::Error OpenvinoBackend::execute( // can change goes to a buffer OpenVINO owns and grows, and is copied // out after inference. if (!owned) { - exr::Error error = create_owned_output( - execution_handle->infer_request, i, output_tensor, &owned); - if (error != exr::Error::Ok) { - return error; + auto result = create_owned_output( + execution_handle->infer_request, i, output_tensor.scalar_type()); + if (!result.ok()) { + return result.error(); } + owned = result.get(); } tensor = owned; } @@ -660,12 +661,11 @@ exr::Error OpenvinoBackend::copy_output( return exr::Error::Ok; } -exr::Error OpenvinoBackend::create_owned_output( +exr::Result OpenvinoBackend::create_owned_output( ov_infer_request_t* infer_request, size_t index, - const exa::Tensor& output, - ov_tensor_t** owned) const { - ov_element_type_e ov_type = convert_to_openvino_type(output.scalar_type()); + exa::ScalarType scalar_type) const { + ov_element_type_e ov_type = convert_to_openvino_type(scalar_type); if (ov_type == OV_ELEMENT_UNDEFINED) { return exr::Error::NotSupported; } @@ -680,11 +680,12 @@ exr::Error OpenvinoBackend::create_owned_output( status = ov_.tensor_get_shape(current, &shape); ov_.tensor_free(current); } + ov_tensor_t* owned = nullptr; if (status == OV_STATUS_OK) { - status = ov_.tensor_create(ov_type, shape, owned); + status = ov_.tensor_create(ov_type, shape, &owned); ov_.shape_free(&shape); } - if (status != OV_STATUS_OK || !*owned) { + if (status != OV_STATUS_OK || !owned) { ET_LOG( Error, "OpenVINO: failed to create a buffer for output %zu (status=%d)", @@ -692,7 +693,7 @@ exr::Error OpenvinoBackend::create_owned_output( status); return exr::Error::Internal; } - return exr::Error::Ok; + return owned; } exr::Result OpenvinoBackend::create_ov_tensor( diff --git a/backends/openvino/runtime/OpenvinoBackend.h b/backends/openvino/runtime/OpenvinoBackend.h index 2ab05fdebc6..d2651e8717d 100644 --- a/backends/openvino/runtime/OpenvinoBackend.h +++ b/backends/openvino/runtime/OpenvinoBackend.h @@ -56,11 +56,10 @@ class OpenvinoBackend final : public ::exr::BackendInterface { exr::Result create_ov_tensor(const exa::Tensor& tensor) const; // Creates a buffer OpenVINO owns and grows for an output whose size can // change. It starts with the shape OpenVINO compiled for that output. - exr::Error create_owned_output( + exr::Result create_owned_output( ov_infer_request_t* infer_request, size_t index, - const exa::Tensor& output, - ov_tensor_t** owned) const; + exa::ScalarType scalar_type) const; exr::Error copy_output( const ov_tensor_t* ov_output, const exa::Tensor& output) const; From d4d99f6796b4b8c4aaefc6cb3971fe17c352c569 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 18:50:02 -0700 Subject: [PATCH 15/29] Test on any device that the OpenVINO runtime takes the last inference_precision The check that the runtime takes the last inference_precision entry ran only on x86 CPUs with bfloat16 units, so CI could skip it. A new test gives an invalid value and then an empty one. Lowering keeps only the empty one, and the runtime has to as well, or it fails to load the program. On the runtime before that fix it fails with "unsupported inference_precision 'f7'". --- backends/openvino/tests/ops/test_addmm.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/backends/openvino/tests/ops/test_addmm.py b/backends/openvino/tests/ops/test_addmm.py index d4052638afc..64b685d2721 100644 --- a/backends/openvino/tests/ops/test_addmm.py +++ b/backends/openvino/tests/ops/test_addmm.py @@ -59,6 +59,17 @@ def test_addmm_computes_in_float32(self): compile_specs=(bf16, CompileSpec("inference_precision", b"")), ) + def test_addmm_takes_the_last_inference_precision(self): + # Lowering keeps only the last value. The runtime must too, or it fails + # to load this invalid first one, on any device. + compile_specs = tuple( + CompileSpec("inference_precision", value) for value in (b"f7", b"") + ) + sample_input = tuple(torch.randn(4, 4) for _ in range(3)) + self.execute_layer_test( + self.create_model(), sample_input, compile_specs=compile_specs + ) + def test_addmm_passes_the_precision_to_openvino(self): if not self.device.startswith("CPU"): self.skipTest("only CPUs compute in float32 by default") From 29e8a4cba35850ce8bfa6c91b75ec5b9905665e1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 14:56:02 -0700 Subject: [PATCH 16/29] Leave bool min, max, clamp and sign to the CPU in the OpenVINO partitioner OpenVINO converts min, max, amin, amax, minimum, maximum, clamp and sign to its ReduceMin, ReduceMax, Minimum, Maximum and Sign operations, and none of those accept bool tensors. The partitioner still gave these ops to OpenVINO when they read bool tensors, so the whole program failed to lower: ReduceMin ... Element type of data input must be numeric. Got: boolean With PyTorch 2.15, MultivariateNormal exports two such mins when it checks scale_tril, so a model that builds one could not use OpenVINO. The partitioner now leaves these ops to the CPU when every tensor they read is bool, and the portable kernels compute them as PyTorch does. An op with a bool and a float input, such as maximum(mask, x), is computed in float, which OpenVINO converts, so it stays with OpenVINO. Other dtypes are not affected. --- backends/openvino/partitioner.py | 26 +++++++++++++ .../openvino/tests/ops/test_reductions.py | 37 +++++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/backends/openvino/partitioner.py b/backends/openvino/partitioner.py index 0d407e33f6e..e2ae8712929 100644 --- a/backends/openvino/partitioner.py +++ b/backends/openvino/partitioner.py @@ -33,11 +33,33 @@ def __init__(self): self.op_types = {} +def _reads_only_bool(node: torch.fx.Node) -> bool: + """Whether OpenVINO would compute the node in bool. It promotes mixed + inputs, so maximum(bool, float) is computed in float and converts fine.""" + vals = [arg.meta.get("val") for arg in node.args if isinstance(arg, torch.fx.Node)] + tensors = [val for val in vals if isinstance(val, torch.Tensor)] + return bool(tensors) and all(t.dtype == torch.bool for t in tensors) + + class OpenvinoOperatorsSupport(OperatorSupportBase): extended_support_dict = { "torch.ops.dim_order_ops._clone_dim_order.default": None, "torch.ops.dim_order_ops._to_dim_order_copy.default": None, } + # OpenVINO converts these to ReduceMin, ReduceMax, Minimum, Maximum or Sign, + # none of which take bool tensors. + bool_unsupported_ops = { + "aten.amax.default", + "aten.amin.default", + "aten.clamp.default", + "aten.max.default", + "aten.max.dim", + "aten.maximum.default", + "aten.min.default", + "aten.min.dim", + "aten.minimum.default", + "aten.sign.default", + } def __init__( self, @@ -87,6 +109,10 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: if op_type == "getitem": return True + if op_type in self.bool_unsupported_ops and _reads_only_bool(node): + print("Op not supported for bool tensors: ", "torch.ops." + str(op_type)) + return False + if "torch.ops." + str(op_type) in supported_ops: return True else: diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index 6f99859dcc8..4195921f672 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -112,6 +112,43 @@ def forward(self, input_ids, offset, mask): ) +class TestBoolMinMax(BaseOpenvinoOpTest): + """OpenVINO has no bool form of min, max, clamp or sign.""" + + def test_bool_ops_run_on_the_cpu(self): + class BoolMinMax(torch.nn.Module): + def forward(self, x): + m = x > 0 + n = m.flip(1) + return ( + m.min(1).values, + m.max(1).values, + m.amin(1), + m.amax(1), + m.min(), + m.max(), + torch.minimum(m, n), + torch.maximum(m, n), + torch.sign(m), + m.clamp(min=0.5), + ) + + # One row all true and one all false, so min and max differ by row. + x = torch.randn(3, 7) + x[0] = x[0].abs() + 0.1 + x[1] = -x[1].abs() - 0.1 + self.execute_layer_test(BoolMinMax(), (x,), expected_partitions=2) + + def test_maximum_of_bool_and_float_stays_with_openvino(self): + class Maximum(torch.nn.Module): + def forward(self, mask, x): + return torch.maximum(mask, x) + + # Computed in float, which OpenVINO converts. + x = torch.randn(2, 7) + self.execute_layer_test(Maximum(), (x > 0, x)) + + class TestFoldConstantCumsumPass(unittest.TestCase): def fold(self, module, inputs): """Edge program module of module, with constant cumsums folded.""" From 872b166f9daeea4962ef6a450052276423ce071c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:44:47 -0700 Subject: [PATCH 17/29] Say why the OpenVINO bf16 addmm check needs the runtime Lowering passes the precision too, so the comment was wrong to say only the runtime does. What makes the check depend on the runtime is that on x86 OpenVINO applies only the precision it gets when it loads the model. --- backends/openvino/tests/ops/test_addmm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/openvino/tests/ops/test_addmm.py b/backends/openvino/tests/ops/test_addmm.py index 64b685d2721..00efb18ae82 100644 --- a/backends/openvino/tests/ops/test_addmm.py +++ b/backends/openvino/tests/ops/test_addmm.py @@ -45,8 +45,8 @@ def test_addmm_computes_in_float32(self): self.atol = self.rtol = 1e-4 sample_input = tuple(torch.randn(256, 256) for _ in range(3)) self.execute_layer_test(self.create_model(), sample_input) - # On x86 only the runtime passes the precision on, so this fails if it - # drops an explicit one. + # On x86 OpenVINO applies only the precision it gets at load, so this + # fails if the runtime drops an explicit one. bf16 = CompileSpec("inference_precision", b"bf16") with self.assertRaisesRegex(AssertionError, "ref_output"): self.execute_layer_test( From 1cd364383de3b0b220edb9ab8abd6a0646d5f097 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:45:52 -0700 Subject: [PATCH 18/29] Keep the OpenVINO input types comment above the converter call The cumsum fold lines had landed between the comment and the decoder call it explains, so the comment read as if it were about the fold. --- backends/openvino/preprocess.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 86038ccd65e..52870fe36cf 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -99,10 +99,10 @@ def preprocess( f"{type(val).__name__}; only tensors and ints are supported" ) - # The converter chooses some lowerings (e.g. boolean-mask indexing) from - # the input types, so they must be known before conversion starts. module = edge_program.module() FoldConstantCumsumPass()(module) + # The converter chooses some lowerings (e.g. boolean-mask indexing) from + # the input types, so they must be known before conversion starts. decoder = TorchFXPythonDecoder( module, input_shapes=input_shapes, input_types=input_types ) From 63d01981fb3a39a18085d60e46cfbee6cc105dff Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:49:17 -0700 Subject: [PATCH 19/29] Evaluate one node at a time in the OpenVINO cumsum fold Since each cumsum is folded on its own, _evaluate is only ever asked for one node. Taking and returning that one value drops the list and dict it no longer needs, and a check that could never be false. --- .../openvino/_passes/fold_constant_cumsum_pass.py | 14 ++++++++------ backends/openvino/tests/ops/test_reductions.py | 6 +++--- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/backends/openvino/_passes/fold_constant_cumsum_pass.py b/backends/openvino/_passes/fold_constant_cumsum_pass.py index f346096c832..690b1cc28d9 100644 --- a/backends/openvino/_passes/fold_constant_cumsum_pass.py +++ b/backends/openvino/_passes/fold_constant_cumsum_pass.py @@ -64,10 +64,12 @@ def _constant_nodes(graph_module: torch.fx.GraphModule) -> set: return constants -def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict: - """Values of targets, computing only what they read, on the CPU.""" +def _evaluate( + graph_module: torch.fx.GraphModule, target: torch.fx.Node +) -> torch.Tensor: + """Value of target, computing only what it reads, on the CPU.""" needed = set() - stack = list(targets) + stack: list = [target] while stack: node = stack.pop() if node not in needed: @@ -98,9 +100,9 @@ def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict: kwargs = {**kwargs, "pin_memory": False} values[node] = node.target(*args, **kwargs) for i in _inputs(node): - if last_use[i] is node and i not in targets: + if last_use[i] is node: del values[i] - return values + return values[target] class FoldConstantCumsumPass(PassBase): @@ -133,7 +135,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # the others from being folded. for node in cumsums: try: - value = _evaluate(graph_module, [node])[node] + value = _evaluate(graph_module, node) except RuntimeError as e: # Some ops have modes the CPU lacks. logging.warning( diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index 4195921f672..82d66d51297 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -199,10 +199,10 @@ def forward(self, x): evaluate = fold_constant_cumsum_pass._evaluate - def fail_on_the_doubled_one(graph_module, targets): - if any("mul" in str(n.args[0].target) for n in targets): + def fail_on_the_doubled_one(graph_module, target): + if "mul" in str(target.args[0].target): raise RuntimeError("not on the CPU") - return evaluate(graph_module, targets) + return evaluate(graph_module, target) x = torch.zeros(3, dtype=torch.int64) with mock.patch.object( From 0dd8ef653f109601964541880b5c3c072893c1f7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:51:28 -0700 Subject: [PATCH 20/29] Test the OpenVINO cumsum fold on an int32 cumsum Only the attention mask test covered an int32 cumsum, and only through an OpenVINO CPU bug that newer OpenVINO releases fix. A fold limited to int64 now fails a pass-level test too. --- .../openvino/tests/ops/test_reductions.py | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index 82d66d51297..fbf26733a82 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -244,17 +244,23 @@ def forward(self, x): def test_cumsum_over_a_split_constant_is_folded(self): class Model(torch.nn.Module): + def __init__(self, dtype): + super().__init__() + self.dtype = dtype + def forward(self, x): part = torch.ones(6, dtype=torch.int64).split(3)[1] - return torch.cumsum(part, 0) + x + return torch.cumsum(part, 0, dtype=self.dtype) + x x = torch.zeros(3, dtype=torch.int64) - folded = self.fold(Model(), (x,)) - self.assertNotIn( - exir_ops.edge.aten.cumsum.default, - [n.target for n in folded.graph.nodes], - ) - torch.testing.assert_close(folded(x), Model()(x)) + for dtype in (torch.int64, torch.int32): + with self.subTest(dtype=dtype): + folded = self.fold(Model(dtype), (x,)) + self.assertNotIn( + exir_ops.edge.aten.cumsum.default, + [n.target for n in folded.graph.nodes], + ) + torch.testing.assert_close(folded(x), Model(dtype)(x)) def test_constants_are_folded_without_a_gpu(self): class Model(torch.nn.Module): From 06bee66fd5114b413f7eec71242be390de0d866c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 17:52:24 -0700 Subject: [PATCH 21/29] Say which x86 CPUs the OpenVINO float32 default slows down Only x86 CPUs with bfloat16 units computed in bfloat16 before, so only they get slower, and only there does asking for bf16 bring the earlier speed back. On other x86 CPUs OpenVINO already computed in float32. --- backends/openvino/README.md | 7 ++++--- docs/source/build-run-openvino.md | 7 ++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/backends/openvino/README.md b/backends/openvino/README.md index 40a245f12a0..9b27ed46424 100644 --- a/backends/openvino/README.md +++ b/backends/openvino/README.md @@ -82,9 +82,10 @@ in float32 unless you lower it with `inference_precision`. On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. So a program lowered without `inference_precision` computes in float32, even -one exported with an ExecuTorch release that did not have this default. That -is slower and can use more memory. Export it again with `inference_precision` -set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +one exported with an ExecuTorch release that did not have this default. On an +x86 CPU with bfloat16 units that is slower and can use more memory; export it +again with `inference_precision` set to `bf16` to get the earlier speed. Other +x86 CPUs already computed such a program in float32. On ARM CPUs, OpenVINO applies the precision when the program is lowered, so a program exported with such an earlier release keeps computing in float16 until it is exported again. diff --git a/docs/source/build-run-openvino.md b/docs/source/build-run-openvino.md index c5f38a2d0d7..ec75d4b96c8 100644 --- a/docs/source/build-run-openvino.md +++ b/docs/source/build-run-openvino.md @@ -136,9 +136,10 @@ in float32 unless you lower it with `inference_precision`. On x86 CPUs, OpenVINO applies the precision when the runtime loads a program. So a program lowered without `inference_precision` computes in float32, even -one exported with an ExecuTorch release that did not have this default. That -is slower and can use more memory. Export it again with `inference_precision` -set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the +one exported with an ExecuTorch release that did not have this default. On an +x86 CPU with bfloat16 units that is slower and can use more memory; export it +again with `inference_precision` set to `bf16` to get the earlier speed. Other +x86 CPUs already computed such a program in float32. On ARM CPUs, OpenVINO applies the precision when the program is lowered, so a program exported with such an earlier release keeps computing in float16 until it is exported again. From a1d3c944ade84eec3103c210705b3f8742f02bd7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 18:07:14 -0700 Subject: [PATCH 22/29] Say that OpenVINO sums int64 values in float32 on CPUs The precision docs said int64 results only go wrong outside the int32 range. Measured through the delegate on a CPU, an int64 sum is computed in float32: 16777216 + 1 gives 16777216, and 123456789 + 987654321 gives 1111111168. cumsum and prod near 2^24 stay exact. --- backends/openvino/README.md | 4 +++- docs/source/build-run-openvino.md | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/backends/openvino/README.md b/backends/openvino/README.md index 9b27ed46424..69e073baa23 100644 --- a/backends/openvino/README.md +++ b/backends/openvino/README.md @@ -92,7 +92,9 @@ earlier release keeps computing in float16 until it is exported again. OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result outside the int32 range wraps around without an error. This applies to every int64 tensor, including the int64 result of `cumsum` and `prod` over smaller -integer types. +integer types. `sum` is less exact still: it adds int64 values in float32, so +a sum above 16,777,216 (2^24) can be off even inside the int32 range, for +example 16,777,216 + 1 gives 16,777,216. ## Directory Structure diff --git a/docs/source/build-run-openvino.md b/docs/source/build-run-openvino.md index ec75d4b96c8..0a9a695cbb5 100644 --- a/docs/source/build-run-openvino.md +++ b/docs/source/build-run-openvino.md @@ -146,7 +146,9 @@ earlier release keeps computing in float16 until it is exported again. OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result outside the int32 range wraps around without an error. This applies to every int64 tensor, including the int64 result of `cumsum` and `prod` over smaller -integer types. +integer types. `sum` is less exact still: it adds int64 values in float32, so +a sum above 16,777,216 (2^24) can be off even inside the int32 range, for +example 16,777,216 + 1 gives 16,777,216. ## Build Instructions for Examples From 1623b2e938ff834adc8e6962ad153a613603d7a6 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 19:41:04 -0700 Subject: [PATCH 23/29] Say which OpenVINO outputs refuse an empty buffer Only an output compiled with a fully fixed shape refuses a buffer of another shape. A partly dynamic output accepts an empty one. Comment only. --- backends/openvino/runtime/OpenvinoBackend.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/openvino/runtime/OpenvinoBackend.cpp b/backends/openvino/runtime/OpenvinoBackend.cpp index 8004a418211..5f15aefee0c 100644 --- a/backends/openvino/runtime/OpenvinoBackend.cpp +++ b/backends/openvino/runtime/OpenvinoBackend.cpp @@ -670,8 +670,8 @@ exr::Result OpenvinoBackend::create_owned_output( return exr::Error::NotSupported; } // The request's own output has the compiled shape, empty along each dynamic - // dimension. OpenVINO refuses any other shape for a dimension it fixed, - // even when ExecuTorch planned that output with a dynamic size. + // dimension. OpenVINO refuses any other shape for an output it compiled + // with a fixed shape, even when ExecuTorch planned it with a dynamic size. ov_tensor_t* current = nullptr; ov_shape_t shape = {}; ov_status_e status = ov_.infer_request_get_output_tensor_by_index( From a19225d07a4ca6c3945ef29cdd7ad904138dcff1 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 21:21:33 -0700 Subject: [PATCH 24/29] Say the OpenVINO cumsum fold is needed from 2025.1 OpenVINO 2025.1 has the same CPU plugin bug with integer cumsums over constants as the later releases, so the fold's docstring now starts the range there. Docstring only. --- backends/openvino/_passes/fold_constant_cumsum_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/openvino/_passes/fold_constant_cumsum_pass.py b/backends/openvino/_passes/fold_constant_cumsum_pass.py index 690b1cc28d9..cae7002f05f 100644 --- a/backends/openvino/_passes/fold_constant_cumsum_pass.py +++ b/backends/openvino/_passes/fold_constant_cumsum_pass.py @@ -108,7 +108,7 @@ def _evaluate( class FoldConstantCumsumPass(PassBase): """Replace an integer cumsum that reads only constants with its result. - With an integer cumsum over a constant, OpenVINO's CPU plugin (2025.4 to + With an integer cumsum over a constant, OpenVINO's CPU plugin (2025.1 to 2026.4) gives wrong results from the second inference on: masks built from it keep only what earlier calls allowed. pi0 builds its attention masks this way. The OpenVINO nightly no longer does this, so the pass can go once From 681995ce3e764d09065fed515ba6fd90df48d653 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 20:56:41 -0700 Subject: [PATCH 25/29] Run the OpenVINO model tests even when an op test fails The OpenVINO CI job runs the op tests and then the model tests under set -e. Now that the test runner exits with code 1 when a test fails, a failing op test stopped the job before the model tests started, so a red run said nothing about the models. The script now runs both suites and exits with code 1 if either failed. --- .ci/scripts/test_openvino.sh | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/.ci/scripts/test_openvino.sh b/.ci/scripts/test_openvino.sh index 2bb2115b1ec..7f5338b3b0a 100755 --- a/.ci/scripts/test_openvino.sh +++ b/.ci/scripts/test_openvino.sh @@ -12,5 +12,8 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh" source openvino/setupvars.sh cd backends/openvino/tests -python test_runner.py --test_type ops -python test_runner.py --test_type models +# Run the model tests even when an op test fails, and still fail the job. +status=0 +python test_runner.py --test_type ops || status=1 +python test_runner.py --test_type models || status=1 +exit "${status}" From e27527cdd6b59dca05cf914cacc682dc1d4c4204 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 21:41:20 -0700 Subject: [PATCH 26/29] Check that the converted OpenVINO model takes every partition input Preprocess put the first partition input back whenever the converted model had one parameter fewer than the partition has inputs. Now it does so only when that first input is really unused, and lowering fails when the counts still differ. The runtime binds inputs by position, so another mismatch would have made every call fail, or swapped two inputs of the same type and shape. A partition can also have no inputs at all, such as a comparison that reads only a buffer next to a bool min on the CPU, and a test covers that case. --- backends/openvino/preprocess.py | 37 +++++++++++++------ .../openvino/tests/ops/test_reductions.py | 12 ++++++ 2 files changed, 38 insertions(+), 11 deletions(-) diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 52870fe36cf..e1901de3611 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -37,6 +37,31 @@ _PRECISIONS = ("f32", "f16", "bf16") +def _match_partition_inputs( + model, module: torch.fx.GraphModule, input_shapes: list, input_types: list +) -> None: + """Make the converted model take the inputs the runtime passes.""" + # The converter takes an unused first input for forward's self and removes + # it, but the runtime passes every partition input. The cumsum fold can + # leave the first input unused. + parameters = list(model.get_parameters()) + placeholders = [n for n in module.graph.nodes if n.op == "placeholder"] + if len(parameters) == len(input_types) - 1 and not placeholders[0].users: + for parameter in parameters: + model.remove_parameter(parameter) + first = opset13.parameter( + input_shapes[0], pt_to_ov_type_map[str(input_types[0])] + ) + model.add_parameters([first, *parameters]) + # The runtime binds inputs by position, so any other mismatch would make + # every call fail, or swap inputs of the same type and shape. + if len(model.get_parameters()) != len(input_types): + raise RuntimeError( + f"OpenVINO backend: expected {len(input_types)} parameters in the " + f"converted model, got {len(model.get_parameters())}" + ) + + @final class OpenvinoBackend(BackendDetails): @@ -108,17 +133,7 @@ def preprocess( ) frontend = FrontEndManager().load_by_framework("pytorch") model = frontend.convert(frontend.load(decoder)) - # The converter takes an unused first input for forward's self and - # removes it, but the runtime passes every partition input. The cumsum - # fold can leave the first input unused. - parameters = list(model.get_parameters()) - if len(parameters) == len(input_types) - 1: - for parameter in parameters: - model.remove_parameter(parameter) - first = opset13.parameter( - input_shapes[0], pt_to_ov_type_map[str(input_types[0])] - ) - model.add_parameters([first, *parameters]) + _match_partition_inputs(model, module, input_shapes, input_types) ppp = PrePostProcessor(model) # 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]. diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index fbf26733a82..f8e3b2d1c51 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -139,6 +139,18 @@ def forward(self, x): x[1] = -x[1].abs() - 0.1 self.execute_layer_test(BoolMinMax(), (x,), expected_partitions=2) + def test_bool_min_of_a_buffer(self): + class BufferMin(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.tensor([1.0, -2.0, 3.0])) + + def forward(self, x): + return x + (self.w > 0).min().float() + + # The comparison reads only the buffer, so its partition has no inputs. + self.execute_layer_test(BufferMin(), (torch.randn(3),), expected_partitions=2) + def test_maximum_of_bool_and_float_stays_with_openvino(self): class Maximum(torch.nn.Module): def forward(self, mask, x): From 79c7b97d798979e64306fb481408988070e24a67 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 21:56:08 -0700 Subject: [PATCH 27/29] Compare the indices of bool min and max in the OpenVINO test TestBoolMinMax compared only the values of min(dim) and max(dim). Their indices are where the order of ties matters, and the all-true and all-false rows are all ties, so the test now compares the indices with eager too. Its docstring also says more exactly why these ops leave OpenVINO. --- backends/openvino/tests/ops/test_reductions.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index f8e3b2d1c51..3e29eeb1d7a 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -113,7 +113,7 @@ def forward(self, input_ids, offset, mask): class TestBoolMinMax(BaseOpenvinoOpTest): - """OpenVINO has no bool form of min, max, clamp or sign.""" + """OpenVINO converts min, max, clamp and sign to ops that take no bool tensors.""" def test_bool_ops_run_on_the_cpu(self): class BoolMinMax(torch.nn.Module): @@ -121,8 +121,8 @@ def forward(self, x): m = x > 0 n = m.flip(1) return ( - m.min(1).values, - m.max(1).values, + *m.min(1), + *m.max(1), m.amin(1), m.amax(1), m.min(), @@ -133,7 +133,8 @@ def forward(self, x): m.clamp(min=0.5), ) - # One row all true and one all false, so min and max differ by row. + # One row all true and one all false, so min and max differ by row and + # their indices must pick the first of the ties. x = torch.randn(3, 7) x[0] = x[0].abs() + 0.1 x[1] = -x[1].abs() - 0.1 From a6a9ba2cc387123069bc3e885d1f61318fd0e775 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 20:56:43 -0700 Subject: [PATCH 28/29] Check the dtype of a folded int32 cumsum in the OpenVINO test x was int64 in both subtests, so the result was int64 whatever dtype the folded cumsum had, and a fold that dropped dtype=torch.int32 still passed. x now has the subtest's dtype, so such a fold fails the int32 subtest. --- backends/openvino/tests/ops/test_reductions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/backends/openvino/tests/ops/test_reductions.py b/backends/openvino/tests/ops/test_reductions.py index 3e29eeb1d7a..ef014df34e4 100644 --- a/backends/openvino/tests/ops/test_reductions.py +++ b/backends/openvino/tests/ops/test_reductions.py @@ -265,9 +265,10 @@ def forward(self, x): part = torch.ones(6, dtype=torch.int64).split(3)[1] return torch.cumsum(part, 0, dtype=self.dtype) + x - x = torch.zeros(3, dtype=torch.int64) for dtype in (torch.int64, torch.int32): with self.subTest(dtype=dtype): + # Of the same dtype, so the result keeps the folded dtype. + x = torch.zeros(3, dtype=dtype) folded = self.fold(Model(dtype), (x,)) self.assertNotIn( exir_ops.edge.aten.cumsum.default, From f2cf062fd08474b620d5e12d910f4608250fd31f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 21:56:08 -0700 Subject: [PATCH 29/29] Say which OpenVINO integer reductions go through float32 The docs said only sum reduces integers in float32. Measured through the delegate on a CPU with OpenVINO 2026.4.1 and 2025.1.0, amax, amin, max and min of int32 and int64 tensors do too: amax over a tensor that holds 16,777,217 gives 16,777,216. Inputs and running totals round as well, so the sum of 16,777,217 and -16,777,216 gives 0. cumsum and maximum stay exact, and so does prod where OpenVINO computes it (2026.4.1). --- backends/openvino/README.md | 8 +++++--- docs/source/build-run-openvino.md | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/backends/openvino/README.md b/backends/openvino/README.md index 69e073baa23..35db7602beb 100644 --- a/backends/openvino/README.md +++ b/backends/openvino/README.md @@ -92,9 +92,11 @@ earlier release keeps computing in float16 until it is exported again. OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result outside the int32 range wraps around without an error. This applies to every int64 tensor, including the int64 result of `cumsum` and `prod` over smaller -integer types. `sum` is less exact still: it adds int64 values in float32, so -a sum above 16,777,216 (2^24) can be off even inside the int32 range, for -example 16,777,216 + 1 gives 16,777,216. +integer types. `sum`, `amax`, `amin`, `max` and `min` are less exact still: +they reduce integers in float32, so a result can be off as soon as any value, +or any running total of a sum, is above 16,777,216 (2^24) or below +-16,777,216, even inside the int32 range. For example, the sum 16,777,216 + 1 +gives 16,777,216, and so does `amax` over a tensor that holds 16,777,217. ## Directory Structure diff --git a/docs/source/build-run-openvino.md b/docs/source/build-run-openvino.md index 0a9a695cbb5..20b1054157b 100644 --- a/docs/source/build-run-openvino.md +++ b/docs/source/build-run-openvino.md @@ -146,9 +146,11 @@ earlier release keeps computing in float16 until it is exported again. OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result outside the int32 range wraps around without an error. This applies to every int64 tensor, including the int64 result of `cumsum` and `prod` over smaller -integer types. `sum` is less exact still: it adds int64 values in float32, so -a sum above 16,777,216 (2^24) can be off even inside the int32 range, for -example 16,777,216 + 1 gives 16,777,216. +integer types. `sum`, `amax`, `amin`, `max` and `min` are less exact still: +they reduce integers in float32, so a result can be off as soon as any value, +or any running total of a sum, is above 16,777,216 (2^24) or below +-16,777,216, even inside the int32 range. For example, the sum 16,777,216 + 1 +gives 16,777,216, and so does `amax` over a tensor that holds 16,777,217. ## Build Instructions for Examples