From d4fbbaf35865684379b752af447b22f49b79a24f Mon Sep 17 00:00:00 2001 From: Moeez Ahmed Date: Fri, 21 Aug 2026 04:13:27 +0530 Subject: [PATCH] fix(graph): deduplicate shared mmap weight files across component graphs Signed-off-by: Moeez Ahmed --- cactus-graph/cactus_graph.h | 13 +++++++- cactus-graph/src/io.cpp | 57 ++++++++++++++++++++++++++++------ cactus-graph/tests/test_io.cpp | 41 ++++++++++++++++++++++++ 3 files changed, 101 insertions(+), 10 deletions(-) diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 9475f4293..28c7c978a 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -864,7 +864,7 @@ class CactusGraph { size_t attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name); static CactusGraph from_serialized(const GraphFile::SerializedGraph& serialized); size_t next_node_id_; - std::vector> mapped_files_; + std::vector> mapped_files_; std::unordered_map weight_cache_; std::unordered_map node_to_mapped_file_; std::vector debug_nodes_; @@ -919,6 +919,17 @@ namespace GraphFile { void save_graph(const CactusGraph& graph, const std::string& filename); void save_node(CactusGraph& graph, size_t node_id, const std::string& filename); + class MappedFile; + + class MappedFileRegistry { + public: + static std::shared_ptr get_or_load(const std::string& filename); + static void clear(); + private: + static std::mutex mutex_; + static std::unordered_map> cache_; + }; + class MappedFile { public: MappedFile(const std::string& filename); diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 2d90133a4..664ab2677 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -437,7 +437,7 @@ CactusGraph CactusGraph::load(const std::string& path) { size_t CactusGraph::mmap_embeddings(const std::string& filename) { const std::string resolved_filename = resolve_quantized_weight_file(filename); - auto mapped_file = std::make_unique(resolved_filename); + auto mapped_file = GraphFile::MappedFileRegistry::get_or_load(resolved_filename); const auto& shape = mapped_file->shape(); if (shape.size() != 2) { @@ -492,9 +492,9 @@ size_t CactusGraph::mmap_embeddings(const std::string& filename) { } size_t file_idx = mapped_files_.size(); - mapped_files_.push_back(std::move(mapped_file)); + mapped_files_.push_back(mapped_file); node_to_mapped_file_[node_id] = file_idx; - weight_cache_[filename] = node_id; + weight_cache_[resolved_filename] = node_id; return node_id; } @@ -505,7 +505,7 @@ size_t CactusGraph::mmap_weights(const std::string& filename) { return it->second; } - auto mapped_file = std::make_unique(resolved_filename); + auto mapped_file = GraphFile::MappedFileRegistry::get_or_load(resolved_filename); const auto& shape = mapped_file->shape(); Precision precision = mapped_file->precision(); @@ -553,7 +553,7 @@ size_t CactusGraph::mmap_weights(const std::string& filename) { } size_t file_idx = mapped_files_.size(); - mapped_files_.push_back(std::move(mapped_file)); + mapped_files_.push_back(mapped_file); node_to_mapped_file_[node_id] = file_idx; weight_cache_[resolved_filename] = node_id; return node_id; @@ -571,7 +571,7 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) throw std::invalid_argument("Can only bind mmap weights to input nodes"); } - auto mapped_file = std::make_unique(resolved_filename); + auto mapped_file = GraphFile::MappedFileRegistry::get_or_load(resolved_filename); const auto& shape = mapped_file->shape(); Precision precision = mapped_file->precision(); auto& buffer = node.output_buffer; @@ -640,7 +640,7 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) } size_t file_idx = mapped_files_.size(); - mapped_files_.push_back(std::move(mapped_file)); + mapped_files_.push_back(mapped_file); node_to_mapped_file_[node_id] = file_idx; weight_cache_[resolved_filename] = node_id; } @@ -666,7 +666,7 @@ void CactusGraph::release_all_weight_pages() { } size_t CactusGraph::embedding(const std::string& filename, size_t indices, ComputeBackend backend) { - auto mapped_file = std::make_unique(filename); + auto mapped_file = GraphFile::MappedFileRegistry::get_or_load(filename); const auto& shape = mapped_file->shape(); if (shape.size() != 2) { @@ -677,7 +677,7 @@ size_t CactusGraph::embedding(const std::string& filename, size_t indices, Compu size_t embeddings_node = input(shape, precision); set_external_input(embeddings_node, const_cast(mapped_file->data()), precision); - mapped_files_.push_back(std::move(mapped_file)); + mapped_files_.push_back(mapped_file); const auto& idx_shape = get_output_buffer(indices).shape; std::vector output_shape = idx_shape; @@ -894,6 +894,45 @@ void save_node(CactusGraph& graph, size_t node_id, const std::string& filename) } } +// MappedFileRegistry implementation + +namespace GraphFile { + +std::mutex MappedFileRegistry::mutex_; +std::unordered_map> MappedFileRegistry::cache_; + +static std::string canonical_file_key(const std::string& filename) { + std::error_code ec; + auto abs_path = std::filesystem::absolute(filename, ec); + if (!ec) { + auto can_path = std::filesystem::canonical(abs_path, ec); + if (!ec) return can_path.string(); + return abs_path.lexically_normal().string(); + } + return filename; +} + +std::shared_ptr MappedFileRegistry::get_or_load(const std::string& filename) { + std::string key = canonical_file_key(filename); + std::lock_guard lock(mutex_); + auto it = cache_.find(key); + if (it != cache_.end()) { + if (auto sp = it->second.lock()) { + return sp; + } + } + auto sp = std::make_shared(filename); + cache_[key] = sp; + return sp; +} + +void MappedFileRegistry::clear() { + std::lock_guard lock(mutex_); + cache_.clear(); +} + +} // namespace GraphFile + // MappedFile implementation MappedFile::MappedFile(const std::string& filename) diff --git a/cactus-graph/tests/test_io.cpp b/cactus-graph/tests/test_io.cpp index d792b6102..40370fdc1 100644 --- a/cactus-graph/tests/test_io.cpp +++ b/cactus-graph/tests/test_io.cpp @@ -453,6 +453,46 @@ bool test_save_load_preserves_kv_cache_num_slots() { } } +bool test_shared_mmap_weights_dedup() { + try { + std::string filename = "test_shared_mmap_dedup.bin"; + { + CactusGraph graph; + size_t in_node = graph.input({2, 4}, Precision::FP16); + std::vector<__fp16> data = {1, 2, 3, 4, 5, 6, 7, 8}; + graph.set_input(in_node, data.data(), Precision::FP16); + graph.execute(); + GraphFile::save_node(graph, in_node, filename); + } + + void* ptr1 = nullptr; + void* ptr2 = nullptr; + + { + CactusGraph graph1; + size_t id1 = graph1.mmap_weights(filename); + graph1.execute(); + ptr1 = graph1.get_output(id1); + + CactusGraph graph2; + size_t id2 = graph2.mmap_weights(filename); + graph2.execute(); + ptr2 = graph2.get_output(id2); + + if (ptr1 == nullptr || ptr2 == nullptr || ptr1 != ptr2) { + std::remove(filename.c_str()); + return false; + } + } + + std::remove(filename.c_str()); + return true; + } catch (const std::exception& e) { + std::cout << "[shared_mmap_weights_dedup] exception: " << e.what() << std::endl; + return false; + } +} + bool run_benchmarks() { const int ITERS = 100; const std::string temp_file = "bench_io_50nodes.cg"; @@ -517,6 +557,7 @@ int main() { test_save_load_preserves_recurrent_cache_persistence()); runner.run_test("Save/Load Preserves KV Cache Num Slots", test_save_load_preserves_kv_cache_num_slots()); + runner.run_test("Shared Mmap Weights Deduplication", test_shared_mmap_weights_dedup()); runner.print_benchmarks_header(); runner.run_bench("benchmarks", run_benchmarks()); runner.print_summary();