feat: add imatrix support (#633)

src/convert.cpp

@@ -76,29 +76,22 @@ static bool load_tensors_for_export(ModelLoader& model_loader,
     return success;
 }
 
-bool convert(const char* input_path,
-             const char* vae_path,
-             const char* output_path,
-             sd_type_t output_type,
-             const char* tensor_type_rules,
-             bool convert_name) {
-    ModelLoader model_loader;
-
-    if (!model_loader.init_from_file(input_path)) {
-        LOG_ERROR("init model loader from file failed: '%s'", input_path);
+static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
+    if (path == nullptr || strlen(path) == 0) {
+        return true;
+    }
+    if (!model_loader.init_from_file(path, prefix)) {
+        LOG_ERROR("init model loader from file failed: '%s'", path);
         return false;
     }
+    loaded_any = true;
+    return true;
+}
 
-    if (vae_path != nullptr && strlen(vae_path) > 0) {
-        if (!model_loader.init_from_file(vae_path, "vae.")) {
-            LOG_ERROR("init model loader from file failed: '%s'", vae_path);
-            return false;
-        }
-    }
-    if (convert_name) {
-        model_loader.convert_tensors_name();
-    }
-
+static bool export_loaded_model(ModelLoader& model_loader,
+                                const char* output_path,
+                                sd_type_t output_type,
+                                const char* tensor_type_rules) {
     ggml_type type             = sd_type_to_ggml_type(output_type);
     bool output_is_safetensors = ends_with(output_path, ".safetensors");
     TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
@@ -136,3 +129,55 @@ bool convert(const char* input_path,
     ggml_free(ggml_ctx);
     return success;
 }
+
+bool convert_with_components(const char* model_path,
+                             const char* clip_l_path,
+                             const char* clip_g_path,
+                             const char* t5xxl_path,
+                             const char* diffusion_model_path,
+                             const char* vae_path,
+                             const char* output_path,
+                             sd_type_t output_type,
+                             const char* tensor_type_rules,
+                             bool convert_name) {
+    ModelLoader model_loader;
+    bool loaded_any = false;
+
+    if (!init_convert_path(model_loader, model_path, "", loaded_any) ||
+        !init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) ||
+        !init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) ||
+        !init_convert_path(model_loader, vae_path, "vae.", loaded_any)) {
+        return false;
+    }
+
+    if (!loaded_any) {
+        LOG_ERROR("no input model path provided for convert");
+        return false;
+    }
+
+    if (convert_name) {
+        model_loader.convert_tensors_name();
+    }
+
+    return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
+}
+
+bool convert(const char* input_path,
+             const char* vae_path,
+             const char* output_path,
+             sd_type_t output_type,
+             const char* tensor_type_rules,
+             bool convert_name) {
+    return convert_with_components(input_path,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   nullptr,
+                                   vae_path,
+                                   output_path,
+                                   output_type,
+                                   tensor_type_rules,
+                                   convert_name);
+}

src/core/ggml_extend.hpp

@@ -2473,7 +2473,10 @@ protected:
             sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
 
-        ggml_status status = ggml_backend_graph_compute(runtime_backend, gf);
+        ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
+                                                                         gf,
+                                                                         sd_get_backend_eval_callback(),
+                                                                         sd_get_backend_eval_callback_data());
         if (status != GGML_STATUS_SUCCESS) {
             LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
             return std::nullopt;

src/core/ggml_extend_backend.cpp

@@ -9,6 +9,7 @@
 #include <vector>
 
 #include "core/util.h"
+#include "ggml/src/ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -364,6 +365,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
     return false;
 }
 
