Repository: ggml-org/llama.cpp
Pull request: metal: optimise `GGML_OP_SUM`

This PR optimises the `GGML_OP_SUM` operation, as implemented in https://github.com/ggml-org/llama.cpp/pull/16539.

The original implementation performed the sum op on one thread as follows:

```cpp
    ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 1, 1, 1);
```

resulting the following `./test-backend-ops perf -o SUM` log:

```
./test-backend-ops perf -o SUM
ggml_metal_library_init: using embedded metal library
ggml_metal_library_init: loaded in 0.007 sec
ggml_metal_device_init: GPU name:   Apple M1 Pro
ggml_metal_device_init: GPU family: MTLGPUFamilyApple7  (1007)
ggml_metal_device_init: GPU family: MTLGPUFamilyCommon3 (3003)
ggml_metal_device_init: GPU family: MTLGPUFamilyMetal4  (5002)
ggml_metal_device_init: simdgroup reduction   = true
ggml_metal_device_init: simdgroup matrix mul. = true
ggml_metal_device_init: has unified memory    = true
ggml_metal_device_init: has bfloat            = true
ggml_metal_device_init: use residency sets    = true
ggml_metal_device_init: use shared buffers    = true
ggml_metal_device_init: recommendedMaxWorkingSetSize  = 12713.12 MB
Testing 3 devices

ggml_metal_init: allocating
ggml_metal_init: found device: Apple M1 Pro
ggml_metal_init: picking default device: Apple M1 Pro
ggml_metal_init: use bfloat         = true
ggml_metal_init: use fusion         = true
ggml_metal_init: use concurrency    = true
ggml_metal_init: use graph optimize = true
Backend 1/3: Metal
  Device description: Apple M1 Pro
  Device memory: 12124 MB (12123 MB free)

ggml_metal_library_compile_pipeline: compiling pipeline: base = 'kernel_op_sum_f32', name = 'kernel_op_sum_f32'
ggml_metal_library_compile_pipeline: loaded kernel_op_sum_f32                             0x102efc6b0 | th_max = 1024 | th_width =   32
  SUM(type=f32,ne=[8192,1,1,1]):              163820 runs -     6.25 us/run -       32 kB/run -    0.02 GB/s
```

### Implementation

To fix this, I've modified the host-side code to launch `nth` threads in one threadgroup, so that each thread sums a strided chunk (similar to `op_sum_rows`).

- `simd_sum(sumf)` is used to perform a partial sum within each SIMD group
- A second `simd_sum(v)`is used to do a reduction across SIMD groups (only SIMD group 0 is involved here)
- The `nth` value is calculated similarly to `op_sum_rows`

Resulting in the following performance log via `./test-backend-ops perf -o SUM`:

```
ggml_metal_library_compile_pipeline: compiling pipeline: base = 'kernel_op_sum_f32', name = 'kernel_op_sum_f32'
ggml_metal_library_compile_pipeline: loaded kernel_op_sum_f32                             0x102efc6b0 | th_max = 1024 | th_width =   32
  SUM(type=f32,ne=[8192,1,1,1]):              163820 runs -     6.25 us/run -       32 kB/run -    4.88 GB/s
  SUM(type=f32,ne=[8192,8192,1,1]):                      128 runs - 23911.46 us/run -   262144 kB/run -   10.54 GB/s
  SUM(type=f32,ne=[128,8192,1,1]):              8191 runs -   296.56 us/run -     4096 kB/run -   13.17 GB/s
  Backend Metal: OK
ggml_metal_free: deallocating
Backend 2/3: BLAS
  Device description: Accelerate
  Device memory: 0 MB (0 MB free)

  SUM(type=f32,ne=[8192,1,1,1]): not supported
  SUM(type=f32,ne=[8192,8192,1,1]): not supported
  SUM(type=f32,ne=[128,8192,1,1]): not supported
  Backend BLAS: OK
Backend 3/3: CPU
  Skipping CPU backend
3/3 backends passed
OK
```

Note that this is still quite a bit slower than `SUM_ROWS`,

```
ggml_metal_library_compile_pipeline: compiling pipeline: base = 'kernel_sum_rows_f32', name = 'kernel_sum_rows_f32'
ggml_metal_library_compile_pipeline: loaded kernel_sum_rows_f32                           0x10552e7f0 | th_max = 1024 | th_width =   32
  SUM_ROWS(type=f32,ne=[8192,1,1,1],permute=0,slice=0):               147438 runs -     6.84 us/run -       32 kB/run -    4.46 GB/s
  SUM_ROWS(type=f32,ne=[8192,8192,1,1],permute=0,slice=0):               768 runs -  1467.22 us/run -   262176 kB/run -  171.74 GB/s
  SUM_ROWS(type=f32,ne=[128,8192,1,1],permute=0,slice=0):              16258 runs -    80.57 us/run -     4128 kB/run -   48.87 GB/s
  Backend Metal: OK
```

So there may be room for improvement in the current kernel implementation.

Changed files:
filename: ggml/src/ggml-cuda/ggml-cuda.cu
patch: @@ -3625,9 +3625,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_CONV_2D_DW:
         case GGML_OP_CONV_TRANSPOSE_2D:
         case GGML_OP_POOL_2D:
-        case GGML_OP_SUM:
         case GGML_OP_ACC:
             return true;
