feat: Re-inplace contiguous slice_copy as zero-copy memory.slice aliases (#10917)

exir/BUCK

@@ -198,6 +198,7 @@ fbcode_target(_kind = runtime.python_library,
         ":tensor",
         "//caffe2:torch",
         "//executorch/exir/operator:convert",
+        "//executorch/exir/passes:replace_view_copy_with_view_pass",
     ],
 )
 

exir/emit/_emitter.py

@@ -1301,6 +1301,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 ``_ViewSpec`` 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,
@@ -1872,6 +1886,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

exir/memory.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 from executorch.exir.delegate import executorch_call_delegate
@@ -89,3 +89,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)

exir/memory_planning.py

@@ -503,6 +503,7 @@ def collect_specs_from_nodes(  # noqa: C901
                 in [
                     memory.alloc,
                     memory.view,
+                    memory.slice,
                     operator.getitem,
                     torch.ops.higher_order.cond,
                     exir_while,
@@ -785,8 +786,10 @@ def get_node_tensor_specs(
     node produces, but every caller that walks specs to place, verify or
     attribute memory needs them.
     """
+    from executorch.exir.passes.replace_view_copy_with_view_pass import _ViewSpec
+
     # get tensor specs
-    if node.target == memory.view:
+    if node.target in (memory.view, memory.slice):
         base = node.args[0]
         assert isinstance(base, torch.fx.Node)
         specs = base.meta.get("spec")
@@ -803,6 +806,12 @@ def get_node_tensor_specs(
             for spec in specs
             if not isinstance(spec, (int, float, bool, str, type(None)))
         ]
+        allocation_specs = []
+        for spec in specs:
+            while isinstance(spec, _ViewSpec):
+                spec = spec._base
+            allocation_specs.append(spec)
+        specs = allocation_specs
 
     if scratch := memory.delegate_scratch_specs(node):
         return list(specs) + scratch

exir/pass_base.py

@@ -1007,7 +1007,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,

exir/passes/BUCK

@@ -29,6 +29,7 @@ fbcode_target(_kind = runtime.python_library,
         ":replace_aten_with_edge_pass",
         ":replace_broken_ops_with_function_ops_pass",
         ":replace_edge_with_backend_pass",
+        ":replace_slice_copy_with_slice_pass",
         ":replace_sym_size_op_pass",
         ":scalar_to_tensor_pass",
         ":spec_prop_pass",
@@ -459,6 +460,20 @@ fbcode_target(_kind = runtime.python_library,
     ],
 )
 
+fbcode_target(_kind = runtime.python_library,
+    name = "replace_slice_copy_with_slice_pass",
+    srcs = [
+        "replace_slice_copy_with_slice_pass.py",
+    ],
+    deps = [
+        ":replace_view_copy_with_view_pass",
+        "//caffe2:torch",
+        "//executorch/exir:memory",
+        "//executorch/exir:tensor",
+        "//executorch/exir/dialects:lib",
+    ],
+)
+
 fbcode_target(_kind = runtime.python_library,
     name = "replace_view_copy_with_view_pass",
     srcs = [

exir/passes/__init__.py

@@ -269,6 +269,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)

exir/passes/replace_slice_copy_with_slice_pass.py

@@ -0,0 +1,169 @@
+# 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
+
+"""Replace safe, static, contiguous slice copies with sub-buffer aliases."""
+
+import logging
+from typing import Any, Optional
+
+import torch
+from executorch.exir import memory
+from executorch.exir.dialects._ops import ops
+from executorch.exir.passes.replace_view_copy_with_view_pass import (
+    _ViewSpec,
+    is_copy_to_view_safe,
+)
+from executorch.exir.tensor import (
+    contiguous_stride_from_shape,
+    dim_order_from_stride,
+    TensorSpec,
+)
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from torch.utils import _pytree as pytree
+
+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_static_slice_argument(value: Any) -> bool:
+    return value is None or isinstance(value, int)
+
+
+def is_contiguous_slice_copy(node: torch.fx.Node) -> bool:
+    """True for static-argument, outermost-dimension, unit-step slices."""
+    if (
+        not _is_slice_copy(node)
+        or node.kwargs
+        or not all(_is_static_slice_argument(arg) for arg in node.args[1:])
+    ):
+        return False
+    dim = node.args[1] if len(node.args) > 1 else 0
+    step = node.args[4] if len(node.args) > 4 else 1
+    base = node.args[0]
+    val = base.meta.get("val") if isinstance(base, torch.fx.Node) else None
+    if not isinstance(dim, int) or step != 1:
+        return False
+    if dim < 0 and isinstance(val, torch.Tensor):
+        dim += val.dim()
+    return dim == 0
+
+
+def _compute_slice_byte_offset(base: TensorSpec, dim: int, start: Optional[int]) -> int:
+    start, _, _ = slice(start, None, 1).indices(base.shape[dim])
+    return start * base.stride[dim] * torch._utils._element_size(base.dtype)
+
+
+def _static_slice_bounds(node: torch.fx.Node) -> Optional[tuple[int, int]]:
+    base = node.args[0]
+    if not isinstance(base, torch.fx.Node) or base.op == "placeholder":
+        return None
+    # Retain copies of aliases, including pending slice candidates.
+    if _is_slice_copy(base) or base.target in (memory.slice, memory.view):
+        return None
+    base_spec = base.meta.get("spec")
+    output_spec = node.meta.get("spec")
+    if (
+        not isinstance(base_spec, TensorSpec)
+        or not isinstance(output_spec, TensorSpec)
+        or not base_spec.is_static_shape_tensor
+        or not output_spec.is_static_shape_tensor
+        or output_spec.dtype != base_spec.dtype
+        or base_spec.const
+        or base_spec.is_sparse
+        or base_spec.layout != torch.strided
+        or base_spec.dim_order
+        != dim_order_from_stride(
+            contiguous_stride_from_shape(torch.Size(base_spec.shape))
+        )
+    ):
+        return None
+    start = node.args[2] if len(node.args) > 2 else None
+    end = node.args[3] if len(node.args) > 3 else None
+    start, end, _ = slice(start, end, 1).indices(base_spec.shape[0])
+    shape = [max(end - start, 0), *base_spec.shape[1:]]
+    # Empty results retain their kernel; do not emit an alias at
+    # the end of (or into an empty) allocation.
+    if 0 in shape or list(output_spec.shape) != shape:
+        return None
+    return start, end
+
+
+class ReplaceSliceCopyWithSlicePass(PassBase):
+    """Replace eligible slice copies after view-copy replacement."""
+
+    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
+            replacements = {}
+            # Analyze consumers first, with already-decided aliasing behavior.
+            # Specs are rebuilt in forward order below so views of slices use
+            # the final base spec, including its byte offset.
+            for node in reversed(module.graph.nodes):
+                if not is_contiguous_slice_copy(node) or any(
+                    u.op == "output" for u in node.users
+                ):
+                    continue
+                bounds = _static_slice_bounds(node)
+                if bounds is None:
+                    continue
+                start, end = bounds
+                base = node.args[0]
+                base_spec = base.meta["spec"]
+                if not is_copy_to_view_safe(node, (memory.view, memory.slice)):
+                    continue
+                replacements[node] = _compute_slice_byte_offset(base_spec, 0, start)
+                node.target = memory.slice
+                node.args = (base, 0, start, end, 1)
+                n_replaced += 1
+
+            updated_specs = set()
+            for node in module.graph.nodes:
+                old_spec = node.meta.get("spec")
+                if node in replacements:
+                    node.meta["spec"] = _ViewSpec(
+                        node.args[0].meta["spec"],
+                        list(node.meta["spec"].shape),
+                        byte_offset=replacements[node],
+                    )
+                    updated_specs.add(node)
+                elif node.target == memory.view and node.args[0] in updated_specs:
+                    node.meta["spec"] = _ViewSpec(
+                        node.args[0].meta["spec"], list(node.meta["spec"].shape)
+                    )
+                    updated_specs.add(node)
+                else:
+                    continue
+                # Preserve specs shared with in-place consumers and outputs.
+                new_spec = node.meta["spec"]
+                for other_node in module.graph.nodes:
+                    if "spec" in other_node.meta:
+                        other_node.meta["spec"] = pytree.tree_map(
+                            lambda spec, old=old_spec, new=new_spec: (
+                                new if spec is old else spec
+                            ),
+                            other_node.meta["spec"],
+                        )
+            module.recompile()
+
+        logger.debug("Replaced %d slice_copy nodes with memory.slice", n_replaced)
+        return PassResult(graph_module, n_replaced > 0)
+
+    def ensures(self, graph_module: torch.fx.GraphModule) -> None:
+        for module in graph_module.modules():
+            if isinstance(module, torch.fx.GraphModule):
+                for node in module.graph.nodes:
+                    if node.target == memory.slice:
+                        assert isinstance(node.meta["spec"], _ViewSpec)