+static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
+    ggml_cgraph cgraph = {
+        /*.size             =*/0,
+        /*.n_nodes          =*/i1 - i0,
+        /*.n_leafs          =*/0,
+        /*.nodes            =*/cgraph0->nodes + i0,
+        /*.grads            =*/nullptr,
+        /*.grad_accs        =*/nullptr,
+        /*.leafs            =*/nullptr,
+        /*.use_counts       =*/cgraph0->use_counts,
+        /*.visited_hash_set =*/cgraph0->visited_hash_set,
+        /*.order            =*/cgraph0->order,
+        /*.uid              =*/0,
+    };
+    return cgraph;
+}
+
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data) {
+    if (callback_eval == nullptr) {
+        return ggml_backend_graph_compute(backend, gf);
+    }
+
+    ggml_status status = GGML_STATUS_SUCCESS;
+    const int n_nodes  = ggml_graph_n_nodes(gf);
+    bool stopped       = false;
+
+    for (int j0 = 0; j0 < n_nodes; ++j0) {
+        ggml_tensor* t = ggml_graph_node(gf, j0);
+        bool need      = callback_eval(t, true, callback_eval_user_data);
+        int j1         = j0;
+
+        while (!need && j1 < n_nodes - 1) {
+            t    = ggml_graph_node(gf, ++j1);
+            need = callback_eval(t, true, callback_eval_user_data);
+        }
+
+        ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
+        status         = ggml_backend_graph_compute_async(backend, &gv);
+        if (status != GGML_STATUS_SUCCESS) {
+            break;
+        }
+
+        ggml_backend_synchronize(backend);
+
+        if (need && !callback_eval(t, false, callback_eval_user_data)) {
+            stopped = true;
+            break;
+        }
+
+        j0 = j1;
+    }
+
+    ggml_backend_synchronize(backend);
+    if (stopped && status == GGML_STATUS_SUCCESS) {
+        status = GGML_STATUS_ABORTED;
+    }
+    return status;
+}
+
 const char* sd_get_system_info() {
     static std::string cache_info = []() -> std::string {
         ggml_backend_load_all_once();

src/core/ggml_extend_backend.h

@@ -9,6 +9,7 @@
 
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "stable-diffusion.h"
 
 enum class SDBackendModule {
     DIFFUSION,
@@ -71,6 +72,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
+ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
+                                                        ggml_cgraph* gf,
+                                                        sd_graph_eval_callback_t callback_eval,
+                                                        void* callback_eval_user_data);
 std::string sd_backend_resolve_name(const std::string& name);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);

src/core/util.cpp

@@ -346,6 +346,9 @@ int sd_preview_interval              = 1;
 bool sd_preview_denoised             = true;
 bool sd_preview_noisy                = false;
 
+static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
+static void* sd_backend_eval_cb_data               = nullptr;
+
 std::u32string utf8_to_utf32(const std::string& utf8_str) {
     std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
     return converter.from_bytes(utf8_str);
@@ -629,6 +632,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b
     sd_preview_noisy    = noisy;
 }
 
+void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
+    sd_backend_eval_cb      = cb;
+    sd_backend_eval_cb_data = data;
+}
+
 sd_preview_cb_t sd_get_preview_callback() {
     return sd_preview_cb;
 }
@@ -649,6 +657,14 @@ bool sd_should_preview_noisy() {
     return sd_preview_noisy;
 }
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback() {
+    return sd_backend_eval_cb;
+}
+
+void* sd_get_backend_eval_callback_data() {
+    return sd_backend_eval_cb_data;
+}
+
 sd_progress_cb_t sd_get_progress_callback() {
     return sd_progress_cb;
 }

src/core/util.h

@@ -98,6 +98,9 @@ int sd_get_preview_interval();
 bool sd_should_preview_denoised();
 bool sd_should_preview_noisy();
 
+sd_graph_eval_callback_t sd_get_backend_eval_callback();
+void* sd_get_backend_eval_callback_data();
+
 // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 

src/model_loader.cpp

@@ -20,6 +20,7 @@
 #include "model_io/torch_legacy_io.h"
 #include "model_io/torch_zip_io.h"
 #include "model_loader.h"
+#include "runtime/imatrix.h"
 #include "stable-diffusion.h"
 
 #include "core/ggml_extend_backend.h"
@@ -156,7 +157,8 @@ void convert_tensor(void* src,
                     void* dst,
                     ggml_type dst_type,
                     int nrows,
-                    int n_per_row) {
+                    int n_per_row,
+                    std::vector<float> imatrix = {}) {
     int n = nrows * n_per_row;
     if (src_type == dst_type) {
         size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
@@ -165,7 +167,7 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
         } else {
-            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            imatrix.resize(n_per_row, 1.0f);
             const float* im = imatrix.data();
             ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im);
         }
