GGML map ops proof of concept.

Author: KerfuffleV2

ggml.c

@@ -1951,7 +1951,7 @@ static void ggml_vec_dot_q4_0(const int n, float * restrict s, const void * rest
     // Initialize accumulator with zeros
     __m256 acc = _mm256_setzero_ps();
 
-    /* Prepare the constants we will need during execution */        
+    /* Prepare the constants we will need during execution */
     const __m256i lowMask = _mm256_set1_epi8( 0xF );
     const __m256i offset_8 = _mm256_set1_epi16( 8 );
 
@@ -1962,60 +1962,60 @@ static void ggml_vec_dot_q4_0(const int n, float * restrict s, const void * rest
     // Main loop
     for (int i = 0; i < nb; i+=UNROLL_COUNT) {
 
-        // This loop will be unrolled by the compiler    
+        // This loop will be unrolled by the compiler
         for (int u=0;u<UNROLL_COUNT;u++)  {
-            /* Compute combined scale for the block */ 
-            const __m256 scale = _mm256_mul_ps( 
-                    _mm256_broadcast_ss( &x[i+u].d ), 
-                    _mm256_broadcast_ss( &y[i+u].d ) ); 
-
-            /* get input from x 
-               Input: 32 Nibbles (16 bytes) at *x[i+u] 
-               Output: 2 vectors with 16 values of type int16_t (x_high_q, x_low_q) */             
-                                      
-            /* Load 16 bytes from memory */  
-            const __m128i tmp_x = _mm_loadu_si128( ( const __m128i* ) x[i+u].qs); 
-            /* Expand bytes into uint16_t values */                                 
-            const __m256i bytes_x = _mm256_cvtepu8_epi16(tmp_x); 
+            /* Compute combined scale for the block */
+            const __m256 scale = _mm256_mul_ps(
+                    _mm256_broadcast_ss( &x[i+u].d ),
+                    _mm256_broadcast_ss( &y[i+u].d ) );
+
+            /* get input from x
+               Input: 32 Nibbles (16 bytes) at *x[i+u]
+               Output: 2 vectors with 16 values of type int16_t (x_high_q, x_low_q) */
+
+            /* Load 16 bytes from memory */
+            const __m128i tmp_x = _mm_loadu_si128( ( const __m128i* ) x[i+u].qs);
+            /* Expand bytes into uint16_t values */
+            const __m256i bytes_x = _mm256_cvtepu8_epi16(tmp_x);
             /* Unpack values into individual bytes */
             __m256i x_low_q = _mm256_and_si256( lowMask, bytes_x );
             const __m256i pre_shift_x_high_q = _mm256_andnot_si256( lowMask, bytes_x );
-            __m256i x_high_q = _mm256_srli_epi16( pre_shift_x_high_q, 4 );            
+            __m256i x_high_q = _mm256_srli_epi16( pre_shift_x_high_q, 4 );
             /* Now we have two vectors with bytes in [ 0 .. 15 ] interval.  Offset them into [ -8 .. +7 ] interval.  */
-            x_high_q = _mm256_sub_epi16( x_high_q, offset_8 ); 
-            x_low_q = _mm256_sub_epi16( x_low_q, offset_8 ); 
+            x_high_q = _mm256_sub_epi16( x_high_q, offset_8 );
+            x_low_q = _mm256_sub_epi16( x_low_q, offset_8 );
 
-            /* get input from y 
-               Input: 32 Nibbles (16 bytes) at *y[i+u] 
-               Output: 2 vectors with 16 values of type int16_t (y_high_q, y_low_q) */             
+            /* get input from y
+               Input: 32 Nibbles (16 bytes) at *y[i+u]
+               Output: 2 vectors with 16 values of type int16_t (y_high_q, y_low_q) */
 
-            /* Load 16 bytes from memory */  
-            const __m128i tmp_y = _mm_loadu_si128( (const __m128i* ) y[i+u].qs); 
-            /* Expand bytes into uint16_t values */     
-            const __m256i bytes_y = _mm256_cvtepu8_epi16(tmp_y); 
+            /* Load 16 bytes from memory */
+            const __m128i tmp_y = _mm_loadu_si128( (const __m128i* ) y[i+u].qs);
+            /* Expand bytes into uint16_t values */
+            const __m256i bytes_y = _mm256_cvtepu8_epi16(tmp_y);
             /* Unpack values into individual bytes */
-            const __m256i pre_shift_y_high_q = _mm256_andnot_si256( lowMask, bytes_y ); 
-            __m256i y_high_q = _mm256_srli_epi16( pre_shift_y_high_q, 4 ); 
-            __m256i y_low_q = _mm256_and_si256( lowMask, bytes_y ); 
+            const __m256i pre_shift_y_high_q = _mm256_andnot_si256( lowMask, bytes_y );
+            __m256i y_high_q = _mm256_srli_epi16( pre_shift_y_high_q, 4 );
+            __m256i y_low_q = _mm256_and_si256( lowMask, bytes_y );
             /* Now we have two vectors with bytes in [ 0 .. 15 ] interval.  Offset them into [ -8 .. +7 ] interval.  */
-            y_high_q = _mm256_sub_epi16( y_high_q, offset_8 ); 
-            y_low_q = _mm256_sub_epi16( y_low_q, offset_8 ); 
+            y_high_q = _mm256_sub_epi16( y_high_q, offset_8 );
+            y_low_q = _mm256_sub_epi16( y_low_q, offset_8 );
 
-            /* Compute products of int16_t integers, add pairwise, store as int32_t */     
-            __m256i xy_high_q = _mm256_madd_epi16( x_high_q, y_high_q ); 
-            __m256i xy_low_q = _mm256_madd_epi16( x_low_q, y_low_q ); 
+            /* Compute products of int16_t integers, add pairwise, store as int32_t */
+            __m256i xy_high_q = _mm256_madd_epi16( x_high_q, y_high_q );
+            __m256i xy_low_q = _mm256_madd_epi16( x_low_q, y_low_q );
 
-            /* Accumulate the products of int32_t integers -> we now have a vector of 8 int_32t */ 
-            __m256i xy_q = _mm256_add_epi32( xy_high_q, xy_low_q ); 
+            /* Accumulate the products of int32_t integers -> we now have a vector of 8 int_32t */
+            __m256i xy_q = _mm256_add_epi32( xy_high_q, xy_low_q );
 
-            /* Convert to vectore of 8 int32_t to 8 floats */ 
-            __m256 q = _mm256_cvtepi32_ps( xy_q ); 
+            /* Convert to vectore of 8 int32_t to 8 floats */
+            __m256 q = _mm256_cvtepi32_ps( xy_q );
 
-            /* Multiply q with scale and accumulate */ 
-            acc = _mm256_fmadd_ps( scale, q, acc );    
+            /* Multiply q with scale and accumulate */
+            acc = _mm256_fmadd_ps( scale, q, acc );
         }
-       
-    }   
+
+    }
 
     // Return horizontal sum of the acc vector
     __m128 res = _mm256_extractf128_ps( acc, 1 );
@@ -2631,9 +2631,12 @@ static const char * GGML_OP_LABEL[GGML_OP_COUNT] = {
 
     "FLASH_ATTN",
     "FLASH_FF",
+
+    "MAP_UNARY",
+    "MAP_BINARY",
 };
 
-static_assert(GGML_OP_COUNT == 35, "GGML_OP_COUNT != 35");
+static_assert(GGML_OP_COUNT == 37, "GGML_OP_COUNT != 37");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -2675,9 +2678,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
 
     "flash_attn(x)",
     "flash_ff(x)",
+
+    "f(x)",
+    "f(x,y)",
 };
 
