Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cpu/ops.cpp (C++)

Code before:
// ggml_compute_forward_pool_1d_ksp
static void ggml_compute_forward_pool_1d_ksp(
        const ggml_compute_params * params,
        const ggml_op_pool op,
        const int k,
        const int s,
        const int p,
        ggml_tensor * dst) {

    const ggml_tensor * src = dst->src[0];

    assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16);

    if (params->ith != 0) {
        return;
    }

    const int64_t IW = src->ne[0];
    const int64_t OW = dst->ne[0];

    const int64_t nr =
        (src->ne[1] > 0 ? src->ne[1] : 1) *
        (src->ne[2] > 0 ? src->ne[2] : 1) *
        (src->ne[3] > 0 ? src->ne[3] : 1);

    for (int64_t ir = params->ith; ir < nr; ir += params->nth) {
        const char * srow_bytes = (const char *) src->data + ir * src->nb[1];
        float      * drow       = (float *) ((char *) dst->data + ir * dst->nb[1]);

        for (int64_t ow = 0; ow < OW; ++ow) {
            float acc;
            if (op == GGML_OP_POOL_AVG) {
                acc = 0.0f;
            } else if (op == GGML_OP_POOL_MAX) {
                acc = -FLT_MAX;
            } else {
                GGML_ABORT("POOL_1D: unsupported op");
            }

            int count = 0;
            const int base = (int) ow * s - p;

            for (int ki = 0; ki < k; ++ki) {
                const int j = base + ki;
                if (j < 0 || j >= (int) IW) {
                    continue;
                }

                float v;
                if (src->type == GGML_TYPE_F32) {
                    v = ((const float *) srow_bytes)[j];

Review comment (suggestion):
The loop is now multi-threaded, but we still exit early above on line 7126 for all threads with `ith != 0`.

There are 2 acceptable fixes:

- Either change this loop back to be single-thread
- Or update the op to use more than 1 thread, remove the early exit in this function and verify the correctness:

https://github.com/ggml-org/llama.cpp/blob/3fc24ad6ba0c50ff1e49a707018c06e754dc7124/ggml/src/ggml-cpu/ggml-cpu.c#L2320-L2326

Code after:
static void ggml_compute_forward_pool_1d_ksp(
        const ggml_compute_params * params,
        const ggml_op_pool op,
        const int k,
        const int s,
        const int p,
        ggml_tensor * dst) {

    const ggml_tensor * src = dst->src[0];

    assert(src->type == GGML_TYPE_F32 || src->type == GGML_TYPE_F16);

    if (params->ith != 0) {
        return;
    }

    const int64_t IW = src->ne[0];
    const int64_t OW = dst->ne[0];

    const int64_t nr = ggml_nrows(src);

    for (int64_t ir = 0; ir < nr; ++ir) {
        const char * srow_bytes =            (const char *) src->data + ir * src->nb[1];
        float      * drow       = (float *) ((      char *) dst->data + ir * dst->nb[1]);

        for (int64_t ow = 0; ow < OW; ++ow) {
            float res = 0;
            switch (op) {
                case GGML_OP_POOL_AVG: res = 0.0f;     break;
                case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
                case GGML_OP_POOL_COUNT: GGML_ABORT("fatal error");
            }

            int count = 0;
            const int base = (int) ow * s - p;

            for (int ki = 0; ki < k; ++ki) {
                const int j = base + ki;
                if (j < 0 || j >= (int) IW) {
                    continue;
                }

                float v;
                if (src->type == GGML_TYPE_F32) {
                    v = ((const float *) srow_bytes)[j];
                } else {
                    v = GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) srow_bytes)[j]);
                }

                switch (op) {
                    case GGML_OP_POOL_AVG: res += v;                break;