Repository navigation
Let a backend share one delegate across calls of the same lowered module - #24127
Conversation
shoumikhin
added
the
release notes: exir
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24127
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 8fcaa95 with merge base 4db778e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
dc8afe1 to
65fe314
Compare
65fe314 to
dc84b5a
Compare
dc84b5a to
897a1e6
Compare
|
why does the TRT backend not just recognize that its already loaded this blob before? |
897a1e6 to
d617bc5
Compare
d617bc5 to
8dc9eec
Compare
|
@JacobSzwejbka It could, but deciding this at export time is simpler and safer. The blob is stored once in the .pte, but each delegate loads it on its own, and with the file or mmap loader each init gets its own copy or mapping, so the pointer does not identify the blob. So a backend cache would need a hash of the blob (or an ID written at export), a reference count across destroy(), and a lock, since a cache shared by several methods or programs can be used from more than one thread. Every backend that wants this would repeat that work, and the file loader would still read the whole blob once for each call when the method loads. The emitter already knows that these calls use one lowered module and run one after another in one method, so it can share one handle before anything is loaded. |
8dc9eec to
dfba074
Compare
dfba074 to
c08cf8c
Compare
c08cf8c to
48f6ea7
Compare
|
@claude review |
|
Claude finished @shoumikhin's task in 3m 24s —— View job This PR lets an opted-in backend share one Detailed reviewPR Review: #24127SummaryWhen a backend opts in, the emitter now emits one API Design
Testing
Code Quality
RecommendationApprove. All findings are minor. The most important change is to write down the cross-version contract in |
When a method calls the same lowered module several times, for example a denoising step that runs 10 times, the emitter writes one BackendDelegate per call. The runtime then runs the backend's init() once per call and keeps every resulting handle alive, even though all of them come from the same blob. A backend can now return PreprocessResult(share_handle_across_calls=True). Every call of that lowered module in a method then points at one BackendDelegate, so init() runs once and only one handle stays in memory. Each call still passes its own arguments. Backends that do not set the flag keep one delegate per call, as today. A delegate that declares scratch is never shared, because the runtime stores scratch on the delegate and scratch is planned per call. Method::init now rejects a program in which a delegate with scratch has more than one call site, which the emitter never writes. The flag is carried like scratch_specs, including through exir save and load, which bumps the serde schema to 5.6. The docs and the execute() comment state what a backend must guarantee before it sets it. The demo test backend opts in, since its runtime keeps no state and reads its arguments on every call.
48f6ea7 to
8fcaa95
Compare
|
Thanks for the review.
|
|
@claude review |
|
Claude finished @shoumikhin's task in 3m 45s —— View job This PR lets a backend that opts in share one Detailed reviewPR Review: #24127SummaryWhen a backend sets Three independent reviewers went through all 14 changed files, and I re-read the code behind each finding. None of us found a correctness bug in:
This review leaves out the points from the previous round that you already answered. Those are: scratch combined with sharing, serde aliasing, and the missing C++ test and Backward Compatibility
RecommendationApprove. The serde point is minor and matches an existing pattern. It's worth fixing if older readers should keep loading files from backends that don't opt in. |
|
Thanks.
|
| 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( |
There was a problem hiding this comment.
AoT we should stop this from happening right?
There was a problem hiding this comment.
Yes, and it already does. The emitter only reuses a delegate when the backend set share_handle_across_calls and the lowered module has no scratch_specs:
if (
lowered_module.share_handle_across_calls
and not lowered_module.scratch_specs
):
self.emitter_state.delegate_cache[lowered_module] = backend_delegate_indexSo a program from to_executorch() never has a delegate with scratch and two call sites. test_a_delegate_with_scratch_is_not_shared_across_calls checks this: two calls give two delegates, each with its own scratch.
This runtime check is only for a .pte that did not come from our emitter, for example one edited by hand or written by another tool. Method::init already checks the program instead of trusting it, like the delegate index check just above. It costs one bool per delegate. Older programs that share a delegate (from before #14564) have no scratch, so they still load.
| # 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 |
There was a problem hiding this comment.
How does this get actually set? My understanding is the partitioner will select N partitions each getting its own preprocessed bytes. How does one bytes then end up serving multiple calls?
There was a problem hiding this comment.
The backend sets it in preprocess(), the same way as scratch_specs:
return PreprocessResult(processed_bytes=blob, share_handle_across_calls=True)to_backend then copies it onto the LoweredBackendModule.
You are right about partitions. N partitions give N lowered modules and N delegates, and this PR does not change that.
Sharing only happens when one lowered module (so one blob) is called more than once in a method. That happens when you lower a piece once and then call it from a bigger model, for example a denoising step in a loop:
step = to_edge_transform_and_lower(export(Step(), x), partitioner=[p]).exported_program().module()
class Denoise(torch.nn.Module):
def __init__(self):
super().__init__()
self.step = step
def forward(self, x):
for _ in range(10):
x = self.step(x)
return xExport unrolls the loop into 10 calls of the same lowered module. Before this PR the emitter wrote 10 delegates for them, all pointing at one blob, so init() ran 10 times. Now, when the backend sets the flag, it writes 1. test_partitioned_module_shares_one_delegate_across_calls covers this path, and test_calls_of_one_lowered_module_share_one_delegate covers the to_backend path.
If you export the whole loop and partition it instead, the partitioner sees the unrolled graph. If every op is supported, the partitioner usually merges the 10 steps into one partition. If a CPU op sits between steps, each step becomes its own partition with its own delegate. This PR keys sharing on the lowered module, not on the bytes, so it does not merge those.
|
@pytorchbot cherry-pick --onto release/1.6 -c fixnewfeature |
Cherry picking #24127The cherry pick PR is at #24207 and it is recommended to link a fixnewfeature cherry pick PR with an issue. Details for Dev Infra teamRaised by workflow job |
…ule (pytorch#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 pytorch#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.
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
BackendDelegateper call. So at load time the runtime calls the backend'sinit()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 oneBackendDelegate, soinit()runs once. Each call still passes its own arguments.Method::initnow 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.scratch_specs, including throughexir.saveandexir.load, so the serde schema moves to version 5.6. The.ptefile format does not change.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
init()calls have no dedicated test.