sd: sync to master-601-eeac950 (#2206)

* sd: sync to master-601-eeac950 * sd: add mmap support
2026-09-18 16:55:14 +02:00 · 2026-05-16 00:23:10 -03:00
parent 3174fccb83
commit f273fd35b9
12 changed files with 514 additions and 108 deletions
@@ -758,16 +758,10 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
    }
 }

-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
-    int64_t process_time_ms = 0;
-    std::atomic<int64_t> read_time_ms(0);
-    std::atomic<int64_t> memcpy_time_ms(0);
-    std::atomic<int64_t> copy_to_backend_time_ms(0);
-    std::atomic<int64_t> convert_time_ms(0);
-    std::atomic<uint64_t> bytes_processed(0);
-
-    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
-    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
+    if (model_files_processed) {
+        return;
+    }

    int64_t start_time = ggml_time_ms();

@@ -779,22 +773,13 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
        processed_tensor_storages.push_back(tensor_storage);
    }

-    process_time_ms = ggml_time_ms() - start_time;
-
-    bool success                          = true;
-    size_t total_tensors_processed        = 0;
-    const size_t total_tensors_to_process = processed_tensor_storages.size();
-    const int64_t t_start                 = ggml_time_ms();
-    int last_n_threads                    = 1;
-
    for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
        std::string file_path = file_paths_[file_index];
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());