-static_assert(GGML_OP_COUNT == 35, "GGML_OP_COUNT != 35");
+static_assert(GGML_OP_COUNT == 37, "GGML_OP_COUNT != 37");
 
 static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN");
 static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN");
@@ -3589,6 +3595,92 @@ struct ggml_tensor * ggml_dup_inplace(
     return ggml_dup_impl(ctx, a, true);
 }
 
+
+// ggml_map_binary
+
+struct ggml_tensor * ggml_map_binary_impl(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        struct ggml_tensor * b,
+        void (*const fun)(int, float *, float *, float *),
+        bool inplace) {
+    GGML_ASSERT(ggml_are_same_shape(a, b));
+
+    bool is_node = false;
+
+    if (!inplace && (a->grad || b->grad)) {
+        is_node = true;
+    }
+
+    struct ggml_tensor * addr_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sizeof(void *) / sizeof(int32_t));
+    *((void **)addr_tensor->data) = fun;
+    struct ggml_tensor *result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
+
+    result->op = GGML_OP_MAP_BINARY;
+    result->grad = is_node ? ggml_dup_tensor(ctx, result) : NULL;
+    result->src0 = a;
+    result->src1 = b;
+    result->opt[0] = addr_tensor;
+
+    return result;
+}
+
+struct ggml_tensor * ggml_map_binary(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        struct ggml_tensor * b,
+        void (*const fun)(int, float *, float *, float *)) {
+    return ggml_map_binary_impl(ctx, a, b, fun, false);
+}
+
+struct ggml_tensor * ggml_map_binary_inplace(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        struct ggml_tensor * b,
+        void (*const fun)(int, float *, float *, float *)) {
+    return ggml_map_binary_impl(ctx, a, b, fun, true);
+}
+
+// ggml_map_unary
+
+struct ggml_tensor * ggml_map_unary_impl(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        void (*const fun)(int, float *, float *),
+        bool inplace) {
+    bool is_node = false;
+
+    if (!inplace && a->grad) {
+        is_node = true;
+    }
+
+    struct ggml_tensor * addr_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sizeof(void *) / sizeof(int32_t));
+    *((void **)addr_tensor->data) = fun;
+    struct ggml_tensor *result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
+
+    result->op = GGML_OP_MAP_UNARY;
+    result->grad = is_node ? ggml_dup_tensor(ctx, result) : NULL;
+    result->src0 = a;
+    result->opt[0] = addr_tensor;
+
+    return result;
+}
+
+struct ggml_tensor * ggml_map_unary(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        void (*const fun)(int, float *, float *)) {
+    return ggml_map_unary_impl(ctx, a, fun, false);
+}
+
+struct ggml_tensor * ggml_map_unary_inplace(
+        struct ggml_context * ctx,
+        struct ggml_tensor * a,
+        void (*const fun)(int, float *, float *)) {
+    return ggml_map_unary_impl(ctx, a, fun, true);
+}
+
+
 // ggml_add
 
 struct ggml_tensor * ggml_add_impl(
@@ -5034,6 +5126,111 @@ static void ggml_compute_forward_dup(
     }
 }
 
