fix(graph): deduplicate shared mmap weight files across component graphs

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<std::unique_ptr<GraphFile::MappedFile>> mapped_files_;
+    std::vector<std::shared_ptr<GraphFile::MappedFile>> mapped_files_;
     std::unordered_map<std::string, size_t> weight_cache_;
     std::unordered_map<size_t, size_t> node_to_mapped_file_;
     std::vector<DebugNodeEntry> 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<MappedFile> get_or_load(const std::string& filename);
+        static void clear();
+    private:
+        static std::mutex mutex_;
+        static std::unordered_map<std::string, std::weak_ptr<MappedFile>> cache_;
+    };
+
     class MappedFile {
     public:
         MappedFile(const std::string& filename);

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<GraphFile::MappedFile>(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<GraphFile::MappedFile>(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<GraphFile::MappedFile>(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<GraphFile::MappedFile>(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<void*>(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<size_t> 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<std::string, std::weak_ptr<MappedFile>> 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<MappedFile> MappedFileRegistry::get_or_load(const std::string& filename) {
+    std::string key = canonical_file_key(filename);
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto it = cache_.find(key);
+    if (it != cache_.end()) {
+        if (auto sp = it->second.lock()) {
+            return sp;
+        }
+    }
+    auto sp = std::make_shared<MappedFile>(filename);
+    cache_[key] = sp;
+    return sp;
+}
+
+void MappedFileRegistry::clear() {
+    std::lock_guard<std::mutex> lock(mutex_);
+    cache_.clear();
+}
+
+} // namespace GraphFile
+
 // MappedFile implementation
 
 MappedFile::MappedFile(const std::string& filename)

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();