diff --git a/backends/xnnpack/_passes/BUCK b/backends/xnnpack/_passes/BUCK index 2edf294fa4e..bf3239953a7 100644 --- a/backends/xnnpack/_passes/BUCK +++ b/backends/xnnpack/_passes/BUCK @@ -18,6 +18,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", "//executorch/exir/passes:const_prop_pass", + "//executorch/exir/passes:constant_prop_pass", "//executorch/exir/passes:memory_format_ops_pass", "//executorch/exir/program:program", "//executorch/backends/transforms:utils", diff --git a/backends/xnnpack/_passes/fold_constant_subgraphs.py b/backends/xnnpack/_passes/fold_constant_subgraphs.py new file mode 100644 index 00000000000..fa60d168a52 --- /dev/null +++ b/backends/xnnpack/_passes/fold_constant_subgraphs.py @@ -0,0 +1,155 @@ +# Copyright 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. + +import operator + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.passes.constant_prop_pass import constant_prop_pass +from torch.export import ExportedProgram + +_CONSTANT_SUBGRAPH_OPS = { + torch.ops.aten.cat.default, + torch.ops.aten.chunk.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.split_with_sizes.default, + operator.getitem, + torch.ops.aten._weight_norm.default, +} + + +def _has_unsafe_tuple_consumer(exported_program: ExportedProgram) -> bool: + """Avoid propagating tuple-valued chunk/split nodes, which + constant_prop_pass cannot replace.""" + tuple_producer_targets = { + torch.ops.aten.chunk.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.split_with_sizes.default, + } + for node in exported_program.graph.nodes: + if ( + node.op == "call_function" + and node.target in tuple_producer_targets + and any( + user.op != "call_function" or user.target is not operator.getitem + for user in node.users + ) + ): + return True + return False + + +def _has_unsafe_mutation(node: torch.fx.Node, schema) -> bool: + for index, argument in enumerate(schema.arguments): + if argument.alias_info is None or not argument.alias_info.is_write: + continue + mutated = ( + node.args[index] + if index < len(node.args) + else node.kwargs.get(argument.name) + ) + producer = mutated if isinstance(mutated, torch.fx.Node) else None + producer_schema = ( + getattr(producer.target, "_schema", None) + if producer is not None and producer.op == "call_function" + else None + ) + if ( + producer_schema is None + or producer_schema.is_mutable + or producer.target in _CONSTANT_SUBGRAPH_OPS + or any(ret.alias_info is not None for ret in producer_schema.returns) + ): + return True + return False + + +def _has_unsafe_node_effect(node: torch.fx.Node) -> bool: + if isinstance(node.target, torch._ops.HigherOrderOperator): + return True + schema = getattr(node.target, "_schema", None) + if schema is not None and schema.is_mutable: + return _has_unsafe_mutation(node, schema) + return node.is_impure() and node.target in _CONSTANT_SUBGRAPH_OPS + + +def _has_unsafe_graph_effects(exported_program: ExportedProgram) -> bool: + """Return whether graph effects make constant folding unsafe.""" + return any( + node.op == "call_function" and _has_unsafe_node_effect(node) + for node in exported_program.graph.nodes + ) + + +class FoldConstantSubgraphsPass: + """ + Fold selected constant subgraphs into direct constants before decomposition. + + This pass folds constant-only subgraphs such as: + + bias_0 bias_1 bias_2 # constant parameters + \\ || // + cat + | + split/chunk + | + getitem + | + addmm / linear + + Although the resulting bias slices are constant, the intermediate + cat/split/getitem chain can prevent later lowering and delegation from + recognizing them as constant inputs. + + Constant propagation evaluates eligible operations ahead of time and replaces + each projection with a direct constant: + + constant bias_0 -> addmm / linear + constant bias_1 -> addmm / linear + constant bias_2 -> addmm / linear + + Folding is skipped when mutation, aliasing, impurity, or higher-order + operations could make replacing a runtime value with persistent constant + data unsafe. + """ + + def __call__(self, exported_program: ExportedProgram) -> ExportedProgram: + # This runs before functionalization, so inspect schemas as well as the + # graph signature for mutation. + if any( + "MUTATION" in str(spec.kind) + for spec in exported_program.graph_signature.output_specs + ): + return exported_program + if _has_unsafe_graph_effects(exported_program): + return exported_program + + if not any( + node.op == "call_function" and node.target in _CONSTANT_SUBGRAPH_OPS + for node in exported_program.graph.nodes + ): + return exported_program + + if _has_unsafe_tuple_consumer(exported_program): + return exported_program + skipped = { + node.target + for node in exported_program.graph.nodes + if node.op == "call_function" and node.target not in _CONSTANT_SUBGRAPH_OPS + } + folded_program = constant_prop_pass( + exported_program, custom_skip_targets=skipped + ) + # Constant propagation inserts an Edge clone for constant outputs, but + # pre-decomposition graphs must remain entirely in the ATen dialect. + replaced_clone = False + for node in folded_program.graph.nodes: + if node.target == exir_ops.edge.aten.clone.default: + node.target = torch.ops.aten.clone.default + replaced_clone = True + if replaced_clone: + folded_program.graph_module.recompile() + return folded_program diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index 238bb992784..fb66eb3c2d9 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 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. @@ -7,18 +8,14 @@ import ctypes import hashlib import logging - from typing import cast, Dict, List, Optional, Tuple import torch from executorch.backends.transforms import get_shape - from executorch.backends.xnnpack._passes.channels_last_tagged_reshape_pass import ( ChannelsLastTaggedReshapePass, ) - from executorch.backends.xnnpack.operators.quant_params import QuantParams - from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import ( ConstantDataOffset, PerChannelGroupQuant, @@ -39,7 +36,6 @@ is_param_node, PERM_NCHW_TO_NHWC, ) - from executorch.backends.xnnpack.utils.xnnpack_constants import ( UINT64_MAX, XNN_INVALID_VALUE_ID, @@ -80,6 +76,17 @@ def get_tensor_value(xvalue: XValue) -> XNNTensorValue: return q_tensor.tensor_value +def _tensor_bytes(value: torch.Tensor) -> bytes: + """Return the tensor's contiguous physical storage as bytes.""" + + assert value.is_contiguous() or value.is_contiguous( + memory_format=torch.channels_last + ), "Tensor must have contiguous storage" + + # data_ptr() accounts for storage_offset(), limiting the read to this tensor. + return ctypes.string_at(value.data_ptr(), value.numel() * value.element_size()) + + class NodeVisitor: """ Node visitor pattern for visiting nodes in an edge IR graph and @@ -288,16 +295,9 @@ def get_quant_params( if quant_params.per_channel_group: scale = scale.to(torch.bfloat16) - # The bytes are read from the start of the storage in memory order. - # Scales are small, so always copy them into storage of their own. - scale = scale.clone(memory_format=torch.contiguous_format) - - num_bytes = scale.untyped_storage().nbytes() - scale_array = ctypes.cast( - scale.untyped_storage().data_ptr(), - ctypes.POINTER(ctypes.c_char * num_bytes), - ).contents - scale_name = hashlib.sha256(bytes(scale_array)).hexdigest() + scale_data = _tensor_bytes(scale.contiguous()) + num_bytes = len(scale_data) + scale_name = hashlib.sha256(scale_data).hexdigest() scale_name = "scale_" + scale_name xnn_graph.constant_data.append( ConstantDataOffset( @@ -309,7 +309,7 @@ def get_quant_params( f"Adding constant data with name, key {scale_name} and external_tag {external_tag} to named_data_store" ) self._named_data_store.add_named_data( - scale_name, bytes(scale_array), CONSTANT_TENSOR_ALIGNMENT, external_tag + scale_name, scale_data, CONSTANT_TENSOR_ALIGNMENT, external_tag ) if quant_params.per_channel_group: @@ -623,26 +623,16 @@ def get_serialized_buffer_index( if quant_params is not None and quant_params.is_qc4w: const_val = self.convert_to_qc4w(const_val) - # The bytes are read from the start of the storage, so a view of a - # larger tensor needs its own copy. - if const_val.nbytes != const_val.untyped_storage().nbytes(): - const_val = const_val.clone() - - size = const_val.untyped_storage().nbytes() - array_type = ctypes.c_char * size - array = ctypes.cast( - const_val.untyped_storage().data_ptr(), - ctypes.POINTER(array_type), - ).contents + data = _tensor_bytes(const_val) + size = len(data) check_or_raise( size > 0, f"Serializing constant data node {tensor} but tensor value has no bytes", ) - sha256_hash = hashlib.sha256(bytes(array)) + sha256_hash = hashlib.sha256(data) named_key = sha256_hash.hexdigest() - size = const_val.untyped_storage().nbytes() xnn_graph.constant_data.append( ConstantDataOffset(offset=UINT64_MAX, size=size, named_key=named_key) ) @@ -657,7 +647,7 @@ def get_serialized_buffer_index( ) self._named_data_store.add_named_data( named_key, - bytes(array), + data, alignment=CONSTANT_TENSOR_ALIGNMENT, external_tag=external_tag, ) diff --git a/backends/xnnpack/partition/BUCK b/backends/xnnpack/partition/BUCK index 2581043d5b2..9b55fdff005 100644 --- a/backends/xnnpack/partition/BUCK +++ b/backends/xnnpack/partition/BUCK @@ -12,6 +12,7 @@ fbcode_target(_kind = runtime.python_library, deps = [ ":partitioner_graphs", "//executorch/backends/xnnpack:xnnpack_preprocess", + "//executorch/backends/xnnpack/_passes:xnnpack_passes", "//executorch/backends/xnnpack/partition/config:xnnpack_partitioner_configs", "//executorch/exir:delegate", "//executorch/exir:lib", diff --git a/backends/xnnpack/partition/xnnpack_partitioner.py b/backends/xnnpack/partition/xnnpack_partitioner.py index 5327121b158..e8ad498c96d 100644 --- a/backends/xnnpack/partition/xnnpack_partitioner.py +++ b/backends/xnnpack/partition/xnnpack_partitioner.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 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. @@ -9,12 +10,14 @@ import logging from typing import List, Optional, Type, Union +from executorch.backends.xnnpack._passes.fold_constant_subgraphs import ( + FoldConstantSubgraphsPass, +) from executorch.backends.xnnpack.partition.config import ALL_PARTITIONER_CONFIGS from executorch.backends.xnnpack.partition.config.xnnpack_config import ( ConfigPrecisionType, XNNPartitionerConfig, ) - from executorch.backends.xnnpack.xnnpack_preprocess import XnnpackBackend from executorch.exir.backend.backend_details import ExportedProgram from executorch.exir.backend.canonical_partitioners.config_partitioner import ( @@ -28,6 +31,12 @@ class XnnpackPartitioner(ConfigerationBasedPartitioner): + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Apply XNNPACK transforms before default ATen decompositions.""" + return FoldConstantSubgraphsPass()(exported_program) + def __init__( self, configs: Optional[List[Type[XNNPartitionerConfig]]] = None, @@ -91,6 +100,7 @@ def partition(self, exported_program): if self._check_if_called_from_to_backend(): logger.warning( "\nDEPRECATION WARNING: You are using the deprecated 'to_edge() + to_backend()' workflow. " + "It does not run pre-decomposition transforms, including constant cat/split folding. " "Please consider migrating to 'to_edge_transform_and_lower()' for better error handling and optimization. " ) diff --git a/backends/xnnpack/test/ops/test_constant_bias.py b/backends/xnnpack/test/ops/test_constant_bias.py new file mode 100644 index 00000000000..0e19f665932 --- /dev/null +++ b/backends/xnnpack/test/ops/test_constant_bias.py @@ -0,0 +1,418 @@ +# Copyright 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. + +"""Regression tests for folding constant cat and split before decomposition.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn.functional as F +from executorch.backends.xnnpack.operators.node_visitor import ( + _tensor_bytes, + NodeVisitor, +) +from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( + XnnpackFloatingPointPartitioner, + XnnpackPartitioner, +) +from executorch.backends.xnnpack.serialization.xnnpack_graph_schema import XNNGraph +from executorch.devtools.backend_debug import get_delegation_info +from executorch.exir import to_edge_transform_and_lower +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) + + +class ProjectionBiases(torch.nn.Module): + def __init__(self, split, projection="linear", runtime_bias=False): + super().__init__() + self.weights = torch.nn.ParameterList( + [torch.nn.Parameter(torch.randn(8, 8)) for _ in range(3)] + ) + bias_shape = (1, 8) if projection == "addmm" else (8,) + self.biases = torch.nn.ParameterList( + [ + torch.nn.Parameter( + (torch.arange(8, dtype=torch.float32) + 10 * index).reshape( + bias_shape + ) + ) + for index in range(3) + ] + ) + self.split = split + self.projection = projection + self.runtime_bias = runtime_bias + + def forward(self, value, bias): + parts = [bias if self.runtime_bias else self.biases[0], *self.biases[1:]] + packed = torch.cat(parts, dim=0) + if self.split == "chunk": + split = packed.chunk(3, dim=0) + elif self.split == "split": + split = packed.split(1 if self.projection == "addmm" else 8, dim=0) + else: + split = packed.split( + [1, 1, 1] if self.projection == "addmm" else [8, 8, 8], + dim=0, + ) + if self.projection == "addmm": + return tuple( + torch.addmm(item, value, weight) + for weight, item in zip(self.weights, split) + ) + return tuple( + F.linear(value, weight, item) for weight, item in zip(self.weights, split) + ) + + +class TestConstantProjectionBiases(unittest.TestCase): + def test_split_forms_delegate_and_match_compiled_runtime(self): + torch.manual_seed(0) + for split in ("chunk", "split", "split_with_sizes"): + for projection in ("linear", "addmm"): + with self.subTest(split=split, projection=projection): + model = ProjectionBiases(split, projection).eval() + inputs = (torch.randn(2, 8), torch.randn(8)) + program = torch.export.export(model, inputs) + edge = to_edge_transform_and_lower( + program, partitioner=[XnnpackPartitioner()] + ) + info = get_delegation_info(edge.exported_program().graph_module) + op = ( + "aten_linear_default" + if projection == "linear" + else "aten_addmm_default" + ) + self.assertEqual(info.delegation_by_operator[op].non_delegated, 0) + self.assertEqual(info.num_delegated_nodes, 3) + pte_buffer = edge.to_executorch().buffer + runtime = _load_for_executorch_from_buffer(pte_buffer) + for actual, expected in zip( + runtime.forward(inputs), model(*inputs) + ): + torch.testing.assert_close( + actual, expected, atol=1e-5, rtol=1e-5 + ) + + def test_inherited_partitioner_folds(self): + model = ProjectionBiases("chunk").eval() + inputs = (torch.randn(2, 8), torch.randn(8)) + edge = to_edge_transform_and_lower( + torch.export.export(model, inputs), + partitioner=[XnnpackFloatingPointPartitioner()], + ) + info = get_delegation_info(edge.exported_program().graph_module) + self.assertEqual( + info.delegation_by_operator["aten_linear_default"].non_delegated, 0 + ) + + def test_serialized_bytes_use_tensor_extent_and_offset(self): + backing = torch.arange(24, dtype=torch.float32) + for view in (backing[4:12], backing[:8]): + with self.subTest(offset=view.storage_offset()): + self.assertEqual(_tensor_bytes(view), view.clone().numpy().tobytes()) + self.assertEqual(len(_tensor_bytes(view)), 8 * view.element_size()) + + def test_quantization_scales_use_tensor_extent_and_offset(self): + backing = torch.arange(16, dtype=torch.float32) / 10 + for offset, external_tag in ((0, None), (4, "scales.bin")): + with self.subTest(offset=offset): + scale = backing[offset : offset + 8] + store = NamedDataStore() + graph = XNNGraph("", [], [], 0, [], [], []) + visitor = NodeVisitor(None, {}, store) + params = SimpleNamespace( + per_channel=True, + per_channel_group=False, + is_dynamic=False, + scale=scale, + axis=0, + ) + visitor.get_quant_params(params, graph, external_tag) + expected = scale.clone().numpy().tobytes() + self.assertEqual(store.buffers, [expected]) + self.assertEqual(graph.constant_data[0].size, len(expected)) + if external_tag is not None: + self.assertIn(external_tag, store.external_data) + + def test_parameter_views_match_compiled_runtime(self): + class ViewBias(torch.nn.Module): + def __init__(self, offset): + super().__init__() + backing = torch.arange(16, dtype=torch.float32) + self.bias = torch.nn.Parameter(backing[offset : offset + 8]) + self.weight = torch.nn.Parameter(torch.randn(8, 8)) + + def forward(self, value): + return F.linear(value, self.weight, self.bias) + + for offset in (0, 4): + with self.subTest(offset=offset): + model = ViewBias(offset).eval() + inputs = (torch.randn(2, 8),) + program = torch.export.export(model, inputs) + edge = to_edge_transform_and_lower( + program, partitioner=[XnnpackPartitioner()] + ) + pte_buffer = edge.to_executorch().buffer + runtime = _load_for_executorch_from_buffer(pte_buffer) + torch.testing.assert_close( + runtime.forward(inputs)[0], model(*inputs), atol=1e-5, rtol=1e-5 + ) + + def test_dynamic_activation_shapes(self): + model = ProjectionBiases("split_with_sizes").eval() + example = (torch.randn(2, 8), torch.randn(8)) + program = torch.export.export( + model, + example, + dynamic_shapes=({0: torch.export.Dim("batch", min=1, max=4)}, None), + ) + edge = to_edge_transform_and_lower(program, partitioner=[XnnpackPartitioner()]) + info = get_delegation_info(edge.exported_program().graph_module) + self.assertEqual( + info.delegation_by_operator["aten_linear_default"].non_delegated, 0 + ) + pte_buffer = edge.to_executorch().buffer + runtime = _load_for_executorch_from_buffer(pte_buffer) + for batch in (2, 3): + inputs = (torch.randn(batch, 8), torch.randn(8)) + for actual, expected in zip(runtime.forward(inputs), model(*inputs)): + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_runtime_bias_and_other_arithmetic_stay_dynamic(self): + model = ProjectionBiases("chunk", runtime_bias=True).eval() + example = (torch.randn(2, 8), torch.randn(8)) + program = torch.export.export(model, example) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertTrue( + any( + node.target == torch.ops.aten.cat.default + for node in transformed.graph.nodes + ) + ) + alternate = (example[0], example[1] + 2) + self.assertFalse( + torch.equal( + transformed.module()(*example)[0], transformed.module()(*alternate)[0] + ) + ) + + class OtherArithmetic(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", torch.arange(8, dtype=torch.float32)) + + def forward(self, value): + return value + self.bias.sin() + + other = torch.export.export(OtherArithmetic(), (torch.randn(8),)) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(other) + self.assertTrue( + any( + node.target == torch.ops.aten.sin.default + for node in transformed.graph.nodes + ) + ) + + class QuantizedConstant(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", torch.arange(8, dtype=torch.float32)) + + def forward(self, value): + quantized = torch.ops.aten.quantize_per_tensor.default( + self.bias, 0.1, 0, torch.quint8 + ) + return value + torch.ops.aten.dequantize.self(quantized) + + quantized = torch.export.export(QuantizedConstant(), (torch.randn(8),)) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(quantized) + self.assertTrue( + any( + node.target == torch.ops.aten.quantize_per_tensor.default + for node in transformed.graph.nodes + ) + ) + + def test_fresh_activation_mutation_does_not_block_constant_bias(self): + class FreshTemporary(ProjectionBiases): + def forward(self, value, bias): + temporary = value * 2 + temporary.add_(1) + return super().forward(temporary, bias) + + model = FreshTemporary("chunk").eval() + inputs = (torch.randn(2, 8), torch.randn(8)) + edge = to_edge_transform_and_lower( + torch.export.export(model, inputs), + partitioner=[XnnpackPartitioner()], + ) + info = get_delegation_info(edge.exported_program().graph_module) + self.assertEqual( + info.delegation_by_operator["aten_linear_default"].non_delegated, 0 + ) + pte_buffer = edge.to_executorch().buffer + runtime = _load_for_executorch_from_buffer(pte_buffer) + for actual, expected in zip(runtime.forward(inputs), model(*inputs)): + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) + + def test_unsafe_tuple_consumer_skips_folding(self): + class MixedTupleConsumers(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8, 8)) + self.unsafe_biases = torch.nn.ParameterList( + [torch.nn.Parameter(torch.randn(8)) for _ in range(3)] + ) + self.safe_biases = torch.nn.ParameterList( + [torch.nn.Parameter(torch.randn(8)) for _ in range(3)] + ) + + def forward(self, value): + unsafe = torch.cat(tuple(self.unsafe_biases)).chunk(3) + safe = torch.cat(tuple(self.safe_biases)).chunk(3) + projected_unsafe = F.linear(value, self.weight, unsafe[0]) + projected_safe = F.linear(value, self.weight, safe[0]) + unsupported = torch.ops.aten._foreach_add.Scalar(unsafe, 1.0) + return projected_unsafe, projected_safe, unsupported[1] + + program = torch.export.export(MixedTupleConsumers(), (torch.randn(2, 8),)) + chunks = [ + node + for node in program.graph.nodes + if node.target == torch.ops.aten.chunk.default + ] + unsafe_chunk = chunks[0] + foreach = next( + node + for node in program.graph.nodes + if node.target == torch.ops.aten._foreach_add.Scalar + ) + foreach.args = (unsafe_chunk, 1.0) + program.graph_module.recompile() + + before = str(program.graph) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertIs(transformed, program) + self.assertEqual(str(transformed.graph), before) + + def test_mutated_cat_split_is_not_folded(self): + class MutatedPacking(torch.nn.Module): + def __init__(self): + super().__init__() + self.bias = torch.nn.Parameter(torch.ones(8)) + + def forward(self, value): + packed = torch.cat((self.bias, self.bias, self.bias)) + packed.add_(1) + return value + packed.chunk(3)[0] + + program = torch.export.export(MutatedPacking(), (torch.zeros(8),)) + before = str(program.graph) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertIs(transformed, program) + self.assertEqual(str(transformed.graph), before) + + def test_cat_used_as_out_argument_is_not_folded(self): + class MutatedOut(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", torch.ones(8)) + + def forward(self, value): + packed = torch.cat((self.bias, self.bias)) + temporary = value * 2 + torch.ops.aten.add.out(temporary, 1, out=packed) + return packed.chunk(2)[0] + + program = torch.export.export(MutatedOut(), (torch.zeros(16),)) + before = str(program.graph) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertIs(transformed, program) + self.assertEqual(str(transformed.graph), before) + edge = to_edge_transform_and_lower(program, partitioner=[XnnpackPartitioner()]) + torch.testing.assert_close( + edge.exported_program().module()(torch.ones(16)), torch.full((8,), 3.0) + ) + + def test_higher_order_control_flow_skips_folding(self): + class Conditional(torch.nn.Module): + def __init__(self): + super().__init__() + self.bias = torch.nn.Parameter(torch.ones(8)) + + def forward(self, predicate, value): + packed = torch.cat((self.bias, self.bias)).chunk(2)[0] + selected = torch.cond( + predicate, lambda item: item + 1, lambda item: item - 1, [value] + ) + return selected + packed + + inputs = (torch.tensor(True), torch.randn(8)) + program = torch.export.export(Conditional(), inputs) + before = str(program.graph) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertIs(transformed, program) + self.assertEqual(str(transformed.graph), before) + + def test_mutable_state_is_not_frozen(self): + class Stateful(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", torch.zeros(8)) + self.register_buffer("other", torch.ones(8)) + + def forward(self, value): + self.bias.add_(1) + packed = torch.cat((self.bias, self.other, self.other)).chunk(3) + return value + packed[0] + + model = Stateful() + program = torch.export.export(model, (torch.zeros(8),)) + before = str(program.graph) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + self.assertIs(transformed, program) + self.assertEqual(str(transformed.graph), before) + edge = to_edge_transform_and_lower( + transformed, partitioner=[XnnpackPartitioner()] + ) + pte_buffer = edge.to_executorch().buffer + runtime = _load_for_executorch_from_buffer(pte_buffer) + first = runtime.forward((torch.zeros(8),))[0] + second = runtime.forward((torch.zeros(8),))[0] + torch.testing.assert_close(first, torch.ones(8)) + torch.testing.assert_close(second, torch.full((8,), 2.0)) + + def test_constant_output_does_not_introduce_edge_op(self): + class ConstantOutput(torch.nn.Module): + def __init__(self): + super().__init__() + self.bias = torch.nn.Parameter(torch.arange(8, dtype=torch.float32)) + + def forward(self, value): + packed = torch.cat((self.bias, self.bias)) + return value + 1, packed + + program = torch.export.export(ConstantOutput().eval(), (torch.randn(8),)) + transformed = XnnpackPartitioner().transform_for_pre_decomposition(program) + edge_ops = [ + node + for node in transformed.graph.nodes + if isinstance(node.target, EdgeOpOverload) + ] + self.assertEqual( + edge_ops, + [], + f"Pre-decomposition transform introduced Edge ops: {edge_ops}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..312f1d0c36a 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 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. @@ -58,6 +59,8 @@ def test_deprecation_warning_for_to_backend_workflow(self): log_contents = log_capture_string.getvalue() self.assertIn("DEPRECATION WARNING", log_contents) self.assertIn("to_edge() + to_backend()", log_contents) + self.assertIn("does not run pre-decomposition transforms", log_contents) + self.assertIn("constant cat/split folding", log_contents) self.assertIn("to_edge_transform_and_lower()", log_contents) def test_no_warning_for_to_edge_transform_and_lower_workflow(self):