+// ggml_compute_forward_map_unary
+
+static void ggml_compute_forward_map_unary_f32(
+        const struct ggml_compute_params * params,
+        const struct ggml_tensor * src0,
+        struct ggml_tensor * dst,
+        void (*const fun)(int, float *, float *)) {
+    GGML_ASSERT(ggml_are_same_shape(src0, dst));
+
+    if (params->type == GGML_TASK_INIT || params->type == GGML_TASK_FINALIZE) {
+        return;
+    }
+
+    const int n  = ggml_nrows(src0);
+    const int nc = src0->ne[0];
+
+    assert( dst->nb[0] == sizeof(float));
+    assert(src0->nb[0] == sizeof(float));
+
+    for (int i = 0; i < n; i++) {
+        fun(nc,
+                (float *) ((char *) dst->data  + i*( dst->nb[1])),
+                (float *) ((char *) src0->data + i*(src0->nb[1])));
+    }
+}
+
+
+static void ggml_compute_forward_map_unary(
+        const struct ggml_compute_params * params,
+        const struct ggml_tensor * src0,
+        struct ggml_tensor * dst,
+        void (*const fun)(int, float *, float *)) {
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_map_unary_f32(params, src0, dst, fun);
+            } break;
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_I8:
+        case GGML_TYPE_I16:
+        case GGML_TYPE_I32:
+        case GGML_TYPE_F16:
+        case GGML_TYPE_COUNT:
+            {
+                GGML_ASSERT(false);
+            } break;
+    }
+}
+
+// ggml_compute_forward_map_binary
+
+static void ggml_compute_forward_map_binary_f32(
+        const struct ggml_compute_params * params,
+        const struct ggml_tensor * src0,
+        const struct ggml_tensor * src1,
+        struct ggml_tensor * dst,
+        void (*const fun)(int, float *, float *, float *)) {
+    assert(params->ith == 0);
+    assert(ggml_are_same_shape(src0, src1) && ggml_are_same_shape(src0, dst));
+
+    if (params->type == GGML_TASK_INIT || params->type == GGML_TASK_FINALIZE) {
+        return;
+    }
+
+    const int n  = ggml_nrows(src0);
+    const int nc = src0->ne[0];
+
+    assert( dst->nb[0] == sizeof(float));
+    assert(src0->nb[0] == sizeof(float));
+    assert(src1->nb[0] == sizeof(float));
+
+    for (int i = 0; i < n; i++) {
+        fun(nc,
+                (float *) ((char *) dst->data  + i*( dst->nb[1])),
+                (float *) ((char *) src0->data + i*(src0->nb[1])),
+                (float *) ((char *) src1->data + i*(src1->nb[1])));
+    }
+}
+
+
+static void ggml_compute_forward_map_binary(
+        const struct ggml_compute_params * params,
+        const struct ggml_tensor * src0,
+        const struct ggml_tensor * src1,
+        struct ggml_tensor * dst,
+        void (*const fun)(int, float *, float *, float *)) {
+    switch (src0->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_map_binary_f32(params, src0, src1, dst, fun);
+            } break;
+        case GGML_TYPE_Q4_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_I8:
+        case GGML_TYPE_I16:
+        case GGML_TYPE_I32:
+        case GGML_TYPE_F16:
+        case GGML_TYPE_COUNT:
+            {
+                GGML_ASSERT(false);
+            } break;
+    }
+}
+
 // ggml_compute_forward_add
 
 static void ggml_compute_forward_add_f32(
@@ -8567,7 +8764,19 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_dup(params, tensor->src0, tensor);
             } break;
-        case GGML_OP_ADD:
+        case GGML_OP_MAP_UNARY:
+            {
+                void (*const fun)(int, float *, float *) = *((void **)tensor->opt[0]->data);
+                ggml_compute_forward_map_unary(params, tensor->src0, tensor, fun);
+            }
+            break;
+        case GGML_OP_MAP_BINARY:
+            {
+                void (*const fun)(int, float *, float *, float *) = *((void **)tensor->opt[0]->data);
+                ggml_compute_forward_map_binary(params, tensor->src0, tensor->src1, tensor, fun);
+            }
+            break;
+            case GGML_OP_ADD:
             {
                 ggml_compute_forward_add(params, tensor->src0, tensor->src1, tensor);
             } break;

ggml.h

@@ -250,6 +250,9 @@ enum ggml_op {
     GGML_OP_FLASH_ATTN,
     GGML_OP_FLASH_FF,
 
+    GGML_OP_MAP_UNARY,
+    GGML_OP_MAP_BINARY,
+
     GGML_OP_COUNT,
 };
 
@@ -416,6 +419,17 @@ struct ggml_tensor * ggml_dup(
         struct ggml_context * ctx,
         struct ggml_tensor  * a);
 
+struct ggml_tensor *ggml_map_unary(
+    struct ggml_context *ctx,
+    struct ggml_tensor *a,
+    void (*const fun)(int, float *, float *));
+
+struct ggml_tensor *ggml_map_binary(
+    struct ggml_context *ctx,
+    struct ggml_tensor *a,
+    struct ggml_tensor *b,
+    void (*const fun)(int, float *, float *, float *));
+
 struct ggml_tensor * ggml_add(
         struct ggml_context * ctx,
         struct ggml_tensor  * a,