Repository: ggml-org/llama.cpp
Pull request: imatrix: fix oob writes if src1 is not contiguous

The imatrix code implicitly assumes that `src1` is contiguous when copying data from a backend to host memory. As a result the vector to which the data is written can end up being resized to a size that is smaller than the amount of data that `ggml_backend_tensor_get` writes, resulting in out-of-bounds writes.

This PR makes it so that the host buffer always has the exact same size as the amount of data that is being copied. Also, if `src1` is not contiguous, then this is considered for calculating the byte addresses of matrix rows.

Unless I'm misunderstanding the code the cases `ne12 > 1` and `ne13 > 1` are also going to result in unexpected behavior but I don't know what the correct fix would be.

Changed files:
filename: tools/imatrix/imatrix.cpp
patch: @@ -46,7 +46,7 @@ class IMatrixCollector {
     common_params                          m_params;
     std::mutex                             m_mutex;
     int                                    m_last_call = 0;
-    std::vector<float>                     m_src1_data;
+    std::vector<char>                      m_src1_data;
     std::vector<char>                      m_ids; // the expert ids from ggml_mul_mat_id
 };
 
@@ -93,11 +93,13 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
     const bool is_host = ggml_backend_buffer_is_host(src1->buffer);
 
     if (!is_host) {
-        m_src1_data.resize(ggml_nelements(src1));
-        ggml_backend_tensor_get(src1, m_src1_data.data(), 0, ggml_nbytes(src1));
+        const size_t src1_nbytes = ggml_nbytes(src1);
+        m_src1_data.resize(src1_nbytes);
+        ggml_backend_tensor_get(src1, m_src1_data.data(), 0, src1_nbytes);
     }
 
-    const float * data = is_host ? (const float *) src1->data : m_src1_data.data();
+    const char * data = is_host ? (const char *) src1->data : m_src1_data.data();
+    GGML_ASSERT(src1->nb[0] == ggml_element_size(src1));
 
     // 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
@@ -144,7 +146,7 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
 
                     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]);
+                    const float * x = (const float *)(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];
@@ -180,7 +182,7 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void *
         ++e.ncall;
         LOG_DBGV(2, "%s[%d]: %32s, %s, %5d x %5d, %d\n", __func__, m_last_call, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[1], (int)src1->type);
         for (int row = 0; row < (int)src1->ne[1]; ++row) {
-            const float * x = data + row * src1->ne[0];
+            const float * x = (const float *) (data + row * src1->nb[1]);
             for (int j = 0; j < (int)src1->ne[0]; ++j) {
                 e.values[j] += x[j]*x[j];
                 e.counts[j]++;