char-fork/llama.cpp

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

Frosty40sycl: fuse rms_norm+mul+add and add+add residual chains (#27610)6703d7894

char/aliceai-t5-moe
29.6 KiB650 linesraw
1#include "binbcast.hpp"
2
3#include <algorithm>
4#include <cstddef>
5#include <cstdint>
6#include <sycl/sycl.hpp>
7
8#include "ggml.h"
9
10template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t>
11static void k_bin_bcast(const src0_t * src0, const src1_t * src1, dst_t * dst,
12        int ne0, int ne1, int ne2, int ne3,
13        int ne10, int ne11, int ne12, int ne13,
14        /*int s0, */ int s1,  int s2,  int s3,
15        int s00, int s01, int s02, int s03,
16        int s10, int s11, int s12, int s13,
17        const sycl::nd_item<3> &item_ct1) {
18    const int i0s = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
19                    item_ct1.get_local_id(2);
20    const int i1 = (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
21                    item_ct1.get_local_id(1));
22    const int i2 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) +
23                    item_ct1.get_local_id(0)) /
24                   ne3;
25    const int i3 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) +
26                    item_ct1.get_local_id(0)) %
27                   ne3;
28
29    if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
30        return;
31    }
32
33    const int i11 = i1 % ne11;
34    const int i12 = i2 % ne12;
35    const int i13 = i3 % ne13;
36
37    const size_t i_src0 =  i3*s03 +  i2*s02 +  i1*s01;
38    const size_t i_src1 = i13*s13 + i12*s12 + i11*s11;
39    const size_t i_dst  =  i3*s3  +  i2*s2  +  i1*s1;
40
41    const src0_t * src0_row = src0 + i_src0;
42    const src1_t * src1_row = src1 + i_src1;
43    dst_t * dst_row = dst + i_dst;
44
45    for (int i0 = i0s; i0 < ne0;
46         i0 += item_ct1.get_local_range(2) * item_ct1.get_group_range(2)) {
47        const int i10 = i0 % ne10;
48        dst_row[i0] = (dst_t)bin_op(src0 ? (float)src0_row[i0*s00] : 0.0f, (float)src1_row[i10*s10]);
49    }
50}
51
52template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t>
53static void k_bin_bcast_unravel(const src0_t * src0, const src1_t * src1, dst_t * dst,
54        int ne0, int ne1, int ne2, int ne3,
55        int ne10, int ne11, int ne12, int ne13,
56        /*int s0, */ int s1,  int s2,  int s3,
57        int s00, int s01, int s02, int s03,
58        int s10, int s11, int s12, int s13,
59        const sycl::nd_item<3> &item_ct1) {
60
61    const int i = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
62                  item_ct1.get_local_id(2);
63
64    const int i3 = i/(ne2*ne1*ne0);
65    const int i2 = (i/(ne1*ne0)) % ne2;
66    const int i1 = (i/ne0) % ne1;
67    const int i0 = i % ne0;
68
69    if (i0 >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
70        return;
71    }
72
73    const int i11 = i1 % ne11;
74    const int i12 = i2 % ne12;
75    const int i13 = i3 % ne13;
76
77    const size_t i_src0 =  i3*s03 +  i2*s02 +  i1*s01;
78    const size_t i_src1 = i13*s13 + i12*s12 + i11*s11;
79    const size_t i_dst  =  i3*s3  +  i2*s2  +  i1*s1;
80
81    const src0_t * src0_row = src0 + i_src0;
82    const src1_t * src1_row = src1 + i_src1;
83    dst_t * dst_row = dst + i_dst;
84
85    const int i10 = i0 % ne10;
86    dst_row[i0] = (dst_t)bin_op(src0 ? (float)src0_row[i0*s00] : 0.0f, (float)src1_row[i10*s10]);
87}
88
89
90template<float (*bin_op)(const float, const float)>
91struct bin_bcast_sycl {
92    template <typename src0_t, typename src1_t, typename dst_t>
93    void operator()(const src0_t * src0_dd, const src1_t * src1_dd, dst_t * dst_dd, const int64_t ne00,
94                    const int64_t ne01, const int64_t ne02, const int64_t ne03, const int64_t ne10, const int64_t ne11,
95                    const int64_t ne12, const int64_t ne13, const int64_t ne0, const int64_t ne1, const int64_t ne2,
96                    const int64_t ne3, const size_t nb00, const size_t nb01, const size_t nb02, const size_t nb03,
97                    const size_t nb10, const size_t nb11, const size_t nb12, const size_t nb13, const size_t nb0,
98                    const size_t nb1, const size_t nb2, const size_t nb3, const bool src0_is_contiguous,
99                    const bool src1_is_contiguous, const bool src0_is_permuted, const bool src1_is_permuted,
100                    queue_ptr stream) {
101        int nr0 = ne10 / ne0;
102        int nr1 = ne11/ne1;
103        int nr2 = ne12/ne2;
104        int nr3 = ne13/ne3;
105
106        int nr[4] = { nr0, nr1, nr2, nr3 };
107
108        // collapse dimensions until first broadcast dimension
109        int64_t cne[] = {ne0, ne1, ne2, ne3};
110        int64_t cne0[] = {ne00, ne01, ne02, ne03};
111        int64_t cne1[] = {ne10, ne11, ne12, ne13};
112        size_t cnb[] = {nb0, nb1, nb2, nb3};
113        size_t cnb0[] = {nb00, nb01, nb02, nb03};
114        size_t cnb1[] = {nb10, nb11, nb12, nb13};
115        auto collapse = [](int64_t cne[]) {
116            cne[0] *= cne[1];
117            cne[1] = cne[2];
118            cne[2] = cne[3];
119            cne[3] = 1;
120        };
121
122        auto collapse_nb = [](size_t cnb[], int64_t cne[]) {
123            cnb[1] *= cne[1];
124            cnb[2] *= cne[2];
125            cnb[3] *= cne[3];
126        };
127
128        if (src0_is_contiguous && src1_is_contiguous && !src0_is_permuted && !src1_is_permuted) {
129            for (int i = 0; i < 4; i++) {
130                if (nr[i] != 1) {
131                    break;
132                }
133                if (i > 0) {
134                    collapse_nb(cnb, cne);
135                    collapse_nb(cnb0, cne0);
136                    collapse_nb(cnb1, cne1);
137                    collapse(cne);
138                    collapse(cne0);
139                    collapse(cne1);
140                }
141            }
142        }
143        {
144            int64_t ne0 = cne[0];
145            int64_t ne1 = cne[1];
146            int64_t ne2 = cne[2];
147            int64_t ne3 = cne[3];
148
149            int64_t ne10 = cne1[0];
150            int64_t ne11 = cne1[1];
151            int64_t ne12 = cne1[2];
152            int64_t ne13 = cne1[3];
153
154            size_t nb0 = cnb[0];
155            size_t nb1 = cnb[1];
156            size_t nb2 = cnb[2];
157            size_t nb3 = cnb[3];
158
159            size_t nb00 = cnb0[0];
160            size_t nb01 = cnb0[1];
161            size_t nb02 = cnb0[2];
162            size_t nb03 = cnb0[3];
163
164            size_t nb10 = cnb1[0];
165            size_t nb11 = cnb1[1];
166            size_t nb12 = cnb1[2];
167            size_t nb13 = cnb1[3];
168
169            // size_t s0 = nb0 / sizeof(dst_t);
170            size_t s1 = nb1 / sizeof(dst_t);
171            size_t s2 = nb2 / sizeof(dst_t);
172            size_t s3 = nb3 / sizeof(dst_t);
173
174            size_t s10 = nb10 / sizeof(src1_t);
175            size_t s11 = nb11 / sizeof(src1_t);
176            size_t s12 = nb12 / sizeof(src1_t);
177            size_t s13 = nb13 / sizeof(src1_t);
178
179            size_t s00 = nb00 / sizeof(src0_t);
180            size_t s01 = nb01 / sizeof(src0_t);
181            size_t s02 = nb02 / sizeof(src0_t);
182            size_t s03 = nb03 / sizeof(src0_t);
183
184            GGML_UNUSED(s00);
185
186            GGML_ASSERT(nb0 % sizeof(dst_t) == 0);
187            GGML_ASSERT(nb1 % sizeof(dst_t) == 0);
188            GGML_ASSERT(nb2 % sizeof(dst_t) == 0);
189            GGML_ASSERT(nb3 % sizeof(dst_t) == 0);
190
191            GGML_ASSERT(nb00 % sizeof(src0_t) == 0);
192            GGML_ASSERT(nb01 % sizeof(src0_t) == 0);
193            GGML_ASSERT(nb02 % sizeof(src0_t) == 0);
194            GGML_ASSERT(nb03 % sizeof(src0_t) == 0);
195
196            GGML_ASSERT(nb10 % sizeof(src1_t) == 0);
197            GGML_ASSERT(nb11 % sizeof(src1_t) == 0);
198            GGML_ASSERT(nb12 % sizeof(src1_t) == 0);
199            GGML_ASSERT(nb13 % sizeof(src1_t) == 0);
200
201            const int block_size = 128;
202
203            int64_t hne0 = std::max(ne0/2LL, 1LL);
204
205            sycl::range<3> block_dims(1, 1, 1);
206            block_dims[2] = std::min<unsigned int>(hne0, block_size);
207            block_dims[1] = std::min<unsigned int>(
208                ne1, block_size / (unsigned int)block_dims[2]);
209            block_dims[0] = std::min(
210                std::min<unsigned int>(
211                    ne2 * ne3, block_size / (unsigned int)block_dims[2] /
212                                   (unsigned int)block_dims[1]),
213                64U);
214
215            sycl::range<3> block_nums(
216                (ne2 * ne3 + block_dims[0] - 1) / block_dims[0],
217                (ne1 + block_dims[1] - 1) / block_dims[1],
218                (hne0 + block_dims[2] - 1) / block_dims[2]);
219
220            if (block_nums[0] > 65535) {
221                // this is the maximum number of blocks in z direction, fallback to 1D grid kernel
222                int block_num = (ne0*ne1*ne2*ne3 + block_size - 1) / block_size;
223                {
224                    dpct::has_capability_or_fail(stream->get_device(),
225                                                 {sycl::aspect::fp16});
226
227                    stream->parallel_for(
228                        sycl::nd_range<3>(sycl::range<3>(1, 1, block_num) *
229                                              sycl::range<3>(1, 1, block_size),
230                                          sycl::range<3>(1, 1, block_size)),
231                        [=](sycl::nd_item<3> item_ct1) {
232                            k_bin_bcast_unravel<bin_op>(
233                                src0_dd, src1_dd, dst_dd, ne0, ne1, ne2, ne3,
234                                ne10, ne11, ne12, ne13, s1, s2, s3, s00, s01, s02,
235                                s03, s10, s11, s12, s13, item_ct1);
236                        });
237                }
238            } else {
239                /*
240                DPCT1049:16: The work-group size passed to the SYCL kernel may
241                exceed the limit. To get the device limit, query
242                info::device::max_work_group_size. Adjust the work-group size if
243                needed.
244                */
245                dpct::has_capability_or_fail(stream->get_device(),
246                                             {sycl::aspect::fp16});
247
248                stream->parallel_for(
249                    sycl::nd_range<3>(block_nums * block_dims, block_dims),
250                    [=](sycl::nd_item<3> item_ct1) {
251                        k_bin_bcast<bin_op>(src0_dd, src1_dd, dst_dd, ne0, ne1,
252                                            ne2, ne3, ne10, ne11, ne12, ne13,
253                                            s1, s2, s3, s00, s01, s02, s03, s10, s11, s12, s13,
254                                            item_ct1);
255                    });
256            }
257        }
258    }
259};
260
261template <class op>
262inline void ggml_sycl_op_bin_bcast(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1,
263                                   ggml_tensor * dst) {
264    dpct::queue_ptr main_stream = ctx.stream();
265    GGML_TENSOR_BINARY_OP_LOCALS
266
267    if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
268        op()((const float *) src0->data, (const float *) src1->data, (float *) dst->data, ne00, ne01, ne02, ne03, ne10,
269             ne11, ne12, ne13, ne0, ne1, ne2, ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1, nb2, nb3,
270             ggml_is_contiguous(src0), ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1), main_stream);
271    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
272        op()((const sycl::half *) src0->data, (const sycl::half *) src1->data, (sycl::half *) dst->data, ne00, ne01,
273             ne02, ne03, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13,
274             nb0, nb1, nb2, nb3, ggml_is_contiguous(src0), ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1),
275             main_stream);
276    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
277        op()((const sycl::half *) src0->data, (const float *) src1->data, (sycl::half *) dst->data, ne00, ne01, ne02,
278             ne03, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1,
279             nb2, nb3, ggml_is_contiguous(src0), ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1),
280             main_stream);
281    } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) {
282        op()((const int32_t *) src0->data, (const int32_t *) src1->data, (int32_t *) dst->data, ne00, ne01, ne02, ne03,
283             ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1, nb2,
284             nb3, ggml_is_contiguous(src0), ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1),
285             main_stream);
286    } else if (src0->type == GGML_TYPE_I16 && src1->type == GGML_TYPE_I16 && dst->type == GGML_TYPE_I16) {
287        op()((const int16_t *) src0->data, (const int16_t *) src1->data, (int16_t *) dst->data, ne00, ne01, ne02, ne03,
288             ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1, nb2,
289             nb3, ggml_is_contiguous(src0), ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1),
290             main_stream);
291#ifdef GGML_SYCL_HAS_BF16
292    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) {
293        op()((const sycl::ext::oneapi::bfloat16 *) src0->data, (const sycl::ext::oneapi::bfloat16 *) src1->data,
294             (sycl::ext::oneapi::bfloat16 *) dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, ne0, ne1, ne2,
295             ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1, nb2, nb3, ggml_is_contiguous(src0),
296             ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1), main_stream);
297    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) {
298        op()((const sycl::ext::oneapi::bfloat16 *) src0->data, (const float *) src1->data,
299             (sycl::ext::oneapi::bfloat16 *) dst->data, ne00, ne01, ne02, ne03, ne10, ne11, ne12, ne13, ne0, ne1, ne2,
300             ne3, nb00, nb01, nb02, nb03, nb10, nb11, nb12, nb13, nb0, nb1, nb2, nb3, ggml_is_contiguous(src0),
301             ggml_is_contiguous(src1), ggml_is_permuted(src0), ggml_is_permuted(src1), main_stream);
302#endif
303    } else {
304        fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__, ggml_type_name(dst->type),
305                ggml_type_name(src0->type), ggml_type_name(src1->type));
306        GGML_ABORT("fatal error");
307    }
308}
309
310inline void ggml_sycl_op_add(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
311
312    ggml_sycl_op_bin_bcast<bin_bcast_sycl<op_add>>(ctx, dst->src[0], dst->src[1], dst);
313}
314
315inline void ggml_sycl_op_sub(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
316
317    ggml_sycl_op_bin_bcast<bin_bcast_sycl<op_sub>>(ctx, dst->src[0], dst->src[1], dst);
318}
319
320inline void ggml_sycl_op_mul(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
321
322    ggml_sycl_op_bin_bcast<bin_bcast_sycl<op_mul>>(ctx, dst->src[0], dst->src[1], dst);
323}
324
325inline void ggml_sycl_op_div(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
326
327    ggml_sycl_op_bin_bcast<bin_bcast_sycl<op_div>>(ctx, dst->src[0], dst->src[1], dst);
328}
329
330inline void ggml_sycl_op_repeat(ggml_backend_sycl_context & ctx, ggml_tensor *dst) {
331    ggml_sycl_op_bin_bcast<bin_bcast_sycl<op_repeat>>(ctx, dst, dst->src[0], dst);
332}
333
334
335void ggml_sycl_add(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
336    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
337    ggml_sycl_op_add(ctx, dst);
338}
339
340void ggml_sycl_sub(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
341    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
342    ggml_sycl_op_sub(ctx, dst);
343}
344
345void ggml_sycl_mul(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
346    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
347    ggml_sycl_op_mul(ctx, dst);
348}
349
350void ggml_sycl_div(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
351    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
352    ggml_sycl_op_div(ctx, dst);
353}
354
355void ggml_sycl_repeat(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
356    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
357    ggml_sycl_op_repeat(ctx, dst);
358}
359
360// fused ADD+ADD: dst = (src0 + src1) + src2. Same indexing as k_bin_bcast, so mixed
361// types, broadcast, and non-contiguous layouts that add() already handles also fuse.
362template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename src2_t, typename dst_t>
363static void k_bin_bcast3(const src0_t * src0, const src1_t * src1, const src2_t * src2, dst_t * dst,
364        int ne0, int ne1, int ne2, int ne3,
365        int ne10, int ne11, int ne12, int ne13,
366        int ne20, int ne21, int ne22, int ne23,
367        int s1, int s2, int s3,
368        int s00, int s01, int s02, int s03,
369        int s10, int s11, int s12, int s13,
370        int s20, int s21, int s22, int s23,
371        const sycl::nd_item<3> & item_ct1) {
372    const int i0s = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
373                    item_ct1.get_local_id(2);
374    const int i1 = (item_ct1.get_local_range(1) * item_ct1.get_group(1) +
375                    item_ct1.get_local_id(1));
376    const int i2 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) +
377                    item_ct1.get_local_id(0)) /
378                   ne3;
379    const int i3 = (item_ct1.get_local_range(0) * item_ct1.get_group(0) +
380                    item_ct1.get_local_id(0)) %
381                   ne3;
382
383    if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
384        return;
385    }
386
387    const int i11 = i1 % ne11;
388    const int i12 = i2 % ne12;
389    const int i13 = i3 % ne13;
390    const int i21 = i1 % ne21;
391    const int i22 = i2 % ne22;
392    const int i23 = i3 % ne23;
393
394    const size_t i_src0 = i3 * s03 + i2 * s02 + i1 * s01;
395    const size_t i_src1 = i13 * s13 + i12 * s12 + i11 * s11;
396    const size_t i_src2 = i23 * s23 + i22 * s22 + i21 * s21;
397    const size_t i_dst  = i3 * s3 + i2 * s2 + i1 * s1;
398
399    const src0_t * src0_row = src0 + i_src0;
400    const src1_t * src1_row = src1 + i_src1;
401    const src2_t * src2_row = src2 + i_src2;
402    dst_t *        dst_row  = dst + i_dst;
403
404    for (int i0 = i0s; i0 < ne0;
405         i0 += item_ct1.get_local_range(2) * item_ct1.get_group_range(2)) {
406        const int   i10 = i0 % ne10;
407        const int   i20 = i0 % ne20;
408        const float acc = bin_op((float) src0_row[i0 * s00], (float) src1_row[i10 * s10]);
409        dst_row[i0]     = (dst_t) bin_op(acc, (float) src2_row[i20 * s20]);
410    }
411}
412
413template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename src2_t, typename dst_t>
414static void k_bin_bcast3_unravel(const src0_t * src0, const src1_t * src1, const src2_t * src2, dst_t * dst,
415        int ne0, int ne1, int ne2, int ne3,
416        int ne10, int ne11, int ne12, int ne13,
417        int ne20, int ne21, int ne22, int ne23,
418        int s1, int s2, int s3,
419        int s00, int s01, int s02, int s03,
420        int s10, int s11, int s12, int s13,
421        int s20, int s21, int s22, int s23,
422        const sycl::nd_item<3> & item_ct1) {
423    const int i = item_ct1.get_local_range(2) * item_ct1.get_group(2) +
424                  item_ct1.get_local_id(2);
425
426    const int i3 = i / (ne2 * ne1 * ne0);
427    const int i2 = (i / (ne1 * ne0)) % ne2;
428    const int i1 = (i / ne0) % ne1;
429    const int i0 = i % ne0;
430
431    if (i0 >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
432        return;
433    }
434
435    const int i11 = i1 % ne11;
436    const int i12 = i2 % ne12;
437    const int i13 = i3 % ne13;
438    const int i21 = i1 % ne21;
439    const int i22 = i2 % ne22;
440    const int i23 = i3 % ne23;
441
442    const size_t i_src0 = i3 * s03 + i2 * s02 + i1 * s01;
443    const size_t i_src1 = i13 * s13 + i12 * s12 + i11 * s11;
444    const size_t i_src2 = i23 * s23 + i22 * s22 + i21 * s21;
445    const size_t i_dst  = i3 * s3 + i2 * s2 + i1 * s1;
446
447    const int   i10 = i0 % ne10;
448    const int   i20 = i0 % ne20;
449    const float acc = bin_op((float) src0[i_src0 + i0 * s00], (float) src1[i_src1 + i10 * s10]);
450    dst[i_dst + i0] = (dst_t) bin_op(acc, (float) src2[i_src2 + i20 * s20]);
451}
452
453template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename src2_t, typename dst_t>
454static void launch_bin_bcast3(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1,
455                              const ggml_tensor * src2, ggml_tensor * dst) {
456    dpct::queue_ptr stream = ctx.stream();
457    SYCL_CHECK(ggml_sycl_set_device(ctx.device));
458
459    GGML_TENSOR_TERNARY_OP_LOCALS
460
461    int nr1[4] = { (int) (ne10 / ne0), (int) (ne11 / ne1), (int) (ne12 / ne2), (int) (ne13 / ne3) };
462    int nr2[4] = { (int) (ne20 / ne0), (int) (ne21 / ne1), (int) (ne22 / ne2), (int) (ne23 / ne3) };
463
464    int64_t cne[]  = { ne0, ne1, ne2, ne3 };
465    int64_t cne0[] = { ne00, ne01, ne02, ne03 };
466    int64_t cne1[] = { ne10, ne11, ne12, ne13 };
467    int64_t cne2[] = { ne20, ne21, ne22, ne23 };
468    size_t  cnb[]  = { nb0, nb1, nb2, nb3 };
469    size_t  cnb0[] = { nb00, nb01, nb02, nb03 };
470    size_t  cnb1[] = { nb10, nb11, nb12, nb13 };
471    size_t  cnb2[] = { nb20, nb21, nb22, nb23 };
472
473    auto collapse = [](int64_t cne[]) {
474        cne[0] *= cne[1];
475        cne[1] = cne[2];
476        cne[2] = cne[3];
477        cne[3] = 1;
478    };
479
480    auto collapse_nb = [](size_t cnb[], int64_t cne[]) {
481        cnb[1] *= cne[1];
482        cnb[2] *= cne[2];
483        cnb[3] *= cne[3];
484    };
485
486    const bool can_collapse = ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(src2) &&
487                              !ggml_is_permuted(src0) && !ggml_is_permuted(src1) && !ggml_is_permuted(src2);
488    if (can_collapse) {
489        for (int i = 0; i < 4; i++) {
490            if (nr1[i] != 1 || nr2[i] != 1) {
491                break;
492            }
493            if (i > 0) {
494                collapse_nb(cnb, cne);
495                collapse_nb(cnb0, cne0);
496                collapse_nb(cnb1, cne1);
497                collapse_nb(cnb2, cne2);
498                collapse(cne);
499                collapse(cne0);
500                collapse(cne1);
501                collapse(cne2);
502            }
503        }
504    }
505
506    {
507        int64_t ne0 = cne[0];
508        int64_t ne1 = cne[1];
509        int64_t ne2 = cne[2];
510        int64_t ne3 = cne[3];
511
512        int64_t ne10 = cne1[0];
513        int64_t ne11 = cne1[1];
514        int64_t ne12 = cne1[2];
515        int64_t ne13 = cne1[3];
516
517        int64_t ne20 = cne2[0];
518        int64_t ne21 = cne2[1];
519        int64_t ne22 = cne2[2];
520        int64_t ne23 = cne2[3];
521
522        size_t s1 = cnb[1] / sizeof(dst_t);
523        size_t s2 = cnb[2] / sizeof(dst_t);
524        size_t s3 = cnb[3] / sizeof(dst_t);
525
526        size_t s00 = cnb0[0] / sizeof(src0_t);
527        size_t s01 = cnb0[1] / sizeof(src0_t);
528        size_t s02 = cnb0[2] / sizeof(src0_t);
529        size_t s03 = cnb0[3] / sizeof(src0_t);
530
531        size_t s10 = cnb1[0] / sizeof(src1_t);
532        size_t s11 = cnb1[1] / sizeof(src1_t);
533        size_t s12 = cnb1[2] / sizeof(src1_t);
534        size_t s13 = cnb1[3] / sizeof(src1_t);
535
536        size_t s20 = cnb2[0] / sizeof(src2_t);
537        size_t s21 = cnb2[1] / sizeof(src2_t);
538        size_t s22 = cnb2[2] / sizeof(src2_t);
539        size_t s23 = cnb2[3] / sizeof(src2_t);
540
541        GGML_ASSERT(cnb[0] % sizeof(dst_t) == 0 && cnb[1] % sizeof(dst_t) == 0 && cnb[2] % sizeof(dst_t) == 0 &&
542                    cnb[3] % sizeof(dst_t) == 0);
543        GGML_ASSERT(cnb0[0] % sizeof(src0_t) == 0 && cnb0[1] % sizeof(src0_t) == 0 && cnb0[2] % sizeof(src0_t) == 0 &&
544                    cnb0[3] % sizeof(src0_t) == 0);
545        GGML_ASSERT(cnb1[0] % sizeof(src1_t) == 0 && cnb1[1] % sizeof(src1_t) == 0 && cnb1[2] % sizeof(src1_t) == 0 &&
546                    cnb1[3] % sizeof(src1_t) == 0);
547        GGML_ASSERT(cnb2[0] % sizeof(src2_t) == 0 && cnb2[1] % sizeof(src2_t) == 0 && cnb2[2] % sizeof(src2_t) == 0 &&
548                    cnb2[3] % sizeof(src2_t) == 0);
549
550        const src0_t * src0_dd = (const src0_t *) src0->data;
551        const src1_t * src1_dd = (const src1_t *) src1->data;
552        const src2_t * src2_dd = (const src2_t *) src2->data;
553        dst_t *        dst_dd  = (dst_t *) dst->data;
554
555        const int block_size = 128;
556        int64_t   hne0       = std::max(ne0 / 2LL, 1LL);
557
558        sycl::range<3> block_dims(1, 1, 1);
559        block_dims[2] = std::min<unsigned int>(hne0, block_size);
560        block_dims[1] = std::min<unsigned int>(ne1, block_size / (unsigned int) block_dims[2]);
561        block_dims[0] = std::min(std::min<unsigned int>(ne2 * ne3,
562                                                        block_size / (unsigned int) block_dims[2] /
563                                                            (unsigned int) block_dims[1]),
564                                 64U);
565
566        sycl::range<3> block_nums((ne2 * ne3 + block_dims[0] - 1) / block_dims[0],
567                                  (ne1 + block_dims[1] - 1) / block_dims[1],
568                                  (hne0 + block_dims[2] - 1) / block_dims[2]);
569
570        dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
571
572        if (block_nums[0] > 65535) {
573            int block_num = (ne0 * ne1 * ne2 * ne3 + block_size - 1) / block_size;
574            stream->parallel_for(
575                sycl::nd_range<3>(sycl::range<3>(1, 1, block_num) * sycl::range<3>(1, 1, block_size),
576                                  sycl::range<3>(1, 1, block_size)),
577                [=](sycl::nd_item<3> item_ct1) {
578                    k_bin_bcast3_unravel<bin_op>(src0_dd, src1_dd, src2_dd, dst_dd, ne0, ne1, ne2, ne3, ne10, ne11,
579                                                 ne12, ne13, ne20, ne21, ne22, ne23, s1, s2, s3, s00, s01, s02, s03,
580                                                 s10, s11, s12, s13, s20, s21, s22, s23, item_ct1);
581                });
582        } else {
583            stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
584                                 [=](sycl::nd_item<3> item_ct1) {
585                                     k_bin_bcast3<bin_op>(src0_dd, src1_dd, src2_dd, dst_dd, ne0, ne1, ne2, ne3, ne10,
586                                                          ne11, ne12, ne13, ne20, ne21, ne22, ne23, s1, s2, s3, s00,
587                                                          s01, s02, s03, s10, s11, s12, s13, s20, s21, s22, s23,
588                                                          item_ct1);
589                                 });
590        }
591    }
592}
593
594void ggml_sycl_op_add_add_fused(ggml_backend_sycl_context & ctx, ggml_tensor * add0, ggml_tensor * add1) {
595    const ggml_tensor * src0 = add0->src[0];
596    const ggml_tensor * src1 = add0->src[1];
597    const ggml_tensor * src2 = add1->src[1];
598    ggml_tensor *       dst  = add1;
599
600    GGML_ASSERT(add1->src[0] == add0);
601    GGML_ASSERT(ggml_sycl_add_kernel_supports(src0->type, src1->type, add0->type));
602    GGML_ASSERT(ggml_sycl_add_kernel_supports(add0->type, src2->type, dst->type));
603
604    if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 &&
605        dst->type == GGML_TYPE_F32) {
606        launch_bin_bcast3<op_add, float, float, float, float>(ctx, src0, src1, src2, dst);
607    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && src2->type == GGML_TYPE_F16 &&
608               dst->type == GGML_TYPE_F16) {
609        launch_bin_bcast3<op_add, sycl::half, sycl::half, sycl::half, sycl::half>(ctx, src0, src1, src2, dst);
610    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 &&
611               dst->type == GGML_TYPE_F16) {
612        launch_bin_bcast3<op_add, sycl::half, float, float, sycl::half>(ctx, src0, src1, src2, dst);
613    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && src2->type == GGML_TYPE_F32 &&
614               dst->type == GGML_TYPE_F16) {
615        launch_bin_bcast3<op_add, sycl::half, sycl::half, float, sycl::half>(ctx, src0, src1, src2, dst);
616    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F16 &&
617               dst->type == GGML_TYPE_F16) {
618        launch_bin_bcast3<op_add, sycl::half, float, sycl::half, sycl::half>(ctx, src0, src1, src2, dst);
619    } else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32 && src2->type == GGML_TYPE_I32 &&
620               dst->type == GGML_TYPE_I32) {
621        launch_bin_bcast3<op_add, int32_t, int32_t, int32_t, int32_t>(ctx, src0, src1, src2, dst);
622    } else if (src0->type == GGML_TYPE_I16 && src1->type == GGML_TYPE_I16 && src2->type == GGML_TYPE_I16 &&
623               dst->type == GGML_TYPE_I16) {
624        launch_bin_bcast3<op_add, int16_t, int16_t, int16_t, int16_t>(ctx, src0, src1, src2, dst);
625#ifdef GGML_SYCL_HAS_BF16
626    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && src2->type == GGML_TYPE_BF16 &&
627               dst->type == GGML_TYPE_BF16) {
628        launch_bin_bcast3<op_add, sycl::ext::oneapi::bfloat16, sycl::ext::oneapi::bfloat16,
629                          sycl::ext::oneapi::bfloat16, sycl::ext::oneapi::bfloat16>(ctx, src0, src1, src2, dst);
630    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 &&
631               dst->type == GGML_TYPE_BF16) {
632        launch_bin_bcast3<op_add, sycl::ext::oneapi::bfloat16, float, float, sycl::ext::oneapi::bfloat16>(
633            ctx, src0, src1, src2, dst);
634    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && src2->type == GGML_TYPE_F32 &&
635               dst->type == GGML_TYPE_BF16) {
636        launch_bin_bcast3<op_add, sycl::ext::oneapi::bfloat16, sycl::ext::oneapi::bfloat16, float,
637                          sycl::ext::oneapi::bfloat16>(ctx, src0, src1, src2, dst);
638    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_BF16 &&
639               dst->type == GGML_TYPE_BF16) {
640        launch_bin_bcast3<op_add, sycl::ext::oneapi::bfloat16, float, sycl::ext::oneapi::bfloat16,
641                          sycl::ext::oneapi::bfloat16>(ctx, src0, src1, src2, dst);
642#endif
643    } else {
644        fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s, src1: %s, src2: %s\n", __func__,
645                ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type),
646                ggml_type_name(src2->type));
647        GGML_ABORT("fatal error");
648    }
649}
650