-        std::vector<const TensorStorage*> file_tensors;
+        std::vector<TensorStorage> file_tensors;
        for (const auto& ts : processed_tensor_storages) {
            if (ts.file_index == file_index) {
-                file_tensors.push_back(&ts);
+                file_tensors.push_back(ts);
            }
        }
        if (file_tensors.empty()) {
@@ -803,21 +788,169 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread

        bool is_zip = false;
        for (auto const& ts : file_tensors) {
-            if (ts->index_in_zip >= 0) {
+            if (ts.index_in_zip >= 0) {
                is_zip = true;
                break;
            }
        }

-        std::unique_ptr<MmapWrapper> mmapped;
+        ModelFileData fdata = {};
+        fdata.path          = file_path;
+        fdata.is_zip        = is_zip;
+        fdata.tensors       = std::move(file_tensors);
+
        if (enable_mmap && !is_zip) {
            LOG_DEBUG("using mmap for I/O");
-            mmapped = MmapWrapper::create(file_path);
-            if (!mmapped) {
-                LOG_WARN("failed to memory-map '%s'", file_path.c_str());
+            std::unique_ptr<MmapWrapper> mmapped = MmapWrapper::create(file_path, writable_mmap);
+            if (mmapped) {
+                uint8_t* mmap_data             = static_cast<uint8_t*>(mmapped->writable_data());
+                ggml_backend_buffer_t buf_mmap = ggml_backend_cpu_buffer_from_ptr(mmap_data, mmapped->size());
+                if (buf_mmap) {
+                    LOG_INFO("using mmap for '%s'", file_path.c_str());
+                    fdata.mmbuffer = std::shared_ptr<struct ggml_backend_buffer>(buf_mmap, ggml_backend_buffer_free);
+                } else {
+                    LOG_WARN("mmap: failed to create backend buffer for file %s", fdata.path.c_str());
+                }
+                fdata.mmapped = std::shared_ptr<MmapWrapper>(std::move(mmapped));
+            } else {
+                LOG_WARN("failed to memory-map '%s' (falling back to read())", file_path.c_str());
            }
+        } else if (!is_zip) {
+            LOG_INFO("NOT using mmap for '%s' (mmap disabled by caller)",
+                     file_path.c_str());
        }

+        file_data.push_back(std::move(fdata));
+    }
+
+    model_files_processed = true;
+
+    int64_t end_time        = ggml_time_ms();
+    int64_t process_time_ms = end_time - start_time;
+
+    LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f);
+}
+
+std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                                       std::set<std::string> ignore_tensors,
+                                                       bool writable_mmap) {
+    process_model_files(true, writable_mmap);
+
+    std::vector<MmapTensorStore> result;
+    uint64_t mapped_bytes = 0;
+    size_t mapped_tensors = 0;
+
+    LOG_DEBUG("memory-mapping tensors...");
+
+    int64_t t_start = ggml_time_ms();
+
+    for (auto& fdata : file_data) {
+        if (!fdata.mmbuffer)
+            continue;
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        size_t file_mapped_bytes   = 0;
+        size_t file_mapped_tensors = 0;
+
+        for (const auto& tensor_storage : file_tensors) {
+            const std::string& name = tensor_storage.name;
+
+            bool is_ignored = false;
+            for (const auto& ignore_prefix : ignore_tensors) {
+                if (starts_with(name, ignore_prefix)) {
+                    is_ignored = true;
+                    break;
+                }
+            }
+            if (is_ignored)
+                continue;
+
+            auto it = tensors.find(name);
+            if (it == tensors.end())
+                continue;
+
+            ggml_tensor* dst_tensor = it->second;
+            if (dst_tensor == nullptr)
+                continue;
+
+            if (tensor_storage.type != dst_tensor->type)
+                continue;
+
+            size_t tensor_size   = tensor_storage.nbytes();
+            size_t tensor_offset = tensor_storage.offset;
+
+            if (tensor_storage.ne[0] != dst_tensor->ne[0] ||
+                tensor_storage.ne[1] != dst_tensor->ne[1] ||
+                tensor_storage.ne[2] != dst_tensor->ne[2] ||
+                tensor_storage.ne[3] != dst_tensor->ne[3] ||
+                tensor_size != ggml_nbytes(dst_tensor)) {
+                // let load_tensors worry about this
+                continue;
+            }
+
+            ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
+            uint8_t* mmap_data             = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
+            dst_tensor->buffer             = buf_mmap;
+            dst_tensor->data               = mmap_data + tensor_offset;
+
+            file_mapped_bytes += tensor_size;
+            file_mapped_tensors++;
+        }
+
+        if (file_mapped_bytes > 0) {
+            mapped_tensors += file_mapped_tensors;
+            mapped_bytes += file_mapped_bytes;
+            result.push_back({fdata.mmapped, fdata.mmbuffer});
+        }
+    }
+
+    int64_t t_end       = ggml_time_ms();
+    int64_t duration_ms = t_end - t_start;
+
+    LOG_INFO("memory-mapped %zu tensors in %zu files (%.2f MB), taking %.2fs",
+             mapped_tensors,
+             result.size(),
+             mapped_bytes / (1024.0 * 1024.0),
+             duration_ms / 1000.0);
+
+    return result;
+}
+
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+    process_model_files(enable_mmap, false);
+
+    std::atomic<int64_t> read_time_ms(0);
+    std::atomic<int64_t> memcpy_time_ms(0);
+    std::atomic<int64_t> copy_to_backend_time_ms(0);
+    std::atomic<int64_t> convert_time_ms(0);
+    std::atomic<uint64_t> bytes_processed(0);
+
+    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+
+    int64_t start_time = ggml_time_ms();
+
+    size_t total_tensors_to_process = 0;
+    for (const auto& fdata : file_data) {
+        total_tensors_to_process += fdata.tensors.size();
+    }
+
+    bool success                   = true;
+    size_t total_tensors_processed = 0;
+    const int64_t t_start          = start_time;
+    int last_n_threads             = 1;
+
+    for (auto& fdata : file_data) {
+        const std::string& file_path = fdata.path;
+        LOG_DEBUG("loading tensors from %s", file_path.c_str());
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        bool is_zip = fdata.is_zip;
+
+        std::shared_ptr<MmapWrapper> mmapped = fdata.mmapped;
+
        int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
        if (n_threads < 1) {
            n_threads = 1;
@@ -858,7 +991,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                        break;
                    }

-                    const TensorStorage& tensor_storage = *file_tensors[idx];
+                    const TensorStorage& tensor_storage = file_tensors[idx];
                    ggml_tensor* dst_tensor             = nullptr;

                    t0 = ggml_time_ms();
@@ -875,6 +1008,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                        continue;
                    }

+                    // skip mmapped tensors
+                    if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
+                        continue;
+                    }
+
                    size_t nbytes_to_read = tensor_storage.nbytes_to_read();

                    auto read_data = [&](char* buf, size_t n) {
@@ -1018,9 +1156,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
    }

    int64_t end_time = ggml_time_ms();
-    LOG_INFO("loading tensors completed, taking %.2fs (process: %.2fs, read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
+    LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
             (end_time - start_time) / 1000.f,
-             process_time_ms / 1000.f,
             (read_time_ms.load() / (float)last_n_threads) / 1000.f,
             (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
             (convert_time_ms.load() / (float)last_n_threads) / 1000.f,