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 new file mode 100644 index 00000000000..4799e98b3d2 --- /dev/null +++ b/exir/passes/replace_slice_copy_with_slice_pass.py @@ -0,0 +1,301 @@ +# 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). + +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 ( + torch.ops.aten.slice_copy.Tensor, + ops.edge.aten.slice_copy.Tensor, + ) + + +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 + + +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 = 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() + + dim = _normalize_dim(dim, rank) + return dim == 0 + + +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. + + ``_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_replaced = 0 + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in module.graph.nodes: + if is_contiguous_slice_copy(node) and all( + u.op != "output" for u in node.users + ): + 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: replaced %d slice_copy node(s) with %s.", + n_replaced, + _SLICE_OP, + ) + 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 new file mode 100644 index 00000000000..77f0425f250 --- /dev/null +++ b/exir/tests/test_replace_slice_copy_with_slice_pass.py @@ -0,0 +1,206 @@ +# 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 +from typing import List + +import torch +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): + 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 _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),)) + 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) + 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): + 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()