char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Aman Karkicuda : add POOL_1D support (#27573)29ea9412a

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.6 KiB85 linesraw
1#include "pool1d.cuh"
2
3static __global__ void pool1d_nchw_kernel(
4        const int iw, const int ow,
5        const int kw, const int sw, const int pw,
6        const int parallel_elements,
7        const float * src, float * dst, const enum ggml_op_pool op) {
8    const int idx = threadIdx.x + blockIdx.x * blockDim.x;
9    if (idx >= parallel_elements) {
10        return;
11    }
12
13    const int nc = idx / ow;
14    const int cur_ow = idx % ow;
15
16    const float * i_ptr = src + nc * iw;
17    float * o_ptr = dst + nc * ow;
18
19    const int start = cur_ow * sw - pw;
20    const int b = max(0, start);
21    const int e = min(iw, start + kw);
22
23    float res;
24    switch (op) {
25        case GGML_OP_POOL_AVG: res = 0.0f;     break;
26        case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
27        default: return;
28    }
29
30    int count = 0;
31    for (int i = b; i < e; i++) {
32#if __CUDA_ARCH__ >= 350
33        float cur = __ldg(i_ptr + i);
34#else
35        float cur = i_ptr[i];
36#endif
37        switch (op) {
38            case GGML_OP_POOL_AVG: res += cur;                break;
39            case GGML_OP_POOL_MAX: res = max(res, cur);       break;
40            default: break;
41        }
42        count++;
43    }
44
45    if (op == GGML_OP_POOL_AVG) {
46        res = (count > 0) ? (res / count) : 0.0f;
47    }
48
49    o_ptr[cur_ow] = res;
50}
51
52static void pool1d_nchw_kernel_f32_f32_cuda(
53        const int iw, const int ow,
54        const int kw, const int sw, const int pw,
55        const int parallel_elements,
56        const float * src, float * dst, const enum ggml_op_pool op,
57        cudaStream_t stream) {
58    const int num_blocks = (parallel_elements + CUDA_POOL1D_BLOCK_SIZE - 1) / CUDA_POOL1D_BLOCK_SIZE;
59    dim3 block_nums(num_blocks);
60    pool1d_nchw_kernel<<<block_nums, CUDA_POOL1D_BLOCK_SIZE, 0, stream>>>(iw, ow, kw, sw, pw, parallel_elements, src, dst, op);
61}
62
63void ggml_cuda_op_pool1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
64    const ggml_tensor * src0 = dst->src[0];
65    const float * src0_d = (const float *)src0->data;
66    float * dst_d = (float *)dst->data;
67    cudaStream_t stream = ctx.stream();
68
69    GGML_ASSERT(src0->type == GGML_TYPE_F32);
70    GGML_ASSERT( dst->type == GGML_TYPE_F32);
71
72    const int32_t * opts = (const int32_t *)dst->op_params;
73    enum ggml_op_pool op = static_cast<ggml_op_pool>(opts[0]);
74    const int k0 = opts[1];
75    const int s0 = opts[2];
76    const int p0 = opts[3];
77
78    const int64_t IW = src0->ne[0];
79    const int64_t OW = dst->ne[0];
80    const int64_t nr = ggml_nrows(src0);
81
82    const int parallel_elements = (int)(nr * OW);
83
84    pool1d_nchw_kernel_f32_f32_cuda(IW, OW, k0, s0, p0, parallel_elements, src0_d, dst_d, op, stream);
85}