cuda: update uncoalesced memory reads in pool2d (#29425)

* cuda: update uncoalesced memory reads in pool2d

* cuda: Added perf and __restrict__ in pool2d

* Update POOL2D_WARP_KERNEL_MIN_WINDOW macro

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

* cuda: Fix compiler bugs

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

ggml/src/ggml-cuda/pool2d.cu

@@ -50,6 +50,65 @@ static  __global__ void pool2d_nchw_kernel(
     o_ptr[cur_oh * ow + cur_ow] = res;
 }
 
+template <typename Ti, typename To>
+static __global__ void pool2d_nchw_kernel_warp(
+        const int ih, const int iw, const int oh, const int ow,
+        const int kh, const int kw, const int sh, const int sw,
+        const int ph, const int pw, const int parallel_elements,
+        const Ti * __restrict__ src, To * __restrict__ dst, const enum ggml_op_pool op) {
+    const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / WARP_SIZE;
+    const int lane     = threadIdx.x % WARP_SIZE;
+    if (warp_id >= parallel_elements) {
+        return;
+    }
+
+    const int I_HW = ih * iw;
+    const int O_HW = oh * ow;
+    const int nc     = warp_id / O_HW;
+    const int cur_oh = warp_id % O_HW / ow;
+    const int cur_ow = warp_id % O_HW % ow;
+    const Ti* i_ptr = src + nc * I_HW;
+
+    const int start_h = cur_oh * sh - ph;
+    const int bh = max(0, start_h);
+    const int eh = min(ih, start_h + kh);
+    const int start_w = cur_ow * sw - pw;
+    const int bw = max(0, start_w);
+    const int ew = min(iw, start_w + kw);
+
+    const int win_w     = ew - bw;
+    const int win_elems = (eh - bh) * win_w;
+    const To scale = 1. / (kh * kw);
+
+    To res;
+    switch (op) {
+        case GGML_OP_POOL_AVG: res = 0; break;
+        case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
+        default: res = 0; assert(false);
+    }
+
+    for (int t = lane; t < win_elems; t += WARP_SIZE) {
+        const int i = bh + t / win_w;
+        const int j = bw + t % win_w;
+        const Ti cur = i_ptr[i * iw + j];
+        switch (op) {
+            case GGML_OP_POOL_AVG: res += cur * scale; break;
+            case GGML_OP_POOL_MAX: res = max(res, (To)cur); break;
+            default: break;
+        }
+    }
+
+#pragma unroll
+    for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
+        const To other = __shfl_xor_sync(0xFFFFFFFF, res, offset, WARP_SIZE);
+        res = (op == GGML_OP_POOL_MAX) ? max(res, other) : res + other;
+    }
+
+    if (lane == 0) {
+        dst[nc * O_HW + cur_oh * ow + cur_ow] = res;
+    }
+}
+
 static void pool2d_nchw_kernel_f32_f32_cuda(
         const int ih, const int iw, const int oh, const int ow,
         const int kh, const int kw, const int sh, const int sw,
@@ -57,6 +116,13 @@ static void pool2d_nchw_kernel_f32_f32_cuda(
         const float * src, float * dst, const enum ggml_op_pool op,
         cudaStream_t stream) {
 
+    if (kh * kw >= CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW) {
+        const int warps_per_block = CUDA_POOL2D_BLOCK_SIZE / WARP_SIZE;
+        const int num_blocks = (parallel_elements + warps_per_block - 1) / warps_per_block;
+        pool2d_nchw_kernel_warp<<<num_blocks, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
+        return;
+    }
+
     const int num_blocks = (parallel_elements + CUDA_POOL2D_BLOCK_SIZE - 1) / CUDA_POOL2D_BLOCK_SIZE;
     dim3 block_nums(num_blocks);
     pool2d_nchw_kernel<<<block_nums, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);

ggml/src/ggml-cuda/pool2d.cuh

@@ -1,5 +1,6 @@
 #include "common.cuh"
 
 #define CUDA_POOL2D_BLOCK_SIZE 256
+#define CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW 32
 
 void ggml_cuda_op_pool2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

tests/test-backend-ops.cpp

@@ -9558,6 +9558,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                     }
                 }
             }
+            // MSA indexer block-max shape: wide 1-D window along ne0, no padding, non-divisible width
+            for (int64_t iw : {4096, 4160}) {
+                for (int blk : {32, 64}) {
+                    test_cases.emplace_back(new test_pool2d(pool_type, type_input, {iw, 2, 1, 1}, blk, 1, blk, 1, 0, 0));
+                }
+            }
         }
     }
 
@@ -11652,6 +11658,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
             GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3}));
     }
 
+    // POOL_2D max over MiniMax-M3 indexer block scores:
+    // sc is [n_ps, n_head=4, n_tokens] and gets ggml_pool_2d(.., MAX, blk=128, 1, blk=128, 1, 0, 0).
+    for (int64_t n_ps : { 8192, 32768 }) {
+        for (int64_t n_tokens : { 1, 512 }) {
+            test_cases.emplace_back(new test_pool2d(
+                GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ps, 4, n_tokens, 1}, 128, 1, 128, 1, 0, 0));
+        }
+    }
+
     // LEAKY_RELU at FFN activation width, for direct comparison with RELU
     for (int64_t n_tokens : {512, 2048}) {
         test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));