@@ -195,7 +197,7 @@ void convert_tensor(void* src,
         if (dst_type == GGML_TYPE_F16) {
             ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
         } else {
-            std::vector<float> imatrix(n_per_row, 1.0f);  // dummy importance matrix
+            imatrix.resize(n_per_row, 1.0f);
             const float* im = imatrix.data();
             ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im);
         }
@@ -970,6 +972,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
     size_t total_tensors_processed = 0;
     const int64_t t_start          = start_time;
     int last_n_threads             = 1;
+    SDVersion imatrix_version      = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
 
     for (size_t file_index = 0; file_index < file_data.size(); ++file_index) {
         auto& fdata                  = file_data[file_index];
@@ -1154,12 +1157,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                             failed = true;
                             return;
                         }
+                        std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version);
+                        std::vector<float> imatrix = get_imatrix_collector().get_values(processed_name);
                         convert_tensor((void*)target_buf,
                                        tensor_storage.type,
                                        convert_buf,
                                        dst_tensor->type,
                                        (int)tensor_storage.nelements() / (int)tensor_storage.ne[0],
-                                       (int)tensor_storage.ne[0]);
+                                       (int)tensor_storage.ne[0],
+                                       std::move(imatrix));
                     } else {
                         convert_buf = read_buf;
                     }

src/runtime/imatrix.cpp

