From 48ecc173b1637e6de7e190f03d9ffdaf217456dd Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 08:59:14 -0700 Subject: [PATCH] Let a backend share one delegate across calls of the same lowered module (#24127) ## What is wrong today A method can call the same lowered module several times, for example a denoising step that is lowered once and called 10 times in a loop. The emitter then writes one `BackendDelegate` per call. So at load time the runtime calls the backend's `init()` once per call and keeps every handle alive, even though all of them come from the same compiled blob. This wastes load time and memory. ## What this change does A backend can now return `PreprocessResult(..., share_handle_across_calls=True)`. Then every call of that lowered module in a method uses one `BackendDelegate`, so `init()` runs once. Each call still passes its own arguments. - It is off by default, so existing backends behave as before. Some backends keep argument pointers on the handle, so sharing is not safe for all of them. - A delegate that declares scratch memory is never shared, because scratch is planned per call and the runtime keeps it on the delegate. `Method::init` now rejects a program where such a delegate has more than one call site. It also checks every delegate index at load, not only at run time. - Only calls of one lowered module share a handle. Separate partitions keep separate delegates, even when their blobs are identical. - The flag is saved and loaded like `scratch_specs`, including through `exir.save` and `exir.load`, so the serde schema moves to version 5.6. The `.pte` file format does not change. - The demo test backend opts in. Before #14564, the emitter shared one delegate between calls with the same blob for every backend. Now it shares by lowered module, and only for backends that ask for it. ## What was tested - New tests cover one delegate for repeated calls with the same output as separate calls, the partitioner path, separate modules keeping their own delegates, save and load, the default, and scratch, including a hand-made program that the runtime rejects. The runtime test for scratch fails without the new runtime check. The load-time index check, calls inside control flow, the multi-method lowering path and the number of `init()` calls have no dedicated test. - On an NVIDIA Jetson, a TensorRT step that a method calls 10 times went from 10 delegates to 1 shared delegate: load time 435 ms to 74 ms, memory 1060 MiB to 284 MiB, with the same outputs and speed. (cherry picked from commit cef2de7db2ccc1d745fd1ce1eb57f20c31655ff7) --- .../compiler-delegate-and-partitioner.md | 30 +++ exir/backend/backend_api.py | 2 + exir/backend/backend_details.py | 10 + exir/backend/test/BUCK | 3 + .../test/backend_with_compiler_demo.py | 2 + exir/backend/test/test_delegate_scratch.py | 46 +++- .../test/test_lowered_backend_module.py | 209 +++++++++++++++++- exir/emit/_emitter.py | 34 ++- exir/lowered_backend_module.py | 11 + exir/serde/schema.py | 3 +- exir/serde/serialize.py | 2 + pytest-windows.ini | 2 + runtime/backend/interface.h | 10 +- runtime/executor/method.cpp | 45 ++-- 14 files changed, 381 insertions(+), 28 deletions(-) diff --git a/docs/source/compiler-delegate-and-partitioner.md b/docs/source/compiler-delegate-and-partitioner.md index 5dec025fb4c..2d8ef530971 100644 --- a/docs/source/compiler-delegate-and-partitioner.md +++ b/docs/source/compiler-delegate-and-partitioner.md @@ -132,6 +132,36 @@ number in the program, lets the planner reuse the bytes for tensors whose lifetimes do not overlap the delegate call, and makes the requirement visible to the memory planning inspection tools. +### Sharing one handle across calls + +A method can call the same lowered module more than once. For example, you lower +a step once with `to_backend`, call it in a loop from a wrapper module, and +export the wrapper again. By default each call gets its own handle (the object +that `init()` returns), so `init()` runs once per call. A backend can ask for +one handle for all of these calls: + +```python +return PreprocessResult( + processed_bytes=blob, + share_handle_across_calls=True, +) +``` + +Then `init()` runs once, and every call runs `execute()` on that handle with its +own arguments. Set it only if `execute()` does all of the following, in every +runtime version of the backend and with any runtime options: + +- It reads its arguments again on every call. +- It keeps nothing on the handle that is tied to one call or that changes later + results, for example saved argument pointers or a cache that each call + updates. +- It writes outputs only into the output tensors it is given. + +Sharing applies to calls of one lowered module. Separate partitions get separate +lowered modules, so they keep separate handles even when their blobs are +identical. A backend that declares scratch always gets one handle per call, +because the runtime keeps scratch on the delegate and each call has its own. + ## Backend Interfaces: Runtime Initialization and Execution During the runtime, the compiled blob from the `preprocess` function will be diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index a3a77e25373..95cebb75ed9 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -125,6 +125,7 @@ def to_backend( compile_specs=compile_specs, named_data_store_output=preprocess_result.data_store_output, scratch_specs=preprocess_result.scratch_specs, + share_handle_across_calls=preprocess_result.share_handle_across_calls, ) lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map @@ -620,6 +621,7 @@ def lower_all_submodules_to_backend( compile_specs=compile_spec, named_data_store_output=preprocess_result.data_store_output, scratch_specs=preprocess_result.scratch_specs, + share_handle_across_calls=preprocess_result.share_handle_across_calls, ) lowered_module.meta = { "debug_handle_map": preprocess_result.debug_handle_map, diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index 838d6fd6eea..ca7d3b6e66a 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -74,6 +74,16 @@ class PreprocessResult: # in this order. They are not arguments of the delegate call. scratch_specs: Optional[List[DelegateScratchSpec]] = None + # Whether every call of this lowered module in a method may run on one + # delegate handle, so the runtime initializes the backend once instead of + # once per call. Set it only if execute(), in every runtime version of the + # backend, reads its arguments again on every call, keeps nothing on the + # handle that is tied to one call or that changes later results, and writes + # outputs only into the output tensors it is given. Ignored when + # scratch_specs is not empty, because the runtime keeps scratch on the + # delegate and each call has its own. + share_handle_across_calls: bool = False + """ How to create a backend (for example, BackendWithCompilerDemo): diff --git a/exir/backend/test/BUCK b/exir/backend/test/BUCK index 4b0574edbc6..99913f65989 100644 --- a/exir/backend/test/BUCK +++ b/exir/backend/test/BUCK @@ -294,9 +294,12 @@ fbcode_target(_kind = runtime.python_test, "fbsource//third-party/pypi/hypothesis:hypothesis", ":backend_with_compiler_demo", ":demo_backend", + ":op_partitioner_demo", "//caffe2:torch", "//executorch/exir:lib", "//executorch/exir:schema", + "//executorch/exir:tensor", + "//executorch/exir/_serialize:lib", "//executorch/exir/backend:backend_api", "//executorch/exir/backend:compile_spec_schema", "//executorch/exir/tests:models", diff --git a/exir/backend/test/backend_with_compiler_demo.py b/exir/backend/test/backend_with_compiler_demo.py index aa60fa2154c..86e5e8d78f4 100644 --- a/exir/backend/test/backend_with_compiler_demo.py +++ b/exir/backend/test/backend_with_compiler_demo.py @@ -139,4 +139,6 @@ def preprocess( ), debug_handle_map=debug_handle_map, _delegate_info_meta="test", + # The runtime keeps no state and reads its arguments on every call. + share_handle_across_calls=True, ) diff --git a/exir/backend/test/test_delegate_scratch.py b/exir/backend/test/test_delegate_scratch.py index 68740982b29..755ef839240 100644 --- a/exir/backend/test/test_delegate_scratch.py +++ b/exir/backend/test/test_delegate_scratch.py @@ -97,6 +97,21 @@ def preprocess( assert TWO_BUFFER_SIZES[1] % ALIGNMENT != 0 +@final +class SharedHandleScratchTestBackend(BackendDetails): + """Asks for one handle across calls and for scratch, which cannot be shared.""" + + @staticmethod + def preprocess( + edge_program: ExportedProgram, compile_specs: List[CompileSpec] + ) -> PreprocessResult: + return PreprocessResult( + processed_bytes=b"shared-handle-scratch-backend-blob", + scratch_specs=[DelegateScratchSpec(nbytes=TWO_BUFFER_SIZES[0])], + share_handle_across_calls=True, + ) + + @final class TwoBufferScratchTestBackend(BackendDetails): """Declares two buffers, to pin down ordering and independent sizing.""" @@ -584,14 +599,14 @@ def test_scratch_is_inserted_when_view_copies_are_kept(self): (delegate_call,) = _delegate_calls(plan) self.assertEqual(_scratch_sizes(delegate_call), [12 * BYTES_PER_ELEMENT]) - def _lowered_single_delegate(self): + def _lowered_single_delegate(self, backend_id="DelegateScratchTestBackend"): exported = torch.export.export( SingleDelegateModule().eval(), (torch.randn(3, 4), torch.randn(3, 4)) ) edge = to_edge_transform_and_lower( exported, compile_config=EdgeCompileConfig(_check_ir_validity=False) ) - return to_backend("DelegateScratchTestBackend", edge.exported_program(), []) + return to_backend(backend_id, edge.exported_program(), []) def test_lowered_module_emits_its_own_program_with_scratch(self): # program() and buffer() run their own small pipeline rather than @@ -602,6 +617,33 @@ def test_lowered_module_emits_its_own_program_with_scratch(self): self.assertEqual(_scratch_sizes(delegate_call), [12 * BYTES_PER_ELEMENT]) self.assertTrue(lowered.buffer()) + def test_a_delegate_with_scratch_is_not_shared_across_calls(self): + # The runtime resolves scratch onto delegates_[delegate_index], so each + # call site that declares scratch needs its own entry even when the + # backend allows one handle across calls. + lowered = self._lowered_single_delegate("SharedHandleScratchTestBackend") + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.lowered = lowered + + def forward(self, x, y): + return self.lowered(self.lowered(x, y), y) + + inputs = (torch.randn(3, 4), torch.randn(3, 4)) + plan = ( + exir.to_edge(torch.export.export(TwoCalls(), inputs)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + + first, second = _delegate_calls(plan) + self.assertEqual(len(plan.delegates), 2) + self.assertNotEqual(first.delegate_index, second.delegate_index) + self.assertEqual(_scratch_sizes(first), [TWO_BUFFER_SIZES[0]]) + self.assertEqual(_scratch_sizes(second), [TWO_BUFFER_SIZES[0]]) + def test_emitting_without_the_pass_is_an_error_not_a_silent_drop(self): # Any emit path that skips the pass must fail. The count comes from the # graph, not the declaration, so otherwise the program ships without diff --git a/exir/backend/test/test_lowered_backend_module.py b/exir/backend/test/test_lowered_backend_module.py index 57575f8a797..b2f2e1cc233 100644 --- a/exir/backend/test/test_lowered_backend_module.py +++ b/exir/backend/test/test_lowered_backend_module.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import copy +import io import operator import unittest from unittest.mock import patch @@ -12,14 +14,17 @@ import torch from executorch import exir -from executorch.exir import to_edge +from executorch.exir import EdgeProgramManager, to_edge, to_edge_transform_and_lower +from executorch.exir._serialize import _PTEFile, _serialize_pte_binary from executorch.exir.backend.backend_api import to_backend from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.exir.backend.test.backend_with_compiler_demo import ( BackendWithCompilerDemo, ) from executorch.exir.backend.test.demo_backend import DemoBackend -from executorch.exir.schema import DelegateCall, Program +from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo +from executorch.exir.schema import DelegateCall, DelegateScratch, Program +from executorch.exir.tensor import make_allocation_info from executorch.extension.pybindings.portable_lib import ( # @manual _load_for_executorch_from_buffer, @@ -107,6 +112,206 @@ def forward(self, x): torch.allclose(model_outputs[0], expected_res, atol=1e-03, rtol=1e-03) ) + def _repeated_sin_calls(self, backend_id, compile_specs, num_calls): + """Lowers sin once and returns it with a module that calls it num_calls times.""" + + class SinModule(torch.nn.Module): + def forward(self, x): + return torch.sin(x) + + model_inputs = (torch.ones(1),) + lowered_sin_module = to_backend( + backend_id, + to_edge(export(SinModule(), model_inputs, strict=True)).exported_program(), + compile_specs, + ) + + class RepeatedCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.sin = lowered_sin_module + + def forward(self, x): + for _ in range(num_calls): + x = self.sin(x) + return x + + edge = to_edge(export(RepeatedCalls(), model_inputs, strict=True)) + return lowered_sin_module, edge, model_inputs + + def _delegate_indices(self, plan): + return [ + instruction.instr_args.delegate_index + for instruction in plan.chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + + def test_calls_of_one_lowered_module_share_one_delegate(self): + lowered_sin_module, edge, _ = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=3, + ) + self.assertTrue(copy.deepcopy(lowered_sin_module).share_handle_across_calls) + + plan = edge.to_executorch().executorch_program.execution_plan[0] + self.assertEqual(len(plan.delegates), 1) + self.assertEqual(self._delegate_indices(plan), [0, 0, 0]) + args = [ + tuple(instruction.instr_args.args) + for instruction in plan.chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + self.assertEqual(len(set(args)), 3) + + def test_shared_delegate_matches_separate_calls_at_runtime(self): + lowered_sin_module, edge, model_inputs = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=3, + ) + program = edge.to_executorch() + self.assertEqual(len(program.executorch_program.execution_plan[0].delegates), 1) + + expected = model_inputs[0] + single_call = _load_for_executorch_from_buffer(lowered_sin_module.buffer()) + for _ in range(3): + expected = single_call.forward([expected])[0] + outputs = _load_for_executorch_from_buffer(program.buffer).forward( + list(model_inputs) + ) + self.assertTrue(torch.equal(outputs[0], expected)) + + def test_sharing_survives_a_serde_round_trip(self): + _, edge, _ = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=2, + ) + buffer = io.BytesIO() + exir.save(edge.exported_program(), buffer) + buffer.seek(0) + + plan = ( + EdgeProgramManager(exir.load(buffer)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + self.assertEqual(self._delegate_indices(plan), [0, 0]) + + def test_runtime_rejects_a_shared_delegate_with_scratch(self): + _, edge, model_inputs = self._repeated_sin_calls( + BackendWithCompilerDemo.__name__, + [CompileSpec("max_value", bytes([1]))], + num_calls=2, + ) + + def delegate_calls(program): + return [ + instruction.instr_args + for instruction in program.execution_plan[0].chains[0].instructions + if isinstance(instruction.instr_args, DelegateCall) + ] + + program = edge.to_executorch().executorch_program + scratch = [DelegateScratch(allocation=make_allocation_info(1, 0), size=4)] + + def run(program): + return _load_for_executorch_from_buffer( + bytes(_serialize_pte_binary(pte_file=_PTEFile(program=program))) + ).forward(list(model_inputs)) + + # The same scratch on two separate delegates is valid, so the errors + # below come from the sharing alone. + separate = copy.deepcopy(program) + delegates = separate.execution_plan[0].delegates + delegates.append(copy.deepcopy(delegates[0])) + delegate_calls(separate)[1].delegate_index = 1 + for call in delegate_calls(separate): + call.scratch = scratch + run(separate) + + # The emitter never writes these, but another producer could. + for calls_with_scratch in ([0, 1], [0], [1]): + shared = copy.deepcopy(program) + for index in calls_with_scratch: + delegate_calls(shared)[index].scratch = scratch + with self.assertRaisesRegex(RuntimeError, "error: 0x23"): + run(shared) + + def test_partitioned_module_shares_one_delegate_across_calls(self): + class AddModule(torch.nn.Module): + def forward(self, x): + return x + x + + model_inputs = (torch.ones(2, 2),) + add = ( + to_edge_transform_and_lower( + export(AddModule(), model_inputs, strict=True), + partitioner=[AddMulPartitionerDemo()], + ) + .exported_program() + .module() + ) + + class TwoCalls(torch.nn.Module): + def __init__(self): + super().__init__() + self.add = add + + def forward(self, x): + return self.add(self.add(x)) + + plan = ( + to_edge(export(TwoCalls(), model_inputs, strict=False)) + .to_executorch() + .executorch_program.execution_plan[0] + ) + self.assertEqual(self._delegate_indices(plan), [0, 0]) + + def test_each_call_uses_the_delegate_of_its_own_module(self): + class SinModule(torch.nn.Module): + def forward(self, x): + return torch.sin(x) + + model_inputs = (torch.ones(1),) + edge_sin = to_edge(export(SinModule(), model_inputs, strict=True)) + compile_specs = [CompileSpec("max_value", bytes([1]))] + first, second = ( + to_backend( + BackendWithCompilerDemo.__name__, + edge_sin.exported_program(), + compile_specs, + ) + for _ in range(2) + ) + + class TwoModules(torch.nn.Module): + def __init__(self): + super().__init__() + self.first = first + self.second = second + + def forward(self, x): + return self.first(self.second(self.first(x))) + + program = ( + to_edge(export(TwoModules(), model_inputs, strict=True)) + .to_executorch() + .executorch_program + ) + # Identical blobs are stored once, but separate modules keep separate + # delegates, and a repeated call goes back to its own module's delegate. + self.assertEqual(len(program.backend_delegate_data), 1) + self.assertEqual(self._delegate_indices(program.execution_plan[0]), [0, 1, 0]) + + def test_calls_keep_their_own_delegate_unless_the_backend_allows_sharing(self): + _, edge, _ = self._repeated_sin_calls(DemoBackend.__name__, [], num_calls=2) + + plan = edge.to_executorch().executorch_program.execution_plan[0] + self.assertEqual(len(plan.delegates), 2) + self.assertEqual(self._delegate_indices(plan), [0, 1]) + def test_emit_lowered_backend_module(self): module_list = [ models.Emformer(), diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index d8db9145232..efdfaab00cd 100644 --- a/exir/emit/_emitter.py +++ b/exir/emit/_emitter.py @@ -153,6 +153,10 @@ class _EmitterState: spec2id_dict: Dict[TensorSpec, int] = field(default_factory=dict) # Only literal arguments with no schema aliasing enter this per-plan pool. constant_cache: Dict[Union[int, Tuple[int, ...]], int] = field(default_factory=dict) + # pyre-ignore[11]: Annotation `LoweredBackendModule` is not defined as a type. + delegate_cache: Dict["LoweredBackendModule", int] = field( # noqa: F821 + default_factory=dict + ) def spec2id(self, spec: TensorSpec) -> int: """Map a TensorSpec to value index in the values array.""" @@ -1499,14 +1503,26 @@ def _emit_delegate( ) ) - backend_delegate = BackendDelegate( - id=lowered_module.backend_id, - processed=BackendDelegateDataReference( - location=DataLocation.INLINE, index=delegate_index - ), - compile_specs=lowered_module.compile_specs, - ) - self.emitter_state.delegates.append(backend_delegate) + backend_delegate_index = self.emitter_state.delegate_cache.get(lowered_module) + if backend_delegate_index is None: + backend_delegate = BackendDelegate( + id=lowered_module.backend_id, + processed=BackendDelegateDataReference( + location=DataLocation.INLINE, index=delegate_index + ), + compile_specs=lowered_module.compile_specs, + ) + self.emitter_state.delegates.append(backend_delegate) + backend_delegate_index = len(self.emitter_state.delegates) - 1 + # Scratch is planned per call site but the runtime keeps it on the + # delegate, so a delegate with scratch cannot be shared. + if ( + lowered_module.share_handle_across_calls + and not lowered_module.scratch_specs + ): + self.emitter_state.delegate_cache[lowered_module] = ( + backend_delegate_index + ) # TODO(angelayi) Will need to emit the kwargs too, in the correct order according to the # function's spec and with default arguments. This requires us to store the function's spec # in to_backend() @@ -1521,7 +1537,7 @@ def _emit_delegate( self.chain.instructions.append( Instruction( DelegateCall( - delegate_index=len(self.emitter_state.delegates) - 1, + delegate_index=backend_delegate_index, args=delegate_args, scratch=self._emit_delegate_scratch(scratch_specs) or None, ) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index dd3b2d8a6ca..8a381c2b0f0 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -71,6 +71,7 @@ class LoweredBackendModule(torch.nn.Module): _scratch_specs: List[ DelegateScratchSpec ] # Scratch buffers the backend needs at execute() time + _share_handle_across_calls: bool # Whether all calls may share one handle meta: Optional[Dict[str, Any]] # Metadata for the lowered module def __init__( @@ -81,6 +82,7 @@ def __init__( compile_specs: List[CompileSpec], named_data_store_output: Optional[NamedDataStoreOutput] = None, scratch_specs: Optional[List[DelegateScratchSpec]] = None, + share_handle_across_calls: bool = False, ) -> None: super().__init__() self._original_exported_program = edge_program @@ -89,6 +91,7 @@ def __init__( self._compile_specs = compile_specs self._named_data_store_output = named_data_store_output self._scratch_specs = scratch_specs or [] + self._share_handle_across_calls = share_handle_across_calls self.meta = None # pyre-ignore @@ -118,6 +121,7 @@ def __deepcopy__(self, memo: Optional[Dict[int, Any]]) -> "LoweredBackendModule" compile_specs=copy.deepcopy(self._compile_specs, memo), named_data_store_output=self._named_data_store_output, scratch_specs=copy.deepcopy(self._scratch_specs, memo), + share_handle_across_calls=self._share_handle_across_calls, ) res.meta = copy.copy(getattr(self, "meta", {})) return res @@ -164,6 +168,13 @@ def scratch_specs(self) -> List[DelegateScratchSpec]: """ return self._scratch_specs + @property + def share_handle_across_calls(self) -> bool: + """ + Returns whether every call of this module in a method may share one delegate handle + """ + return self._share_handle_across_calls + # TODO(chenlai): consolidate the seriailization config with serialize_to_flatbuffer api def buffer( self, diff --git a/exir/serde/schema.py b/exir/serde/schema.py index f2347af7367..ff4cf26efb2 100644 --- a/exir/serde/schema.py +++ b/exir/serde/schema.py @@ -17,7 +17,7 @@ from executorch.exir.serde.union import _Union # NOTE: Please update this value if any modifications are made to the schema -SCHEMA_VERSION = (5, 5) +SCHEMA_VERSION = (5, 6) TREESPEC_VERSION = 1 @@ -415,3 +415,4 @@ class LoweredBackendModule: named_data_store: Optional[bytes] = None # The delegate scratch buffers, in declaration order. scratch_specs: Optional[List[DelegateScratchSpec]] = None + share_handle_across_calls: bool = False diff --git a/exir/serde/serialize.py b/exir/serde/serialize.py index d9cd9a700a7..677ac072ee0 100644 --- a/exir/serde/serialize.py +++ b/exir/serde/serialize.py @@ -314,6 +314,7 @@ def serialize_bytes(b: bytes) -> str: SerdeDelegateScratchSpec(nbytes=spec.nbytes, label=spec.label) for spec in lowered_module.scratch_specs ], + share_handle_across_calls=lowered_module.share_handle_across_calls, ) json_lowered_module = json.dumps( @@ -636,6 +637,7 @@ def deserialize_lowered_module( DelegateScratchSpec(nbytes=spec.nbytes, label=spec.label) for spec in (serialized_lowered_module.scratch_specs or []) ], + serialized_lowered_module.share_handle_across_calls, ) self.module.register_module(serialized_lowered_module_arg.name, lowered_module) return self.graph.get_attr(serialized_lowered_module_arg.name) diff --git a/pytest-windows.ini b/pytest-windows.ini index 39f238cdc00..26a0ee6dabc 100644 --- a/pytest-windows.ini +++ b/pytest-windows.ini @@ -103,6 +103,8 @@ addopts = --deselect=exir/backend/test/test_compatibility.py::TestCompatibility::test_compatibility_in_runtime --deselect=exir/backend/test/test_compatibility.py::TestCompatibility::test_compatibility_in_runtime_edge_program_manager --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_lowered_backend_module_end_to_end + --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_shared_delegate_matches_separate_calls_at_runtime + --deselect=exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_runtime_rejects_a_shared_delegate_with_scratch --deselect=exir/backend/test/test_to_backend_multi_method.py::TestToBackendMultiMethod::test_multi_method_end_to_end --deselect=extension/llm/custom_ops/test_sdpa_with_kv_cache.py::SDPATestForSpeculativeDecode::test_sdpa_with_cache_seq_len_130 --deselect=devtools/inspector/tests/inspector_test.py::TestInspector::test_etrecord_populates_correct_edge_dialect_aot_intermediate_outputs diff --git a/runtime/backend/interface.h b/runtime/backend/interface.h index 921d9ed324d..d0e54c3218d 100644 --- a/runtime/backend/interface.h +++ b/runtime/backend/interface.h @@ -92,7 +92,15 @@ class BackendInterface { * * @param[in] handle An opaque handle returned by `init()`. Usually a backend * executable unit. This executable unit should be ready to execute the - * delegate blobs. + * delegate blobs. One handle can serve several calls in a method, each + * with its own `args`. The emitter does this when the backend set + * `PreprocessResult.share_handle_across_calls` and declared no scratch, + * and programs from ExecuTorch 0.7 and earlier can do it for any + * backend. Such programs rely on `execute()` reading `args` again on + * every call, keeping nothing on the handle that is tied to one call or + * that changes later results, and writing outputs only into the output + * tensors in `args`. Setting the flag promises this for every runtime + * version of the backend that may load the program. * @param[in] args The method’s inputs and outputs. * @retval Error::Ok if successful. */ diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index c56a55f35bd..9b8b9e34008 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -116,6 +116,7 @@ class BackendDelegate final { out->backend_ = backend; out->handle_ = nullptr; out->scratch_ = Span>(); + out->has_call_site_ = false; // Pass a pointer to this buffer to the backend. It's safe for the backend // to point its handle to this object, since it will outlive the backend. new (&out->segment_) FreeableBuffer(std::move(processed_data.get())); @@ -154,6 +155,15 @@ class BackendDelegate final { scratch_ = scratch; } + /// Whether an instruction already calls this delegate. + bool has_call_site() const { + return has_call_site_; + } + + void add_call_site() { + has_call_site_ = true; + } + Error Execute( BackendExecutionContext& backend_execution_context, Span args) const { @@ -240,6 +250,7 @@ class BackendDelegate final { const BackendInterface* backend_; DelegateHandle* handle_; Span> scratch_; + bool has_call_site_; }; /** @@ -1143,19 +1154,27 @@ Error Method::init( } chain_instruction_arg_lists[instr_idx] = res.get(); + const auto delegate_index = delegate_call->delegate_index(); + ET_CHECK_OR_RETURN_ERROR( + delegate_index >= 0 && + static_cast(delegate_index) < n_delegate_, + InvalidProgram, + "Delegate index %" PRId32 " >= num delegates %" ET_PRIsize_t, + delegate_index, + n_delegate_); + BackendDelegate& delegate = delegates_[delegate_index]; const auto scratch = delegate_call->scratch(); - if (scratch != nullptr && scratch->size() > 0) { - // The emitter appends one BackendDelegate per delegate call, so - // the index identifies the call site and the buffers can hang - // off the delegate itself. - const auto delegate_index = delegate_call->delegate_index(); - ET_CHECK_OR_RETURN_ERROR( - delegate_index >= 0 && - static_cast(delegate_index) < n_delegate_, - InvalidProgram, - "Delegate index %" PRId32 " >= num delegates %" ET_PRIsize_t, - delegate_index, - n_delegate_); + const bool has_scratch = scratch != nullptr && scratch->size() > 0; + // Scratch is planned for one call site and kept on the delegate, so + // a delegate with scratch must not serve any other call site. + ET_CHECK_OR_RETURN_ERROR( + !delegate.has_call_site() || + (!has_scratch && delegate.scratch().empty()), + InvalidProgram, + "Delegate %" PRId32 " has scratch and more than one call site", + delegate_index); + delegate.add_call_site(); + if (has_scratch) { Result>> buffers = gen_delegate_scratch( method_allocator, memory_manager_->planned_memory(), scratch); if (!buffers.ok()) { @@ -1167,7 +1186,7 @@ Error Method::init( instr_idx); return buffers.error(); } - delegates_[delegate_index].set_scratch(buffers.get()); + delegate.set_scratch(buffers.get()); } } break; case executorch_flatbuffer::InstructionArguments::JumpFalseCall: {