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.

Author: shoumikhin

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

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,

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

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",

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

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

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(),