@@ -0,0 +1,308 @@
+#include "runtime/imatrix.h"
+
+/* Adapted from llama.cpp (credits: Kawrakow). */
+
+#include "core/util.h"
+#include "ggml-backend.h"
+#include "ggml.h"
+#include "stable-diffusion.h"
+
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+
+static IMatrixCollector imatrix_collector;
+
+IMatrixCollector& get_imatrix_collector() {
+    return imatrix_collector;
+}
+
+// remove any prefix and suffixes from the name
+// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight
+static std::string filter_tensor_name(const char* name) {
+    std::string wname;
+    const char* p = strchr(name, '#');
+    if (p != NULL) {
+        p             = p + 1;
+        const char* q = strchr(p, '#');
+        if (q != NULL) {
+            wname = std::string(p, q - p);
+        } else {
+            wname = p;
+        }
+    } else {
+        wname = name;
+    }
+    return wname;
+}
+
+bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
+    GGML_UNUSED(user_data);
+    if (t == nullptr) {
+        return false;
+    }
+    if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) {
+        return false;
+    }
+
+    const struct ggml_tensor* src0 = t->src[0];
+    const struct ggml_tensor* src1 = t->src[1];
+    if (src0 == nullptr || src1 == nullptr) {
+        return false;
+    }
+    std::string wname = filter_tensor_name(src0->name);
+
+    // when ask is true, the scheduler wants to know if we are interested in data from this tensor
+    // if we return true, a follow-up call will be made with ask=false in which we can do the actual collection
+    if (ask) {
+        if (t->op == GGML_OP_MUL_MAT_ID) {
+            return true;  // collect all indirect matrix multiplications
+        }
+        // why are small batches ignored (<16 tokens)?
+        // if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false;
+        if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) {
+            return false;
+        }
+        return true;
+    }
+    std::lock_guard<std::mutex> lock(mutex_);
+
+    // copy the data from the GPU memory if needed
+    const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer);
+
+    if (!is_host) {
+        src1_data_.resize(ggml_nelements(src1));
+        ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1));
+    }
+
+    const float* data = is_host ? (const float*)src1->data : src1_data_.data();
+
+    // this has been adapted to the new format of storing merged experts in a single 3d tensor
+    // ref: https://github.com/ggml-org/llama.cpp/pull/6387
+    if (t->op == GGML_OP_MUL_MAT_ID) {
+        //   ids  -> [n_experts_used, n_tokens]
+        //   src1 -> [cols, n_expert_used, n_tokens]
+        const ggml_tensor* ids = t->src[2];
+        const int n_as         = static_cast<int>(src0->ne[2]);
+        const int n_ids        = static_cast<int>(ids->ne[0]);
+
+        // the top-k selected expert ids are stored in the ids tensor
+        // for simplicity, always copy ids to host, because it is small
+        // take into account that ids is not contiguous!
+
+        GGML_ASSERT(ids->ne[1] == src1->ne[2]);
+
+        ids_.resize(ggml_nbytes(ids));
+        ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids));
+
+        auto& e = stats_[wname];
+
+        ++e.ncall;
+
+        if (e.values.empty()) {
+            e.values.resize(src1->ne[0] * n_as, 0);
+            e.counts.resize(src1->ne[0] * n_as, 0);
+        } else if (e.values.size() != (size_t)src1->ne[0] * n_as) {
+            LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as);
+            exit(1);  // GGML_ABORT("fatal error");
+        }
+        // loop over all possible experts, regardless if they are used or not in the batch
+        for (int ex = 0; ex < n_as; ++ex) {
+            size_t e_start = ex * src1->ne[0];
+
+            for (int idx = 0; idx < n_ids; ++idx) {
+                for (int row = 0; row < (int)src1->ne[2]; ++row) {
+                    const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]);
+
+                    GGML_ASSERT(excur >= 0 && excur < n_as);  // sanity check
+
+                    if (excur != ex)
+                        continue;
+
+                    const int64_t i11 = idx % src1->ne[1];
+                    const int64_t i12 = row;
+                    const float* x    = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]);
+
+                    for (int j = 0; j < (int)src1->ne[0]; ++j) {
+                        e.values[e_start + j] += x[j] * x[j];
+                        e.counts[e_start + j]++;
+                        if (!std::isfinite(e.values[e_start + j])) {
+                            LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str());
+                            exit(1);
+                        }
+                    }
+                }
+            }
+        }
+    } else {
+        auto& e = stats_[wname];
+        if (e.values.empty()) {
+            e.values.resize(src1->ne[0], 0);
+            e.counts.resize(src1->ne[0], 0);
+        } else if (e.values.size() != (size_t)src1->ne[0]) {
+            LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]);
+            exit(1);  // GGML_ABORT("fatal error");
+        }
+
+        ++e.ncall;
+        for (int row = 0; row < (int)src1->ne[1]; ++row) {
+            const float* x = data + row * src1->ne[0];
+            for (int j = 0; j < (int)src1->ne[0]; ++j) {
+                if (std::isfinite(x[j])) {
+                    e.values[j] += x[j] * x[j];
+                    e.counts[j]++;
+                    if (!std::isfinite(e.values[j])) {
+                        LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str());
+                        exit(1);
+                    }
+                } else {
+                    // Likely something from an attention mask?
+                }
+            }
+        }
+    }
+    return true;
+}
+
+bool load_imatrix(const char* imatrix_path) {
+    return imatrix_collector.load_imatrix(imatrix_path);
+}
+
+void save_imatrix(const char* imatrix_path) {
+    imatrix_collector.save_imatrix(imatrix_path);
+}
+
+static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
+    return imatrix_collector.collect_imatrix(t, ask, user_data);
+}
+
+void enable_imatrix_collection() {
+    sd_set_backend_eval_callback(collect_imatrix, nullptr);
+}
+
+void disable_imatrix_collection() {
+    sd_set_backend_eval_callback(nullptr, nullptr);
+}
+
+void IMatrixCollector::save_imatrix(std::string fname, int ncall) const {
+    if (ncall > 0) {
+        fname += ".at_";
+        fname += std::to_string(ncall);
+    }
+    // avoid writing imatrix entries that do not have full data
+    // this can happen with MoE models where some of the experts end up not being exercised by the provided training data
+
+    int n_entries = 0;
+    std::vector<std::string> to_store;
+
+    for (const auto& kv : stats_) {
+        const int n_all = static_cast<int>(kv.second.counts.size());
+
+        if (n_all == 0) {
+            continue;
+        }
+
+        int n_zeros = 0;
+        for (const int c : kv.second.counts) {
+            if (c == 0) {
+                n_zeros++;
+            }
+        }
+
+        if (n_zeros == n_all) {
+            LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str());
+            continue;
+        }
+
+        if (n_zeros > 0) {
+            LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all);
+            continue;
+        }
+
+        n_entries++;
+        to_store.push_back(kv.first);
+    }
+
+    if (to_store.size() < stats_.size()) {
+        LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size());
+    }
+
+    std::ofstream out(fname, std::ios::binary);
+    out.write((const char*)&n_entries, sizeof(n_entries));
+    for (const auto& name : to_store) {
+        const auto& stat = stats_.at(name);
+        int len          = static_cast<int>(name.size());
+        out.write((const char*)&len, sizeof(len));
+        out.write(name.c_str(), len);
+        out.write((const char*)&stat.ncall, sizeof(stat.ncall));
+        int nval = static_cast<int>(stat.values.size());
+        out.write((const char*)&nval, sizeof(nval));
+        if (nval > 0) {
+            std::vector<float> tmp(nval);
+            for (int i = 0; i < nval; i++) {
+                tmp[i] = (stat.values[i] / static_cast<float>(stat.counts[i])) * static_cast<float>(stat.ncall);
+            }
+            out.write((const char*)tmp.data(), nval * sizeof(float));
+        }
+    }
+
+    // Write the number of call the matrix was computed with
+    out.write((const char*)&last_call_, sizeof(last_call_));
+}
+
+bool IMatrixCollector::load_imatrix(const char* fname) {
+    std::ifstream in(fname, std::ios::binary);
+    if (!in) {
+        LOG_ERROR("failed to open %s\n", fname);
+        return false;
+    }
+    int n_entries;
+    in.read((char*)&n_entries, sizeof(n_entries));
+    if (in.fail() || n_entries < 1) {
+        LOG_ERROR("no data in file %s\n", fname);
+        return false;
+    }
+    for (int i = 0; i < n_entries; ++i) {
+        int len;
+        in.read((char*)&len, sizeof(len));
+        std::vector<char> name_as_vec(len + 1);
+        in.read((char*)name_as_vec.data(), len);
+        if (in.fail()) {
+            LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname);
+            return false;
+        }
+        name_as_vec[len] = 0;
+        std::string name{name_as_vec.data()};
+        auto& e = stats_[std::move(name)];
+        int ncall;
+        in.read((char*)&ncall, sizeof(ncall));
+        int nval;
+        in.read((char*)&nval, sizeof(nval));
+        if (in.fail() || nval < 1) {
+            LOG_ERROR("failed reading number of values for entry %d\n", i);
+            stats_ = {};
+            return false;
+        }
+
+        if (e.values.empty()) {
+            e.values.resize(nval, 0);
+            e.counts.resize(nval, 0);
+        }
+
+        std::vector<float> tmp(nval);
+        in.read((char*)tmp.data(), nval * sizeof(float));
+        if (in.fail()) {
+            LOG_ERROR("failed reading data for entry %d\n", i);
+            stats_ = {};
+            return false;
+        }
+
+        // Recreate the state as expected by save_imatrix(), and correct for weighted sum.
+        for (int i = 0; i < nval; i++) {
+            e.values[i] += tmp[i];
+            e.counts[i] += ncall;
+        }
+        e.ncall += ncall;
+    }
+    return true;
+}