+        case GGML_OP_SUM:
+            return ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_ARGSORT:
             // TODO: Support arbitrary column width
             return op->src[0]->ne[0] <= 1024;

filename: ggml/src/ggml-metal/ggml-metal-device.m
patch: @@ -662,6 +662,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_LOG:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_SUM:
+            return has_simdgroup_reduction && ggml_is_contiguous(op->src[0]);
         case GGML_OP_SUM_ROWS:
         case GGML_OP_MEAN:
         case GGML_OP_SOFT_MAX:

filename: ggml/src/ggml-metal/ggml-metal-ops.cpp
patch: @@ -866,12 +866,25 @@ int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) {
 
     ggml_metal_pipeline_t pipeline = ggml_metal_library_get_pipeline_sum(lib, op);
 
+    int nth = 32; // SIMD width
+
+    while (nth < (int) n && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) {
+        nth *= 2;
+    }
+
+    nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
+    nth = std::min(nth, (int) n);
+
+    const int nsg = (nth + 31) / 32;
+
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         2);
 
-    ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, 1, 1, 1);
+    ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0);
+
+    ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1);
 
     return 1;
 }

filename: ggml/src/ggml-metal/ggml-metal.metal
patch: @@ -1727,18 +1727,48 @@ kernel void kernel_op_sum_f32(
         constant ggml_metal_kargs_sum & args,
         device const float * src0,
         device       float * dst,
-        ushort  tiitg[[thread_index_in_threadgroup]]) {
+        threadgroup  float * shmem_f32 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort3 tpitg[[thread_position_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]],
+        ushort  tiisg[[thread_index_in_simdgroup]],
+        ushort3   ntg[[threads_per_threadgroup]]) {
 
-    if (tiitg != 0) {
+    if (args.np == 0) {
         return;
     }
 
-    float acc = 0.0f;
-    for (ulong i = 0; i < args.np; ++i) {
-        acc += src0[i];
+    const uint nsg = (ntg.x + 31) / 32;
+
+    float sumf = 0;
+
+    for (int64_t i0 = tpitg.x; i0 < args.np; i0 += ntg.x) {
+        sumf += src0[i0];
     }
 
-    dst[0] = acc;
+    sumf = simd_sum(sumf);
+
+    if (tiisg == 0) {
+        shmem_f32[sgitg] = sumf;
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float total = 0;
+
+    if (sgitg == 0) {
+        float v = 0;
+
+        if (tpitg.x < nsg) {
+            v = shmem_f32[tpitg.x];
+        }
+
+        total = simd_sum(v);
+
+        if (tpitg.x == 0) {
+            dst[0] = total;
+        }
+    }
 }
 
 template <bool norm>

filename: tests/test-backend-ops.cpp
patch: @@ -4588,20 +4588,31 @@ struct test_topk_moe: public test_case {
 struct test_sum : public test_case {
     const ggml_type type;
     const std::array<int64_t, 4> ne;
+    const std::array<int64_t, 4> permute;
+    bool _use_permute;
 
     std::string vars() override {
-        return VARS_TO_STR2(type, ne);
+        std::string v = VARS_TO_STR2(type, ne);
+        if (_use_permute) v += "," + VAR_TO_STR(permute);
+        return v;
     }
 
     test_sum(ggml_type type = GGML_TYPE_F32,
-            std::array<int64_t, 4> ne = {10, 5, 4, 3})
-        : type(type), ne(ne) {}
+            std::array<int64_t, 4> ne = {10, 5, 4, 3},
+            std::array<int64_t, 4> permute = {0, 0, 0, 0})
+        : type(type), ne(ne), permute(permute),
+            _use_permute(permute[0] + permute[1] + permute[2] + permute[3] > 0) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data());
         ggml_set_param(a);
         ggml_set_name(a, "a");
 
+        if (_use_permute) {
+            a = ggml_permute(ctx, a, permute[0], permute[1], permute[2], permute[3]);
+            ggml_set_name(a, "a_permuted");
+        }
+
         ggml_tensor * out = ggml_sum(ctx, a);
         ggml_set_name(out, "out");
 
@@ -6724,6 +6735,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     test_cases.emplace_back(new test_sum());
     test_cases.emplace_back(new test_sum_rows());
+    test_cases.emplace_back(new test_sum(GGML_TYPE_F32, {11, 5, 6, 3}, {0, 2, 1, 3}));  // row-contiguous but non-contiguous
+    test_cases.emplace_back(new test_sum(GGML_TYPE_F32, {11, 5, 6, 3}, {0, 3, 2, 1}));
+    test_cases.emplace_back(new test_sum(GGML_TYPE_F32, {11, 5, 6, 3}, {0, 1, 3, 2}));
     test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, false));
     test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, false, true));
     test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 11, 5, 6, 3 }, true, true));
@@ -6734,6 +6748,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 1024, 1, 1 }));
     test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 1024, 1, 1 }));
     test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 }));
+    test_cases.emplace_back(new test_sum(GGML_TYPE_F32, { 33, 256, 1, 1 }, { 1, 0, 2, 3 })); // sum dst not-contiguous
     test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, { 33, 256, 1, 1 }));
     test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 33, 256, 1, 1 }));
     test_cases.emplace_back(new test_mean(GGML_TYPE_F32, { 32769, 1, 1, 1 }));