Skip to content

Add Method.set_output and write unplanned input mutations in place in pybindings - #23183

Merged
nil-is-all merged 18 commits into
pytorch:mainfrom
IlyasMoutawwakil:pybindings-method-set-output
Oct 28, 2026
Merged

nil-is-all merged 18 commits into
pytorch:mainfrom
IlyasMoutawwakil:pybindings-method-set-output

Conversation

@IlyasMoutawwakil

@IlyasMoutawwakil IlyasMoutawwakil commented Sep 28, 2026 •

Copy link
Copy Markdown
Contributor

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 (#23089). Only PyMethod changes.

The MLX index_copy changes from the first version of this PR are in #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

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani @nil-is-all @metascroy

Copilot AI lite review requested due to automatic review settings September 28, 2026 14:29
@pytorch-bot

pytorch-bot Bot commented Sep 28, 2026 •

Copy link
Copy Markdown

🔗 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 Failures

As of commit c38d1cf with merge base 0549e1e (image):
💚 Looks good so far! There are no failures yet. 💚

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

@meta-cla

meta-cla Bot commented Sep 28, 2026

Copy link
Copy Markdown

Hi @IlyasMoutawwakil!

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!

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 Medium severity

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.

Comment thread extension/pybindings/pybindings.cpp Outdated
Comment on lines +1393 to +1398
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);
Comment thread extension/pybindings/pybindings.cpp Outdated

// Retrieve outputs
return get_outputs_as_py_list(result, clone_outputs);
py::list outputs = get_outputs_as_py_list(result, clone_outputs);
@executorch-triage executorch-triage Bot added the community: contribution PRs coming from community (excluding hardware partners) label Sep 28, 2026
@nil-is-all nil-is-all added the module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon label Sep 28, 2026
@nil-is-all

Copy link
Copy Markdown
Contributor

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.

@metascroy

Copy link
Copy Markdown
Contributor

Can we separate the MLX parts and pybindings into separate PRs?

I'll review this PR for pybinding material.

Comment thread extension/pybindings/pybindings.cpp Outdated
Comment on lines +1553 to +1569
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>();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread extension/pybindings/pybindings.cpp Outdated
// 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) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)?

Comment thread extension/pybindings/pybindings.cpp Outdated
Comment on lines +1405 to +1408
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>();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread extension/pybindings/pybindings.cpp Outdated
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) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread runtime/__init__.py Outdated
from types import ModuleType
from typing import Any, BinaryIO, Dict, List, Optional, Sequence, Set, Union

import torch

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: torch is only used in an annotation here. Could this go under TYPE_CHECKING, or use from __future__ import annotations?

Comment thread runtime/__init__.py Outdated
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread runtime/test/test_runtime.py Outdated
Comment on lines +19 to +20
from executorch.exir import ExecutorchBackendConfig, to_edge
from executorch.exir.passes import MemoryPlanningPass

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread runtime/test/test_runtime.py Outdated
return cache.sum(0, keepdim=True)


def _load_unplanned_io_method(module, inputs):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add tests for:

  • feeding a clone_outputs=False output back in as an input (would catch the aliasing bug);
  • binding a tensor of the wrong dtype;
  • set_output on an output that aliases an input;
  • a dynamic-shape output smaller than its bound tensor (the narrow branch).

@nil-is-all

Copy link
Copy Markdown
Contributor

@IlyasMoutawwakil do rebase with main and ping us for review after addressing the above comments.

@nil-is-all nil-is-all added the need-user-input The issue needs more information from the reporter before moving forward label Oct 5, 2026
IlyasMoutawwakil and others added 3 commits October 6, 2026 12:26
… 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.
@nil-is-all

Copy link
Copy Markdown
Contributor

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

@IlyasMoutawwakil

Copy link
Copy Markdown
Contributor Author

@nil-is-all done

@nil-is-all

Copy link
Copy Markdown
Contributor

@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>
@IlyasMoutawwakil

Copy link
Copy Markdown
Contributor Author

@nil-is-all done

@metascroy

metascroy commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

LGTM, running CI now

One question on BC change here regarding output cloning (cc @JacobSzwejbka for his thoughts):

Consider:

def forward(self, x):
    return x, x + 1

After lowering with MemoryPlanningPass(alloc_graph_input=False), and running:

x = torch.tensor([1.0])
y, z = method.execute((x,))

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.

@nil-is-all

Copy link
Copy Markdown
Contributor

@IlyasMoutawwakil could you check the failing CI tests and sign CLA?

