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 @angelayi @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.

❌ 1 Unclassified Failure

As of commit e5abcb8 with merge base ffb13bf ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• pull / test-arm-backend-vkml (test_runtime_vgf) / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command docker exec -t 2a8cbd708dc3b40bdd376d6ddf2cde6c6670f2404e6278edf47b071dcaa8c099 /exec failed with exit code 1

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, 1420cee, 1b6d0b9, 6c3dc21, a18eb2a, faf764d)

[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.

[iRAFEEK]

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

[nil-is-all]

Hi @iRAFEEK, extremely sorry for the delay. Could you update with latest main branch and ping us? We will run CI and review right away. Appreciate your contribution!

[iRAFEEK]

@nil-is-all Hi! I updated the branch with the latest main and fixed the slice-copy cases that were failing CI. The PR is ready for CI again—could you please approve and run the workflows when you have a chance? Thank you!

[iRAFEEK]

@nil-is-all, @JacobSzwejbka, could you please run the Tests?!

[nil-is-all]

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

[iRAFEEK]

@nil-is-all Hi! I pushed a formatting-only follow-up commit to address the lintrunner feedback. I also checked every failed job: 139 stop during dependency setup because torchao==0.18.0.dev20260729 is unavailable, before any tests run. Could you please rerun CI once that dependency issue is resolved? Thank you!

[nil-is-all]

@nil-is-all Hi! I pushed a formatting-only follow-up commit to address the lintrunner feedback. I also checked every failed job: 139 stop during dependency setup because torchao==0.18.0.dev20260729 is unavailable, before any tests run. Could you please rerun CI once that dependency issue is resolved? Thank you!

Thanks for addressing CI. We've fixed the dependency issue now. Retriggered CI

[iRAFEEK]

@nil-is-all Hi! I pushed a formatting-only follow-up commit to address the lintrunner feedback. I also checked every failed job: 139 stop during dependency setup because torchao==0.18.0.dev20260729 is unavailable, before any tests run. Could you please rerun CI once that dependency issue is resolved? Thank you!

Thanks for addressing CI. We've fixed the dependency issue now. Retriggered CI

@nil-is-all Thank you for rerunning CI. I checked the rerun and it is still failing during dependency installation with No matching distribution found for torchao==0.18.0.dev20260729, before any build or tests start.

I also updated the branch afterward, so the new head now needs approval. Could you please take another look at the dependency issue and, once it is resolved, approve and rerun CI on the updated branch? Thank you!

[nil-is-all]

@nil-is-all Hi! I pushed a formatting-only follow-up commit to address the lintrunner feedback. I also checked every failed job: 139 stop during dependency setup because torchao==0.18.0.dev20260729 is unavailable, before any tests run. Could you please rerun CI once that dependency issue is resolved? Thank you!

Thanks for addressing CI. We've fixed the dependency issue now. Retriggered CI

@nil-is-all Thank you for rerunning CI. I checked the rerun and it is still failing during dependency installation with No matching distribution found for torchao==0.18.0.dev20260729, before any build or tests start.

I also updated the branch afterward, so the new head now needs approval. Could you please take another look at the dependency issue and, once it is resolved, approve and rerun CI on the updated branch? Thank you!

could you rebase with the main branch and give us a ping for review? We've pushed a fix for dependency issues which caused the above CI failures.

[iRAFEEK]

@nil-is-all Hi! I rebased the branch onto the latest main and pushed the updated history. Could you please take a look when you have a chance? Thank you!

[nil-is-all]

@nil-is-all Hi! I rebased the branch onto the latest main and pushed the updated history. Could you please take a look when you have a chance? Thank you!

Thanks, running CI on latest commit.

[iRAFEEK]

@nil-is-all Hi! I fixed the shared failure in the four macOS/Linux jobs: the quant-fusion test was expecting a zero-copy alias for a slice of an input tensor, but input slices correctly remain slice_copy. The test-only correction is now pushed on the latest commit. Could you please approve and rerun CI when you have a chance? Thank you!

[iRAFEEK]

@nil-is-all Hi! The four unittest jobs, lintrunner, and CLA now pass on the latest commit. The remaining failed job, test-arm-backend-vkml (test_runtime_vgf) / linux-job, stops before running tests with: "Unknown test suite: test_runtime_vgf" from backends/arm/test/test_arm_backend.sh. Since this PR only changes exir slice logic and tests, could you please take a look at that backend CI configuration or rerun it when appropriate? Thank you!

[JacobSzwejbka]

Instead of adding a new spec can you just relax the invariant on viewspec

[JacobSzwejbka]

Reviewing this made me realize we dont compose well with the reinplace pass.py today #23279 I want to land that and have you integrate it.

The other issue is PyTorch for some reason allows out of bound slices

x[2:99] # tensor([2, 3])
x[99:100] # tensor([])
x[-99:2] # tensor([0, 1])

I think you need to add some clamping to the pass. Or reject these slices outright

[iRAFEEK]

@JacobSzwejbka Thanks for the review and for flagging both issues. I will wait for #23279 to land, then rebase/integrate the updated reinplace behavior and add safe handling for out-of-bounds slice bounds before requesting another review. Thank you!