Convert contiguous select_copy to zero-copy view in ReplaceViewCopyWithViewPass (pytorch#19198)

Summary:

Extends the ReplaceViewCopyWithViewPass to convert `select_copy` ops to zero-copy `memory.select` views when the output is a contiguous sub-region of the base tensor. This is the same pattern used for `view_copy` -> `memory.view`, but for select operations.

The pass checks that the base is densely packed, non-constant, and static, and that the selected output forms a dense packing. It uses the base spec dim_order to compute actual memory strides, so this works for any contiguous layout (C-contiguous, channels-last, etc.).

For static memory-planned subviews, the emitter elides the op entirely (no runtime instruction) by serializing tensor metadata with `mem_offset = base_offset + byte_delta`. For dynamic shapes, a new `executorch_prim::et_select` runtime op sets the output data pointer to `self.data_ptr + offset`.

Changes:
- `exir/memory.py`: Added `memory.select` function
- `exir/passes/replace_view_copy_with_view_pass.py`: Extended `_ViewSpec` with `byte_offset`, `stride`, `dim_order` params; added contiguity check using `stride_from_dim_order`; added select_copy handling in the pass
- Pipeline integration: memory planner, to_out_var skiplist, emitter, serialization
- `kernels/prim_ops/et_select.{h,cpp}`: C++ runtime op for dynamic select views
- Tests: 5 new Python tests + 1 C++ test

Authored with Claude.

Differential Revision: D102396195

Author: JacobSzwejbka

exir/emit/_emitter.py

@@ -1259,6 +1259,42 @@ def _emit_view(self, args: Tuple[_Argument, ...]) -> _EmitterValue:
         self.chain.instructions.append(kernel)
         return out_arg
 
+    def _emit_subview(self, args: Tuple[_Argument, ...]) -> _EmitterValue:
+        spec = self.node.meta["spec"]
+        is_static = spec.is_static_shape_tensor
+        is_memory_planned = (spec.mem_id is not None) and (spec.mem_offset is not None)
+        if is_static and is_memory_planned:
+            return self._emit_spec(spec)
+
+        out_arg = self._emit_argument(
+            self._emit_spec(self.node.meta["spec"]), torch.TensorType  # pyre-ignore[6]
+        )
+
+        if self.node.target == memory.select:
+            self_arg = self._emit_argument(args[0], torch.TensorType)  # pyre-ignore[6]
+            dim_arg = self._emit_argument(args[1], torch.IntType)
+            index_arg = self._emit_argument(args[2], torch.IntType)
+            op_idx, op = self._get_operator(
+                name="executorch_prim::et_select",
+                overload="default",
+            )
+            kernel = Instruction(
+                KernelCall(
+                    op_idx,
+                    args=[self_arg.id, dim_arg.id, index_arg.id, out_arg.id],
+                )
+            )
+        else:
+            raise InternalError(
+                self._emit_node_specific_error(
+                    self.node,
+                    f"Unsupported subview target: {self.node.target}",
+                )
+            )
+
+        self.chain.instructions.append(kernel)
+        return out_arg
+
     def _add_debug_handle(
         self,
         emitter_id: int,
@@ -1758,6 +1794,9 @@ def call_function(  # pyre-fixme[14]
         elif target == memory.view:
             return self._emit_view(args)
 
+        elif target == memory.select:
+            return self._emit_subview(args)
+
         elif target == memory.free:
             assert len(args) == 1
             # pyre-ignore

exir/memory.py

@@ -48,3 +48,12 @@ def view(base: torch.Tensor, size: List[int]) -> torch.Tensor:
     It is used to elide view_copy nodes.
     """
     return base.view(size)
+
+
+def select(base: torch.Tensor, dim: int, index: int) -> torch.Tensor:
+    """
+    This function mimics torch.ops.aten.select.int.
+
+    It is used to elide select_copy nodes when the result is contiguous.
+    """
+    return base.select(dim, index)

exir/memory_planning.py

@@ -486,6 +486,7 @@ def collect_specs_from_nodes(  # noqa: C901
                 in [
                     memory.alloc,
                     memory.view,
+                    memory.select,
                     operator.getitem,
                     torch.ops.higher_order.cond,
                     exir_while,
@@ -763,7 +764,7 @@ def get_node_tensor_specs(
     has no tensor specs.
     """
     # get tensor specs
-    if node.target == memory.view:
+    if node.target in (memory.view, memory.select):
         base = node.args[0]
         assert isinstance(base, torch.fx.Node)
         specs = base.meta.get("spec")

exir/passes/__init__.py

@@ -261,6 +261,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.select,
     executorch_call_delegate,
 }
 to_out_var_skiplist.update(_EXECUTORCH_SYM_OPS)

exir/passes/replace_view_copy_with_view_pass.py

@@ -33,7 +33,27 @@ def _is_view_copy(node: torch.fx.Node) -> bool:
     )
 
 
+def _is_select_copy(node: torch.fx.Node) -> bool:
+    return node.op == "call_function" and node.target in (
+        torch.ops.aten.select_copy.int,
+        ops.edge.aten.select_copy.int,
+    )
+
+
 _VIEW_OP = memory.view
+_SELECT_OP = memory.select
+
+
+def _is_contiguous_base(spec: TensorSpec) -> bool:
+    if not isinstance(spec, TensorSpec):
+        return False
+    if spec.const:
+        return False
+    shape = list(spec.shape)
+    if tuple(spec.dim_order) != tuple(range(len(shape))):
+        return False
+    expected_stride = contiguous_stride_from_shape(torch.Size(shape))
+    return tuple(spec.stride) == tuple(expected_stride)
 
 
 class _Guard:
@@ -54,7 +74,12 @@ def __call__(self, view_spec) -> None:  # pyre-ignore[2]
 
 
 class _ViewSpec(TensorSpec):
-    def __init__(self, base: TensorSpec, shape: List[int]) -> None:
+    def __init__(
+        self,
+        base: TensorSpec,
+        shape: List[int],
+        byte_offset: int = 0,
+    ) -> None:
         """
         A _ViewSpec is TensorSpec that shares non-size related fields with its base.
         The size-related fields are: shape, stride, dim_order, and shape_dynamism.
@@ -65,7 +90,11 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None:
 
         A _ViewSpec can only be created from a non-sparse, strided TensorSpec.
         On creation, a _ViewSpec must be compatible with its base with respect to
-        shape_dynamism, dtype, and nbytes.
+        shape_dynamism, dtype, and nbytes (when byte_offset is 0).
+
+        When byte_offset is non-zero (used for select/slice sub-views), the view
+        describes a contiguous sub-region of the base at the given byte offset.
+        In this case, nbytes may differ from the base and rank may change.
 
         A _ViewSpec contains _guards that are evaluated on every __getattribute__ call.
         The purpose of the guards is to make sure the _ViewSpec is still compatible
@@ -119,6 +148,7 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None:
 
         self._guards: List[_Guard] = []
         self._unguarded_access = False
+        self._byte_offset: int = byte_offset
 
         # Make sure base is not sparse and add a guard
         if base.is_sparse:
@@ -183,15 +213,19 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None:
             _Guard("dtype", lambda view_spec: view_spec.dtype, base.dtype)
         )
 
-        # We do not guard nbytes because dynamic symints are replaced by upper bounds.
-        # We do guard on rank, though
-        if self.nbytes() != base.nbytes():
-            raise Exception(
-                f"_ViewSpec is incompatible with its base on creation.  It has nbytes={self.nbytes()}, but its base has nbytes={base.nbytes()}."
+        # For traditional views (same nbytes, zero offset), rank is guarded.
+        # For sub-views (select/slice), the output is a contiguous subset so
+        # nbytes will differ and rank may change.
+        is_full_view = byte_offset == 0 and self.nbytes() == base.nbytes()
+        if is_full_view:
+            self._guards.append(
+                _Guard("rank", lambda view_spec: len(view_spec.shape), len(shape))
             )
-        self._guards.append(
-            _Guard("rank", lambda view_spec: len(view_spec.shape), len(shape))
-        )
+        else:
+            if self.nbytes() + byte_offset > base.nbytes():
+                raise Exception(
+                    f"_ViewSpec sub-view extends beyond base.  Sub-view needs {self.nbytes()} bytes at offset {byte_offset}, but base has {base.nbytes()} bytes."
+                )
 
     def _run_guards(self) -> None:
         unguarded_access = self._unguarded_access
@@ -211,6 +245,7 @@ def __getattribute__(self, name: str):  # pyre-ignore
             "_guards",
             "_unguarded_access",
             "_run_guards",
+            "_byte_offset",
         ]:
             return object.__getattribute__(self, name)
 
@@ -219,6 +254,11 @@ def __getattribute__(self, name: str):  # pyre-ignore
             val = object.__getattribute__(self, name)
         elif name in self._base_fields:
             val = object.__getattribute__(self._base, name)
+            # For sub-views (select/slice), adjust mem_offset by byte_offset
+            if name == "mem_offset" and val is not None:
+                byte_offset = object.__getattribute__(self, "_byte_offset")
+                if byte_offset != 0:
+                    val = val + byte_offset
         else:
             if len(name) > 0 and name[0] != "_":
                 logger.warning(
@@ -239,6 +279,7 @@ def __setattr__(self, name: str, val) -> None:  # pyre-ignore
             "_guards",
             "_unguarded_access",
             "_run_guards",
+            "_byte_offset",
         ]:
             object.__setattr__(self, name, val)
             return
@@ -271,6 +312,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         During memory planning, view nodes share the same storage as their base.
         """
 
+        n_replaced = 0
         n_replaced = 0
         for module in graph_module.modules():
             if not isinstance(module, torch.fx.GraphModule):
@@ -293,11 +335,50 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
                     n_replaced += 1
 
+                elif _is_select_copy(node) and all(
+                    u.op != "output" for u in node.users
+                ):
+                    replaced = self._try_replace_select(node)
+                    if replaced:
+                        n_replaced += 1
+
             module.recompile()
 
         logger.debug(f"Replaced {n_replaced} view_copy nodes with {_VIEW_OP} nodes.")
+        logger.debug(
+            f"Replaced {n_replaced} select_copy nodes with {_SELECT_OP} nodes."
+        )
         return PassResult(graph_module, n_replaced > 0)
 
+    def _try_replace_select(self, node: torch.fx.Node) -> bool:
+        base = node.args[0]
+        assert isinstance(base, torch.fx.Node)
+        base_spec = base.meta["spec"]
+        if not _is_contiguous_base(base_spec):
+            return False
+
+        dim: int = node.args[1]
+        index: int = node.args[2]
+        base_shape = [int(s) for s in base_spec.shape]
+
+        if dim < 0:
+            dim += len(base_shape)
+
+        if any(base_shape[i] != 1 for i in range(dim)):
+            return False
+
+        if index < 0:
+            index += base_shape[dim]
+
+        out_shape = list(node.meta["spec"].shape)
+        base_stride = contiguous_stride_from_shape(torch.Size(base_shape))
+        element_size = torch._utils._element_size(base_spec.dtype)
+        byte_offset = index * base_stride[dim] * element_size
+
+        node.target = _SELECT_OP
+        node.meta["spec"] = _ViewSpec(base_spec, out_shape, byte_offset)
+        return True
+
     def ensures(self, graph_module: torch.fx.GraphModule) -> None:
         for module in graph_module.modules():
             if not isinstance(module, torch.fx.GraphModule):
@@ -310,6 +391,8 @@ def ensures(self, graph_module: torch.fx.GraphModule) -> None:
                 )
                 if node.op == "call_function" and node.target == _VIEW_OP:
                     assert isinstance(node.meta["spec"], _ViewSpec)
+                if node.op == "call_function" and node.target == _SELECT_OP:
+                    assert isinstance(node.meta["spec"], _ViewSpec)
 
     def requires(self, graph_module: torch.fx.GraphModule) -> None:
         """

exir/serde/export_serialize.py

@@ -209,6 +209,7 @@ def _reverse_map(d: Dict[Any, Enum]):
 
 _KNOWN_FUNCTIONS = {
     exir.memory.view,
+    exir.memory.select,
 }
 
 

exir/serde/serialize.py

@@ -376,6 +376,7 @@ def serialize(
 
 _KNOWN_FUNCTIONS_MAP = {
     "executorch.exir.memory.view": exir.memory.view,
+    "executorch.exir.memory.select": exir.memory.select,
 }