@IlyasMoutawwakil

Copy link
Copy Markdown
Contributor Author

@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>
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 14, 2026
@meta-cla

meta-cla Bot commented Oct 14, 2026

Copy link
Copy Markdown

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>
@IlyasMoutawwakil

Copy link
Copy Markdown
Contributor Author

@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 return x, x + 1 example, or the write back of a mutated input) is cloned like before, and you only get the caller's tensor back with clone_outputs=False. the mutation itself is still written into the caller's unplanned input. I left outputs bound with set_output uncloned since binding is already the opt-in (and Method.execute doesn't take clone_outputs), but happy to change that if you'd prefer. added test_returned_input_is_cloned_by_default with your example and updated the pyi, docstrings and description.

@nil-is-all nil-is-all removed the need-user-input The issue needs more information from the reporter before moving forward label Oct 16, 2026
@metascroy

Copy link
Copy Markdown
Contributor

Looks great! Can you fix the merge conflicts @IlyasMoutawwakil?

@shoumikhin

Copy link
Copy Markdown
Contributor

@claude /pr-review

@claude

claude Bot commented Oct 27, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 33m 47s —— View job


PR Review: #23183

  • Gather PR context and diff
  • Run three independent reviewers
  • Consolidate findings
  • Fact-check findings
  • Post review

Request Changes: after this PR was approved, the merge of main changed the clone_outputs=False behavior that #23089 tested, and the untested USE_ATEN_LIB path makes outputs returned before a set_output read the bound tensor, or freed memory once the output is rebound.

Detailed review

Summary

This PR adds Method.set_output(tensor, index) so the caller can supply the storage for an unplanned output, and under clone_outputs=False it returns an unplanned input that is also an output as the caller's own tensor.

  • The merge of main made after approval changes that clone_outputs=False behavior for existing callers.
  • The ATen half is untested and behaves differently from portable.
  • Several of the new tests cannot fail.

I read the code against c45d76f9, the main parent of merge 0ad204b9, because the shallow checkout has no merge base with origin/main. No tests were run.

Backward Compatibility

  • clone_outputs=False no longer returns [PyBindings] Remove ATen dependency from portable bindings #23089's detached alias for an output that is an unplanned input (extension/pybindings/pybindings.cpp:2377-2419, extension/pybindings/test/test_pybindings.py:935-973).

    On main (portable, Torch output mode) this output comes from aliased_torch_output, a new detach().as_strided(...) tensor. test_borrowed_external_output_does_not_retain_previous_outputs asserted that it has requires_grad=False and that the previous output is freed. Merge 0ad204b9, made after b28bb4bb was approved, returns the caller's own object when its shape/stride() equal the written ones, and falls back to main's behavior otherwise. Consequences:

    • Grad tensors and Parameters come back as themselves. A requires_grad=True leaf or a Parameter that is fed in is now returned unchanged. out.add_(1) then raises "a leaf Variable that requires grad is being used in an in-place operation", and autograd tracks the output. main returned a detached tensor.
    • Which object comes back depends on stride details:
      • use_layout (599-606) canonicalizes the strides of size-1 dims. So a contiguous torch.ones(3, 1).t() (stride (1, 1)) gets the detached alias, while torch.ones(1, 3) gets itself.
      • Changing the input's metadata after set_inputs also gets the alias.
      • A Parameter or subclass whose strides differ gets a frombuffer result instead, because aliased_torch_output requires an exact torch.Tensor (2579).
      • Non-Torch output modes keep main's result (2384-2387).
    • The retention check is gone. The test was renamed and its assertions flipped (assertIs(output, previous), assertTrue(output.requires_grad)). Both subtests now take the matching-stride path, so nothing checks that the fallback paths don't keep previous outputs alive.
    • The stub docstring is wrong for these cases. The ExecuTorchMethod docstring in pybindings.pyi says such an output "is the tensor passed in only with clone_outputs=False". That is false for the fallback cases above.
    • The PR description's rationale is stale. The "Previously … read zeros" paragraph no longer matches main. main already gives these outputs no storage (the is_input loop in allocate_output_storages) and does not repoint them (the input_data_ptrs check in setup_output_storage), so the mutation already lands in the caller's tensor.

    Pick one rule (always the caller's object, or always main's detached alias) and document it. Keep a retention test for every path that still builds an owner. This changed after approval, so ask the approver and the [PyBindings] Remove ATen dependency from portable bindings #23089 author to re-review.

API Design

  • Whether set_output succeeds depends on call order (pybindings.cpp:2510-2549).
    • Aliases are only computed after the first successful set_inputs, because the deprecated Method::get_input marks the input as set (runtime/executor/method.cpp:1896-1899).
    • So binding an output that is an input succeeds before the first call. That first call then installs the inputs, removes the binding and raises, and the Method.set_output docstring has to explain all of this.
    • The value indices that answer "is output i input j" without side effects already exist, but they are private (Method::get_input_index/get_output_index, runtime/executor/method.h:384-385).
    • A public, side-effect-free query on Method or MethodMeta would let set_output reject these bindings immediately and drop the remove-and-raise path.
  • Errors don't match the rest of the file (pybindings.cpp:2236-2291).
    • Every set_output rejection is a RuntimeError, while the equivalent input checks (validate_tensor_input, TorchTensorView, validate_device) raise ValueError. test_set_output_rejects_bad_bindings now locks in RuntimeError for a new public API.
    • "dtype and device must match the bound tensor" (2276) names neither the expected nor the actual dtype/device. The input checks print both.
    • aliased_input_message (2555) says the output "is returned as the tensor passed in for it". Under the default clone_outputs=True it is a copy.
  • Overlapping bindings are accepted (pybindings.cpp:2243-2293).
    • The duplicate check compares output EValues, not the bound storage. So set_output(t, 0); set_output(t, 1), or two views of one buffer, makes two outputs write the same memory.
    • Binding a tensor that is also fed as an input makes an out-variant kernel write into memory it is still reading.
    • Neither case is rejected or documented.
  • A binding cannot be removed (pybindings.cpp:2224, runtime/__init__.py:171). set_output only replaces a binding: it doesn't accept None, and there is no clear_output. Program.load_method caches the Method (runtime/__init__.py:259-263), so every later user of that method writes into the bound tensor, and the tensor stays alive as long as the program does.

Testing

  • test_bound_output_is_not_cloned cannot fail (runtime/test/test_runtime.py:246-253). Portable clones a CPU output with a memcpy in PyExecuTorchResult (pybindings_executorch_result.h:108-114), not with aten::clone, and executorch.runtime only loads portable. The identity check in test_set_output_binds_caller_storage already covers this.
  • The "wrong device" case never reaches the device comparison (test_runtime.py:284). It uses device="meta", which the TorchTensorView constructor rejects (pybindings.cpp:538-541) before the comparison at 2270-2271 runs.
  • The last assertion of test_set_output_rejects_mutated_input is vacuous (test_runtime.py:321). It checks cache[1], which the first execute at line 310 already set to ones, so it passes even if the final call writes nothing.
  • test_failed_set_inputs_returns_installed_input doesn't exercise the reset it targets (test_runtime.py:195-207).
    • It uses the default get_outputs(), which clones and skips the caller_outputs_ path (2379-2381). So it passes with or without the reset at pybindings.cpp:2195.
    • assertIsNot(outputs[0], cache) holds for any clone.
    • Use clone_outputs=False and assert that the output shares new_cache's storage.
  • Nothing guards why alias detection is deferred (comment at pybindings.cpp:2510-2512). Add a test that execute() without set_inputs still raises, both with and without an earlier set_output. That would catch a refactor that moves find_inputs_aliased_by_outputs() into set_output.

Performance

  • Every execute() re-runs all of set_output for each bound output (pybindings.cpp:2305-2320).

    • That repeats checks that cannot change after binding: index, tag, alias, the scan over all outputs, memory planning, and dim order.
    • It also makes several Python attribute calls (is_contiguous, nbytes twice) and builds a TorchTensorView twice: once in set_output (2269) and again at 2315.
    • In the decode-style loops this API targets, that is Python overhead on every step.

    Check the static properties once in set_output, and on each call re-check only what the caller can change (storage, size, dtype, device), using a single view.

Code Quality

  • The USE_ATEN_LIB branches are untested, and three of the new tests would fail on them (pybindings.cpp:2264-2267, 2309-2313, 2408-2418). executorch.runtime only imports portable_lib (runtime/__init__.py:129-140), so test_runtime.py never runs the ATen code.

    • Earlier outputs read the bound tensor. ATen's set_tensor_data swaps the data pointer of the output's StorageImpl in place (runtime/core/exec_aten/util/tensor_util_aten.cpp:206-207). An output returned earlier with clone_outputs=False shares that TensorImpl, so after set_output + execute it reads the bound tensor, and test_set_output_keeps_earlier_outputs_alive would fail.
    • They can also read freed memory. The DataPtr does not own the bound tensor's memory. After set_output(other, i) and dropping the old bound tensor, the earlier output reads freed memory. The comment at 2222-2223 promises the opposite.
    • Inputs returned as outputs are not checked. For an output that is an input, ATen returns caller without checking its storage or its metadata (2410-2412). So test_aliased_output_rejects_replaced_input_storage and test_aliased_output_preserves_written_metadata would fail.

    Move these tests to extension/pybindings/test/test_pybindings.py, which runs against whichever kernel library is built (kernel_mode). Then either make ATen match portable or document the difference.

  • fed_inputs_ duplicates existing state (pybindings.cpp:2068-2079, 2493).

    • In portable builds, torch_inputs_ already holds the owner of every torch input.
    • In both builds, fed_inputs_ keeps every torch input alive until the next set_inputs, including memory-planned inputs that were copied in. In ATen, that keeps a large planned input alive after the call for no reason.
    • The only entries ever read are those for inputs in inputs_aliased_by_outputs_.

Recommendation

Request Changes

  • The post-approval merge changed [PyBindings] Remove ATen dependency from portable bindings #23089's tested clone_outputs=False behavior for existing portable callers and rewrote that test to match. That needs one consistent rule and a re-review.
  • In the untested ATen build, earlier outputs read the bound tensor, or freed memory after a rebind.
  • Several new tests cannot fail.

The other items are minor.

IlyasMoutawwakil and others added 3 commits October 27, 2026 21:59
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>
@IlyasMoutawwakil

IlyasMoutawwakil commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Addressed in 7b552a3 and b0b8f9c:

  • clone_outputs=False for an output that is an input: it now always follows main ([PyBindings] Remove ATen dependency from portable bindings #23089): the detached alias of the caller's storage, whatever its strides. Only an output bound with set_output comes back as the caller's own tensor. test_borrowed_external_output_does_not_retain_previous_outputs is back to main's version, and the per-input bookkeeping (fed_inputs_) the identity path needed is gone.
  • ATen: set_output now refuses ATen builds, since swapping the output's data pointer in place would make earlier outputs read the bound tensor, or freed memory after a rebind. With nothing bound, the ATen path is unchanged from main.
  • Overlapping bindings: a binding that shares memory with another bound output, or with an input, raises.
  • Errors: rejected bindings raise ValueError like the input checks, a dtype/device mismatch names both sides, and the alias message no longer claims the output is returned as the caller's tensor.
  • Tests: dropped the ones that could not fail (test_bound_output_is_not_cloned, the meta "wrong device" case), the last check of test_set_output_rejects_mutated_input now runs on a fresh cache, and new tests cover shared memory and execute() without inputs after a binding (alias detection being deferred). The failed set_inputs test was rewritten in the main merge: the method refuses to run until a full set is installed.
  • Description: rewritten; mutated unplanned inputs are handled as on main, so it now only describes set_output.

Left for a follow-up: unbinding (set_output(None, i)), checking a binding's static properties once instead of on every call, and making set_output independent of call order (needs a public way to ask whether output i is input j without marking the input as set).

@metascroy the main merge changed how a mutated unplanned input is returned with clone_outputs=False (now main's detached alias rather than the tensor itself), and set_output now raises ValueError and refuses ATen builds. Could you take another look?

runtime/test/test_runtime.py and extension/pybindings/test/test_pybindings.py pass locally on a portable build, except test_quantized_ops / test_method_quantized_ops, which need the quantized kernels library that build doesn't load.

IlyasMoutawwakil and others added 3 commits October 27, 2026 23:44
- 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>
@IlyasMoutawwakil

Copy link
Copy Markdown
Contributor Author

CI on b0b8f9c, for the record:

  • lintrunner: three clang-format line wraps in pybindings.cpp, applied.
  • unittest-buck, test-qnn-buck-build-linux: "Attempted to register target root//backends/transforms:fuse_batch_norm_with_linear twice", from the main snapshot merged earlier; Remove duplicate fuse_batch_norm_with_linear target #24213 removed the duplicate on main.

I merged main again, which brings no changes to extension/pybindings or runtime.

@nil-is-all
nil-is-all merged commit 1935655 into pytorch:main Oct 28, 2026
464 checks passed
Erik-Lundell pushed a commit to Erik-Lundell/executorch that referenced this pull request Oct 28, 2026
… 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. community: contribution PRs coming from community (excluding hardware partners) module: arm Issues related to arm backend module: mlx Issues related to MLX Backend: Metal-accelerated inference on Apple Silicon

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants