char-fork/llama.cpp

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

Neo Zhang[SYCL] support OPs: conv_2d, conv_2d_dw, conv2d_transpose (#24600)6f1034b32

char/aliceai-t5-moe
5.6 KiB150 linesraw
1#include "conv2d.hpp"
2#include "convert.hpp"
3
4struct conv2d_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 conv2d_kernel_bounds {
17    int64_t y_min, y_max;
18    int64_t x_min, x_max;
19};
20
21static inline int64_t conv2d_max64(int64_t a, int64_t b) {
22    return (a > b) ? a : b;
23}
24
25static inline int64_t conv2d_min64(int64_t a, int64_t b) {
26    return (a < b) ? a : b;
27}
28
29static inline conv2d_kernel_bounds calculate_kernel_bounds(int64_t out_x, int64_t out_y, const conv2d_params & P) {
30    conv2d_kernel_bounds bounds;
31    bounds.y_min = conv2d_max64(0, (P.PD_Y - out_y * P.ST_Y + P.DL_Y - 1) / P.DL_Y);
32    bounds.y_max = conv2d_min64(P.KH, (P.IH + P.PD_Y - out_y * P.ST_Y + P.DL_Y - 1) / P.DL_Y);
33    bounds.x_min = conv2d_max64(0, (P.PD_X - out_x * P.ST_X + P.DL_X - 1) / P.DL_X);
34    bounds.x_max = conv2d_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
38static inline int calculate_input_coord(int64_t out_coord, int64_t kern_coord, int64_t stride,
39                                        int64_t dilation, int64_t padding) {
40    return out_coord * stride + kern_coord * dilation - padding;
41}
42
43// whcn layout helpers (matching ggml tensor memory order)
44static inline int64_t whcn_input_index(int64_t n, int64_t c, int64_t y, int64_t x, const conv2d_params & P) {
45    return n * (P.IC * P.IW * P.IH) + c * P.IW * P.IH + y * P.IW + x;
46}
47
48static inline int64_t whcn_kernel_index(int64_t c_out, int64_t c_in, int64_t ky, int64_t kx, const conv2d_params & P) {
49    return c_out * (P.IC * P.KH * P.KW) + c_in * (P.KH * P.KW) + ky * P.KW + kx;
50}
51
52static inline int64_t whcn_output_index(int64_t n, int64_t c, int64_t y, int64_t x, const conv2d_params & P) {
53    return n * (P.OC * P.OW * P.OH) + c * P.OW * P.OH + y * P.OW + x;
54}
55
56template <typename T>
57static void conv2d_kernel(const float * input, const T * kernel, float * output,
58                          const conv2d_params P, const sycl::nd_item<3> & item_ct1) {
59    const int64_t global_idx = item_ct1.get_local_id(2) +
60                               item_ct1.get_group(2) * item_ct1.get_local_range(2);
61
62    if (global_idx >= P.TOTAL) {
63        return;
64    }
65
66    const int64_t out_x  = global_idx % P.OW;
67    const int64_t out_y  = (global_idx / P.OW) % P.OH;
68    const int64_t c_out  = (global_idx / (P.OW * P.OH)) % P.OC;
69    const int64_t n      = global_idx / (P.OW * P.OH * P.OC);
70
71    float acc = 0.0f;
72
73    const conv2d_kernel_bounds bounds = calculate_kernel_bounds(out_x, out_y, P);
74
75    for (int64_t c_in = 0; c_in < P.IC; ++c_in) {
76        for (int64_t ky = bounds.y_min; ky < bounds.y_max; ++ky) {
77            const int64_t in_y = calculate_input_coord(out_y, ky, P.ST_Y, P.DL_Y, P.PD_Y);
78            for (int64_t kx = bounds.x_min; kx < bounds.x_max; ++kx) {
79                const int64_t in_x = calculate_input_coord(out_x, kx, P.ST_X, P.DL_X, P.PD_X);
80                const float input_val  = input[whcn_input_index(n, c_in, in_y, in_x, P)];
81                const T     kernel_val = kernel[whcn_kernel_index(c_out, c_in, ky, kx, P)];
82                acc += input_val * ggml_sycl_cast<float>(kernel_val);
83            }
84        }
85    }
86
87    output[whcn_output_index(n, c_out, out_y, out_x, P)] = acc;
88}
89
90template <typename T>
91static void conv2d_sycl(const float * X_D, const T * K_D, float * Y_D,
92                        const conv2d_params P, const queue_ptr & stream) {
93    const int num_blocks = (P.TOTAL + SYCL_CONV2D_BLOCK_SIZE - 1) / SYCL_CONV2D_BLOCK_SIZE;
94    const sycl::range<3> block_dims(1, 1, SYCL_CONV2D_BLOCK_SIZE);
95    const sycl::range<3> block_nums(1, 1, num_blocks);
96    stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
97        [=](sycl::nd_item<3> item_ct1) {
98            conv2d_kernel<T>(X_D, K_D, Y_D, P, item_ct1);
99        });
100}
101
102void ggml_sycl_op_conv2d(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
103    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
104
105    const ggml_tensor * kernel = dst->src[0];
106    const ggml_tensor * input  = dst->src[1];
107    const float *       K_D    = (const float *) kernel->data;
108    const float *       X_D    = (const float *) input->data;
109    float *             Y_D    = (float *) dst->data;
110
111    GGML_ASSERT(ggml_is_contiguous(kernel));
112    GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32);
113    GGML_ASSERT(input->type == GGML_TYPE_F32);
114    GGML_ASSERT(dst->type == GGML_TYPE_F32);
115
116    // same number of input channels
117    GGML_ASSERT(input->ne[2] == kernel->ne[2]);
118
119    const queue_ptr stream = ctx.stream();
120
121    const int32_t * p    = (const int32_t *) dst->op_params;
122    const int       ST_X = p[0];
123    const int       ST_Y = p[1];
124    const int       PD_X = p[2];
125    const int       PD_Y = p[3];
126    const int       DL_X = p[4];
127    const int       DL_Y = p[5];
128
129    // no cwhn layout support
130    GGML_ASSERT(p[6] == 0);
131
132    const int IW = input->ne[0];
133    const int IH = input->ne[1];
134    const int OW = dst->ne[0];
135    const int OH = dst->ne[1];
136    const int KW = kernel->ne[0];
137    const int KH = kernel->ne[1];
138    const int IC = input->ne[2];
139    const int OC = kernel->ne[3];
140    const int B  = input->ne[3];
141
142    const int64_t     total  = (int64_t) B * OC * OH * OW;
143    const conv2d_params params = { IW, IH, OW, OH, KW, KH, ST_X, ST_Y, PD_X, PD_Y, DL_X, DL_Y, IC, OC, B, total };
144
145    if (kernel->type == GGML_TYPE_F16) {
146        conv2d_sycl<sycl::half>(X_D, (const sycl::half *) K_D, Y_D, params, stream);
147    } else {
148        conv2d_sycl<float>(X_D, K_D, Y_D, params, stream);
149    }
150}