src/runtime/imatrix.h

@@ -0,0 +1,45 @@
+#ifndef __SD_RUNTIME_IMATRIX_H__
+#define __SD_RUNTIME_IMATRIX_H__
+
+#include <fstream>
+#include <mutex>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+/* Adapted from llama.cpp (credits: Kawrakow). */
+
+struct ggml_tensor;
+
+struct IMatrixStats {
+    std::vector<float> values{};
+    std::vector<int> counts{};
+    int ncall = 0;
+};
+
+class IMatrixCollector {
+private:
+    std::unordered_map<std::string, IMatrixStats> stats_ = {};
+    std::mutex mutex_;
+    int last_call_ = 0;
+    std::vector<float> src1_data_;
+    std::vector<char> ids_;  // the expert ids from ggml_mul_mat_id
+
+public:
+    IMatrixCollector() = default;
+    bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data);
+    void save_imatrix(std::string fname, int ncall = -1) const;
+    bool load_imatrix(const char* fname);
+    std::vector<float> get_values(const std::string& key) const {
+        auto it = stats_.find(key);
+        if (it != stats_.end()) {
+            return it->second.values;
+        } else {
+            return {};
+        }
+    }
+};
+
+IMatrixCollector& get_imatrix_collector();
+
+#endif  // __SD_RUNTIME_IMATRIX_H__