metal : refactor diag_mask_inf parameters into a struct

Author: alexju

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -341,4 +341,10 @@ typedef struct {
     uint32_t n_head_log2;
 } ggml_metal_kargs_soft_max;
 
+typedef struct {
+    int64_t  ne00;
+    int64_t  ne01;
+    int      n_past;
+} ggml_metal_kargs_diag_mask_inf;
+
 #endif // GGML_METAL_IMPL

ggml/src/ggml-metal/ggml-metal.m

@@ -2061,13 +2061,16 @@ static void ggml_metal_encode_node(
                     pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_DIAG_MASK_INF].pipeline;
                 }
 
-                // TODO: add ggml_metal_kargs struct
+                ggml_metal_kargs_diag_mask_inf args = {
+                    /*.ne00 =*/ ne00,
+                    /*.ne01 =*/ ne01,
+                    /*.n_past =*/ n_past,
+                };
+
                 [encoder setComputePipelineState:pipeline];
                 [encoder setBuffer:id_src0 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 setBytes:&args  length:sizeof(args) atIndex:2];
 
                 if (ne00%8 == 0) {
                     [encoder dispatchThreadgroups:MTLSizeMake(ne00*ne01*ne02/8, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];

ggml/src/ggml-metal/ggml-metal.metal

@@ -1175,43 +1175,39 @@ template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kerne
 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,
+        constant ggml_metal_kargs_diag_mask_inf & args,
         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;
+    if (i00 > args.n_past + i01) {
+        dst[i02*args.ne01*args.ne00 + i01*args.ne00 + i00] = -INFINITY;
     } else {
-        dst[i02*ne01*ne00 + i01*ne00 + i00] = src0[i02*ne01*ne00 + i01*ne00 + i00];
+        dst[i02*args.ne01*args.ne00 + i01*args.ne00 + i00] = src0[i02*args.ne01*args.ne00 + i01*args.ne00 + i00];
     }
 }
 
 kernel void kernel_diag_mask_inf_8(
         device const float4 * src0,
         device       float4 * dst,
-        constant    int64_t & ne00,
-        constant    int64_t & ne01,
-        constant        int & n_past,
+        constant ggml_metal_kargs_diag_mask_inf & args,
         uint3 tpig[[thread_position_in_grid]]) {
 
     const int64_t i = 2*tpig[0];
 
     dst[i+0] = src0[i+0];
     dst[i+1] = src0[i+1];
     int64_t i4 = 4*i;
-    const int64_t i02 = i4/(ne00*ne01); i4 -= i02*ne00*ne01;
-    const int64_t i01 = i4/(ne00);      i4 -= i01*ne00;
+    const int64_t i02 = i4/(args.ne00*args.ne01); i4 -= i02*args.ne00*args.ne01;
+    const int64_t i01 = i4/(args.ne00);      i4 -= i01*args.ne00;
     const int64_t i00 = i4;
     for (int k = 3; k >= 0; --k) {
-        if (i00 + 4 + k <= n_past + i01) {
+        if (i00 + 4 + k <= args.n_past + i01) {
             break;
         }
         dst[i+1][k] = -INFINITY;
-        if (i00 + k > n_past + i01) {
+        if (i00 + k > args.n_past + i01) {
             dst[i][k] = -INFINITY;
         }
     }