cuda : add POOL_1D support (#27573)

* cuda : add POOL_1D support

* fix: add missing trailing newline for editorconfig compliance
This commit is contained in:
Aman Karki
2026-08-23 10:37:32 +03:00
committed by GitHub
parent 70adb1b4ce
commit 29ea9412a6
3 changed files with 95 additions and 0 deletions
+5
View File
@@ -38,6 +38,7 @@
#include "ggml-cuda/out-prod.cuh"
#include "ggml-cuda/pad.cuh"
#include "ggml-cuda/pool2d.cuh"
#include "ggml-cuda/pool1d.cuh"
#include "ggml-cuda/quantize.cuh"
#include "ggml-cuda/rope.cuh"
#include "ggml-cuda/roll.cuh"
@@ -2326,6 +2327,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_POOL_2D:
ggml_cuda_op_pool2d(ctx, dst);
break;
case GGML_OP_POOL_1D:
ggml_cuda_op_pool1d(ctx, dst);
break;
case GGML_OP_SUM:
ggml_cuda_op_sum(ctx, dst);
break;
@@ -5245,6 +5249,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
case GGML_OP_CONV_2D_DW:
return op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_CONV_TRANSPOSE_2D:
case GGML_OP_POOL_1D:
case GGML_OP_POOL_2D:
return true;
case GGML_OP_ACC:
+85
View File
@@ -0,0 +1,85 @@
#include "pool1d.cuh"
static __global__ void pool1d_nchw_kernel(
const int iw, const int ow,
const int kw, const int sw, const int pw,
const int parallel_elements,
const float * src, float * dst, const enum ggml_op_pool op) {
const int idx = threadIdx.x + blockIdx.x * blockDim.x;
if (idx >= parallel_elements) {
return;
}
const int nc = idx / ow;
const int cur_ow = idx % ow;
const float * i_ptr = src + nc * iw;
float * o_ptr = dst + nc * ow;
const int start = cur_ow * sw - pw;
const int b = max(0, start);
const int e = min(iw, start + kw);
float res;
switch (op) {
case GGML_OP_POOL_AVG: res = 0.0f; break;
case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
default: return;
}
int count = 0;
for (int i = b; i < e; i++) {
#if __CUDA_ARCH__ >= 350
float cur = __ldg(i_ptr + i);
#else
float cur = i_ptr[i];
#endif
switch (op) {
case GGML_OP_POOL_AVG: res += cur; break;
case GGML_OP_POOL_MAX: res = max(res, cur); break;
default: break;
}
count++;
}
if (op == GGML_OP_POOL_AVG) {
res = (count > 0) ? (res / count) : 0.0f;
}
o_ptr[cur_ow] = res;
}
static void pool1d_nchw_kernel_f32_f32_cuda(
const int iw, const int ow,
const int kw, const int sw, const int pw,
const int parallel_elements,
const float * src, float * dst, const enum ggml_op_pool op,
cudaStream_t stream) {
const int num_blocks = (parallel_elements + CUDA_POOL1D_BLOCK_SIZE - 1) / CUDA_POOL1D_BLOCK_SIZE;
dim3 block_nums(num_blocks);
pool1d_nchw_kernel<<<block_nums, CUDA_POOL1D_BLOCK_SIZE, 0, stream>>>(iw, ow, kw, sw, pw, parallel_elements, src, dst, op);
}
void ggml_cuda_op_pool1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const float * src0_d = (const float *)src0->data;
float * dst_d = (float *)dst->data;
cudaStream_t stream = ctx.stream();
GGML_ASSERT(src0->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F32);
const int32_t * opts = (const int32_t *)dst->op_params;
enum ggml_op_pool op = static_cast<ggml_op_pool>(opts[0]);
const int k0 = opts[1];
const int s0 = opts[2];
const int p0 = opts[3];
const int64_t IW = src0->ne[0];
const int64_t OW = dst->ne[0];
const int64_t nr = ggml_nrows(src0);
const int parallel_elements = (int)(nr * OW);
pool1d_nchw_kernel_f32_f32_cuda(IW, OW, k0, s0, p0, parallel_elements, src0_d, dst_d, op, stream);
}
+5
View File
@@ -0,0 +1,5 @@
#include "common.cuh"
#define CUDA_POOL1D_BLOCK_SIZE 256
void ggml_cuda_op_pool1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);