diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a8a1c09ca..b4c128141 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -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: diff --git a/ggml/src/ggml-cuda/pool1d.cu b/ggml/src/ggml-cuda/pool1d.cu new file mode 100644 index 000000000..ac6fb0cbd --- /dev/null +++ b/ggml/src/ggml-cuda/pool1d.cu @@ -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<<>>(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(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); +} diff --git a/ggml/src/ggml-cuda/pool1d.cuh b/ggml/src/ggml-cuda/pool1d.cuh new file mode 100644 index 000000000..c79461dd8 --- /dev/null +++ b/ggml/src/ggml-cuda/pool1d.cuh @@ -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);