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

[iRAFEEK]

Summary

A contiguous slice (e.g. x[1:3] on a contiguous input) was always emitted as a full-copy aten::slice_copy kernel, even though it can alias a sub-region of the base buffer — the same idea ReplaceViewCopyWithViewPass already applies to view_copy. On a memory-constrained device that is a wasted allocation and a wasted copy on every inference.

This implements #10917: ReplaceSliceCopyWithSlicePass rewrites eligible contiguous slice_copy nodes into memory.slice aliases, so no copy kernel is emitted.

How it works

• _SliceSpec shares the base's mem_id and computes mem_offset = base.mem_offset + start * base.stride[0] * elem_size. AllocationDetails already carries (memory_id, memory_offset) and mem_offset reaches it through make_allocation_info, so no schema or runtime change is needed. _ViewSpec can't be reused here because it requires nbytes == base.nbytes() — a view aliases the whole buffer, a slice only part of it.
• Memory planning treats memory.slice like memory.view: it's on the collect_specs_from_nodes skip-list, and get_node_tensor_specs returns the base spec. Since update_all_tensors_lifetime walks chain([node], node.args, ...), the base's lifetime is extended to cover the slice's consumers, so the buffer can't be reused while the alias is live.
• Emission mirrors _emit_view's elide path — static and memory-planned specs go straight through _emit_spec, so no runtime kernel is required.

Eligibility. Restricted to dim-0, step == 1, 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, slice-of-view) has no concrete allocation to offset from. Everything outside those gates falls back to slice_copy unchanged, so this is opt-in by construction.

Per @JacobSzwejbka's note on the issue that other dim orders can be a follow-up, v1 keeps to the default layout.

Fixes #10917

Test plan

Copy elimination — x[1:3] + 1.0 lowered with to_edge(...).to_executorch():

before:  aten::slice_copy Tensor_out,  aten::add out
after:   aten::add out

Unit tests (exir/tests/test_replace_slice_copy_with_slice_pass.py, 9 tests) cover eligibility classification, negative-dim resolution, the rewrite itself, non-default dim_order skipped, negative start skipped, chained slices falling back, base-outlives-slice lifetime, and end-to-end parity against eager. The end-to-end test asserts the copy was elided, not just that the numbers match — a fallback to copy would pass a numerical check alone.

Runtime verification against the executorch wheel, comparing _load_for_executorch_from_buffer output to eager with torch.arange inputs so a wrong offset is visible:

[PASS] dim0 start=1                 elide=True
[PASS] dim0 start=0                 elide=True
[PASS] dim0 last rows               elide=True
[PASS] two slices in one graph      elide=True
[PASS] base reused after slice      (lifetime)
[PASS] slice consumed late          (lifetime)
[PASS] inner-dim slice              elide=False  -> falls back
[PASS] strided step=2               elide=False  -> falls back
[PASS] negative start               elide=False  -> falls back
[PASS] slice of slice               -> falls back
[PASS] 3-D tensor
11/11 numerically correct

No regressions. exir/tests, exir/emit, and exir/backend/test were run against both this change and a pristine executorch install; the failure sets are identical (the pre-existing failures are missing quantized out-variants and backend runtime pieces in the wheel). exir/tests/test_memory_planning.py 37 passed, exir/emit/test 69 passed.

lintrunner (including MYPY) is clean with no patch to apply.

Checklist

• Contiguous dim-0 slices emit no copy kernel
• Output matches eager, including non-zero offsets
• Base buffer stays live across the alias's lifetime
• Non-default dim_order falls back to copy
• Negative start falls back to copy
• Aliasing bases (slice-of-slice / slice-of-view) fall back to copy
• Inner-dim and strided slices unchanged
• No regressions vs a pristine baseline
• lintrunner clean
• exir/passes/BUCK + exir/tests/targets.bzl registration — happy to add once the approach is confirmed
• Other dim orders — follow-up per Re-inplace slice_copy with slice #10917 discussion

cc @JacobSzwejbka @metascroy

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21552

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 13 Awaiting Approval

As of commit a18eb2a with merge base 0bbb6f3 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: iRAFEEK / name: iRAFEEK (5a3e4a3, a3fef4f, a5fda43, c24c7dd, db374f2, 1b6d0b9)

[meta-cla]

Hi @iRAFEEK!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[iRAFEEK]

@pytorchbot label "release notes: none"

[iRAFEEK]

@metascroy — CI workflows are awaiting maintainer approval to run. Could you approve them when you get a chance? Thanks!

[nil-is-all]

Thanks for the PR, @iRAFEEK. Running CI now.

[iRAFEEK]

@nil-is-all — pushed a formatting fix (db374f2) for the lintrunner failure. lintrunner -a reports no issues locally, and the unit tests pass. Could you approve CI when you get a chance? Thanks!

[iRAFEEK]

@nil-is-all @JacobSzwejbka could you please check it out now , thank you so much

[nil-is-all]

@nil-is-all — pushed a formatting fix (db374f2) for the lintrunner failure. lintrunner -a reports no issues locally, and the unit tests pass. Could you approve CI when you get a chance? Thanks!

Sure, thanks. Running CI again

[iRAFEEK]

@JacobSzwejbka — design check before I build out the runtime side of this.

The blocker looks structural: _ViewSpec can't represent a slice, since it raises when nbytes != base.nbytes() (replace_view_copy_with_view_pass.py#L188) — a view aliases the whole base at offset 0, while a slice aliases a sub-region at a non-zero byte offset. So it needs its own spec type. But the surrounding infrastructure already supports offsets, so I don't think anything new is needed at the format or runtime layer. What I'd propose:

1. _SliceSpec — shares the base's mem_id, with mem_offset = base.mem_offset + start * base.stride[0] * elem_size. AllocationDetails already carries (memory_id, memory_offset), and mem_offset reaches it through make_allocation_info (_emitter.py#L359), so no schema change.

2. Memory planning — add memory.slice alongside memory.view in the collect_specs_from_nodes skip-list and in get_node_tensor_specs (returning the base spec). Since update_all_tensors_lifetime walks chain([node], node.args, ...), that extends the base's lifetime to cover the slice's consumers, so the buffer can't be reused while the slice is live.

3. Emission — mirror _emit_view's elide path: static + memory-planned → return self._emit_spec(spec). No new kernel, since the base's producer has already written those bytes.

4. Eligibility (v1) — dim-0, step == 1, non-negative start, and a base with default dim_order; anything else falls back to slice_copy. Per your earlier note that other dim orders can be a follow-up, I'd keep v1 to the default — forcing contiguous stride on a channels-last base would silently produce wrong values rather than fail, so I'd rather gate it explicitly.

Two things I'd like a steer on:

• _emit_view elides when static+planned and otherwise emits executorch_prim::et_view. For v1 I'd implement only the elide path, leaving dynamic/non-planned slices as slice_copy and deferring an et_slice kernel entirely. Reasonable, or would you rather have the kernel fallback up front?
• Skip aliasing bases (slice-of-slice, slice-of-view) in v1, or add a NormalizeSliceCopyBasePass mirroring NormalizeViewCopyBasePass?

Unless you object, I'll build it as described above and update this PR. I have the pass and memory-planning wiring prototyped locally and am filling in the correctness tests now (numerical parity against eager, plus negative cases for strided/inner-dim/channels-last bases falling back to copy).

[iRAFEEK]

@nil-is-all @JacobSzwejbka could you please check it out now? Thank you so much

[iRAFEEK]

@nil-is-all @JacobSzwejbka Just following up if u can run the tests, please.