From 68a871e19090432aa27d0d5c99c915388fd99302 Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Sun, 13 Sep 2026 10:13:00 +0200 Subject: [PATCH 1/3] Fix use-after-free of unpacked constants in the XNNPACK weights cache 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, 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. --- backends/xnnpack/runtime/XNNCompiler.cpp | 9 ++++++++- backends/xnnpack/runtime/XNNWeightsCache.cpp | 21 ++++++++++++++++---- backends/xnnpack/runtime/XNNWeightsCache.h | 3 ++- 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 2b8ab9b3181..0a1aab09f22 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -2310,7 +2310,14 @@ ET_NODISCARD Error XNNCompiler::compileModel( std::vector packed_weights_names; if (use_weight_cache) { - auto packed_weights_names_result = weights_cache->finalize_for_runtime(); + // Constants XNNPACK did not pack come back here: the subgraph still points + // into them, so the executor has to own them for as long as the runtime. + std::vector retained_unpacked; + auto packed_weights_names_result = + weights_cache->finalize_for_runtime(&retained_unpacked); + for (FreeableBuffer& buffer : retained_unpacked) { + executor->unpacked_buffers_.push_back(std::move(buffer)); + } ET_CHECK_OR_RETURN_ERROR( packed_weights_names_result.ok(), Internal, diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index f39d003e2c0..ed4fc7df66d 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -296,13 +296,26 @@ Error XNNWeightsCache::initialize_for_runtime( return Error::Ok; } -Result> XNNWeightsCache::finalize_for_runtime() { +Result> XNNWeightsCache::finalize_for_runtime( + std::vector* retained_unpacked) { is_finalized_ = true; - // All data has been packed by create_runtime - // so we clear the unpacked data as it is no longer needed + // Most of this data was packed by create_runtime, which copied it into the + // packed region, so the unpacked copies can go. Not all of it was: operators + // like PReLU take their constants unpacked, and the subgraph keeps a pointer + // into this memory for the life of the runtime. Freeing those here leaves + // the runtime reading freed memory, so hand them back to the caller to own + // instead. This mirrors the non-weight-cache path, which parks them in + // XNNExecutor::unpacked_buffers_. for (FreeableBuffer& buffer : unpacked_data_) { - buffer.Free(); + auto name_entry = unpacked_data_to_name_.find(buffer.data()); + const bool was_packed = name_entry != unpacked_data_to_name_.end() && + name_to_packed_data_metadata_.count(name_entry->second) > 0; + if (was_packed || retained_unpacked == nullptr) { + buffer.Free(); + } else { + retained_unpacked->push_back(std::move(buffer)); + } } unpacked_data_.clear(); unpacked_data_to_name_.clear(); diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index d0459965b98..cb4518a97c4 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -89,7 +89,8 @@ class XNNWeightsCache { * This should only be called after creating the runtime. Returns * the name of all the packed weights used by this runtime */ - Result> finalize_for_runtime(); + Result> finalize_for_runtime( + std::vector* retained_unpacked = nullptr); // Taken from XNN_ALLOCATION_ALIGNMENT in xnnpack/common.h static const size_t kPackedAllocationAlignment = 64; From 8b2c23c1a405caf4190235802a12341e0b7c31f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 15 Sep 2026 17:06:37 +0200 Subject: [PATCH 2/3] Retain unpacked constants per value, matching the non-cache path 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. --- backends/xnnpack/runtime/XNNCompiler.cpp | 19 ++++--- backends/xnnpack/runtime/XNNWeightsCache.cpp | 54 ++++++++++++++------ backends/xnnpack/runtime/XNNWeightsCache.h | 23 ++++++++- 3 files changed, 69 insertions(+), 27 deletions(-) diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 0a1aab09f22..3ec1a33554c 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -2234,6 +2234,10 @@ ET_NODISCARD Error XNNCompiler::compileModel( Error err = Error::Ok; for (auto value : *flatbuffer_graph->xvalues()) { size_t prev_buffers = unpacked_buffers.size(); + // With the weights cache the buffers land in the cache rather than in + // unpacked_buffers, so track its list too. + size_t prev_cached_buffers = + use_weight_cache ? weights_cache->get_num_unpacked_data() : 0; err = defineTensor( subgraph.get(), remapped_ids, @@ -2262,6 +2266,10 @@ ET_NODISCARD Error XNNCompiler::compileModel( executor->unpacked_buffers_.push_back(std::move(unpacked_buffers[i])); } unpacked_buffers.resize(prev_buffers); + if (use_weight_cache) { + weights_cache->take_unpacked_data_from( + prev_cached_buffers, executor->unpacked_buffers_); + } } } @@ -2310,14 +2318,9 @@ ET_NODISCARD Error XNNCompiler::compileModel( std::vector packed_weights_names; if (use_weight_cache) { - // Constants XNNPACK did not pack come back here: the subgraph still points - // into them, so the executor has to own them for as long as the runtime. - std::vector retained_unpacked; - auto packed_weights_names_result = - weights_cache->finalize_for_runtime(&retained_unpacked); - for (FreeableBuffer& buffer : retained_unpacked) { - executor->unpacked_buffers_.push_back(std::move(buffer)); - } + // Constants XNNPACK does not pack were already moved to the executor in + // the value loop above, so everything left here is safe to free. + auto packed_weights_names_result = weights_cache->finalize_for_runtime(); ET_CHECK_OR_RETURN_ERROR( packed_weights_names_result.ok(), Internal, diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index ed4fc7df66d..36d6b850cb5 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -215,6 +215,17 @@ Error XNNWeightsCache::initialize_for_runtime( named_data_map_ = named_data_map; is_finalized_ = false; + // An earlier compile can bail out between its first load_unpacked_data and + // its finalize_for_runtime, leaving its constants here. This instance is + // shared by every model in the process, so without this the next model to + // compile successfully would adopt those buffers and hold them until it is + // unloaded. The model that loaded them is gone, so free them. + for (FreeableBuffer& buffer : unpacked_data_) { + buffer.Free(); + } + unpacked_data_.clear(); + unpacked_data_to_name_.clear(); + #ifndef _WIN32 if (packed_cache_path_.empty() || packed_file_fd_ >= 0) { return Error::Ok; @@ -296,26 +307,35 @@ Error XNNWeightsCache::initialize_for_runtime( return Error::Ok; } -Result> XNNWeightsCache::finalize_for_runtime( - std::vector* retained_unpacked) { +void XNNWeightsCache::take_unpacked_data_from( + size_t first_index, + std::vector& out) { + if (first_index >= unpacked_data_.size()) { + return; + } + for (size_t i = first_index; i < unpacked_data_.size(); i++) { + // The name map is keyed on the data pointer, which a move preserves, so + // look_up_or_insert keeps naming packed entries correctly after this. + out.push_back(std::move(unpacked_data_[i])); + } + // The moved-from entries stay in the list rather than being erased: + // FreeableBuffer deletes move assignment, so vector::erase does not compile, + // and a moved-from buffer holds a null pointer, which makes the Free() in + // finalize_for_runtime a no-op. Leaving them also keeps the indices handed + // out by get_num_unpacked_data() stable across calls. +} + +Result> XNNWeightsCache::finalize_for_runtime() { is_finalized_ = true; - // Most of this data was packed by create_runtime, which copied it into the - // packed region, so the unpacked copies can go. Not all of it was: operators - // like PReLU take their constants unpacked, and the subgraph keeps a pointer - // into this memory for the life of the runtime. Freeing those here leaves - // the runtime reading freed memory, so hand them back to the caller to own - // instead. This mirrors the non-weight-cache path, which parks them in - // XNNExecutor::unpacked_buffers_. + // Everything still here belongs to a value XNNPACK packed, so create_runtime + // has copied it into the packed region and the unpacked copy can go. Buffers + // for values XNNPACK does not pack (PReLU slopes, for example) were handed to + // the executor by take_unpacked_data_from while the graph was being built; + // the subgraph keeps pointers into those, so freeing them here would leave + // the runtime reading freed memory. for (FreeableBuffer& buffer : unpacked_data_) { - auto name_entry = unpacked_data_to_name_.find(buffer.data()); - const bool was_packed = name_entry != unpacked_data_to_name_.end() && - name_to_packed_data_metadata_.count(name_entry->second) > 0; - if (was_packed || retained_unpacked == nullptr) { - buffer.Free(); - } else { - retained_unpacked->push_back(std::move(buffer)); - } + buffer.Free(); } unpacked_data_.clear(); unpacked_data_to_name_.clear(); diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index cb4518a97c4..e9055acd26b 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -89,8 +89,27 @@ class XNNWeightsCache { * This should only be called after creating the runtime. Returns * the name of all the packed weights used by this runtime */ - Result> finalize_for_runtime( - std::vector* retained_unpacked = nullptr); + Result> finalize_for_runtime(); + + /** + * Transfers ownership of the unpacked buffers loaded since `first_index` + * out of this cache. finalize_for_runtime() will not free them. + * + * For values XNNPACK does not pack (PReLU slopes, for example) the subgraph + * keeps a pointer into the unpacked memory for the life of the runtime, so + * something has to keep those buffers alive past finalize_for_runtime(). + * Callers pair this with get_num_unpacked_data() taken before the value was + * defined, which is how the non-weights-cache path in XNNCompiler decides + * what to retain. + * + * @param[in] first_index Index into the unpacked buffer list, from + * get_num_unpacked_data() before the value was defined. + * @param[out] out Receives the buffers. The caller owns them and must keep + * them alive for at least as long as the runtime. + */ + void take_unpacked_data_from( + size_t first_index, + std::vector& out); // Taken from XNN_ALLOCATION_ALIGNMENT in xnnpack/common.h static const size_t kPackedAllocationAlignment = 64; From d14cb9845605e79a022d24dfc40ec58759f246e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Thu, 17 Sep 2026 11:55:36 +0200 Subject: [PATCH 3/3] Address review: stale comment, dead guard, and add the missing tests 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. --- backends/xnnpack/runtime/XNNWeightsCache.cpp | 17 ++++--- .../test/runtime/test_xnn_weights_cache.cpp | 48 +++++++++++++++++++ 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index 36d6b850cb5..3d1c5fec199 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -216,10 +216,11 @@ Error XNNWeightsCache::initialize_for_runtime( is_finalized_ = false; // An earlier compile can bail out between its first load_unpacked_data and - // its finalize_for_runtime, leaving its constants here. This instance is - // shared by every model in the process, so without this the next model to - // compile successfully would adopt those buffers and hold them until it is - // unloaded. The model that loaded them is gone, so free them. + // its finalize_for_runtime, leaving its constants here. A successful compile + // frees whatever is left in finalize_for_runtime, but a failed one never + // reaches it, so back to back failures would pile up on an instance that + // lives as long as the process. The model that loaded them is gone, so free + // them here. for (FreeableBuffer& buffer : unpacked_data_) { buffer.Free(); } @@ -310,9 +311,11 @@ Error XNNWeightsCache::initialize_for_runtime( void XNNWeightsCache::take_unpacked_data_from( size_t first_index, std::vector& out) { - if (first_index >= unpacked_data_.size()) { - return; - } + // No bounds check: the loop below is empty when first_index is at or past + // the end, and the list only grows during a compile so that cannot happen + // anyway. An index that is too small cannot be detected here, and would hand + // the executor buffers of values XNNPACK did pack; callers must pass the + // get_num_unpacked_data() taken immediately before the value was defined. for (size_t i = first_index; i < unpacked_data_.size(); i++) { // The name map is keyed on the data pointer, which a move preserves, so // look_up_or_insert keeps naming packed entries correctly after this. diff --git a/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp b/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp index ec409ee0a16..fc388bd679c 100644 --- a/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp +++ b/backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp @@ -320,6 +320,54 @@ TEST_F(XNNWeightsCacheTest, ReusePackedWeights) { ASSERT_EQ(packed_data_names.size(), 0); } +TEST_F(XNNWeightsCacheTest, TakenUnpackedDataSurvivesFinalize) { + // XNNPACK does not pack every constant it is given: PReLU slopes, for one, + // stay as the unpacked buffer the subgraph points at. Those buffers are taken + // out of the cache while the graph is built, and finalize_for_runtime must + // leave them alone while still freeing the ones it packed. + XNNWeightsCache cache; + cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()); + + Result packed = cache.load_unpacked_data("weight"); + ASSERT_EQ(packed.error(), Error::Ok); + // The index the caller would take immediately before defining the value it + // wants to keep, which is what XNNCompiler passes down. + size_t first_index = cache.get_num_unpacked_data(); + Result retained_load = cache.load_unpacked_data("bias"); + ASSERT_EQ(retained_load.error(), Error::Ok); + ASSERT_EQ(cache.get_num_unpacked_data(), first_index + 1); + + std::vector retained; + cache.take_unpacked_data_from(first_index, retained); + ASSERT_EQ(retained.size(), 1u); + ASSERT_EQ(retained[0].size(), static_cast(kSegmentSizes[1])); + + Result> names = cache.finalize_for_runtime(); + ASSERT_EQ(names.error(), Error::Ok); + + // Still readable, and still the "bias" segment, which SetUp fills with 2s. + ASSERT_NE(retained[0].data(), nullptr); + const uint8_t* data = static_cast(retained[0].data()); + for (size_t i = 0; i < retained[0].size(); i++) { + ASSERT_EQ(data[i], 2) << "byte " << i << " was freed or overwritten"; + } +} + +TEST_F(XNNWeightsCacheTest, InitializeForRuntimeClearsLeftoverUnpackedData) { + // A compile that bails out after loading a constant never reaches + // finalize_for_runtime. This instance outlives any one model, so the next + // initialize has to drop what the failed compile left behind. + XNNWeightsCache cache; + cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()); + Result loaded = cache.load_unpacked_data("weight"); + ASSERT_EQ(loaded.error(), Error::Ok); + ASSERT_EQ(cache.get_num_unpacked_data(), 1u); + + // No finalize_for_runtime(): this is the failed-compile path. + cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()); + ASSERT_EQ(cache.get_num_unpacked_data(), 0u); +} + #ifndef _WIN32 // Verify pack-and-run works when packed weight allocations go to a // MAP_SHARED file instead of heap. The cache path is unique per test so