metal : refactor timestep_embedding parameters into a struct

Author: alexju

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

@@ -485,6 +485,12 @@ typedef struct {
     int32_t  p1;
 } ggml_metal_kargs_pad_reflect_1d;
 
+typedef struct {
+    uint64_t nb1;
+    int      dim;
+    int      max_period;
+} ggml_metal_kargs_timestep_embedding;
+
 typedef struct {
     int64_t  ne0;
     float    start;

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

@@ -3553,13 +3553,16 @@ static void ggml_metal_encode_node(
 
                 id<MTLComputePipelineState> pipeline = ctx->kernels[GGML_METAL_KERNEL_TYPE_TIMESTEP_EMBEDDING_F32].pipeline;
 
-                // TODO: add ggml_metal_kargs struct
+                ggml_metal_kargs_timestep_embedding args = {
+                    /*.nb1 =*/ nb1,
+                    /*.dim =*/ dim,
+                    /*.max_period =*/ max_period
+                };
+
                 [encoder setComputePipelineState:pipeline];
                 [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
                 [encoder setBuffer:id_dst  offset:offs_dst  atIndex:1];
-                [encoder setBytes:&nb1   length:sizeof(nb1) atIndex:2];
-                [encoder setBytes:&dim   length:sizeof(dim) atIndex:3];
-                [encoder setBytes:&max_period length:sizeof(max_period) atIndex:4];
+                [encoder setBytes:&args length:sizeof(args) atIndex:2];
 
                 const int nth = MIN(1024, half);
 

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

@@ -2770,27 +2770,25 @@ kernel void kernel_arange_f32(
 kernel void kernel_timestep_embedding_f32(
     device  const char * src0,
     device        char * dst,
-    constant  uint64_t & nb1,
-    constant  int      & dim,
-    constant  int      & max_period,
+    constant  ggml_metal_kargs_timestep_embedding & args,
     uint3 tgpig[[threadgroup_position_in_grid]],
     uint3 tpitg[[thread_position_in_threadgroup]],
     uint3   ntg[[threads_per_threadgroup]]) {
 
     int i = tgpig.x;
-    device float * embed_data = (device float *)(dst +  i*nb1);
+    device float * embed_data = (device float *)(dst + i*args.nb1);
 
-    int half_ = dim / 2;
+    int half_ = args.dim / 2;
     for (int j = tpitg.x; j < half_; j += ntg.x) {
         float timestep = ((device float *)src0)[i];
-        float freq = (float)exp(-log((float)max_period) * j / half_);
+        float freq = (float)exp(-log((float)args.max_period) * j / half_);
         float arg = timestep * freq;
         embed_data[j        ] = cos(arg);
         embed_data[j + half_] = sin(arg);
     }
 
-    if (dim % 2 != 0 && tpitg.x == 0) {
-        embed_data[dim] = 0.f;
+    if (args.dim % 2 != 0 && tpitg.x == 0) {
+        embed_data[args.dim] = 0.f;
     }
 }