Skip to content

Commit 096a20a

Browse files
authored
Persist seed-mismatch re-packs so the packed cache converges after an XNNPACK upgrade (pytorch#22183)
Differential Revision: D117550758 Pull Request resolved: pytorch#22183
1 parent 7224da6 commit 096a20a

3 files changed

Lines changed: 187 additions & 26 deletions

File tree

‎backends/xnnpack/runtime/XNNWeightsCache.cpp‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,7 @@ size_t XNNWeightsCache::look_up(
338338
const xnn_weights_cache_look_up_key* cache_key) {
339339
const void* unpacked_weights_ptr = cache_key->kernel;
340340
const void* unpacked_bias_ptr = cache_key->bias;
341+
context->last_lookup_seed_mismatch_ = false;
341342
auto entry = context->unpacked_data_to_name_.find(unpacked_weights_ptr);
342343
if (entry == context->unpacked_data_to_name_.end()) {
343344
context->last_lookup_unnamed_ = true;
@@ -369,6 +370,7 @@ size_t XNNWeightsCache::look_up(
369370
weight_bias_name.c_str(),
370371
packed_weight_entry->second.seed,
371372
cache_key->seed);
373+
context->last_lookup_seed_mismatch_ = true;
372374
return SIZE_MAX;
373375
}
374376
packed_weight_entry->second.in_current_runtime = true;
@@ -377,7 +379,21 @@ size_t XNNWeightsCache::look_up(
377379

378380
void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) {
379381
#ifndef _WIN32
380-
if (context->last_lookup_unnamed_ || context->loaded_from_disk_) {
382+
// XNNPACK also re-packs for a different runtime context by calling
383+
// reserve_space with no preceding look_up, so consume the flag here:
384+
// only the re-pack directly following a seed-mismatch look_up may take
385+
// the file-backed path.
386+
const bool seed_mismatch_repack = context->last_lookup_seed_mismatch_;
387+
context->last_lookup_seed_mismatch_ = false;
388+
389+
// Unnamed constants never reload by name, and an incidental re-pack of an
390+
// already-valid loaded cache needn't persist — both go to heap. A
391+
// seed-mismatch re-pack is different: the loaded entry is stale (XNNPACK
392+
// upgrade), so it must take the file-backed path below, letting
393+
// save_packed_index persist the refreshed seed so the next launch hits
394+
// instead of re-packing into heap (dirty memory) every time.
395+
if (context->last_lookup_unnamed_ ||
396+
(context->loaded_from_disk_ && !seed_mismatch_repack)) {
381397
return context->reserve_space_heap(n);
382398
}
383399
if (context->packed_file_fd_ >= 0) {
@@ -670,6 +686,13 @@ bool XNNWeightsCache::load_packed_cache() {
670686
const size_t index_region_end = file_size - 20;
671687

672688
void* map = mmap(nullptr, file_size, PROT_READ, MAP_SHARED, fd, 0);
689+
// The live MAP_SHARED mapping keeps this open file description — and with
690+
// it the flock — alive for as long as the mapping exists, so close() alone
691+
// does NOT drop the lock. Without the explicit LOCK_UN the O_RDWR reopen in
692+
// initialize_for_runtime always fails with EWOULDBLOCK, leaving
693+
// packed_file_fd_ == -1, which silently disables every write path (new
694+
// packs, seed-mismatch re-packs, save_packed_index) after a warm load.
695+
flock(fd, LOCK_UN);
673696
close(fd);
674697
if (map == MAP_FAILED) {
675698
return false;

‎backends/xnnpack/runtime/XNNWeightsCache.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,16 @@ class XNNWeightsCache {
236236
// different runtime context) that doesn't need to persist.
237237
bool loaded_from_disk_{false};
238238

239+
// Set by look_up when a named entry is present but its cached seed does not
240+
// match the current ukernel's seed (i.e. an XNNPACK upgrade invalidated the
241+
// loaded packing). Unlike an incidental re-pack of an already-valid loaded
242+
// cache, this stale entry MUST be re-packed to the file and persisted;
243+
// otherwise every launch re-packs it into heap (anonymous dirty memory that
244+
// can OOM/jetsam the app in the background) and the cache never converges.
245+
// Consumed (and cleared) by reserve_space so the file-backed path is taken
246+
// only by the re-pack that directly follows the seed-mismatch look_up.
247+
bool last_lookup_seed_mismatch_{false};
248+
239249
// See mutex() for the locking contract — caller-owned, no internal
240250
// use within this class.
241251
std::mutex instance_mutex_;

‎backends/xnnpack/test/runtime/test_xnn_weights_cache.cpp‎

Lines changed: 153 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,36 @@ class XNNWeightsCacheTest : public ::testing::Test {
226226
ASSERT_EQ(status, xnn_status_success);
227227
}
228228

229+
// One "app launch": a fresh cache instance bound to cache_path loads the
230+
// file if it exists, packs-and-runs the graph, and persists the index.
231+
void RunLaunchWithCacheFile(
232+
const std::string& cache_path,
233+
std::vector<float>& output) {
234+
const std::vector<size_t> batches{1, 2, 3};
235+
const size_t num_batches = 1 * 2 * 3;
236+
const size_t padding = 32;
237+
std::vector<float> input(
238+
num_batches * kLaunchInputChannels + padding, 1.0f);
239+
output.assign(num_batches * kLaunchOutputChannels, 0.0f);
240+
241+
XNNWeightsCache cache;
242+
cache.set_packed_cache_path(cache_path);
243+
ASSERT_EQ(
244+
cache.initialize_for_runtime(memory_allocator_.get(), data_map_.get()),
245+
Error::Ok);
246+
BuildAndRunGraphWithWeightsCache(
247+
cache,
248+
batches,
249+
kLaunchInputChannels,
250+
kLaunchOutputChannels,
251+
input.data(),
252+
output.data());
253+
ASSERT_EQ(cache.save_packed_index(), Error::Ok);
254+
}
255+
256+
static constexpr size_t kLaunchInputChannels = 3;
257+
static constexpr size_t kLaunchOutputChannels = 4;
258+
229259
// Program builder constants.
230260
static constexpr int kSegmentAlignment = 16;
231261
static constexpr std::array<int, 2> kSegmentSizes{384, 128};
@@ -736,6 +766,57 @@ void write_le_u64(std::ostream& f, uint64_t v) {
736766
}
737767
}
738768

769+
// Locate the 4-byte seed field of the first index entry by walking the
770+
// footer [index_start:u64][entry_count:u32][magic:u32][version:u32] at the
771+
// end of the file, then the entry record
772+
// [name_len:u32][name][file_offset:u64][data_size:u64][seed:u32].
773+
// Returns 0 when the file is too small or carries no entries.
774+
size_t find_first_entry_seed_offset(std::istream& f) {
775+
f.seekg(0, std::ios::end);
776+
const size_t file_size = f.tellg();
777+
if (file_size < 24) {
778+
return 0;
779+
}
780+
uint8_t footer[20];
781+
f.seekg(file_size - 20);
782+
f.read(reinterpret_cast<char*>(footer), 20);
783+
if (read_le_u32(footer + 8) == 0) {
784+
return 0;
785+
}
786+
const uint64_t index_start = read_le_u64(footer);
787+
f.seekg(index_start);
788+
uint8_t name_len_buf[4];
789+
f.read(reinterpret_cast<char*>(name_len_buf), 4);
790+
return index_start + 4 + read_le_u32(name_len_buf) + 8 + 8;
791+
}
792+
793+
uint32_t read_first_entry_seed(const std::string& path) {
794+
std::ifstream f(path, std::ios::binary);
795+
const size_t seed_offset = find_first_entry_seed_offset(f);
796+
if (seed_offset == 0) {
797+
return 0;
798+
}
799+
uint8_t seed_buf[4];
800+
f.seekg(seed_offset);
801+
f.read(reinterpret_cast<char*>(seed_buf), 4);
802+
return read_le_u32(seed_buf);
803+
}
804+
805+
void write_first_entry_seed(const std::string& path, uint32_t seed) {
806+
std::fstream f(path, std::ios::binary | std::ios::in | std::ios::out);
807+
const size_t seed_offset = find_first_entry_seed_offset(f);
808+
if (seed_offset == 0) {
809+
return;
810+
}
811+
f.seekp(seed_offset);
812+
write_le_u32(f, seed);
813+
}
814+
815+
off_t file_size_of(const std::string& path) {
816+
struct stat st {};
817+
return ::stat(path.c_str(), &st) == 0 ? st.st_size : -1;
818+
}
819+
739820
} // namespace
740821

741822
// A cache file written by older code (kCacheVersion=1) carries no per-entry
@@ -889,31 +970,8 @@ TEST_F(
889970

890971
// Corrupt the seed field of the first entry to a value no real ukernel
891972
// would emit (0xDEADBEEF).
892-
{
893-
std::fstream f(cache_path, std::ios::binary | std::ios::in | std::ios::out);
894-
ASSERT_TRUE(f.is_open());
895-
f.seekg(0, std::ios::end);
896-
size_t file_size = f.tellg();
897-
ASSERT_GE(file_size, 24u);
898-
899-
uint8_t footer_buf[20];
900-
f.seekg(file_size - 20);
901-
f.read(reinterpret_cast<char*>(footer_buf), 20);
902-
uint64_t index_start = read_le_u64(footer_buf);
903-
uint32_t entry_count = read_le_u32(footer_buf + 8);
904-
ASSERT_GT(entry_count, 0u);
905-
906-
f.seekg(index_start);
907-
uint8_t name_len_buf[4];
908-
f.read(reinterpret_cast<char*>(name_len_buf), 4);
909-
uint32_t name_len = read_le_u32(name_len_buf);
910-
911-
size_t seed_offset = index_start + 4 + name_len + 8 + 8;
912-
f.seekp(seed_offset);
913-
uint32_t corrupted = 0xDEADBEEFu;
914-
f.write(reinterpret_cast<const char*>(&corrupted), 4);
915-
f.close();
916-
}
973+
write_first_entry_seed(cache_path, 0xDEADBEEFu);
974+
ASSERT_EQ(read_first_entry_seed(cache_path), 0xDEADBEEFu);
917975

918976
// Reload and run. Output must still match baseline.
919977
std::vector<float> after_corruption(num_batches * output_channels, 0.0f);
@@ -936,6 +994,76 @@ TEST_F(
936994
::unlink(cache_path.c_str());
937995
}
938996

997+
// An XNNPACK upgrade changes the per-ukernel seed, so every loaded entry
998+
// misses look_up and is re-packed. That re-pack must land in the cache file,
999+
// not on the heap: only the file-backed path lets save_packed_index write
1000+
// back the refreshed seed. On the heap path the file keeps the stale seed
1001+
// and the same re-pack repeats into anonymous dirty memory every launch.
1002+
TEST_F(XNNWeightsCacheTest, SeedMismatch_RepackIsPersistedToFile) {
1003+
const std::string cache_path = std::string("/tmp/xnn_weights_cache_reseed_") +
1004+
std::to_string(::getpid()) + ".packed_cache";
1005+
::unlink(cache_path.c_str());
1006+
1007+
std::vector<float> output;
1008+
RunLaunchWithCacheFile(cache_path, output);
1009+
const off_t size_after_first_launch = file_size_of(cache_path);
1010+
ASSERT_GT(size_after_first_launch, 0);
1011+
1012+
const uint32_t real_seed = read_first_entry_seed(cache_path);
1013+
ASSERT_NE(real_seed, 0u);
1014+
1015+
// Simulate the upgrade: the persisted seed no longer matches what the
1016+
// current ukernel reports for the same weights.
1017+
write_first_entry_seed(cache_path, 0xDEADBEEFu);
1018+
ASSERT_EQ(read_first_entry_seed(cache_path), 0xDEADBEEFu);
1019+
1020+
std::vector<float> output_after_upgrade;
1021+
RunLaunchWithCacheFile(cache_path, output_after_upgrade);
1022+
1023+
EXPECT_EQ(read_first_entry_seed(cache_path), real_seed)
1024+
<< "seed-mismatch re-pack must be written back to the cache file";
1025+
EXPECT_GT(file_size_of(cache_path), size_after_first_launch)
1026+
<< "the re-packed weights must be appended to the file, not left on heap";
1027+
EXPECT_EQ(output_after_upgrade, output);
1028+
1029+
::unlink(cache_path.c_str());
1030+
}
1031+
1032+
// The point of persisting the re-pack is convergence: exactly one launch
1033+
// pays for the rebuild, and every launch after it is a plain cache hit —
1034+
// no new packing, no file growth. Before the fix each launch re-packed into
1035+
// heap and the file never stopped being stale.
1036+
TEST_F(XNNWeightsCacheTest, SeedMismatch_CacheConvergesAfterOneRebuild) {
1037+
const std::string cache_path =
1038+
std::string("/tmp/xnn_weights_cache_converge_") +
1039+
std::to_string(::getpid()) + ".packed_cache";
1040+
::unlink(cache_path.c_str());
1041+
1042+
std::vector<float> baseline;
1043+
RunLaunchWithCacheFile(cache_path, baseline);
1044+
const uint32_t real_seed = read_first_entry_seed(cache_path);
1045+
ASSERT_NE(real_seed, 0u);
1046+
1047+
write_first_entry_seed(cache_path, 0xDEADBEEFu);
1048+
1049+
// Rebuild launch: absorbs the one-time re-pack.
1050+
std::vector<float> output;
1051+
RunLaunchWithCacheFile(cache_path, output);
1052+
const off_t size_after_rebuild = file_size_of(cache_path);
1053+
ASSERT_GT(size_after_rebuild, 0);
1054+
1055+
// Steady state: three more launches must be pure hits.
1056+
for (int launch = 0; launch < 3; ++launch) {
1057+
RunLaunchWithCacheFile(cache_path, output);
1058+
EXPECT_EQ(file_size_of(cache_path), size_after_rebuild)
1059+
<< "cache did not converge; launch " << launch << " re-packed";
1060+
EXPECT_EQ(read_first_entry_seed(cache_path), real_seed);
1061+
EXPECT_EQ(output, baseline);
1062+
}
1063+
1064+
::unlink(cache_path.c_str());
1065+
}
1066+
9391067
// save_packed_index must be a true no-op when no new reserve_space happened
9401068
// since the last save — same content but writing would still bump mtime,
9411069
// making the cache file look modified on every model load.

0 commit comments

Comments
 (0)