char-fork/llama.cpp

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

pmaybankTest support for alternative conv layout (#25617)958d9c0b6

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.4 KiB167 linesraw
1#include "conv2d.cuh"
2#include "convert.cuh"
3
4struct conv_params {
5    const int64_t IW, IH;
6    const int64_t OW, OH;
7    const int64_t KW, KH;
8    const int64_t ST_X, ST_Y;
9    const int64_t PD_X, PD_Y;
10    const int64_t DL_X, DL_Y;
11    const int64_t IC, OC;
12    const int64_t B;
13    const int64_t TOTAL;
14};
15
16struct kernel_bounds {
17    int64_t y_min, y_max;
18    int64_t x_min, x_max;
19};
20
21__device__ __forceinline__ int64_t max64(int64_t a, int64_t b) {
22    return (a > b) ? a : b;
23}
24
25__device__ __forceinline__ int64_t min64(int64_t a, int64_t b) {
26    return (a < b) ? a : b;
27}
28
29__device__ __forceinline__ kernel_bounds calculate_kernel_bounds(int64_t out_x, int64_t out_y, const conv_params & P) {
30    kernel_bounds bounds;
31    bounds.y_min = max64(0, (P.PD_Y - out_y * P.ST_Y + P.DL_Y - 1) / P.DL_Y);
32    bounds.y_max = min64(P.KH, (P.IH + P.PD_Y - out_y * P.ST_Y + P.DL_Y - 1) / P.DL_Y);
33    bounds.x_min = max64(0, (P.PD_X - out_x * P.ST_X + P.DL_X - 1) / P.DL_X);
34    bounds.x_max = min64(P.KW, (P.IW + P.PD_X - out_x * P.ST_X + P.DL_X - 1) / P.DL_X);
35    return bounds;
36}
37
38__device__ __forceinline__ int calculate_input_coord(int64_t out_coord,
39                                                     int64_t kern_coord,
40                                                     int64_t stride,
41                                                     int64_t dilation,
42                                                     int64_t padding) {
43    return out_coord * stride + kern_coord * dilation - padding;
44}
45
46struct whcn_layout {
47    __device__ static int64_t input_index(int64_t n, int64_t c, int64_t y, int64_t x, const conv_params & P) {
48        return n * (P.IC * P.IW * P.IH) + c * P.IW * P.IH + y * P.IW + x;
49    }
50
51    __device__ static int64_t kernel_index(int64_t c_out, int64_t c_in, int64_t ky, int64_t kx, const conv_params & P) {
52        return c_out * (P.IC * P.KH * P.KW) + c_in * (P.KH * P.KW) + ky * P.KW + kx;
53    }
54
55    __device__ static int64_t output_index(int64_t n, int64_t c, int64_t y, int64_t x, const conv_params & P) {
56        return n * (P.OC * P.OW * P.OH) + c * P.OW * P.OH + y * P.OW + x;
57    }
58
59    __device__ static void unpack_indices(int64_t             global_idx,
60                                          const conv_params & P,
61                                          int64_t &           n,
62                                          int64_t &           c,
63                                          int64_t &           out_y,
64                                          int64_t &           out_x) {
65        out_x = global_idx % P.OW;
66        out_y = (global_idx / P.OW) % P.OH;
67        c     = (global_idx / (P.OW * P.OH)) % P.OC;
68        n     = global_idx / (P.OW * P.OH * P.OC);
69    }
70};
71
72template <typename T, typename Layout>
73static __global__ void conv2d_kernel(const float * __restrict__ input,
74                                     const T * __restrict__ kernel,
75                                     float * __restrict__ output,
76                                     const conv_params P) {
77    const int64_t global_idx = blockIdx.x * blockDim.x + threadIdx.x;
78
79    if (global_idx >= P.TOTAL) {
80        return;
81    }
82
83    int64_t n, c_out, out_y, out_x;
84    Layout::unpack_indices(global_idx, P, n, c_out, out_y, out_x);
85
86    float acc = 0.0f;
87
88    for (int64_t c_in = 0; c_in < P.IC; ++c_in) {
89        kernel_bounds bounds = calculate_kernel_bounds(out_x, out_y, P);
90
91        for (int64_t ky = bounds.y_min; ky < bounds.y_max; ++ky) {
92            const int64_t in_y = calculate_input_coord(out_y, ky, P.ST_Y, P.DL_Y, P.PD_Y);
93
94            for (int64_t kx = bounds.x_min; kx < bounds.x_max; ++kx) {
95                const int64_t in_x = calculate_input_coord(out_x, kx, P.ST_X, P.DL_X, P.PD_X);
96
97                const float input_val = input[Layout::input_index(n, c_in, in_y, in_x, P)];
98                const T kernel_val = kernel[Layout::kernel_index(c_out, c_in, ky, kx, P)];
99                acc += (input_val * ggml_cuda_cast<float>(kernel_val));
100            }
101        }
102    }
103
104    // [N, OC, OH, OW]
105    output[Layout::output_index(n, c_out, out_y, out_x, P)] = acc;
106}
107
108template <typename T>
109static void conv2d_cuda(const float * X_D, const T * K_D, float * Y_D, const conv_params P, cudaStream_t st) {
110    const int blocks = (P.TOTAL + CUDA_CONV2D_BLOCK_SIZE - 1) / CUDA_CONV2D_BLOCK_SIZE;
111    conv2d_kernel<T, whcn_layout><<<blocks, CUDA_CONV2D_BLOCK_SIZE, 0, st>>>(X_D, K_D, Y_D, P);
112}
113
114static void conv2d_cuda_f16(const float * X_D, const half * K_D, float * Y_D, const conv_params P, cudaStream_t st) {
115    conv2d_cuda<half>(X_D, K_D, Y_D, P, st);
116}
117
118static void conv2d_cuda_f32(const float * X_D, const float * K_D, float * Y_D, const conv_params P, cudaStream_t st) {
119    conv2d_cuda<float>(X_D, K_D, Y_D, P, st);
120}
121
122void ggml_cuda_op_conv2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
123    const ggml_tensor * kernel = dst->src[0];
124    const ggml_tensor * input  = dst->src[1];
125    float *             K_D    = (float *) kernel->data;
126    const float *       X_D    = (const float *) input->data;
127    float *             Y_D    = (float *) dst->data;
128
129    GGML_ASSERT(ggml_is_contiguous(input));
130    GGML_ASSERT(ggml_is_contiguous(kernel));
131    GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32);
132
133    // same number of input channels
134    GGML_ASSERT(input->ne[2] == kernel->ne[2]);
135
136    cudaStream_t st = ctx.stream();
137
138    const int32_t * p    = (const int32_t *) dst->op_params;
139    const int       ST_X = p[0];  // stride_x
140    const int       ST_Y = p[1];  // stride_y
141    const int       PD_X = p[2];  // padding_x
142    const int       PD_Y = p[3];  // padding_y
143    const int       DL_X = p[4];  // dilation_x
144    const int       DL_Y = p[5];  // dilation_y
145
146    // No cwhn
147    GGML_ASSERT(p[6] == false);
148
149    const int IW = input->ne[0];   // input_w
150    const int IH = input->ne[1];   // input_h
151    const int OW = dst->ne[0];     // output_w
152    const int OH = dst->ne[1];     // output_h
153    const int KW = kernel->ne[0];  // kernel_w
154    const int KH = kernel->ne[1];  // kernel_h
155    const int IC = input->ne[2];   // input_channels
156    const int OC = kernel->ne[3];  // ouptut_chanles
157    const int B  = input->ne[3];   // n_batches
158
159    const int64_t total  = B * OC * OH * OW;
160    conv_params   params = { IW, IH, OW, OH, KW, KH, ST_X, ST_Y, PD_X, PD_Y, DL_X, DL_Y, IC, OC, B, total };
161
162    if (kernel->type == GGML_TYPE_F16) {
163        conv2d_cuda_f16(X_D, (half *) K_D, Y_D, params, st);
164    } else {
165        conv2d_cuda_f32(X_D, K_D, Y_D, params, st);
166    }
167}