Repository navigation
Metal backend: keep non-packed views in their parent's buffer - #22984
Open
abdelaziz-mahdy wants to merge 45 commits into
Open
abdelaziz-mahdy wants to merge 45 commits into
abdelaziz-mahdy wants to merge 45 commits into
Conversation
…set view aoti_torch__reinterpret_tensor read views that start partway into a buffer before the GPU had written them. A non-packed view is copied on the CPU, and the wait for pending GPU work was guarded by metal_is_device_pointer(src). src is the offset pointer, and only a buffer's base address is registered, so the wait was skipped whenever the storage offset was non-zero. A packed view gets its own no-copy MTLBuffer over the parent's memory; Metal tracks hazards per buffer object, so nothing ordered work reading it after the pending work that writes the parent. Decide on the wait from the base pointer, and wait before aliasing a packed view. Adds two linear-then-chunk modules that fed the second chunk to another linear and came back as zeros.
…aliasing it aoti_torch__reinterpret_tensor gave every packed view with a storage offset its own no-copy MTLBuffer over the parent's memory. Metal treats the two buffers as unrelated: a write through one followed by a read of the other in the same serial compute encoder sees stale data every time. Inductor does exactly that when it fills the result of a cat by writing through views of it, and the opposite when an op reads a chunk of a buffer another op just wrote, so models built on split/cat (YOLO's C2f blocks) ran but were wrong. Register such a view against the buffer it lives in and bind that buffer at the view's offset, in compute encoders and in blits. This replaces the wait-before-aliasing from the previous commit, which only covered reads. MPSGraphTensorData cannot address into a buffer, so MPSGraph ops still get an alias for an offset view, now with the stream synchronized on both sides of the graph. Adds cat variants of the chunk tests. The last-dim one is skipped: its slices are non-packed views, which reinterpret_tensor materializes into a copy, so writes through them are still lost.
Inductor flattens nested views into a single reinterpret of the base buffer with the combined offset, so this exercises a larger packed offset rather than a view of a view, but it pins that behavior down.
A view registration is counted, but only aoti_torch__reinterpret_tensor with a non-zero offset took a count. aoti_torch_new_tensor_handle on a view, and a reinterpret of a view at offset 0, both make another handle at the same address without one, while aoti_torch_delete_tensor_object gives one back for every handle. Deleting either handle therefore unregistered the view for the other: a kernel writing through the surviving handle fell back to a temporary and the write never reached the buffer, and MPSGraph ops failed to find the tensor's Metal buffer. metal_retain_view takes a count for an address that is a registered view and does nothing otherwise; both paths call it. Tests cover both ways of making the second handle, in both deletion orders. They need the test_metal_memory target, hence the merge of main.
get_mtl_buffer asked the stream to wait after the next graph before it knew the alias MTLBuffer could be created. On failure it throws, and the flag would have made the next, unrelated graph wait. Also reword a comment in materialize_packed that claimed only base addresses are known as device pointers; constant sub-buffers are mapped at interior addresses too.
aoti_torch__reinterpret_tensor replaced a view that is not densely packed with a packed copy. Kernels generated by inductor do not look at the tensor they are handed: they index with the strides they were compiled with, and they write through such views as well as read them. Inductor's default channels-last layout makes this the common case, since a chunk along C is non-packed there and a cat is filled through four views of it. Against a packed copy those kernels lose their writes, so the cat buffer stays empty, and index up to `parent numel` into a buffer of `view numel`, which is an out-of-bounds GPU access: a C2f block returns wrong values, and yolov8n crashes inside MPSGraph once the heap has been written over. A non-packed view of a Metal buffer now stays where it is and is registered as a view of its parent, like a packed one. The tensor cannot describe its strides, so it carries packed ones and the real ones are recorded (metal_record_strided_view). aoti_torch_mps_set_arg_tensor binds the view in place. Hand-written ops need dense input, so get_mtl_buffer and ETMetalKernelFunction::setArg hand them a packed copy made when the op runs; an MPSGraph op asked to write its result through such a view, and aoti_torch_copy_ on one, fail instead of writing to the wrong place. Non-packed views of CPU memory are materialized as before. Adds pointwise_c2f, the YOLO C2f block reduced to 1x1 convs, and enables linear_chunk_cat_last_dim, which was skipped for this reason.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22984
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New FailuresAs of commit cb8b453 with merge base 5c788db ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
A critical encoder-lifecycle issue and a moderate view-reference lifetime issue remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (2)
What changed in this PR
Updates the Metal backend to preserve non-packed views in parent buffers while supplying dense copies to operations that require them.
Changes:
- Adds Metal view registration, binding, and lifetime tracking.
- Adds packed-copy handling for dense-only operations.
- Adds regression and view-lifetime tests.
| File | Reviewed change |
|---|---|
backends/apple/metal/tests/test_modules.py |
Adds view-layout regressions. Nit (1 vote): update the misleading docstring about materialization. |
backends/apple/metal/runtime/test/test_memory.cpp |
Tests view-handle lifetime tracking. |
backends/apple/metal/runtime/shims/shim_mps.mm |
Binds generated kernels directly to view storage. |
backends/apple/metal/runtime/shims/memory.cpp |
Preserves views and validates copies. Moderate (3 votes): owning-buffer references for offset views are not released. |
backends/apple/metal/runtime/shims/et_metal.mm |
Implements view resolution and packed copies. Critical (3 votes): synchronization can end the active encoder before packed-buffer binding and dispatch. |
backends/apple/metal/runtime/shims/et_metal.h |
Exposes view and synchronization APIs. |
backends/apple/metal/runtime/ops/common.mm |
Provides dense buffers for MPSGraph operations. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Two problems in the previous commit, both found in review. The packed copy that hand-written ops get of a strided view was made on the CPU after a COMMIT_AND_WAIT. ETMetalKernelFunction::setArg runs after the op has started encoding, and synchronize() ends the stream's encoder, so the op went on binding and dispatching through an encoder that had been ended. The copy is now a gather kernel encoded on the stream's current encoder: nothing is committed or waited for, the copy is ordered after the work that fills the view, and the op's own pipeline state is set again afterwards. It binds at the top of the argument table so the arguments the op has already set stay put. A view at another address than its parent takes a count on the parent's memory, but deleting the view only unregistered it, so the parent's count never came back down and its buffer was not freed until cleanup. The shims now remember which parent each such view counts towards and give the count back when the view is deleted. This predates the previous commit; that commit makes more views take this path. Adds test_metal_strided_view (the view stays in its parent, the GPU gather returns its elements, a copied handle is a strided view too) and two tests to test_metal_memory that fail without the refcount change.
abdelaziz-mahdy
requested review from
kirklandsign and
larryliu0820
as code owners
September 22, 2026 00:28
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical issues remain in packed-view fallback/output handling and nested or copied-view lifetime ownership.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 4
Open (5)
Resolved since last review (2)
…memory alive Review of the previous commit found four problems, and a fifth came up checking the rest of the change. - setArg packed every strided argument, including ones the kernel writes, and when packing failed it bound the view with the packed strides it carries. setArg now takes how the kernel uses the argument: a hand-written kernel's input gets a packed copy, an argument it writes is refused if it is a strided view (writes to a copy would be lost), and a kernel generated by inductor gets the view in place. A view that cannot be packed is an error. The outputs of the hand-written kernels are allocated by the ops and are never views; gated_delta_rule's state, updated in place, is the one caller tensor they write. - A view of a view took no count on the allocation it lives in, and a handle copied from a view took none either, so deleting the base and the first handle could free memory a remaining handle still pointed into. Every handle into memory the runtime owns now holds a count on that allocation, found through the handle it was made from. - aoti_torch_copy_ refused strided views, which broke models whose output is one: the backend copies outputs out with it. A strided source is now copied through the same GPU gather. Adds linear_chunk_last_dim_output, tests for nested views and copied view handles in test_metal_memory, and tests for the argument rules in test_metal_strided_view. Test docstrings no longer describe views as materialized.
Contributor
Author
|
Pushed b48b581 for the five review comments above. Going over the rest of the change I also found a regression the reviews did not cover: 4648c9a made |
executorch-triage
Bot
added
the
community: contribution
…iews working Review of 89dbe34 found two problems, and a third came up checking the memory accounting around them. - The wait that an aliasing buffer needs was a flag on the stream, armed by get_mtl_buffer and consumed by whichever graph ran next. Another thread's graph could consume it, and an op failing between get_mtl_buffer and its graph left it armed for an unrelated one. get_mtl_buffer now only reports that it made an alias, the op passes that to its own executeMPSGraph call, and the stream waits before and after encoding that graph inside the same block on its serial queue, so nothing else can come in between. The stream no longer holds any such state. - A packed offset view of CPU memory lost the no-copy buffer it used to get, so MPSGraph ops could not find it and kernels wrote to a temporary. It gets one again. Unlike before, the buffer is counted with the view's handles and released with the last of them, so that it does not outlive the CPU memory under it; a longer view at the same address gets a buffer that covers it. - A view at an offset took a count on its parent allocation that deleting it never gave back, and a view of a view, or a handle copied from a view, took none. The parent was then either never freed, or freed while such a handle still pointed into it. Every handle into memory the runtime owns now holds a count on the allocation it lives in, found through the handle it was made from, and gives it back when deleted. test_metal_memory: 9 new tests. All but AliasedGraphSettlesItsOwnWork fail on 89dbe34; that one covers what the change must keep working.
…-views Brings in the per-graph alias wait, the counted buffers of CPU-backed views, and the ownership accounting this branch already had. Two things needed changing for them to work together: - A view of CPU memory that has a Metal buffer of its own counts as a device pointer, but it cannot hold a non-packed view in place: its buffer covers only itself. Such a view is materialized, like any view of CPU memory. - materialize_packed waits for the GPU again before reading. Kernels can now write CPU memory through the buffers of views of it, and a copy made afterwards has to see those writes. test_metal_strided_view gains a test for each; both fail without the change.
The strides a reinterpreted tensor carries are packed ones whether the view stays in its parent as a strided view or is materialized, so the check that they fit a tensor now comes before either.
An empty view divides by zero in the gather, and negative strides wrap in its unsigned index math. metal_record_strided_view now refuses both, and size/stride lists of different lengths; reinterpret refuses a non-packed view with negative strides before registering anything.
…ed memory aoti_torch_mps_free now forgets the views in any buffer it frees, not only in constants buffers with nested ones, and retires the handles still in it: deleting one later only drops the handle, so it cannot unregister or uncount a later allocation at its address, and shims refuse it. It refuses memory a tensor owns and a constant's own buffer, which go with the tensor and with the buffer of all constants. The CPU region stamp no longer records a stream, which the one-stream model does not have. topk freed the memory of its values handle on later errors, and SDPA leaked its output handle and buffers; both now delete the handle once it owns the memory. topk refuses a negative k and returns empty tensors for an empty result, which copy_ now copies out whatever strides either side has. topk returns its scratch to the pool instead of going through mps_free.
Retiring the handles in freed memory also forgets their strided-view records (FreeingABufferForgetsTheStridedViewsInIt), and copy_ checks for retired handles before refusing strided destinations. Strided views are now checked when they are made against everything the gather needs (metal_can_pack_strided_view: at most 16 dims, no empty dim, no negative stride, under 2^32 elements in the view and in the memory it spans), so recording one cannot fail afterwards. FailingToPackIsAnError recorded a view the gather could not index; recording now refuses it, and the case moved to ViewsTheGatherCannotPackAreNotRecorded.
Views reaching past 2^32 elements into their buffer (chunks of large logits) were refused when made; the gather now steps through the memory a view spans in 64 bits, and only packing more than 2^32 elements at once fails, when an op needs it. Views it cannot pack at all (negative strides, more than 16 dims) are still refused when made. setArg recorded every argument of every kernel, with a retain or a heap copy each, although only a hand-written kernel reading a strided view runs the gather, and the gather only overwrites slots 28-30. Only those are recorded now. The overlap check of a strided copy is its own function, so a test can check it without depending on the order the GPU runs threads in.
clearBindings keeps the records it resets, so binding a kernel again does not allocate them. The temporary buffer for a large CPU argument is released right after it is bound: the stream's command buffers retain what their encoders bind. The header comments now say when the packed copy fails, and sit on the functions they describe.
Contributor
Author
|
Merged main to pick up the refreshed torchao pin (#23184); the Sep 28 runs failed at dependency install on the pruned |
Contributor
|
Similarly, @abdelaziz-mahdy please update branch and ping again, will trigger CI and review sooner |
Contributor
Author
|
@nil-is-all updated as well: main and the latest #22957 are merged in, and the gtests and Metal module suite pass. Ready for review. It's stacked on #22957, so its diff shrinks to its own commits once that lands. |
An empty view is bound at its parent's start, so that one at the end of its buffer is not taken for a view of the next allocation. That also dropped its storage offset: reinterpreting an empty view at offset 8 of an 8-element allocation as 4 elements at offset 0 addressed elements 0-3 instead of being refused past the end. The address an empty view stands for is now recorded (empty_view_address) and views of it, and of its copies, are offset from there and bounds-checked against it.
# Conflicts: # backends/apple/metal/runtime/shims/memory.cpp # backends/apple/metal/runtime/test/test_memory.cpp
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
A critical pointer-classification issue can release a registered CPU region while an active handle still depends on it, requiring human assessment of the memory-lifetime behavior.
No open issues 🎉
1 resolved since last review
Review effort: Lite
Contributor
Author
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #22983
Stacked on #22957, which adds the view registration this builds on, and merges it as it changes. Once #22957 lands, the diff reduces to this PR's own commits.
Stack
destroyfrees only its own delegate's memoryaoti_torch__reinterpret_tensorreplaced a view that is not densely packed with a packed copy. Kernels generated by inductor index with the strides they were compiled with and write through such views as well as read them, so against a packed copy they lose their writes and index up toparent numelinto a buffer ofview numel. Inductor's default channels-last layout makes this the common case: a chunk along C is non-packed there, and a cat is filled through views of it.make_tensor_ptrrejects them), so it carries packed ones and the real ones are recorded withmetal_record_strided_view.ETMetalKernelFunction::setArgtakes how the kernel uses the argument (ArgAccess).aoti_torch_mps_set_arg_tensorpasseskStridedInPlace: a generated kernel gets the view where it is, which is all it needs.get_mtl_buffer(mm, addmm, bmm, convolution, topk) andsetArgwithkReadhand them a packed copy, made by a small gather kernel encoded on the stream when the op runs. Nothing is committed or waited for, and the op's own encoder stays open. An argument a hand-written kernel writes (kWrite) is refused if it is a strided view.[1, T, V]logits past 2^32 elements) are packed too. A view it cannot pack at all (negative strides, more than 16 dims) is refused when it is made. An empty view is packed whatever its strides, so it never needs the gather.aoti_torch_copy_into one, now fail instead of writing to the wrong place.aoti_torch_copy_from one copies through the same gather; the backend uses it to copy a model's outputs out, so a model can return such a view.materialize_packedloses its device path, which nothing reaches any more.This resolves the "Not fixed here" item of #22957.
Test plan
pointwise_c2fadded (the YOLO C2f block reduced to 1x1 convs),linear_chunk_last_dim_outputadded (the model's output is such a view), andlinear_chunk_cat_last_dimenabled; it was skipped for this reason.pointwise_c2ffails withmax_atol=0.586(float32) and0.755(bfloat16),linear_chunk_cat_last_dimwithmax_atol=0.518.Ran 174 tests ... OK.test_metal_memory94 of 94 and newtest_metal_strided_view14 of 14; each test added for a review finding fails without its change.DISABLED_PackingAViewSpanningMoreThan32Bitsneeds a 6.4 GB buffer, so it runs only on request (--gtest_also_run_disabled_tests); it passes, and fails with 32-bit strides.Ran 166 tests ... OK), all-ones input against eager throughexecutor_runner:yolov8nexported with the default layout went from crashing 10 runs out of 10 tomax abs diff 0.0018over 5 runs; withlayout_optimization=Falseit stays at 0.0019.yolo11nwith the default layout still stops inaoti_torch_mps_bmm_out, which rejects a permuted input ([2,64,400], strides[1,800,2]); that is a separate limitation.lintrunnerclean.