add opt-step-adamw kernel for metal

Author: cern1710

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

@@ -773,4 +773,15 @@ typedef struct {
     uint64_t nb01;
 } ggml_metal_kargs_argmax;
 
+typedef struct {
+    float    alpha;
+    float    beta1;
+    float    beta2;
+    float    eps;
+    float    wd;
+    float    beta1h;
+    float    beta2h;
+    int64_t  np;
+} ggml_metal_kargs_opt_step_adamw;
+
 #endif // GGML_METAL_IMPL

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -3407,5 +3407,44 @@ int ggml_metal_op_leaky_relu(ggml_metal_op_t ctx, int idx) {
 }
 
 int ggml_metal_op_opt_step_adamw(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
+    GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb);
+    GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
+    GGML_TENSOR_LOCALS(uint32_t, nb,  op,         nb);
+
+    ggml_metal_pipeline_t pipeline = ggml_metal_library_get_pipeline_opt_step_adamw(lib, op);
+
+    const int64_t np     = ggml_nelements(op->src[0]);
+    const float * params = (const float *) op->src[4]->data;
+    ggml_metal_kargs_opt_step_adamw args = {
+        /*.alpha  =*/ params[0],
+        /*.beta1  =*/ params[1],
+        /*.beta2  =*/ params[2],
+        /*.eps    =*/ params[3],
+        /*.wd     =*/ params[4],
+        /*.beta1h =*/ params[5],
+        /*.beta2h =*/ params[6],
+        /*.np     =*/ np,
+    };
+
+    int ida = 0;
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), ida++);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), ida++);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[1]), ida++);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[2]), ida++);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[3]), ida++);
+
+    const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne0);
+    const int64_t n = (np + nth - 1) / nth;
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, nth, 1, 1);
+
     return 1;
 }

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

@@ -8754,3 +8754,36 @@ kernel void kernel_pool_2d_avg_f32(
 
     o_ptr[cur_oh * args.OW + cur_ow] = res;
 }
+
+kernel void kernel_opt_step_adamw_f32(
+        constant    ggml_metal_kargs_opt_step_adamw & args,
+        device       float * x,
+        device const float * g,
+        device       float * g_m,
+        device       float * g_v,
+        uint        gid[[thread_position_in_grid]]) {
+
+    if (gid >= args.np) {
+        return;
+    }
+
+    const float alpha  = args.alpha;
+    const float beta1  = args.beta1;
+    const float beta2  = args.beta2;
+    const float eps    = args.eps;
+    const float wd     = args.wd;
+    const float beta1h = args.beta1h;
+    const float beta2h = args.beta2h;
+
+    const float gi = g[gid];
+    const float gmi = g_m[gid] * beta1 +      gi * (1.0f - beta1);
+    const float gvi = g_v[gid] * beta2 + gi * gi * (1.0f - beta2);
+
+    g_m[gid] = gmi;
+    g_v[gid] = gvi;
+
+    const float mh =      gmi * beta1h;
+    const float vh = sqrt(gvi * beta2h) + eps;
+
+    x[gid] = x[gid] * (1.0f - alpha * wd) - alpha * mh / vh;
+}