Repository navigation
Add Method.set_output and write unplanned input mutations in place in pybindings - #23183
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23183
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit c38d1cf with merge base 0549e1e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Thank you for your pull request and welcome to our community. Action RequiredIn 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. ProcessIn 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 If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
This PR needs a
|
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Several moderate pybindings issues remain unresolved around dtype validation, input tracking, copying, allocation, and alias rebinding.
Review effort: Lite
Findings: 2
Open (2)
What changed in this PR
This PR adds Python support for in-place input mutation and output-buffer binding, plus MLX lowering for functional index_copy.
Changes:
- Tracks caller-owned inputs and adds
Method.set_output. - Avoids unnecessary copies for mutated or bound outputs.
- Restricts MLX in-place patterns and adds functional lowering.
- Adds runtime and MLX tests.
| File | Summary | Review notes |
|---|---|---|
runtime/test/test_runtime.py |
Adds mutation and output-binding tests. | — |
runtime/__init__.py |
Exposes Method.set_output. |
— |
extension/pybindings/pybindings.pyi |
Updates Python type declarations. | — |
extension/pybindings/pybindings.cpp |
Implements input tracking, aliasing, binding, and output handling. | Five moderate issues remain: validate bound-output dtype (4 votes); commit fed-input tracking only after successful set_inputs (1 vote); bypass cloning for caller-backed outputs (2 votes); skip allocation for pre-bound outputs (1 vote); rebind aliased outputs when inputs change (1 vote). |
backends/mlx/test/test_ops.py |
Tests user-input index_copy. |
— |
backends/mlx/patterns.py |
Restricts in-place matching to delegate-owned buffers. | — |
backends/mlx/ops.py |
Adds functional index_copy lowering. |
— |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| if (!tensor.is_contiguous()) { | ||
| throw std::runtime_error( | ||
| "output " + std::to_string(index) + " must be bound to a contiguous tensor"); | ||
| } | ||
| Error status = method_->set_output_data_ptr( | ||
| mutable_tensor_data_ptr_no_cow(tensor), tensor.nbytes(), index); |
|
|
||
| // Retrieve outputs | ||
| return get_outputs_as_py_list(result, clone_outputs); | ||
| py::list outputs = get_outputs_as_py_list(result, clone_outputs); |
|
Hi @IlyasMoutawwakil, 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. |
|
Can we separate the MLX parts and pybindings into separate PRs? I'll review this PR for pybinding material. |
| std::optional<at::Tensor> fed_input_aliased_by(size_t index) { | ||
| const bool known = index < aliases_input_.size() && aliases_input_[index]; | ||
| if (!known && | ||
| (index >= output_storages_.size() || output_storages_[index].empty())) { | ||
| return std::nullopt; // planned, or not a tensor | ||
| } | ||
| const auto& output = method_->get_output(index); | ||
| if (!output.isTensor()) { | ||
| return std::nullopt; | ||
| } | ||
| const void* data = output.toTensor().const_data_ptr(); | ||
| for (const auto& fed : fed_inputs_) { | ||
| if (data != nullptr && fed.const_data_ptr() == data) { | ||
| if (!known) { | ||
| aliases_input_.resize(output_storages_.size(), false); | ||
| aliases_input_[index] = true; | ||
| output_storages_[index] = std::vector<uint8_t>(); |
There was a problem hiding this comment.
Comparing data pointers can match an output that doesn't actually alias an input. If a caller gets outputs with clone_outputs=False and feeds one back in as an input on the next call (e.g., recurrent state), the output's current pointer is output_storages_[i]. That equals the fed input's pointer, so this marks the output as aliasing an input and frees output_storages_[i]. The fed input and the tensor the caller got from the previous call then both point at freed memory. aliases_input_[i] also stays true forever, so later calls with fresh inputs keep writing into that freed buffer.
An output that carries an input mutation is the same value as that input: the copy_ added by insert_write_back_for_buffers_pass writes into the input placeholder. Can we detect it structurally, e.g. &method_->get_output(i) == &method_->get_input(j), computed once in the constructor? That removes the pointer comparison, the aliases_input_ cache and the storage freeing here.
| // execute() on, instead of into storage this method allocates. The binding | ||
| // persists across calls and the tensor is kept alive until it is replaced. | ||
| // Binding an output to the input it mutates makes that write in place. | ||
| void set_output(at::Tensor tensor, size_t index) { |
There was a problem hiding this comment.
This checks contiguity and size, but not dtype or device. A float64 or int tensor bound to a float32 output gets the raw bits written into it and is returned as-is. A CUDA tensor hands CPU kernels a device pointer. set_inputs already rejects unsupported devices. Can we check both against method_meta().output_tensor_meta(index)?
| bound_outputs_[index] = tensor; | ||
| // Storage allocated for this output before it was bound is never used again. | ||
| if (index < output_storages_.size()) { | ||
| output_storages_[index] = std::vector<uint8_t>(); |
There was a problem hiding this comment.
Freeing output_storages_[index] here leaves any tensor already returned for this output with clone_outputs=False pointing at freed memory. Either say that set_output must be called before the first execute(), or keep the old storage alive.
| allocate_output_storages(); | ||
| std::vector<Span<uint8_t>> output_storage_spans(num_outputs); | ||
| caller_outputs_.assign(num_outputs, std::nullopt); | ||
| for (int i = 0; i < output_storages_.size(); ++i) { |
There was a problem hiding this comment.
nit: this loop compares a signed int with size_t. It was already here, but the new code is inside it, so worth switching to size_t.
| // An output bound with set_output() is pointed at the caller's tensor | ||
| // again on every call: one that is also an input was just repointed at | ||
| // the tensor passed in by set_inputs(). | ||
| if (is_bound_output(i)) { |
There was a problem hiding this comment.
If the bound output is also an input (a mutated input), this runs after set_inputs and re-points the input's value at the bound tensor. So if I bind output 0 to tensor A and then call with cache B, the method reads and writes A and silently ignores B. Since mutated unplanned inputs are now handled automatically (below), should set_output just reject outputs that alias an input? That's easy once aliasing is detected structurally.
| from types import ModuleType | ||
| from typing import Any, BinaryIO, Dict, List, Optional, Sequence, Set, Union | ||
|
|
||
| import torch |
There was a problem hiding this comment.
nit: torch is only used in an annotation here. Could this go under TYPE_CHECKING, or use from __future__ import annotations?
| The method then writes that output straight into ``tensor`` and returns | ||
| ``tensor`` itself, uncloned, instead of a copy of its own storage. Binding | ||
| the output that carries an input mutation to the input tensor it mutates | ||
| makes the mutation in place, with no copies. The output must not be |
There was a problem hiding this comment.
Binding a mutated input's output to that input isn't needed anymore, because it happens automatically. And binding it to a different tensor overrides the input the caller passes (see the comment in pybindings.cpp on is_bound_output). Can we drop this sentence? Please also spell out the dtype/device requirements, and the fact that the returned tensor is overwritten by later calls.
| from executorch.exir import ExecutorchBackendConfig, to_edge | ||
| from executorch.exir.passes import MemoryPlanningPass |
There was a problem hiding this comment.
These new imports need //executorch/exir:lib (and the passes target) added to test_runtime's deps in runtime/test/targets.bzl. Right now it only builds through make_test's transitive deps.
| return cache.sum(0, keepdim=True) | ||
|
|
||
|
|
||
| def _load_unplanned_io_method(module, inputs): |
There was a problem hiding this comment.
nit: several lines in this helper and the tests are over 88 columns, so ufmt/lintrunner will flag them.
| self.assertIs(outputs[1], summed) | ||
| self.assertTrue(torch.equal(summed, cache.sum(0, keepdim=True))) | ||
|
|
||
| def test_set_output_rejects_bad_bindings(self): |
There was a problem hiding this comment.
Could you add tests for:
- feeding a
clone_outputs=Falseoutput back in as an input (would catch the aliasing bug); - binding a tensor of the wrong dtype;
set_outputon an output that aliases an input;- a dynamic-shape output smaller than its bound tensor (the
narrowbranch).
|
@IlyasMoutawwakil do rebase with main and ping us for review after addressing the above comments. |
… pybindings ExecuTorchMethod.set_output(tensor, index) binds an unplanned output to caller storage, checked for dtype, device, contiguity and capacity. The write-back of a mutated input exported with alloc_graph_input=False is detected structurally and returned as the caller's tensor, so the mutation lands in place without copies. Only PyMethod changes; PyModule::run_method keeps its current behavior. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
set_output rejects an output returned at more than one index, a binding made before the first call to an output that turns out to be an input is removed with an error instead of breaking every later call, and inputs are only kept alive once set_inputs succeeds. Docs now cover any unplanned input returned as an output, and set_output is listed in the runtime API reference. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The MLX changes moved to their own PR; this branch keeps only the pybindings and runtime changes.
|
@IlyasMoutawwakil just in case, could you rebase with main and ping for review? There were some fixes put up which cleans up unrelated CI fixtures and I'd like to avoid the noise. |
|
@nil-is-all done |
|
@IlyasMoutawwakil could you check the CI failures? |
Since in-place results are planned onto the spec of the input they update (pytorch#23563), a mutated input shares its spec with the output that carries the mutation, so with outputs left unplanned the input is unplanned as well, and the method then rightly returns the caller's tensor. Plan both, and check the input really is planned. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@nil-is-all done |
|
LGTM, running CI now One question on BC change here regarding output cloning (cc @JacobSzwejbka for his thoughts): Consider: After lowering with MemoryPlanningPass(alloc_graph_input=False), and running: Before this PR, default output cloning makes y an independent copy of x, whereas before it was not. This happens without calling the new set_output API. Perhaps we should not change the behavior here. If you really want no copy, you need to specify clone_outputs=False explicitly when using set_output. |
|
@IlyasMoutawwakil could you check the failing CI tests and sign CLA? |
|
@nil-is-all don't look related to me, correct me if i'm wrong |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
An unplanned input returned as an output, e.g. `return x, x + 1` exported with alloc_graph_input=False, was a copy with the default clone_outputs=True before this PR, and became the caller's own tensor with it. Keep the old default: such an output, including the write-back of a mutated input, is cloned like any other, and is the tensor passed in only with clone_outputs=False. The mutation is still written into the caller's tensor, and an output bound with set_output is still returned as the bound tensor, since binding it is the opt-in. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
@nil-is-all those 4 failures weren't from this PR: qnn-optimum (dit), samsung-models and arduino-library all died with "Executing the custom container implementation failed" on the runners, and the bert coreml macos log had already expired. they pass on main. I merged main to rerun everything. there's a new qwen3 coreml macos failure on this run but it's the same kind of thing, the runner couldn't reach github.com while cloning the sljit submodule, so it just needs a rerun. CLA is signed now. @metascroy agreed, I'd rather not change the default either. fixed in b28bb4b: with clone_outputs=True an output that's also an input (your |
|
Looks great! Can you fix the merge conflicts @IlyasMoutawwakil? |
|
@claude /pr-review |
|
Claude finished @shoumikhin's task in 33m 47s —— View job PR Review: #23183
Request Changes: after this PR was approved, the merge of Detailed reviewSummaryThis PR adds
I read the code against Backward Compatibility
API Design
Testing
Performance
Code Quality
RecommendationRequest Changes
The other items are minor. |
Resolve the conflicts with the torch-free pybindings changes on main: - set_inputs: read inputs through main's read_input/scalar_input; record the tensor the method writes (the layout-converted copy, if any) for outputs that alias an input. - A failed set_inputs follows main: the method refuses to run until a full set is installed, so the earlier generation is no longer kept alive. - Output setup keeps bound and input-aliased outputs, on main's device buffer accessor; TorchTensorView now takes the output index. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With clone_outputs=False, an output that is an unplanned input now always comes back as main's detached alias of the caller's storage (pytorch#23089), instead of the caller's own object only when its strides happened to match. Only outputs bound with set_output are returned as the caller's tensor. This restores main's retention test and drops the per-input bookkeeping the identity path needed. Also drop tests that could not fail: the profiler clone check, the "wrong device" binding (rejected before the device check), and a final assertion on a row the first call had already written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Thanks for the review. Addressed in 7b552a3 and b0b8f9c:
Left for a follow-up: unbinding ( @metascroy the main merge changed how a mutated unplanned input is returned with
|
- The ATen bindings swap the data pointer of the output's storage in place, so an output returned earlier would read the bound tensor, or freed memory after a rebind. set_output now refuses ATen builds. - A binding whose memory overlaps another bound output, or an input, is rejected: two outputs written into one buffer, or an output written over an input a kernel is still reading, would corrupt each other silently. - Rejected bindings raise ValueError, like the input checks, and a dtype/device mismatch names both sides. - Tests for shared memory, and for execute() without inputs after set_output (alias detection is deferred to the first set_inputs()). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
CI on b0b8f9c, for the record:
I merged main again, which brings no changes to |
… pybindings (pytorch#23183) ## Summary Adds `ExecuTorchMethod.set_output(tensor, index)`, exposed as `executorch.runtime.Method.set_output`. It binds an output that is not memory planned to caller storage: the method writes that output straight into `tensor` on every later call, with no copy, and returns it as that tensor (viewed at the written shape when a dynamic output is smaller), uncloned, since binding is the opt-in. Later calls overwrite it. `set_output` rejects, with `ValueError`, a tensor that doesn't match the output's dtype and device, isn't contiguous or is too small for the output's largest shape, or shares memory with another bound output or an input; and an output that is memory planned, returned at more than one index, not contiguous, or also an input. Storage the output used before binding is kept alive, so tensors already returned with `clone_outputs=False` stay valid. The ATen bindings don't support it: they swap the data pointer of the output's storage in place, so earlier outputs would read the bound tensor, or freed memory after a rebind. An output that is also an input is found by checking that both are the same value in the method (`&get_output(i) == &get_input(j)`), not by comparing data pointers, which could match the wrong tensor. The check runs once, after the first successful `set_inputs()`, not in the constructor: `get_input()` marks an input as set, which would disable `execute()`'s "Input N has not been set" check. So binding such an output before the first call can't be rejected up front: the first call removes that binding and raises, and later calls work normally. Unplanned inputs are handled as on main: a mutation is written into the caller's tensor, and an output that is that input is a clone by default and main's detached alias of the caller's storage with `clone_outputs=False` (pytorch#23089). Only `PyMethod` changes. The MLX `index_copy` changes from the first version of this PR are in pytorch#23508. ## Test plan New tests in `runtime/test/test_runtime.py` cover binding, rebinding, and dynamic outputs smaller than the bound tensor; rejection of bad dtype, size, layout, planned, duplicated, input-aliasing and memory-sharing bindings; recovery from a binding made before the first call; `execute()` without inputs after a binding; mutated unplanned inputs (updated in place, returned as a copy by default and sharing the caller's storage with `clone_outputs=False`) and planned ones (not updated); an uncloned output fed back in as the next input; and a failed `set_inputs()`. pytest runtime/test/test_runtime.py extension/pybindings/test This PR was authored with Claude Code. 🤖 Generated with [Claude Code](https://claude.com/claude-code) cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all @metascroy --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Summary
Adds
ExecuTorchMethod.set_output(tensor, index), exposed asexecutorch.runtime.Method.set_output. It binds an output that is not memory planned to caller storage: the method writes that output straight intotensoron every later call, with no copy, and returns it as that tensor (viewed at the written shape when a dynamic output is smaller), uncloned, since binding is the opt-in. Later calls overwrite it.set_outputrejects, withValueError, a tensor that doesn't match the output's dtype and device, isn't contiguous or is too small for the output's largest shape, or shares memory with another bound output or an input; and an output that is memory planned, returned at more than one index, not contiguous, or also an input. Storage the output used before binding is kept alive, so tensors already returned withclone_outputs=Falsestay valid. The ATen bindings don't support it: they swap the data pointer of the output's storage in place, so earlier outputs would read the bound tensor, or freed memory after a rebind.An output that is also an input is found by checking that both are the same value in the method (
&get_output(i) == &get_input(j)), not by comparing data pointers, which could match the wrong tensor. The check runs once, after the first successfulset_inputs(), not in the constructor:get_input()marks an input as set, which would disableexecute()'s "Input N has not been set" check. So binding such an output before the first call can't be rejected up front: the first call removes that binding and raises, and later calls work normally.Unplanned inputs are handled as on main: a mutation is written into the caller's tensor, and an output that is that input is a clone by default and main's detached alias of the caller's storage with
clone_outputs=False(#23089). OnlyPyMethodchanges.The MLX
index_copychanges from the first version of this PR are in #23508.Test plan
New tests in
runtime/test/test_runtime.pycover binding, rebinding, and dynamic outputs smaller than the bound tensor; rejection of bad dtype, size, layout, planned, duplicated, input-aliasing and memory-sharing bindings; recovery from a binding made before the first call;execute()without inputs after a binding; mutated unplanned inputs (updated in place, returned as a copy by default and sharing the caller's storage withclone_outputs=False) and planned ones (not updated); an uncloned output fed back in as the next input; and a failedset_inputs().This PR was authored with Claude Code.
🤖 Generated with Claude Code
cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all @metascroy