From 6207bdfd38b4d98837bae054821a9a3a03543d21 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 08:54:41 -0700 Subject: [PATCH 01/20] 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 a8ccc45f6a12dab266967da2917eb00370551284 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 20:09:53 -0700 Subject: [PATCH 02/20] Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check (#24299) ## What is wrong today Robot policies and follow-up checks on #24222 found these bugs in the CUDA backend and the runtime: 1. **CUDA mask indexing:** `x[mask]` and `x[mask] = values` have no C shims, so lowering fails with "missing fallback kernels". This blocks SmolVLA. 2. **CUDA graphs:** a channels-last output turns CUDA graphs off, although a raw copy is correct. 3. **CUDA graph shapes:** when the captured call has a new batch size, it and every replay return the previous output shape. 4. **Output buffers:** `Method::set_output_data_ptr` checks a buffer against the output's current size. After a dynamic output shrinks, a too-small buffer is accepted and a later, larger call writes past it. ## What this change does One commit per fix: 1. New shims for index and index_put follow ATen's indexing rules. The Windows import library gains the new names. Inside a CUDA graph capture, a mask index fails with a clear error. 2. The check compares output strides with the strides the ETensor's dim order gives. 3. The capture step resizes each output to the captured shape. 4. The buffer is checked against the output's upper bound. Code that binds a smaller buffer after a call now gets `InvalidArgument`. ## Tests Each new test fails without its fix. - **CUDA (A100):** mask reads and writes, also inside a capture; the raw-copy check; a capture call with a new batch size. - **Runtime:** a too-small output buffer is refused after a dynamic output shrinks; the runtime tests pass. On LeRobot policies, SmolVLA now lowers to CUDA and runs in 43 ms per call on an A100, against 504 ms in eager PyTorch (cosine similarity 0.99996). (cherry picked from commit ff1bab35c7c9500113026fbb843f29b406d919f5) --- backends/aoti/common_shims_slim.cpp | 13 + backends/aoti/common_shims_slim.h | 5 + backends/cuda/CMakeLists.txt | 4 +- backends/cuda/cuda_backend.py | 2 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 43802 -> 47046 bytes backends/cuda/runtime/aoti_cuda_shims.lib.md | 2 +- backends/cuda/runtime/cuda_backend.cpp | 32 + backends/cuda/runtime/shims/index.cu | 1062 +++++++++++++++++ backends/cuda/runtime/shims/index.h | 82 ++ backends/cuda/runtime/targets.bzl | 2 + .../runtime/test/test_cuda_output_copy.cpp | 72 +- backends/cuda/runtime/utils.h | 37 + backends/cuda/tests/test_cuda_export.py | 59 + .../tests/test_output_stride_rearrange.py | 126 ++ backends/cuda/tests/test_sort_shim.py | 2 + runtime/executor/method.cpp | 7 +- runtime/executor/method.h | 4 +- runtime/executor/test/method_test.cpp | 27 + 18 files changed, 1529 insertions(+), 9 deletions(-) create mode 100644 backends/cuda/runtime/shims/index.cu create mode 100644 backends/cuda/runtime/shims/index.h diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index d9e255bf220..e92c9d319ae 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -36,6 +36,19 @@ AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) { return Error::Ok; } +AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size) { + if (tensor == nullptr || ret_size == nullptr) { + return Error::InvalidArgument; + } + const auto ndim = static_cast(tensor->dim()); + if (d < -ndim || d >= ndim) { + return Error::InvalidArgument; + } + *ret_size = tensor->size(d); + return Error::Ok; +} + AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) { if (tensor == nullptr || ret_strides == nullptr) { return Error::InvalidArgument; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index 0d60f578605..dea91ff8e67 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -39,6 +39,11 @@ aoti_torch_get_data_ptr(Tensor* tensor, void** ret_data_ptr); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes); +// Generated code reads the size of a result whose shape depends on the data +// through this, one dim at a time. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size); + AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides); diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 5f2879a7b32..fc89ab581bf 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -255,8 +255,8 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp # Only build CUDA-only operator shims when the CUDA toolchain is available. if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER) - list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu - runtime/shims/sort.cu runtime/shims/rand.cu + list(APPEND _aoti_cuda_shim_sources runtime/shims/index.cu + runtime/shims/int4mm.cu runtime/shims/sort.cu runtime/shims/rand.cu ) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 787d2502b63..8a6805596ae 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -773,6 +773,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "at::_ops::sort_stable::call": None, "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, + "aoti_torch_cuda_index_Tensor": None, + "aoti_torch_index_put_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index c0d61c611008c3897350d7b0950a8e0229607ba3..3294d1ec2629c60f221598224c005977bdd04da4 100644 GIT binary patch delta 1971 zcmZXUX-rgC6oBu2z<@9_m0=A#1}*3?@3G9l1jL0(ML=0AwGeg`RMsI{#e!HHpi-fS z;#L6(R;z76PTJ5`v6{von${($+Qg)`P1`@NZJKmRmviS5gY8Yexo^%r=iIZsclf4q z>7mkZ??|%SBAlz-lQGuBnBQ-V1%56GKE;?ItVI%;&zK-eFIjYnF+t3f#Jq(ufx1Q# zpT?NLs!5Wg855+0NYY$87)x_yZ25PR^oxuMGR7pCU5p8`D<$j5u^@kuq=;+<>nUi# zhP#r|PZ$%F4@s)(857jTOX|o*u({&~{SXbKlI8}+1P+^I8z~81zA0&+VNB35De3BD zOt7~^(oJ@PUJ6FgN4kRkS0%8TF~N{gl7H|4V}c`>CBp|96C5L_f^mvka3WQ5ih>cm zNBV+^VaZfGV}cJ8AtERU6eSNP#)3R6TFY7;Ep^p3)o?yI9PWffVYf1fjzut)?~ylqvoOX+!_Y8rvMm z4Y}MiNz-PB2_}tx*ktt4=~^{Sv&KQQF+{P(%Pm1}iL?dc9^WEYmKYD@6I6_jE99^# z*$Bg$KaR$e*^f3A!`G7M+!PZ`y87Z_E7^Z(SMlK>NwqW+wA<%nZ*n?^_fzArHkyp) zukfAQtjVw$Xv$IHo9Er0ZrXplh|SAJI6O-CgHSgwY)VPv*s!vT!?g?(bh;TaH;ppN zSv6a~=}Z$utTdvwf~3~farn1ApK}vFPA}zP%90UsAH1GXLr!y4JhHQZb6@;8vzWv8 zxtgrl+M;K`;O+_6bUL^da}P#xa^XsWCXN{aMw*WDE}ic>j+m1&qeN~LP4 z^1*<~_Q&oj z8;1?`)Z$9A9B=T!g4$$`ky~k?!eK?NA4b(#$#NL}<8ZNs8e2)8&D#uc zx>1kk8sa&;+ctY&*bW1HG0%Yen#umic9||bXctmHtaZ>BZry3YuiMDO`Y*-@GP&GnQ=E_VEq$^SmQBwUjgGo9M`KIXrm~GS4)|ao1a1uE;%E;I zlNt7V`QyCT>Eh1@Rr%U3K4b)=-V3h2;f}ZasBZe5BJUrWzJHrOZ z((ADg5~x#0Y?X2M2xC<>p%fl;>1BD<{Zl4@avex+54dQxiLQV z1q1PFt!Yu`J(2$ZD+g!nA#nZp>;xnLHw2*hZ7N#tNed456mlhz!8(H<{xTfPU98Dg zh%AeB4kx$5(Xn8hIcDQ7wm2(Xx5ql)Z^g#of*};&8eht}4K55?kiSDsOTwZPR#Ht8 Ls!MSCWTo;i#E^fN delta 1512 zcmXZbYf#ix6aeshHdqn%zrgZ{P-J0%c*Oqyte~Jg92CNMXfA@11|lO#3ZglgF!B&_ zH3;kwB!qxMF$R|Y(dAzZd@23>b3n6X2r?|cmLMjLk4*tgu+Lbz9Bpsenj*bf<9qU&bYJ`wZyh^$FhY-@(EoIz|-77*d z()2ujZStpUrjJcD<$2nSKfV=4&6Xm-xp8muv!;P5C6)4h4Kw{2wA}*)MwslhKlCtA zR<8>%FI2-DrT_}IX8;~qWl%;GM;&1r-H%KFDvoFcc!JJFuLZP5>GW3&e|tO9ie+)Z zR2!QEcq2yU>;c-85D$1YPN#1_;5E<1TIq2jmr?tSP7OLo)+AW*aiR}~d1yFQ@&Nx! zwBUmz7Z&IS9dd31e3%%Bb8G$R*VXxeSx%iQQi{0^Lq-GX=QX^UGsQv2^t_hepz?0K?WBs;!AM8{YD?QhRnLpgGf^4ZVEAaJmiWIl3<=1L*#SlYs85 zK-{(2Lb-Vy`Enyi_GJfQMu88- z0rwUdi*17qyGlbaYnd5+9lzs1sf9ns;cvw&Fn^1&I0p#eCDAlhQUDmcZE5?tts1^w z7DN|zZUpq)Zcyhbp4erVFBX>j(1~4qH~*9x6n>tg?h;=rEaws3E7R%G1V>(BZL;JxX`|eBzk9*3(YLI;=Gqe-&Lmo9@uNpzXKc%tqh=nnl!-VeZ12E z?=)W(fX%gLsypZe^4!nkxK>P`r#s+>f%U*R*os-Y|Fl8=dX3;-t15>ogzA z9Ok=t__j`0MqFS@6gZ+q(9@T40dEfIivtA*T_5Mc^j}yiuV|1T#Dq5uz0U=;jx1gD z%3+OOnBsK@M%B*@k}zU=6<&(|93Nk4#?1G0nw;jI=f^bCuEz1xzwDr@k19?714@qf Apa1{> diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib.md b/backends/cuda/runtime/aoti_cuda_shims.lib.md index af2fb65ca13..2a561c057c8 100644 --- a/backends/cuda/runtime/aoti_cuda_shims.lib.md +++ b/backends/cuda/runtime/aoti_cuda_shims.lib.md @@ -10,7 +10,7 @@ pin moves and the generated wrapper starts calling something new: ``` nm --defined-only aoti_cuda_shims.lib \ - | sed -n 's/.* T _\?\(aoti_torch_[a-z_0-9]*\)$/\1/p' \ + | sed -n 's/.* T _\?\(aoti_torch_[A-Za-z_0-9]*\)$/\1/p' \ | sort -u > exports.txt # add the new names to exports.txt, then { echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 921c15ae7af..7b89702f2b7 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -1053,6 +1053,38 @@ 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; + } + // Replay copies raw bytes and never resizes, so the ETensor takes the + // captured shape here and keeps it for every replay. + const Error resize_err = + executorch::ET_RUNTIME_NAMESPACE::resize_tensor( + args[i + n_inputs]->toTensor(), slim_outputs[i]->sizes()); + if (resize_err != Error::Ok) { + std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); + ET_LOG( + Error, + "CUDA graph: cannot resize output %zu of '%s' to its captured " + "shape; this call fails and later calls run without a graph", + i, + handle->method_name.c_str()); + return resize_err; + } + } + gerr = cudaGraphInstantiate( &handle->cuda_graph_state.graph_exec, handle->cuda_graph_state.graph, diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu new file mode 100644 index 00000000000..97335a75ea0 --- /dev/null +++ b/backends/cuda/runtime/shims/index.cu @@ -0,0 +1,1062 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +namespace c10_slim = executorch::backends::aoti::slim::c10; +using c10_slim::ScalarType; +using executorch::runtime::Error; + +namespace { + +constexpr int kMaxDims = 16; +constexpr int kThreads = 256; + +// One indexed dimension of the tensor being indexed. Its index is an int64 or +// int32 tensor, or the row-major positions of a mask's true elements, from +// which this dimension's coordinate is (position / mask_divisor) % mask_size. +struct DimIndex { + const void* data; + int64_t strides[kMaxDims]; // along the iteration dims, 0 where broadcast + int64_t size; // of the indexed dimension + int64_t stride; // of the indexed dimension + int64_t mask_divisor; // 0 for an index tensor + int64_t mask_size; + bool is_int32; +}; + +// The iteration space of an indexing call is the shape of the indexed result. +// Each point maps to an element of the indexed tensor and to one of the other +// tensor: the output of a gather, the values of a put. The dims the indices +// produce are [broadcast_start, broadcast_start + broadcast_ndim). +struct IndexPlan { + int ndim; + int64_t sizes[kMaxDims]; + int64_t indexed_strides[kMaxDims]; // 0 along the dims the indices produce + int64_t other_strides[kMaxDims]; + int broadcast_start; + int broadcast_ndim; + int64_t broadcast_numel; + int64_t after_numel; // of the dims after the ones the indices produce + bool from_mask[kMaxDims]; // size is a mask's count, unknown at export + int num_indexed; + DimIndex indexed[kMaxDims]; +}; + +// Adds the offsets of `linear` within dims [first, end). +__device__ __forceinline__ void add_offsets( + const IndexPlan& plan, + int first, + int end, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + for (int d = end - 1; d >= first; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + indexed_offset += coordinate * plan.indexed_strides[d]; + other_offset += coordinate * plan.other_strides[d]; + } +} + +// Adds the offsets of a point of the dims the indices produce, including the +// element the indices pick. +__device__ __forceinline__ void add_index_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + int64_t index_offsets[kMaxDims]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] = 0; + } + for (int d = plan.broadcast_start + plan.broadcast_ndim - 1; + d >= plan.broadcast_start; + --d) { + const int64_t coordinate = point % plan.sizes[d]; + point /= plan.sizes[d]; + other_offset += coordinate * plan.other_strides[d]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] += coordinate * plan.indexed[i].strides[d]; + } + } + for (int i = 0; i < plan.num_indexed; ++i) { + const DimIndex& index = plan.indexed[i]; + int64_t position = index.is_int32 + ? static_cast(index.data)[index_offsets[i]] + : static_cast(index.data)[index_offsets[i]]; + if (index.mask_divisor != 0) { + position = position / index.mask_divisor % index.mask_size; + } + if (position < 0) { + position += index.size; + } + // Checked on the device, as ATen and Inductor's generated code do, since + // a check on the host would have to wait for the stream. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < index.size, "index out of range"); + indexed_offset += position * index.stride; + } +} + +// Adds the offsets of a point of the other dims, the slice each index picks. +__device__ __forceinline__ void add_slice_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + const int end = plan.broadcast_start + plan.broadcast_ndim; + add_offsets( + plan, + end, + plan.ndim, + point % plan.after_numel, + indexed_offset, + other_offset); + add_offsets( + plan, + 0, + plan.broadcast_start, + point / plan.after_numel, + indexed_offset, + other_offset); +} + +__device__ __forceinline__ void locate( + const IndexPlan& plan, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + indexed_offset = 0; + other_offset = 0; + const int64_t rest = linear / plan.after_numel; + add_index_offsets( + plan, rest % plan.broadcast_numel, indexed_offset, other_offset); + add_slice_offsets( + plan, + rest / plan.broadcast_numel * plan.after_numel + + linear % plan.after_numel, + indexed_offset, + other_offset); +} + +template +__global__ void gather_kernel( + const IndexPlan plan, + int64_t numel, + const T* __restrict__ self, + T* __restrict__ out) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t self_offset; + int64_t out_offset; + locate(plan, linear, self_offset, out_offset); + out[out_offset] = self[self_offset]; + } +} + +// Adds in AccT, ATen's accumulation type, and rounds once. +template +__device__ __forceinline__ void store(T* out, AccT sum) { + *out = kAccumulate ? static_cast(static_cast(*out) + sum) + : static_cast(sum); +} + +// For indices that cannot repeat an element, which is any mix of masks and +// full slices. +template +__global__ void put_kernel( + const IndexPlan plan, + int64_t numel, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset; + int64_t values_offset; + locate(plan, linear, out_offset, values_offset); + store( + out + out_offset, static_cast(values[values_offset])); + } +} + +__global__ void index_offsets_kernel( + const IndexPlan plan, + int64_t points, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ values_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset = 0; + int64_t values_offset = 0; + add_index_offsets(plan, point, out_offset, values_offset); + out_offsets[point] = out_offset; + values_offsets[point] = values_offset; + } +} + +// After a stable sort of the index points by the element they pick, the +// points that pick one element are a run in index order. Each thread takes +// one element of the slice of one run, adds the run's values in that order +// or keeps the last, and writes once, so the result does not depend on +// scheduling. +template +__global__ void put_sorted_kernel( + const IndexPlan plan, + int64_t points, + int64_t slice_numel, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ values_offsets, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < points * slice_numel; + linear += static_cast(gridDim.x) * blockDim.x) { + const int64_t i = linear / slice_numel; + const int64_t offset = out_offsets[i]; + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + int64_t out_offset = offset; + int64_t values_slice_offset = 0; + add_slice_offsets( + plan, linear % slice_numel, out_offset, values_slice_offset); + // Starts from the element itself, so the additions associate as they do + // one after another. + AccT result = kAccumulate ? static_cast(out[out_offset]) : AccT{}; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j) { + const auto value = + static_cast(values[values_offsets[j] + values_slice_offset]); + result = kAccumulate ? static_cast(result + value) : value; + } + out[out_offset] = static_cast(result); + } +} + +// A mask's element at a row-major position is set. +struct MaskIsSet { + const uint8_t* data; + int ndim; + int64_t sizes[kMaxDims]; + int64_t strides[kMaxDims]; + + __device__ bool operator()(int64_t position) const { + int64_t offset = 0; + for (int d = ndim - 1; d >= 0; --d) { + offset += position % sizes[d] * strides[d]; + position /= sizes[d]; + } + return data[offset] != 0; + } +}; + +// Stream-ordered scratch, freed on the same stream when it goes out of scope. +class Scratch { + public: + Scratch(int32_t device, cudaStream_t stream) + : device_(device), stream_(stream) {} + + Scratch(const Scratch&) = delete; + Scratch& operator=(const Scratch&) = delete; + + ~Scratch() { + for (void* block : blocks_) { + (void)CudaAllocator::deallocate_stream_ordered(block, device_, stream_); + } + } + + void* allocate(size_t nbytes) { + auto block = CudaAllocator::allocate_stream_ordered( + std::max(nbytes, 1), device_, stream_); + if (!block.ok()) { + return nullptr; + } + blocks_.push_back(block.get()); + return block.get(); + } + + private: + int32_t device_; + cudaStream_t stream_; + std::vector blocks_; +}; + +// Writes the row-major positions of a mask's true elements to scratch and +// reads their count back, which waits for the stream. +Error mask_positions( + const Tensor& mask, + Scratch& scratch, + cudaStream_t stream, + const int64_t** positions, + int64_t* count) { + const auto numel = static_cast(mask.numel()); + *positions = nullptr; + *count = 0; + if (numel == 0) { + return Error::Ok; + } + // A graph capture cannot record the host read of the count below, so refuse + // before allocating anything. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); + ET_CHECK_OR_RETURN_ERROR( + capture == cudaStreamCaptureStatusNone, + NotSupported, + "index: a mask index cannot run inside a CUDA graph capture"); + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + numel <= INT_MAX, + NotSupported, + "index: masks with more than %d elements are not supported", + INT_MAX); + MaskIsSet is_set{static_cast(mask.data_ptr())}; + is_set.ndim = static_cast(mask.dim()); + for (int d = 0; d < is_set.ndim; ++d) { + is_set.sizes[d] = mask.size(d); + is_set.strides[d] = mask.stride(d); + } + auto* selected = + static_cast(scratch.allocate(numel * sizeof(int64_t))); + auto* selected_count = + static_cast(scratch.allocate(sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + selected != nullptr && selected_count != nullptr, + MemoryAllocationFailed, + "index: allocating the mask positions failed"); + thrust::counting_iterator first(0); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + nullptr, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index: allocating the mask scan scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + temp, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( + count, selected_count, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(stream)); + *positions = selected; + return Error::Ok; +} + +struct Advanced { + int64_t dim; // of the indexed tensor + bool from_mask; + DimIndex index; + std::vector sizes; + std::vector strides; +}; + +// Lays out an indexing call the way ATen's index and index_put do: masks +// become one index per dimension they cover, all indices broadcast together, +// and the dimensions they produce replace the indexed ones if those are +// adjacent and go first otherwise. Sets *may_repeat when two points can land +// on one element, which only an index tensor allows. +Error plan_indexing( + const Tensor& indexed, + Tensor* const* indices, + int64_t num_indices, + Scratch& scratch, + cudaStream_t stream, + IndexPlan* plan, + bool* may_repeat) { + const auto ndim = static_cast(indexed.dim()); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "index: tensors with more than %d dims are not supported", + kMaxDims); + ET_CHECK_OR_RETURN_ERROR( + num_indices <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + *may_repeat = false; + std::vector advanced; + int64_t dim = 0; + for (int64_t i = 0; i < num_indices; ++i) { + const Tensor* index = indices[i]; + if (index == nullptr) { + ++dim; + continue; + } + ET_CHECK_OR_RETURN_ERROR( + index->is_cuda(), + NotSupported, + "index: index %" PRId64 " is not on the CUDA device", + i); + const auto dtype = index->dtype(); + if (dtype == ScalarType::Bool || dtype == ScalarType::Byte) { + const auto mask_ndim = static_cast(index->dim()); + // A 0-dim mask covers no dims, so it adds no index, as in ATen. + if (mask_ndim == 0) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + dim + mask_ndim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + for (int64_t j = 0; j < mask_ndim; ++j) { + ET_CHECK_OR_RETURN_ERROR( + index->size(j) == indexed.size(dim + j), + InvalidArgument, + "index: mask %" PRId64 + " does not match the indexed shape at dim %" PRId64, + i, + dim + j); + } + const int64_t* positions = nullptr; + int64_t count = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + mask_positions(*index, scratch, stream, &positions, &count)); + // Row-major position p has coordinate p / divisor % size along dim j. + std::vector divisors(mask_ndim, 1); + for (int64_t j = mask_ndim - 1; j > 0; --j) { + divisors[j - 1] = divisors[j] * std::max(index->size(j), 1); + } + for (int64_t j = 0; j < mask_ndim; ++j) { + Advanced entry{dim + j, true}; + entry.index.data = positions; + entry.index.mask_divisor = divisors[j]; + entry.index.mask_size = index->size(j); + entry.sizes = {count}; + entry.strides = {1}; + advanced.push_back(entry); + } + dim += mask_ndim; + } else if (dtype == ScalarType::Long || dtype == ScalarType::Int) { + ET_CHECK_OR_RETURN_ERROR( + dim < ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + ET_CHECK_OR_RETURN_ERROR( + static_cast(index->dim()) <= kMaxDims, + NotSupported, + "index: index tensors with more than %d dims are not supported", + kMaxDims); + Advanced entry{dim, false}; + entry.index.data = index->data_ptr(); + entry.index.is_int32 = dtype == ScalarType::Int; + entry.sizes.assign(index->sizes().begin(), index->sizes().end()); + entry.strides.assign(index->strides().begin(), index->strides().end()); + advanced.push_back(entry); + *may_repeat = true; + ++dim; + } else { + ET_LOG( + Error, + "index: index %" PRId64 + " has dtype %d; indices must be int64, int32, bool or uint8", + i, + static_cast(dtype)); + return Error::InvalidArgument; + } + } + ET_CHECK_OR_RETURN_ERROR( + dim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + // ATen's advanced indexing rejects a call with only None and 0-dim masks; + // its plan would cover the whole tensor. + ET_CHECK_OR_RETURN_ERROR( + !advanced.empty(), + InvalidArgument, + "index: at least one index must be a tensor that is not a 0-dim mask"); + + size_t broadcast_ndim = 0; + for (const auto& entry : advanced) { + broadcast_ndim = std::max(broadcast_ndim, entry.sizes.size()); + } + std::vector broadcast(broadcast_ndim, 1); + std::vector masked(broadcast_ndim, false); + std::vector fixed(broadcast_ndim, false); + for (const auto& entry : advanced) { + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.from_mask) { + masked[lead + d] = true; + } else if (entry.sizes[d] != 1) { + fixed[lead + d] = true; + } + int64_t& size = broadcast[lead + d]; + ET_CHECK_OR_RETURN_ERROR( + entry.sizes[d] == 1 || size == 1 || entry.sizes[d] == size, + InvalidArgument, + "index: the indices cannot be broadcast together"); + if (size == 1) { + size = entry.sizes[d]; + } + } + } + + const bool adjacent = advanced.back().dim - advanced.front().dim + 1 == + static_cast(advanced.size()); + const int64_t first = advanced.front().dim; + std::vector is_indexed(ndim, false); + for (const auto& entry : advanced) { + is_indexed[entry.dim] = true; + } + std::vector plain_dims; + for (int64_t d = 0; d < ndim; ++d) { + if (!is_indexed[d]) { + plain_dims.push_back(d); + } + } + // Dims of the indexed tensor in iteration order, -1 for each broadcast dim. + std::vector order; + for (int64_t d : plain_dims) { + if (adjacent && d < first) { + order.push_back(d); + } + } + const auto broadcast_start = static_cast(order.size()); + order.insert(order.end(), broadcast_ndim, -1); + for (int64_t d : plain_dims) { + if (!adjacent || d > first) { + order.push_back(d); + } + } + ET_CHECK_OR_RETURN_ERROR( + order.size() <= static_cast(kMaxDims), + NotSupported, + "index: results with more than %d dims are not supported", + kMaxDims); + + *plan = IndexPlan{}; + plan->ndim = static_cast(order.size()); + plan->broadcast_start = static_cast(broadcast_start); + plan->broadcast_ndim = static_cast(broadcast_ndim); + plan->broadcast_numel = 1; + for (int64_t size : broadcast) { + plan->broadcast_numel *= size; + } + // As in ATen, an index into a dim of size 0 is out of range unless there are + // no index points. The result is empty, so no kernel would assert it. + for (const auto& entry : advanced) { + ET_CHECK_OR_RETURN_ERROR( + plan->broadcast_numel == 0 || indexed.size(entry.dim) > 0, + InvalidArgument, + "index: index out of range for dim %" PRId64 " of size 0", + entry.dim); + } + plan->after_numel = 1; + for (size_t k = static_cast(broadcast_start) + broadcast_ndim; + k < order.size(); + ++k) { + plan->after_numel *= indexed.size(order[k]); + } + for (size_t k = 0; k < order.size(); ++k) { + if (order[k] >= 0) { + plan->sizes[k] = indexed.size(order[k]); + plan->indexed_strides[k] = indexed.stride(order[k]); + } else { + plan->sizes[k] = broadcast[k - broadcast_start]; + plan->from_mask[k] = + masked[k - broadcast_start] && !fixed[k - broadcast_start]; + } + } + plan->num_indexed = static_cast(advanced.size()); + for (size_t i = 0; i < advanced.size(); ++i) { + const auto& entry = advanced[i]; + DimIndex& index = plan->indexed[i]; + index = entry.index; + index.size = indexed.size(entry.dim); + index.stride = indexed.stride(entry.dim); + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.sizes[d] != 1) { + index.strides[broadcast_start + lead + d] = entry.strides[d]; + } + } + } + return Error::Ok; +} + +int64_t plan_numel(const IndexPlan& plan) { + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + numel *= plan.sizes[d]; + } + return numel; +} + +unsigned int grid_size(int64_t numel) { + return static_cast( + std::min((numel + kThreads - 1) / kThreads, 65535)); +} + +// The meta function decided the layout with each mask's count unknown, so a +// count is never taken to be 1, and the result is only known to have fewer +// than 2 elements when another size is 0. +bool is_one(const IndexPlan& plan, int d) { + return !plan.from_mask[d] && plan.sizes[d] == 1; +} + +bool is_contiguous(const IndexPlan& plan, const int64_t* strides) { + bool any_unknown = false; + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + if (!plan.from_mask[d] && plan.sizes[d] == 0) { + return true; + } + any_unknown |= plan.from_mask[d]; + numel *= plan.sizes[d]; + } + if (!any_unknown && numel < 2) { + return true; + } + int64_t expected = 1; + for (int d = plan.ndim - 1; d >= 0; --d) { + if (is_one(plan, d)) { + continue; + } + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + return true; +} + +bool is_channels_last_contiguous( + const IndexPlan& plan, + const int64_t* strides) { + if (plan.ndim != 4) { + return false; + } + int64_t expected = 1; + for (int d : {1, 3, 2, 0}) { + if (!is_one(plan, d)) { + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + } + return true; +} + +// The output strides index.Tensor's meta function gives, which Inductor bakes +// into the code that reads the result: those of an elementwise output for an +// input with `strides` (compute_elementwise_output_logical_to_physical_perm in +// torch/_prims_common). +std::vector output_strides(const IndexPlan& plan, bool self_is_empty) { + const int64_t* sizes = plan.sizes; + const int64_t* strides = plan.indexed_strides; + const int ndim = plan.ndim; + std::vector perm(ndim); + std::iota(perm.begin(), perm.end(), 0); + const bool contiguous = is_contiguous(plan, strides); + const bool channels_last = is_channels_last_contiguous(plan, strides); + if (self_is_empty || ndim < 2 || (contiguous && !channels_last)) { + // Row-major. + } else if (channels_last && !contiguous) { + perm = {0, 2, 3, 1}; + } else { + auto ge = [](int64_t a, int64_t b) { + return b == 0 || (a != 0 && (a >= b || a % b == 0)); + }; + auto should_swap = [&](int a, int b) { + const int64_t stride_a = strides[a]; + const int64_t stride_b = strides[b]; + if (stride_a == 0 || stride_b == 0) { + return 0; + } + if (stride_a == stride_b) { + return ge(sizes[b], sizes[a]) ? 0 : 1; + } + if (ge(stride_b, stride_a)) { + return -1; + } + return ge(stride_a, stride_b) ? 1 : 0; + }; + std::reverse(perm.begin(), perm.end()); + for (int i = 1; i < ndim; ++i) { + int dim1 = i; + for (int dim0 = i - 1; dim0 >= 0; --dim0) { + const int comparison = should_swap(perm[dim0], perm[dim1]); + if (comparison > 0) { + std::swap(perm[dim0], perm[dim1]); + dim1 = dim0; + } else if (comparison < 0) { + break; + } + } + } + std::reverse(perm.begin(), perm.end()); + } + std::vector result(ndim); + int64_t expected = 1; + for (int k = ndim - 1; k >= 0; --k) { + result[perm[k]] = expected; + expected *= std::max(sizes[perm[k]], 1); + } + return result; +} + +template +Error gather( + const IndexPlan& plan, + int64_t numel, + const void* self, + void* out, + cudaStream_t stream) { + gather_kernel<<>>( + plan, numel, static_cast(self), static_cast(out)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +template +Error put( + const IndexPlan& plan, + int64_t numel, + bool may_repeat, + void* out_data, + const void* values_data, + Scratch& scratch, + cudaStream_t stream) { + auto* out = static_cast(out_data); + const auto* values = static_cast(values_data); + if (!may_repeat) { + put_kernel + <<>>(plan, numel, out, values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; + } + // One key per index point, so the scratch and the sort grow with the + // number of indices, not with the elements they write. + const int64_t points = plan.broadcast_numel; + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "index_put: more than %d index points are not supported", + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "index_put: allocating the offsets failed"); + int64_t* out_offsets = offsets; + int64_t* values_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_values_offsets = offsets + 3 * points; + index_offsets_kernel<<>>( + plan, points, out_offsets, values_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + const auto count = static_cast(points); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index_put: allocating the sort scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + put_sorted_kernel + <<>>( + plan, + points, + numel / points, + sorted_out_offsets, + sorted_values_offsets, + out, + values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +} // namespace + +#ifdef __cplusplus +extern "C" { +#endif + +AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && ret0 != nullptr && + (indices != nullptr || indices_len_ == 0), + InvalidArgument, + "aoti_torch_cuda_index_Tensor: null argument"); + ET_CHECK_OR_RETURN_ERROR( + self->is_cuda(), + NotSupported, + "aoti_torch_cuda_index_Tensor: self is not on the CUDA device"); + const auto device = static_cast(self->device_index()); + auto stream_result = getCurrentCUDAStream(self->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_cuda_index_Tensor: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + std::vector index_tensors(indices_len_); + for (int64_t i = 0; i < indices_len_; ++i) { + index_tensors[i] = indices[i] == nullptr ? nullptr : *indices[i]; + } + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *self, + index_tensors.data(), + indices_len_, + scratch, + stream, + &plan, + &may_repeat)); + + const std::vector sizes(plan.sizes, plan.sizes + plan.ndim); + const std::vector strides = output_strides(plan, self->numel() == 0); + Tensor* result = nullptr; + ET_CHECK_OK_OR_RETURN_ERROR(aoti_torch_empty_strided( + plan.ndim, + sizes.data(), + strides.data(), + static_cast(self->dtype()), + static_cast(c10_slim::DeviceType::CUDA), + device, + &result)); + // The caller takes ownership of *ret0 only when this call succeeds. + std::unique_ptr owner(result, [](Tensor* tensor) { + (void)aoti_torch_delete_tensor_object(tensor); + }); + std::copy(strides.begin(), strides.end(), plan.other_strides); + + const int64_t numel = plan_numel(plan); + if (numel > 0) { + const void* src = self->data_ptr(); + void* dst = result->data_ptr(); + switch (self->itemsize()) { + case 1: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 2: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 4: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 8: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + default: + ET_LOG( + Error, + "aoti_torch_cuda_index_Tensor: unsupported element size %zu", + self->itemsize()); + return Error::NotSupported; + } + } + *ret0 = owner.release(); + return Error::Ok; +} + +AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && values != nullptr && + (indices != nullptr || num_indices == 0), + InvalidArgument, + "aoti_torch_index_put_out: null argument"); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "aoti_torch_index_put_out: out must be self, as AOTInductor calls it"); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && values->is_cuda(), + NotSupported, + "aoti_torch_index_put_out: tensors must be on the CUDA device"); + ET_CHECK_OR_RETURN_ERROR( + values->dtype() == out->dtype(), + InvalidArgument, + "aoti_torch_index_put_out: values and out must have the same dtype"); + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_index_put_out: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *out, indices, num_indices, scratch, stream, &plan, &may_repeat)); + + const auto values_ndim = static_cast(values->dim()); + ET_CHECK_OR_RETURN_ERROR( + values_ndim <= plan.ndim, + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + const int lead = plan.ndim - values_ndim; + for (int d = 0; d < values_ndim; ++d) { + const int64_t size = values->size(d); + ET_CHECK_OR_RETURN_ERROR( + size == 1 || size == plan.sizes[lead + d], + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + plan.other_strides[lead + d] = size == 1 ? 0 : values->stride(d); + } + + const int64_t numel = plan_numel(plan); + if (numel == 0) { + return Error::Ok; + } + void* dst = out->data_ptr(); + const void* src = values->data_ptr(); + if (!accumulate) { + switch (out->itemsize()) { + case 1: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 2: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 4: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 8: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported element size %zu", + out->itemsize()); + return Error::NotSupported; + } + } + // The accumulation types are ATen's. + switch (out->dtype()) { + case ScalarType::Float: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Half: + return put<__half, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::BFloat16: + return put<__nv_bfloat16, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Long: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Int: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Short: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Char: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Byte: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Bool: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported dtype %d", + static_cast(out->dtype())); + return Error::NotSupported; + } +} + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h new file mode 100644 index 00000000000..3476eaa8c88 --- /dev/null +++ b/backends/cuda/runtime/shims/index.h @@ -0,0 +1,82 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::cuda { + +using executorch::backends::aoti::AOTITorchError; +using executorch::backends::aoti::Tensor; + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Advanced indexing, self[indices], with the semantics of aten::index.Tensor. + * + * Inductor calls this fallback when an index is a mask, so the number of true + * elements decides the output size. Reading that count waits for the stream, + * so a CUDA graph cannot capture this call. As in ATen, an index out of range + * triggers a device-side assert, and an index into a dimension of size 0 is + * an error unless the indices select nothing. + * + * @param self Tensor to index (CUDA device). It, each index and the result + * can have up to 16 dims. + * @param indices indices_len_ pointers to index tensors, null for a full + * slice. An index is int64 or int32, or a bool or uint8 mask that covers as + * many dimensions of self as it has. At least one entry must be an index that + * is not a 0-dim mask. + * @param indices_len_ Number of entries in indices + * @param ret0 Output: new tensor, strided the way Inductor expects + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0); + +/** + * out[indices] = values, or out[indices] += values when accumulate is set, + * with the semantics of aten::index_put_. + * + * Inductor calls this fallback, with out and self the same tensor, for an + * index that is a mask, for accumulation into a dtype its generated code + * cannot add atomically, and in deterministic mode. Writes that land on the + * same element are applied in index order, and half and bfloat16 values are + * summed in float, so the result is deterministic and small terms are not + * lost. A mask makes this call wait for the stream. Indices are checked as + * for aoti_torch_cuda_index_Tensor. + * + * @param out Tensor to write (CUDA device, up to 16 dims, as for + * aoti_torch_cuda_index_Tensor) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param indices num_indices index tensors, null for a full slice, with the + * same dtypes as for aoti_torch_cuda_index_Tensor + * @param num_indices Number of entries in indices + * @param values Values to write, broadcastable to the indexed shape, same + * dtype as out + * @param accumulate If true, add values instead of writing them + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate); + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 5ff2b15867c..d5635100b62 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(is_fbcode = False): name = "runtime_shims", srcs = [ "shims/cuda_guard.cpp", + "shims/index.cu", "shims/int4mm.cu", "shims/memory.cpp", "shims/rand.cu", @@ -43,6 +44,7 @@ def define_common_targets(is_fbcode = False): ], headers = [ "shims/cuda_guard.h", + "shims/index.h", "shims/int4mm.cuh", "shims/int4mm.h", "shims/memory.h", diff --git a/backends/cuda/runtime/test/test_cuda_output_copy.cpp b/backends/cuda/runtime/test/test_cuda_output_copy.cpp index 34f36934f73..04d55455adb 100644 --- a/backends/cuda/runtime/test/test_cuda_output_copy.cpp +++ b/backends/cuda/runtime/test/test_cuda_output_copy.cpp @@ -12,7 +12,7 @@ namespace slim = ::executorch::backends::aoti::slim; using namespace ::executorch::runtime; -using ::executorch::backends::cuda::copy_slimtensor_to_device_etensor_async; +using namespace ::executorch::backends::cuda; namespace { @@ -96,4 +96,74 @@ TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { Error::InvalidArgument); } +// Whether an output with these sizes and strides can be raw-copied into an +// ETensor that has this shape, dim order and strides from an earlier call. +bool raw_copyable( + std::initializer_list output_sizes, + std::initializer_list output_strides, + std::vector sizes, + std::vector dim_order, + std::vector strides) { + std::vector data(192); + slim::SlimTensor output = slim::from_blob( + data.data(), output_sizes, output_strides, slim::c10::ScalarType::Float); + etensor::TensorImpl impl( + etensor::ScalarType::Float, + static_cast(sizes.size()), + sizes.data(), + data.data(), + dim_order.data(), + strides.data()); + etensor::Tensor etensor(&impl); + return output_is_raw_copyable(&output, &etensor); +} + +TEST(CudaOutputRawCopyTest, AcceptsAnOutputLaidOutLikeItsETensor) { + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + // The ETensor's strides are stale from a smaller shape. + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 1, 12, 3}, {1, 3, 2, 2}, {0, 2, 3, 1}, {12, 1, 6, 3})); + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 16, 4, 1}, {1, 3, 2, 2}, {0, 1, 2, 3}, {12, 4, 2, 1})); + // A size-1 dim may have any stride, as with a batch of one. + EXPECT_TRUE(raw_copyable( + {1, 3, 4, 4}, {1, 1, 12, 3}, {1, 3, 4, 4}, {0, 2, 3, 1}, {48, 1, 12, 3})); + // An empty output has nothing to copy, whatever its strides. + EXPECT_TRUE(raw_copyable( + {2, 0, 4, 4}, + {16, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); +} + +TEST(CudaOutputRawCopyTest, RejectsAnOutputLaidOutUnlikeItsETensor) { + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + // Every row is followed by a gap, as in a slice of a wider tensor. + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {96, 32, 8, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + EXPECT_FALSE(raw_copyable( + {3, 4, 4}, {16, 4, 1}, {1, 3, 4, 4}, {0, 1, 2, 3}, {48, 16, 4, 1})); +} + } // namespace diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index 30a79500c16..ec800689804 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -269,6 +269,43 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( } } // namespace +/** + * Whether copying an output's raw bytes into its ETensor puts every element + * where the ETensor expects it, which is how a CUDA graph replay lands it. + * + * The ETensor may still have its previous shape, so its strides can be stale + * for dynamic shapes. Its expected strides come from its dim order and the + * output's sizes instead, the way resizing it computes them. Size-1 dims never + * move an element, so their strides may differ. + * + * @param output Pointer to the output SlimTensor (must not be null). + * @param etensor Pointer to the destination ETensor (must not be null). + * @return True if the raw copy is correct, including when there is nothing + * to copy. + */ +inline bool output_is_raw_copyable( + const executorch::backends::aoti::slim::SlimTensor* output, + const executorch::runtime::etensor::Tensor* etensor) { + const size_t ndim = output->dim(); + if (ndim != static_cast(etensor->dim())) { + return false; + } + if (output->numel() == 0) { + return true; + } + const auto dim_order = etensor->dim_order(); + int64_t expected_stride = 1; + for (size_t i = ndim; i > 0; --i) { + const int64_t d = dim_order[i - 1]; + const int64_t size = output->size(d); + if (size != 1 && output->stride(d) != expected_stride) { + return false; + } + expected_stride *= size; + } + return true; +} + /** * Copies data from a SlimTensor to an ETensor asynchronously. * diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8329f42e522..8f295ea10e2 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -6,6 +6,7 @@ import contextlib import os +import tempfile import threading import unittest from typing import Tuple @@ -562,6 +563,64 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) + def test_missing_scatter_kernel_fails_at_export(self): + """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" + + class Int64ScatterAdd(torch.nn.Module): + def forward(self, counts: torch.Tensor, idx: torch.Tensor): + return counts.scatter_add(0, idx, torch.ones_like(idx)) + + inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_mask_index_and_index_put(self): + """A mask index makes the result size depend on the data, and Inductor + has no atomic add for int64, so these run in the index shims.""" + from executorch.runtime import Runtime + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + model = MaskIndex() + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + edge = self._export_to_cuda_with_lower( + model, (x, masks[0], column, idx, counts) + ) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + edge.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for mask in masks: + inputs = (x, mask, column, idx, counts) + outputs = method.execute(list(inputs)) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) + 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_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 24df64d6de9..431bb56a88f 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -256,6 +256,132 @@ 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 _Double(nn.Module): + def __init__(self, memory_format): + super().__init__() + self.memory_format = memory_format + + def forward(self, x): + return (x * 2).contiguous(memory_format=self.memory_format) + + +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) + + def test_output_shape_changes_at_capture(self): + torch.manual_seed(0) + # The capture call has a larger batch than the warmup calls before it. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + batch = torch.export.Dim("batch", min=1, max=8) + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({0: batch},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + self._assert_matches_eager(module, x, out) + if __name__ == "__main__": unittest.main() diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 99cba99cfc5..d81d9e9e187 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -41,6 +41,8 @@ "at::_ops::sort_stable::call", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", } ) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 9b8b9e34008..90618c7cf93 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1494,13 +1494,14 @@ Method::set_output_data_ptr(void* buffer, size_t size, size_t output_idx) { return Error::InvalidArgument; } + // A later call can grow a dynamic output back to its upper bound. ET_CHECK_OR_RETURN_ERROR( - t.nbytes() <= size, + tensor_meta->nbytes() <= size, InvalidArgument, "buffer size: %" ET_PRIsize_t - " is smaller then expected tensor size: %" ET_PRIsize_t, + " is smaller than the largest output size: %" ET_PRIsize_t, size, - t.nbytes()); + tensor_meta->nbytes()); // Set data return internal::set_tensor_data(t, buffer, size); diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 834325f008c..501569285f9 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -154,8 +154,8 @@ class Method final { * * @param[in] buffer The block of memory to point the specified tensor at. * - * @param[in] size the length of buffer in bytes, must be >= the nbytes of the - * specified tensor. + * @param[in] size the length of buffer in bytes, must be >= the largest + * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index 0310c5d648b..ef97f8db657 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -317,6 +317,33 @@ TEST_F(MethodTest, AliasedIOTest) { } } +TEST_F(MethodTest, SetOutputDataPtrChecksUpperBound) { + // The "cat" model's output is at most (4, 4) and is not memory planned. + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + float input[4] = {}; + int32_t sizes[2] = {1, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, 2, sizes, input, dim_order, strides); + ASSERT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), Error::Ok); + float output[16]; + ASSERT_EQ(method->set_output_data_ptr(output, sizeof(output), 0), Error::Ok); + ASSERT_EQ(method->execute(), Error::Ok); + ASSERT_EQ(method->get_output(0).toTensor().nbytes(), 8 * sizeof(float)); + + // A later call can grow the output back to (4, 4), so a buffer that only + // fits the current (2, 4) output must be refused. + float small_output[8]; + EXPECT_EQ( + method->set_output_data_ptr(small_output, sizeof(small_output), 0), + Error::InvalidArgument); +} + TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input. // set_input validates numel/nbytes overflow before resize_tensor. From ee53b6fb6d993b97689a55268ddd8c9ccb761e11 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 20:09:53 -0700 Subject: [PATCH 03/20] Fix CUDA mask indexing, CUDA graph outputs and the runtime output-buffer check (#24299) ## What is wrong today Robot policies and follow-up checks on #24222 found these bugs in the CUDA backend and the runtime: 1. **CUDA mask indexing:** `x[mask]` and `x[mask] = values` have no C shims, so lowering fails with "missing fallback kernels". This blocks SmolVLA. 2. **CUDA graphs:** a channels-last output turns CUDA graphs off, although a raw copy is correct. 3. **CUDA graph shapes:** when the captured call has a new batch size, it and every replay return the previous output shape. 4. **Output buffers:** `Method::set_output_data_ptr` checks a buffer against the output's current size. After a dynamic output shrinks, a too-small buffer is accepted and a later, larger call writes past it. ## What this change does One commit per fix: 1. New shims for index and index_put follow ATen's indexing rules. The Windows import library gains the new names. Inside a CUDA graph capture, a mask index fails with a clear error. 2. The check compares output strides with the strides the ETensor's dim order gives. 3. The capture step resizes each output to the captured shape. 4. The buffer is checked against the output's upper bound. Code that binds a smaller buffer after a call now gets `InvalidArgument`. ## Tests Each new test fails without its fix. - **CUDA (A100):** mask reads and writes, also inside a capture; the raw-copy check; a capture call with a new batch size. - **Runtime:** a too-small output buffer is refused after a dynamic output shrinks; the runtime tests pass. On LeRobot policies, SmolVLA now lowers to CUDA and runs in 43 ms per call on an A100, against 504 ms in eager PyTorch (cosine similarity 0.99996). --- backends/aoti/common_shims_slim.cpp | 13 + backends/aoti/common_shims_slim.h | 5 + backends/cuda/CMakeLists.txt | 4 +- backends/cuda/cuda_backend.py | 2 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 44602 -> 47046 bytes backends/cuda/runtime/aoti_cuda_shims.lib.md | 2 +- backends/cuda/runtime/cuda_backend.cpp | 37 +- backends/cuda/runtime/shims/index.cu | 1062 +++++++++++++++++ backends/cuda/runtime/shims/index.h | 82 ++ backends/cuda/runtime/targets.bzl | 2 + .../runtime/test/test_cuda_output_copy.cpp | 72 +- backends/cuda/runtime/utils.h | 37 + backends/cuda/tests/test_cuda_export.py | 65 +- .../tests/test_output_stride_rearrange.py | 23 + backends/cuda/tests/test_sort_shim.py | 2 + runtime/executor/method.cpp | 7 +- runtime/executor/method.h | 4 +- runtime/executor/test/method_test.cpp | 27 + 18 files changed, 1403 insertions(+), 43 deletions(-) create mode 100644 backends/cuda/runtime/shims/index.cu create mode 100644 backends/cuda/runtime/shims/index.h diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index d9e255bf220..e92c9d319ae 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -36,6 +36,19 @@ AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) { return Error::Ok; } +AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size) { + if (tensor == nullptr || ret_size == nullptr) { + return Error::InvalidArgument; + } + const auto ndim = static_cast(tensor->dim()); + if (d < -ndim || d >= ndim) { + return Error::InvalidArgument; + } + *ret_size = tensor->size(d); + return Error::Ok; +} + AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) { if (tensor == nullptr || ret_strides == nullptr) { return Error::InvalidArgument; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index 0d60f578605..dea91ff8e67 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -39,6 +39,11 @@ aoti_torch_get_data_ptr(Tensor* tensor, void** ret_data_ptr); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes); +// Generated code reads the size of a result whose shape depends on the data +// through this, one dim at a time. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size); + AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides); diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 5f2879a7b32..fc89ab581bf 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -255,8 +255,8 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp # Only build CUDA-only operator shims when the CUDA toolchain is available. if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER) - list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu - runtime/shims/sort.cu runtime/shims/rand.cu + list(APPEND _aoti_cuda_shim_sources runtime/shims/index.cu + runtime/shims/int4mm.cu runtime/shims/sort.cu runtime/shims/rand.cu ) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 0bd1599cd5a..1f55a6099a9 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -839,6 +839,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "at::_ops::sort_stable::call": None, "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, + "aoti_torch_cuda_index_Tensor": None, + "aoti_torch_index_put_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 9f48cd086077a24be2d6cff506f6c665f6ba9b88..3294d1ec2629c60f221598224c005977bdd04da4 100644 GIT binary patch delta 1602 zcmZwGeNa?Y6aesZ7k5EmS=(I|5K3kk#zxq;EW3-2u%L(`65*nviGt!{m@Hq5qNtDo zQh}}<0yMvM5>YEXIn!8bIp&m2jp>ZzjAkEymYMlS+T@I#yVo3>dNaTG?z^Aoop4H<+GN-yHhOWv;0H_BT`v4&Jf-+$QKzg=ON%8_n$ek%1b=4Gl0I9QF z>E?0Ln;eYP%YCI_hk^wF(m;&j9l8S`y>m&~GXx;*=SihQ9JO>fOBvx{q+{G)dVh~{ zasz;L%1L)cb(nsIuccxhz8F6q_vVC0hTeSsrn*LK@H~fGbCNOAV8Ju`DD=-zWwKj;9jd< zIB>ctLFdYl58=XGQykr}=1rom$zFl=9_!PQTuwTR!4eH073#y)wEq_{oyhr6V5>fg zf*y~c_ht+V88_8W&lc31KO~SNr;?c-f&K!G60=JM-koj6i3Krae}QWk%+bg+#I@Ip z%vf$1L*~Mn0v*M6`ec(=2rF(kB+?&6(*&C4+NpCg_w4h=;hho#-cP-a=St?{sCgRY z%;koz`Bp3nH{(t1S8j;L8>KOLt~4KSFL=C$f1wG_EU?h<{5*l5FG{C`GXC&_#qrd= zkYhPeuF=PvN`;8S@52+RdU2M(H5D4w9puR$R9ev<5ka|2T>`V0+v!Lx*M3@Mp{QjU z0>}6`h$gP-tG3|RDV&}8Xmi@Hl1>PyzWE)IGVZIEpSG? zoknW8=0byoOm%X$CZ5g5ktE@^(XsShgG1nt&33iIinZ}{ERZSis}`HeUD0|Y?r0fH zU2F3>{jE0ovxOt^YNeB0PMe@hZ=j&TZl|E_J$?LNdkj?dmQ7&U zHk-~lg+A*^6d3Gu;`N{zui$w{ET~J*4_&Cp=lsc~gCZ_3s|(ar1nL^AUa45=3*a~X z33#tR2{Zd_bgqXN{|L>Lyp^+He*5t*>L0M67Hy;l!5q%bpe^+0bT+`H(Ft@7Ir1-d zsw86Jt{B`k7)8eiG6bI7tu}+Oe#_4Yu#s~`#X7=hx?7VX!kg(-8D<#^n)7Z zHhD$7eVEP(XgWBSECbh++t@NWpn delta 1260 zcmXZZZ%kWN6aesZD-=Te3MnvIWsxk~WLfrJ!43vO!6CS%(6NCmAPWmJnjv%GM#O-0 zAb-k)jr433Xb>38HNcqTs2_+XGnz#cm?nNOCX2y8@xwwkqJA+`EW~KfAor(j1CV-@MfTtG~M-V{TQKO{g0!TamRCX@`Na+*GzDfW| z&r~w@07$MBB|8B?%8gR;Hv9m28vq5X${{W+6%Q$gn*k&bCzZTT<(LgXsx&E84*{g= zFO}M90IBYRQlAGPo#r!=?}5_fpXW=nv+YVtA%JwAH%V>lO8Zv;QpbwYbqPS~Zd7{t zj5NRvNrSvs8oI1p;pU~FOS#5Vka9ym0!S10l*tbPq-kzQnsqB5*a4)-rm_&hk!KvV z(^UEtEkzHQ@%w8#=%Upt@O~tY5^Oobl$AlAc(0lMNG?jJ!qMXb0|iOA6m6ycLvBI8 z7KQ}AQ(~dU;(S406paZydBj2`hk0$TgxB&*Ewsnu7T9~lpt2V#1e!f|OgoxDvEG*i zHk2Cl>^Mh%E3?yj`O5-5M>FW#+g#^xymorSt;xs{)5oRl`sB^qiSSRNZi zwRL#{KRs#ClSSU^uCq~keUZS>DJM;}a>P-eM5fc30-IkmC_KlJM-7SewJ%?wtC3HI zIr51wk-|+mg3M>e1T~*MAn3c-#|7oKWC^_fhLfgSx#F&tZ4`UnDX^v0Nq>hq^81^K zwAOY|VD<&4%DvL2;h9d2uD53kjCo5}V+*u9usJT4Iy;I5{_wU==Q6o!O^1dz;!?=h z$d)1aUH6+-Mnqh~vLyAKE?zpe)C>bB$NfSFQz`DX9-7<9Rht4H;wV2dT1{^{ol z)b;7)3vkx;emkuW>H-f240=4wk(EI^P7OyB4(WoPTnY-RxZ)D@`EZbrSYnn2gR)pR8%-^6 zTg`Jk>P5cnqZ=B%y2RPsHx2scHfOsV(QyB=9cOochX3BVh~{~NZY}ddU_m48qrC!G L7j exports.txt # add the new names to exports.txt, then { echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 712ed5ff545..ce429a9a65a 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -87,28 +87,6 @@ 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() { @@ -1090,6 +1068,21 @@ class ET_EXPERIMENTAL CudaBackend final handle->method_name.c_str()); return Error::NotSupported; } + // Replay copies raw bytes and never resizes, so the ETensor takes the + // captured shape here and keeps it for every replay. + const Error resize_err = + executorch::ET_RUNTIME_NAMESPACE::resize_tensor( + args[i + n_inputs]->toTensor(), slim_outputs[i]->sizes()); + if (resize_err != Error::Ok) { + std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); + ET_LOG( + Error, + "CUDA graph: cannot resize output %zu of '%s' to its captured " + "shape; this call fails and later calls run without a graph", + i, + handle->method_name.c_str()); + return resize_err; + } } gerr = cudaGraphInstantiate( diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu new file mode 100644 index 00000000000..97335a75ea0 --- /dev/null +++ b/backends/cuda/runtime/shims/index.cu @@ -0,0 +1,1062 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +namespace c10_slim = executorch::backends::aoti::slim::c10; +using c10_slim::ScalarType; +using executorch::runtime::Error; + +namespace { + +constexpr int kMaxDims = 16; +constexpr int kThreads = 256; + +// One indexed dimension of the tensor being indexed. Its index is an int64 or +// int32 tensor, or the row-major positions of a mask's true elements, from +// which this dimension's coordinate is (position / mask_divisor) % mask_size. +struct DimIndex { + const void* data; + int64_t strides[kMaxDims]; // along the iteration dims, 0 where broadcast + int64_t size; // of the indexed dimension + int64_t stride; // of the indexed dimension + int64_t mask_divisor; // 0 for an index tensor + int64_t mask_size; + bool is_int32; +}; + +// The iteration space of an indexing call is the shape of the indexed result. +// Each point maps to an element of the indexed tensor and to one of the other +// tensor: the output of a gather, the values of a put. The dims the indices +// produce are [broadcast_start, broadcast_start + broadcast_ndim). +struct IndexPlan { + int ndim; + int64_t sizes[kMaxDims]; + int64_t indexed_strides[kMaxDims]; // 0 along the dims the indices produce + int64_t other_strides[kMaxDims]; + int broadcast_start; + int broadcast_ndim; + int64_t broadcast_numel; + int64_t after_numel; // of the dims after the ones the indices produce + bool from_mask[kMaxDims]; // size is a mask's count, unknown at export + int num_indexed; + DimIndex indexed[kMaxDims]; +}; + +// Adds the offsets of `linear` within dims [first, end). +__device__ __forceinline__ void add_offsets( + const IndexPlan& plan, + int first, + int end, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + for (int d = end - 1; d >= first; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + indexed_offset += coordinate * plan.indexed_strides[d]; + other_offset += coordinate * plan.other_strides[d]; + } +} + +// Adds the offsets of a point of the dims the indices produce, including the +// element the indices pick. +__device__ __forceinline__ void add_index_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + int64_t index_offsets[kMaxDims]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] = 0; + } + for (int d = plan.broadcast_start + plan.broadcast_ndim - 1; + d >= plan.broadcast_start; + --d) { + const int64_t coordinate = point % plan.sizes[d]; + point /= plan.sizes[d]; + other_offset += coordinate * plan.other_strides[d]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] += coordinate * plan.indexed[i].strides[d]; + } + } + for (int i = 0; i < plan.num_indexed; ++i) { + const DimIndex& index = plan.indexed[i]; + int64_t position = index.is_int32 + ? static_cast(index.data)[index_offsets[i]] + : static_cast(index.data)[index_offsets[i]]; + if (index.mask_divisor != 0) { + position = position / index.mask_divisor % index.mask_size; + } + if (position < 0) { + position += index.size; + } + // Checked on the device, as ATen and Inductor's generated code do, since + // a check on the host would have to wait for the stream. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < index.size, "index out of range"); + indexed_offset += position * index.stride; + } +} + +// Adds the offsets of a point of the other dims, the slice each index picks. +__device__ __forceinline__ void add_slice_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + const int end = plan.broadcast_start + plan.broadcast_ndim; + add_offsets( + plan, + end, + plan.ndim, + point % plan.after_numel, + indexed_offset, + other_offset); + add_offsets( + plan, + 0, + plan.broadcast_start, + point / plan.after_numel, + indexed_offset, + other_offset); +} + +__device__ __forceinline__ void locate( + const IndexPlan& plan, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + indexed_offset = 0; + other_offset = 0; + const int64_t rest = linear / plan.after_numel; + add_index_offsets( + plan, rest % plan.broadcast_numel, indexed_offset, other_offset); + add_slice_offsets( + plan, + rest / plan.broadcast_numel * plan.after_numel + + linear % plan.after_numel, + indexed_offset, + other_offset); +} + +template +__global__ void gather_kernel( + const IndexPlan plan, + int64_t numel, + const T* __restrict__ self, + T* __restrict__ out) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t self_offset; + int64_t out_offset; + locate(plan, linear, self_offset, out_offset); + out[out_offset] = self[self_offset]; + } +} + +// Adds in AccT, ATen's accumulation type, and rounds once. +template +__device__ __forceinline__ void store(T* out, AccT sum) { + *out = kAccumulate ? static_cast(static_cast(*out) + sum) + : static_cast(sum); +} + +// For indices that cannot repeat an element, which is any mix of masks and +// full slices. +template +__global__ void put_kernel( + const IndexPlan plan, + int64_t numel, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset; + int64_t values_offset; + locate(plan, linear, out_offset, values_offset); + store( + out + out_offset, static_cast(values[values_offset])); + } +} + +__global__ void index_offsets_kernel( + const IndexPlan plan, + int64_t points, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ values_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset = 0; + int64_t values_offset = 0; + add_index_offsets(plan, point, out_offset, values_offset); + out_offsets[point] = out_offset; + values_offsets[point] = values_offset; + } +} + +// After a stable sort of the index points by the element they pick, the +// points that pick one element are a run in index order. Each thread takes +// one element of the slice of one run, adds the run's values in that order +// or keeps the last, and writes once, so the result does not depend on +// scheduling. +template +__global__ void put_sorted_kernel( + const IndexPlan plan, + int64_t points, + int64_t slice_numel, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ values_offsets, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < points * slice_numel; + linear += static_cast(gridDim.x) * blockDim.x) { + const int64_t i = linear / slice_numel; + const int64_t offset = out_offsets[i]; + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + int64_t out_offset = offset; + int64_t values_slice_offset = 0; + add_slice_offsets( + plan, linear % slice_numel, out_offset, values_slice_offset); + // Starts from the element itself, so the additions associate as they do + // one after another. + AccT result = kAccumulate ? static_cast(out[out_offset]) : AccT{}; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j) { + const auto value = + static_cast(values[values_offsets[j] + values_slice_offset]); + result = kAccumulate ? static_cast(result + value) : value; + } + out[out_offset] = static_cast(result); + } +} + +// A mask's element at a row-major position is set. +struct MaskIsSet { + const uint8_t* data; + int ndim; + int64_t sizes[kMaxDims]; + int64_t strides[kMaxDims]; + + __device__ bool operator()(int64_t position) const { + int64_t offset = 0; + for (int d = ndim - 1; d >= 0; --d) { + offset += position % sizes[d] * strides[d]; + position /= sizes[d]; + } + return data[offset] != 0; + } +}; + +// Stream-ordered scratch, freed on the same stream when it goes out of scope. +class Scratch { + public: + Scratch(int32_t device, cudaStream_t stream) + : device_(device), stream_(stream) {} + + Scratch(const Scratch&) = delete; + Scratch& operator=(const Scratch&) = delete; + + ~Scratch() { + for (void* block : blocks_) { + (void)CudaAllocator::deallocate_stream_ordered(block, device_, stream_); + } + } + + void* allocate(size_t nbytes) { + auto block = CudaAllocator::allocate_stream_ordered( + std::max(nbytes, 1), device_, stream_); + if (!block.ok()) { + return nullptr; + } + blocks_.push_back(block.get()); + return block.get(); + } + + private: + int32_t device_; + cudaStream_t stream_; + std::vector blocks_; +}; + +// Writes the row-major positions of a mask's true elements to scratch and +// reads their count back, which waits for the stream. +Error mask_positions( + const Tensor& mask, + Scratch& scratch, + cudaStream_t stream, + const int64_t** positions, + int64_t* count) { + const auto numel = static_cast(mask.numel()); + *positions = nullptr; + *count = 0; + if (numel == 0) { + return Error::Ok; + } + // A graph capture cannot record the host read of the count below, so refuse + // before allocating anything. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); + ET_CHECK_OR_RETURN_ERROR( + capture == cudaStreamCaptureStatusNone, + NotSupported, + "index: a mask index cannot run inside a CUDA graph capture"); + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + numel <= INT_MAX, + NotSupported, + "index: masks with more than %d elements are not supported", + INT_MAX); + MaskIsSet is_set{static_cast(mask.data_ptr())}; + is_set.ndim = static_cast(mask.dim()); + for (int d = 0; d < is_set.ndim; ++d) { + is_set.sizes[d] = mask.size(d); + is_set.strides[d] = mask.stride(d); + } + auto* selected = + static_cast(scratch.allocate(numel * sizeof(int64_t))); + auto* selected_count = + static_cast(scratch.allocate(sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + selected != nullptr && selected_count != nullptr, + MemoryAllocationFailed, + "index: allocating the mask positions failed"); + thrust::counting_iterator first(0); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + nullptr, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index: allocating the mask scan scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + temp, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( + count, selected_count, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(stream)); + *positions = selected; + return Error::Ok; +} + +struct Advanced { + int64_t dim; // of the indexed tensor + bool from_mask; + DimIndex index; + std::vector sizes; + std::vector strides; +}; + +// Lays out an indexing call the way ATen's index and index_put do: masks +// become one index per dimension they cover, all indices broadcast together, +// and the dimensions they produce replace the indexed ones if those are +// adjacent and go first otherwise. Sets *may_repeat when two points can land +// on one element, which only an index tensor allows. +Error plan_indexing( + const Tensor& indexed, + Tensor* const* indices, + int64_t num_indices, + Scratch& scratch, + cudaStream_t stream, + IndexPlan* plan, + bool* may_repeat) { + const auto ndim = static_cast(indexed.dim()); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "index: tensors with more than %d dims are not supported", + kMaxDims); + ET_CHECK_OR_RETURN_ERROR( + num_indices <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + *may_repeat = false; + std::vector advanced; + int64_t dim = 0; + for (int64_t i = 0; i < num_indices; ++i) { + const Tensor* index = indices[i]; + if (index == nullptr) { + ++dim; + continue; + } + ET_CHECK_OR_RETURN_ERROR( + index->is_cuda(), + NotSupported, + "index: index %" PRId64 " is not on the CUDA device", + i); + const auto dtype = index->dtype(); + if (dtype == ScalarType::Bool || dtype == ScalarType::Byte) { + const auto mask_ndim = static_cast(index->dim()); + // A 0-dim mask covers no dims, so it adds no index, as in ATen. + if (mask_ndim == 0) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + dim + mask_ndim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + for (int64_t j = 0; j < mask_ndim; ++j) { + ET_CHECK_OR_RETURN_ERROR( + index->size(j) == indexed.size(dim + j), + InvalidArgument, + "index: mask %" PRId64 + " does not match the indexed shape at dim %" PRId64, + i, + dim + j); + } + const int64_t* positions = nullptr; + int64_t count = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + mask_positions(*index, scratch, stream, &positions, &count)); + // Row-major position p has coordinate p / divisor % size along dim j. + std::vector divisors(mask_ndim, 1); + for (int64_t j = mask_ndim - 1; j > 0; --j) { + divisors[j - 1] = divisors[j] * std::max(index->size(j), 1); + } + for (int64_t j = 0; j < mask_ndim; ++j) { + Advanced entry{dim + j, true}; + entry.index.data = positions; + entry.index.mask_divisor = divisors[j]; + entry.index.mask_size = index->size(j); + entry.sizes = {count}; + entry.strides = {1}; + advanced.push_back(entry); + } + dim += mask_ndim; + } else if (dtype == ScalarType::Long || dtype == ScalarType::Int) { + ET_CHECK_OR_RETURN_ERROR( + dim < ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + ET_CHECK_OR_RETURN_ERROR( + static_cast(index->dim()) <= kMaxDims, + NotSupported, + "index: index tensors with more than %d dims are not supported", + kMaxDims); + Advanced entry{dim, false}; + entry.index.data = index->data_ptr(); + entry.index.is_int32 = dtype == ScalarType::Int; + entry.sizes.assign(index->sizes().begin(), index->sizes().end()); + entry.strides.assign(index->strides().begin(), index->strides().end()); + advanced.push_back(entry); + *may_repeat = true; + ++dim; + } else { + ET_LOG( + Error, + "index: index %" PRId64 + " has dtype %d; indices must be int64, int32, bool or uint8", + i, + static_cast(dtype)); + return Error::InvalidArgument; + } + } + ET_CHECK_OR_RETURN_ERROR( + dim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + // ATen's advanced indexing rejects a call with only None and 0-dim masks; + // its plan would cover the whole tensor. + ET_CHECK_OR_RETURN_ERROR( + !advanced.empty(), + InvalidArgument, + "index: at least one index must be a tensor that is not a 0-dim mask"); + + size_t broadcast_ndim = 0; + for (const auto& entry : advanced) { + broadcast_ndim = std::max(broadcast_ndim, entry.sizes.size()); + } + std::vector broadcast(broadcast_ndim, 1); + std::vector masked(broadcast_ndim, false); + std::vector fixed(broadcast_ndim, false); + for (const auto& entry : advanced) { + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.from_mask) { + masked[lead + d] = true; + } else if (entry.sizes[d] != 1) { + fixed[lead + d] = true; + } + int64_t& size = broadcast[lead + d]; + ET_CHECK_OR_RETURN_ERROR( + entry.sizes[d] == 1 || size == 1 || entry.sizes[d] == size, + InvalidArgument, + "index: the indices cannot be broadcast together"); + if (size == 1) { + size = entry.sizes[d]; + } + } + } + + const bool adjacent = advanced.back().dim - advanced.front().dim + 1 == + static_cast(advanced.size()); + const int64_t first = advanced.front().dim; + std::vector is_indexed(ndim, false); + for (const auto& entry : advanced) { + is_indexed[entry.dim] = true; + } + std::vector plain_dims; + for (int64_t d = 0; d < ndim; ++d) { + if (!is_indexed[d]) { + plain_dims.push_back(d); + } + } + // Dims of the indexed tensor in iteration order, -1 for each broadcast dim. + std::vector order; + for (int64_t d : plain_dims) { + if (adjacent && d < first) { + order.push_back(d); + } + } + const auto broadcast_start = static_cast(order.size()); + order.insert(order.end(), broadcast_ndim, -1); + for (int64_t d : plain_dims) { + if (!adjacent || d > first) { + order.push_back(d); + } + } + ET_CHECK_OR_RETURN_ERROR( + order.size() <= static_cast(kMaxDims), + NotSupported, + "index: results with more than %d dims are not supported", + kMaxDims); + + *plan = IndexPlan{}; + plan->ndim = static_cast(order.size()); + plan->broadcast_start = static_cast(broadcast_start); + plan->broadcast_ndim = static_cast(broadcast_ndim); + plan->broadcast_numel = 1; + for (int64_t size : broadcast) { + plan->broadcast_numel *= size; + } + // As in ATen, an index into a dim of size 0 is out of range unless there are + // no index points. The result is empty, so no kernel would assert it. + for (const auto& entry : advanced) { + ET_CHECK_OR_RETURN_ERROR( + plan->broadcast_numel == 0 || indexed.size(entry.dim) > 0, + InvalidArgument, + "index: index out of range for dim %" PRId64 " of size 0", + entry.dim); + } + plan->after_numel = 1; + for (size_t k = static_cast(broadcast_start) + broadcast_ndim; + k < order.size(); + ++k) { + plan->after_numel *= indexed.size(order[k]); + } + for (size_t k = 0; k < order.size(); ++k) { + if (order[k] >= 0) { + plan->sizes[k] = indexed.size(order[k]); + plan->indexed_strides[k] = indexed.stride(order[k]); + } else { + plan->sizes[k] = broadcast[k - broadcast_start]; + plan->from_mask[k] = + masked[k - broadcast_start] && !fixed[k - broadcast_start]; + } + } + plan->num_indexed = static_cast(advanced.size()); + for (size_t i = 0; i < advanced.size(); ++i) { + const auto& entry = advanced[i]; + DimIndex& index = plan->indexed[i]; + index = entry.index; + index.size = indexed.size(entry.dim); + index.stride = indexed.stride(entry.dim); + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.sizes[d] != 1) { + index.strides[broadcast_start + lead + d] = entry.strides[d]; + } + } + } + return Error::Ok; +} + +int64_t plan_numel(const IndexPlan& plan) { + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + numel *= plan.sizes[d]; + } + return numel; +} + +unsigned int grid_size(int64_t numel) { + return static_cast( + std::min((numel + kThreads - 1) / kThreads, 65535)); +} + +// The meta function decided the layout with each mask's count unknown, so a +// count is never taken to be 1, and the result is only known to have fewer +// than 2 elements when another size is 0. +bool is_one(const IndexPlan& plan, int d) { + return !plan.from_mask[d] && plan.sizes[d] == 1; +} + +bool is_contiguous(const IndexPlan& plan, const int64_t* strides) { + bool any_unknown = false; + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + if (!plan.from_mask[d] && plan.sizes[d] == 0) { + return true; + } + any_unknown |= plan.from_mask[d]; + numel *= plan.sizes[d]; + } + if (!any_unknown && numel < 2) { + return true; + } + int64_t expected = 1; + for (int d = plan.ndim - 1; d >= 0; --d) { + if (is_one(plan, d)) { + continue; + } + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + return true; +} + +bool is_channels_last_contiguous( + const IndexPlan& plan, + const int64_t* strides) { + if (plan.ndim != 4) { + return false; + } + int64_t expected = 1; + for (int d : {1, 3, 2, 0}) { + if (!is_one(plan, d)) { + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + } + return true; +} + +// The output strides index.Tensor's meta function gives, which Inductor bakes +// into the code that reads the result: those of an elementwise output for an +// input with `strides` (compute_elementwise_output_logical_to_physical_perm in +// torch/_prims_common). +std::vector output_strides(const IndexPlan& plan, bool self_is_empty) { + const int64_t* sizes = plan.sizes; + const int64_t* strides = plan.indexed_strides; + const int ndim = plan.ndim; + std::vector perm(ndim); + std::iota(perm.begin(), perm.end(), 0); + const bool contiguous = is_contiguous(plan, strides); + const bool channels_last = is_channels_last_contiguous(plan, strides); + if (self_is_empty || ndim < 2 || (contiguous && !channels_last)) { + // Row-major. + } else if (channels_last && !contiguous) { + perm = {0, 2, 3, 1}; + } else { + auto ge = [](int64_t a, int64_t b) { + return b == 0 || (a != 0 && (a >= b || a % b == 0)); + }; + auto should_swap = [&](int a, int b) { + const int64_t stride_a = strides[a]; + const int64_t stride_b = strides[b]; + if (stride_a == 0 || stride_b == 0) { + return 0; + } + if (stride_a == stride_b) { + return ge(sizes[b], sizes[a]) ? 0 : 1; + } + if (ge(stride_b, stride_a)) { + return -1; + } + return ge(stride_a, stride_b) ? 1 : 0; + }; + std::reverse(perm.begin(), perm.end()); + for (int i = 1; i < ndim; ++i) { + int dim1 = i; + for (int dim0 = i - 1; dim0 >= 0; --dim0) { + const int comparison = should_swap(perm[dim0], perm[dim1]); + if (comparison > 0) { + std::swap(perm[dim0], perm[dim1]); + dim1 = dim0; + } else if (comparison < 0) { + break; + } + } + } + std::reverse(perm.begin(), perm.end()); + } + std::vector result(ndim); + int64_t expected = 1; + for (int k = ndim - 1; k >= 0; --k) { + result[perm[k]] = expected; + expected *= std::max(sizes[perm[k]], 1); + } + return result; +} + +template +Error gather( + const IndexPlan& plan, + int64_t numel, + const void* self, + void* out, + cudaStream_t stream) { + gather_kernel<<>>( + plan, numel, static_cast(self), static_cast(out)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +template +Error put( + const IndexPlan& plan, + int64_t numel, + bool may_repeat, + void* out_data, + const void* values_data, + Scratch& scratch, + cudaStream_t stream) { + auto* out = static_cast(out_data); + const auto* values = static_cast(values_data); + if (!may_repeat) { + put_kernel + <<>>(plan, numel, out, values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; + } + // One key per index point, so the scratch and the sort grow with the + // number of indices, not with the elements they write. + const int64_t points = plan.broadcast_numel; + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "index_put: more than %d index points are not supported", + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "index_put: allocating the offsets failed"); + int64_t* out_offsets = offsets; + int64_t* values_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_values_offsets = offsets + 3 * points; + index_offsets_kernel<<>>( + plan, points, out_offsets, values_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + const auto count = static_cast(points); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index_put: allocating the sort scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + put_sorted_kernel + <<>>( + plan, + points, + numel / points, + sorted_out_offsets, + sorted_values_offsets, + out, + values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +} // namespace + +#ifdef __cplusplus +extern "C" { +#endif + +AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && ret0 != nullptr && + (indices != nullptr || indices_len_ == 0), + InvalidArgument, + "aoti_torch_cuda_index_Tensor: null argument"); + ET_CHECK_OR_RETURN_ERROR( + self->is_cuda(), + NotSupported, + "aoti_torch_cuda_index_Tensor: self is not on the CUDA device"); + const auto device = static_cast(self->device_index()); + auto stream_result = getCurrentCUDAStream(self->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_cuda_index_Tensor: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + std::vector index_tensors(indices_len_); + for (int64_t i = 0; i < indices_len_; ++i) { + index_tensors[i] = indices[i] == nullptr ? nullptr : *indices[i]; + } + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *self, + index_tensors.data(), + indices_len_, + scratch, + stream, + &plan, + &may_repeat)); + + const std::vector sizes(plan.sizes, plan.sizes + plan.ndim); + const std::vector strides = output_strides(plan, self->numel() == 0); + Tensor* result = nullptr; + ET_CHECK_OK_OR_RETURN_ERROR(aoti_torch_empty_strided( + plan.ndim, + sizes.data(), + strides.data(), + static_cast(self->dtype()), + static_cast(c10_slim::DeviceType::CUDA), + device, + &result)); + // The caller takes ownership of *ret0 only when this call succeeds. + std::unique_ptr owner(result, [](Tensor* tensor) { + (void)aoti_torch_delete_tensor_object(tensor); + }); + std::copy(strides.begin(), strides.end(), plan.other_strides); + + const int64_t numel = plan_numel(plan); + if (numel > 0) { + const void* src = self->data_ptr(); + void* dst = result->data_ptr(); + switch (self->itemsize()) { + case 1: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 2: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 4: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 8: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + default: + ET_LOG( + Error, + "aoti_torch_cuda_index_Tensor: unsupported element size %zu", + self->itemsize()); + return Error::NotSupported; + } + } + *ret0 = owner.release(); + return Error::Ok; +} + +AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && values != nullptr && + (indices != nullptr || num_indices == 0), + InvalidArgument, + "aoti_torch_index_put_out: null argument"); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "aoti_torch_index_put_out: out must be self, as AOTInductor calls it"); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && values->is_cuda(), + NotSupported, + "aoti_torch_index_put_out: tensors must be on the CUDA device"); + ET_CHECK_OR_RETURN_ERROR( + values->dtype() == out->dtype(), + InvalidArgument, + "aoti_torch_index_put_out: values and out must have the same dtype"); + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_index_put_out: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *out, indices, num_indices, scratch, stream, &plan, &may_repeat)); + + const auto values_ndim = static_cast(values->dim()); + ET_CHECK_OR_RETURN_ERROR( + values_ndim <= plan.ndim, + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + const int lead = plan.ndim - values_ndim; + for (int d = 0; d < values_ndim; ++d) { + const int64_t size = values->size(d); + ET_CHECK_OR_RETURN_ERROR( + size == 1 || size == plan.sizes[lead + d], + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + plan.other_strides[lead + d] = size == 1 ? 0 : values->stride(d); + } + + const int64_t numel = plan_numel(plan); + if (numel == 0) { + return Error::Ok; + } + void* dst = out->data_ptr(); + const void* src = values->data_ptr(); + if (!accumulate) { + switch (out->itemsize()) { + case 1: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 2: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 4: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 8: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported element size %zu", + out->itemsize()); + return Error::NotSupported; + } + } + // The accumulation types are ATen's. + switch (out->dtype()) { + case ScalarType::Float: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Half: + return put<__half, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::BFloat16: + return put<__nv_bfloat16, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Long: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Int: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Short: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Char: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Byte: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Bool: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported dtype %d", + static_cast(out->dtype())); + return Error::NotSupported; + } +} + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h new file mode 100644 index 00000000000..3476eaa8c88 --- /dev/null +++ b/backends/cuda/runtime/shims/index.h @@ -0,0 +1,82 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::cuda { + +using executorch::backends::aoti::AOTITorchError; +using executorch::backends::aoti::Tensor; + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Advanced indexing, self[indices], with the semantics of aten::index.Tensor. + * + * Inductor calls this fallback when an index is a mask, so the number of true + * elements decides the output size. Reading that count waits for the stream, + * so a CUDA graph cannot capture this call. As in ATen, an index out of range + * triggers a device-side assert, and an index into a dimension of size 0 is + * an error unless the indices select nothing. + * + * @param self Tensor to index (CUDA device). It, each index and the result + * can have up to 16 dims. + * @param indices indices_len_ pointers to index tensors, null for a full + * slice. An index is int64 or int32, or a bool or uint8 mask that covers as + * many dimensions of self as it has. At least one entry must be an index that + * is not a 0-dim mask. + * @param indices_len_ Number of entries in indices + * @param ret0 Output: new tensor, strided the way Inductor expects + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0); + +/** + * out[indices] = values, or out[indices] += values when accumulate is set, + * with the semantics of aten::index_put_. + * + * Inductor calls this fallback, with out and self the same tensor, for an + * index that is a mask, for accumulation into a dtype its generated code + * cannot add atomically, and in deterministic mode. Writes that land on the + * same element are applied in index order, and half and bfloat16 values are + * summed in float, so the result is deterministic and small terms are not + * lost. A mask makes this call wait for the stream. Indices are checked as + * for aoti_torch_cuda_index_Tensor. + * + * @param out Tensor to write (CUDA device, up to 16 dims, as for + * aoti_torch_cuda_index_Tensor) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param indices num_indices index tensors, null for a full slice, with the + * same dtypes as for aoti_torch_cuda_index_Tensor + * @param num_indices Number of entries in indices + * @param values Values to write, broadcastable to the indexed shape, same + * dtype as out + * @param accumulate If true, add values instead of writing them + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate); + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 5ff2b15867c..d5635100b62 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(is_fbcode = False): name = "runtime_shims", srcs = [ "shims/cuda_guard.cpp", + "shims/index.cu", "shims/int4mm.cu", "shims/memory.cpp", "shims/rand.cu", @@ -43,6 +44,7 @@ def define_common_targets(is_fbcode = False): ], headers = [ "shims/cuda_guard.h", + "shims/index.h", "shims/int4mm.cuh", "shims/int4mm.h", "shims/memory.h", diff --git a/backends/cuda/runtime/test/test_cuda_output_copy.cpp b/backends/cuda/runtime/test/test_cuda_output_copy.cpp index 1952b23a18f..e1b72bd2f17 100644 --- a/backends/cuda/runtime/test/test_cuda_output_copy.cpp +++ b/backends/cuda/runtime/test/test_cuda_output_copy.cpp @@ -12,7 +12,7 @@ namespace slim = ::executorch::backends::aoti::slim; using namespace ::executorch::runtime; -using ::executorch::backends::cuda::copy_slimtensor_to_device_etensor_async; +using namespace ::executorch::backends::cuda; namespace { @@ -115,4 +115,74 @@ TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { Error::InvalidArgument); } +// Whether an output with these sizes and strides can be raw-copied into an +// ETensor that has this shape, dim order and strides from an earlier call. +bool raw_copyable( + std::initializer_list output_sizes, + std::initializer_list output_strides, + std::vector sizes, + std::vector dim_order, + std::vector strides) { + std::vector data(192); + slim::SlimTensor output = slim::from_blob( + data.data(), output_sizes, output_strides, slim::c10::ScalarType::Float); + etensor::TensorImpl impl( + etensor::ScalarType::Float, + static_cast(sizes.size()), + sizes.data(), + data.data(), + dim_order.data(), + strides.data()); + etensor::Tensor etensor(&impl); + return output_is_raw_copyable(&output, &etensor); +} + +TEST(CudaOutputRawCopyTest, AcceptsAnOutputLaidOutLikeItsETensor) { + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + // The ETensor's strides are stale from a smaller shape. + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 1, 12, 3}, {1, 3, 2, 2}, {0, 2, 3, 1}, {12, 1, 6, 3})); + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 16, 4, 1}, {1, 3, 2, 2}, {0, 1, 2, 3}, {12, 4, 2, 1})); + // A size-1 dim may have any stride, as with a batch of one. + EXPECT_TRUE(raw_copyable( + {1, 3, 4, 4}, {1, 1, 12, 3}, {1, 3, 4, 4}, {0, 2, 3, 1}, {48, 1, 12, 3})); + // An empty output has nothing to copy, whatever its strides. + EXPECT_TRUE(raw_copyable( + {2, 0, 4, 4}, + {16, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); +} + +TEST(CudaOutputRawCopyTest, RejectsAnOutputLaidOutUnlikeItsETensor) { + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + // Every row is followed by a gap, as in a slice of a wider tensor. + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {96, 32, 8, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + EXPECT_FALSE(raw_copyable( + {3, 4, 4}, {16, 4, 1}, {1, 3, 4, 4}, {0, 1, 2, 3}, {48, 16, 4, 1})); +} + } // namespace diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index 6ca1c64d493..6dd860e11e1 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -299,6 +299,43 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( } } // namespace +/** + * Whether copying an output's raw bytes into its ETensor puts every element + * where the ETensor expects it, which is how a CUDA graph replay lands it. + * + * The ETensor may still have its previous shape, so its strides can be stale + * for dynamic shapes. Its expected strides come from its dim order and the + * output's sizes instead, the way resizing it computes them. Size-1 dims never + * move an element, so their strides may differ. + * + * @param output Pointer to the output SlimTensor (must not be null). + * @param etensor Pointer to the destination ETensor (must not be null). + * @return True if the raw copy is correct, including when there is nothing + * to copy. + */ +inline bool output_is_raw_copyable( + const executorch::backends::aoti::slim::SlimTensor* output, + const executorch::runtime::etensor::Tensor* etensor) { + const size_t ndim = output->dim(); + if (ndim != static_cast(etensor->dim())) { + return false; + } + if (output->numel() == 0) { + return true; + } + const auto dim_order = etensor->dim_order(); + int64_t expected_stride = 1; + for (size_t i = ndim; i > 0; --i) { + const int64_t d = dim_order[i - 1]; + const int64_t size = output->size(d); + if (size != 1 && output->stride(d) != expected_stride) { + return false; + } + expected_stride *= size; + } + return true; +} + /** * Copies data from a SlimTensor to an ETensor asynchronously. * diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index d962db09117..96f637a829c 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -6,6 +6,7 @@ import contextlib import os +import tempfile import threading import unittest from typing import Tuple @@ -653,23 +654,63 @@ 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.""" + def test_missing_scatter_kernel_fails_at_export(self): + """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" - class Int64Accumulate(torch.nn.Module): + class Int64ScatterAdd(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), - ) + return counts.scatter_add(0, idx, torch.ones_like(idx)) 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) + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_mask_index_and_index_put(self): + """A mask index makes the result size depend on the data, and Inductor + has no atomic add for int64, so these run in the index shims.""" + from executorch.runtime import Runtime + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + model = MaskIndex() + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + edge = self._export_to_cuda_with_lower( + model, (x, masks[0], column, idx, counts) + ) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + edge.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for mask in masks: + inputs = (x, mask, column, idx, counts) + outputs = method.execute(list(inputs)) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) def test_device_info_propagated_to_cuda_delegate_outputs(self): """ diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index b1bf91638a7..431bb56a88f 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -282,6 +282,15 @@ def forward(self, x): return self.proj(x)[:, :, :14] +class _Double(nn.Module): + def __init__(self, memory_format): + super().__init__() + self.memory_format = memory_format + + def forward(self, x): + return (x * 2).contiguous(memory_format=self.memory_format) + + class TestStridedOutputWithCudaGraph(unittest.TestCase): """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" @@ -359,6 +368,20 @@ def test_strided_output_first_seen_at_capture(self): else: self._assert_matches_eager(module, x, out) + def test_output_shape_changes_at_capture(self): + torch.manual_seed(0) + # The capture call has a larger batch than the warmup calls before it. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + batch = torch.export.Dim("batch", min=1, max=8) + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({0: batch},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + self._assert_matches_eager(module, x, out) + if __name__ == "__main__": unittest.main() diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 99cba99cfc5..d81d9e9e187 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -41,6 +41,8 @@ "at::_ops::sort_stable::call", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", } ) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 9b8b9e34008..90618c7cf93 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1494,13 +1494,14 @@ Method::set_output_data_ptr(void* buffer, size_t size, size_t output_idx) { return Error::InvalidArgument; } + // A later call can grow a dynamic output back to its upper bound. ET_CHECK_OR_RETURN_ERROR( - t.nbytes() <= size, + tensor_meta->nbytes() <= size, InvalidArgument, "buffer size: %" ET_PRIsize_t - " is smaller then expected tensor size: %" ET_PRIsize_t, + " is smaller than the largest output size: %" ET_PRIsize_t, size, - t.nbytes()); + tensor_meta->nbytes()); // Set data return internal::set_tensor_data(t, buffer, size); diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 834325f008c..501569285f9 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -154,8 +154,8 @@ class Method final { * * @param[in] buffer The block of memory to point the specified tensor at. * - * @param[in] size the length of buffer in bytes, must be >= the nbytes of the - * specified tensor. + * @param[in] size the length of buffer in bytes, must be >= the largest + * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index 0310c5d648b..ef97f8db657 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -317,6 +317,33 @@ TEST_F(MethodTest, AliasedIOTest) { } } +TEST_F(MethodTest, SetOutputDataPtrChecksUpperBound) { + // The "cat" model's output is at most (4, 4) and is not memory planned. + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + float input[4] = {}; + int32_t sizes[2] = {1, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, 2, sizes, input, dim_order, strides); + ASSERT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), Error::Ok); + float output[16]; + ASSERT_EQ(method->set_output_data_ptr(output, sizeof(output), 0), Error::Ok); + ASSERT_EQ(method->execute(), Error::Ok); + ASSERT_EQ(method->get_output(0).toTensor().nbytes(), 8 * sizeof(float)); + + // A later call can grow the output back to (4, 4), so a buffer that only + // fits the current (2, 4) output must be refused. + float small_output[8]; + EXPECT_EQ( + method->set_output_data_ptr(small_output, sizeof(small_output), 0), + Error::InvalidArgument); +} + TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input. // set_input validates numel/nbytes overflow before resize_tensor. From 0dceaca26e05f6e8a32f587274e5cdaae66722b2 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:06:15 -0700 Subject: [PATCH 04/20] CUDA backend: per-program weights file names and antialiased resize (#24314) This change fixes several problems in the CUDA backend. Each fix is its own commit. ## Name the weights file of each program The backend saves every program's weights to `aoti_cuda_blob.ptd`. Exporting several programs into one folder overwrites the weights, and the earlier programs then fail to load. Renaming each file after export works, but is easy to miss. A new option names the file: ```python partitioner = CudaPartitioner([ CudaBackend.generate_method_name_compile_spec("forward"), CudaBackend.generate_data_name_compile_spec("encoder"), ]) # writes encoder.ptd ``` Without the option nothing changes. The name is checked: up to 200 ASCII letters, digits, `_`, `-` and `.`, not starting with `.` or `-`, not ending with `.ptd`, and not a Windows device name such as `NUL`. All methods of a program must share it. A backend that keeps its data inside the `.pte`, such as Metal, refuses it. ## Lower antialiased resizes `F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)` failed to lower with "missing fallback kernels". The backend now rewrites it into plain tensor math, with PyTorch's window and weight formulas, for dynamic sizes and float32, float16 and bfloat16. The results are not bit-identical: in float32 they stay within about 2e-5 of PyTorch's kernel for inputs between 0 and 1. ## Tests Ran on an A100 GPU (Linux x86_64): `test_cuda_export.py`, `test_antialiased_resize.py`, `test_missing_weights_blob.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py`. Each test added for a fix fails when that fix is removed. (cherry picked from commit 7586c16e8a1d3e8d2de375f47ba3fa383e0455d3) --- backends/aoti/aoti_backend.py | 54 ++++- backends/cuda/cuda_backend.py | 206 ++++++++++++++++++ backends/cuda/cuda_weight_collector.py | 16 +- backends/cuda/merge_ptes.py | 9 +- .../cuda/tests/test_antialiased_resize.py | 173 +++++++++++++++ backends/cuda/tests/test_cuda_export.py | 54 +++++ backends/cuda/tests/test_merge_ptes.py | 15 +- .../cuda/tests/test_missing_weights_blob.py | 72 +++++- docs/source/backends/cuda/cuda-overview.md | 19 ++ .../backends/tensorrt/tensorrt-overview.md | 11 +- 10 files changed, 612 insertions(+), 17 deletions(-) create mode 100644 backends/cuda/tests/test_antialiased_resize.py diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 14f3b7cbd11..0655081035a 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -27,6 +27,29 @@ class COMPILE_SPEC_KEYS(Enum): METHOD_NAME = "method_name" + DATA_NAME = "data_name" + + +# The name becomes a file name on every platform, so keep it short, to portable +# characters, and clear of Windows device names such as NUL. Without re.ASCII, +# IGNORECASE lets [A-Za-z] match letters like the Kelvin sign. A leading '-' +# reads as an option on a command line. Some writers add .ptd only when it is +# missing, so "encoder" and "encoder.ptd" could both write encoder.ptd. +_DATA_NAME_PATTERN = re.compile( + r"(?!(?:CON|PRN|AUX|NUL|COM\d|LPT\d)(?:\.|$))(?!.*\.ptd$)" + r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,199}", + re.IGNORECASE | re.ASCII, +) + + +def _checked_data_name(data_name: str) -> str: + if not _DATA_NAME_PATTERN.fullmatch(data_name): + raise ValueError( + f"Invalid data name {data_name!r}: use up to 200 ASCII letters, " + "digits, '_', '-' and '.', do not start with '.' or '-', do not end " + "with '.ptd', and avoid Windows device names such as NUL" + ) + return data_name @experimental( @@ -330,7 +353,9 @@ def preprocess( # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file external_tag = ( - f"aoti_{device_name}_blob" if cls.save_data_externally() else None + cls.data_name_from_compile_specs(compile_specs) + if cls.save_data_externally() + else None ) named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) @@ -366,6 +391,33 @@ def generate_method_name_compile_spec( method_name.encode("utf-8"), ) + @classmethod + def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec: + """ + Generate a CompileSpec that names the external data file, written as + ``.ptd``, so several programs can share a folder. + """ + if not cls.save_data_externally(): + raise ValueError( + f"{cls.__name__} keeps its data in the .pte file, so it takes no " + "data name" + ) + return CompileSpec( + COMPILE_SPEC_KEYS.DATA_NAME.value, + _checked_data_name(data_name).encode("utf-8"), + ) + + @classmethod + def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str: + """ + The external data file name, without ``.ptd``, set by + generate_data_name_compile_spec, or ``aoti__blob`` by default. + """ + for spec in compile_specs: + if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value: + return _checked_data_name(spec.value.decode("utf-8")) + return f"aoti_{cls.get_device_name()}_blob" + @classmethod def method_name_from_compile_specs( cls, diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 787d2502b63..20203972bb1 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -9,6 +9,7 @@ import copy import functools import logging +import math import os import shutil import threading @@ -50,7 +51,9 @@ from executorch.exir._warnings import experimental from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec +from torch._decomp import decomposition_table from torch._inductor.decomposition import conv1d_to_conv2d +from torch._prims_common import suggest_memory_format from torch.nn.attention import SDPBackend # --------------------------------------------------------------------------- @@ -482,6 +485,183 @@ 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) + + +def _linear_kernel(x: torch.Tensor) -> torch.Tensor: + return torch.clamp(1 - x.abs(), min=0) + + +def _cubic_kernel(x: torch.Tensor) -> torch.Tensor: + # Keys' cubic with a = -0.5, the one ATen's antialiased bicubic uses. + a = -0.5 + x = x.abs() + near = ((a + 2) * x - (a + 3)) * x * x + 1 + far = (((x - 5) * x + 8) * x - 4) * a + return torch.where(x < 1, near, torch.where(x < 2, far, torch.zeros_like(x))) + + +def _antialias_taps( + in_size: int, + out_size: int, + align_corners: bool, + scale: Optional[float], + support: int, + kernel: typing.Callable[[torch.Tensor], torch.Tensor], + device: torch.device, +) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Taps of an antialiased resize along one dim, computed as ATen does: + output o reads input index[o, k] with weight[o, k], and inside[o, k] + marks the taps in its window. The window is the kernel's support widened + by the downscale factor, and the weights in it are normalized. + """ + if align_corners or scale is None or scale <= 0: + if align_corners: + # ATen uses a ratio of 0 for an output size of 1. min and max give + # that without a conditional, which torch.compile cannot handle for + # dynamic sizes. + num = (in_size - 1) * torch.sym_min(out_size - 1, 1) + den = torch.sym_max(out_size - 1, 1) + else: + num, den = in_size, out_size + ratio = num / den + # The tap count is a shape that both the host and the kernel work out. + # In float32 the kernel can round num / den just past a whole number + # and count more taps than the host allocated, so use integers. + reach = torch.sym_max((support * num + den - 1) // den, support) + else: + ratio = 1.0 / scale + reach = math.ceil(support * max(ratio, 1.0)) + widen = torch.sym_max(ratio, 1.0) + dtype = torch.float32 + center = ratio * (torch.arange(out_size, dtype=dtype, device=device) + 0.5) + first = torch.clamp(torch.trunc(center - support * widen + 0.5), min=0) + end = torch.clamp(torch.trunc(center + support * widen + 0.5), max=in_size) + taps = 2 * reach + 1 + tap = torch.arange(taps, dtype=dtype, device=device) + position = first[:, None] + tap + inside = position < end[:, None] + weights = kernel((tap + (first - center)[:, None] + 0.5) * (1.0 / widen)) + weights = torch.where(inside, weights, 0) + total = weights.sum(1, keepdim=True) + weights = torch.where(total != 0, weights / total, weights) + # Clamp after the cast: in float32, in_size - 1 rounds up to in_size for + # very long axes. + index = position.to(torch.int64).clamp(max=in_size - 1) + return index, weights, inside + + +def _antialiased_resize( + support: int, + kernel: typing.Callable[[torch.Tensor], torch.Tensor], + input: torch.Tensor, + output_size: List[int], + align_corners: bool, + scales_h: Optional[float] = None, + scales_w: Optional[float] = None, +) -> torch.Tensor: + """ + Rewrite an antialiased bilinear or bicubic resize as a weighted sum over + each output's taps, one dim at a time. Inductor has no lowering for these + resizes, so it would call ATen's kernel, which this backend does not ship. + """ + # The backend's tensors have no float64, and ATen's CUDA kernel has no + # integer version. + if input.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return NotImplemented + # Half types are computed in float and rounded once, so they can differ + # slightly from ATen, which keeps its weights in the input type. + result = input.float() + for dim, scale in ((3, scales_w), (2, scales_h)): + index, weights, inside = _antialias_taps( + input.shape[dim], + output_size[dim - 2], + align_corners, + scale, + support, + kernel, + input.device, + ) + rows, taps = index.shape + trailing = [1] * (3 - dim) + # Plain indexing lets Inductor fold the gather into the sum instead of + # keeping every tap of every row in memory. + gathered = result[(slice(None),) * dim + (index,)] + # Taps outside the window already have zero weight, but NaN * 0 is NaN, + # so mask them out to keep a NaN or inf local. + weighted = torch.where( + inside.view(rows, taps, *trailing), + gathered * weights.view(rows, taps, *trailing), + 0, + ) + result = weighted.sum(dim + 1) + return result.to(input.dtype).contiguous(memory_format=suggest_memory_format(input)) + + @final @experimental( "This API and all of cuda backend related functionality are experimental." @@ -638,6 +818,7 @@ def _preprocess_with_weight_collector( cls.get_device_name(), target_sm=target_sm, ptx_compute=ptx_compute, + external_tag=cls.data_name_from_compile_specs(compile_specs), ) return result @@ -660,6 +841,17 @@ def preprocess_multimethod( compile_specs: Dict[str, List[List[CompileSpec]]], ) -> Dict[str, List[PreprocessResult]]: """Compile every method against one global ``(device, FQN)`` store.""" + # The methods share one weights file, so they must agree on its name. + data_names = { + cls.data_name_from_compile_specs(specs) + for method_specs in compile_specs.values() + for specs in method_specs + } + if len(data_names) > 1: + raise ValueError( + "Every method of a program must use the same data name, got " + f"{sorted(data_names)}" + ) collector = CudaWeightCollector() results: Dict[str, List[PreprocessResult]] = {} for method_name, programs in edge_programs.items(): @@ -780,6 +972,20 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: return { torch.ops.aten.conv1d.default: conv1d_to_conv2d, } + for resize, support, kernel in ( + (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel), + (torch.ops.aten._upsample_bicubic2d_aa, 2, _cubic_kernel), + ): + table[resize.default] = functools.partial( + _antialiased_resize, support, kernel + ) + # The vec form works out the output size and calls the default form. + table[resize.vec] = decomposition_table[resize.vec] + # 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..a56e0c7e348 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -568,6 +568,7 @@ def add_preprocess_result( device_name: str, target_sm: Optional[int] = None, ptx_compute: int = 0, + external_tag: Optional[str] = None, ) -> None: if result.data_store_output is None: raise RuntimeError("CUDA AOTI preprocess returned no named data") @@ -587,7 +588,7 @@ def add_preprocess_result( keep_compatibility_blob=not artifact.entries, ) - external_tag = f"aoti_{device_name}_blob" + external_tag = external_tag or f"aoti_{device_name}_blob" serialized_entries = [] for entry in artifact.entries: data = artifact.storages.get(entry.storage_key) @@ -599,6 +600,19 @@ def add_preprocess_result( ), ) if data is not None: + # Programs can write their 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: diff --git a/backends/cuda/merge_ptes.py b/backends/cuda/merge_ptes.py index 14456dda54f..184bbfd3b86 100644 --- a/backends/cuda/merge_ptes.py +++ b/backends/cuda/merge_ptes.py @@ -208,6 +208,11 @@ def _load_artifact(source: CudaPteInput) -> _Artifact: return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates) +# Compile specs that only shape how a variant was exported, so variants that +# differ in them still merge. +_EXPORT_ONLY_SPEC_KEYS = ("cuda_include_ptx", "data_name") + + def _normalized_program(program: Program) -> Program: normalized = copy.deepcopy(program) payloads = [] @@ -219,7 +224,7 @@ def _normalized_program(program: Program) -> Program: delegate.compile_specs = [ spec for spec in delegate.compile_specs - if spec.key != "cuda_include_ptx" + if spec.key not in _EXPORT_ONLY_SPEC_KEYS ] delegate.processed = BackendDelegateDataReference( location=DataLocation.INLINE, index=len(payloads) @@ -534,7 +539,7 @@ def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStor delegate.compile_specs = [ spec for spec in delegate.compile_specs - if spec.key != "cuda_include_ptx" + if spec.key not in _EXPORT_ONLY_SPEC_KEYS ] merged_store = NamedDataStore() diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py new file mode 100644 index 00000000000..963444f80e7 --- /dev/null +++ b/backends/cuda/tests/test_antialiased_resize.py @@ -0,0 +1,173 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Antialiased bilinear and bicubic resizes on the CUDA backend. + +Inductor has no lowering for these resizes, so the backend rewrites them into +plain tensor math. These tests lower real programs, run them through the Python +runtime, and compare with eager PyTorch. +""" + +import functools +import os +import tempfile +import unittest +from typing import Optional, Sequence, Tuple + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from torch.export import Dim, export + + +class Resize(torch.nn.Module): + """Resizes x by scale_factor, or else to like's height and a width of 56.""" + + def __init__( + self, + mode: str, + align_corners: bool = False, + scale_factor: Optional[Tuple[float, float]] = None, + ) -> None: + super().__init__() + self.mode = mode + self.align_corners = align_corners + self.scale_factor = scale_factor + + def forward(self, x: torch.Tensor, like: torch.Tensor) -> torch.Tensor: + size = None if self.scale_factor else (like.shape[2], 56) + return torch.nn.functional.interpolate( + x, + size=size, + scale_factor=self.scale_factor, + mode=self.mode, + align_corners=self.align_corners, + antialias=True, + ) + + +# Dynamic heights let one program shrink, grow, and resize to 1. +DYNAMIC_HEIGHTS = ( + {2: Dim("in_height", min=4, max=64)}, + {2: Dim("out_height", min=1, max=32)}, +) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class TestAntialiasedResize(unittest.TestCase): + def _run( + self, + model: Resize, + cases: Sequence[Tuple[torch.Tensor, torch.Tensor]], + dynamic_shapes=None, + **tolerance: float, + ) -> None: + """Lower model for cases[0], then run every case and compare with eager.""" + exported = export(model, cases[0], dynamic_shapes=dynamic_shapes) + spec = CudaBackend.generate_method_name_compile_spec("forward") + lowered = to_edge_transform_and_lower( + exported, + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + lowered.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for x, like in cases: + (output,) = method.execute([x, like]) + # The rewrite stands in for ATen's CUDA kernel. + expected = model(x.cuda().float(), like.cuda()).to(x.dtype).cpu() + torch.testing.assert_close( + output.cpu(), expected, equal_nan=True, **tolerance + ) + + def _check_dynamic_heights(self, mode: str, align_corners: bool) -> None: + cases = [] + # 21 to 7 and 22 to 8 give whole ratios, with and without align_corners. + heights = ((48, 20), (64, 20), (10, 20), (10, 1), (21, 7), (22, 8)) + for in_height, out_height in heights: + x = torch.randn(1, 3, in_height, 32) + # As in ATen, a NaN reaches only the outputs whose window has it. + x[0, 0, -1, -1] = float("nan") + cases.append((x, torch.empty(1, 1, out_height))) + self._run( + Resize(mode, align_corners), + cases, + DYNAMIC_HEIGHTS, + atol=1e-5, + rtol=1e-4, + ) + + def test_bilinear_antialias_resize(self): + self._check_dynamic_heights("bilinear", align_corners=False) + + def test_bicubic_antialias_resize(self): + self._check_dynamic_heights("bicubic", align_corners=True) + + def test_antialias_resize_with_uneven_scale_factor(self): + # 37 * 1.75 and 51 * 0.45 are not whole, so the scale the resize uses + # differs from output size / input size. + x = torch.randn(1, 3, 37, 51) + self._run( + Resize("bicubic", scale_factor=(1.75, 0.45)), + [(x, torch.empty(0))], + atol=1e-5, + rtol=1e-4, + ) + + def test_half_precision_antialias_resize(self): + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + x = torch.randn(1, 3, 48, 32).to( + dtype=dtype, memory_format=torch.channels_last + ) + self._run( + Resize("bilinear"), + [(x, torch.empty(1, 1, 20))], + atol=1e-2, + rtol=1e-2, + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestAntialiasedResizeDecomposition(unittest.TestCase): + def test_keeps_memory_format(self): + # The program's output layout is fixed at export, so only a direct call + # shows the rewrite's own layout. + resize = CudaBackend.get_decomposition_table()[ + torch.ops.aten._upsample_bilinear2d_aa.default + ] + x = torch.randn(1, 3, 48, 32, device="cuda") + output = resize(x.to(memory_format=torch.channels_last), [20, 56], False) + self.assertTrue(output.is_contiguous(memory_format=torch.channels_last)) + + def test_align_corners_compiles_with_dynamic_sizes(self): + from torch._inductor.compile_fx import compile_fx + + model = Resize("bicubic", align_corners=True) + compiled = torch.compile( + model, + backend=functools.partial( + compile_fx, decompositions=CudaBackend.get_decomposition_table() + ), + dynamic=True, + fullgraph=True, + ) + for in_height, out_height in ((48, 20), (10, 1)): + x = torch.randn(1, 3, in_height, 32, device="cuda") + like = torch.empty(1, 1, out_height, device="cuda") + torch.testing.assert_close( + compiled(x, like), model(x, like), atol=1e-5, rtol=1e-4 + ) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8329f42e522..6585328f63e 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -236,6 +236,60 @@ def test_invalid_cuda_include_ptx_compile_spec(self): [CompileSpec(key="cuda_include_ptx", value=b"MAYBE")] ) + def test_invalid_data_name_compile_spec(self): + for name in ( + "", + "..", + ".hidden", + "../outside", + "a/b", + "a\\b", + "a/../../outside", + "C:model", + "model?", + "nul", + "COM1.txt", + "a" * 201, + # Non-ASCII letters that IGNORECASE alone would let through. + "\u212a" * 200, + "\u017fizes", + "\u0130d", + "\u0131d", + "-enc", + # One writer turns "encoder.ptd" into the same file as "encoder". + # The check ignores case, like the rest of the pattern. + "encoder.ptd", + "encoder.PTD", + ): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.generate_data_name_compile_spec(name) + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.data_name_from_compile_specs( + [CompileSpec(key="data_name", value=name.encode())] + ) + # A backend that keeps its data in the .pte would ignore the name. + with patch.object(CudaBackend, "save_data_externally", return_value=False): + with self.assertRaisesRegex(ValueError, "takes no data name"): + CudaBackend.generate_data_name_compile_spec("encoder") + + def test_valid_data_name_compile_spec(self): + self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") + for name in ("encoder", "model_sm80", "v1.2-final", "a" * 200, "x.ptd.bak"): + with self.subTest(name=name): + spec = CudaBackend.generate_data_name_compile_spec(name) + self.assertEqual(CudaBackend.data_name_from_compile_specs([spec]), name) + + def test_methods_of_one_program_must_share_a_data_name(self): + def specs(name): + return [[CudaBackend.generate_data_name_compile_spec(name)]] + + with self.assertRaisesRegex(ValueError, "same data name"): + CudaBackend.preprocess_multimethod( + {"forward": [None], "second": [None]}, + {"forward": specs("first"), "second": specs("second")}, + ) + class TestCudaExport(unittest.TestCase): """Test CUDA export functionality for various operations using to_edge_transform_and_lower.""" diff --git a/backends/cuda/tests/test_merge_ptes.py b/backends/cuda/tests/test_merge_ptes.py index f0bb773070f..8ca8b8ecee5 100644 --- a/backends/cuda/tests/test_merge_ptes.py +++ b/backends/cuda/tests/test_merge_ptes.py @@ -58,6 +58,7 @@ def _write_artifact( fqn: str = "model.weight", weight_key: str = "cuda_fqn_weight:cuda:model.weight", ptx_compute: int | None = None, + data_name: str | None = None, ) -> CudaPteInput: if ptx_compute is None: ptx_compute = target_sm @@ -79,6 +80,8 @@ def _write_artifact( compile_specs.append( CompileSpec("cuda_include_ptx", b"ON" if ptx_compute else b"OFF") ) + if data_name is not None: + compile_specs.append(CompileSpec("data_name", data_name.encode())) delegate = BackendDelegate( id="CudaBackend", processed=BackendDelegateDataReference(DataLocation.INLINE, 0), @@ -148,8 +151,12 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: sm80.mkdir() sm120.mkdir() inputs = [ - self._write_artifact(sm80, 80, b"sm80-so", b"shared-weight"), - self._write_artifact(sm120, 120, b"sm120-so", b"shared-weight"), + self._write_artifact( + sm80, 80, b"sm80-so", b"shared-weight", data_name="model_sm80" + ), + self._write_artifact( + sm120, 120, b"sm120-so", b"shared-weight", data_name="model_sm120" + ), ] merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) @@ -176,14 +183,14 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: hashlib.sha256(b"sm120-so").hexdigest() + "_so_blob", }, ) - self.assertNotIn( - "cuda_include_ptx", + self.assertEqual( { spec.key for spec in merged.program.execution_plan[0] .delegates[0] .compile_specs }, + {"method_name"}, ) def test_uses_ptx_only_from_explicit_fallback(self) -> None: diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index 410184ee238..d4e65edddd9 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -22,6 +22,7 @@ import os import tempfile import unittest +from typing import Optional import torch from executorch.backends.cuda.cuda_backend import CudaBackend @@ -51,14 +52,25 @@ 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) - module = ModuleWithConstants().eval().cuda() + def _lower( + self, + outdir: str, + seed: int = 0, + name: str = "model.pte", + data_name: Optional[str] = None, + module: Optional[torch.nn.Module] = None, + ) -> str: + torch.manual_seed(seed) + if module is None: + module = ModuleWithConstants() + module = module.eval().cuda() inputs = (torch.randn(4, 64, device="cuda"),) exported = export(module, inputs) - spec = CudaBackend.generate_method_name_compile_spec("forward") + specs = [CudaBackend.generate_method_name_compile_spec("forward")] + if data_name is not None: + specs.append(CudaBackend.generate_data_name_compile_spec(data_name)) lowered = to_edge_transform_and_lower( - exported, partitioner=[CudaPartitioner([spec])] + exported, partitioner=[CudaPartitioner(specs)] ) program = lowered.to_executorch() path = os.path.join(outdir, "model.pte") @@ -160,6 +172,56 @@ 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_programs_with_their_own_data_names_share_a_folder(self) -> None: + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + for seed, name in ((0, "first"), (1, "second")): + self._lower(outdir, seed=seed, name=f"{name}.pte", data_name=name) + # A program with no constants still writes a data file. + self._lower( + outdir, name="third.pte", data_name="third", module=torch.nn.ReLU() + ) + self.assertEqual( + sorted(f for f in os.listdir(outdir) if f.endswith(".ptd")), + ["first.ptd", "second.ptd", "third.ptd"], + ) + for seed, name in ((0, "first"), (1, "second")): + # The same draws _lower made, so the same weights and input. + torch.manual_seed(seed) + module = ModuleWithConstants().eval().cuda() + value = torch.randn(4, 64, device="cuda") + method = ( + Runtime.get() + .load_program( + os.path.join(outdir, f"{name}.pte"), + data_path=os.path.join(outdir, f"{name}.ptd"), + ) + .load_method("forward") + ) + out = method.execute([value.cpu()]) + torch.testing.assert_close( + (out[0] if isinstance(out, (list, tuple)) else out).cpu(), + module(value).cpu(), + ) + 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/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index 518073101d9..a78a4b831fa 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -81,6 +81,25 @@ save_pte_program(exec_program, model_name, "./output_dir") ``` This generates `.pte` and `.ptd` files that can be executed on CUDA devices. +By default, the weights go to a data file named `aoti_cuda_blob.ptd`. To export several programs into one folder, give each program its own +data file name. This partitioner writes `encoder.ptd`: + +```python +partitioner = CudaPartitioner([ + CudaBackend.generate_method_name_compile_spec("forward"), + CudaBackend.generate_data_name_compile_spec("encoder"), +]) +``` + +The name can have up to 200 ASCII letters, digits, `_`, `-` and `.`. It cannot start with `.` or `-`, cannot end with `.ptd` +(the export adds it), and cannot be a Windows device name such as `NUL`. Load each program with its own data file: + +```python +from executorch.runtime import Runtime + +program = Runtime.get().load_program("encoder.pte", data_path="encoder.ptd") +``` + For a complete working example, see the [CUDA export script](https://github.com/pytorch/executorch/blob/main/examples/cuda/scripts/export.py). diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index 411cc262961..9af3a20a2da 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -170,9 +170,11 @@ 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. +without it. By default 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 then fails to load. To keep +several models in one directory, give each its own data file name with +`CudaBackend.generate_data_name_compile_spec`, as the +[CUDA backend page](../cuda/cuda-overview.md) shows, or save each model into its own directory. ### Keeping inputs and outputs on the GPU @@ -432,7 +434,8 @@ shows how to set the options that apply at run time, from C++ or from Python. Pass `external_engine_data="engines"` to `torch_tensorrt.save` to write the engines to `engines.ptd` instead of the `.pte`, and pass that file when you load the program. The Python loader takes one data file, so for a program that also uses the CUDA backend, pass - `external_engine_data="aoti_cuda_blob"` to share its file. + the name of the CUDA data file to share it: `external_engine_data="aoti_cuda_blob"` by default, + or the name given to `CudaBackend.generate_data_name_compile_spec`. - [**Optimization profiles.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#optimization-profiles) One engine can be tuned for several ranges of input shapes, for example one for decoding one token at a time and one for a long prompt. A C++ application picks the profile for each call From 6c27e1cb5819cf4c7a1d76e0237f0a473f658826 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:06:15 -0700 Subject: [PATCH 05/20] CUDA backend: per-program weights file names and antialiased resize (#24314) This change fixes several problems in the CUDA backend. Each fix is its own commit. ## Name the weights file of each program The backend saves every program's weights to `aoti_cuda_blob.ptd`. Exporting several programs into one folder overwrites the weights, and the earlier programs then fail to load. Renaming each file after export works, but is easy to miss. A new option names the file: ```python partitioner = CudaPartitioner([ CudaBackend.generate_method_name_compile_spec("forward"), CudaBackend.generate_data_name_compile_spec("encoder"), ]) # writes encoder.ptd ``` Without the option nothing changes. The name is checked: up to 200 ASCII letters, digits, `_`, `-` and `.`, not starting with `.` or `-`, not ending with `.ptd`, and not a Windows device name such as `NUL`. All methods of a program must share it. A backend that keeps its data inside the `.pte`, such as Metal, refuses it. ## Lower antialiased resizes `F.interpolate(x, mode="bilinear" or "bicubic", antialias=True)` failed to lower with "missing fallback kernels". The backend now rewrites it into plain tensor math, with PyTorch's window and weight formulas, for dynamic sizes and float32, float16 and bfloat16. The results are not bit-identical: in float32 they stay within about 2e-5 of PyTorch's kernel for inputs between 0 and 1. ## Tests Ran on an A100 GPU (Linux x86_64): `test_cuda_export.py`, `test_antialiased_resize.py`, `test_missing_weights_blob.py`, `test_merge_ptes.py` and `test_cuda_weight_metadata.py`. Each test added for a fix fails when that fix is removed. --- backends/aoti/aoti_backend.py | 54 +++++- backends/cuda/cuda_backend.py | 135 ++++++++++++++ backends/cuda/cuda_weight_collector.py | 5 +- backends/cuda/merge_ptes.py | 9 +- .../cuda/tests/test_antialiased_resize.py | 173 ++++++++++++++++++ backends/cuda/tests/test_cuda_export.py | 54 ++++++ backends/cuda/tests/test_merge_ptes.py | 15 +- .../cuda/tests/test_missing_weights_blob.py | 53 +++++- docs/source/backends/cuda/cuda-overview.md | 19 ++ .../backends/tensorrt/tensorrt-overview.md | 11 +- 10 files changed, 511 insertions(+), 17 deletions(-) create mode 100644 backends/cuda/tests/test_antialiased_resize.py diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 4a2ed0be168..0d0670b7318 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -28,6 +28,29 @@ class COMPILE_SPEC_KEYS(Enum): METHOD_NAME = "method_name" + DATA_NAME = "data_name" + + +# The name becomes a file name on every platform, so keep it short, to portable +# characters, and clear of Windows device names such as NUL. Without re.ASCII, +# IGNORECASE lets [A-Za-z] match letters like the Kelvin sign. A leading '-' +# reads as an option on a command line. Some writers add .ptd only when it is +# missing, so "encoder" and "encoder.ptd" could both write encoder.ptd. +_DATA_NAME_PATTERN = re.compile( + r"(?!(?:CON|PRN|AUX|NUL|COM\d|LPT\d)(?:\.|$))(?!.*\.ptd$)" + r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,199}", + re.IGNORECASE | re.ASCII, +) + + +def _checked_data_name(data_name: str) -> str: + if not _DATA_NAME_PATTERN.fullmatch(data_name): + raise ValueError( + f"Invalid data name {data_name!r}: use up to 200 ASCII letters, " + "digits, '_', '-' and '.', do not start with '.' or '-', do not end " + "with '.ptd', and avoid Windows device names such as NUL" + ) + return data_name @experimental( @@ -368,7 +391,9 @@ def preprocess( # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file external_tag = ( - f"aoti_{device_name}_blob" if cls.save_data_externally() else None + cls.data_name_from_compile_specs(compile_specs) + if cls.save_data_externally() + else None ) named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) @@ -404,6 +429,33 @@ def generate_method_name_compile_spec( method_name.encode("utf-8"), ) + @classmethod + def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec: + """ + Generate a CompileSpec that names the external data file, written as + ``.ptd``, so several programs can share a folder. + """ + if not cls.save_data_externally(): + raise ValueError( + f"{cls.__name__} keeps its data in the .pte file, so it takes no " + "data name" + ) + return CompileSpec( + COMPILE_SPEC_KEYS.DATA_NAME.value, + _checked_data_name(data_name).encode("utf-8"), + ) + + @classmethod + def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str: + """ + The external data file name, without ``.ptd``, set by + generate_data_name_compile_spec, or ``aoti__blob`` by default. + """ + for spec in compile_specs: + if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value: + return _checked_data_name(spec.value.decode("utf-8")) + return f"aoti_{cls.get_device_name()}_blob" + @classmethod def method_name_from_compile_specs( cls, diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1f55a6099a9..891fb81bd06 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -9,6 +9,7 @@ import copy import functools import logging +import math import os import shutil import threading @@ -50,7 +51,9 @@ from executorch.exir._warnings import experimental from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec +from torch._decomp import decomposition_table from torch._inductor.decomposition import conv1d_to_conv2d +from torch._prims_common import suggest_memory_format from torch.nn.attention import SDPBackend # --------------------------------------------------------------------------- @@ -548,6 +551,117 @@ def _conv_transpose_as_conv( return conv(x, weight, bias, [1] * n, [0] * n, dilation, groups) +def _linear_kernel(x: torch.Tensor) -> torch.Tensor: + return torch.clamp(1 - x.abs(), min=0) + + +def _cubic_kernel(x: torch.Tensor) -> torch.Tensor: + # Keys' cubic with a = -0.5, the one ATen's antialiased bicubic uses. + a = -0.5 + x = x.abs() + near = ((a + 2) * x - (a + 3)) * x * x + 1 + far = (((x - 5) * x + 8) * x - 4) * a + return torch.where(x < 1, near, torch.where(x < 2, far, torch.zeros_like(x))) + + +def _antialias_taps( + in_size: int, + out_size: int, + align_corners: bool, + scale: Optional[float], + support: int, + kernel: typing.Callable[[torch.Tensor], torch.Tensor], + device: torch.device, +) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Taps of an antialiased resize along one dim, computed as ATen does: + output o reads input index[o, k] with weight[o, k], and inside[o, k] + marks the taps in its window. The window is the kernel's support widened + by the downscale factor, and the weights in it are normalized. + """ + if align_corners or scale is None or scale <= 0: + if align_corners: + # ATen uses a ratio of 0 for an output size of 1. min and max give + # that without a conditional, which torch.compile cannot handle for + # dynamic sizes. + num = (in_size - 1) * torch.sym_min(out_size - 1, 1) + den = torch.sym_max(out_size - 1, 1) + else: + num, den = in_size, out_size + ratio = num / den + # The tap count is a shape that both the host and the kernel work out. + # In float32 the kernel can round num / den just past a whole number + # and count more taps than the host allocated, so use integers. + reach = torch.sym_max((support * num + den - 1) // den, support) + else: + ratio = 1.0 / scale + reach = math.ceil(support * max(ratio, 1.0)) + widen = torch.sym_max(ratio, 1.0) + dtype = torch.float32 + center = ratio * (torch.arange(out_size, dtype=dtype, device=device) + 0.5) + first = torch.clamp(torch.trunc(center - support * widen + 0.5), min=0) + end = torch.clamp(torch.trunc(center + support * widen + 0.5), max=in_size) + taps = 2 * reach + 1 + tap = torch.arange(taps, dtype=dtype, device=device) + position = first[:, None] + tap + inside = position < end[:, None] + weights = kernel((tap + (first - center)[:, None] + 0.5) * (1.0 / widen)) + weights = torch.where(inside, weights, 0) + total = weights.sum(1, keepdim=True) + weights = torch.where(total != 0, weights / total, weights) + # Clamp after the cast: in float32, in_size - 1 rounds up to in_size for + # very long axes. + index = position.to(torch.int64).clamp(max=in_size - 1) + return index, weights, inside + + +def _antialiased_resize( + support: int, + kernel: typing.Callable[[torch.Tensor], torch.Tensor], + input: torch.Tensor, + output_size: List[int], + align_corners: bool, + scales_h: Optional[float] = None, + scales_w: Optional[float] = None, +) -> torch.Tensor: + """ + Rewrite an antialiased bilinear or bicubic resize as a weighted sum over + each output's taps, one dim at a time. Inductor has no lowering for these + resizes, so it would call ATen's kernel, which this backend does not ship. + """ + # The backend's tensors have no float64, and ATen's CUDA kernel has no + # integer version. + if input.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return NotImplemented + # Half types are computed in float and rounded once, so they can differ + # slightly from ATen, which keeps its weights in the input type. + result = input.float() + for dim, scale in ((3, scales_w), (2, scales_h)): + index, weights, inside = _antialias_taps( + input.shape[dim], + output_size[dim - 2], + align_corners, + scale, + support, + kernel, + input.device, + ) + rows, taps = index.shape + trailing = [1] * (3 - dim) + # Plain indexing lets Inductor fold the gather into the sum instead of + # keeping every tap of every row in memory. + gathered = result[(slice(None),) * dim + (index,)] + # Taps outside the window already have zero weight, but NaN * 0 is NaN, + # so mask them out to keep a NaN or inf local. + weighted = torch.where( + inside.view(rows, taps, *trailing), + gathered * weights.view(rows, taps, *trailing), + 0, + ) + result = weighted.sum(dim + 1) + return result.to(input.dtype).contiguous(memory_format=suggest_memory_format(input)) + + @final @experimental( "This API and all of cuda backend related functionality are experimental." @@ -704,6 +818,7 @@ def _preprocess_with_weight_collector( cls.get_device_name(), target_sm=target_sm, ptx_compute=ptx_compute, + external_tag=cls.data_name_from_compile_specs(compile_specs), ) return result @@ -726,6 +841,17 @@ def preprocess_multimethod( compile_specs: Dict[str, List[List[CompileSpec]]], ) -> Dict[str, List[PreprocessResult]]: """Compile every method against one global ``(device, FQN)`` store.""" + # The methods share one weights file, so they must agree on its name. + data_names = { + cls.data_name_from_compile_specs(specs) + for method_specs in compile_specs.values() + for specs in method_specs + } + if len(data_names) > 1: + raise ValueError( + "Every method of a program must use the same data name, got " + f"{sorted(data_names)}" + ) collector = CudaWeightCollector() results: Dict[str, List[PreprocessResult]] = {} for method_name, programs in edge_programs.items(): @@ -852,6 +978,15 @@ def get_decomposition_table(cls) -> Dict[Any, Any]: # for it and the shim no fallback kernel. torch.ops.aten._is_all_true.default: torch.ops.aten.all.default, } + for resize, support, kernel in ( + (torch.ops.aten._upsample_bilinear2d_aa, 1, _linear_kernel), + (torch.ops.aten._upsample_bicubic2d_aa, 2, _cubic_kernel), + ): + table[resize.default] = functools.partial( + _antialiased_resize, support, kernel + ) + # The vec form works out the output size and calls the default form. + table[resize.vec] = decomposition_table[resize.vec] # Inductor disables its transposed conv2d Triton template only on CUDA # (https://github.com/pytorch/pytorch/issues/187081). if torch.version.hip is None: diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index 19e9f8c8bb6..f826a34d718 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -570,6 +570,7 @@ def add_preprocess_result( device_name: str, target_sm: Optional[int] = None, ptx_compute: int = 0, + external_tag: Optional[str] = None, ) -> None: if result.data_store_output is None: raise RuntimeError("CUDA AOTI preprocess returned no named data") @@ -589,7 +590,7 @@ def add_preprocess_result( keep_compatibility_blob=not artifact.entries, ) - external_tag = f"aoti_{device_name}_blob" + external_tag = external_tag or f"aoti_{device_name}_blob" serialized_entries = [] for entry in artifact.entries: data = artifact.storages.get(entry.storage_key) @@ -601,7 +602,7 @@ def add_preprocess_result( ), ) if data is not None: - # Every program writes its weights to the same file name, so the + # Programs can write their 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()}" diff --git a/backends/cuda/merge_ptes.py b/backends/cuda/merge_ptes.py index 14456dda54f..184bbfd3b86 100644 --- a/backends/cuda/merge_ptes.py +++ b/backends/cuda/merge_ptes.py @@ -208,6 +208,11 @@ def _load_artifact(source: CudaPteInput) -> _Artifact: return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates) +# Compile specs that only shape how a variant was exported, so variants that +# differ in them still merge. +_EXPORT_ONLY_SPEC_KEYS = ("cuda_include_ptx", "data_name") + + def _normalized_program(program: Program) -> Program: normalized = copy.deepcopy(program) payloads = [] @@ -219,7 +224,7 @@ def _normalized_program(program: Program) -> Program: delegate.compile_specs = [ spec for spec in delegate.compile_specs - if spec.key != "cuda_include_ptx" + if spec.key not in _EXPORT_ONLY_SPEC_KEYS ] delegate.processed = BackendDelegateDataReference( location=DataLocation.INLINE, index=len(payloads) @@ -534,7 +539,7 @@ def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStor delegate.compile_specs = [ spec for spec in delegate.compile_specs - if spec.key != "cuda_include_ptx" + if spec.key not in _EXPORT_ONLY_SPEC_KEYS ] merged_store = NamedDataStore() diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py new file mode 100644 index 00000000000..963444f80e7 --- /dev/null +++ b/backends/cuda/tests/test_antialiased_resize.py @@ -0,0 +1,173 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Antialiased bilinear and bicubic resizes on the CUDA backend. + +Inductor has no lowering for these resizes, so the backend rewrites them into +plain tensor math. These tests lower real programs, run them through the Python +runtime, and compare with eager PyTorch. +""" + +import functools +import os +import tempfile +import unittest +from typing import Optional, Sequence, Tuple + +import torch +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_partitioner import CudaPartitioner +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from torch.export import Dim, export + + +class Resize(torch.nn.Module): + """Resizes x by scale_factor, or else to like's height and a width of 56.""" + + def __init__( + self, + mode: str, + align_corners: bool = False, + scale_factor: Optional[Tuple[float, float]] = None, + ) -> None: + super().__init__() + self.mode = mode + self.align_corners = align_corners + self.scale_factor = scale_factor + + def forward(self, x: torch.Tensor, like: torch.Tensor) -> torch.Tensor: + size = None if self.scale_factor else (like.shape[2], 56) + return torch.nn.functional.interpolate( + x, + size=size, + scale_factor=self.scale_factor, + mode=self.mode, + align_corners=self.align_corners, + antialias=True, + ) + + +# Dynamic heights let one program shrink, grow, and resize to 1. +DYNAMIC_HEIGHTS = ( + {2: Dim("in_height", min=4, max=64)}, + {2: Dim("out_height", min=1, max=32)}, +) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +@unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", +) +class TestAntialiasedResize(unittest.TestCase): + def _run( + self, + model: Resize, + cases: Sequence[Tuple[torch.Tensor, torch.Tensor]], + dynamic_shapes=None, + **tolerance: float, + ) -> None: + """Lower model for cases[0], then run every case and compare with eager.""" + exported = export(model, cases[0], dynamic_shapes=dynamic_shapes) + spec = CudaBackend.generate_method_name_compile_spec("forward") + lowered = to_edge_transform_and_lower( + exported, + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + lowered.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for x, like in cases: + (output,) = method.execute([x, like]) + # The rewrite stands in for ATen's CUDA kernel. + expected = model(x.cuda().float(), like.cuda()).to(x.dtype).cpu() + torch.testing.assert_close( + output.cpu(), expected, equal_nan=True, **tolerance + ) + + def _check_dynamic_heights(self, mode: str, align_corners: bool) -> None: + cases = [] + # 21 to 7 and 22 to 8 give whole ratios, with and without align_corners. + heights = ((48, 20), (64, 20), (10, 20), (10, 1), (21, 7), (22, 8)) + for in_height, out_height in heights: + x = torch.randn(1, 3, in_height, 32) + # As in ATen, a NaN reaches only the outputs whose window has it. + x[0, 0, -1, -1] = float("nan") + cases.append((x, torch.empty(1, 1, out_height))) + self._run( + Resize(mode, align_corners), + cases, + DYNAMIC_HEIGHTS, + atol=1e-5, + rtol=1e-4, + ) + + def test_bilinear_antialias_resize(self): + self._check_dynamic_heights("bilinear", align_corners=False) + + def test_bicubic_antialias_resize(self): + self._check_dynamic_heights("bicubic", align_corners=True) + + def test_antialias_resize_with_uneven_scale_factor(self): + # 37 * 1.75 and 51 * 0.45 are not whole, so the scale the resize uses + # differs from output size / input size. + x = torch.randn(1, 3, 37, 51) + self._run( + Resize("bicubic", scale_factor=(1.75, 0.45)), + [(x, torch.empty(0))], + atol=1e-5, + rtol=1e-4, + ) + + def test_half_precision_antialias_resize(self): + for dtype in (torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + x = torch.randn(1, 3, 48, 32).to( + dtype=dtype, memory_format=torch.channels_last + ) + self._run( + Resize("bilinear"), + [(x, torch.empty(1, 1, 20))], + atol=1e-2, + rtol=1e-2, + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestAntialiasedResizeDecomposition(unittest.TestCase): + def test_keeps_memory_format(self): + # The program's output layout is fixed at export, so only a direct call + # shows the rewrite's own layout. + resize = CudaBackend.get_decomposition_table()[ + torch.ops.aten._upsample_bilinear2d_aa.default + ] + x = torch.randn(1, 3, 48, 32, device="cuda") + output = resize(x.to(memory_format=torch.channels_last), [20, 56], False) + self.assertTrue(output.is_contiguous(memory_format=torch.channels_last)) + + def test_align_corners_compiles_with_dynamic_sizes(self): + from torch._inductor.compile_fx import compile_fx + + model = Resize("bicubic", align_corners=True) + compiled = torch.compile( + model, + backend=functools.partial( + compile_fx, decompositions=CudaBackend.get_decomposition_table() + ), + dynamic=True, + fullgraph=True, + ) + for in_height, out_height in ((48, 20), (10, 1)): + x = torch.randn(1, 3, in_height, 32, device="cuda") + like = torch.empty(1, 1, out_height, device="cuda") + torch.testing.assert_close( + compiled(x, like), model(x, like), atol=1e-5, rtol=1e-4 + ) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 96f637a829c..637ff4036b1 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -237,6 +237,60 @@ def test_invalid_cuda_include_ptx_compile_spec(self): [CompileSpec(key="cuda_include_ptx", value=b"MAYBE")] ) + def test_invalid_data_name_compile_spec(self): + for name in ( + "", + "..", + ".hidden", + "../outside", + "a/b", + "a\\b", + "a/../../outside", + "C:model", + "model?", + "nul", + "COM1.txt", + "a" * 201, + # Non-ASCII letters that IGNORECASE alone would let through. + "\u212a" * 200, + "\u017fizes", + "\u0130d", + "\u0131d", + "-enc", + # One writer turns "encoder.ptd" into the same file as "encoder". + # The check ignores case, like the rest of the pattern. + "encoder.ptd", + "encoder.PTD", + ): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.generate_data_name_compile_spec(name) + with self.assertRaisesRegex(ValueError, "Invalid data name"): + CudaBackend.data_name_from_compile_specs( + [CompileSpec(key="data_name", value=name.encode())] + ) + # A backend that keeps its data in the .pte would ignore the name. + with patch.object(CudaBackend, "save_data_externally", return_value=False): + with self.assertRaisesRegex(ValueError, "takes no data name"): + CudaBackend.generate_data_name_compile_spec("encoder") + + def test_valid_data_name_compile_spec(self): + self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") + for name in ("encoder", "model_sm80", "v1.2-final", "a" * 200, "x.ptd.bak"): + with self.subTest(name=name): + spec = CudaBackend.generate_data_name_compile_spec(name) + self.assertEqual(CudaBackend.data_name_from_compile_specs([spec]), name) + + def test_methods_of_one_program_must_share_a_data_name(self): + def specs(name): + return [[CudaBackend.generate_data_name_compile_spec(name)]] + + with self.assertRaisesRegex(ValueError, "same data name"): + CudaBackend.preprocess_multimethod( + {"forward": [None], "second": [None]}, + {"forward": specs("first"), "second": specs("second")}, + ) + class TestCudaExport(unittest.TestCase): """Test CUDA export functionality for various operations using to_edge_transform_and_lower.""" diff --git a/backends/cuda/tests/test_merge_ptes.py b/backends/cuda/tests/test_merge_ptes.py index f0bb773070f..8ca8b8ecee5 100644 --- a/backends/cuda/tests/test_merge_ptes.py +++ b/backends/cuda/tests/test_merge_ptes.py @@ -58,6 +58,7 @@ def _write_artifact( fqn: str = "model.weight", weight_key: str = "cuda_fqn_weight:cuda:model.weight", ptx_compute: int | None = None, + data_name: str | None = None, ) -> CudaPteInput: if ptx_compute is None: ptx_compute = target_sm @@ -79,6 +80,8 @@ def _write_artifact( compile_specs.append( CompileSpec("cuda_include_ptx", b"ON" if ptx_compute else b"OFF") ) + if data_name is not None: + compile_specs.append(CompileSpec("data_name", data_name.encode())) delegate = BackendDelegate( id="CudaBackend", processed=BackendDelegateDataReference(DataLocation.INLINE, 0), @@ -148,8 +151,12 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: sm80.mkdir() sm120.mkdir() inputs = [ - self._write_artifact(sm80, 80, b"sm80-so", b"shared-weight"), - self._write_artifact(sm120, 120, b"sm120-so", b"shared-weight"), + self._write_artifact( + sm80, 80, b"sm80-so", b"shared-weight", data_name="model_sm80" + ), + self._write_artifact( + sm120, 120, b"sm120-so", b"shared-weight", data_name="model_sm120" + ), ] merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) @@ -176,14 +183,14 @@ def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: hashlib.sha256(b"sm120-so").hexdigest() + "_so_blob", }, ) - self.assertNotIn( - "cuda_include_ptx", + self.assertEqual( { spec.key for spec in merged.program.execution_plan[0] .delegates[0] .compile_specs }, + {"method_name"}, ) def test_uses_ptx_only_from_explicit_fallback(self) -> None: diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index fb2556cb3f8..6f1308b20c3 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -22,6 +22,7 @@ import os import tempfile import unittest +from typing import Optional import torch from executorch.backends.cuda.cuda_backend import CudaBackend @@ -51,14 +52,25 @@ 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, seed: int = 0, name: str = "model.pte") -> str: + def _lower( + self, + outdir: str, + seed: int = 0, + name: str = "model.pte", + data_name: Optional[str] = None, + module: Optional[torch.nn.Module] = None, + ) -> str: torch.manual_seed(seed) - module = ModuleWithConstants().eval().cuda() + if module is None: + module = ModuleWithConstants() + module = module.eval().cuda() inputs = (torch.randn(4, 64, device="cuda"),) exported = export(module, inputs) - spec = CudaBackend.generate_method_name_compile_spec("forward") + specs = [CudaBackend.generate_method_name_compile_spec("forward")] + if data_name is not None: + specs.append(CudaBackend.generate_data_name_compile_spec(data_name)) lowered = to_edge_transform_and_lower( - exported, partitioner=[CudaPartitioner([spec])] + exported, partitioner=[CudaPartitioner(specs)] ) program = lowered.to_executorch() path = os.path.join(outdir, name) @@ -177,6 +189,39 @@ def test_load_fails_with_another_programs_blob(self) -> None: ): program.load_method("forward") + def test_programs_with_their_own_data_names_share_a_folder(self) -> None: + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + for seed, name in ((0, "first"), (1, "second")): + self._lower(outdir, seed=seed, name=f"{name}.pte", data_name=name) + # A program with no constants still writes a data file. + self._lower( + outdir, name="third.pte", data_name="third", module=torch.nn.ReLU() + ) + self.assertEqual( + sorted(f for f in os.listdir(outdir) if f.endswith(".ptd")), + ["first.ptd", "second.ptd", "third.ptd"], + ) + for seed, name in ((0, "first"), (1, "second")): + # The same draws _lower made, so the same weights and input. + torch.manual_seed(seed) + module = ModuleWithConstants().eval().cuda() + value = torch.randn(4, 64, device="cuda") + method = ( + Runtime.get() + .load_program( + os.path.join(outdir, f"{name}.pte"), + data_path=os.path.join(outdir, f"{name}.ptd"), + ) + .load_method("forward") + ) + out = method.execute([value.cpu()]) + torch.testing.assert_close( + (out[0] if isinstance(out, (list, tuple)) else out).cpu(), + module(value).cpu(), + ) + 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/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index 518073101d9..a78a4b831fa 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -81,6 +81,25 @@ save_pte_program(exec_program, model_name, "./output_dir") ``` This generates `.pte` and `.ptd` files that can be executed on CUDA devices. +By default, the weights go to a data file named `aoti_cuda_blob.ptd`. To export several programs into one folder, give each program its own +data file name. This partitioner writes `encoder.ptd`: + +```python +partitioner = CudaPartitioner([ + CudaBackend.generate_method_name_compile_spec("forward"), + CudaBackend.generate_data_name_compile_spec("encoder"), +]) +``` + +The name can have up to 200 ASCII letters, digits, `_`, `-` and `.`. It cannot start with `.` or `-`, cannot end with `.ptd` +(the export adds it), and cannot be a Windows device name such as `NUL`. Load each program with its own data file: + +```python +from executorch.runtime import Runtime + +program = Runtime.get().load_program("encoder.pte", data_path="encoder.ptd") +``` + For a complete working example, see the [CUDA export script](https://github.com/pytorch/executorch/blob/main/examples/cuda/scripts/export.py). diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index 411cc262961..9af3a20a2da 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -170,9 +170,11 @@ 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. +without it. By default 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 then fails to load. To keep +several models in one directory, give each its own data file name with +`CudaBackend.generate_data_name_compile_spec`, as the +[CUDA backend page](../cuda/cuda-overview.md) shows, or save each model into its own directory. ### Keeping inputs and outputs on the GPU @@ -432,7 +434,8 @@ shows how to set the options that apply at run time, from C++ or from Python. Pass `external_engine_data="engines"` to `torch_tensorrt.save` to write the engines to `engines.ptd` instead of the `.pte`, and pass that file when you load the program. The Python loader takes one data file, so for a program that also uses the CUDA backend, pass - `external_engine_data="aoti_cuda_blob"` to share its file. + the name of the CUDA data file to share it: `external_engine_data="aoti_cuda_blob"` by default, + or the name given to `CudaBackend.generate_data_name_compile_spec`. - [**Optimization profiles.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#optimization-profiles) One engine can be tuned for several ranges of input shapes, for example one for decoding one token at a time and one for a long prompt. A C++ application picks the profile for each call From d5b0e87afd54ac87917dd3fb9933e1b462735e23 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:45:22 -0700 Subject: [PATCH 06/20] Refuse a CUDA graph replay when an input changes shape A CUDA graph replays the shapes it captured. Before each replay, the backend copied the new input into the captured input buffer, and it checked only that the byte count matched. An input can change shape without changing its byte count, for example from (2, 3, 4, 4) to (2, 4, 3, 4). Such a call replayed the captured shape and returned an output with the old shape and wrong values, with no error. The backend now records each input's shape at capture, and a replay with a different shape fails with InvalidArgument, the same error a different byte count already gave. test_output_stride_rearrange.py captures with (2, 3, 4, 4) inputs, then calls with (2, 4, 3, 4), for contiguous and channels-last outputs. Without this change those calls return the wrong shape. --- backends/cuda/runtime/cuda_backend.cpp | 15 +++++++----- backends/cuda/runtime/cuda_delegate_handle.h | 6 +++++ .../tests/test_output_stride_rearrange.py | 23 +++++++++++++++++++ 3 files changed, 38 insertions(+), 6 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index ce429a9a65a..b13605a907f 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -768,14 +768,14 @@ class ET_EXPERIMENTAL CudaBackend final // Copy new input data (GPU-resident) into static input buffers (D2D) for (size_t i = 0; i < n_inputs; i++) { auto* et_input = &(args[i]->toTensor()); + const auto sizes = et_input->sizes(); + const auto& captured = handle->cuda_graph_state.static_input_sizes[i]; ET_CHECK_OR_RETURN_ERROR( - et_input->nbytes() == - handle->cuda_graph_state.static_input_nbytes[i], + std::equal( + sizes.begin(), sizes.end(), captured.begin(), captured.end()), InvalidArgument, - "CUDA graph replay: input %zu size mismatch (expected %zu, got %zu)", - i, - handle->cuda_graph_state.static_input_nbytes[i], - et_input->nbytes()); + "CUDA graph replay: input %zu has a different shape than at capture", + i); ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( handle->cuda_graph_state.static_input_ptrs[i], et_input->const_data_ptr(), @@ -851,6 +851,7 @@ class ET_EXPERIMENTAL CudaBackend final state_->static_output_ptrs.clear(); state_->static_input_nbytes.clear(); state_->static_output_nbytes.clear(); + state_->static_input_sizes.clear(); // Same order as ~CudaGraphState: the exec depends on the graph. if (state_->graph_exec != nullptr) { (void)cudaGraphExecDestroy(state_->graph_exec); @@ -903,6 +904,8 @@ class ET_EXPERIMENTAL CudaBackend final // through the guard instead of leaking this allocation. handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); + handle->cuda_graph_state.static_input_sizes.emplace_back( + et_input->sizes().begin(), et_input->sizes().end()); ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy( static_ptr, diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index f42f0f2382f..605f50bf5d6 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -141,6 +141,9 @@ struct CudaGraphState { std::vector static_output_ptrs; std::vector static_input_nbytes; std::vector static_output_nbytes; + // An input with another shape of the same byte count would replay with the + // captured shape and return a wrong result with no error. + std::vector> static_input_sizes; // Allocations the graph makes and never frees -- the outputs AOTI allocates // while being captured. AutoFreeOnLaunch frees one launch's before the next @@ -227,6 +230,7 @@ struct CudaGraphState { static_output_ptrs.clear(); static_input_nbytes.clear(); static_output_nbytes.clear(); + static_input_sizes.clear(); } // Drops the captured graph so a new one is captured, for when memory the @@ -253,6 +257,7 @@ struct CudaGraphState { static_output_ptrs(std::move(other.static_output_ptrs)), static_input_nbytes(std::move(other.static_input_nbytes)), static_output_nbytes(std::move(other.static_output_nbytes)), + static_input_sizes(std::move(other.static_input_sizes)), graph_allocations(std::move(other.graph_allocations)) { other.graph = nullptr; other.graph_exec = nullptr; @@ -270,6 +275,7 @@ struct CudaGraphState { static_output_ptrs = std::move(other.static_output_ptrs); static_input_nbytes = std::move(other.static_input_nbytes); static_output_nbytes = std::move(other.static_output_nbytes); + static_input_sizes = std::move(other.static_input_sizes); graph_allocations = std::move(other.graph_allocations); other.graph = nullptr; diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index 431bb56a88f..908b58f4286 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -368,6 +368,29 @@ def test_strided_output_first_seen_at_capture(self): else: self._assert_matches_eager(module, x, out) + def test_input_shape_changes_after_capture(self): + torch.manual_seed(0) + # The last calls swap two dims, which keeps the byte count, so only the + # shape tells them apart from the captured input. + shapes = [(2, 3, 4, 4)] * 4 + [(2, 4, 3, 4)] * 2 + inputs = [torch.randn(shape) for shape in shapes] + dims = { + 1: torch.export.Dim("c", min=2, max=8), + 2: torch.export.Dim("h", min=2, max=8), + } + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[0].cuda(),), inputs, dynamic_shapes=(dims,) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + if i < 4: + self._assert_matches_eager(module, x, out) + else: + # InvalidArgument, from the replay's input shape check. + self.assertIn("error 0x12", str(out)) + def test_output_shape_changes_at_capture(self): torch.manual_seed(0) # The capture call has a larger batch than the warmup calls before it. From d7322be2d20ef7cabf4892b22918a3841e2c9648 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 21:45:23 -0700 Subject: [PATCH 07/20] List the CUDA index shims in the wheel's shared library check The wheel check names every kernel shim that libaoti_cuda_shims must define, so a library built without some of its CUDA sources fails the check. The mask index and index_put shims were missing from that list. --- .ci/scripts/wheel/test_shared_libraries.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index a1d8c791b83..17964dd103b 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -103,6 +103,8 @@ "aoti_torch_cuda_randint_low_out", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_guard_set_index", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( From b315e3f542586bbc825944c3853334e299d0e946 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 22:58:53 -0700 Subject: [PATCH 08/20] Keep Inductor's decompositions in the antialiased resize compile test The test gave Inductor only the CUDA backend's decomposition table, which replaced Inductor's own table. With torch 2.14 and CI set, Inductor then rejects the fallback for aten.arange and the test fails. Real lowering applies the backend's table at export and keeps Inductor's defaults, so the test now merges the two. --- backends/cuda/tests/test_antialiased_resize.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py index 963444f80e7..0415a057c0c 100644 --- a/backends/cuda/tests/test_antialiased_resize.py +++ b/backends/cuda/tests/test_antialiased_resize.py @@ -155,13 +155,17 @@ def test_keeps_memory_format(self): def test_align_corners_compiles_with_dynamic_sizes(self): from torch._inductor.compile_fx import compile_fx + from torch._inductor.decomposition import select_decomp_table model = Resize("bicubic", align_corners=True) + # Lowering keeps Inductor's own decompositions after the backend's. + decompositions = { + **select_decomp_table(), + **CudaBackend.get_decomposition_table(), + } compiled = torch.compile( model, - backend=functools.partial( - compile_fx, decompositions=CudaBackend.get_decomposition_table() - ), + backend=functools.partial(compile_fx, decompositions=decompositions), dynamic=True, fullgraph=True, ) From c9a926a78a5c5fd68b075aaacbe9da234d035fc3 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:59:47 -0700 Subject: [PATCH 09/20] Delete the antialiased resize torch.compile test test_align_corners_compiles_with_dynamic_sizes could not fail. It passed with no rewrite at all, since torch.compile then calls ATen's kernel, and it passed with the min and max in the align_corners ratio replaced by a plain if, since torch.compile treats an output size of 1 as a constant and compiles again. With torch 2.14 and CI set it also failed for a reason unrelated to the resize. test_bicubic_antialias_resize already runs an exported program at an output height of 1, which is where a plain if breaks. --- .../cuda/tests/test_antialiased_resize.py | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/backends/cuda/tests/test_antialiased_resize.py b/backends/cuda/tests/test_antialiased_resize.py index 0415a057c0c..be48e7fefe1 100644 --- a/backends/cuda/tests/test_antialiased_resize.py +++ b/backends/cuda/tests/test_antialiased_resize.py @@ -11,7 +11,6 @@ runtime, and compare with eager PyTorch. """ -import functools import os import tempfile import unittest @@ -152,26 +151,3 @@ def test_keeps_memory_format(self): x = torch.randn(1, 3, 48, 32, device="cuda") output = resize(x.to(memory_format=torch.channels_last), [20, 56], False) self.assertTrue(output.is_contiguous(memory_format=torch.channels_last)) - - def test_align_corners_compiles_with_dynamic_sizes(self): - from torch._inductor.compile_fx import compile_fx - from torch._inductor.decomposition import select_decomp_table - - model = Resize("bicubic", align_corners=True) - # Lowering keeps Inductor's own decompositions after the backend's. - decompositions = { - **select_decomp_table(), - **CudaBackend.get_decomposition_table(), - } - compiled = torch.compile( - model, - backend=functools.partial(compile_fx, decompositions=decompositions), - dynamic=True, - fullgraph=True, - ) - for in_height, out_height in ((48, 20), (10, 1)): - x = torch.randn(1, 3, in_height, 32, device="cuda") - like = torch.empty(1, 1, out_height, device="cuda") - torch.testing.assert_close( - compiled(x, like), model(x, like), atol=1e-5, rtol=1e-4 - ) From d9083763698bdbb7db720e12a8a499c5b5a04bf4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 29 Oct 2026 23:59:48 -0700 Subject: [PATCH 10/20] Say why the antialiased resize avoids a conditional The comment on the min and max in the align_corners ratio blamed torch.compile, which handles a conditional by compiling again. The real reason is export: an exported program keeps the branch it was traced with, so a plain if divides by zero at an output size of 1 when the program was traced with a larger one. --- backends/cuda/cuda_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 891fb81bd06..546d9c31b6f 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -582,8 +582,8 @@ def _antialias_taps( if align_corners or scale is None or scale <= 0: if align_corners: # ATen uses a ratio of 0 for an output size of 1. min and max give - # that without a conditional, which torch.compile cannot handle for - # dynamic sizes. + # that without a conditional: an exported program keeps the branch + # it was traced with, so a plain if would divide by zero there. num = (in_size - 1) * torch.sym_min(out_size - 1, 1) den = torch.sym_max(out_size - 1, 1) else: From 93879d50eeef96b2a101783aaca1215db78f634a Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 00:00:21 -0700 Subject: [PATCH 11/20] Refuse a hand-built data name in AOTI backends that keep data in the .pte Only generate_data_name_compile_spec refused a data name for a backend that keeps its data in the .pte, such as Metal. A data_name compile spec built by hand went through preprocess with no error and no name check, even for a name like "../../outside". data_name_from_compile_specs now refuses it too, and preprocess always calls it, so both paths behave the same. --- backends/aoti/aoti_backend.py | 22 ++++++++++++---------- backends/cuda/tests/test_cuda_export.py | 7 ++++++- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 0d0670b7318..d6b78d66000 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -390,11 +390,8 @@ def preprocess( named_data_store.add_named_data(so_blob_key, so_data, 1, None) # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file - external_tag = ( - cls.data_name_from_compile_specs(compile_specs) - if cls.save_data_externally() - else None - ) + data_name = cls.data_name_from_compile_specs(compile_specs) + external_tag = data_name if cls.save_data_externally() else None named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) @@ -435,11 +432,7 @@ def generate_data_name_compile_spec(cls, data_name: str) -> CompileSpec: Generate a CompileSpec that names the external data file, written as ``.ptd``, so several programs can share a folder. """ - if not cls.save_data_externally(): - raise ValueError( - f"{cls.__name__} keeps its data in the .pte file, so it takes no " - "data name" - ) + cls._check_saves_data_externally() return CompileSpec( COMPILE_SPEC_KEYS.DATA_NAME.value, _checked_data_name(data_name).encode("utf-8"), @@ -453,9 +446,18 @@ def data_name_from_compile_specs(cls, compile_specs: List[CompileSpec]) -> str: """ for spec in compile_specs: if spec.key == COMPILE_SPEC_KEYS.DATA_NAME.value: + cls._check_saves_data_externally() return _checked_data_name(spec.value.decode("utf-8")) return f"aoti_{cls.get_device_name()}_blob" + @classmethod + def _check_saves_data_externally(cls) -> None: + if not cls.save_data_externally(): + raise ValueError( + f"{cls.__name__} keeps its data in the .pte file, so it takes no " + "data name" + ) + @classmethod def method_name_from_compile_specs( cls, diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 637ff4036b1..58ba305ea72 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -269,10 +269,15 @@ def test_invalid_data_name_compile_spec(self): CudaBackend.data_name_from_compile_specs( [CompileSpec(key="data_name", value=name.encode())] ) - # A backend that keeps its data in the .pte would ignore the name. + # A backend that keeps its data in the .pte refuses a name, also one + # built by hand. with patch.object(CudaBackend, "save_data_externally", return_value=False): with self.assertRaisesRegex(ValueError, "takes no data name"): CudaBackend.generate_data_name_compile_spec("encoder") + with self.assertRaisesRegex(ValueError, "takes no data name"): + CudaBackend.data_name_from_compile_specs( + [CompileSpec(key="data_name", value=b"encoder")] + ) def test_valid_data_name_compile_spec(self): self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") From 4ae4daaec0d348c7b4418dadd2053e1732ba889e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 00:00:22 -0700 Subject: [PATCH 12/20] Warn that CUDA data names differing only in case can share a file Windows and macOS usually treat Encoder.ptd and encoder.ptd as one file, so two programs exported into one folder with those data names overwrite each other's weights, and the first one then fails to load. The name check allows both, as both are valid names elsewhere, so the docs now say to pick names that differ in more than case. --- docs/source/backends/cuda/cuda-overview.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index a78a4b831fa..56dccf5813b 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -92,7 +92,9 @@ partitioner = CudaPartitioner([ ``` The name can have up to 200 ASCII letters, digits, `_`, `-` and `.`. It cannot start with `.` or `-`, cannot end with `.ptd` -(the export adds it), and cannot be a Windows device name such as `NUL`. Load each program with its own data file: +(the export adds it), and cannot be a Windows device name such as `NUL`. Windows and macOS usually treat names that differ +only in case, such as `Encoder` and `encoder`, as the same file, so give programs that share a folder names that differ in +more than case. Load each program with its own data file: ```python from executorch.runtime import Runtime From 92c4baac2bf59eb37eda1a2e86747796c77ef6a4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:34:21 -0700 Subject: [PATCH 13/20] Correct comments on the CUDA index shims and set_output_data_ptr Some comments added with the CUDA mask indexing change say the wrong thing: - index.h described accumulation as out[indices] += values. With a repeated index, += keeps only one of the values, while index_put_ with accumulate adds every value, and so does the shim. - index.cu said CUB counts items in an int for the index_put sort. That is true for the mask path, but the radix sort takes a wider count. The limit there comes from the int the code passes to the sort. - index.cu said the accumulation types are ATen's. That holds for half and bfloat16, which add in float. Narrower integers add in int64 here, which gives the same wrapped result. - method.h named the output index i in the set_output_data_ptr docs, but the parameter is output_idx. Only comments change. --- backends/cuda/runtime/shims/index.cu | 5 +++-- backends/cuda/runtime/shims/index.h | 4 ++-- runtime/executor/method.h | 3 ++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu index 97335a75ea0..c7838946a8a 100644 --- a/backends/cuda/runtime/shims/index.cu +++ b/backends/cuda/runtime/shims/index.cu @@ -781,7 +781,7 @@ Error put( // One key per index point, so the scratch and the sort grow with the // number of indices, not with the elements they write. const int64_t points = plan.broadcast_numel; - // CUB counts items in an int on CUDA 12. + // The sort below takes the count as an int. ET_CHECK_OR_RETURN_ERROR( points <= INT_MAX, NotSupported, @@ -1017,7 +1017,8 @@ AOTITorchError aoti_torch_index_put_out( return Error::NotSupported; } } - // The accumulation types are ATen's. + // Half and bfloat16 add in float, as ATen does. Narrower integers add in + // int64, which wraps to the same result as adding in their own type. switch (out->dtype()) { case ScalarType::Float: return put( diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h index 3476eaa8c88..25f015cf4c2 100644 --- a/backends/cuda/runtime/shims/index.h +++ b/backends/cuda/runtime/shims/index.h @@ -45,8 +45,8 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( Tensor** ret0); /** - * out[indices] = values, or out[indices] += values when accumulate is set, - * with the semantics of aten::index_put_. + * Writes values into out at indices, with the semantics of aten::index_put_. + * When accumulate is set, every value is added, also when an index repeats. * * Inductor calls this fallback, with out and self the same tensor, for an * index that is a mask, for accumulation into a dtype its generated code diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 501569285f9..9432e2cc891 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -155,7 +155,8 @@ class Method final { * @param[in] buffer The block of memory to point the specified tensor at. * * @param[in] size the length of buffer in bytes, must be >= the largest - * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. + * nbytes the output can have, + * `MethodMeta::output_tensor_meta(output_idx)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer From 1819fdfa9023417772991ecf6257dfb0750018c4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 06:51:06 -0700 Subject: [PATCH 14/20] Check the AOTI data name before compiling A backend that keeps its data in the .pte, such as Metal, refuses a data_name compile spec. preprocess only checked it after moving the program to the device and compiling it with AOTInductor, which took about 18 seconds and left the compiled library on disk. It now checks the name first. The data name test calls the shared preprocess with such a spec, so moving the check back behind the compile, or dropping it, fails the test. --- backends/aoti/aoti_backend.py | 4 +++- backends/cuda/tests/test_cuda_export.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index d6b78d66000..87a15498de7 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -286,6 +286,9 @@ def preprocess( Preprocess the edge program and compile it using AOTInductor. Weights are always separated from the SO file. """ + # First, so that a backend keeping its data in the .pte refuses a data + # name before anything is compiled. + data_name = cls.data_name_from_compile_specs(compile_specs) device_name = cls.get_device_name() decomposition_table = cls.get_decomposition_table() options = cls.get_aoti_compile_options(compile_specs) @@ -390,7 +393,6 @@ def preprocess( named_data_store.add_named_data(so_blob_key, so_data, 1, None) # Determine whether to save named data externally based on backend setting # External: save to separate .ptd file, otherwise merge with .pte file - data_name = cls.data_name_from_compile_specs(compile_specs) external_tag = data_name if cls.save_data_externally() else None named_data_store.add_named_data(weights_blob_key, blob_data, 1, external_tag) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 58ba305ea72..1f61ad38836 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -14,6 +14,7 @@ from unittest.mock import patch import torch +from executorch.backends.aoti.aoti_backend import AotiBackend from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.examples.models.toy_model import SdpaModule @@ -278,6 +279,11 @@ def test_invalid_data_name_compile_spec(self): CudaBackend.data_name_from_compile_specs( [CompileSpec(key="data_name", value=b"encoder")] ) + # The shared preprocess, which Metal uses, refuses it before compiling. + with self.assertRaisesRegex(ValueError, "takes no data name"): + AotiBackend.preprocess( + None, [CompileSpec(key="data_name", value=b"encoder")] + ) def test_valid_data_name_compile_spec(self): self.assertEqual(CudaBackend.data_name_from_compile_specs([]), "aoti_cuda_blob") From cdb488e94b5bd6685e80adca727fe7207347d196 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 09:23:16 -0700 Subject: [PATCH 15/20] Restore release/1.6 Windows Voxtral CI files that the #24299 cherry-pick reverted The bot pick took the older main versions of these two files, which undid the release/1.6 fix from #24289. --- .ci/scripts/test_model_e2e_windows.ps1 | 6 +++--- .github/workflows/cuda-windows.yml | 16 ++++++++++------ 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/.ci/scripts/test_model_e2e_windows.ps1 b/.ci/scripts/test_model_e2e_windows.ps1 index fd01194fa51..7b15db1f3c9 100644 --- a/.ci/scripts/test_model_e2e_windows.ps1 +++ b/.ci/scripts/test_model_e2e_windows.ps1 @@ -89,12 +89,12 @@ switch ($HfModel) { $runnerTarget = "voxtral_realtime_runner" $runnerPath = "voxtral_realtime" $runnerPreset = "voxtral-realtime-cuda" - $expectedOutput = "Quilter" + $expectedOutput = "Phoebe" $preprocessor = "preprocessor.pte" $tokenizerUrl = "" $tokenizerFile = "tekken.json" - $audioUrl = "https://github.com/voxserv/audio_quality_testing_samples/raw/refs/heads/master/testaudio/16000/test01_20s.wav" - $audioFile = "poem.wav" + $audioUrl = "https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav" + $audioFile = "test_audio.wav" } "facebook/dinov2-small-imagenet1k-1-layer" { $runnerTarget = "dinov2_runner" diff --git a/.github/workflows/cuda-windows.yml b/.github/workflows/cuda-windows.yml index e145fc7daaf..83f34c4d114 100644 --- a/.github/workflows/cuda-windows.yml +++ b/.github/workflows/cuda-windows.yml @@ -17,6 +17,8 @@ on: - backends/cuda/** - backends/aoti/** - extension/cuda/** + - .ci/scripts/export_model_artifact.sh + - .ci/scripts/test_model_e2e_windows.ps1 # Only unittest-cuda-runtime-windows runs for these alone; the other jobs # filter on their own paths and stay skipped. - .ci/scripts/test_cuda_runtime_windows.ps1 @@ -61,6 +63,8 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || contains(needs.changed-files.outputs.changed-files, 'extension/cuda') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda-windows.yml') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e_windows.ps1') || needs.run-decision.outputs.is-full-run == 'true' ) uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main @@ -92,6 +96,7 @@ jobs: - model_repo: "mistralai" model_name: "Voxtral-Mini-4B-Realtime-2602" quant: "quantized-int4-tile-packed" + mode: "vr-offline" - model_repo: "facebook" model_name: "dinov2-small-imagenet1k-1-layer" quant: "non-quantized" @@ -147,11 +152,7 @@ jobs: echo "::endgroup::" fi - VR_MODE="" - if [ "${{ matrix.model_name }}" = "Voxtral-Mini-4B-Realtime-2602" ]; then - VR_MODE="vr-offline" - fi - source .ci/scripts/export_model_artifact.sh cuda-windows "${{ matrix.model_repo }}/${{ matrix.model_name }}" "${{ matrix.quant }}" "${RUNNER_ARTIFACT_DIR}" "${VR_MODE}" + source .ci/scripts/export_model_artifact.sh cuda-windows "${{ matrix.model_repo }}/${{ matrix.model_name }}" "${{ matrix.quant }}" "${RUNNER_ARTIFACT_DIR}" "${{ matrix.mode }}" test-model-cuda-windows-e2e: name: test-model-cuda-windows-e2e @@ -168,6 +169,8 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || contains(needs.changed-files.outputs.changed-files, 'extension/cuda') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda-windows.yml') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || + contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e_windows.ps1') || needs.run-decision.outputs.is-full-run == 'true' ) uses: pytorch/test-infra/.github/workflows/windows_job.yml@main @@ -195,6 +198,7 @@ jobs: - model_repo: "mistralai" model_name: "Voxtral-Mini-4B-Realtime-2602" quant: "quantized-int4-tile-packed" + mode: "vr-offline" - model_repo: "facebook" model_name: "dinov2-small-imagenet1k-1-layer" quant: "non-quantized" @@ -225,7 +229,7 @@ jobs: throw 'RUNNER_ARTIFACT_DIR is empty. Ensure download-artifact is configured for windows_job.yml.' } - .ci/scripts/test_model_e2e_windows.ps1 -Device cuda-windows -HfModel '${{ matrix.model_repo }}/${{ matrix.model_name }}' -QuantName '${{ matrix.quant }}' -ModelDir \$artifactDir -ExpectedCudaVersion '13.2' + .ci/scripts/test_model_e2e_windows.ps1 -Device cuda-windows -HfModel '${{ matrix.model_repo }}/${{ matrix.model_name }}' -QuantName '${{ matrix.quant }}' -ModelDir \$artifactDir -ExpectedCudaVersion '13.2' -Mode '${{ matrix.mode }}' }" # The counterpart of cuda.yml's unittest-cuda-runtime on Windows. The Windows GPU runners From 07002ff26b6111fa4a500fd6498bc0447a48d52f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 30 Oct 2026 11:26:42 -0700 Subject: [PATCH 16/20] Regenerate aoti_cuda_shims.lib with aoti_torch_dtype_element_size Two release/1.6 changes each regenerated this import library from the same base. This one adds the CUDA index shims, aoti_torch_as_strided and aoti_torch_get_size. The other adds aoti_torch_dtype_element_size. Git cannot merge a binary file, so whichever lands second would conflict, and keeping either side alone drops names the other needs. List the union of both sets, regenerated per aoti_cuda_shims.lib.md. The archive has no member timestamps, and rebuilding it gives the same bytes. --- backends/cuda/runtime/aoti_cuda_shims.lib | Bin 47046 -> 47886 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 3294d1ec2629c60f221598224c005977bdd04da4..aa97264e523c2155b7d27a28f75db138096369e8 100644 GIT binary patch delta 1147 zcmaLTX>3eU6ae6JOGj0wbV^m#UXh^hy_D&qF)E0yN?Js$4XP308(Dx1=(}s z1?CqR^uwDj@Ka#~&B*6SVzEIj1sv(w5l(?caRLK(8uZ;OJMA*logFCxpY7IE@6~&v z=+v$Rfj1o*ElJK0A_mtv#*t&sT!DV4PI+-$H@!3r&jt^r?0u;M50q)t=j3Srato%E z1yNe5Rp5pSoyIousV^!cDY|@)z$BMOH5HuKbI^hxss_=|%4C7x4(XKb;;5tALcIrd zK{snW0-I}0gw>kBDz`@OYxsYVCyX+VOcS^Wb?SC=r1#h`8s*`fu)4oFo_Y&rh6K{f zV;ue1pi@d6Pf*x20v|LCAb-7$i@Z9GY2@m(6X8_Rv_PP@SyS7HZV9EeCuZ@DT6IcG z=F0wU!!gf1h-@wM1m?DD)YZzB?@pW1XO5srZ5*|o(beW|bcW(XvzZ2V@XiOE)78#h z=R@gXXOh4NzP~f3T{PoiUof3KA1kosl1|@!yybXWdXW!w`ZaoVDMNU!m@(FG!M2Il h@mE(VCUnP=&%a!VL=so5qF;hx#juzuw7Gj@z;9LXdrtrW delta 1043 zcmXZZTS!z<6b9h62N`C(oJ`X$idje`#5t4Nq?MC&qr%cAqUeTML|!UOQM`^$nl+-C zi-u9Nv_wg?u~`m94;B$V)I(Na&_hK%bc1>-%1Bzf?Sbz+th4so|K6wVr)lE5sXfEB zI>Vm$|H%slfD(Y1UjPdpD{(^rQo>nfi4#Cd3Rjj*07%IJWkoT7WNcSb901ZPo07g9 zKw2BGI6_w zY_-{D|1CyF#!z5RiP%rYeK?seXh(6U zps^CWp!m~1LGDt!z%OTP^r_4#u(&jqa?1^Ycgt4NXa!f+lt*J;sF~s_(*!nD=>PW+ zOX+!auE3lsopNe<`egMY9Cumhaor|?iM0m3O4%($5?*#Ek-s5ZptWA7r_OE*-g&BJdf`cIRG+p=6&;AbNEQTg*N4 zd{HhA9MMB~FgGx>Y@4q~p)p4r>0;lV*G z*7k?d`J0;se&{o3x}PT;yE(r@3#ZcqTzt}RsNa%35P_LP8tuBp7aknc>0=p3u;0?K zSPP}ZVLn*%kU@1LJis!n(VN@cw{yfG;|TXn-j2YYz#IzPS Date: Sat, 31 Oct 2026 09:02:27 -0700 Subject: [PATCH 17/20] 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 0ec89ad61fb9c34dfff5c4c711e5971b3cd7807e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 10:57:37 -0700 Subject: [PATCH 18/20] 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 187f1e3796a8d682a0fc3bf328d2ce02307150a4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:39:25 -0700 Subject: [PATCH 19/20] 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 f714c0871eeff40d1876da0956429482f1bc4eb4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 31 Oct 2026 13:37:57 -0700 Subject: [PATCH 20/20] 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):