From c24c7ddc881b625d53260bbaaadabd9104d07145 Mon Sep 17 00:00:00 2001 From: iRAFEEK <182501111+iRAFEEK@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:46:21 -0700 Subject: [PATCH 1/5] feat: add ReplaceSliceCopyWithSlicePass with contiguity detection (#10917) Slice analog of ReplaceViewCopyWithViewPass. Detects contiguous (outermost-dim, unit-step) slice_copy nodes eligible to be re-inplaced as zero-copy slices. Rewrite is gated behind offset-based sub-buffer aliasing support in memory planning (pending design discussion), so the pass currently runs as a safe no-op. --- .../replace_slice_copy_with_slice_pass.py | 123 ++++++++++++++++++ 1 file changed, 123 insertions(+) create mode 100644 exir/passes/replace_slice_copy_with_slice_pass.py diff --git a/exir/passes/replace_slice_copy_with_slice_pass.py b/exir/passes/replace_slice_copy_with_slice_pass.py new file mode 100644 index 00000000000..d0c53089a99 --- /dev/null +++ b/exir/passes/replace_slice_copy_with_slice_pass.py @@ -0,0 +1,123 @@ +# 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. + +# pyre-strict + +"""Re-inplace contiguous ``slice_copy`` nodes as lightweight slices (#10917). + +This is the slice analog of :class:`ReplaceViewCopyWithViewPass`. A +``slice_copy`` addresses a sub-region of its input's storage; when that +sub-region is *contiguous* it can, in principle, be re-inplaced as a zero-copy +view into the base buffer instead of emitting a full-copy ``slice_copy`` kernel. + +Scope note (see #10917): + ``ReplaceViewCopyWithViewPass`` can reuse ``memory.view`` because a view + aliases the *entire* base buffer -- same ``nbytes`` and offset ``0`` (the + ``_ViewSpec`` guards ``nbytes == base.nbytes``). A slice aliases only a + *sub-region* at a non-zero byte offset with fewer bytes than the base, and + ExecuTorch has no offset-based aliasing mechanism in memory planning today. + Fully eliminating the copy therefore requires (a) memory-planning support + for offset sub-buffer aliasing and (b) a lightweight runtime op + (``et_slice``) mirroring ``et_view``. That runtime design is under + discussion with the maintainer. + + This pass implements the piece that is well-defined regardless of that + design decision: correctly identifying which ``slice_copy`` nodes are + *eligible* (contiguous) for re-inplacing. The rewrite is gated behind the + offset-aliasing support and is a no-op until it lands. +""" + +import logging + +import torch +from executorch.exir.dialects._ops import ops +from torch.fx.passes.infra.pass_base import PassBase, PassResult + +logger: logging.Logger = logging.getLogger(__name__) + + +def _is_slice_copy(node: torch.fx.Node) -> bool: + return node.op == "call_function" and node.target in ( + torch.ops.aten.slice_copy.Tensor, + ops.edge.aten.slice_copy.Tensor, + ) + + +def is_contiguous_slice_copy(node: torch.fx.Node) -> bool: + """Return True if ``node`` is a ``slice_copy`` whose result is a contiguous + sub-region of a contiguous input, and is therefore eligible to be + re-inplaced as a zero-copy slice. + + A slice ``self[start:end:step]`` along ``dim`` is a contiguous sub-buffer of + a contiguous input only when it is taken along the outermost (first) storage + dimension with unit step. Slicing an inner dimension, or using ``step > 1``, + produces a strided (non-contiguous) result that cannot alias the base buffer + without a copy. + + Signature: ``slice_copy.Tensor(self, dim=0, start=None, end=None, step=1)``. + """ + if not _is_slice_copy(node): + return False + + args = node.args + self_arg = args[0] + + # dim defaults to 0; normalize negatives against the input rank. + dim = args[1] if len(args) > 1 else 0 + step = args[4] if len(args) > 4 else 1 + + if step != 1: + return False + + rank = None + self_val = self_arg.meta.get("val") if isinstance(self_arg, torch.fx.Node) else None + if self_val is not None and hasattr(self_val, "dim"): + rank = self_val.dim() + + if isinstance(dim, int) and dim < 0: + if rank is None: + # Cannot resolve a negative dim to the outermost dim without rank. + return False + dim = dim + rank + + # Only an outermost-dim, unit-step slice of a contiguous input is a + # contiguous sub-buffer that could alias the base storage. + return dim == 0 + + +class ReplaceSliceCopyWithSlicePass(PassBase): + """Re-inplace eligible (contiguous) ``slice_copy`` nodes as lightweight + slices. + + Until offset-based sub-buffer aliasing lands in memory planning (see the + module docstring and #10917), this pass only *identifies* eligible nodes and + performs no graph mutation, so it is safe to run in the pipeline. + """ + + def __init__(self) -> None: + super().__init__() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + n_eligible = 0 + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in module.graph.nodes: + # A slice feeding the graph output can have its pointer modified + # at runtime, mirroring the view_copy pass's output guard. + if is_contiguous_slice_copy(node) and all( + u.op != "output" for u in node.users + ): + n_eligible += 1 + + logger.debug( + "ReplaceSliceCopyWithSlicePass: %d contiguous slice_copy node(s) " + "eligible for re-inplacing (rewrite pending offset-aliasing support, " + "#10917).", + n_eligible, + ) + # No mutation yet -> report unchanged. + return PassResult(graph_module, False) From a3fef4f0dd44f190b43a522eee16c21d02eea860 Mon Sep 17 00:00:00 2001 From: iRAFEEK <182501111+iRAFEEK@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:46:21 -0700 Subject: [PATCH 2/5] test: add unit tests for slice_copy contiguity detection (#10917) Covers outermost-dim/unit-step eligibility, negative-dim resolution, strided/inner-dim rejection, and that the pass is a safe no-op until the offset-aliasing rewrite lands. --- ...test_replace_slice_copy_with_slice_pass.py | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 exir/tests/test_replace_slice_copy_with_slice_pass.py diff --git a/exir/tests/test_replace_slice_copy_with_slice_pass.py b/exir/tests/test_replace_slice_copy_with_slice_pass.py new file mode 100644 index 00000000000..4e2e468a5d7 --- /dev/null +++ b/exir/tests/test_replace_slice_copy_with_slice_pass.py @@ -0,0 +1,87 @@ +# 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. + +# pyre-strict + +import unittest + +import torch +from executorch.exir import to_edge +from executorch.exir.passes.replace_slice_copy_with_slice_pass import ( + _is_slice_copy, + is_contiguous_slice_copy, + ReplaceSliceCopyWithSlicePass, +) +from torch.export import export + + +class TestReplaceSliceCopyWithSlicePass(unittest.TestCase): + def _edge_graph_module( + self, module: torch.nn.Module, inputs: tuple + ) -> torch.fx.GraphModule: + ep = export(module.eval(), inputs, strict=True) + return to_edge(ep).exported_program().graph_module + + def test_contiguity_classification(self) -> None: + """A unit-step slice along the outermost dim is contiguous (eligible); + inner-dim or strided slices are not.""" + + class M(torch.nn.Module): + def forward(self, x): + a = x[0:2] # dim 0, step 1 -> contiguous (eligible) + b = x[:, 1:3] # dim 1 -> strided (not eligible) + c = x[0:4:2] # dim 0, step 2 -> strided (not eligible) + return a.sum() + b.sum() + c.sum() + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + slice_nodes = [n for n in gm.graph.nodes if _is_slice_copy(n)] + eligible = [n for n in slice_nodes if is_contiguous_slice_copy(n)] + + self.assertEqual(len(slice_nodes), 3) + self.assertEqual(len(eligible), 1) + + def test_negative_outermost_dim_is_contiguous(self) -> None: + """A negative dim that resolves to the outermost dim is still eligible.""" + + class M(torch.nn.Module): + def forward(self, x): + # dim=-2 on a rank-2 tensor resolves to dim 0. + return torch.ops.aten.slice_copy.Tensor(x, -2, 0, 2).sum() + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + eligible = [ + n for n in gm.graph.nodes if is_contiguous_slice_copy(n) + ] + self.assertEqual(len(eligible), 1) + + def test_pass_is_safe_noop_until_offset_aliasing_lands(self) -> None: + """The pass must run cleanly and not mutate the graph while the + offset-aliasing rewrite is still gated (see #10917).""" + + class M(torch.nn.Module): + def forward(self, x): + return x[0:2] + 1.0 + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + before = gm.code + result = ReplaceSliceCopyWithSlicePass()(gm) + self.assertIsNotNone(result) + self.assertFalse(result.modified) + self.assertEqual(before, result.graph_module.code) + + def test_non_slice_nodes_are_ignored(self) -> None: + class M(torch.nn.Module): + def forward(self, x): + return (x + 1.0).relu() + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + self.assertEqual( + [n for n in gm.graph.nodes if is_contiguous_slice_copy(n)], [] + ) + + +if __name__ == "__main__": + unittest.main() From 5a3e4a3782209500b8cbf155d227e15ec59e9d97 Mon Sep 17 00:00:00 2001 From: iRAFEEK <182501111+iRAFEEK@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:44:08 -0700 Subject: [PATCH 3/5] style: collapse single-line assertEqual to satisfy lintrunner --- exir/tests/test_replace_slice_copy_with_slice_pass.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/exir/tests/test_replace_slice_copy_with_slice_pass.py b/exir/tests/test_replace_slice_copy_with_slice_pass.py index 4e2e468a5d7..2721206ef1e 100644 --- a/exir/tests/test_replace_slice_copy_with_slice_pass.py +++ b/exir/tests/test_replace_slice_copy_with_slice_pass.py @@ -78,9 +78,7 @@ def forward(self, x): return (x + 1.0).relu() gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) - self.assertEqual( - [n for n in gm.graph.nodes if is_contiguous_slice_copy(n)], [] - ) + self.assertEqual([n for n in gm.graph.nodes if is_contiguous_slice_copy(n)], []) if __name__ == "__main__": From db374f25ee34dafe95830b98c657386279feb704 Mon Sep 17 00:00:00 2001 From: iRAFEEK <182501111+iRAFEEK@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:42:30 -0700 Subject: [PATCH 4/5] style: format slice copy contiguity test --- exir/tests/test_replace_slice_copy_with_slice_pass.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/exir/tests/test_replace_slice_copy_with_slice_pass.py b/exir/tests/test_replace_slice_copy_with_slice_pass.py index 2721206ef1e..95751eea6ca 100644 --- a/exir/tests/test_replace_slice_copy_with_slice_pass.py +++ b/exir/tests/test_replace_slice_copy_with_slice_pass.py @@ -52,9 +52,7 @@ def forward(self, x): return torch.ops.aten.slice_copy.Tensor(x, -2, 0, 2).sum() gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) - eligible = [ - n for n in gm.graph.nodes if is_contiguous_slice_copy(n) - ] + eligible = [n for n in gm.graph.nodes if is_contiguous_slice_copy(n)] self.assertEqual(len(eligible), 1) def test_pass_is_safe_noop_until_offset_aliasing_lands(self) -> None: From a5fda433b080c5d63dc6bcfb6f52599fa08f82aa Mon Sep 17 00:00:00 2001 From: iRAFEEK <182501111+iRAFEEK@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:21:10 -0700 Subject: [PATCH 5/5] feat: implement zero-copy sub-buffer aliasing for contiguous slice_copy Replaces eligible contiguous slice_copy nodes with a memory.slice alias so the emitted program does not pay for a full tensor copy. _SliceSpec shares the base's mem_id and computes mem_offset = base.mem_offset + start * base.stride[0] * elem_size. The .pte format already carries (memory_id, memory_offset) via AllocationDetails, so no schema change is required. Memory planning handles memory.slice like memory.view -- the base spec is returned from get_node_tensor_specs, which extends the base's lifetime over the slice's consumers so the buffer is not reused while the alias is live. Emission mirrors _emit_view's elide path, needing no runtime kernel. Eligibility is gated to dim-0, unit-step slices with a non-negative start on a base that has the default dim order and its own allocation. Non-default layouts would otherwise be silently reinterpreted by the contiguous output stride, and an aliasing base (slice-of-slice or slice-of-view) has no concrete allocation to offset from. Everything outside those gates falls back to slice_copy unchanged. Also declares inplace_base on _SliceSpec, which the greedy memory planning algorithm reads. Verified locally against the executorch wheel runtime: - contiguous slices emit no slice_copy kernel (only aten::add) - outputs match eager for offset/lifetime/chained/3-D cases - ineligible slices still fall back to copy and stay correct - no regressions: exir/tests, exir/emit, exir/backend failure sets are identical to a pristine baseline --- exir/emit/_emitter.py | 17 + exir/memory.py | 15 +- exir/memory_planning.py | 5 + exir/pass_base.py | 2 +- exir/passes/__init__.py | 1 + .../replace_slice_copy_with_slice_pass.py | 292 ++++++++++++++---- exir/program/_program.py | 4 + exir/serde/export_serialize.py | 1 + exir/serde/serialize.py | 1 + ...test_replace_slice_copy_with_slice_pass.py | 135 +++++++- 10 files changed, 408 insertions(+), 65 deletions(-) diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index 4e3baf26592..6029206c2e1 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -1287,6 +1287,20 @@ def _emit_view(self, args: Tuple[_Argument, ...]) -> _EmitterValue: self.chain.instructions.append(kernel) return out_arg + def _emit_slice(self, args: Tuple[_Argument, ...]) -> _EmitterValue: + """Emit a statically memory-planned slice as a sub-buffer alias. + + ``ReplaceSliceCopyWithSlicePass`` creates ``_SliceSpec`` values whose + ``mem_offset`` is the byte offset into their base allocation. No kernel + is needed for a static planned slice: the tensor value can be emitted + directly from that specification. + """ + assert 4 <= len(args) <= 5 + spec = self.node.meta["spec"] + assert spec.is_static_shape_tensor + assert spec.mem_id is not None and spec.mem_offset is not None + return self._emit_spec(spec) + def _add_debug_handle( self, emitter_id: int, @@ -1786,6 +1800,9 @@ def call_function( # pyre-fixme[14] elif target == memory.view: return self._emit_view(args) + elif target == memory.slice: + return self._emit_slice(args) + elif target == memory.free: assert len(args) == 1 # pyre-ignore diff --git a/exir/memory.py b/exir/memory.py index 36a244bc02f..8295218cae3 100644 --- a/exir/memory.py +++ b/exir/memory.py @@ -6,7 +6,7 @@ # pyre-strict -from typing import List, Tuple, Union +from typing import List, Optional, Tuple, Union import torch from executorch.exir.sym_util import eval_shape @@ -48,3 +48,16 @@ def view(base: torch.Tensor, size: List[int]) -> torch.Tensor: It is used to elide view_copy nodes. """ return base.view(size) + + +def slice( # noqa: A001 + base: torch.Tensor, + dim: int = 0, + start: Optional[int] = None, + end: Optional[int] = None, + step: int = 1, +) -> torch.Tensor: + """ + Mimics ``aten.slice.Tensor`` for eliding contiguous ``slice_copy`` nodes. + """ + return torch.ops.aten.slice.Tensor(base, dim, start, end, step) diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 012cf8dd144..79b992fd680 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -627,6 +627,7 @@ def collect_specs_from_nodes( # noqa: C901 in [ memory.alloc, memory.view, + memory.slice, operator.getitem, torch.ops.higher_order.cond, exir_while, @@ -908,6 +909,10 @@ def get_node_tensor_specs( base = node.args[0] assert isinstance(base, torch.fx.Node) specs = base.meta.get("spec") + elif node.target == memory.slice: + base = node.args[0] + assert isinstance(base, torch.fx.Node) + specs = base.meta.get("spec") else: specs = node.meta.get("spec") diff --git a/exir/pass_base.py b/exir/pass_base.py index c657ac53a91..0f3e10365c0 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -863,7 +863,7 @@ def call_function( # TODO according to zhengxu ExportPassBase should not be aware of # memory.alloc. Check this comment: # https://www.internalfb.com/diff/D42758019?dst_version_fbid=5906016402813292&transaction_fbid=1104713900200176 - elif target == memory.alloc: + elif target in (memory.alloc, memory.slice): return self.callback._fx( "call_function", target, diff --git a/exir/passes/__init__.py b/exir/passes/__init__.py index 51ae9055ec7..9e1b01ac807 100644 --- a/exir/passes/__init__.py +++ b/exir/passes/__init__.py @@ -266,6 +266,7 @@ def callWithLoggerEnabled(self, graph_module: torch.fx.GraphModule) -> None: # it's retraced after running to_out_variant with the first trace. memory.alloc, memory.view, + memory.slice, executorch_call_delegate, } to_out_var_skiplist.update(_EXECUTORCH_SYM_OPS) diff --git a/exir/passes/replace_slice_copy_with_slice_pass.py b/exir/passes/replace_slice_copy_with_slice_pass.py index d0c53089a99..4799e98b3d2 100644 --- a/exir/passes/replace_slice_copy_with_slice_pass.py +++ b/exir/passes/replace_slice_copy_with_slice_pass.py @@ -8,36 +8,31 @@ """Re-inplace contiguous ``slice_copy`` nodes as lightweight slices (#10917). -This is the slice analog of :class:`ReplaceViewCopyWithViewPass`. A -``slice_copy`` addresses a sub-region of its input's storage; when that -sub-region is *contiguous* it can, in principle, be re-inplaced as a zero-copy -view into the base buffer instead of emitting a full-copy ``slice_copy`` kernel. - -Scope note (see #10917): - ``ReplaceViewCopyWithViewPass`` can reuse ``memory.view`` because a view - aliases the *entire* base buffer -- same ``nbytes`` and offset ``0`` (the - ``_ViewSpec`` guards ``nbytes == base.nbytes``). A slice aliases only a - *sub-region* at a non-zero byte offset with fewer bytes than the base, and - ExecuTorch has no offset-based aliasing mechanism in memory planning today. - Fully eliminating the copy therefore requires (a) memory-planning support - for offset sub-buffer aliasing and (b) a lightweight runtime op - (``et_slice``) mirroring ``et_view``. That runtime design is under - discussion with the maintainer. - - This pass implements the piece that is well-defined regardless of that - design decision: correctly identifying which ``slice_copy`` nodes are - *eligible* (contiguous) for re-inplacing. The rewrite is gated behind the - offset-aliasing support and is a no-op until it lands. +Slice analog of :class:`ReplaceViewCopyWithViewPass`. Contiguous slices taken +along the outermost dimension with unit step alias a sub-region of the base +buffer and can be represented with :class:`_SliceSpec`, which shares the +base's ``mem_id`` but uses a computed byte ``mem_offset``. """ import logging +from typing import Any, List, Optional import torch +from executorch.exir import memory from executorch.exir.dialects._ops import ops +from executorch.exir.sym_util import eval_shape +from executorch.exir.tensor import ( + contiguous_stride_from_shape, + determine_tensor_dynanism, + dim_order_from_stride, + TensorSpec, +) from torch.fx.passes.infra.pass_base import PassBase, PassResult logger: logging.Logger = logging.getLogger(__name__) +_SLICE_OP = memory.slice + def _is_slice_copy(node: torch.fx.Node) -> bool: return node.op == "call_function" and node.target in ( @@ -46,26 +41,21 @@ def _is_slice_copy(node: torch.fx.Node) -> bool: ) -def is_contiguous_slice_copy(node: torch.fx.Node) -> bool: - """Return True if ``node`` is a ``slice_copy`` whose result is a contiguous - sub-region of a contiguous input, and is therefore eligible to be - re-inplaced as a zero-copy slice. +def _normalize_dim(dim: int, rank: Optional[int]) -> Optional[int]: + if isinstance(dim, int) and dim < 0: + if rank is None: + return None + dim = dim + rank + return dim - A slice ``self[start:end:step]`` along ``dim`` is a contiguous sub-buffer of - a contiguous input only when it is taken along the outermost (first) storage - dimension with unit step. Slicing an inner dimension, or using ``step > 1``, - produces a strided (non-contiguous) result that cannot alias the base buffer - without a copy. - Signature: ``slice_copy.Tensor(self, dim=0, start=None, end=None, step=1)``. - """ +def is_contiguous_slice_copy(node: torch.fx.Node) -> bool: + """True if ``node`` is an outermost-dim, unit-step ``slice_copy``.""" if not _is_slice_copy(node): return False args = node.args self_arg = args[0] - - # dim defaults to 0; normalize negatives against the input rank. dim = args[1] if len(args) > 1 else 0 step = args[4] if len(args) > 4 else 1 @@ -77,47 +67,235 @@ def is_contiguous_slice_copy(node: torch.fx.Node) -> bool: if self_val is not None and hasattr(self_val, "dim"): rank = self_val.dim() - if isinstance(dim, int) and dim < 0: - if rank is None: - # Cannot resolve a negative dim to the outermost dim without rank. - return False - dim = dim + rank - - # Only an outermost-dim, unit-step slice of a contiguous input is a - # contiguous sub-buffer that could alias the base storage. + dim = _normalize_dim(dim, rank) return dim == 0 -class ReplaceSliceCopyWithSlicePass(PassBase): - """Re-inplace eligible (contiguous) ``slice_copy`` nodes as lightweight - slices. +def _slice_start_as_int(start: Any) -> int: + if start is None: + return 0 + if isinstance(start, int): + return start + if isinstance(start, torch.SymInt): + return int(eval_shape([start])[0]) + return int(start) + + +def _compute_slice_byte_offset(base: TensorSpec, dim: int, start: Any) -> int: + start_int = _slice_start_as_int(start) + if start_int < 0: + raise ValueError("memory.slice does not support negative slice starts.") + elem_size = torch._utils._element_size(base.dtype) + return start_int * base.stride[dim] * elem_size + + +def _is_aliasing_base(base: torch.fx.Node) -> bool: + """Whether ``base`` is itself an alias rather than a real allocation. + + ``memory.slice`` and ``memory.view`` nodes do not own storage, so a slice + taken from one has no concrete ``mem_offset`` to build on until the chain is + normalized down to the first real allocation. + """ + return base.op == "call_function" and base.target in (memory.slice, memory.view) + + +def _has_default_dim_order(spec: TensorSpec) -> bool: + """Whether ``spec`` has the standard contiguous dimension ordering. - Until offset-based sub-buffer aliasing lands in memory planning (see the - module docstring and #10917), this pass only *identifies* eligible nodes and - performs no graph mutation, so it is safe to run in the pipeline. + ``_SliceSpec`` computes a contiguous output stride. That is only a valid + alias for a dim-0 slice when the base itself has the default dim order. """ + return spec.dim_order == dim_order_from_stride( + contiguous_stride_from_shape(torch.Size(spec.shape)) + ) + + +class _SliceSpec(TensorSpec): + """TensorSpec for a zero-copy slice into a contiguous base buffer.""" + + def __init__( + self, + base: TensorSpec, + shape: List[int], + dim: int, + start: Any, + ) -> None: + if base.is_sparse: + raise Exception( + "_SliceSpec can only be created from non-sparse TensorSpec." + ) + if base.layout != torch.strided: + raise Exception(f"_SliceSpec requires strided layout, got {base.layout}.") + + self._base = base + self._byte_offset = _compute_slice_byte_offset(base, dim, start) + self._unguarded_access = False + + self._self_fields = [ + "debug", + "__repr__", + "shape", + "stride", + "dim_order", + "shape_dynamism", + "nbytes", + "allocated_memory", + "is_dynamic_shape_tensor", + "is_static_shape_tensor", + "is_upper_bound_tensor", + "is_dynamic_unbound_tensor", + "mem_offset", + ] + self._base_fields = [ + "scalar_type", + "const", + "alignment", + "storage", + "requires_grad", + "layout", + "is_sparse", + "init_mem_planning_fields", + "realign", + "from_tensor", + "lifetime", + "mem_id", + "mem_obj_id", + "dtype", + "extra_tensor_info", + "device", + "device_index", + # Read by the memory planning algorithms (e.g. ``greedy``). A slice + # is never itself an in-place target, so it defers to its base. + "inplace_base", + ] + + self.shape = list(shape) + self.stride = contiguous_stride_from_shape(torch.Size(self.shape)) + self.dim_order = dim_order_from_stride(self.stride) + self.shape_dynamism = determine_tensor_dynanism(torch.Size(self.shape)) + + if self.shape_dynamism != base.shape_dynamism: + raise Exception( + f"_SliceSpec shape_dynamism {self.shape_dynamism} != base {base.shape_dynamism}" + ) + if self.dtype != base.dtype: + raise Exception(f"_SliceSpec dtype {self.dtype} != base {base.dtype}") + + def __getattribute__(self, name: str): # pyre-ignore + if name in [ + "_base", + "_self_fields", + "_base_fields", + "_byte_offset", + "_unguarded_access", + ]: + return object.__getattribute__(self, name) + + self_fields = object.__getattribute__(self, "_self_fields") + base_fields = object.__getattribute__(self, "_base_fields") + + if name == "mem_offset": + base = object.__getattribute__(self, "_base") + base_offset = base.mem_offset + if base_offset is None: + return None + byte_offset = object.__getattribute__(self, "_byte_offset") + return base_offset + byte_offset + + if name in self_fields: + if name in ("nbytes", "allocated_memory"): + return TensorSpec.__getattribute__(self, name) + return object.__getattribute__(self, name) + + if name in base_fields: + base = object.__getattribute__(self, "_base") + return object.__getattribute__(base, name) + + return object.__getattribute__(self, name) + + def __setattr__(self, name: str, val) -> None: # pyre-ignore + if name in [ + "_base", + "_self_fields", + "_base_fields", + "_byte_offset", + "_unguarded_access", + ]: + object.__setattr__(self, name, val) + return + + if hasattr(self, "_self_fields") and name in self._self_fields: + if name == "mem_offset": + raise Exception("_SliceSpec.mem_offset is computed from the base.") + object.__setattr__(self, name, val) + return + + if hasattr(self, "_base_fields") and name in self._base_fields: + object.__setattr__(self._base, name, val) + return + + object.__setattr__(self, name, val) + + +class ReplaceSliceCopyWithSlicePass(PassBase): + """Replace eligible contiguous ``slice_copy`` nodes with ``memory.slice``.""" def __init__(self) -> None: super().__init__() def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - n_eligible = 0 + n_replaced = 0 for module in graph_module.modules(): if not isinstance(module, torch.fx.GraphModule): continue for node in module.graph.nodes: - # A slice feeding the graph output can have its pointer modified - # at runtime, mirroring the view_copy pass's output guard. if is_contiguous_slice_copy(node) and all( u.op != "output" for u in node.users ): - n_eligible += 1 + base = node.args[0] + if ( + not isinstance(base, torch.fx.Node) + or "spec" not in base.meta + or not base.meta["spec"].is_static_shape_tensor + or not _has_default_dim_order(base.meta["spec"]) + ): + # Specs are populated by the lowering pipeline before this + # pass. Skip bare FX graphs so the pass remains safe to use + # in isolation as well. + continue + if _is_aliasing_base(base): + # The base is itself an alias (a slice or a view), so it + # has no allocation of its own to offset from. Chaining + # offsets through it would require normalizing to the + # first real allocation first, so leave this as a copy. + continue + dim = node.args[1] if len(node.args) > 1 else 0 + start = node.args[2] if len(node.args) > 2 else None + if _slice_start_as_int(start) < 0: + # Negative starts are relative to the end of the + # dimension. They cannot be expressed as a static + # offset without normalizing against the base shape. + continue + node.target = _SLICE_OP + shape = node.meta["val"].shape + node.meta["spec"] = _SliceSpec( + base.meta["spec"], list(shape), dim, start + ) + n_replaced += 1 + + module.recompile() logger.debug( - "ReplaceSliceCopyWithSlicePass: %d contiguous slice_copy node(s) " - "eligible for re-inplacing (rewrite pending offset-aliasing support, " - "#10917).", - n_eligible, + "ReplaceSliceCopyWithSlicePass: replaced %d slice_copy node(s) with %s.", + n_replaced, + _SLICE_OP, ) - # No mutation yet -> report unchanged. - return PassResult(graph_module, False) + return PassResult(graph_module, n_replaced > 0) + + def ensures(self, graph_module: torch.fx.GraphModule) -> None: + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in module.graph.nodes: + if node.op == "call_function" and node.target == _SLICE_OP: + assert isinstance(node.meta["spec"], _SliceSpec) diff --git a/exir/program/_program.py b/exir/program/_program.py index 0e63266e663..56346a005cd 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -70,6 +70,9 @@ ) from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge +from executorch.exir.passes.replace_slice_copy_with_slice_pass import ( + ReplaceSliceCopyWithSlicePass, +) from executorch.exir.passes.replace_view_copy_with_view_pass import ( ReplaceViewCopyWithViewPass, ) @@ -748,6 +751,7 @@ def pre_memory_planning_passes( NormalizeViewCopyBasePass(), dead_code_elimination_pass, ReplaceViewCopyWithViewPass(), + ReplaceSliceCopyWithSlicePass(), sym_shape_eval_pass, config.to_out_var_pass, ] diff --git a/exir/serde/export_serialize.py b/exir/serde/export_serialize.py index 572d87f2dec..954cefc1905 100644 --- a/exir/serde/export_serialize.py +++ b/exir/serde/export_serialize.py @@ -212,6 +212,7 @@ def _reverse_map(d: Dict[Any, Enum]): _KNOWN_FUNCTIONS = { exir.memory.view, + exir.memory.slice, } diff --git a/exir/serde/serialize.py b/exir/serde/serialize.py index a2eb2491067..849f7d425ad 100644 --- a/exir/serde/serialize.py +++ b/exir/serde/serialize.py @@ -376,6 +376,7 @@ def serialize( _KNOWN_FUNCTIONS_MAP = { "executorch.exir.memory.view": exir.memory.view, + "executorch.exir.memory.slice": exir.memory.slice, } diff --git a/exir/tests/test_replace_slice_copy_with_slice_pass.py b/exir/tests/test_replace_slice_copy_with_slice_pass.py index 95751eea6ca..77f0425f250 100644 --- a/exir/tests/test_replace_slice_copy_with_slice_pass.py +++ b/exir/tests/test_replace_slice_copy_with_slice_pass.py @@ -7,15 +7,22 @@ # pyre-strict import unittest +from typing import List import torch -from executorch.exir import to_edge +from executorch.exir import memory, to_edge from executorch.exir.passes.replace_slice_copy_with_slice_pass import ( + _compute_slice_byte_offset, _is_slice_copy, is_contiguous_slice_copy, ReplaceSliceCopyWithSlicePass, ) +from executorch.exir.tensor import TensorSpec +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) from torch.export import export +from torch.testing import assert_close class TestReplaceSliceCopyWithSlicePass(unittest.TestCase): @@ -55,20 +62,136 @@ def forward(self, x): eligible = [n for n in gm.graph.nodes if is_contiguous_slice_copy(n)] self.assertEqual(len(eligible), 1) - def test_pass_is_safe_noop_until_offset_aliasing_lands(self) -> None: - """The pass must run cleanly and not mutate the graph while the - offset-aliasing rewrite is still gated (see #10917).""" + def _annotate_input_spec(self, gm: torch.fx.GraphModule) -> None: + """Populate ``spec`` on every tensor placeholder. + + The lowering pipeline normally does this before the pass runs. Note + that ``to_edge`` lifts scalar constants to their own placeholders, so + annotating only the first placeholder would miss the real input. + """ + for node in gm.graph.nodes: + if node.op != "placeholder": + continue + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + node.meta["spec"] = TensorSpec.from_tensor(val) + + def test_pass_replaces_annotated_contiguous_slice(self) -> None: + """A statically annotated dim-0 slice becomes a memory alias.""" class M(torch.nn.Module): def forward(self, x): return x[0:2] + 1.0 gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) - before = gm.code + self._annotate_input_spec(gm) result = ReplaceSliceCopyWithSlicePass()(gm) self.assertIsNotNone(result) + self.assertTrue(result.modified) + self.assertEqual( + len( + [ + n + for n in result.graph_module.graph.nodes + if n.op == "call_function" and n.target == memory.slice + ] + ), + 1, + ) + + def test_pass_skips_nondefault_base_dim_order(self) -> None: + """Avoid aliases that would reinterpret a non-contiguous base layout.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[0:2] + 1.0 + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + self._annotate_input_spec(gm) + # Mutate the layout of the slice's own base, not just any placeholder. + slice_node = next(n for n in gm.graph.nodes if _is_slice_copy(n)) + slice_node.args[0].meta["spec"].dim_order = (1, 0) + + result = ReplaceSliceCopyWithSlicePass()(gm) + self.assertFalse(result.modified) + + def test_pass_skips_negative_start(self) -> None: + """Negative starts need shape-dependent normalization, so keep copying.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[-2:] + 1.0 + + gm = self._edge_graph_module(M(), (torch.randn(4, 8),)) + self._annotate_input_spec(gm) + + result = ReplaceSliceCopyWithSlicePass()(gm) self.assertFalse(result.modified) - self.assertEqual(before, result.graph_module.code) + with self.assertRaises(ValueError): + _compute_slice_byte_offset( + next(n for n in gm.graph.nodes if n.op == "placeholder").meta["spec"], + 0, + -2, + ) + + def _emitted_operators(self, program) -> List[str]: + return [ + str(op) for op in program.executorch_program.execution_plan[0].operators + ] + + def test_lowered_program_matches_eager_output(self) -> None: + """The emitted sub-buffer alias executes with the original semantics.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[1:3] + 1.0 + + model = M().eval() + example_input = torch.arange(32, dtype=torch.float32).reshape(4, 8) + et = to_edge(export(model, (example_input,), strict=True)).to_executorch() + + # The slice must be aliased away, not merely produce the right answer -- + # falling back to a copy would also pass a numerical check alone. + self.assertFalse( + any("slice_copy" in op for op in self._emitted_operators(et)), + "expected the contiguous slice to be elided, but slice_copy was emitted", + ) + + runtime_module = _load_for_executorch_from_buffer(et.buffer) + assert_close(runtime_module.forward((example_input,))[0], model(example_input)) + + def test_base_outlives_slice_when_reused(self) -> None: + """The base buffer must not be reused while the alias is still live.""" + + class M(torch.nn.Module): + def forward(self, x): + sliced = x[1:3] + 1.0 + # ``x`` is consumed *after* the slice, so the planner has to keep + # the base alive across the alias's lifetime. + return sliced.sum() + x.sum() + + model = M().eval() + example_input = torch.arange(32, dtype=torch.float32).reshape(4, 8) + et = to_edge(export(model, (example_input,), strict=True)).to_executorch() + runtime_module = _load_for_executorch_from_buffer(et.buffer) + + assert_close(runtime_module.forward((example_input,))[0], model(example_input)) + + def test_chained_slice_falls_back_to_copy(self) -> None: + """A slice of a slice has no concrete base allocation to offset from.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[0:3][1:2] + 1.0 + + model = M().eval() + example_input = torch.arange(32, dtype=torch.float32).reshape(4, 8) + # Must lower and execute correctly rather than tripping over an + # aliasing base during memory planning. + et = to_edge(export(model, (example_input,), strict=True)).to_executorch() + runtime_module = _load_for_executorch_from_buffer(et.buffer) + + assert_close(runtime_module.forward((example_input,))[0], model(example_input)) def test_non_slice_nodes_are_ignored(self) -> None: class M(torch.nn.Module):