From fc99f94586c373292a8c2c8e9aacfd1150379ad5 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 28 Sep 2026 16:11:34 +0200 Subject: [PATCH 01/10] allow inplace io binding in mlx and add python runtime interface --- backends/mlx/ops.py | 28 +++++++ backends/mlx/patterns.py | 10 ++- backends/mlx/test/test_ops.py | 72 +++++++++++++++++ extension/pybindings/pybindings.cpp | 117 +++++++++++++++++++++++++++- extension/pybindings/pybindings.pyi | 2 + runtime/__init__.py | 18 +++++ runtime/test/test_runtime.py | 62 +++++++++++++++ 7 files changed, 304 insertions(+), 5 deletions(-) diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 3973f1c80d7..0cd379846b4 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -94,6 +94,7 @@ GreaterEqualNode, GreaterNode, IdCopyNode, + IndexCopyNode, IntOrVid, IntOrVidOrTid, ItemIntNode, @@ -2406,6 +2407,33 @@ def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +@REGISTRY.register(target=[torch.ops.aten.index_copy.default]) +def _index_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """index_copy into a new tensor, leaving `dst` as it was. + + The INDEX_COPY pattern writes in place when `dst` is a buffer the method + mutates; anything else, a mutated user input included, lands here and the + caller writes the result back. + """ + args = P.args(n) + require_args(args, 4, 4, "aten.index_copy.default") + require_kwargs(P.kwargs(n), set(), "aten.index_copy.default") + dst, axis, indices, update = args + if not isinstance(axis, int): + raise ValueError(f"aten.index_copy.default requires a literal int dim, got {axis}") + out = P.make_or_get_slot(n) + P.emit( + IndexCopyNode( + dst=P.slot_to_tid(dst), + update=P.slot_to_tid(update), + indices=P.slot_to_tid(indices), + out=P.slot_to_tid(out), + axis=axis, + ) + ) + return out + + @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: args = P.args(n) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 8abd8933e4f..c1d682f9a1c 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -123,6 +123,10 @@ def _int4_qparams_lower(qdata_node: Node, scale_node: Node, group_size: int) -> class IndexCopyHandler(PatternHandler): """ Pattern for index-based updates on mutable buffers. + + Only a buffer the method mutates is written in place: the delegate owns it + across calls. A mutated user input is not the delegate's to keep, so its + index_copy is left to the functional handler and written back by the caller. """ def __init__( @@ -148,10 +152,8 @@ def maybe_create( # noqa: C901 if not match_target(index_copy_node, torch.ops.aten.index_copy.default): return None - # index_copy should write to a mutable input/buffer to be an index update. - if (index_copy_node.name not in ep.graph_signature.buffers_to_mutate) and ( - index_copy_node.name not in ep.graph_signature.user_inputs_to_mutate - ): + # index_copy should write to a mutable buffer to be an index update. + if index_copy_node.name not in ep.graph_signature.buffers_to_mutate: return None # index_copy(dst, axis, indices, update) diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 6942cd7ae1b..9b71fc1fbfb 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -4750,6 +4750,78 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (indices, update) +class IndexCopyInputModel(nn.Module): + """Model that performs index_copy on a tensor passed in as an input. + + Unlike a buffer, a user input is not the delegate's to own: the MLX graph + returns the updated tensor and the method writes it back into the input. + """ + + def __init__(self, axis: int = 0): + super().__init__() + self.axis = axis + + def forward( + self, data: torch.Tensor, indices: torch.Tensor, update: torch.Tensor + ) -> torch.Tensor: + data.index_copy_(self.axis, indices, update) + return data.sum(self.axis) + + +@register_test +class IndexCopyInputTest(OpTestCase): + """Test case for index_copy into a mutated user input (e.g. a static KV + cache passed to the method), lowered to a functional IndexCopyNode.""" + + name = "index_copy_input" + rtol = 1e-5 + atol = 1e-5 + expected_node_counts = {"IndexCopyNode": 1} + + def __init__( + self, + buffer_size: int = 128, + feature_dim: int = 64, + num_indices: int = 8, + axis: int = 0, + ): + self.buffer_size = buffer_size + self.feature_dim = feature_dim + self.num_indices = num_indices + self.axis = axis + self.name = ( + f"index_copy_input_axis{axis}_{buffer_size}x{feature_dim}_idx{num_indices}" + ) + + @classmethod + def get_test_configs(cls) -> List["IndexCopyInputTest"]: + return [ + cls(buffer_size=128, feature_dim=64, num_indices=8, axis=0), + cls(buffer_size=64, feature_dim=32, num_indices=8, axis=1), + ] + + def create_model(self) -> nn.Module: + return IndexCopyInputModel(axis=self.axis) + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + shape = ( + (self.buffer_size, self.feature_dim) + if self.axis == 0 + else (self.feature_dim, self.buffer_size) + ) + update_shape = list(shape) + update_shape[self.axis] = self.num_indices + indices = torch.randperm(self.buffer_size)[: self.num_indices].to(torch.int64) + return (torch.randn(shape), indices, torch.randn(update_shape)) + + def compute_expected_outputs(self, model, test_inputs): + # The method returns the mutated input first, then the model's own + # output. Mutate a copy, so the saved inputs stay as they were fed. + data, indices, update = test_inputs + data = data.clone() + return [data, model(data, indices, update)] + + class SplitWithSizesModel(nn.Module): """Model that splits a tensor into chunks with specified sizes.""" diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 27255b3b8ac..5df44aa64d2 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1284,12 +1284,14 @@ struct PyMethod final { input_tensors.reserve(inputs_size); #endif + fed_inputs_.clear(); // Convert python objects into EValues. for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; const std::string& type_str = py::str(python_input.get_type()); if (type_str == "") { auto at_tensor = python_input.cast(); + fed_inputs_.push_back(at_tensor); #ifdef USE_ATEN_LIB EValue evalue(at_tensor); @@ -1377,11 +1379,61 @@ struct PyMethod final { static_cast(set_inputs_status)); } + // Writes output `index` straight into `tensor`'s storage from the next + // 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) { + const auto num_outputs = method_->outputs_size(); + if (index >= num_outputs) { + throw std::out_of_range( + "output index " + std::to_string(index) + " is out of range for " + + std::to_string(num_outputs) + " outputs"); + } + 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); + THROW_IF_ERROR( + status, + "set_output_data_ptr failed for output %zu with error 0x%" PRIx32, + index, + static_cast(status)); + bound_outputs_.resize(num_outputs); + 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(); + } + } + void execute() { const auto num_outputs = method_->outputs_size(); allocate_output_storages(); std::vector> output_storage_spans(num_outputs); + caller_outputs_.assign(num_outputs, std::nullopt); for (int i = 0; i < output_storages_.size(); ++i) { + // 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)) { + auto& bound = bound_outputs_[i].value(); + output_storage_spans[i] = Span( + static_cast(mutable_tensor_data_ptr_no_cow(bound)), + bound.nbytes()); + caller_outputs_[i] = bound; + continue; + } + // An unplanned output that is one of the inputs (the write-back of an + // input the method mutates) already points at the tensor fed for it. + // Giving it storage here would repoint that input too, so the method + // would read and write this storage instead of the caller's tensor. + if (auto fed = fed_input_aliased_by(i)) { + caller_outputs_[i] = fed; + continue; + } output_storage_spans[i] = Span(output_storages_[i].data(), output_storages_[i].size()); } @@ -1418,7 +1470,25 @@ struct PyMethod final { static_cast(get_outputs_status)); // Retrieve outputs - return get_outputs_as_py_list(result, clone_outputs); + py::list outputs = get_outputs_as_py_list(result, clone_outputs); + // An output bound with set_output(), or one that is a fed input, already + // lives in the caller's tensor, so it is returned as that tensor, never + // cloned. A dynamic output may have been resized smaller than that tensor, + // so view it at its written size. + for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { + if (!caller_outputs_[i].has_value() || !result[i].isTensor()) { + continue; + } + const auto& bound = caller_outputs_[i].value(); + const auto& written = result[i].toTensor(); + std::vector sizes(written.sizes().begin(), written.sizes().end()); + if (bound.sizes().vec() == sizes) { + outputs[i] = py::cast(bound); + } else { + outputs[i] = py::cast(bound.view({-1}).narrow(0, 0, written.numel()).view(sizes)); + } + } + return outputs; } py::list call(const py::sequence& inputs, bool clone_outputs = true) { @@ -1464,6 +1534,45 @@ struct PyMethod final { // Need to keep-alive output storages until they can be compared in case of // bundled programs. std::vector> output_storages_; + // Caller tensors bound as output storage with set_output(), by output index. + std::vector> bound_outputs_; + // The tensors passed to the last set_inputs(), in order. + std::vector fed_inputs_; + // Per output of the last execute(), the caller tensor it lives in, if any. + std::vector> caller_outputs_; + // Outputs found to be one of the method's inputs. + std::vector aliases_input_; + + bool is_bound_output(size_t index) const { + return index < bound_outputs_.size() && bound_outputs_[index].has_value(); + } + + // The fed input whose memory unplanned output `index` points at, if any: + // the method returns a mutated input as that very input. That holds for + // every call once seen, so its unused storage is released then. + std::optional 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(); + } + return fed; + } + } + return std::nullopt; + } void allocate_output_storages() { const auto num_outputs = method_->outputs_size(); @@ -2012,6 +2121,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { py::class_(m, "ExecuTorchMethod") .def("set_inputs", &PyMethod::set_inputs, py::arg("inputs"), call_guard) .def("execute", &PyMethod::execute, call_guard) + .def( + "set_output", + &PyMethod::set_output, + py::arg("tensor"), + py::arg("index"), + call_guard) .def( "get_outputs", &PyMethod::get_outputs, diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index f2a037ce816..5832a94146e 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -9,6 +9,7 @@ from __future__ import annotations from typing import Any, Dict, Enum, List, Optional, Sequence, Tuple +import torch from executorch.exir._warnings import experimental @experimental("This API is experimental and subject to change without notice.") @@ -101,6 +102,7 @@ class ExecuTorchMethod: # pyre-ignore[2]: "Any" in parameter type annotations. def set_inputs(self, inputs: Sequence[Any]) -> None: ... def execute(self) -> None: ... + def set_output(self, tensor: torch.Tensor, index: int) -> None: ... # pyre-ignore[3]: "Any" in return type annotations. def get_outputs(self, clone_outputs: bool = True) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. diff --git a/runtime/__init__.py b/runtime/__init__.py index ba83ff54856..c8ee11aa410 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -113,6 +113,8 @@ from types import ModuleType from typing import Any, BinaryIO, Dict, List, Optional, Sequence, Set, Union +import torch + try: from executorch.extension.pybindings.portable_lib import ( # type: ignore[import-not-found] ExecuTorchMethod, @@ -146,6 +148,22 @@ def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: """ return self._method(inputs) + def set_output(self, tensor: torch.Tensor, index: int) -> None: + """Binds output ``index`` to ``tensor``'s storage for every later call. + + 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 + memory planned (export with ``alloc_graph_output=False``) and ``tensor`` + must be contiguous and at least as large as the output. + + Args: + tensor: The tensor that stores the output from now on. + index: The index of the output to bind. + """ + self._method.set_output(tensor, index) + @property def metadata(self) -> MethodMeta: """Gets the metadata for the method. diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 784ac35a1fe..f4bf7aae9da 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -16,9 +16,36 @@ ModuleAdd, ModuleMulti, ) +from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass from executorch.runtime import Runtime, Verification +class ModuleCacheUpdate(torch.nn.Module): + """Writes one row of a cache passed in as an input, like a static KV cache.""" + + def forward(self, x, cache, pos): + cache.index_copy_(0, pos, x) + return cache.sum(0, keepdim=True) + + +def _load_unplanned_io_method(module, inputs): + """Loads `forward` exported with neither its inputs nor its outputs memory + planned. Returns the program buffer too: the loaded program reads it in + place, so it must outlive the method.""" + config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ) + ) + buffer = to_edge(torch.export.export(module, inputs)).to_executorch(config).buffer + runtime = Runtime.get() + method = runtime.load_program( + buffer, verification=Verification.Minimal + ).load_method("forward") + return method, buffer + + class RuntimeTest(unittest.TestCase): def test_smoke(self): ep, inputs = create_program(ModuleAdd()) @@ -31,6 +58,41 @@ def test_smoke(self): outputs = method.execute(inputs) self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + def test_mutated_input_is_updated_in_place(self): + # An unplanned input the method mutates is written where it lives, and + # its write-back output is that very tensor rather than a copy. + x, cache, pos = torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) + method, _buffer = _load_unplanned_io_method(ModuleCacheUpdate(), (x, cache, pos)) + for step, row in enumerate((1, 2)): + outputs = method.execute((x * (step + 1), cache, torch.tensor([row]))) + self.assertIs(outputs[0], cache) + expected = torch.zeros(3, 4) + expected[1], expected[2] = 1.0, 2.0 + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.equal(outputs[1], expected.sum(0, keepdim=True))) + + def test_set_output_binds_caller_storage(self): + x, cache, pos = torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) + method, _buffer = _load_unplanned_io_method(ModuleCacheUpdate(), (x, cache, pos)) + # Outputs are [the mutated cache, the sum]: bind the sum to a caller + # tensor, which the method then writes and returns uncloned. + summed = torch.empty(1, 4) + method.set_output(summed, 1) + for row in (1, 2): + outputs = method.execute((x, cache, torch.tensor([row]))) + self.assertIs(outputs[1], summed) + self.assertTrue(torch.equal(summed, cache.sum(0, keepdim=True))) + + def test_set_output_rejects_bad_bindings(self): + x, cache, pos = torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) + method, _buffer = _load_unplanned_io_method(ModuleCacheUpdate(), (x, cache, pos)) + with self.assertRaises(IndexError): + method.set_output(torch.empty(1, 4), 2) + with self.assertRaises(RuntimeError): + method.set_output(torch.zeros(4, 3).t(), 0) + with self.assertRaises(RuntimeError): + method.set_output(torch.empty(1, 2), 1) + def test_module_with_multiple_method_names(self): ep, inputs = create_program(ModuleMulti()) runtime = Runtime.get() From 84f37056c5554a520e13fdc829530ec61f114a17 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 6 Oct 2026 12:26:03 +0200 Subject: [PATCH 02/10] Add Method.set_output and write unplanned input mutations in place in 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 --- extension/pybindings/pybindings.cpp | 171 +++++++++++++++++++++++++++- extension/pybindings/pybindings.pyi | 18 ++- runtime/__init__.py | 41 ++++++- runtime/test/targets.bzl | 2 + runtime/test/test_runtime.py | 141 +++++++++++++++++++++++ 5 files changed, 367 insertions(+), 6 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 27255b3b8ac..5a18f2d099d 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -108,6 +108,8 @@ using torch::executor::ETDumpGen; #ifndef USE_ATEN_LIB using ::executorch::extension::alias_attensor_to_etensor; using ::executorch::extension::alias_etensor_to_attensor; +using ::executorch::extension::executorch_to_torch_device; +using ::executorch::extension::executorch_to_torch_scalar_type; using ::executorch::extension::torch_to_executorch_device; using ::executorch::extension::torch_to_executorch_scalar_type; #endif // !USE_ATEN_LIB @@ -1284,12 +1286,14 @@ struct PyMethod final { input_tensors.reserve(inputs_size); #endif + fed_inputs_.assign(inputs_size, std::nullopt); // Convert python objects into EValues. for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; const std::string& type_str = py::str(python_input.get_type()); if (type_str == "") { auto at_tensor = python_input.cast(); + fed_inputs_[i] = at_tensor; #ifdef USE_ATEN_LIB EValue evalue(at_tensor); @@ -1375,13 +1379,94 @@ struct PyMethod final { "method->set_inputs() for method '%s' failed with error 0x%" PRIx32, method_->method_meta().name(), static_cast(set_inputs_status)); + find_inputs_aliased_by_outputs(); + } + + // Writes output `index` straight into `tensor`'s storage from the next + // execute() on, instead of into storage this method allocates. The binding + // persists across calls and the tensor is kept alive until it is replaced. + // Storage the output used before is kept alive, so tensors already returned + // for it with clone_outputs=False stay valid. + void set_output(at::Tensor tensor, size_t index) { + const auto num_outputs = method_->outputs_size(); + if (index >= num_outputs) { + throw std::out_of_range( + "output index " + std::to_string(index) + " is out of range for " + + std::to_string(num_outputs) + " outputs"); + } + const auto& output = method_->get_output(index); + if (!output.isTensor()) { + throw std::runtime_error( + "output " + std::to_string(index) + " is not a tensor"); + } + if (aliased_input(index).has_value()) { + throw std::runtime_error( + "output " + std::to_string(index) + + " is an input the method mutates, which is already written in place"); + } + const auto meta = method_->method_meta().output_tensor_meta(index); + THROW_IF_ERROR( + meta.error(), "Failed to get output tensor meta for output %zu", index); + if (meta->is_memory_planned()) { + throw std::runtime_error( + "output " + std::to_string(index) + + " is memory planned and cannot be bound"); + } +#ifdef USE_ATEN_LIB + const auto dtype = meta->scalar_type(); + const auto device = output.toTensor().device(); +#else + const auto dtype = executorch_to_torch_scalar_type(meta->scalar_type()); + const auto device = executorch_to_torch_device(output.toTensor().device()); +#endif + if (tensor.scalar_type() != dtype || tensor.device() != device) { + throw std::runtime_error( + "output " + std::to_string(index) + " is " + c10::toString(dtype) + + " on " + device.str() + ", so it cannot be bound to a " + + c10::toString(tensor.scalar_type()) + " tensor on " + + tensor.device().str()); + } + if (!tensor.is_contiguous()) { + throw std::runtime_error( + "output " + std::to_string(index) + + " must be bound to a contiguous tensor"); + } + if (tensor.nbytes() < meta->nbytes()) { + throw std::runtime_error( + "output " + std::to_string(index) + " needs " + + std::to_string(meta->nbytes()) + " bytes, but the tensor has " + + std::to_string(tensor.nbytes())); + } + bound_outputs_.resize(num_outputs); + bound_outputs_[index] = tensor; } void execute() { const auto num_outputs = method_->outputs_size(); allocate_output_storages(); std::vector> output_storage_spans(num_outputs); - for (int i = 0; i < output_storages_.size(); ++i) { + caller_outputs_.assign(num_outputs, std::nullopt); + for (size_t i = 0; i < output_storages_.size(); ++i) { + if (is_bound_output(i)) { + if (aliased_input(i).has_value()) { + throw std::runtime_error( + "output " + std::to_string(i) + + " is an input the method mutates and cannot be bound"); + } + auto& bound = bound_outputs_[i].value(); + output_storage_spans[i] = Span( + static_cast(mutable_tensor_data_ptr_no_cow(bound)), + bound.nbytes()); + caller_outputs_[i] = bound; + continue; + } + // The write-back of an unplanned input the method mutates is that input, + // already pointing at the tensor fed for it. Giving it storage would + // repoint the input too. + if (const auto input = aliased_input(i)) { + caller_outputs_[i] = fed_inputs_[*input]; + continue; + } output_storage_spans[i] = Span(output_storages_[i].data(), output_storages_[i].size()); } @@ -1418,7 +1503,28 @@ struct PyMethod final { static_cast(get_outputs_status)); // Retrieve outputs - return get_outputs_as_py_list(result, clone_outputs); + py::list outputs = get_outputs_as_py_list(result, clone_outputs); + // An output bound with set_output(), or an input the method mutates, + // already lives in the caller's tensor, so it is returned as that tensor, + // never cloned. A dynamic bound output may have been resized smaller than + // its (contiguous) tensor, so view it at its written size. A mutated input + // keeps the sizes it was fed with, and may be channels-last. + for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { + if (!caller_outputs_[i].has_value() || !result[i].isTensor()) { + continue; + } + const auto& caller = caller_outputs_[i].value(); + const auto& written = result[i].toTensor(); + std::vector sizes( + written.sizes().begin(), written.sizes().end()); + if (!is_bound_output(i) || caller.sizes().vec() == sizes) { + outputs[i] = py::cast(caller); + } else { + outputs[i] = py::cast( + caller.view({-1}).narrow(0, 0, written.numel()).view(sizes)); + } + } + return outputs; } py::list call(const py::sequence& inputs, bool clone_outputs = true) { @@ -1464,6 +1570,55 @@ struct PyMethod final { // Need to keep-alive output storages until they can be compared in case of // bundled programs. std::vector> output_storages_; + // Caller tensors bound as output storage with set_output(), by output index. + std::vector> bound_outputs_; + // The tensors passed to the last set_inputs(), by input index. They stay + // alive until the next set_inputs(), even if the caller drops them (e.g. a + // large KV cache). + std::vector> fed_inputs_; + // Per output of the last execute(), the caller tensor it lives in, if any. + std::vector> caller_outputs_; + // Per output, the unplanned input it is, if any. Empty until the first + // set_inputs(). + std::vector> inputs_aliased_by_outputs_; + + bool is_bound_output(size_t index) const { + return index < bound_outputs_.size() && bound_outputs_[index].has_value(); + } + + std::optional aliased_input(size_t index) const { + return index < inputs_aliased_by_outputs_.size() + ? inputs_aliased_by_outputs_[index] + : std::nullopt; + } + + // An output that carries an input mutation is the same value as that input. + // get_input() marks the input as set, so this only runs once set_inputs() + // has set them all, leaving execute()'s check for unset inputs intact. + void find_inputs_aliased_by_outputs() { + const auto num_outputs = method_->outputs_size(); + if (inputs_aliased_by_outputs_.size() == num_outputs) { + return; + } + inputs_aliased_by_outputs_.assign(num_outputs, std::nullopt); + const auto meta = method_->method_meta(); + for (size_t i = 0; i < num_outputs; ++i) { + const auto output_meta = meta.output_tensor_meta(i); + if (!output_meta.ok() || output_meta->is_memory_planned()) { + continue; + } + for (size_t j = 0; j < method_->inputs_size(); ++j) { +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + const bool same = &method_->get_output(i) == &method_->get_input(j); +#pragma GCC diagnostic pop + if (same) { + inputs_aliased_by_outputs_[i] = j; + break; + } + } + } + } void allocate_output_storages() { const auto num_outputs = method_->outputs_size(); @@ -1490,8 +1645,10 @@ struct PyMethod final { output_tensor_meta.error(), "Failed to get output tensor meta for output %zu", i); - if (output_tensor_meta.get().is_memory_planned()) { - // Skip allocating storage for planned memory outputs. + if (output_tensor_meta.get().is_memory_planned() || is_bound_output(i) || + aliased_input(i).has_value()) { + // Skip allocating storage for planned memory outputs, and for outputs + // that live in a caller's tensor. output_storages_.emplace_back(); continue; } @@ -2012,6 +2169,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { py::class_(m, "ExecuTorchMethod") .def("set_inputs", &PyMethod::set_inputs, py::arg("inputs"), call_guard) .def("execute", &PyMethod::execute, call_guard) + .def( + "set_output", + &PyMethod::set_output, + py::arg("tensor"), + py::arg("index"), + call_guard) .def( "get_outputs", &PyMethod::get_outputs, diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index f2a037ce816..02ac38471a2 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -9,6 +9,7 @@ from __future__ import annotations from typing import Any, Dict, Enum, List, Optional, Sequence, Tuple +import torch from executorch.exir._warnings import experimental @experimental("This API is experimental and subject to change without notice.") @@ -93,6 +94,17 @@ class ExecuTorchProgram: class ExecuTorchMethod: """ExecuTorchMethod is a Python wrapper around a loaded C++ method. + Outputs are cloned when ``clone_outputs`` is True, except for outputs that + already live in a caller's tensor, which are returned as that tensor: + + - an output bound with ``set_output``; + - the write-back of an input the method mutates, when that input is not + memory planned (exported with ``alloc_graph_input=False``). The mutation + is written into the tensor passed in for it. A memory planned input is + copied into the method, so the caller's tensor is not updated. + + Later calls overwrite these tensors. + .. warning:: This API is experimental and subject to change without notice. @@ -101,8 +113,12 @@ class ExecuTorchMethod: # pyre-ignore[2]: "Any" in parameter type annotations. def set_inputs(self, inputs: Sequence[Any]) -> None: ... def execute(self) -> None: ... + def set_output(self, tensor: torch.Tensor, index: int) -> None: ... # pyre-ignore[3]: "Any" in return type annotations. - def get_outputs(self, clone_outputs: bool = True) -> List[Any]: ... + def get_outputs(self, clone_outputs: bool = True) -> List[Any]: + """Returns the outputs of the last ``execute``. Outputs living in a + caller's tensor are never cloned, see ``ExecuTorchMethod``.""" + ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def call( self, diff --git a/runtime/__init__.py b/runtime/__init__.py index ba83ff54856..f6cd967099b 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -111,7 +111,20 @@ import functools from pathlib import Path from types import ModuleType -from typing import Any, BinaryIO, Dict, List, Optional, Sequence, Set, Union +from typing import ( + Any, + BinaryIO, + Dict, + List, + Optional, + Sequence, + Set, + TYPE_CHECKING, + Union, +) + +if TYPE_CHECKING: + import torch try: from executorch.extension.pybindings.portable_lib import ( # type: ignore[import-not-found] @@ -141,11 +154,37 @@ def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: Args: inputs: A sequence of input values, typically torch.Tensor objects. + Outputs are copies, except an output bound with :meth:`set_output` and + the write-back of an input the method mutates when that input is not + memory planned (exported with ``alloc_graph_input=False``): these are + returned as the caller's tensor, which the mutation is written into + and later calls overwrite. A memory planned input is copied into the + method, so the tensor passed in for it is not updated. + Returns: A list of output values, typically torch.Tensor objects. """ return self._method(inputs) + def set_output(self, tensor: "torch.Tensor", index: int) -> None: + """Binds output ``index`` to ``tensor``'s storage for every later call. + + The method then writes that output straight into ``tensor`` and returns + ``tensor`` itself, uncloned, instead of a copy of its own storage. Every + later call overwrites ``tensor``, so clone a returned output to keep it. + + The output must not be memory planned (export with + ``alloc_graph_output=False``) and must not be an input the method + mutates, which is already written in place. ``tensor`` must have the + output's dtype and device, be contiguous, and hold at least as many + bytes as the output's largest shape. + + Args: + tensor: The tensor that stores the output from now on. + index: The index of the output to bind. + """ + self._method.set_output(tensor, index) + @property def metadata(self) -> MethodMeta: """Gets the metadata for the method. diff --git a/runtime/test/targets.bzl b/runtime/test/targets.bzl index 74309a59e01..0ef92d4be23 100644 --- a/runtime/test/targets.bzl +++ b/runtime/test/targets.bzl @@ -11,6 +11,8 @@ def define_common_targets(is_fbcode = False): "//executorch/extension/pybindings/test:make_test", "//executorch/runtime:runtime", "//executorch/devtools/etdump:serialize", + "//executorch/exir:lib", + "//executorch/exir/passes:lib", ], ) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 19f25df6ab5..13167b6e0c3 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -12,6 +12,8 @@ from pathlib import Path import torch +from executorch.exir import ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass from executorch.extension.pybindings.test.make_test import ( create_program, @@ -21,6 +23,47 @@ from executorch.runtime import Runtime, Verification +class ModuleCacheUpdate(torch.nn.Module): + """Writes one row of a cache passed in as an input, like a static KV cache.""" + + def forward(self, x, cache, pos): + cache.index_copy_(0, pos, x) + return cache.sum(0, keepdim=True) + + +class ModuleRecurrent(torch.nn.Module): + def forward(self, x, state): + return state + x + + +class ModuleDouble(torch.nn.Module): + def forward(self, x): + return x * 2 + + +def _cache_inputs(): + return torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) + + +def _load_unplanned_io_method(module, inputs, dynamic_shapes=None, plan_inputs=False): + """Loads `forward` exported with neither its inputs (unless `plan_inputs`) + nor its outputs memory planned. Returns the program buffer too: the loaded + program reads it in place, so it must outlive the method.""" + config = ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=plan_inputs, alloc_graph_output=False + ) + ) + exported = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + buffer = to_edge(exported).to_executorch(config).buffer + program = Runtime.get().load_program(buffer, verification=Verification.Minimal) + return program.load_method("forward"), buffer + + +def _load_cache_update_method(): + return _load_unplanned_io_method(ModuleCacheUpdate(), _cache_inputs()) + + class RuntimeTest(unittest.TestCase): def test_runtime_import_does_not_load_the_exporter(self): # A fresh interpreter, since this process has already imported exir. @@ -44,6 +87,104 @@ def test_smoke(self): outputs = method.execute(inputs) self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + def test_mutated_input_is_updated_in_place(self): + # An unplanned input the method mutates is written where it lives, and + # its write-back output is that very tensor rather than a copy. + x, cache, _ = _cache_inputs() + method, _buffer = _load_cache_update_method() + for step, row in enumerate((1, 2)): + outputs = method.execute((x * (step + 1), cache, torch.tensor([row]))) + self.assertIs(outputs[0], cache) + expected = torch.zeros(3, 4) + expected[1], expected[2] = 1.0, 2.0 + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.equal(outputs[1], expected.sum(0, keepdim=True))) + + def test_mutated_planned_input_is_not_updated(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_unplanned_io_method( + ModuleCacheUpdate(), _cache_inputs(), plan_inputs=True + ) + outputs = method.execute((x, cache, pos)) + self.assertIsNot(outputs[0], cache) + self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) + self.assertTrue(torch.equal(outputs[0][1], torch.ones(4))) + + def test_output_fed_back_as_input_is_not_that_input(self): + # An uncloned output fed back in shares the input's memory, but is not + # an input the method mutates. + x, state = torch.ones(2), torch.zeros(2) + method, _buffer = _load_unplanned_io_method(ModuleRecurrent(), (x, state)) + for step in (1, 2): + fed = state + state = method._method((x, fed), clone_outputs=False)[0] + self.assertIsNot(state, fed) + self.assertTrue(torch.equal(state, torch.full((2,), float(step)))) + fresh = torch.zeros(2) + outputs = method._method((x, fresh), clone_outputs=False) + self.assertIsNot(outputs[0], fresh) + self.assertTrue(torch.equal(fresh, torch.zeros(2))) + self.assertTrue(torch.equal(outputs[0], torch.ones(2))) + + def test_set_output_binds_caller_storage(self): + x, cache, _ = _cache_inputs() + method, _buffer = _load_cache_update_method() + # Outputs are [the mutated cache, the sum]: bind the sum to a caller + # tensor, which the method then writes and returns uncloned. + summed = torch.empty(1, 4) + method.set_output(summed, 1) + for row in (1, 2): + outputs = method.execute((x, cache, torch.tensor([row]))) + self.assertIs(outputs[1], summed) + self.assertTrue(torch.equal(summed, cache.sum(0, keepdim=True))) + + def test_set_output_keeps_earlier_outputs_alive(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + earlier = method._method((x, cache, pos), clone_outputs=False)[1] + expected = earlier.clone() + method.set_output(torch.empty(1, 4), 1) + method.execute((x * 2, cache, torch.tensor([2]))) + self.assertTrue(torch.equal(earlier, expected)) + + def test_set_output_narrows_dynamic_output(self): + dynamic_shapes = ({0: torch.export.Dim("n", max=8)},) + method, _buffer = _load_unplanned_io_method( + ModuleDouble(), (torch.ones(4, 3),), dynamic_shapes + ) + bound = torch.empty(8, 3) + method.set_output(bound, 0) + x = torch.arange(6.0).reshape(2, 3) + output = method.execute((x,))[0] + self.assertEqual(output.shape, (2, 3)) + self.assertEqual(output.data_ptr(), bound.data_ptr()) + self.assertTrue(torch.equal(output, x * 2)) + + def test_set_output_rejects_bad_bindings(self): + method, _buffer = _load_cache_update_method() + bad_bindings = { + "out of range": (IndexError, torch.empty(1, 4), 2), + "non contiguous": (RuntimeError, torch.empty(1, 8)[:, ::2], 1), + "too small": (RuntimeError, torch.empty(1, 2), 1), + "wrong dtype": (RuntimeError, torch.empty(1, 4, dtype=torch.float64), 1), + "wrong device": (RuntimeError, torch.empty(1, 4, device="meta"), 1), + } + for name, (error, tensor, index) in bad_bindings.items(): + with self.subTest(name), self.assertRaises(error): + method.set_output(tensor, index) + + def test_set_output_rejects_mutated_input(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method.execute((x, cache, pos)) + with self.assertRaises(RuntimeError): + method.set_output(torch.zeros(3, 4), 0) + # Bound before any inputs are set, the binding fails on the first call. + method, _buffer = _load_cache_update_method() + method.set_output(torch.zeros(3, 4), 0) + with self.assertRaises(RuntimeError): + method.execute((x, cache, pos)) + def test_module_with_multiple_method_names(self): ep, inputs = create_program(ModuleMulti()) runtime = Runtime.get() From 75fa40a6debd6649a940f234ebfc9f06ff211bac Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 6 Oct 2026 15:16:05 +0200 Subject: [PATCH 03/10] Harden PyMethod output bindings and document caller-owned outputs 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 --- docs/source/runtime-python-api-reference.rst | 2 +- extension/pybindings/pybindings.cpp | 49 ++++++++++++----- extension/pybindings/pybindings.pyi | 18 +++---- runtime/__init__.py | 31 ++++++----- runtime/test/test_runtime.py | 57 +++++++++++++++++--- 5 files changed, 112 insertions(+), 45 deletions(-) diff --git a/docs/source/runtime-python-api-reference.rst b/docs/source/runtime-python-api-reference.rst index 270cab9b61e..69ea9f2ea22 100644 --- a/docs/source/runtime-python-api-reference.rst +++ b/docs/source/runtime-python-api-reference.rst @@ -15,4 +15,4 @@ For detailed information on how APIs evolve and the deprecation process, please :members: method_names, load_method .. autoclass:: Method - :members: execute, metadata + :members: execute, set_output, metadata diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 5a18f2d099d..14b3b7a9fa8 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1286,14 +1286,14 @@ struct PyMethod final { input_tensors.reserve(inputs_size); #endif - fed_inputs_.assign(inputs_size, std::nullopt); + std::vector> fed_inputs(inputs_size); // Convert python objects into EValues. for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; const std::string& type_str = py::str(python_input.get_type()); if (type_str == "") { auto at_tensor = python_input.cast(); - fed_inputs_[i] = at_tensor; + fed_inputs[i] = at_tensor; #ifdef USE_ATEN_LIB EValue evalue(at_tensor); @@ -1379,6 +1379,7 @@ struct PyMethod final { "method->set_inputs() for method '%s' failed with error 0x%" PRIx32, method_->method_meta().name(), static_cast(set_inputs_status)); + fed_inputs_ = std::move(fed_inputs); find_inputs_aliased_by_outputs(); } @@ -1399,10 +1400,16 @@ struct PyMethod final { throw std::runtime_error( "output " + std::to_string(index) + " is not a tensor"); } - if (aliased_input(index).has_value()) { - throw std::runtime_error( - "output " + std::to_string(index) + - " is an input the method mutates, which is already written in place"); + if (const auto input = aliased_input(index)) { + throw std::runtime_error(aliased_input_message(index, *input)); + } + for (size_t i = 0; i < num_outputs; ++i) { + if (i != index && &method_->get_output(i) == &output) { + throw std::runtime_error( + "output " + std::to_string(index) + + " is the same value as output " + std::to_string(i) + + " and cannot be bound"); + } } const auto meta = method_->method_meta().output_tensor_meta(index); THROW_IF_ERROR( @@ -1448,11 +1455,6 @@ struct PyMethod final { caller_outputs_.assign(num_outputs, std::nullopt); for (size_t i = 0; i < output_storages_.size(); ++i) { if (is_bound_output(i)) { - if (aliased_input(i).has_value()) { - throw std::runtime_error( - "output " + std::to_string(i) + - " is an input the method mutates and cannot be bound"); - } auto& bound = bound_outputs_[i].value(); output_storage_spans[i] = Span( static_cast(mutable_tensor_data_ptr_no_cow(bound)), @@ -1460,9 +1462,9 @@ struct PyMethod final { caller_outputs_[i] = bound; continue; } - // The write-back of an unplanned input the method mutates is that input, - // already pointing at the tensor fed for it. Giving it storage would - // repoint the input too. + // An output that is an unplanned input (e.g. the write-back of a mutated + // input) already points at the tensor fed for it. Giving it storage + // would repoint the input too. if (const auto input = aliased_input(i)) { caller_outputs_[i] = fed_inputs_[*input]; continue; @@ -1618,6 +1620,25 @@ struct PyMethod final { } } } + // Only possible for a binding made before the first set_inputs(), when + // aliases were not known yet. + std::string error; + for (size_t i = 0; i < num_outputs; ++i) { + if (const auto input = aliased_input(i); input && is_bound_output(i)) { + bound_outputs_[i].reset(); + error += + aliased_input_message(i, *input) + ", so its binding was removed. "; + } + } + if (!error.empty()) { + throw std::runtime_error(error); + } + } + + static std::string aliased_input_message(size_t output, size_t input) { + return "output " + std::to_string(output) + " is input " + + std::to_string(input) + + ", which is returned as the tensor passed in for it"; } void allocate_output_storages() { diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 02ac38471a2..23d2273d0b2 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -97,13 +97,12 @@ class ExecuTorchMethod: Outputs are cloned when ``clone_outputs`` is True, except for outputs that already live in a caller's tensor, which are returned as that tensor: - - an output bound with ``set_output``; - - the write-back of an input the method mutates, when that input is not - memory planned (exported with ``alloc_graph_input=False``). The mutation - is written into the tensor passed in for it. A memory planned input is - copied into the method, so the caller's tensor is not updated. - - Later calls overwrite these tensors. + - an output bound with ``set_output``, which later calls overwrite; + - an output that is an input not memory planned (e.g. exported with + ``alloc_graph_input=False``), such as the write-back of a mutated input. + It is the tensor passed in for that input, mutation included. A memory + planned input is copied into the method, so the caller's tensor is not + updated. .. warning:: @@ -115,10 +114,7 @@ class ExecuTorchMethod: def execute(self) -> None: ... def set_output(self, tensor: torch.Tensor, index: int) -> None: ... # pyre-ignore[3]: "Any" in return type annotations. - def get_outputs(self, clone_outputs: bool = True) -> List[Any]: - """Returns the outputs of the last ``execute``. Outputs living in a - caller's tensor are never cloned, see ``ExecuTorchMethod``.""" - ... + def get_outputs(self, clone_outputs: bool = True) -> List[Any]: ... # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def call( self, diff --git a/runtime/__init__.py b/runtime/__init__.py index f6cd967099b..fc47e456718 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -151,16 +151,17 @@ def __init__(self, method: ExecuTorchMethod) -> None: def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: """Executes the method with the given inputs. + Outputs are copies, except for outputs that live in a caller's tensor: + an output bound with :meth:`set_output` (overwritten by later calls), + and an output that is an input not memory planned (e.g. exported with + ``alloc_graph_input=False``), such as the write-back of a mutated input, + which is returned as the tensor passed in for it, mutation included. A + memory planned input is copied into the method, so the tensor passed in + for it is not updated. + Args: inputs: A sequence of input values, typically torch.Tensor objects. - Outputs are copies, except an output bound with :meth:`set_output` and - the write-back of an input the method mutates when that input is not - memory planned (exported with ``alloc_graph_input=False``): these are - returned as the caller's tensor, which the mutation is written into - and later calls overwrite. A memory planned input is copied into the - method, so the tensor passed in for it is not updated. - Returns: A list of output values, typically torch.Tensor objects. """ @@ -170,14 +171,18 @@ def set_output(self, tensor: "torch.Tensor", index: int) -> None: """Binds output ``index`` to ``tensor``'s storage for every later call. The method then writes that output straight into ``tensor`` and returns - ``tensor`` itself, uncloned, instead of a copy of its own storage. Every - later call overwrites ``tensor``, so clone a returned output to keep it. + it uncloned (viewed at the written shape when that differs) instead of + a copy of its own storage. Every later call overwrites ``tensor``, so + clone a returned output to keep it. The output must not be memory planned (export with - ``alloc_graph_output=False``) and must not be an input the method - mutates, which is already written in place. ``tensor`` must have the - output's dtype and device, be contiguous, and hold at least as many - bytes as the output's largest shape. + ``alloc_graph_output=False``), must not be returned at another output + index too, and must not be an input, which is already returned as the + tensor passed in for it. That last case is only known once inputs have + been set: a binding made before the first call is removed by that call, + which raises. ``tensor`` must have the output's dtype and device, be + contiguous, and hold at least as many bytes as the output's largest + shape. Args: tensor: The tensor that stores the output from now on. diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 13167b6e0c3..13dfd3b7abb 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -41,17 +41,25 @@ def forward(self, x): return x * 2 +class ModuleDoubleTwice(torch.nn.Module): + def forward(self, x): + y = x * 2 + return y, y + + def _cache_inputs(): return torch.ones(1, 4), torch.zeros(3, 4), torch.tensor([1]) -def _load_unplanned_io_method(module, inputs, dynamic_shapes=None, plan_inputs=False): - """Loads `forward` exported with neither its inputs (unless `plan_inputs`) - nor its outputs memory planned. Returns the program buffer too: the loaded - program reads it in place, so it must outlive the method.""" +def _load_unplanned_io_method( + module, inputs, dynamic_shapes=None, plan_inputs=False, plan_outputs=False +): + """Loads `forward` exported with its inputs and outputs not memory planned, + unless asked. Returns the program buffer too: the loaded program reads it + in place, so it must outlive the method.""" config = ExecutorchBackendConfig( memory_planning_pass=MemoryPlanningPass( - alloc_graph_input=plan_inputs, alloc_graph_output=False + alloc_graph_input=plan_inputs, alloc_graph_output=plan_outputs ) ) exported = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) @@ -138,6 +146,28 @@ def test_set_output_binds_caller_storage(self): self.assertIs(outputs[1], summed) self.assertTrue(torch.equal(summed, cache.sum(0, keepdim=True))) + def test_set_output_rebinds(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + first, second = torch.zeros(1, 4), torch.zeros(1, 4) + method.set_output(first, 1) + method.set_output(second, 1) + outputs = method.execute((x, cache, pos)) + self.assertIs(outputs[1], second) + self.assertTrue(torch.equal(second, torch.ones(1, 4))) + self.assertTrue(torch.equal(first, torch.zeros(1, 4))) + + def test_failed_set_inputs_keeps_previous_inputs(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method._method.set_inputs((x, cache, pos)) + with self.assertRaises(RuntimeError): + method._method.set_inputs((x, torch.zeros(3, 4))) + method._method.execute() + outputs = method._method.get_outputs() + self.assertIs(outputs[0], cache) + self.assertTrue(torch.equal(cache[1], torch.ones(4))) + def test_set_output_keeps_earlier_outputs_alive(self): x, cache, pos = _cache_inputs() method, _buffer = _load_cache_update_method() @@ -173,17 +203,32 @@ def test_set_output_rejects_bad_bindings(self): with self.subTest(name), self.assertRaises(error): method.set_output(tensor, index) + def test_set_output_rejects_outputs_it_cannot_own(self): + x = torch.ones(2) + method, _buffer = _load_unplanned_io_method( + ModuleDouble(), (x,), plan_outputs=True + ) + with self.subTest("memory planned"), self.assertRaises(RuntimeError): + method.set_output(torch.empty(2), 0) + method, _buffer = _load_unplanned_io_method(ModuleDoubleTwice(), (x,)) + with self.subTest("returned twice"), self.assertRaises(RuntimeError): + method.set_output(torch.empty(2), 0) + def test_set_output_rejects_mutated_input(self): x, cache, pos = _cache_inputs() method, _buffer = _load_cache_update_method() method.execute((x, cache, pos)) with self.assertRaises(RuntimeError): method.set_output(torch.zeros(3, 4), 0) - # Bound before any inputs are set, the binding fails on the first call. + # Bound before any inputs are set, the first call removes the binding + # and raises, and later calls work. method, _buffer = _load_cache_update_method() method.set_output(torch.zeros(3, 4), 0) with self.assertRaises(RuntimeError): method.execute((x, cache, pos)) + outputs = method.execute((x, cache, pos)) + self.assertIs(outputs[0], cache) + self.assertTrue(torch.equal(cache[1], torch.ones(4))) def test_module_with_multiple_method_names(self): ep, inputs = create_program(ModuleMulti()) From a08f44ed5b7c585c6c35e7d645dda0bfe1f1d4cf Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 6 Oct 2026 21:41:38 +0200 Subject: [PATCH 04/10] Refuse execute after a failed set_inputs, reject non-contiguous bound outputs, and never clone caller-backed outputs Method::set_inputs does not roll back the inputs it set before failing, so execute() now refuses to run until a set_inputs() succeeds instead of writing into the new tensors while returning the old ones. set_output rejects outputs whose dim order is not contiguous, since only the data pointer is swapped and a channels-last output would be returned in the wrong order. Outputs that live in a caller's tensor are taken out of the result before cloning, so clone_outputs=True no longer copies a KV cache only to discard the copy. Authored with Claude Code. Co-Authored-By: Claude Opus 5.5 --- extension/pybindings/pybindings.cpp | 45 +++++++++++++++++++++++------ runtime/__init__.py | 2 +- runtime/test/test_runtime.py | 35 ++++++++++++++++++---- 3 files changed, 66 insertions(+), 16 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 14b3b7a9fa8..52c53c5d225 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -1286,6 +1286,9 @@ struct PyMethod final { input_tensors.reserve(inputs_size); #endif + // Method::set_inputs() does not roll back the inputs it set before + // failing, so execute() is refused until a set_inputs() succeeds. + inputs_valid_ = false; std::vector> fed_inputs(inputs_size); // Convert python objects into EValues. for (size_t i = 0; i < inputs_size; ++i) { @@ -1380,6 +1383,7 @@ struct PyMethod final { method_->method_meta().name(), static_cast(set_inputs_status)); fed_inputs_ = std::move(fed_inputs); + inputs_valid_ = true; find_inputs_aliased_by_outputs(); } @@ -1433,6 +1437,14 @@ struct PyMethod final { c10::toString(tensor.scalar_type()) + " tensor on " + tensor.device().str()); } + const auto dim_order = meta->dim_order(); + for (size_t d = 0; d < dim_order.size(); ++d) { + if (dim_order[d] != d) { + throw std::runtime_error( + "output " + std::to_string(index) + + " is not contiguous (e.g. channels-last) and cannot be bound"); + } + } if (!tensor.is_contiguous()) { throw std::runtime_error( "output " + std::to_string(index) + @@ -1449,6 +1461,11 @@ struct PyMethod final { } void execute() { + if (!inputs_valid_) { + throw std::runtime_error( + "The last set_inputs() failed, so inputs must be set again before " + "execute()"); + } const auto num_outputs = method_->outputs_size(); allocate_output_storages(); std::vector> output_storage_spans(num_outputs); @@ -1504,13 +1521,13 @@ struct PyMethod final { method_->method_meta().name(), static_cast(get_outputs_status)); - // Retrieve outputs - py::list outputs = get_outputs_as_py_list(result, clone_outputs); - // An output bound with set_output(), or an input the method mutates, - // already lives in the caller's tensor, so it is returned as that tensor, - // never cloned. A dynamic bound output may have been resized smaller than - // its (contiguous) tensor, so view it at its written size. A mutated input - // keeps the sizes it was fed with, and may be channels-last. + // An output bound with set_output(), or an output that is an input, + // already lives in the caller's tensor, so it is returned as that tensor + // and taken out of `result` before it could be cloned. A dynamic bound + // output may have been resized smaller than its (contiguous) tensor, so + // view it at its written size. An input keeps the sizes it was fed with, + // and may be channels-last. + std::vector caller_backed(result.size()); for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { if (!caller_outputs_[i].has_value() || !result[i].isTensor()) { continue; @@ -1520,11 +1537,20 @@ struct PyMethod final { std::vector sizes( written.sizes().begin(), written.sizes().end()); if (!is_bound_output(i) || caller.sizes().vec() == sizes) { - outputs[i] = py::cast(caller); + caller_backed[i] = py::cast(caller); } else { - outputs[i] = py::cast( + caller_backed[i] = py::cast( caller.view({-1}).narrow(0, 0, written.numel()).view(sizes)); } + result[i] = EValue(); + } + + // Retrieve outputs + py::list outputs = get_outputs_as_py_list(result, clone_outputs); + for (size_t i = 0; i < caller_backed.size(); ++i) { + if (caller_backed[i]) { + outputs[i] = caller_backed[i]; + } } return outputs; } @@ -1583,6 +1609,7 @@ struct PyMethod final { // Per output, the unplanned input it is, if any. Empty until the first // set_inputs(). std::vector> inputs_aliased_by_outputs_; + bool inputs_valid_ = true; bool is_bound_output(size_t index) const { return index < bound_outputs_.size() && bound_outputs_[index].has_value(); diff --git a/runtime/__init__.py b/runtime/__init__.py index fc47e456718..76504f597eb 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -182,7 +182,7 @@ def set_output(self, tensor: "torch.Tensor", index: int) -> None: been set: a binding made before the first call is removed by that call, which raises. ``tensor`` must have the output's dtype and device, be contiguous, and hold at least as many bytes as the output's largest - shape. + shape. The output itself must be contiguous, not e.g. channels-last. Args: tensor: The tensor that stores the output from now on. diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 13dfd3b7abb..058deb737d9 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -41,6 +41,11 @@ def forward(self, x): return x * 2 +class ModuleDoubleChannelsLast(torch.nn.Module): + def forward(self, x): + return (x * 2).contiguous(memory_format=torch.channels_last) + + class ModuleDoubleTwice(torch.nn.Module): def forward(self, x): y = x * 2 @@ -157,16 +162,28 @@ def test_set_output_rebinds(self): self.assertTrue(torch.equal(second, torch.ones(1, 4))) self.assertTrue(torch.equal(first, torch.zeros(1, 4))) - def test_failed_set_inputs_keeps_previous_inputs(self): + def test_execute_refused_after_failed_set_inputs(self): x, cache, pos = _cache_inputs() method, _buffer = _load_cache_update_method() method._method.set_inputs((x, cache, pos)) + # Fails on the position, after the new cache was already set. + new_cache = torch.zeros(3, 4) with self.assertRaises(RuntimeError): - method._method.set_inputs((x, torch.zeros(3, 4))) - method._method.execute() - outputs = method._method.get_outputs() - self.assertIs(outputs[0], cache) - self.assertTrue(torch.equal(cache[1], torch.ones(4))) + method._method.set_inputs((x, new_cache, torch.tensor([1.0]))) + with self.assertRaises(RuntimeError): + method._method.execute() + outputs = method.execute((x, new_cache, pos)) + self.assertIs(outputs[0], new_cache) + self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) + + def test_caller_backed_outputs_are_not_cloned(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method.set_output(torch.empty(1, 4), 1) + with torch.profiler.profile() as profile: + method.execute((x, cache, pos)) + clones = [e for e in profile.events() if e.name == "aten::clone"] + self.assertEqual(clones, []) def test_set_output_keeps_earlier_outputs_alive(self): x, cache, pos = _cache_inputs() @@ -213,6 +230,12 @@ def test_set_output_rejects_outputs_it_cannot_own(self): method, _buffer = _load_unplanned_io_method(ModuleDoubleTwice(), (x,)) with self.subTest("returned twice"), self.assertRaises(RuntimeError): method.set_output(torch.empty(2), 0) + image = torch.ones(1, 2, 3, 4) + method, _buffer = _load_unplanned_io_method( + ModuleDoubleChannelsLast(), (image,) + ) + with self.subTest("channels last"), self.assertRaises(RuntimeError): + method.set_output(torch.empty(1, 2, 3, 4), 0) def test_set_output_rejects_mutated_input(self): x, cache, pos = _cache_inputs() From a5765827ace07e196cce4737a7c8352cea9b7d1d Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Mon, 12 Oct 2026 09:02:50 +0200 Subject: [PATCH 05/10] Plan outputs too when testing a memory planned mutated input Since in-place results are planned onto the spec of the input they update (#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 --- runtime/test/test_runtime.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 3a5b341c27e..432006b6020 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -115,9 +115,12 @@ def test_mutated_input_is_updated_in_place(self): def test_mutated_planned_input_is_not_updated(self): x, cache, pos = _cache_inputs() + # The update is planned onto the input's own spec, which is also an + # output, so both must be planned for the cache to be planned. method, _buffer = _load_unplanned_io_method( - ModuleCacheUpdate(), _cache_inputs(), plan_inputs=True + ModuleCacheUpdate(), _cache_inputs(), plan_inputs=True, plan_outputs=True ) + self.assertTrue(method.metadata.input_tensor_meta(1).is_memory_planned()) outputs = method.execute((x, cache, pos)) self.assertIsNot(outputs[0], cache) self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) From b28bb4bb37d74cfc12e9e5954e0d69317f8080d0 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Wed, 14 Oct 2026 14:00:21 -0700 Subject: [PATCH 06/10] Clone outputs that are inputs unless clone_outputs=False 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 --- extension/pybindings/pybindings.cpp | 17 ++++++---- extension/pybindings/pybindings.pyi | 20 ++++++----- runtime/__init__.py | 21 ++++++------ runtime/test/test_runtime.py | 52 ++++++++++++++++++++++------- 4 files changed, 72 insertions(+), 38 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 68396e69f51..90ac6ad9a96 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -2033,15 +2033,18 @@ struct PyMethod final { method_->method_meta().name(), static_cast(get_outputs_status)); - // An output bound with set_output(), or an output that is an input, - // already lives in the caller's tensor, so it is returned as that tensor - // and taken out of `result` before it could be cloned. A dynamic bound - // output may have been resized smaller than its (contiguous) tensor, so - // view it at its written size. An input keeps the sizes it was fed with, - // and may be channels-last. + // An output bound with set_output() already lives in the caller's tensor, + // so it is returned as that tensor and taken out of `result` before it + // could be cloned. An output that is an input lives in the tensor passed + // in for it, which is returned only with clone_outputs=False: by default + // it is cloned like any other output. A dynamic bound output may have + // been resized smaller than its (contiguous) tensor, so view it at its + // written size. An input keeps the sizes it was fed with, and may be + // channels-last. std::vector caller_backed(result.size()); for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { - if (!caller_outputs_[i].has_value() || !result[i].isTensor()) { + if (!caller_outputs_[i].has_value() || !result[i].isTensor() || + (clone_outputs && !is_bound_output(i))) { continue; } const auto& caller = caller_outputs_[i].value(); diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 72cbb0a86cd..cbc6fc96d74 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -94,15 +94,17 @@ class ExecuTorchProgram: class ExecuTorchMethod: """ExecuTorchMethod is a Python wrapper around a loaded C++ method. - Outputs are cloned when ``clone_outputs`` is True, except for outputs that - already live in a caller's tensor, which are returned as that tensor: - - - an output bound with ``set_output``, which later calls overwrite; - - an output that is an input not memory planned (e.g. exported with - ``alloc_graph_input=False``), such as the write-back of a mutated input. - It is the tensor passed in for that input, mutation included. A memory - planned input is copied into the method, so the caller's tensor is not - updated. + Outputs are cloned when ``clone_outputs`` is True, except for an output + bound with ``set_output``, which is returned as the bound tensor and + overwritten by later calls. + + An input that is not memory planned (e.g. exported with + ``alloc_graph_input=False``) is used in place, so a mutation of it is + written into the tensor passed in for it. When such an input is also an + output (e.g. the write-back of a mutated input, or an input returned + unchanged), that output is a clone like any other by default, and is the + tensor passed in only with ``clone_outputs=False``. A memory planned input + is copied into the method, so the tensor passed in for it is not updated. .. warning:: diff --git a/runtime/__init__.py b/runtime/__init__.py index 84fbc073e1b..5ce0eca09d0 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -151,13 +151,14 @@ def __init__(self, method: ExecuTorchMethod) -> None: def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: """Executes the method with the given inputs. - Outputs are copies, except for outputs that live in a caller's tensor: - an output bound with :meth:`set_output` (overwritten by later calls), - and an output that is an input not memory planned (e.g. exported with - ``alloc_graph_input=False``), such as the write-back of a mutated input, - which is returned as the tensor passed in for it, mutation included. A - memory planned input is copied into the method, so the tensor passed in - for it is not updated. + Outputs are copies, except for an output bound with + :meth:`set_output`, which is returned as the bound tensor and + overwritten by later calls. An input that is not memory planned (e.g. + exported with ``alloc_graph_input=False``) is used in place, so a + mutation of it is written into the tensor passed in for it, while an + output that is that input is still returned as a copy. A memory planned + input is copied into the method, so the tensor passed in for it is not + updated. Args: inputs: A sequence of input values, typically torch.Tensor objects. @@ -177,9 +178,9 @@ def set_output(self, tensor: "torch.Tensor", index: int) -> None: The output must not be memory planned (export with ``alloc_graph_output=False``), must not be returned at another output - index too, and must not be an input, which is already returned as the - tensor passed in for it. That last case is only known once inputs have - been set: a binding made before the first call is removed by that call, + index too, and must not be an input, which already lives in the tensor + passed in for it. That last case is only known once inputs have been + set: a binding made before the first call is removed by that call, which raises. ``tensor`` must have the output's dtype and device, be contiguous, and hold at least as many bytes as the output's largest shape. The output itself must be contiguous, not e.g. channels-last. diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 432006b6020..0cea6fad129 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -36,6 +36,11 @@ def forward(self, x, state): return state + x +class ModuleInputAndIncrement(torch.nn.Module): + def forward(self, x): + return x, x + 1 + + class ModuleDouble(torch.nn.Module): def forward(self, x): return x * 2 @@ -101,18 +106,41 @@ def test_smoke(self): self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) def test_mutated_input_is_updated_in_place(self): - # An unplanned input the method mutates is written where it lives, and - # its write-back output is that very tensor rather than a copy. + # An unplanned input the method mutates is written where it lives. Its + # write-back output is a copy by default, and that very tensor only + # with clone_outputs=False. x, cache, _ = _cache_inputs() method, _buffer = _load_cache_update_method() - for step, row in enumerate((1, 2)): - outputs = method.execute((x * (step + 1), cache, torch.tensor([row]))) - self.assertIs(outputs[0], cache) expected = torch.zeros(3, 4) - expected[1], expected[2] = 1.0, 2.0 + expected[1] = 1.0 + outputs = method.execute((x, cache, torch.tensor([1]))) + self.assertIsNot(outputs[0], cache) + self.assertNotEqual(outputs[0].data_ptr(), cache.data_ptr()) + self.assertTrue(torch.equal(outputs[0], expected)) + self.assertTrue(torch.equal(cache, expected)) + outputs = method._method((x * 2, cache, torch.tensor([2])), clone_outputs=False) + self.assertIs(outputs[0], cache) + expected[2] = 2.0 self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.equal(outputs[1], expected.sum(0, keepdim=True))) + def test_returned_input_is_cloned_by_default(self): + # An unplanned input returned as an output is an independent copy with + # the default clone_outputs=True, as before set_output existed. + method, _buffer = _load_unplanned_io_method( + ModuleInputAndIncrement(), (torch.tensor([1.0]),) + ) + x = torch.tensor([1.0]) + y, z = method.execute((x,)) + self.assertIsNot(y, x) + self.assertNotEqual(y.data_ptr(), x.data_ptr()) + x.fill_(5.0) + self.assertTrue(torch.equal(y, torch.tensor([1.0]))) + self.assertTrue(torch.equal(z, torch.tensor([2.0]))) + y, z = method._method((x,), clone_outputs=False) + self.assertIs(y, x) + self.assertTrue(torch.equal(z, torch.tensor([6.0]))) + def test_mutated_planned_input_is_not_updated(self): x, cache, pos = _cache_inputs() # The update is planned onto the input's own spec, which is also an @@ -179,12 +207,12 @@ def test_failed_set_inputs_returns_installed_input(self): self.assertTrue(torch.equal(outputs[0], new_cache)) self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) - def test_caller_backed_outputs_are_not_cloned(self): - x, cache, pos = _cache_inputs() - method, _buffer = _load_cache_update_method() - method.set_output(torch.empty(1, 4), 1) + def test_bound_output_is_not_cloned(self): + x = torch.ones(2) + method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) + method.set_output(torch.empty(2), 0) with torch.profiler.profile() as profile: - method.execute((x, cache, pos)) + method.execute((x,)) clones = [e for e in profile.events() if e.name == "aten::clone"] self.assertEqual(clones, []) @@ -252,7 +280,7 @@ def test_set_output_rejects_mutated_input(self): method.set_output(torch.zeros(3, 4), 0) with self.assertRaises(RuntimeError): method.execute((x, cache, pos)) - outputs = method.execute((x, cache, pos)) + outputs = method._method((x, cache, pos), clone_outputs=False) self.assertIs(outputs[0], cache) self.assertTrue(torch.equal(cache[1], torch.ones(4))) From 23ecda0ba1b47f1b69d51d87ce3efe401db81a61 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 22:17:59 +0100 Subject: [PATCH 07/10] Match main's message for a replaced aliased input Co-Authored-By: Claude Opus 5.5 --- runtime/test/test_runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 2fff03f1c06..53fc52356f0 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -234,7 +234,7 @@ def test_aliased_output_rejects_replaced_input_storage(self): method._method.set_inputs((x,)) method._method.execute() x.data = torch.zeros_like(x) - with self.assertRaisesRegex(ValueError, "Input storage changed"): + with self.assertRaisesRegex(ValueError, "storage changed"): method._method.get_outputs(clone_outputs=False) def test_aliased_output_preserves_written_metadata(self): From 7b552a3c2bcc283cbcec224d34f5433362d920a5 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:15:53 +0100 Subject: [PATCH 08/10] Return an input that is also an output as main does 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 (#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 --- extension/pybindings/pybindings.cpp | 53 +++----------------- extension/pybindings/pybindings.pyi | 7 +-- extension/pybindings/test/test_pybindings.py | 9 ++-- runtime/test/test_runtime.py | 29 ++++------- 4 files changed, 27 insertions(+), 71 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 5836cfa33b5..a39a3d2e5b5 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -3406,7 +3406,6 @@ struct PyMethod final { input_tensors.reserve(inputs_size); #endif - std::vector fed_inputs(inputs_size); // Convert python objects into EValues. for (size_t i = 0; i < inputs_size; ++i) { auto python_input = inputs[i]; @@ -3429,7 +3428,6 @@ struct PyMethod final { at_tensor.scalar_type(), tensor_sizes); borrowed_inputs.emplace_back(i, python_input); - fed_inputs[i] = py::reinterpret_borrow(python_input); (void)mutable_tensor_data_ptr_no_cow(at_tensor); cpp_inputs.emplace_back(at_tensor); #else @@ -3453,8 +3451,6 @@ struct PyMethod final { // Watch the tensor the method reads: after a layout conversion that // is the copy, which an output aliasing the input can expose. borrowed_inputs.emplace_back(i, view->owner()); - // An output aliasing this input returns the tensor the method writes. - fed_inputs[i] = view->owner(); auto tensor = for_blob(view->data(), view->sizes(), view->scalar_type()) .strides(view->strides()) .dim_order(view->dim_order()) @@ -3592,7 +3588,6 @@ struct PyMethod final { torch_inputs_ = std::move(torch_inputs); torch_tensor_ptrs_ = std::move(input_tensors); #endif - fed_inputs_ = std::move(fed_inputs); find_inputs_aliased_by_outputs(); } @@ -3739,8 +3734,7 @@ struct PyMethod final { // An output that is an unplanned input (e.g. the write-back of a mutated // input) already points at the tensor fed for it. Giving it storage // would repoint the input too. - if (const auto input = aliased_input(i)) { - caller_outputs_[i] = fed_inputs_[*input]; + if (aliased_input(i)) { continue; } const auto& device_buffer = storage_->device_output_storages[i]; @@ -3798,47 +3792,20 @@ struct PyMethod final { // An output bound with set_output() already lives in the caller's tensor, // so it is returned as that tensor and taken out of `result` before it - // could be cloned. An output that is an input lives in the tensor passed - // in for it, which is returned only with clone_outputs=False: by default - // it is cloned like any other output. A dynamic bound output may have - // been resized smaller than its (contiguous) tensor, so view it at its - // written size. An input keeps the sizes it was fed with, and may be - // channels-last. + // could be cloned. A dynamic bound output may have been resized smaller + // than its (contiguous) tensor, so view it at its written size. An output + // that is an input is returned like any other: cloned by default, and as + // an alias of the caller's tensor with clone_outputs=False. std::vector caller_backed(result.size()); for (size_t i = 0; i < caller_outputs_.size() && i < result.size(); ++i) { - if (!caller_outputs_[i] || !result[i].isTensor() || - (clone_outputs && !is_bound_output(i))) { + if (!caller_outputs_[i] || !result[i].isTensor()) { continue; } -#ifndef USE_ATEN_LIB - if (!is_bound_output(i) && - portable_tensor_output() != PortableTensorOutput::Torch) { - continue; - } -#endif const auto& caller = caller_outputs_[i]; const auto& written = result[i].toTensor(); -#ifndef USE_ATEN_LIB - if (!is_bound_output(i)) { - auto alias = aliased_torch_output(written); - const std::vector sizes( - written.sizes().begin(), written.sizes().end()); - const std::vector strides( - written.strides().begin(), written.strides().end()); - if (caller.attr("shape").cast>() != sizes || - caller.attr("stride")().cast>() != strides) { - if (!alias.is_none()) { - caller_backed[i] = std::move(alias); - result[i] = EValue(); - } - continue; - } - } -#endif std::vector sizes( written.sizes().begin(), written.sizes().end()); - if (!is_bound_output(i) || - caller.attr("shape").cast>() == sizes) { + if (caller.attr("shape").cast>() == sizes) { caller_backed[i] = caller; } else { caller_backed[i] = caller.attr("view")(-1) @@ -3946,10 +3913,6 @@ struct PyMethod final { std::vector borrowed_inputs_; // Caller tensors bound as output storage with set_output(), by output index. std::vector bound_outputs_; - // The tensors passed to the last set_inputs(), by input index. They stay - // alive until the next set_inputs(), even if the caller drops them (e.g. a - // large KV cache). - std::vector fed_inputs_; // Per output of the last execute(), the caller tensor it lives in, if any. std::vector caller_outputs_; // Per output, the unplanned input it is, if any. Empty until the first @@ -4011,7 +3974,7 @@ struct PyMethod final { static std::string aliased_input_message(size_t output, size_t input) { return "output " + std::to_string(output) + " is input " + std::to_string(input) + - ", which is returned as the tensor passed in for it"; + ", which the method already writes in place"; } #ifndef USE_ATEN_LIB std::vector> dlpack_inputs_; diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 298ce177858..0cc6bd784e6 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -145,9 +145,10 @@ class ExecuTorchMethod: ``alloc_graph_input=False``) is used in place, so a mutation of it is written into the tensor passed in for it. When such an input is also an output (e.g. the write-back of a mutated input, or an input returned - unchanged), that output is a clone like any other by default, and is the - tensor passed in only with ``clone_outputs=False``. A memory planned input - is copied into the method, so the tensor passed in for it is not updated. + unchanged), that output is a clone like any other by default, and shares + the storage of the tensor passed in with ``clone_outputs=False``. A memory + planned input is copied into the method, so the tensor passed in for it is + not updated. .. warning:: diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 4325d2d7c2b..1b661995501 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/extension/pybindings/test/test_pybindings.py @@ -18,6 +18,7 @@ import numpy as np import torch + from executorch.exir import ExecutorchBackendConfig, to_edge from executorch.exir.backend.test.device_util import DeviceAwarePartitioner from executorch.exir.passes import MemoryPlanningPass @@ -1498,7 +1499,7 @@ def forward(self, x): gc.collect() self.assertIsNone(value_ref()) - def test_unplanned_input_output_preserves_caller_identity(self): + def test_borrowed_external_output_does_not_retain_previous_outputs(self): if self.kernel_mode != "portable": self.skipTest("only portable bindings retain Python tensor owners") @@ -1528,15 +1529,15 @@ def forward(self, x): strides = value.stride() offset = value.storage_offset() for _ in range(32): - previous = output + previous = weakref.ref(output) output = method(output, clone_outputs=False)[0] method.set_inputs(output) gc.collect() - self.assertIs(output, previous) + self.assertIsNone(previous()) self.assertEqual(output.data_ptr(), address) self.assertEqual(output.stride(), strides) self.assertEqual(output.storage_offset(), offset) - self.assertTrue(output.requires_grad) + self.assertFalse(output.requires_grad) del value, method torch.testing.assert_close(output, expected) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index 53fc52356f0..f701ad657cb 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -114,8 +114,8 @@ def test_smoke(self): def test_mutated_input_is_updated_in_place(self): # An unplanned input the method mutates is written where it lives. Its - # write-back output is a copy by default, and that very tensor only - # with clone_outputs=False. + # write-back output is a copy by default, and shares the caller's + # storage with clone_outputs=False. x, cache, _ = _cache_inputs() method, _buffer = _load_cache_update_method() expected = torch.zeros(3, 4) @@ -126,7 +126,7 @@ def test_mutated_input_is_updated_in_place(self): self.assertTrue(torch.equal(outputs[0], expected)) self.assertTrue(torch.equal(cache, expected)) outputs = method._method((x * 2, cache, torch.tensor([2])), clone_outputs=False) - self.assertIs(outputs[0], cache) + self.assertEqual(outputs[0].data_ptr(), cache.data_ptr()) expected[2] = 2.0 self.assertTrue(torch.equal(cache, expected)) self.assertTrue(torch.equal(outputs[1], expected.sum(0, keepdim=True))) @@ -145,7 +145,7 @@ def test_returned_input_is_cloned_by_default(self): self.assertTrue(torch.equal(y, torch.tensor([1.0]))) self.assertTrue(torch.equal(z, torch.tensor([2.0]))) y, z = method._method((x,), clone_outputs=False) - self.assertIs(y, x) + self.assertEqual(y.data_ptr(), x.data_ptr()) self.assertTrue(torch.equal(z, torch.tensor([6.0]))) def test_mutated_planned_input_is_not_updated(self): @@ -213,7 +213,7 @@ def test_failed_set_inputs_then_full_set_aliases_new_input(self): with self.assertRaises(RuntimeError): method._method.execute() outputs = method._method((x, new_cache, pos), clone_outputs=False) - self.assertIs(outputs[0], new_cache) + self.assertEqual(outputs[0].data_ptr(), new_cache.data_ptr()) self.assertTrue(torch.equal(new_cache[1], torch.ones(4))) self.assertTrue(torch.equal(cache, torch.zeros(3, 4))) @@ -254,15 +254,6 @@ def test_aliased_output_preserves_written_metadata(self): self.assertEqual(output.stride(), (3, 1)) self.assertEqual(output.data_ptr(), x.data_ptr()) - def test_bound_output_is_not_cloned(self): - x = torch.ones(2) - method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) - method.set_output(torch.empty(2), 0) - with torch.profiler.profile() as profile: - method.execute((x,)) - clones = [e for e in profile.events() if e.name == "aten::clone"] - self.assertEqual(clones, []) - def test_set_output_keeps_earlier_outputs_alive(self): x, cache, pos = _cache_inputs() method, _buffer = _load_cache_update_method() @@ -292,7 +283,6 @@ def test_set_output_rejects_bad_bindings(self): "non contiguous": (RuntimeError, torch.empty(1, 8)[:, ::2], 1), "too small": (RuntimeError, torch.empty(1, 2), 1), "wrong dtype": (RuntimeError, torch.empty(1, 4, dtype=torch.float64), 1), - "wrong device": (RuntimeError, torch.empty(1, 4, device="meta"), 1), } for name, (error, tensor, index) in bad_bindings.items(): with self.subTest(name), self.assertRaises(error): @@ -325,11 +315,12 @@ def test_set_output_rejects_mutated_input(self): # and raises, and later calls work. method, _buffer = _load_cache_update_method() method.set_output(torch.zeros(3, 4), 0) + fresh = torch.zeros(3, 4) with self.assertRaises(RuntimeError): - method.execute((x, cache, pos)) - outputs = method._method((x, cache, pos), clone_outputs=False) - self.assertIs(outputs[0], cache) - self.assertTrue(torch.equal(cache[1], torch.ones(4))) + method.execute((x, fresh, pos)) + outputs = method._method((x, fresh, pos), clone_outputs=False) + self.assertEqual(outputs[0].data_ptr(), fresh.data_ptr()) + self.assertTrue(torch.equal(fresh[1], torch.ones(4))) def test_execute_can_borrow_outputs(self): ep, inputs = create_program(ModuleAdd()) From b0b8f9cb71ef5547b39e3778654fac38b2f5f1e1 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Tue, 27 Oct 2026 23:44:09 +0100 Subject: [PATCH 09/10] set_output: refuse ATen builds and shared memory, raise ValueError - 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 --- extension/pybindings/pybindings.cpp | 92 +++++++++++++++++++---------- runtime/__init__.py | 6 +- runtime/test/test_runtime.py | 48 ++++++++++++--- 3 files changed, 104 insertions(+), 42 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index a39a3d2e5b5..22722a0148e 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -3598,6 +3598,12 @@ struct PyMethod final { // for it with clone_outputs=False stay valid. void set_output(py::object tensor, size_t index) { ObjectLock lock(mutex_); +#ifdef USE_ATEN_LIB + // ATen swaps the data pointer of the output's storage in place, so an + // output returned earlier would read the bound tensor, or freed memory once + // it is rebound. + throw std::runtime_error("set_output is not supported by the ATen bindings"); +#else if (!is_torch_tensor(tensor)) { throw py::type_error("Output storage must be a torch tensor"); } @@ -3609,15 +3615,15 @@ struct PyMethod final { } const auto& output = storage_->method->get_output(index); if (!output.isTensor()) { - throw std::runtime_error( + throw py::value_error( "output " + std::to_string(index) + " is not a tensor"); } if (const auto input = aliased_input(index)) { - throw std::runtime_error(aliased_input_message(index, *input)); + throw py::value_error(aliased_input_message(index, *input)); } for (size_t i = 0; i < num_outputs; ++i) { if (i != index && &storage_->method->get_output(i) == &output) { - throw std::runtime_error( + throw py::value_error( "output " + std::to_string(index) + " is the same value as output " + std::to_string(i) + " and cannot be bound"); @@ -3627,45 +3633,72 @@ struct PyMethod final { THROW_IF_ERROR( meta.error(), "Failed to get output tensor meta for output %zu", index); if (meta->is_memory_planned()) { - throw std::runtime_error( + throw py::value_error( "output " + std::to_string(index) + " is memory planned and cannot be bound"); } if (!tensor.attr("is_contiguous")().cast()) { - throw std::runtime_error( + throw py::value_error( "output " + std::to_string(index) + " must be bound to a contiguous tensor"); } -#ifdef USE_ATEN_LIB - const auto at_tensor = tensor.cast(); - const bool matches = at_tensor.scalar_type() == meta->scalar_type() && - at_tensor.device() == output.toTensor().device(); -#else const TorchTensorView view(tensor, index); - const bool matches = view.scalar_type() == meta->scalar_type() && - view.device() == output.toTensor().device(); -#endif - if (!matches) { - throw std::runtime_error( - "output " + std::to_string(index) + - " dtype and device must match the bound tensor"); + const auto device = output.toTensor().device(); + if (view.scalar_type() != meta->scalar_type() || view.device() != device) { + const auto name = [](executorch::aten::Device value) { + return value.is_cpu() ? std::string("cpu") + : "cuda:" + std::to_string(value.index()); + }; + throw py::value_error( + "output " + std::to_string(index) + " is " + + runtime::toString(meta->scalar_type()) + " on " + name(device) + + ", but the bound tensor is " + + runtime::toString(view.scalar_type()) + " on " + name(view.device())); } const auto dim_order = meta->dim_order(); for (size_t d = 0; d < dim_order.size(); ++d) { if (dim_order[d] != d) { - throw std::runtime_error( + throw py::value_error( "output " + std::to_string(index) + " is not contiguous (e.g. channels-last) and cannot be bound"); } } - if (tensor.attr("nbytes").cast() < meta->nbytes()) { - throw std::runtime_error( + if (view.nbytes() < meta->nbytes()) { + throw py::value_error( "output " + std::to_string(index) + " needs " + std::to_string(meta->nbytes()) + " bytes, but the tensor has " + - std::to_string(tensor.attr("nbytes").cast())); + std::to_string(view.nbytes())); + } + // Two outputs written into one buffer, or an output written over an input + // a kernel is still reading, would silently corrupt each other. + const auto overlaps = [&](const void* data, size_t nbytes) { + const auto start = reinterpret_cast(view.data()); + const auto other = reinterpret_cast(data); + return nbytes != 0 && view.nbytes() != 0 && start < other + nbytes && + other < start + view.nbytes(); + }; + for (size_t i = 0; i < bound_outputs_.size(); ++i) { + if (i == index || !bound_outputs_[i]) { + continue; + } + const TorchTensorView bound(bound_outputs_[i], i); + if (overlaps(bound.data(), bound.nbytes())) { + throw py::value_error( + "output " + std::to_string(index) + + " would share memory with output " + std::to_string(i) + + ", which is already bound"); + } + } + for (const auto& input : torch_inputs_) { + if (overlaps(input->data(), input->nbytes())) { + throw py::value_error( + "output " + std::to_string(index) + + " would share memory with input " + std::to_string(input->index())); + } } bound_outputs_.resize(num_outputs); bound_outputs_[index] = tensor; +#endif } void execute() { @@ -3714,23 +3747,18 @@ struct PyMethod final { std::vector> output_storage_spans(num_outputs); caller_outputs_.assign(num_outputs, py::object()); for (size_t i = 0; i < storage_->output_storages.size(); ++i) { +#ifndef USE_ATEN_LIB + // set_output() refuses to bind anything in ATen builds. if (is_bound_output(i)) { auto bound = bound_outputs_[i]; set_output(bound, i); -#ifdef USE_ATEN_LIB - const auto at_tensor = bound.cast(); - auto* data = - static_cast(mutable_tensor_data_ptr_no_cow(at_tensor)); - const auto nbytes = at_tensor.nbytes(); -#else const TorchTensorView view(bound, i); - auto* data = static_cast(view.data()); - const auto nbytes = view.nbytes(); -#endif - output_storage_spans[i] = Span(data, nbytes); + output_storage_spans[i] = + Span(static_cast(view.data()), view.nbytes()); caller_outputs_[i] = bound; continue; } +#endif // An output that is an unplanned input (e.g. the write-back of a mutated // input) already points at the tensor fed for it. Giving it storage // would repoint the input too. @@ -3967,7 +3995,7 @@ struct PyMethod final { } } if (!error.empty()) { - throw std::runtime_error(error); + throw py::value_error(error); } } diff --git a/runtime/__init__.py b/runtime/__init__.py index 80b3b438d62..0aa0d89655c 100644 --- a/runtime/__init__.py +++ b/runtime/__init__.py @@ -191,8 +191,10 @@ def set_output(self, tensor: "torch.Tensor", index: int) -> None: passed in for it. That last case is only known once inputs have been set: a binding made before the first call is removed by that call, which raises. ``tensor`` must have the output's dtype and device, be - contiguous, and hold at least as many bytes as the output's largest - shape. The output itself must be contiguous, not e.g. channels-last. + contiguous, hold at least as many bytes as the output's largest shape, + and not share memory with another bound output or with an input. The + output itself must be contiguous, not e.g. channels-last. A rejected + binding raises ``ValueError``. The ATen bindings don't support it. Args: tensor: The tensor that stores the output from now on. diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py index f701ad657cb..1fadea62692 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -50,6 +50,11 @@ def forward(self, x): return (x * 2).contiguous(memory_format=torch.channels_last) +class ModuleTwoOutputs(torch.nn.Module): + def forward(self, x): + return x * 2, x + 1 + + class ModuleDoubleTwice(torch.nn.Module): def forward(self, x): y = x * 2 @@ -280,9 +285,9 @@ def test_set_output_rejects_bad_bindings(self): method, _buffer = _load_cache_update_method() bad_bindings = { "out of range": (IndexError, torch.empty(1, 4), 2), - "non contiguous": (RuntimeError, torch.empty(1, 8)[:, ::2], 1), - "too small": (RuntimeError, torch.empty(1, 2), 1), - "wrong dtype": (RuntimeError, torch.empty(1, 4, dtype=torch.float64), 1), + "non contiguous": (ValueError, torch.empty(1, 8)[:, ::2], 1), + "too small": (ValueError, torch.empty(1, 2), 1), + "wrong dtype": (ValueError, torch.empty(1, 4, dtype=torch.float64), 1), } for name, (error, tensor, index) in bad_bindings.items(): with self.subTest(name), self.assertRaises(error): @@ -293,35 +298,62 @@ def test_set_output_rejects_outputs_it_cannot_own(self): method, _buffer = _load_unplanned_io_method( ModuleDouble(), (x,), plan_outputs=True ) - with self.subTest("memory planned"), self.assertRaises(RuntimeError): + with self.subTest("memory planned"), self.assertRaises(ValueError): method.set_output(torch.empty(2), 0) method, _buffer = _load_unplanned_io_method(ModuleDoubleTwice(), (x,)) - with self.subTest("returned twice"), self.assertRaises(RuntimeError): + with self.subTest("returned twice"), self.assertRaises(ValueError): method.set_output(torch.empty(2), 0) image = torch.ones(1, 2, 3, 4) method, _buffer = _load_unplanned_io_method( ModuleDoubleChannelsLast(), (image,) ) - with self.subTest("channels last"), self.assertRaises(RuntimeError): + with self.subTest("channels last"), self.assertRaises(ValueError): method.set_output(torch.empty(1, 2, 3, 4), 0) def test_set_output_rejects_mutated_input(self): x, cache, pos = _cache_inputs() method, _buffer = _load_cache_update_method() method.execute((x, cache, pos)) - with self.assertRaises(RuntimeError): + with self.assertRaises(ValueError): method.set_output(torch.zeros(3, 4), 0) # Bound before any inputs are set, the first call removes the binding # and raises, and later calls work. method, _buffer = _load_cache_update_method() method.set_output(torch.zeros(3, 4), 0) fresh = torch.zeros(3, 4) - with self.assertRaises(RuntimeError): + with self.assertRaises(ValueError): method.execute((x, fresh, pos)) outputs = method._method((x, fresh, pos), clone_outputs=False) self.assertEqual(outputs[0].data_ptr(), fresh.data_ptr()) self.assertTrue(torch.equal(fresh[1], torch.ones(4))) + def test_set_output_rejects_shared_memory(self): + x = torch.ones(2) + method, _buffer = _load_unplanned_io_method(ModuleTwoOutputs(), (x,)) + storage = torch.empty(4) + method.set_output(storage[:2], 0) + with self.subTest("another bound output"), self.assertRaises(ValueError): + method.set_output(storage[1:3], 1) + method.set_output(storage[2:], 1) + # Inputs are only known once set, so an output bound over an input is + # refused by the call that sets it. + method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) + method.set_output(x, 0) + with self.subTest("an input"), self.assertRaises(ValueError): + method.execute((x,)) + + def test_execute_without_inputs_raises_after_set_output(self): + # Aliases are found only after inputs are set, since reading an input + # marks it as set; binding first must not hide the missing inputs. + x = torch.ones(2) + for bind in (False, True): + with self.subTest(bind=bind): + method, _buffer = _load_unplanned_io_method(ModuleDouble(), (x,)) + if bind: + method.set_output(torch.empty(2), 0) + with self.assertRaisesRegex(RuntimeError, "no inputs yet"): + method._method.execute() + def test_execute_can_borrow_outputs(self): ep, inputs = create_program(ModuleAdd()) method = Runtime.get().load_program(ep.buffer).load_method("forward") From d8b663c1cff73da432a6a260aa1393599fb91670 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Wed, 28 Oct 2026 00:24:36 +0100 Subject: [PATCH 10/10] Apply clang-format Co-Authored-By: Claude Opus 5.5 --- extension/pybindings/pybindings.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 22722a0148e..39ae5082bef 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -3602,7 +3602,8 @@ struct PyMethod final { // ATen swaps the data pointer of the output's storage in place, so an // output returned earlier would read the bound tensor, or freed memory once // it is rebound. - throw std::runtime_error("set_output is not supported by the ATen bindings"); + throw std::runtime_error( + "set_output is not supported by the ATen bindings"); #else if (!is_torch_tensor(tensor)) { throw py::type_error("Output storage must be a torch tensor"); @@ -3652,8 +3653,8 @@ struct PyMethod final { throw py::value_error( "output " + std::to_string(index) + " is " + runtime::toString(meta->scalar_type()) + " on " + name(device) + - ", but the bound tensor is " + - runtime::toString(view.scalar_type()) + " on " + name(view.device())); + ", but the bound tensor is " + runtime::toString(view.scalar_type()) + + " on " + name(view.device())); } const auto dim_order = meta->dim_order(); for (size_t d = 0; d < dim_order.size(); ++d) { @@ -4001,8 +4002,7 @@ struct PyMethod final { static std::string aliased_input_message(size_t output, size_t input) { return "output " + std::to_string(output) + " is input " + - std::to_string(input) + - ", which the method already writes in place"; + std::to_string(input) + ", which the method already writes in place"; } #ifndef USE_ATEN_LIB std::vector> dlpack_inputs_;