From cce0234f5184afcf169d30ba7edee7bb58fd84a2 Mon Sep 17 00:00:00 2001 From: Thomas Jannaud Date: Sat, 24 Oct 2026 09:10:35 -0700 Subject: [PATCH 1/2] Resolve address-backed external constants safely Summary: Problem External constants can be represented by an integer address rather than a host pointer. Treating both representations as a directly accessible pointer can dereference an invalid address or accept a truncated/overflowing extent. Solution Distinguish pointer-backed and address-backed `FreeableBuffer` storage in the tensor parser, validate requested ranges and resolve address-backed storage through the loader. Keep fbcode/xplat mirrors and dual-address loader regressions together. Dependencies and priority Generic loader correctness change; independent of the mask, embedding splitting and hybrid optimizations. Shared resident views consume it. P0 for address-backed loading correctness. History provenance Resolve local external tensor data through checked pointer or integer addresses. Preserve both mirrored parsers and keep pointer-backed, overflow, short-storage, and shared `FlatTensor` loader regressions together. Differential Revision: D125508895 --- runtime/executor/tensor_parser_exec_aten.cpp | 25 +++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index e8b8bf5b6ab..c32f88d37a3 100644 --- a/runtime/executor/tensor_parser_exec_aten.cpp +++ b/runtime/executor/tensor_parser_exec_aten.cpp @@ -8,6 +8,9 @@ #include +#include +#include + #include #include #include @@ -218,7 +221,27 @@ ET_NODISCARD Result getTensorDataPtr( if (allocation_info == nullptr) { NamedData* data = get_data_by_key(fqn, external_constants); if (data != nullptr) { - return const_cast(data->buffer.data()); + auto pointer = data->buffer.data_safe(); + if (pointer.ok()) { + return const_cast(pointer.get()); + } + auto address = data->buffer.data_uint64_type(); + if (!address.ok()) { + return address.error(); + } + ET_CHECK_OR_RETURN_ERROR( + nbytes <= data->buffer.size_uint64(), + InvalidExternalData, + "External tensor %s exceeds its buffer size", + fqn); + constexpr auto max_address = std::numeric_limits::max(); + ET_CHECK_OR_RETURN_ERROR( + address.get() <= max_address && + (nbytes == 0 || nbytes - 1 <= max_address - address.get()), + NotSupported, + "External tensor %s is not locally addressable", + fqn); + return reinterpret_cast(static_cast(address.get())); } // Should never reach here; these tensors are resolved in // Method::parse_external_constants. Any errors should be caught there. From 63a1c003aa2d9f6cf672aeda0927afe30fbeac8b Mon Sep 17 00:00:00 2001 From: Thomas Jannaud Date: Sat, 24 Oct 2026 09:10:35 -0700 Subject: [PATCH 2/2] Preserve distinct mutation output names in graph signatures (#24105) Summary: Problem `ExportGraphSignature.buffers_to_mutate` and `user_inputs_to_mutate` index writes by output node name. No-op folding and CSE can merge two writes into one node, so one of the mutation entries is lost from the signature. Solution Rebuild the signature from `output_specs` (`_get_updated_graph_signature` in `exir/_program_utils.py`). Copy a computed result only when two writes collide in the same mutation map, and keep positional user outputs and ordinary graph uses shared. Select ATen, edge or dim-order cloning from the program's verifier. Preserve value/debug metadata and FX provenance on the copies without inheriting allocation or custom-kernel placement annotations. Mutation output copies keep independent storage, as required by the existing output pointer contract. This is the ExecuTorch half of D125489522, split out so that ExecuTorch owners can review it on its own. The Helios half (removing the Helios-side duplicated-output splitter that this change makes unnecessary) stays in D125489522, stacked on top. fbcode and xplat copies are identical. {F2001970137} Differential Revision: D127072149 --- exir/BUCK | 1 + exir/_program_utils.py | 101 ++++++++- exir/pass_manager.py | 15 +- exir/program/BUCK | 1 + exir/program/_program.py | 78 ++++++- exir/tests/targets.bzl | 2 + exir/tests/test_passes.py | 437 +++++++++++++++++++++++++++++++++++++- 7 files changed, 619 insertions(+), 16 deletions(-) diff --git a/exir/BUCK b/exir/BUCK index de6c355c7aa..c29925914b1 100644 --- a/exir/BUCK +++ b/exir/BUCK @@ -296,6 +296,7 @@ fbcode_target(_kind = runtime.python_library, ], deps = [ "//caffe2:torch", + "//executorch/exir/dialects:lib", ], ) diff --git a/exir/_program_utils.py b/exir/_program_utils.py index 6673f2a8cc8..a0e846b2982 100644 --- a/exir/_program_utils.py +++ b/exir/_program_utils.py @@ -6,13 +6,105 @@ # pyre-unsafe +from collections.abc import Callable + import torch +from executorch.exir.dialects._ops import ops as exir_ops from torch.export.exported_program import ( ConstantArgument, ExportGraphSignature, InputSpec, OutputSpec, ) +from torch.export.graph_signature import OutputKind +from torch.fx.traceback import NodeSource, NodeSourceAction + +_MUTATION_OUTPUT_CLONE_OP = "_mutation_output_clone_op" + + +def _get_mutation_output_clone_op( + program: torch.export.ExportedProgram, + clone_op: Callable[..., torch.Tensor] | None = None, +) -> Callable[..., torch.Tensor] | None: + if clone_op is not None: + return clone_op + if program.dialect == "ATEN": + return torch.ops.aten.clone.default + kind = program.graph_module.meta.get(_MUTATION_OUTPUT_CLONE_OP) + if kind == "aten": + return torch.ops.aten.clone.default + if kind == "edge": + return exir_ops.edge.aten.clone.default + if kind == "edge_dim_order": + return exir_ops.edge.dim_order_ops._clone_dim_order.default + return None + + +def _set_mutation_output_clone_op( + gm: torch.fx.GraphModule, clone_op: Callable[..., torch.Tensor] | None +) -> None: + """Record the clone dialect and layout as metadata preserved by deepcopy.""" + if clone_op is None: + return + if clone_op is torch.ops.aten.clone.default: + kind = "aten" + elif clone_op is exir_ops.edge.aten.clone.default: + kind = "edge" + elif clone_op is exir_ops.edge.dim_order_ops._clone_dim_order.default: + kind = "edge_dim_order" + else: + raise ValueError(f"Unsupported mutation output clone operator: {clone_op}") + gm.meta[_MUTATION_OUTPUT_CLONE_OP] = kind + + +def _disambiguate_mutation_outputs( + signature: ExportGraphSignature, + gm: torch.fx.GraphModule, + *, + clone_op: Callable[..., torch.Tensor] | None, +) -> bool: + """Copy colliding mutation outputs in-place, preserving positional user outputs.""" + output_node = gm.graph.output_node() + output_args = output_node.args[0] + assert isinstance(output_args, (tuple, list)) + assert len(output_args) == len( + signature.output_specs + ), "Number of outputs changed after transformation" + seen_mutations: set[tuple[OutputKind, torch.fx.Node]] = set() + new_output_args = list(output_args) + for i, (node, spec) in enumerate(zip(output_args, signature.output_specs)): + if spec.kind not in ( + OutputKind.BUFFER_MUTATION, + OutputKind.USER_INPUT_MUTATION, + ): + continue + assert isinstance(node, torch.fx.Node) + mutation = (spec.kind, node) + if mutation in seen_mutations: + if clone_op is None: + raise RuntimeError( + "Mutation output repair requires an explicit clone operator " + "for programs without an Edge compile configuration" + ) + # Mutation dictionaries are keyed by node name, while user outputs + # are positional and may share a node. Copy only colliding writes. + with gm.graph.inserting_before(output_node): + clone = gm.graph.call_function(clone_op, args=(node,)) + clone.meta = { + key: node.meta[key] + for key in ("val", "stack_trace", "nn_module_stack", "debug_handle") + if key in node.meta + } + clone.meta["from_node"] = [ + NodeSource(node, "_disambiguate_mutation_outputs", NodeSourceAction.CREATE) + ] + new_output_args[i] = clone + seen_mutations.add(mutation) + if new_output_args == list(output_args): + return False + output_node.args = (tuple(new_output_args),) + output_node.args[1:] + gm.recompile() + return True def _get_updated_range_constraints(gm, old_range_constraints=None): @@ -55,9 +147,7 @@ def _get_updated_graph_signature( old_signature: ExportGraphSignature, new_gm: torch.fx.GraphModule, ) -> ExportGraphSignature: - """ - Update the graph signature's user_input/user_outputs. - """ + """Rebuild the positional signature without modifying the graph module.""" new_input_specs = [] i = 0 for node in new_gm.graph.nodes: @@ -87,6 +177,11 @@ def _get_updated_graph_signature( output_node = new_gm.graph.output_node() assert output_node.op == "output" + output_args = output_node.args[0] + assert isinstance(output_args, (tuple, list)) + assert len(output_args) == len( + old_signature.output_specs + ), "Number of outputs changed after transformation" new_output_specs = [] for i, node in enumerate(output_node.args[0]): assert i < len( diff --git a/exir/pass_manager.py b/exir/pass_manager.py index b999547b51b..4f480e6b1c4 100644 --- a/exir/pass_manager.py +++ b/exir/pass_manager.py @@ -15,8 +15,11 @@ import torch.fx.passes.infra.pass_manager as fx import torch.utils._pytree as pytree from executorch.exir._program_utils import ( + _disambiguate_mutation_outputs, + _get_mutation_output_clone_op, _get_updated_graph_signature, _get_updated_range_constraints, + _set_mutation_output_clone_op, ) from executorch.exir.error import ExportError, ExportErrorType from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult @@ -179,6 +182,8 @@ def __call__( # noqa: C901 self, exported_program: ExportedProgram, override_verifiers: Optional[list[Type[Verifier]]] = None, + *, + mutation_output_clone_op: Callable[..., torch.Tensor] | None = None, ) -> ExportedProgramPassResult: """ Runs passes on an ExportedProgram. @@ -195,6 +200,9 @@ def __call__( # noqa: C901 self.solve_constraints() exported_program = copy.copy(exported_program) + clone_op = _get_mutation_output_clone_op( + exported_program, mutation_output_clone_op + ) if override_verifiers: exported_program._verifiers = override_verifiers @@ -222,9 +230,14 @@ def __call__( # noqa: C901 # possible that the verifier will fail upon new ExportedProgram construction, # and we should only run verification after each pass if # run_checks_after_each_pass is True. - new_graph_signature = _get_updated_graph_signature( + _disambiguate_mutation_outputs( exported_program.graph_signature, res.graph_module, + clone_op=clone_op, + ) + _set_mutation_output_clone_op(res.graph_module, clone_op) + new_graph_signature = _get_updated_graph_signature( + exported_program.graph_signature, res.graph_module ) # ExportedProgram.__init__ runs this, but this path mutates # the program in place and never reconstructs it. Duplicate diff --git a/exir/program/BUCK b/exir/program/BUCK index be1a6fc76a0..f5dfed272c3 100644 --- a/exir/program/BUCK +++ b/exir/program/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/exir/backend:partitioner", "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib", "//executorch/exir/capture:config", + "//executorch/exir/dialects:lib", "//executorch/exir/emit:emit", "//executorch/exir/emit:lib", "//executorch/exir/operator:util", diff --git a/exir/program/_program.py b/exir/program/_program.py index ed55b345362..b1d6c49f487 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -11,7 +11,7 @@ import logging import os from collections import defaultdict -from typing import Any, Dict, List, Optional, Sequence, Set, TextIO, Type, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Set, TextIO, Type, Union import torch import torch._export @@ -30,6 +30,7 @@ from executorch.exir.backend.partitioner import Partitioner from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig from executorch.exir.delegate import executorch_call_delegate, is_lowered_module +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.emit import emit_program, EmitterOutput from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap from executorch.exir.error import ExportError @@ -132,8 +133,11 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: from executorch.exir._program_utils import ( # noqa: E402 + _disambiguate_mutation_outputs, + _get_mutation_output_clone_op, _get_updated_graph_signature, _get_updated_range_constraints, + _set_mutation_output_clone_op, ) @@ -141,6 +145,7 @@ def _transform( self, *passes: PassType, override_verifiers: None | list[Type[Verifier]] = None, + mutation_output_clone_op: Callable[..., torch.Tensor] | None = None, ) -> "ExportedProgram": """ Transforms the program according to the provided passes. @@ -150,6 +155,8 @@ def _transform( *passes: A sequence of passes to apply to the program override_verifiers: Optional list of verifier classes to use instead of the default verifiers. This is needed if the transforms yields illegal graph that the default verifier cannot handle. + mutation_output_clone_op: Clone operator for the transformed graph's dialect and layout mode. + Defaults to ATen cloning or the contract recorded during Edge conversion. Returns: ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made @@ -160,7 +167,10 @@ def _transform( ), f"Expected all passes to be of PassType, not list or Verifier. Use override_verifiers kwarg instead. Got: {list(passes)}" return _transform_with_pass_manager( - self, ExportedProgramPassManager(list(passes)), override_verifiers + self, + ExportedProgramPassManager(list(passes)), + override_verifiers, + mutation_output_clone_op=mutation_output_clone_op, ) @@ -168,6 +178,8 @@ def _transform_with_pass_manager( self: ExportedProgram, pass_manager: Union[ExportedProgramPassManager, GraphModulePassManager], override_verifiers: None | list[Type[Verifier]] = None, + *, + mutation_output_clone_op: Callable[..., torch.Tensor] | None = None, ) -> "ExportedProgram": """ Transforms the program using the provided pass_manager. @@ -184,7 +196,11 @@ def _transform_with_pass_manager( ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made """ if isinstance(pass_manager, ExportedProgramPassManager): - res = pass_manager(self, override_verifiers) + res = pass_manager( + self, + override_verifiers, + mutation_output_clone_op=mutation_output_clone_op, + ) if not res.modified: return self return res.exported_program @@ -193,7 +209,10 @@ def _transform_with_pass_manager( if not res.modified: return self return _update_exported_program_graph_module( - self, res.graph_module, override_verifiers + self, + res.graph_module, + override_verifiers, + mutation_output_clone_op=mutation_output_clone_op, ) @@ -201,12 +220,19 @@ def _update_exported_program_graph_module( exported_program: ExportedProgram, gm: torch.fx.GraphModule, override_verifiers: None | list[Type[Verifier]] = None, + *, + mutation_output_clone_op: Callable[..., torch.Tensor] | None = None, ) -> "ExportedProgram": + clone_op = _get_mutation_output_clone_op(exported_program, mutation_output_clone_op) + _disambiguate_mutation_outputs( + exported_program.graph_signature, gm, clone_op=clone_op + ) transformed_ep = ExportedProgram( root=gm, graph=gm.graph, graph_signature=_get_updated_graph_signature( - exported_program.graph_signature, gm + exported_program.graph_signature, + gm, ), state_dict=exported_program.state_dict, range_constraints=_get_updated_range_constraints( @@ -219,6 +245,7 @@ def _update_exported_program_graph_module( ) transformed_ep.graph_module.meta.update(exported_program.graph_module.meta) transformed_ep.graph_module.meta.update(gm.meta) + _set_mutation_output_clone_op(transformed_ep.graph_module, clone_op) return transformed_ep @@ -707,11 +734,21 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram": assert new_gm_res is not None new_gm = new_gm_res.graph_module + clone_op = ( + exir_ops.edge.aten.clone.default + if config._skip_dim_order + else exir_ops.edge.dim_order_ops._clone_dim_order.default + ) + _disambiguate_mutation_outputs( + new_ep.exported_program.graph_signature, new_gm, clone_op=clone_op + ) + _set_mutation_output_clone_op(new_gm, clone_op) new_ep.exported_program = ExportedProgram( root=new_gm, graph=new_gm.graph, graph_signature=_get_updated_graph_signature( - new_ep.exported_program.graph_signature, new_gm + new_ep.exported_program.graph_signature, + new_gm, ), state_dict=new_ep.exported_program.state_dict, range_constraints=new_ep.exported_program.range_constraints, @@ -848,10 +885,20 @@ def _generate_edge_program( assert gm_res is not None gm = gm_res.graph_module + clone_op = ( + exir_ops.edge.aten.clone.default + if config._skip_dim_order + else exir_ops.edge.dim_order_ops._clone_dim_order.default + ) + _disambiguate_mutation_outputs(program.graph_signature, gm, clone_op=clone_op) + _set_mutation_output_clone_op(gm, clone_op) edge_program = ExportedProgram( root=gm, graph=gm.graph, - graph_signature=_get_updated_graph_signature(program.graph_signature, gm), + graph_signature=_get_updated_graph_signature( + program.graph_signature, + gm, + ), state_dict=program.state_dict, range_constraints=program.range_constraints, module_call_graph=program.module_call_graph, @@ -1612,6 +1659,11 @@ def transform( """ compile_config = compile_config or self.compile_config + clone_op = ( + exir_ops.edge.aten.clone.default + if compile_config._skip_dim_order + else exir_ops.edge.dim_order_ops._clone_dim_order.default + ) new_programs: Dict[str, ExportedProgram] = {} # Cast passes parameter upfront. @@ -1640,11 +1692,17 @@ def transform( # Depending on the passes parameter, call the corresponding transform function. if passes_seq is not None: - new_programs[name] = _transform(program, *passes_seq) + new_programs[name] = _transform( + program, *passes_seq, mutation_output_clone_op=clone_op + ) elif passes_dict is not None: - new_programs[name] = _transform(program, *passes_dict[name]) + new_programs[name] = _transform( + program, *passes_dict[name], mutation_output_clone_op=clone_op + ) elif pass_manager is not None: - new_programs[name] = _transform_with_pass_manager(program, pass_manager) + new_programs[name] = _transform_with_pass_manager( + program, pass_manager, mutation_output_clone_op=clone_op + ) # Verify the correctness of model graph after each transformation. EXIREdgeDialectVerifier(edge_compile_config=compile_config)( diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index 81444c14fa8..b9cf6ffb1dc 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -220,6 +220,7 @@ def define_common_targets(is_fbcode = False): ":models", "//caffe2:torch", "//caffe2/functorch:functorch_src", + "//executorch/exir:_program_utils", "//executorch/exir:graph_module", "//executorch/exir:lib", "//executorch/exir:memory", @@ -246,6 +247,7 @@ def define_common_targets(is_fbcode = False): "//executorch/exir/passes:spec_prop_pass", "//executorch/exir/passes:sym_to_tensor_pass", "//executorch/exir/program:program", + "//executorch/exir/verification:verifier", "//executorch/extension/pybindings:portable_lib", # @manual "//executorch/backends/xnnpack/partition:xnnpack_partitioner", "//executorch/backends/xnnpack/quantizer:xnnpack_quantizer", diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 1afaece0f3a..0662a93bf74 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -12,6 +12,7 @@ import tempfile import unittest from typing import Callable, List, Optional, Tuple +from unittest import mock import executorch.exir as exir @@ -37,6 +38,11 @@ to_edge, to_edge_transform_and_lower, ) +from executorch.exir._program_utils import ( + _disambiguate_mutation_outputs, + _get_updated_graph_signature, + _MUTATION_OUTPUT_CLONE_OP, +) from executorch.exir.dialects._ops import bind_pattern_to_op, ops, ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.emit import emit_program @@ -80,24 +86,38 @@ from executorch.exir.passes.spec_prop_pass import SpecPropPass from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass from executorch.exir.passes.sym_to_tensor_pass import SymToTensorPass -from executorch.exir.program._program import lift_constant_tensor_pass +from executorch.exir.program._program import ( + _transform, + _transform_with_pass_manager, + lift_constant_tensor_pass, +) from executorch.exir.schema import TensorShapeDynamism from executorch.exir.sym_util import eval_upper_bound from executorch.exir.tensor import TensorSpec from executorch.exir.tests.common import register_additional_test_aten_ops from executorch.exir.tests.control_flow_models import FTCondDeadCode, FTMapBasic from executorch.exir.tests.models import FeedForwardBlock, MLP, Mul +from executorch.exir.verification.verifier import EXIREdgeDialectVerifier from functorch.experimental import control_flow from torch import nn +from torch._export.verifier import Verifier from torch._prims_common import ELEMENTWISE_TYPE_PROMOTION_KIND # Import passes from torch._subclasses.fake_tensor import FakeTensorMode from torch.export import export -from torch.export.graph_signature import InputKind, InputSpec, TensorArgument +from torch.export.graph_signature import ( + ExportGraphSignature, + InputKind, + InputSpec, + OutputKind, + OutputSpec, + TensorArgument, +) from torch.fx import GraphModule, subgraph_rewriter from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager from torch.library import impl, Library from torch.testing import FileCheck from torch.utils import _pytree as pytree @@ -302,6 +322,419 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if node.op == "call_function": self.assertNotEqual(node.target, torch.ops.aten.to.dtype) + def test_remove_noop_folds_repeated_user_outputs(self) -> None: + for target in (torch.ops.aten.to.dtype, torch.ops.aten.slice_copy.Tensor): + for both_outputs in (False, True): + with self.subTest(target=target, both_outputs=both_outputs): + graph = torch.fx.Graph() + value = torch.ones(2, 3) + x = graph.placeholder("x") + x.meta["val"] = value + args = ( + (x, torch.float32) + if target == torch.ops.aten.to.dtype + else (x, 0, 0, 2) + ) + noop = graph.call_function(target, args) + noop.meta["val"] = value + graph.output((x, noop) if both_outputs else (noop,)) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + + result = RemoveNoopPass()(module) + self.assertIsNotNone(result) + outputs = module.graph.output_node().all_input_nodes + self.assertEqual(len(outputs), 1) + self.assertIs(outputs[0], x) + module.recompile() + for output in module(value): + torch.testing.assert_close(output, value) + + def test_mutation_output_repair_preserves_sharing_of_user_outputs(self) -> None: + for use_edge_ops, use_dim_order in ( + (False, False), + (True, False), + (True, True), + ): + for kind in (OutputKind.BUFFER_MUTATION, OutputKind.USER_INPUT_MUTATION): + with self.subTest( + use_edge_ops=use_edge_ops, use_dim_order=use_dim_order, kind=kind + ): + graph = torch.fx.Graph() + x = graph.placeholder("x") + result = graph.call_function(torch.ops.aten.mm.default, (x, x)) + result.meta["val"] = torch.empty(4, 4) + result.meta["stack_trace"] = "matrix multiplication" + result.meta["custom_kernel_tag"] = 7 + result.meta["spec"] = object() + result.meta["_share_alloc_with_arg_idx"] = 0 + graph.output((result, result, result, result, result)) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None) + ], + output_specs=[ + OutputSpec(kind, TensorArgument(result.name), f"state_{i}") + for i in range(3) + ] + + [ + OutputSpec( + OutputKind.USER_OUTPUT, + TensorArgument(result.name), + None, + ) + for _ in range(2) + ], + ) + + clone_op = ( + ( + exir_ops.edge.dim_order_ops._clone_dim_order.default + if use_dim_order + else exir_ops.edge.aten.clone.default + ) + if use_edge_ops + else torch.ops.aten.clone.default + ) + self.assertTrue( + _disambiguate_mutation_outputs( + signature, + module, + clone_op=clone_op, + ) + ) + updated = _get_updated_graph_signature(signature, module) + + mutations = ( + updated.buffers_to_mutate + if kind == OutputKind.BUFFER_MUTATION + else updated.user_inputs_to_mutate + ) + self.assertEqual( + list(mutations.values()), [f"state_{i}" for i in range(3)] + ) + self.assertEqual( + len( + graph.find_nodes( + op="call_function", target=torch.ops.aten.mm.default + ) + ), + 1, + ) + outputs = graph.output_node().args[0] + self.assertIs(outputs[3], result) + self.assertIs(outputs[4], result) + self.assertEqual(outputs[1].args, (result,)) + self.assertEqual(outputs[2].args, (result,)) + self.assertIs(outputs[1].meta["val"], result.meta["val"]) + self.assertEqual( + outputs[1].meta["stack_trace"], "matrix multiplication" + ) + for copy_node in outputs[1:3]: + self.assertIs(copy_node.target, clone_op) + for key in ( + "custom_kernel_tag", + "spec", + "_share_alloc_with_arg_idx", + ): + self.assertNotIn(key, copy_node.meta) + values = module(torch.arange(16, dtype=torch.float32).view(4, 4)) + for value in values: + torch.testing.assert_close(value, values[0]) + self.assertEqual(len({value.data_ptr() for value in values[:3]}), 3) + self.assertEqual(values[3].data_ptr(), values[4].data_ptr()) + node_count = len(list(graph.nodes)) + self.assertFalse(_disambiguate_mutation_outputs( + updated, + module, + clone_op=clone_op, + )) + self.assertEqual(len(list(graph.nodes)), node_count) + + def test_mutation_output_repair_handles_inputs_and_attributes(self) -> None: + for source in ("placeholder", "get_attr"): + with self.subTest(source=source): + graph = torch.fx.Graph() + root = torch.nn.Module() + value = torch.tensor([7], dtype=torch.int32) + root.register_buffer("state", value) + node = ( + graph.placeholder("state") + if source == "placeholder" + else graph.get_attr("state") + ) + node.meta["val"] = value + graph.output((node, node)) + module = torch.fx.GraphModule(root, graph) + signature = ExportGraphSignature( + input_specs=( + [ + InputSpec( + InputKind.BUFFER, + TensorArgument("state"), + "state", + persistent=True, + ) + ] + if source == "placeholder" + else [] + ), + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(node.name), + f"state_{i}", + ) + for i in range(2) + ], + ) + + _disambiguate_mutation_outputs( + signature, module, clone_op=torch.ops.aten.clone.default + ) + updated = _get_updated_graph_signature(signature, module) + + self.assertEqual(len(graph.find_nodes(op=source)), 1) + self.assertEqual(len(updated.input_specs), len(signature.input_specs)) + self.assertEqual(len(updated.buffers_to_mutate), 2) + outputs = graph.output_node().args[0] + self.assertEqual(outputs[1].meta["from_node"][0].name, node.name) + values = module(value) if source == "placeholder" else module() + torch.testing.assert_close(values[0], value) + torch.testing.assert_close(values[1], value) + self.assertNotEqual(values[0].data_ptr(), values[1].data_ptr()) + + def test_signature_update_does_not_copy_large_repeated_user_outputs(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + graph.output((x, x)) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None)], + output_specs=[ + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument("x"), None) + for _ in range(2) + ], + ) + value = torch.ones(1, 2, 1024, 256) + + updated = _get_updated_graph_signature(signature, module) + + self.assertEqual(updated.user_outputs, ("x", "x")) + self.assertFalse(any(node.op == "call_function" for node in graph.nodes)) + first, second = module(value) + self.assertEqual(first.data_ptr(), value.data_ptr()) + self.assertEqual(second.data_ptr(), value.data_ptr()) + + def test_signature_update_does_not_modify_colliding_mutation_outputs(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + graph.output((x, x)) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None)], + output_specs=[ + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument("x"), name) + for name in ("first", "second") + ], + ) + original_code = module.code + original_nodes = tuple(graph.nodes) + with mock.patch.object(module, "recompile") as recompile: + updated = _get_updated_graph_signature(signature, module) + + recompile.assert_not_called() + self.assertEqual(tuple(graph.nodes), original_nodes) + self.assertEqual(module.code, original_code) + self.assertEqual(graph.output_node().args[0], (x, x)) + self.assertEqual([spec.target for spec in updated.output_specs], ["first", "second"]) + + def test_signature_update_rejects_changed_output_count(self) -> None: + graph = torch.fx.Graph() + graph.output(()) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[], + output_specs=[ + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument("missing"), None) + ], + ) + + with self.assertRaisesRegex(AssertionError, "Number of outputs changed"): + _get_updated_graph_signature(signature, module) + + def test_mutations_in_separate_name_maps_can_share_a_result(self) -> None: + graph = torch.fx.Graph() + graph.placeholder("buffer") + x = graph.placeholder("x") + graph.output((x, x)) + module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[ + InputSpec( + InputKind.BUFFER, + TensorArgument("buffer"), + "buffer", + persistent=True, + ), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument("x"), "buffer"), + OutputSpec(OutputKind.USER_INPUT_MUTATION, TensorArgument("x"), "x"), + ], + ) + + self.assertFalse( + _disambiguate_mutation_outputs(signature, module, clone_op=None) + ) + updated = _get_updated_graph_signature(signature, module) + + self.assertEqual(updated.buffers_to_mutate, {"x": "buffer"}) + self.assertEqual(updated.user_inputs_to_mutate, {"x": "x"}) + self.assertFalse(any(node.op == "call_function" for node in graph.nodes)) + + def test_repeated_buffer_mutations_survive_aten_and_edge_transforms(self) -> None: + class StateModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("first", torch.zeros(4)) + self.register_buffer("second", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.first.copy_(x + x) + self.second.copy_(x + x) + return x * 3 + + inputs = (torch.arange(4, dtype=torch.float32),) + program = torch.export.export(StateModel(), inputs).run_decompositions({}) + aten = _transform(program, CSEPass()) + for use_dim_order in (False, True): + edge = to_edge( + copy.deepcopy(aten), + compile_config=EdgeCompileConfig(_skip_dim_order=not use_dim_order), + ) + edge = edge.transform([CSEPass()]) + for stage in (aten, edge.exported_program()): + with self.subTest(dialect=stage.dialect, use_dim_order=use_dim_order): + stage.validate() + self.assertCountEqual( + stage.graph_signature.buffers_to_mutate.values(), + ["first", "second"], + ) + torch.testing.assert_close(stage.module()(*inputs), inputs[0] * 3) + for name in ("first", "second"): + torch.testing.assert_close( + stage.state_dict[name], inputs[0] * 2 + ) + executable = edge.to_executorch() + self.assertGreater(len(executable.buffer), 0) + + def test_mutation_output_repair_with_custom_verifier(self) -> None: + class ShareStateOutputsPass(ExportPass): + def call(self, gm: torch.fx.GraphModule) -> PassResult: + output = gm.graph.output_node() + first, _, user_output = output.args[0] + output.args = ((first, first, user_output),) + gm.graph.eliminate_dead_code() + return PassResult(gm, True) + + class StateModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("first", torch.zeros(4)) + self.register_buffer("second", torch.zeros(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.first.copy_(x + x) + self.second.copy_(x + x) + return x * 3 + + inputs = (torch.arange(4, dtype=torch.float32),) + program = torch.export.export(StateModel(), inputs).run_decompositions({}) + for use_dim_order in (False, True): + config = EdgeCompileConfig(_skip_dim_order=not use_dim_order) + standard_verifier = EXIREdgeDialectVerifier(edge_compile_config=config) + + class CustomEdgeVerifier(Verifier): + dialect = "EDGE" + + def allowed_op_types(self): + return standard_verifier.allowed_op_types() + + def allowed_getattr_types(self): + return standard_verifier.allowed_getattr_types() + + def check_valid_op(self, op): + standard_verifier.check_valid_op(op) + + def check_additional(self, gm: torch.fx.GraphModule) -> None: + standard_verifier.check_additional(gm) + + self.assertNotIn("use_dim_order", vars(CustomEdgeVerifier())) + edge = to_edge(copy.deepcopy(program), compile_config=config) + expected_clone = ( + exir_ops.edge.dim_order_ops._clone_dim_order.default + if use_dim_order + else exir_ops.edge.aten.clone.default + ) + unconfigured = copy.deepcopy(edge.exported_program()) + unconfigured.graph_module.meta.pop(_MUTATION_OUTPUT_CLONE_OP) + with self.assertRaisesRegex(RuntimeError, "explicit clone operator"): + _transform_with_pass_manager( + unconfigured, + GraphModulePassManager([ShareStateOutputsPass()]), + [CustomEdgeVerifier], + ) + for graph_module_only, explicit_clone in itertools.product( + (False, True), (False, True) + ): + with self.subTest( + use_dim_order=use_dim_order, + graph_module_only=graph_module_only, + explicit_clone=explicit_clone, + ): + source = copy.deepcopy(edge.exported_program()) + if explicit_clone: + source.graph_module.meta.pop(_MUTATION_OUTPUT_CLONE_OP) + clone_op = expected_clone if explicit_clone else None + if graph_module_only: + transformed = _transform_with_pass_manager( + source, + GraphModulePassManager([ShareStateOutputsPass()]), + [CustomEdgeVerifier], + mutation_output_clone_op=clone_op, + ) + else: + transformed = _transform( + source, + ShareStateOutputsPass(), + override_verifiers=[CustomEdgeVerifier], + mutation_output_clone_op=clone_op, + ) + transformed.validate() + self.assertIs(transformed.verifier, CustomEdgeVerifier) + self.assertEqual( + transformed.graph_module.meta[_MUTATION_OUTPUT_CLONE_OP], + "edge_dim_order" if use_dim_order else "edge", + ) + self.assertCountEqual( + transformed.graph_signature.buffers_to_mutate.values(), + ["first", "second"], + ) + self.assertTrue( + transformed.graph.find_nodes( + op="call_function", target=expected_clone + ) + ) + torch.testing.assert_close( + transformed.module()(*inputs), inputs[0] * 3 + ) + for name in ("first", "second"): + torch.testing.assert_close( + transformed.state_dict[name], inputs[0] * 2 + ) + def test_redundant_slice_copy_removal(self) -> None: class FooWithNoSlice(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: