mtl : add diag_mask_inf kernel

examples/mtl/mtl.cpp

@@ -41,18 +41,15 @@ int main(int argc, char ** argv) {
 
     // TODO: tmp to match the input used when creating the cgraph
     {
-        const int n_ctx   = 128;
+        const int n_past  = 128;
         const int n_batch = 32;
 
         const std::vector<int> tmp(n_batch, 1); // BOS
 
-        struct ggml_tensor * input = ggml_graph_get_tensor(&gf, "embd");
+        // the actual inference happens here
-        memcpy(input->data, tmp.data(), tmp.size() * sizeof(int));
+        llama_mtl_eval(ctx_mtl, &gf, tmp.data(), tmp.size(), n_past);
     }
 
-    // the actual inference happens here
-    llama_mtl_eval(ctx_mtl, &gf);
-
     llama_mtl_free(ctx_mtl);
 
     ggml_free(ctx_work);

examples/mtl/mtl.h

@@ -20,7 +20,10 @@ void llama_mtl_free(struct ggml_mtl_context * ctx);
 // return 0 on success
 int llama_mtl_eval(
         struct ggml_mtl_context * ctx,
-        struct ggml_cgraph      * gf);
+             struct ggml_cgraph * gf,
+                      const int * tokens,
+                            int   n_tokens,
+                            int   n_past);
 
 #ifdef __cplusplus
 }

examples/mtl/mtl.m

@@ -36,6 +36,9 @@ struct ggml_mtl_context {
     id<MTLFunction>             function_soft_max;
     id<MTLComputePipelineState> pipeline_soft_max;
 
+    id<MTLFunction>             function_diag_mask_inf;
+    id<MTLComputePipelineState> pipeline_diag_mask_inf;
+
     id<MTLFunction>             function_get_rows_q4_0;
     id<MTLComputePipelineState> pipeline_get_rows_q4_0;
 
@@ -150,6 +153,10 @@ struct ggml_mtl_context * llama_mtl_init(
         ctx->pipeline_soft_max = [ctx->device newComputePipelineStateWithFunction:ctx->function_soft_max error:nil];
         fprintf(stderr, "%s: loaded kernel_soft_max: %p\n", __func__, (void *) ctx->pipeline_soft_max);
 
+        ctx->function_diag_mask_inf = [ctx->library newFunctionWithName:@"kernel_diag_mask_inf" constantValues:constants error:nil];
+        ctx->pipeline_diag_mask_inf = [ctx->device newComputePipelineStateWithFunction:ctx->function_diag_mask_inf error:nil];
+        fprintf(stderr, "%s: loaded kernel_diag_mask_inf: %p\n", __func__, (void *) ctx->pipeline_diag_mask_inf);
+
         ctx->function_get_rows_q4_0 = [ctx->library newFunctionWithName:@"kernel_get_rows_q4_0"];
         ctx->pipeline_get_rows_q4_0 = [ctx->device newComputePipelineStateWithFunction:ctx->function_get_rows_q4_0 error:nil];
         fprintf(stderr, "%s: loaded kernel_get_rows_q4_0: %p\n", __func__, (void *) ctx->pipeline_get_rows_q4_0);
@@ -248,8 +255,14 @@ id<MTLBuffer> llama_mtl_get_buffer(struct ggml_mtl_context * ctx, struct ggml_te
 
 int llama_mtl_eval(
         struct ggml_mtl_context * ctx,
-        struct ggml_cgraph      * gf) {
+             struct ggml_cgraph * gf,
-    fprintf(stderr, "%s: evaluating\n", __func__);
+                      const int * tokens,
+                            int   n_tokens,
+                            int   n_past) {
+    fprintf(stderr, "%s: evaluating, n_tokens = %d, n_past = %d\n", __func__, n_tokens, n_past);
+
+    struct ggml_tensor * input = ggml_graph_get_tensor(gf, "embd");
+    memcpy(input->data, tokens, n_tokens * sizeof(int));
 
     id<MTLCommandBuffer> command_buffer  = [ctx->queue commandBuffer];
     id<MTLComputeCommandEncoder> encoder = nil;
@@ -371,6 +384,28 @@ int llama_mtl_eval(
 
                     [encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
                 } break;
+            case GGML_OP_DIAG_MASK_INF:
+                {
+                    if (encoder == nil) {
+                        encoder = [command_buffer computeCommandEncoder];
+                    }
+
+                    id<MTLBuffer> id_src = llama_mtl_get_buffer(ctx, gf->nodes[i]->src0, &offs_src0);
+                    id<MTLBuffer> id_dst = llama_mtl_get_buffer(ctx, gf->nodes[i],       &offs_dst);
+
+                    const int64_t ne00 = gf->nodes[i]->src0->ne[0];
+                    const int64_t ne01 = gf->nodes[i]->src0->ne[1];
+                    const int64_t ne02 = gf->nodes[i]->src0->ne[2];
+
+                    [encoder setComputePipelineState:ctx->pipeline_diag_mask_inf];
+                    [encoder setBuffer:id_src offset:offs_src0 atIndex:0];
+                    [encoder setBuffer:id_dst offset:offs_dst  atIndex:1];
+                    [encoder setBytes:&ne00   length:sizeof(ne00) atIndex:2];
+                    [encoder setBytes:&ne01   length:sizeof(ne01) atIndex:3];
+                    [encoder setBytes:&n_past length:sizeof(int)  atIndex:4];
+
+                    [encoder dispatchThreadgroups:MTLSizeMake(ne00, ne01, ne02) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
+                } break;
             case GGML_OP_MUL_MAT:
                 {
                     id<MTLBuffer> id_src0 = llama_mtl_get_buffer(ctx, gf->nodes[i]->src0, &offs_src0);
@@ -550,7 +585,7 @@ int llama_mtl_eval(
                     const uint64_t nb2 = gf->nodes[i]->nb[2];
                     const uint64_t nb3 = gf->nodes[i]->nb[3];
 
-                    const int n_past = ((int32_t *) gf->nodes[i]->src1->data)[0]; // TODO: TMP !!!!!
+                    //const int n_past = ((int32_t *) gf->nodes[i]->src1->data)[0]; // TODO: TMP !!!!!
                     const int n_dims = ((int32_t *) gf->nodes[i]->src1->data)[1];
                     const int mode   = ((int32_t *) gf->nodes[i]->src1->data)[2];
 
@@ -697,17 +732,15 @@ int llama_mtl_eval(
         if (t->type == GGML_TYPE_F32) {
             const const float * data = (float *) ctx->out.contents;
             printf("data: ");
-            int n = ggml_nelements(t);
+            for (int i = 0; i < (int) t->ne[0]; i++) {
-            if (n > 10) {
-                n = 10;
-            }
-            for (int i = 0; i < n; i++) {
                 printf("%f ", data[i]);
             }
             printf("\n");
             double sum = 0.0;
             for (int i = 0; i < ggml_nelements(t); i++) {
-                sum += data[i];
+                double cur = data[i];
+                if (isinf(cur)) continue;
+                sum += cur;
             }
             printf("sum:  %f\n", sum);
         } else if (t->type == GGML_TYPE_F16) {

examples/mtl/mtl.metal

@@ -86,6 +86,42 @@ kernel void kernel_soft_max(
     }
 }
 
+//const int n  = ggml_nrows(src0);
+//const int nc = src0->ne[0];
+//const int nr = src0->ne[1];
+//const int nz = n/nr;
+//
+//assert( dst->nb[0] == sizeof(float));
+//assert(src0->nb[0] == sizeof(float));
+//
+//for (int k = 0; k < nz; k++) {
+//    for (int j = ith; j < nr; j += nth) {
+//        for (int i = n_past; i < nc; i++) {
+//            if (i > n_past + j) {
+//                *(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value;
+//            }
+//        }
+//    }
+//}
+
+kernel void kernel_diag_mask_inf(
+        device const float * src0,
+        device       float * dst,
+        constant   int64_t & ne00,
+        constant   int64_t & ne01,
+        constant       int & n_past,
+        uint3 tpig[[thread_position_in_grid]]) {
+    const int64_t i02 = tpig[2];
+    const int64_t i01 = tpig[1];
+    const int64_t i00 = tpig[0];
+
+    if (i00 > n_past + i01) {
+        dst[i02*ne01*ne00 + i01*ne00 + i00] = -INFINITY;
+    } else {
+        dst[i02*ne01*ne00 + i01*ne00 + i00] = src0[i02*ne01*ne00 + i01*ne00 + i00];
+    }
+}
+
 kernel void kernel_get_rows_q4_0(
         device const  void * src0,
         device const   int * src1,

ggml.c

@@ -15061,7 +15061,6 @@ struct ggml_cgraph ggml_graph_import(const char * fname, struct ggml_context **
                             memcpy(&offs, args[2]->data, sizeof(offs));
 
                             tensor->data = ((char *) tensor->data) + offs;
-                            printf("xxxxxx offs: %zu\n", offs);
                         } break;
                     case GGML_OP_TRANSPOSE:
                         {

llama.cpp

@@ -1332,14 +1332,14 @@ static bool llama_eval_internal(
             // KQ_scaled shape [n_past + N, N, n_head, 1]
             struct ggml_tensor * KQ_scaled = ggml_scale_inplace(ctx0, KQ, KQ_scale);
             ggml_set_name(KQ_scaled, "KQ_scaled");
-            // TODO: TMP !!!!
-            if (il == 0) {
-                ggml_set_name(KQ_scaled, "mtl-check");
-            }
 
             // KQ_masked = mask_past(KQ_scaled)
             struct ggml_tensor * KQ_masked = ggml_diag_mask_inf_inplace(ctx0, KQ_scaled, n_past);
             ggml_set_name(KQ_masked, "KQ_masked");
+            // TODO: TMP !!!!
+            if (il == 0) {
+                ggml_set_name(KQ_masked, "mtl-check");
+            }
 
             // KQ = soft_max(KQ_masked)
             struct ggml_tensor * KQ_soft_max = ggml_soft_max_inplace(ctx0, KQ_masked);
@@ -1464,12 +1464,14 @@ static bool llama_eval_internal(
         auto print_t_f32 = [&](struct ggml_tensor * t) {
             float * data = (float *)t->data;
             printf("data: ");
-            for (int i = 0; i < std::min((int) t->ne[0], 10); i++) {
+            for (int i = 0; i < (int) t->ne[0]; i++) {
                 printf("%f ", data[i]);
             }
             printf("\n");
             double sum = 0.0;
             for (int i = 0; i < ggml_nelements(t); i++) {
+                double cur = data[i];
+                if (isinf(cur)) continue;
                 sum += data[i];
             }
             printf("sum:  %f\n", sum);