Metal: add opt_step_adamw and op_sum

[cern1710]

Part of #14909:

This PR adds Metal backend for the OPT_STEP_ADAMW operator.

Implementation

• Added kernel_opt_step_adamw_f32 kernel in ggml-metal.metal
• Implemented argument struct for parameters, ggml_metal_kargs_opt_step_adamw
• Storing parameters and np in the struct as a passed constant & args
• Threadgroup size calculation, similar to ggml_metal_op_upscale

I've not written test cases for this operator as one already exists in test-backend-ops.cpp.

Additional changes

• Due to test-opts failing, the GGML_OP_SUM operator is implemented in Metal (currently unoptimised)

[cern1710]

I've tried commenting out most of the code inside ggml_metal_op_opt_step_adamw() (L3410 - L3440) as well as the entire kernel, which did not resolve the CI failure. Below I've copied the same CI error on my own machine by running ./test-opt:

pre-allocated tensor (loss_sum) in a buffer (Metal) that cannot run the operation (SUM)
(lldb) process attach --pid 84566
Process 84566 stopped
* thread #1, queue = 'com.apple.main-thread', stop reason = signal SIGSTOP
    frame #0: 0x000000018342d42c libsystem_kernel.dylib`__wait4 + 8
libsystem_kernel.dylib`__wait4:
->  0x18342d42c <+8>:  b.lo   0x18342d44c    ; <+40>
    0x18342d430 <+12>: pacibsp 
    0x18342d434 <+16>: stp    x29, x30, [sp, #-0x10]!
    0x18342d438 <+20>: mov    x29, sp
Executable module set to "/Users/samuel/Desktop/Personal/llama.cpp/build/bin/test-opt".
Architecture set to: arm64-apple-macosx-.
(lldb) bt
* thread #1, queue = 'com.apple.main-thread', stop reason = signal SIGSTOP
  * frame #0: 0x000000018342d42c libsystem_kernel.dylib`__wait4 + 8
    frame #1: 0x000000010115146c libggml-base.dylib`ggml_abort + 156
    frame #2: 0x000000010116b724 libggml-base.dylib`ggml_backend_sched_backend_id_from_cur(ggml_backend_sched*, ggml_tensor*) + 896
    frame #3: 0x0000000101169304 libggml-base.dylib`ggml_backend_sched_split_graph + 444
    frame #4: 0x000000010116c2f0 libggml-base.dylib`ggml_backend_sched_alloc_graph + 96
    frame #5: 0x0000000101170d68 libggml-base.dylib`ggml_opt_alloc + 872
    frame #6: 0x0000000100f3397c test-opt`test_backend(ggml_backend_sched*, ggml_backend*, ggml_opt_optimizer_type) + 1772
    frame #7: 0x0000000100f32ff4 test-opt`main + 1300
    frame #8: 0x00000001830a9d54 dyld`start + 7184
(lldb) quit
[1]    84566 abort      ./test-opt

The culprit is this line right here, but I'm unsure as to whether this has something to do with my kernel or whether this is another issue.

llama.cpp/ggml/src/ggml-backend.cpp

Line 796 in c7be9fe

GGML_ABORT("pre-allocated tensor (%s) in a buffer (%s) that cannot run the operation (%s)", tensor->name, ggml_backend_buffer_name(buffer), ggml_op_name(tensor->op));

[ggerganov]

We need to implement the GGML_OP_SUM operator in Metal as well.

[cern1710]

I see, thanks for the follow up. Would you like me to implement GGML_OP_SUM for Metal as part of this PR, or should that be done in a follow-up one?

[ggerganov]

Yes, it should be in this PR so that the CI becomes green.

[cern1710]

I have a working GGML_OP_SUM implementation for Metal which is very simple, but fixes the CI abort.
Would you prefer if I continue optimising it or just leave it as is for now?

P.S. For an optimised implementation, would we be looking at reducing the lane for each SIMD group using simd_sum(), then a shared memory reduction to combine the SIMD group partials?

[ggerganov]

Yes, we can optimize it in a separate PR. For now just correctness is OK.

P.S. For an optimised implementation, would we be looking at reducing the lane for each SIMD group using simd_sum(), then a shared memory reduction to combine the SIMD group partials?

Yes, exactly.

[cern1710]

Right, this change fixed CI / ggml-ci-mac-metal, but somehow CI / macOS-latest-cmake-arm64 is still failing.

Is it possible that it is too restrictive to tie SUM to has_simdgroup_reduction?

        case GGML_OP_SUM:
        case GGML_OP_SUM_ROWS:
        case GGML_OP_MEAN:
        case GGML_OP_SOFT_MAX:
        case GGML_OP_GROUP_NORM:
            return has_simdgroup_reduction && ggml_is_contiguous_rows(op->src[0]);

I think It might be better to do something like this for now:

        case GGML_OP_SUM:
            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;

[ggerganov]

It might be better to add has_simdgroup_reduction to OPT_STEP_ADAMW so that both ops would end up being not supported on the virtualized device of the CI. This way GGML_OP_SUM would keep the simdgroup reduction requirement, which we will need later for optimization.

[cern1710]

Yes, that makes sense. I've made that change in the newest commit

I also think the previous commit has passed the failing CI