Test ggml compute chunks (ggml-org#16)

* chunked RMS and mulmat for testing
* linux compilation fix - not super clean

Author: cmp-nct

ggml.c

@@ -44,9 +44,10 @@
 #if defined(_WIN32)
 
 #include <windows.h>
-
+#ifndef atomic_int
 typedef volatile LONG atomic_int;
 typedef atomic_int atomic_bool;
+#endif
 
 static void atomic_store(atomic_int* ptr, LONG val) {
     InterlockedExchange(ptr, val);
@@ -3761,9 +3762,6 @@ struct ggml_context_container {
     struct ggml_context context;
 };
 
-//
-// ggml state
-//
 
 struct ggml_state {
     struct ggml_context_container contexts[GGML_MAX_CONTEXTS];
@@ -9939,18 +9937,20 @@ static void ggml_compute_forward_rms_norm_f32(
     GGML_ASSERT(ggml_are_same_shape(src0, dst));
 
     if (params->type == GGML_TASK_INIT || params->type == GGML_TASK_FINALIZE) {
+        atomic_store(params->aic, 0);
+
         return;
     }
 
     GGML_ASSERT(src0->nb[0] == sizeof(float));
 
-    const int ith = params->ith;
+    const int ith = params->ith; UNUSED(ith);
     const int nth = params->nth;
 
     const int64_t ne00 = src0->ne[0];
     const int64_t ne01 = src0->ne[1];
     const int64_t ne02 = src0->ne[2];
-    const int64_t ne03 = src0->ne[3];
+    const int64_t ne03 = src0->ne[3]; UNUSED(ne03);
 
     const size_t nb01 = src0->nb[1];
     const size_t nb02 = src0->nb[2];
@@ -9962,30 +9962,45 @@ static void ggml_compute_forward_rms_norm_f32(
 
     const float eps = 1e-6f; // TODO: make this a parameter
 
-    // TODO: optimize
-    for (int64_t i03 = 0; i03 < ne03; i03++) {
-        for (int64_t i02 = 0; i02 < ne02; i02++) {
-            for (int64_t i01 = ith; i01 < ne01; i01 += nth) {
-                const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
-
-                ggml_float sum = 0.0;
-                for (int64_t i00 = 0; i00 < ne00; i00++) {
-                    sum += (ggml_float)(x[i00] * x[i00]);
-                }
+    const int nr = ggml_nrows(src0);
+    const int dr = (nr + 8*nth - 1)/(8*nth);
 
-                const float mean = sum/ne00;
+    while (true) {
+        const int ir0 = atomic_fetch_add(params->aic, dr);
 
-                float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
+        for (int ir = ir0; ir < ir0 + dr; ++ir) {
+            if (ir >= nr) {
+                break;
+            }
 
-                memcpy(y, x, ne00 * sizeof(float));
-                // for (int i00 = 0; i00 < ne00; i00++) {
-                //     y[i00] = x[i00];
-                // }
+            // src0 indices
+            const int i03 = ir/(ne02*ne01);
+            const int i02 = (ir - i03*ne02*ne01)/ne01;
+            const int i01 = (ir - i03*ne02*ne01 - i02*ne01);
 
-                const float scale = 1.0f/sqrtf(mean + eps);
+            const float * x = (float *) ((char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03);
 
-                ggml_vec_scale_f32(ne00, y, scale);
+            ggml_float sum = 0.0;
+            for (int64_t i00 = 0; i00 < ne00; i00++) {
+                sum += (ggml_float)(x[i00] * x[i00]);
             }
+
+            float mean = sum/ne00;
+
+            float * y = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3);
+
+            memcpy(y, x, ne00 * sizeof(float));
+            // for (int i00 = 0; i00 < ne00; i00++) {
+            //     y[i00] = x[i00];
+            // }
+
+            const float scale = 1.0f/sqrtf(mean + eps);
+
+            ggml_vec_scale_f32(ne00, y, scale);
+        }
+
+        if (ir0 + dr >= nr) {
+            break;
         }
     }
 }
@@ -10635,7 +10650,7 @@ static void ggml_compute_forward_mul_mat_q_f32(
     const int nb2  = dst->nb[2];
     const int nb3  = dst->nb[3];
 
-    const int ith = params->ith;
+    const int ith = params->ith; UNUSED(ith);
     const int nth = params->nth;
 
     GGML_ASSERT(ne02 == ne12);
@@ -10737,50 +10752,57 @@ static void ggml_compute_forward_mul_mat_q_f32(
             }
         }
 
+        atomic_store(params->aic, 0);
+
         return;
     }
 
     if (params->type == GGML_TASK_FINALIZE) {
         return;
     }
 
+    void * wdata = params->wdata;
+    const size_t row_size = ne00*GGML_TYPE_SIZE[vec_dot_type]/GGML_BLCK_SIZE[vec_dot_type];
+
     // parallelize by src0 rows using ggml_vec_dot_q
 
-    // total rows in src0
-    const int nr = ne01*ne02*ne03;
+    const int nr = ggml_nrows(src0); // nr stands for number of rows
+    const int dr = (nr + 8*nth - 1)/(8*nth); // dr stands for delta rows
 
-    // rows per thread
-    const int dr = (nr + nth - 1)/nth;
+    while (true) {
+        const int ir0 = atomic_fetch_add(params->aic, dr);
 
-    // row range for this thread
-    const int ir0 = dr*ith;
-    const int ir1 = MIN(ir0 + dr, nr);
+        for (int ir = ir0; ir < ir0 + dr; ++ir) {
+            if (ir >= nr) {
+                break;
+            }
 
-    void * wdata = params->wdata;
-    const size_t row_size = ne00*GGML_TYPE_SIZE[vec_dot_type]/GGML_BLCK_SIZE[vec_dot_type];
+            // src0 indices
+            const int i03 = ir/(ne02*ne01);
+            const int i02 = (ir - i03*ne02*ne01)/ne01;
+            const int i01 = (ir - i03*ne02*ne01 - i02*ne01);
 
-    for (int ir = ir0; ir < ir1; ++ir) {
-        // src0 indices
-        const int i03 = ir/(ne02*ne01);
-        const int i02 = (ir - i03*ne02*ne01)/ne01;
-        const int i01 = (ir - i03*ne02*ne01 - i02*ne01);
+            const int i13 = i03;
+            const int i12 = i02;
 
-        const int i13 = i03;
-        const int i12 = i02;
+            const int i0 = i01;
+            const int i2 = i02;
+            const int i3 = i03;
 
-        const int i0 = i01;
-        const int i2 = i02;
-        const int i3 = i03;
+            void * src0_row = (void *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
+            char * src1_col =          ((char *)      wdata + (      (0 + i12*ne11 + i13*ne12*ne11)*row_size));
 
-        void * src0_row = (void *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
-        char * src1_col =          ((char *)      wdata + (      (0 + i12*ne11 + i13*ne12*ne11)*row_size));
+            float * dst_col = (float *) ((char *) dst->data + (i0*nb0 + 0*nb1 + i2*nb2 + i3*nb3));
 
-        float * dst_col = (float *) ((char *) dst->data + (i0*nb0 + 0*nb1 + i2*nb2 + i3*nb3));
+            assert(ne00 % 32 == 0);
 
-        assert(ne00 % 32 == 0);
+            for (int64_t ic = 0; ic < ne11; ++ic) {
+                vec_dot_q(ne00, &dst_col[ic*ne0], src0_row, (void *) (src1_col + ic*row_size));
+            }
+        }
 
-        for (int64_t ic = 0; ic < ne11; ++ic) {
-            vec_dot_q(ne00, &dst_col[ic*ne0], src0_row, (void *) (src1_col + ic*row_size));
+        if (ir0 + dr >= nr) {
+            break;
         }
     }
 
@@ -15879,6 +15901,7 @@ struct ggml_compute_state_shared {
 
     // synchronization primitives
     atomic_int  n_ready;
+    atomic_int  aic;
     atomic_bool has_work;
     atomic_bool stop; // stop all threads
 };
@@ -15947,6 +15970,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
         /*.spin      =*/ GGML_LOCK_INITIALIZER,
         /*.n_threads =*/ n_threads,
         /*.n_ready   =*/ 0,
+        /*.aic       =*/ 0,
         /*.has_work  =*/ false,
         /*.stop      =*/ false,
     };
@@ -15967,6 +15991,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
                     .nth   = n_threads,
                     .wsize = cgraph->work ? ggml_nbytes(cgraph->work) : 0,
                     .wdata = cgraph->work ? cgraph->work->data : NULL,
+                    .aic   = &state_shared.aic,
                 },
                 .node   = NULL,
                 .shared = &state_shared,
@@ -16308,6 +16333,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
             /*.nth   =*/ node->n_tasks,
             /*.wsize =*/ cgraph->work ? ggml_nbytes(cgraph->work) : 0,
             /*.wdata =*/ cgraph->work ? cgraph->work->data : NULL,
+            /*.aic   =*/ &state_shared.aic,
         };
 
         ggml_compute_forward(&params, node);
@@ -16331,6 +16357,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
                     .nth   = node->n_tasks,
                     .wsize = cgraph->work ? ggml_nbytes(cgraph->work) : 0,
                     .wdata = cgraph->work ? cgraph->work->data : NULL,
+                    .aic   = &state_shared.aic,
                 };
                 workers[j].node = node;
             }
@@ -16346,6 +16373,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
         }
 
         params.type = GGML_TASK_COMPUTE;
+        params.aic  = &state_shared.aic;
         ggml_compute_forward(&params, node);
 
         // wait for thread pool
@@ -16386,6 +16414,7 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
                     .nth   = node->n_tasks,
                     .wsize = cgraph->work ? ggml_nbytes(cgraph->work) : 0,
                     .wdata = cgraph->work ? cgraph->work->data : NULL,
+                    .aic   = &state_shared.aic,
                 };
                 workers[j].node = node;
             }

ggml.h

@@ -220,6 +220,18 @@ extern "C" {
     typedef uint16_t ggml_fp16_t;
 #endif
 
+// we need atomic support for the param struct - this needs to be improved
+#if defined(_WIN32) 
+    typedef volatile long atomic_int;
+    typedef atomic_int atomic_bool;
+#else
+    #ifndef __cplusplus
+    #include <stdatomic.h>
+    #else
+    typedef volatile long atomic_int;
+    typedef atomic_int atomic_bool;
+    #endif
+#endif
     // convert FP16 <-> FP32
     GGML_API float       ggml_fp16_to_fp32(ggml_fp16_t x);
     GGML_API ggml_fp16_t ggml_fp32_to_fp16(float x);
@@ -454,8 +466,12 @@ extern "C" {
         // work buffer for all threads
         size_t wsize;
         void * wdata;
+
+        // atomic counter used to distribute chunks of work
+        atomic_int * aic;
     };
 
+
     // misc
 
     GGML_API void    ggml_time_init(void); // call this once at the beginning of the program

libfalcon.cpp

@@ -45,6 +45,11 @@
 #include <sstream>
 #include <numeric>
 
+#if defined(_MSC_VER)
+// disable "possible loss of data" 
+#pragma warning(disable: 4244 4267)
+#endif
+
 #define LLAMA_USE_SCRATCH
 #define LLAMA_MAX_SCRATCH_BUFFERS 16