Fix use-after-free of unpacked constants in the XNNPACK weights cache - #22779
msluszniak wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22779
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 New Failures, 3 Unclassified FailuresAs of commit d14cb98 with merge base 500849b ( NEW FAILURES - The following jobs have failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
… (#22779) finalize_for_runtime() frees every unpacked buffer on the grounds that "All data has been packed by create_runtime". That is not true for operators XNNPACK does not pack: PReLU reads its slope straight out of that memory for the life of the runtime, so the subgraph is left pointing at freed memory. #21480 fixed this for the path that does not use the weights cache, by parking the FreeableBuffers in XNNExecutor::unpacked_buffers_. The weights-cache path was never given the same treatment, and the android and apple presets both set EXECUTORCH_XNNPACK_ENABLE_WEIGHT_CACHE=ON, so our Android libs hit it. finalize_for_runtime() now hands back the buffers whose data was never packed, and XNNCompiler gives them to the executor, matching the non-cache path. Found via react-native-executorch: MediaPipe Face Mesh (23 PReLUs) returned all-NaN through the RNE runtime on a Galaxy S26 Ultra while the same .pte was exact under executor_runner. Minimal repro is prelu(prelu(x - 10) - 10) with slopes 0.25 and 0.5, which should give -6.25 for a zero input and gave -7.25e+30. After this change both are exact, and stay exact with 2 MB allocated between init and execute. Upstream PR: pytorch/executorch#22779
|
@JakeStevens I think you should take a look. |
|
@pytorchbot label "module: xnnpack" |
pytorch-bot
Bot
added
the
module: xnnpack
|
Hmm, all checks fail on some dependency error with |
finalize_for_runtime() frees every unpacked buffer on the grounds that "All data has been packed by create_runtime". That is not true for operators XNNPACK does not pack: PReLU reads its slope straight out of that memory for the life of the runtime, so the subgraph is left pointing at freed memory. pytorch#21480 fixed this for the path that does not use the weights cache, by parking the FreeableBuffers in XNNExecutor::unpacked_buffers_. The weights-cache path was never given the same treatment, so any build with EXECUTORCH_XNNPACK_ENABLE_WEIGHT_CACHE=ON still frees them. That is the default for the android and apple presets. finalize_for_runtime() now hands back the buffers whose data was never packed, and XNNCompiler gives them to the executor, matching the non-cache path. Being undefined behaviour it reads as environment-dependent. On a Galaxy S26 Ultra, a model that applies prelu(x - 10) twice with slopes 0.25 and 0.5 should return -6.25 for a zero input. Before this change it returned -0.0982864, and -7.25e+30 in a busy app process; MediaPipe Face Mesh, which has 23 PReLUs, returned all-NaN. After it, both are exact, and stay exact when the process allocates 2 MB between init and execute.
fc17ac0 to
68a871e
Compare
JakeStevens
left a comment
JakeStevens
left a comment
There was a problem hiding this comment.
Some notes inline.
These are about lines the change does not touch, so they could not be attached to one.
In backends/xnnpack/runtime/XNNWeightsCache.h, around line 85:
The doc block above this declaration was left as it was. It still describes only the return value, so nothing tells a reader what the new parameter is for, that the caller takes ownership of what comes back, or that those buffers have to outlive the runtime. The method just below documents its parameters in the param in and param out style, so the convention is already here. Please add a line for the new parameter.
Review feedback. The first version decided what to keep inside finalize_for_runtime by looking the buffer's name up in the packed-data map. That was wrong in three ways: look_up_or_insert stores an entry under the kernel and bias names joined together, so neither bare name was ever found and almost every constant stayed resident; the map holds every name the shared cache ever packed, not only this runtime's, so a stale entry could free a buffer this runtime still reads; and scale names never appear as packed keys at all, so scales of packed weights were kept when the other path frees them. Use the same signal the non-weights-cache path uses instead. compileModel already builds packed_value_ids, and the value loop already retains the buffers of values XNNPACK does not pack, scales included. Hand the cache's buffers for those values to the executor there, via take_unpacked_data_from, so both paths answer the question the same way. finalize_for_runtime goes back to taking no arguments: everything left in the cache by then belongs to a packed value, so it frees all of it. There is no longer a defaulted pointer out-parameter, and no null that means free-everything. Also free any unpacked data left over at the top of initialize_for_runtime. A compile that bails out between its first load and its finalize used to leave its constants in this process-wide cache for the next model to adopt.
| // out by get_num_unpacked_data() stable across calls. | ||
| } | ||
|
|
||
| Result<std::vector<std::string>> XNNWeightsCache::finalize_for_runtime() { |
There was a problem hiding this comment.
The description and the repro table still describe the first version, where finalize handed the buffers back. The version in the tree frees everything left and retains per value instead, so the numbers no longer match the code. Please update them.
| return Error::Ok; | ||
| } | ||
|
|
||
| void XNNWeightsCache::take_unpacked_data_from( |
There was a problem hiding this comment.
Some tests can still be provided. Load two named constants, take the count between them, take from that index, finalize, then check the taken buffer is still readable and the other is gone. A second one for the new cleanup: load a constant, skip finalize, initialize again, check the list is empty. Both fail on the base branch.
One warning so you are not surprised. The three weights cache test files are wired into the buck targets only. The cmake test target lists the executor test, the workspace manager test and the threadpool test, so nothing in the open source CI builds them today. A case added there is still worth having, but it will not turn CI red on a regression until that target is fixed.
|
some nits mostly around comment drift and then LGTM, thanks for the back and forth! |
Will address them tomorrow, thank you for your review :)) |
The comment above the cleanup loop in initialize_for_runtime described the first version's hazard. finalize_for_runtime now frees everything left, so the next model cannot adopt these buffers; the loop earns its place because back to back failed compiles never reach a finalize at all. Drop the bounds check in take_unpacked_data_from. The loop is already empty for an index at or past the end, and the list only grows during a compile, so it could not fire. The dangerous index is one that is too small, which cannot be detected here, so say that in the contract instead of implying it is checked. Add the two cases the review asked for: a taken buffer stays readable across finalize_for_runtime while the packed one is freed, and a load with no finalize is cleared by the next initialize_for_runtime. Both fail on the base branch.
|
Addressed code and PR description comments. |
3 participants
Summary
XNNWeightsCache::finalize_for_runtime()frees every unpacked buffer, on the grounds thatThat holds for weights XNNPACK packs, but not for operators that take their constants unpacked. PReLU reads its slope straight out of that memory for the life of the runtime, so the subgraph is left pointing at freed memory.
#21480 fixed exactly this for the path that does not use the weights cache, by parking the
FreeableBuffers inXNNExecutor::unpacked_buffers_. The weights-cache branch ofgetConstantDataPtr()was never given the same treatment, so a build withEXECUTORCH_XNNPACK_ENABLE_WEIGHT_CACHE=ONstill frees them. That is the default intools/cmake/preset/android.cmakeandapple_common.cmake.XNNCompilernow moves those buffers to the executor as each value is defined, taking everything the cache loaded for that value via the indexget_num_unpacked_data()returned just before it, which mirrors how the non-cache path decides what to retain.finalize_for_runtime()then frees whatever is left, all of which really was packed.initialize_for_runtime()also drops anything a previous compile left behind, since a compile that bails out never reaches a finalize and the cache outlives any one model.Repro
Two chained PReLUs, arranged so the slope is readable from the output:
prelu(prelu(x - 10) - 10)with slopes 0.25 and 0.5. For a zero input the answer is-6.25. On a Galaxy S26 Ultra, arm64, weights cache on:executor_runner41.2182, 37.6233, 159.611, 170.073The after column was measured on the first revision of this patch, which held the buffers back inside
finalize_for_runtime(). The current revision retains the same set at value-definition time instead, so the numbers have not been re-taken on the code as it now stands.It is undefined behaviour, so it presents as environment-dependent: one build of
executor_runnerreturned the right answer, and the same source rebuilt with unrelated flags did not. A single PReLU usually survives, since nothing has reused the block yet; two or more do not.Test plan
Verified on device as above, plus two cases in
test_xnn_weights_cache.cpp: a taken buffer stays readable acrossfinalize_for_runtime()while the packed one is freed, and a load with no finalize is cleared by the nextinitialize_for_runtime(). Both fail on the base branch. Note they only build under buck today, so they will not turn OSS CI red until the cmake target lists that file.An end-to-end test is still out of reach:
backends/xnnpack/test/ops/test_prelu.pyruns through the pybindings, which load the whole.pteinto one buffer, and thereFreeableBuffer::Free()is a no-op, so the dangling read still lands on valid memory.Note
getConstantDataPtr()has a sibling case I did not touch: for inline constants it returnsconstant_data_ptr + offsetintoprocessed, whichXNNPACKBackendfrees right aftercompileModel. That looks like the same defect for.ptes without named data, but I could not produce one that reaches it, so I left it alone rather than change it blind.cc @GregoryComer @digantdesai @cbilgin @JakeStevens