char-fork/llama.cpp

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

Yihao WangCUDA & CPU: support F32 kernel type for `CONV_TRANSPOSE_2D` (#17094)0a524f240

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
4.7 KiB115 linesraw
1#include "conv2d-transpose.cuh"
2#include "convert.cuh"
3
4template <typename kernel_t>
5static __global__ void conv2d_transpose_kernel(const float * __restrict__ input,
6                                               const kernel_t * __restrict__ kernel,
7                                               float * __restrict__ output,
8                                               const int in_w,
9                                               const int in_h,
10                                               const int out_w,
11                                               const int out_h,
12                                               const int kernel_w,
13                                               const int kernel_h,
14                                               const int stride,
15                                               const int c_in,
16                                               const int c_out,
17                                               const int batches) {
18    const int global_idx = blockIdx.x * blockDim.x + threadIdx.x;
19
20    const int total_elements = out_w * out_h * c_out * batches;
21
22    if (global_idx >= total_elements) {
23        return;
24    }
25
26    const int out_x_idx = global_idx % out_w;
27    const int out_y_idx = (global_idx / out_w) % out_h;
28    const int c_idx     = (global_idx / (out_w * out_h)) % c_out;
29    const int n_idx     = global_idx / (out_w * out_h * c_out);
30
31    float accumulator = 0;
32    // For each output idx, find the inputs that contribute to it by checking stride alignment and bounds
33
34    for (int c_in_idx = 0; c_in_idx < c_in; c_in_idx++) {
35        for (int kh = 0; kh < kernel_h; ++kh) {
36            int in_y = out_y_idx - kh;
37            if (in_y < 0 || in_y % stride) {
38                continue;
39            }
40            in_y /= stride;
41            if (in_y >= in_h) {
42                continue;
43            }
44
45            for (int kw = 0; kw < kernel_w; ++kw) {
46                int in_x = out_x_idx - kw;
47                if (in_x < 0 || in_x % stride) {
48                    continue;
49                }
50                in_x /= stride;
51                if (in_x >= in_w) {
52                    continue;
53                }
54
55                const int input_idx = (in_w * in_h * c_in) * n_idx + (in_w * in_h) * c_in_idx + (in_w) *in_y + in_x;
56                const int kernel_idx =
57                    (kernel_h * kernel_w * c_out) * c_in_idx + (kernel_h * kernel_w) * c_idx + (kernel_w) *kh + kw;
58
59                float    input_val = input[input_idx];
60                kernel_t kern_val  = kernel[kernel_idx];
61
62                accumulator += input_val * ggml_cuda_cast<float>(kern_val);
63            }
64        }
65    }
66
67    output[(out_w * out_h * c_out) * n_idx + (out_w * out_h) * c_idx + (out_w) *out_y_idx + out_x_idx] = accumulator;
68}
69
70//input is (W, H, C_in, N), Kernel is (W, H, C_out, C_in)
71void ggml_cuda_conv_2d_transpose_p0(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
72    const ggml_tensor * kernel = dst->src[0];
73    const ggml_tensor * input  = dst->src[1];
74
75    GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32);
76    GGML_ASSERT(input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
77
78    const float * input_data  = (const float *) input->data;
79    float *       output_data = (float *) dst->data;
80    const void *  kernel_data = kernel->data;
81
82    const int input_w      = input->ne[0];
83    const int input_h      = input->ne[1];
84    const int output_w     = dst->ne[0];
85    const int output_h     = dst->ne[1];
86    const int channels_in  = input->ne[2];
87    const int channels_out = kernel->ne[2];
88    const int kernel_w     = kernel->ne[0];
89    const int kernel_h     = kernel->ne[1];
90    const int stride       = dst->op_params[0];
91    const int batches      = input->ne[3];
92
93    GGML_ASSERT(channels_in == kernel->ne[3]);
94    GGML_ASSERT(stride > 0);
95
96    cudaStream_t st = ctx.stream();
97
98    GGML_ASSERT(ggml_is_contiguous(input));
99    GGML_ASSERT(ggml_is_contiguous(kernel));
100    GGML_ASSERT(ggml_is_contiguous(dst));
101
102    const int total  = output_w * output_h * channels_out * batches;
103    const int blocks = (total + CUDA_CONV2D_TRANSPOSE_BLOCK_SIZE - 1) / CUDA_CONV2D_TRANSPOSE_BLOCK_SIZE;
104
105    if (kernel->type == GGML_TYPE_F16) {
106        conv2d_transpose_kernel<half><<<blocks, CUDA_CONV2D_TRANSPOSE_BLOCK_SIZE, 0, st>>>(
107            input_data, (const half *) kernel_data, output_data, input_w, input_h, output_w, output_h, kernel_w,
108            kernel_h, stride, channels_in, channels_out, batches);
109
110    } else {
111        conv2d_transpose_kernel<float><<<blocks, CUDA_CONV2D_TRANSPOSE_BLOCK_SIZE, 0, st>>>(
112            input_data, (const float *) kernel_data, output_data, input_w, input_h, output_w, output_h, kernel_w,
113            kernel_h, stride, channels_in, channels_out, batches);
114    }
115}