Repository navigation
Isolate cached delegates by Module lifetime (#24462) - #24462
alirezamonfared wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24462
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 2f58db1 with merge base b1701a4 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
|
@alirezamonfared has exported this pull request. If you are a Meta employee, you can view the originating Diff in D122851825. |
This PR needs a
|
9179591 to
362a88f
Compare
Summary: Pull Request resolved: pytorch#24462 ## Summary The QNN backend caches a `QnnManager` so methods in one ExecuTorch `Module` can share an expensive QNN context. The previous cache key used only the exported context signature, so independent modules loading equivalent programs could share a manager that was placement-allocated from one module's runtime allocator. Key the cache by both context signature and runtime allocator. Methods in one module continue to share a context, while independent modules receive managers with independent lifetimes. The change also: - serializes cache lookup, insertion, membership checks, and removal; - avoids allocating a manager on cache hits; - removes cache state before running the non-trivial `QnnManager` destructor. ## Scope This fixes sequential independent-module lifetime isolation. It does not change exported program contents, backend selection, or model behavior. Broader same-module concurrent execution and initialization races remain outside this change. Differential Revision: D122851825
|
@alirezamonfared another suggestion while we're here: if InitBackend/InitContext/AllocateTensor fails below, we return without calling ~QnnManager(). Since init failed, ExecuTorch never calls destroy() for this handle, so the QNN backend/context handles leak. Matters for apps that retry or fall back to another backend. A scope guard right after the placement new, dismissed just before add_cached_delegate, would cover all three error paths. wdyt ? |
qti-horodnic
left a comment
There was a problem hiding this comment.
Thanks @alirezamonfared, left a few comments. As I mentioned in one of them, for modules that load the same blob, memory usage might increase since there will multiple contexts. Can you confirm that no FCB flow loads the same blob bytes into several Modules (e.g. weight-sharing bins split across .pte files)? I think it should be fine but I want to make sure there are no FCB regressions
| std::lock_guard<std::mutex> guard(mutex_); | ||
| delegate_map_[signature] = handle; | ||
| delegate_map_rev_[handle] = signature; | ||
| delegate_map_[key] = handle; |
There was a problem hiding this comment.
Two init calls with the same key can both miss the cache. The second call overwrites delegate_map_[key], but delegate_map_rev_ keeps both handles, so a later erase removes the wrong entry. Can we use try_emplace instead? And if the key exists, destroy the new manager and return the cached handle.
| // destructible, we must call the destructor manually in destroy(). | ||
| new (qnn_manager) QnnManager(qnn_executorch_options, qnn_context_blob); | ||
|
|
||
| ET_CHECK_OR_RETURN_ERROR( |
There was a problem hiding this comment.
This is a pre-existing issue, but if InitBackend, InitContext, CompileDlc or AllocateTensor fails, ~QnnManager() doesn't run. Can we call qnn_manager->~QnnManager() before each error return or use a scope guard?
|
|
||
| auto iter = delegate_map_.find(signature); | ||
| if (iter != delegate_map_.end()) { | ||
| const DelegateKey key{signature, runtime_allocator}; |
There was a problem hiding this comment.
Modules that load the same blob no longer share one QNN context, so HTP memory can increase. Also, destroy() now runs the full destructor, so the tensor wrappers are freed. Please mention this in the PR description.
| QnnManager* qnn_manager = static_cast<QnnManager*>(handle); | ||
| qnn_manager->Destroy(); | ||
| erase_cached_delegate(handle); | ||
| if (handle != nullptr && erase_cached_delegate(handle)) { |
There was a problem hiding this comment.
Can we add a test for this branch? A unit test that loads the same .pte in 2 programs, deletes the first, and executes the second. Also can we check that 2 methods of one Module still get the same handle?
There was a problem hiding this comment.
Also, looking a bit more at this, the methods of one Module share a handle but nothing counts the users. So there can be case where unload_method("prefill") destroys the manager while decode is loaded, then decode fails with "DelegateHandle has been deleted". This is another existing issue which not caused by this pr, do you want to handle this here too? Maybe keeping a user count per cache entry and run ~QnnManager() only when the count is 0? This will make the solution more robust, but I'm fine if you want to keep the scope more contained.
| // NOTE: Since we use placement new and since this type is not trivially | ||
| // destructible, we must call the destructor manually in destroy(). | ||
| new (qnn_manager) QnnManager(qnn_executorch_options, qnn_context_blob); | ||
| // TODO: this is a temporal solution for multi-graph support, will be |
There was a problem hiding this comment.
Minor: pre-existing typo, I guess should be temporary, not temporal.
Summary: The QNN backend caches a `QnnManager` so methods in one ExecuTorch `Module` can share a QNN context. A context signature alone also matched independent modules loading identical program bytes, even though the manager was placement-allocated from only one module's runtime allocator. Key the cache by both signature and allocator to give independent modules independent manager lifetimes. Count each loaded method's reference to its cached manager and destroy the manager after its last method releases it. Protect cache operations with the mutex and use `try_emplace` to preserve the first manager published for a key. A competing initializer returns that manager and destroys its own candidate. A destructor-only ownership guard cleans up every initialization error path, including backend, context, DLC compilation, and tensor allocation failures. Scope HTP spill-fill groups by runtime allocator and configured buffer size. Each context receives its group handle explicitly, replacing the process-global static handle. Keep the first group manager alive until all group members and pending initializations are gone, and free each member before the first context. Serialize initialization within a group so concurrent loads select one first context. Separate live modules now allocate separate QNN contexts, which can increase HTP memory use when they load identical blobs. Sharing between methods within one module remains intact. Full `QnnManager` destruction also releases its tensor wrappers and other non-trivial members. FCB export contents and selection are unchanged; FCB weight sharing within an artifact does not provide shared runtime ownership across independent modules. Export the lifecycle tests so they are visible in the PR. They cover both module destruction orders, same-module sharing, cleanup/reload, and method unload/reload on the HTP simulator. C++ tests cover allocator isolation, reference accounting, concurrent cache publication, spill-fill group isolation, cleanup/reload, and owner destruction order. The simulator lifecycle suite also runs with multi-context enabled and a nonzero spill-fill size; target SDK registration still requires device verification. Concurrent execution through one module and execution racing destruction still require client coordination. Authored with AI assistance from Codex. Differential Revision: D122851825 Signed-off-by: Ali Monfared <alimonfared@meta.com>
362a88f to
2f58db1
Compare
|
@alirezamonfared both CLA & linters need a fix |
|
@claude can you review mine and Arik's comments on this PR and check if they all have been addressd |
|
Claude finished @psiddh's task in 2m 15s —— View job All of @qti-horodnic's and @psiddh's code comments are addressed in the current revision, but lintrunner is still failing (psiddh's last comment), and the CLA check still needs the author to link their commit email. Comment-by-comment status@qti-horodnic (review, CHANGES_REQUESTED)
@psiddh
Remaining items
|
Summary:
The QNN backend caches a
QnnManagerso methods in one ExecuTorchModulecan share a QNN context. A context signature alone also matched independent modules loading identical program bytes, even though the manager was placement-allocated from only one module's runtime allocator. Key the cache by both signature and allocator to give independent modules independent manager lifetimes.Count each loaded method's reference to its cached manager and destroy the manager after its last method releases it. Protect cache operations with the mutex and use
try_emplaceto preserve the first manager published for a key. A competing initializer returns that manager and destroys its own candidate. A destructor-only ownership guard cleans up every initialization error path, including backend, context, DLC compilation, and tensor allocation failures.Scope HTP spill-fill groups by runtime allocator and configured buffer size. Each context receives its group handle explicitly, replacing the process-global static handle. Keep the first group manager alive until all group members and pending initializations are gone, and free each member before the first context. Serialize initialization within a group so concurrent loads select one first context.
Separate live modules now allocate separate QNN contexts, which can increase HTP memory use when they load identical blobs. Sharing between methods within one module remains intact. Full
QnnManagerdestruction also releases its tensor wrappers and other non-trivial members. FCB export contents and selection are unchanged; FCB weight sharing within an artifact does not provide shared runtime ownership across independent modules.Export the lifecycle tests so they are visible in the PR. They cover both module destruction orders, same-module sharing, cleanup/reload, and method unload/reload on the HTP simulator. C++ tests cover allocator isolation, reference accounting, concurrent cache publication, spill-fill group isolation, cleanup/reload, and owner destruction order. The simulator lifecycle suite also runs with multi-context enabled and a nonzero spill-fill size; target SDK registration still requires device verification. Concurrent execution through one module and execution racing destruction still require client coordination.
Authored with AI assistance from Codex.
Differential Revision: D122851825