KBaba7/llama.cpp

1#include "sumrows.cuh"2 3static __global__ void k_sum_rows_f32(const float * x, float * dst, const int ncols) {4    const int row = blockIdx.x;5    const int col = threadIdx.x;6 7    float sum = 0.0f;8    for (int i = col; i < ncols; i += blockDim.x) {9        sum += x[row * ncols + i];10    }11 12    sum = warp_reduce_sum(sum);13 14    if (col == 0) {15        dst[row] = sum;16    }17}18 19void sum_rows_f32_cuda(const float * x, float * dst, const int ncols, const int nrows, cudaStream_t stream) {20    const dim3 block_dims(WARP_SIZE, 1, 1);21    const dim3 block_nums(nrows, 1, 1);22    k_sum_rows_f32<<<block_nums, block_dims, 0, stream>>>(x, dst, ncols);23}24 25void ggml_cuda_op_sum_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {26    const ggml_tensor * src0 = dst->src[0];27    const float * src0_d = (const float *)src0->data;28    float * dst_d = (float *)dst->data;29    cudaStream_t stream = ctx.stream();30 31    GGML_ASSERT(src0->type == GGML_TYPE_F32);32    GGML_ASSERT( dst->type == GGML_TYPE_F32);33    GGML_ASSERT(ggml_is_contiguous(src0));34 35    const int64_t ncols = src0->ne[0];36    const int64_t nrows = ggml_nrows(src0);37 38    sum_rows_f32_cuda(src0_d, dst_d, ncols, nrows, stream);39}40