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 c0d61c61100..9f48cd08607 100644 Binary files a/backends/cuda/runtime/aoti_cuda_shims.lib and b/backends/cuda/runtime/aoti_cuda_shims.lib differ 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 9d6e0f747fd..e8745524512 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -172,6 +172,8 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): """Return annotation pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ RemoveRedundancy, + # Retraced autocast regions may emit ops that the passes below decompose. + DecomposeWrapWithAutocast, RecomposePixelUnshuffle, RecomposeRmsNorm, ReplaceArangeArgs, @@ -194,7 +196,6 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): DecomposeTriu, DecomposeTrunc, DecomposeVar, - DecomposeWrapWithAutocast, DecomposeEinsum, DecomposeEmpty, DecomposeExpM1, @@ -218,6 +219,7 @@ def get_annotation_passes(cls, convert_linear_to_conv2d: bool = False): def get_export_passes(cls): """Return export pipeline pass classes. Override in subclasses to add backend-specific passes.""" passes = [ + DecomposeWrapWithAutocast, ConvertComplexStft, DecomposeAsStrided, DecomposeBinaryAlpha, @@ -239,7 +241,6 @@ def get_export_passes(cls): # We keep div and floor operations in floating-point to maintain precision. # This pass is needed before to_edge pipeline to avoid mixed type for div operator with RemoveMixedTypeOperators pass. DecomposeFloorDivide, - DecomposeWrapWithAutocast, ConvertSquareToPow, LiftConstantScalarOperands, InsertReshapeForReduceOps, diff --git a/backends/qualcomm/_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 f29ac21b7f6..640d6fbf478 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -14,7 +14,10 @@ PER_CHANNEL_GROUP_ENCODING, ) from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader -from executorch.backends.qualcomm.builders.utils import is_parameter +from executorch.backends.qualcomm.builders.utils import ( + is_mutable_buffer_input, + is_parameter, +) from executorch.backends.qualcomm.qnn_preprocess import QnnBackend from executorch.backends.qualcomm.serialization.qc_schema import ( HtpArch, @@ -57,6 +60,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, ) @@ -122,7 +127,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 @@ -169,6 +174,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. @@ -253,7 +270,14 @@ def generate_partitions( edge_program.graph_module, op_support=self.op_support_checker, ) - return self._drop_orphaned_parameter_dequants(partitions, edge_program) + partitions = self._drop_orphaned_parameter_dequants(partitions, edge_program) + # LPAI places its CPU fallback boundaries before partitioning, and GPU + # accepts graphs without inputs, so only HTP partitions are dropped. + if self.backend == QnnExecuTorchBackendType.kHtpBackend: + partitions = self._drop_partitions_without_runtime_inputs( + partitions, edge_program + ) + return partitions @staticmethod def _drop_orphaned_parameter_dequants( @@ -286,6 +310,32 @@ def _drop_orphaned_parameter_dequants( ) return [partition for partition in partitions if partition.nodes] + @staticmethod + def _drop_partitions_without_runtime_inputs( + partitions: List[Partition], edge_program: torch.export.ExportedProgram + ) -> List[Partition]: + """ + HTP refuses a graph that reads nothing at runtime, e.g. a subgraph fed + only by constants once a CPU op splits it from the rest. Such a + partition stays on CPU. + """ + + def is_runtime_input(node: torch.fx.Node) -> bool: + return not is_parameter(node, edge_program) or is_mutable_buffer_input( + node, edge_program + ) + + return [ + partition + for partition in partitions + if any( + is_runtime_input(arg) + for node in partition.nodes + for arg in node.all_input_nodes + if arg not in partition.nodes + ) + ] + def tag_nodes( self, partitions: List[Partition], edge_program: torch.export.ExportedProgram ) -> None: @@ -298,6 +348,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 d8c5c32c103..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,151 @@ 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 ): 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 01c74110da6..68681688457 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -3049,6 +3049,19 @@ def forward(self, x, y): self.lower_module_and_test_output(high_dim_module, sample_inputs) + def test_vulkan_backend_high_dimensional_constant(self): + class HighDimConstantModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("c", torch.rand(1, 2, 3, 4, 5)) + + def forward(self, x): + return x * self.c + + sample_inputs = (torch.rand(size=(1, 2, 3, 4, 5), dtype=torch.float32),) + + self.lower_module_and_test_output(HighDimConstantModule(), sample_inputs) + def test_vulkan_backend_torchao_wo_quantized_linear(self): in_features = 1024 out_features = 512 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 e36f075a7b0..8e465561b50 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -2537,10 +2537,7 @@ def forward(self, x): model, (test_inputs[0],), dynamic_shapes=dynamic_shapes ) - edge_program = to_edge( - exported_program, - compile_config=exir.EdgeCompileConfig(_check_ir_validity=False), - ) + edge_program = to_edge(exported_program) et_program = edge_program.to_executorch() program_buffer = et_program.buffer et_module = _load_for_executorch_from_buffer(program_buffer) diff --git a/exir/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 39c649c822d..f4e5f17e974 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -450,14 +450,11 @@ def make_val( try: assert self.fake_tensor_mode is not None - # TODO we should allocate static shapes - # for param/buffer values - if isinstance(x, torch.nn.Parameter): - fake_tensor = self.fake_tensor_mode.from_tensor( - x, static_shapes=True - ) - else: - fake_tensor = self.fake_tensor_mode.from_tensor(x) + # A real tensor here is constant data, so its sizes + # are known and must not become new symbols. + fake_tensor = self.fake_tensor_mode.from_tensor( + x, static_shapes=True + ) except UnsupportedFakeTensorException: # TODO: This is just a workaround to get over the # x.as_subclass error diff --git a/exir/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 168485c2905..e33cb1b68ba 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -502,6 +502,28 @@ def test_export_pass_preserves_symbolic_input_metadata(self) -> None: self.assertEqual(self._symbolic_input_shape(new_input), original_snapshot) + def test_export_pass_keeps_static_shapes_for_tensor_constants(self) -> None: + class AddConstant(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.constant = torch.randn(9, 2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.constant + + graph_module = export(AddConstant(), (torch.randn(9, 2),), strict=True).module( + check_guards=False + ) + (constant,) = [n for n in graph_module.graph.nodes if n.op == "get_attr"] + # A constant made by an earlier pass, e.g. constant folding, is a + # tensor the fake mode has not seen before. + setattr(graph_module, str(constant.target), torch.randn(9, 2)) + + new_graph_module = ExportPass()(graph_module).graph_module + (constant,) = [n for n in new_graph_module.graph.nodes if n.op == "get_attr"] + + self.assertEqual([type(dim) for dim in constant.meta["val"].shape], [int, int]) + def test_export_pass_tracks_replayed_unbacked_symbols(self) -> None: def data_dependent_slice( tensor: torch.Tensor, length: torch.Tensor diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index eb917f55490..81c5298aab1 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2915,6 +2915,69 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: module = new_ep.module() self.assertFalse(torch.equal(module(x), module(x))) + def test_constant_prop_pass_keeps_tensor_written_in_place(self) -> None: + """ + torch.export keeps in-place ops. A tensor or buffer that one writes, + directly, through a view, as an out= argument, in a list or inside a + no_grad or autocast block, has to stay a runtime tensor, or the folded + value is wrong and decomposing can fail. + """ + + class WrittenInPlace(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + for name in "bcdef": + self.register_buffer(name, torch.zeros(6)) + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + t = torch.zeros(2, 6) + rows = t.unbind(0) + rows[0].add_(x) + out = torch.zeros(6) + torch.mul(x, 2, out=out) + listed = torch.zeros(6) + torch._foreach_add_([listed], [x]) + self.b[1:3].copy_(x[:2]) + no_grad = torch.zeros(6) + with torch.no_grad(): + no_grad[1:3] = x[:2] + self.c[1:3].copy_(x[:2]) + autocast = torch.zeros(6) + with torch.autocast("cpu", enabled=False): + autocast[2:4] = x[:2] + nested = torch.zeros(6) + with torch.no_grad(): + with torch.autocast("cpu", enabled=False): + nested[1:3] = x[:2] + torch.atleast_2d(self.d)[0].add_(x) + torch.meshgrid(self.e, x[:1], indexing="ij")[0][:, 0].add_(x) + torch.broadcast_tensors(self.f, x.view(1, 6))[0][0].add_(x) + written = rows[0] + rows[1] + out + listed + no_grad + autocast + nested + buffers = self.b + self.c + self.d + self.e + self.f + return written + buffers * 2, self.b + + x = torch.randn(6) + ep = constant_prop_pass(export(WrittenInPlace(), (x,), strict=False)) + self.assertEqual(len(ep.constants), 0) + torch.testing.assert_close(ep.module()(x), WrittenInPlace()(x)) + + # A saved and loaded program no longer shares storage between views. + class WrittenThroughSlice(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("b", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.b[:2].copy_(x) + return self.b * 2 + + x = torch.randn(2) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "written.pt2") + torch.export.save(export(WrittenThroughSlice(), (x,), strict=False), path) + ep = constant_prop_pass(torch.export.load(path)) + torch.testing.assert_close(ep.module()(x), WrittenThroughSlice()(x)) + def test_constant_prop_pass_detaches_folded_parameters(self) -> None: class M(torch.nn.Module): def __init__(self): 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 36e79f0dc43..38bc6440c80 100644 --- a/kernels/aten/functions.yaml +++ b/kernels/aten/functions.yaml @@ -133,6 +133,8 @@ - op: cosh.out +- op: cummin.out + - op: cumsum.out - op: dequantize.self_out diff --git a/kernels/portable/cpu/op_cummin.cpp b/kernels/portable/cpu/op_cummin.cpp new file mode 100644 index 00000000000..4e7bb44f47a --- /dev/null +++ b/kernels/portable/cpu/op_cummin.cpp @@ -0,0 +1,94 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using Tensor = executorch::aten::Tensor; +using ScalarType = executorch::aten::ScalarType; + +/** + * Computes the running minimum of self along dim, and the index of that + * minimum. Like PyTorch, a tie moves the index to the later element and a NaN + * becomes the minimum from where it appears on. + */ +std::tuple cummin_out( + KernelRuntimeContext& ctx, + const Tensor& self, + int64_t dim, + Tensor& values, + Tensor& indices) { + std::tuple ret_val(values, indices); + + ET_KERNEL_CHECK( + ctx, dim_is_valid(dim, nonzero_dim(self)), InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dtype(self, values), InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, indices.scalar_type() == ScalarType::Long, InvalidArgument, ret_val); + ET_KERNEL_CHECK( + ctx, + tensor_is_default_dim_order(self) && + tensor_is_default_dim_order(values) && + tensor_is_default_dim_order(indices), + InvalidArgument, + ret_val); + ET_KERNEL_CHECK( + ctx, + resize_tensor(values, self.sizes()) == Error::Ok && + resize_tensor(indices, self.sizes()) == Error::Ok, + InvalidArgument, + ret_val); + + if (self.numel() == 0) { + return ret_val; + } + + dim = dim < 0 ? dim + nonzero_dim(self) : dim; + const size_t dim_size = nonempty_size(self, dim); + const size_t leading_dims = getLeadingDims(self, dim); + const size_t trailing_dims = self.dim() == 0 ? 1 : getTrailingDims(self, dim); + + ET_SWITCH_REALHBBF16_TYPES( + self.scalar_type(), ctx, "cummin.out", CTYPE, [&]() { + const CTYPE* in_data = self.const_data_ptr(); + CTYPE* values_data = values.mutable_data_ptr(); + int64_t* indices_data = indices.mutable_data_ptr(); + + for (size_t i = 0; i < leading_dims; i++) { + const size_t start = i * dim_size * trailing_dims; + for (size_t j = 0; j < dim_size; j++) { + const size_t cur = start + j * trailing_dims; + const size_t prev = cur - trailing_dims; + for (size_t k = 0; k < trailing_dims; k++) { + const CTYPE v = in_data[cur + k]; + if (j == 0 || utils::isnan_override(v) || + (!utils::isnan_override(values_data[prev + k]) && + v <= values_data[prev + k])) { + values_data[cur + k] = v; + indices_data[cur + k] = static_cast(j); + } else { + values_data[cur + k] = values_data[prev + k]; + indices_data[cur + k] = indices_data[prev + k]; + } + } + } + } + }); + + return ret_val; +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index 97f60f75c02..e67976fac40 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -323,6 +323,11 @@ - arg_meta: null kernel_name: torch::executor::cosh_out +- op: cummin.out + kernels: + - arg_meta: null + kernel_name: torch::executor::cummin_out + - op: cumsum.out kernels: - arg_meta: null diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index 86f3aa07319..029c9252998 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -197,6 +197,7 @@ set(all_test_sources "op_copy_test.cpp" "op_cos_test.cpp" "op_cosh_test.cpp" + "op_cummin_test.cpp" "op_cumsum_test.cpp" "op_detach_copy_test.cpp" "op_diagonal_copy_test.cpp" diff --git a/kernels/test/op_cummin_test.cpp b/kernels/test/op_cummin_test.cpp new file mode 100644 index 00000000000..7ba11762ccb --- /dev/null +++ b/kernels/test/op_cummin_test.cpp @@ -0,0 +1,102 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include // Declares the operator +#include +#include +#include +#include + +using namespace executorch::aten; +using torch::executor::testing::TensorFactory; + +class OpCumminOutTest : public OperatorTest { + protected: + std::tuple op_cummin_out( + const Tensor& self, + int64_t dim, + Tensor& values, + Tensor& indices) { + return torch::executor::aten::cummin_outf( + context_, self, dim, values, indices); + } + + template + void test_dtype() { + TensorFactory tf; + TensorFactory tf_long; + // clang-format off + Tensor in = tf.make({2, 4}, { + 3, 1, 2, 1, + 0, 5, 0, 4}); + // clang-format on + Tensor values = tf.zeros({2, 4}); + Tensor indices = tf_long.zeros({2, 4}); + + op_cummin_out(in, /*dim=*/1, values, indices); + // A tie moves the index to the later element, as in PyTorch. + EXPECT_TENSOR_EQ(values, tf.make({2, 4}, {3, 1, 1, 1, 0, 0, 0, 0})); + EXPECT_TENSOR_EQ(indices, tf_long.make({2, 4}, {0, 1, 1, 3, 0, 0, 2, 2})); + + op_cummin_out(in, /*dim=*/-2, values, indices); + EXPECT_TENSOR_EQ(values, tf.make({2, 4}, {3, 1, 2, 1, 0, 1, 0, 1})); + EXPECT_TENSOR_EQ(indices, tf_long.make({2, 4}, {0, 0, 0, 0, 1, 0, 1, 0})); + } +}; + +TEST_F(OpCumminOutTest, AllRealDtypesSupported) { +#define TEST_ENTRY(ctype, dtype) test_dtype(); + ET_FORALL_REALHBF16_TYPES(TEST_ENTRY); +#undef TEST_ENTRY +} + +TEST_F(OpCumminOutTest, NanBecomesTheMinimum) { + TensorFactory tf; + TensorFactory tf_long; + const float nan = NAN; + Tensor in = tf.make({5}, {2, nan, 1, nan, 0}); + Tensor values = tf.zeros({5}); + Tensor indices = tf_long.zeros({5}); + + op_cummin_out(in, /*dim=*/0, values, indices); + EXPECT_TENSOR_CLOSE(values, tf.make({5}, {2, nan, nan, nan, nan})); + EXPECT_TENSOR_EQ(indices, tf_long.make({5}, {0, 1, 1, 3, 3})); +} + +TEST_F(OpCumminOutTest, Bool) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.make({4}, {true, false, true, false}); + Tensor values = tf.zeros({4}); + Tensor indices = tf_long.zeros({4}); + + op_cummin_out(in, /*dim=*/0, values, indices); + EXPECT_TENSOR_EQ(values, tf.make({4}, {true, false, false, false})); + EXPECT_TENSOR_EQ(indices, tf_long.make({4}, {0, 1, 1, 3})); +} + +TEST_F(OpCumminOutTest, ZeroDim) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.make({}, {2}); + Tensor values = tf.zeros({}); + Tensor indices = tf_long.ones({}); + + op_cummin_out(in, /*dim=*/-1, values, indices); + EXPECT_TENSOR_EQ(values, in); + EXPECT_TENSOR_EQ(indices, tf_long.zeros({})); +} + +TEST_F(OpCumminOutTest, InvalidDimDies) { + TensorFactory tf; + TensorFactory tf_long; + Tensor in = tf.ones({2, 3}); + Tensor values = tf.zeros({2, 3}); + Tensor indices = tf_long.zeros({2, 3}); + ET_EXPECT_KERNEL_FAILURE(context_, op_cummin_out(in, 2, values, indices)); +} diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 7555c14c4a8..8289a1ddd5f 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -248,6 +248,7 @@ def define_common_targets(): _common_op_test("op_copy_test", ["aten", "portable"]) _common_op_test("op_cos_test", ["aten", "portable"]) _common_op_test("op_cosh_test", ["aten", "portable"]) + _common_op_test("op_cummin_test", ["aten", "portable"]) _common_op_test("op_cumsum_test", ["aten", "portable"]) _common_op_test("op_detach_copy_test", ["aten", "portable"]) _common_op_test("op_diagonal_copy_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index 82d0c40064e..b0e7f7e7065 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -112,6 +112,7 @@ PORTABLE_KERNELS_SRCS = [ "kernels/portable/cpu/op_copy.cpp", "kernels/portable/cpu/op_cos.cpp", "kernels/portable/cpu/op_cosh.cpp", + "kernels/portable/cpu/op_cummin.cpp", "kernels/portable/cpu/op_cumsum.cpp", "kernels/portable/cpu/op_detach_copy.cpp", "kernels/portable/cpu/op_diagonal_copy.cpp", diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index f3d5739a862..733eed5a780 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -502,6 +502,12 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/pattern:pattern", ], ), + op_target( + name = "op_cummin", + deps = [ + "//executorch/kernels/portable/cpu/util:math_util", + ], + ), op_target( name = "op_cumsum", deps = [