diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 2b8ab9b3181..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,6 +2318,8 @@ ET_NODISCARD Error XNNCompiler::compileModel( std::vector packed_weights_names; if (use_weight_cache) { + // 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(), diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index f39d003e2c0..3d1c5fec199 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -215,6 +215,18 @@ 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. 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(); + } + unpacked_data_.clear(); + unpacked_data_to_name_.clear(); + #ifndef _WIN32 if (packed_cache_path_.empty() || packed_file_fd_ >= 0) { return Error::Ok; @@ -296,11 +308,35 @@ Error XNNWeightsCache::initialize_for_runtime( return Error::Ok; } +void XNNWeightsCache::take_unpacked_data_from( + size_t first_index, + std::vector& out) { + // 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. + 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; - // All data has been packed by create_runtime - // so we clear the unpacked data as it is no longer needed + // 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_) { buffer.Free(); } diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index d0459965b98..e9055acd26b 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -91,6 +91,26 @@ class XNNWeightsCache { */ 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; 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