diff --git a/docs/source/runtime-python-api-reference.rst b/docs/source/runtime-python-api-reference.rst index 9345f7af53c..6473e5536bf 100644 --- a/docs/source/runtime-python-api-reference.rst +++ b/docs/source/runtime-python-api-reference.rst @@ -18,7 +18,7 @@ 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 .. autoclass:: ExecuTorchResult :members: shape, strides, dtype, device, nbytes, to_host, from_host diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 25fb2abe576..39ae5082bef 100644 --- a/extension/pybindings/pybindings.cpp +++ b/extension/pybindings/pybindings.cpp @@ -3587,6 +3587,118 @@ struct PyMethod final { dlpack_tensor_ptrs_ = std::move(dlpack_tensor_ptrs); torch_inputs_ = std::move(torch_inputs); torch_tensor_ptrs_ = std::move(input_tensors); +#endif + 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(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"); + } + const auto num_outputs = storage_->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 = storage_->method->get_output(index); + if (!output.isTensor()) { + throw py::value_error( + "output " + std::to_string(index) + " is not a tensor"); + } + if (const auto input = aliased_input(index)) { + 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 py::value_error( + "output " + std::to_string(index) + + " is the same value as output " + std::to_string(i) + + " and cannot be bound"); + } + } + const auto meta = storage_->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 py::value_error( + "output " + std::to_string(index) + + " is memory planned and cannot be bound"); + } + if (!tensor.attr("is_contiguous")().cast()) { + throw py::value_error( + "output " + std::to_string(index) + + " must be bound to a contiguous tensor"); + } + const TorchTensorView view(tensor, index); + 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 py::value_error( + "output " + std::to_string(index) + + " is not contiguous (e.g. channels-last) and cannot be bound"); + } + } + 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(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 } @@ -3634,7 +3746,26 @@ struct PyMethod final { const auto num_outputs = storage_->method->outputs_size(); allocate_output_storages(); std::vector> output_storage_spans(num_outputs); - for (int i = 0; i < storage_->output_storages.size(); ++i) { + 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); + const TorchTensorView view(bound, i); + 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. + if (aliased_input(i)) { + continue; + } const auto& device_buffer = storage_->device_output_storages[i]; output_storage_spans[i] = device_buffer ? device_buffer->as_span() @@ -3688,8 +3819,39 @@ struct PyMethod final { storage_->method->method_meta().name(), static_cast(get_outputs_status)); + // 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. 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()) { + continue; + } + const auto& caller = caller_outputs_[i]; + const auto& written = result[i].toTensor(); + std::vector sizes( + written.sizes().begin(), written.sizes().end()); + if (caller.attr("shape").cast>() == sizes) { + caller_backed[i] = caller; + } else { + caller_backed[i] = caller.attr("view")(-1) + .attr("narrow")(0, 0, written.numel()) + .attr("view")(py::cast(sizes)); + } + result[i] = EValue(); + } + // Retrieve outputs - return get_outputs_as_py_list(result, clone_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; } py::list call(const py::object& inputs, bool clone_outputs = true) { @@ -3778,6 +3940,70 @@ struct PyMethod final { std::vector> buffer_inputs_; std::vector buffer_tensor_ptrs_; std::vector borrowed_inputs_; + // Caller tensors bound as output storage with set_output(), by output index. + std::vector bound_outputs_; + // 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() && bool(bound_outputs_[index]); + } + + 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 = storage_->method->outputs_size(); + if (inputs_aliased_by_outputs_.size() == num_outputs) { + return; + } + inputs_aliased_by_outputs_.assign(num_outputs, std::nullopt); + const auto meta = storage_->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 < storage_->method->inputs_size(); ++j) { +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" + const bool same = + &storage_->method->get_output(i) == &storage_->method->get_input(j); +#pragma GCC diagnostic pop + if (same) { + inputs_aliased_by_outputs_[i] = j; + break; + } + } + } + // 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] = py::object(); + error += + aliased_input_message(i, *input) + ", so its binding was removed. "; + } + } + if (!error.empty()) { + throw py::value_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 the method already writes in place"; + } #ifndef USE_ATEN_LIB std::vector> dlpack_inputs_; std::vector dlpack_tensor_ptrs_; @@ -3997,17 +4223,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. - output_storages.emplace_back(); - continue; - } - const auto& output = storage_->method->get_output(i); - bool is_input = false; - for (size_t j = 0; j < storage_->method->inputs_size(); ++j) { - is_input |= &output == &storage_->method->get_input(j); - } - if (is_input) { + 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; } @@ -4776,6 +4995,12 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) { py::class_(m, "ExecuTorchMethod", py::module_local()) .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 3810cc45b8f..0cc6bd784e6 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -10,6 +10,8 @@ from __future__ import annotations from enum import Enum from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +import torch + _T = TypeVar("_T") def experimental(message: str) -> Callable[[_T], _T]: ... @@ -135,6 +137,19 @@ class ExecuTorchProgram: class ExecuTorchMethod: """ExecuTorchMethod is a Python wrapper around a loaded C++ method. + 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 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:: This API is experimental and subject to change without notice. @@ -143,6 +158,7 @@ class ExecuTorchMethod: # pyre-ignore[2]: "Any" in parameter type annotations. def set_inputs(self, inputs: 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 a4ebf9a5e72..0aa0d89655c 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] # noqa: F401 @@ -141,6 +154,15 @@ def execute( ) -> Sequence[Any]: """Executes the method with the given inputs. + 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. Without PyTorch, pass NumPy arrays, other buffers such as @@ -155,6 +177,31 @@ def execute( """ return self._method(inputs, clone_outputs=clone_outputs) + 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 + 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``), must not be returned at another output + 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, 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. + 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 1b1dcf3cd48..1fadea62692 100644 --- a/runtime/test/test_runtime.py +++ b/runtime/test/test_runtime.py @@ -12,7 +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, ModuleAdd, @@ -21,6 +22,70 @@ 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 ModuleInputAndIncrement(torch.nn.Module): + def forward(self, x): + return x, x + 1 + + +class ModuleDouble(torch.nn.Module): + 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 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 + 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, 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=plan_outputs + ) + ) + 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. @@ -52,6 +117,243 @@ 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. Its + # 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) + 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.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))) + + 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.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): + 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, 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))) + 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_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_then_full_set_aliases_new_input(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: the + # method refuses to run until a full set is installed, and the mutated + # input it then writes is the new one. + new_cache = torch.zeros(3, 4) + with self.assertRaises(RuntimeError): + method._method.set_inputs((x, new_cache, 1)) + with self.assertRaises(RuntimeError): + method._method.execute() + outputs = method._method((x, new_cache, pos), clone_outputs=False) + 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))) + + def test_failed_input_prevalidation_preserves_previous_inputs(self): + x, cache, pos = _cache_inputs() + method, _buffer = _load_cache_update_method() + method._method.set_inputs((x, cache, pos)) + new_cache = torch.zeros_like(cache) + with self.assertRaises(ValueError): + method._method.set_inputs((x, new_cache, torch.tensor([1.0]))) + method._method.execute() + self.assertTrue(torch.equal(cache[1], x[0])) + self.assertTrue(torch.equal(new_cache, torch.zeros_like(cache))) + + def test_aliased_output_rejects_replaced_input_storage(self): + x = torch.ones(2, 2) + method, _buffer = _load_unplanned_io_method(ModuleInputAndIncrement(), (x,)) + method._method.set_inputs((x,)) + method._method.execute() + x.data = torch.zeros_like(x) + with self.assertRaisesRegex(ValueError, "storage changed"): + method._method.get_outputs(clone_outputs=False) + + def test_aliased_output_preserves_written_metadata(self): + for parameter in (False, True): + with self.subTest(parameter=parameter): + x = torch.ones(2, 3) + if parameter: + x = torch.nn.Parameter(x, requires_grad=False) + method, _buffer = _load_unplanned_io_method( + ModuleInputAndIncrement(), (x,) + ) + method._method.set_inputs((x,)) + method._method.execute() + x.transpose_(0, 1) + output = method._method.get_outputs(clone_outputs=False)[0] + self.assertEqual(output.shape, (2, 3)) + self.assertEqual(output.stride(), (3, 1)) + self.assertEqual(output.data_ptr(), x.data_ptr()) + + 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": (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): + 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(ValueError): + method.set_output(torch.empty(2), 0) + method, _buffer = _load_unplanned_io_method(ModuleDoubleTwice(), (x,)) + 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(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(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(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")