char-fork/llama.cpp

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

fairydreamingggml : address integer overflows in binary ops CUDA implementation (#24706)f728adab6

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
24.2 KiB574 linesraw
1#include "binbcast.cuh"
2#include <cstdint>
3#include <utility>
4
5template<typename T, size_t>
6using type_for_index = T;
7
8static __device__ __forceinline__ float op_repeat(const float a, const float b) {
9    return b;
10    GGML_UNUSED(a);
11}
12
13static __device__ __forceinline__ float op_add(const float a, const float b) {
14    return a + b;
15}
16
17static __device__ __forceinline__ float op_sub(const float a, const float b) {
18    return a - b;
19}
20
21static __device__ __forceinline__ float op_mul(const float a, const float b) {
22    return a * b;
23}
24
25static __device__ __forceinline__ float op_div(const float a, const float b) {
26    return a / b;
27}
28
29template <float (*bin_op)(const float, const float),
30          typename src0_t,
31          typename src1_t,
32          typename dst_t,
33          typename... src1_ptrs>
34static __global__ void k_bin_bcast(const src0_t *         src0,
35                                   const src1_t *         src1,
36                                   dst_t *                dst,
37                                   const uint32_t         ne0,
38                                   const uint32_t         ne1,
39                                   const uint32_t         ne2,
40                                   const uint3            ne3,
41                                   const uint3            ne10,
42                                   const uint3            ne11,
43                                   const uint3            ne12,
44                                   const uint3            ne13,
45                                 /*const uint32_t         s0,*/
46                                   const uint32_t         s1,
47                                   const uint32_t         s2,
48                                   const uint32_t         s3,
49                                   const uint32_t         s00,
50                                   const uint32_t         s01,
51                                   const uint32_t         s02,
52                                   const uint32_t         s03,
53                                   const uint32_t         s10,
54                                   const uint32_t         s11,
55                                   const uint32_t         s12,
56                                   const uint32_t         s13,
57                                   src1_ptrs... src1s) {
58    ggml_cuda_pdl_lc();
59    const uint32_t i0s = blockDim.x * blockIdx.x + threadIdx.x;
60    const uint32_t i1  = (blockDim.y * blockIdx.y + threadIdx.y);
61    const uint32_t i2  = fastdiv((blockDim.z * blockIdx.z + threadIdx.z), ne3);
62    const uint32_t i3  = (blockDim.z * blockIdx.z + threadIdx.z) - (i2 * ne3.z);
63
64    if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3.z) {
65        return;
66    }
67
68    const uint32_t i11 = fastmodulo(i1, ne11);
69    const uint32_t i12 = fastmodulo(i2, ne12);
70    const uint32_t i13 = fastmodulo(i3, ne13);
71
72    const size_t i_src0 = size_t( i3)*s03 + size_t( i2)*s02 + size_t( i1)*s01;
73    const size_t i_src1 = size_t(i13)*s13 + size_t(i12)*s12 + size_t(i11)*s11;
74    const size_t i_dst  = size_t( i3)*s3  + size_t( i2)*s2  + size_t( i1)*s1;
75
76    const src0_t * src0_row = src0 ? (src0 + i_src0) : nullptr;
77    dst_t * dst_row = dst + i_dst;
78
79    const uint32_t s0 = blockDim.x * gridDim.x;
80
81    ggml_cuda_pdl_sync();
82    for (uint32_t i0 = i0s; i0 < ne0; i0 += s0) {
83        const uint32_t i10 = fastmodulo(i0, ne10);
84
85        float result = src0_row ? (float) src0_row[size_t(i0)*s00] : 0.0f;
86        if constexpr (sizeof...(src1_ptrs) > 0) {
87            result = (..., (result = bin_op(result, (float)src1s[i_src1 + size_t(i10)*s10])));
88        } else {
89            result = bin_op(result, (float)src1[i_src1 + size_t(i10)*s10]);
90        }
91
92        dst_row[i0] = (dst_t) result;
93
94        // protect i0 from overflow
95        if (ne0 - i0 <= s0) {
96           break;
97        }
98    }
99}
100
101template <float (*bin_op)(const float, const float),
102          typename src0_t,
103          typename src1_t,
104          typename dst_t,
105          typename... src1_ptrs>
106static __global__ void k_bin_bcast_unravel(const src0_t *         src0,
107                                           const src1_t *         src1,
108                                           dst_t *                dst,
109                                           const uint3            ne0,
110                                           const uint3            ne1,
111                                           const uint3            ne2,
112                                           const uint32_t         ne3,
113                                           const uint3            prod_012,
114                                           const uint3            prod_01,
115                                           const uint3            ne10,
116                                           const uint3            ne11,
117                                           const uint3            ne12,
118                                           const uint3            ne13,
119                                         /*const int              s0,*/
120                                           const uint32_t         s1,
121                                           const uint32_t         s2,
122                                           const uint32_t         s3,
123                                           const uint32_t         s00,
124                                           const uint32_t         s01,
125                                           const uint32_t         s02,
126                                           const uint32_t         s03,
127                                           const uint32_t         s10,
128                                           const uint32_t         s11,
129                                           const uint32_t         s12,
130                                           const uint32_t         s13,
131                                           src1_ptrs... src1s) {
132    const uint32_t i  = blockDim.x*blockIdx.x + threadIdx.x;
133
134    const uint32_t i3 = fastdiv(i, prod_012);
135    const uint32_t i2 = fastdiv(i - i3 * prod_012.z, prod_01);
136    const uint32_t i1 = fastdiv(i - i3 * prod_012.z - i2 * prod_01.z, ne0);
137    const uint32_t i0 = i - i3 * prod_012.z - i2 * prod_01.z - i1 * ne0.z;
138
139    if (i0 >= ne0.z || i1 >= ne1.z || i2 >= ne2.z || i3 >= ne3) {
140        return;
141    }
142
143    const uint32_t i11 = fastmodulo(i1, ne11);
144    const uint32_t i12 = fastmodulo(i2, ne12);
145    const uint32_t i13 = fastmodulo(i3, ne13);
146
147    const size_t i_src0 = size_t( i3)*s03 + size_t( i2)*s02 + size_t( i1)*s01;
148    const size_t i_src1 = size_t(i13)*s13 + size_t(i12)*s12 + size_t(i11)*s11;
149    const size_t i_dst  = size_t( i3)*s3  + size_t( i2)*s2  + size_t( i1)*s1;
150
151    const src0_t * src0_row = src0 ? (src0 + i_src0) : nullptr;
152    dst_t * dst_row = dst + i_dst;
153
154    const uint32_t i10 = fastmodulo(i0, ne10);
155
156    ggml_cuda_pdl_sync();
157    float result = src0_row ? (float) src0_row[size_t(i0)*s00] : 0.0f;
158    if constexpr (sizeof...(src1_ptrs) > 0) {
159        result = (..., (result = bin_op(result, (float)src1s[i_src1 + size_t(i10)*s10])));
160    } else {
161        result = bin_op(result, (float)src1[i_src1 + size_t(i10)*s10]);
162    }
163
164    dst_row[i0] = (dst_t) result;
165}
166
167template <float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t, size_t... I>
168static void launch_bin_bcast_pack(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
169                                  const src0_t * src0_dd, const src1_t * src1_dd, dst_t * dst_dd,
170                                  cudaStream_t stream, std::index_sequence<I...>) {
171    GGML_TENSOR_BINARY_OP_LOCALS
172
173    int nr0 = ne10 / ne0;
174    int nr1 = ne11 / ne1;
175    int nr2 = ne12 / ne2;
176    int nr3 = ne13 / ne3;
177
178    int nr[4] = { nr0, nr1, nr2, nr3 };
179
180    int64_t cne[]  = { ne0, ne1, ne2, ne3 };
181    int64_t cne0[] = { ne00, ne01, ne02, ne03 };
182    int64_t cne1[] = { ne10, ne11, ne12, ne13 };
183
184    size_t cnb[]  = { nb0, nb1, nb2, nb3 };
185    size_t cnb0[] = { nb00, nb01, nb02, nb03 };
186    size_t cnb1[] = { nb10, nb11, nb12, nb13 };
187
188    auto collapse = [](int64_t cne[]) {
189        cne[0] *= cne[1];
190        cne[1] = cne[2];
191        cne[2] = cne[3];
192        cne[3] = 1;
193    };
194
195    auto collapse_nb = [](size_t cnb[], const int64_t cne[]) {
196        cnb[1] *= cne[1];
197        cnb[2] *= cne[2];
198        cnb[3] *= cne[3];
199    };
200
201    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && !ggml_is_permuted(src0) && !ggml_is_permuted(src1)) {
202        for (int i = 0; i < 4; i++) {
203            if (nr[i] != 1) {
204                break;
205            }
206            if (i > 0) {
207                collapse_nb(cnb, cne);
208                collapse_nb(cnb0, cne0);
209                collapse_nb(cnb1, cne1);
210                collapse(cne);
211                collapse(cne0);
212                collapse(cne1);
213            }
214        }
215    }
216
217    {
218        int64_t ne0 = cne[0];
219        int64_t ne1 = cne[1];
220        int64_t ne2 = cne[2];
221        int64_t ne3 = cne[3];
222
223        //int64_t ne00 = cne0[0]; GGML_UNUSED(ne00);
224        //int64_t ne01 = cne0[1]; GGML_UNUSED(ne01);
225        //int64_t ne02 = cne0[2]; GGML_UNUSED(ne02);
226        //int64_t ne03 = cne0[3]; GGML_UNUSED(ne03);
227
228        size_t nb0 = cnb[0];
229        size_t nb1 = cnb[1];
230        size_t nb2 = cnb[2];
231        size_t nb3 = cnb[3];
232
233        size_t nb00 = cnb0[0];
234        size_t nb01 = cnb0[1];
235        size_t nb02 = cnb0[2];
236        size_t nb03 = cnb0[3];
237
238        size_t nb10 = cnb1[0];
239        size_t nb11 = cnb1[1];
240        size_t nb12 = cnb1[2];
241        size_t nb13 = cnb1[3];
242
243      //size_t s0 = nb0 / sizeof(dst_t);
244        size_t s1 = nb1 / sizeof(dst_t);
245        size_t s2 = nb2 / sizeof(dst_t);
246        size_t s3 = nb3 / sizeof(dst_t);
247
248        size_t s10 = nb10 / sizeof(src1_t);
249        size_t s11 = nb11 / sizeof(src1_t);
250        size_t s12 = nb12 / sizeof(src1_t);
251        size_t s13 = nb13 / sizeof(src1_t);
252
253        size_t s00 = nb00 / sizeof(src0_t);
254        size_t s01 = nb01 / sizeof(src0_t);
255        size_t s02 = nb02 / sizeof(src0_t);
256        size_t s03 = nb03 / sizeof(src0_t);
257
258        GGML_ASSERT(ne0 <= std::numeric_limits<uint32_t>::max());
259        GGML_ASSERT(ne1 <= std::numeric_limits<uint32_t>::max());
260        GGML_ASSERT(ne2 <= std::numeric_limits<uint32_t>::max());
261        GGML_ASSERT(ne3 <= std::numeric_limits<uint32_t>::max());
262
263      //GGML_ASSERT(s0  <= std::numeric_limits<uint32_t>::max());
264        GGML_ASSERT(s1  <= std::numeric_limits<uint32_t>::max());
265        GGML_ASSERT(s2  <= std::numeric_limits<uint32_t>::max());
266        GGML_ASSERT(s3  <= std::numeric_limits<uint32_t>::max());
267
268        GGML_ASSERT(s00 <= std::numeric_limits<uint32_t>::max());
269        GGML_ASSERT(s01 <= std::numeric_limits<uint32_t>::max());
270        GGML_ASSERT(s02 <= std::numeric_limits<uint32_t>::max());
271        GGML_ASSERT(s03 <= std::numeric_limits<uint32_t>::max());
272
273        GGML_ASSERT(s10 <= std::numeric_limits<uint32_t>::max());
274        GGML_ASSERT(s11 <= std::numeric_limits<uint32_t>::max());
275        GGML_ASSERT(s12 <= std::numeric_limits<uint32_t>::max());
276        GGML_ASSERT(s13 <= std::numeric_limits<uint32_t>::max());
277
278        GGML_ASSERT(cne1[0] <= std::numeric_limits<uint32_t>::max());
279        GGML_ASSERT(cne1[1] <= std::numeric_limits<uint32_t>::max());
280        GGML_ASSERT(cne1[2] <= std::numeric_limits<uint32_t>::max());
281        GGML_ASSERT(cne1[3] <= std::numeric_limits<uint32_t>::max());
282
283        GGML_ASSERT(nb0 % sizeof(dst_t) == 0);
284        GGML_ASSERT(nb1 % sizeof(dst_t) == 0);
285        GGML_ASSERT(nb2 % sizeof(dst_t) == 0);
286        GGML_ASSERT(nb3 % sizeof(dst_t) == 0);
287
288        GGML_ASSERT(nb00 % sizeof(src0_t) == 0);
289        GGML_ASSERT(nb01 % sizeof(src0_t) == 0);
290        GGML_ASSERT(nb02 % sizeof(src0_t) == 0);
291        GGML_ASSERT(nb03 % sizeof(src0_t) == 0);
292
293        GGML_ASSERT(nb10 % sizeof(src1_t) == 0);
294        GGML_ASSERT(nb11 % sizeof(src1_t) == 0);
295        GGML_ASSERT(nb12 % sizeof(src1_t) == 0);
296        GGML_ASSERT(nb13 % sizeof(src1_t) == 0);
297
298        GGML_ASSERT(ne2 * ne3 <= std::numeric_limits<unsigned int>::max());
299
300        const int block_size = 128;
301
302        int64_t hne0 = std::max(ne0 / 2LL, 1LL);
303
304        dim3 block_dims;
305        block_dims.x = std::min<unsigned int>(hne0, block_size);
306        block_dims.y = std::min<unsigned int>(ne1, block_size / block_dims.x);
307        block_dims.z = std::min(std::min<unsigned int>(ne2 * ne3, block_size / block_dims.x / block_dims.y), 64U);
308
309        dim3 block_nums((hne0 + block_dims.x - 1) / block_dims.x, (ne1 + block_dims.y - 1) / block_dims.y,
310                        (ne2 * ne3 + block_dims.z - 1) / block_dims.z);
311
312        const uint3 ne10 = init_fastdiv_values((uint32_t) cne1[0]);
313        const uint3 ne11 = init_fastdiv_values((uint32_t) cne1[1]);
314        const uint3 ne12 = init_fastdiv_values((uint32_t) cne1[2]);
315        const uint3 ne13 = init_fastdiv_values((uint32_t) cne1[3]);
316
317        if (block_nums.z > 65535 || block_nums.y > 65535) {
318            int64_t     block_num   = (ne0 * ne1 * ne2 * ne3 + block_size - 1) / block_size;
319
320            GGML_ASSERT(block_num              <= std::numeric_limits<uint32_t>::max());
321            GGML_ASSERT(block_num * block_size <= std::numeric_limits<uint32_t>::max());
322            GGML_ASSERT(ne0 * ne1              <= std::numeric_limits<uint32_t>::max());
323            GGML_ASSERT(ne0 * ne1 * ne2        <= std::numeric_limits<uint32_t>::max());
324
325            const uint3 prod_012    = init_fastdiv_values((uint32_t) (ne0 * ne1 * ne2));
326            const uint3 prod_01     = init_fastdiv_values((uint32_t) (ne0 * ne1));
327            const uint3 ne0_fastdiv = init_fastdiv_values((uint32_t) ne0);
328            const uint3 ne1_fastdiv = init_fastdiv_values((uint32_t) ne1);
329            const uint3 ne2_fastdiv = init_fastdiv_values((uint32_t) ne2);
330
331            {
332                const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)block_num, block_size, 0, stream);
333                ggml_cuda_kernel_launch(k_bin_bcast_unravel<bin_op, src0_t, src1_t, dst_t, type_for_index<const src1_t *, I>...>, launch_params,
334                    src0_dd, src1_dd, dst_dd, ne0_fastdiv, ne1_fastdiv, ne2_fastdiv, ne3, prod_012, prod_01, ne10, ne11,
335                    ne12, ne13,
336                  /*s0,*/ s1,  s2,  s3,
337                    s00, s01, s02, s03,
338                    s10, s11, s12, s13, (const src1_t *) dst->src[I + 1]->data...);
339            }
340        } else {
341            GGML_ASSERT(int64_t(block_nums.x) * block_dims.x <= std::numeric_limits<uint32_t>::max());
342            GGML_ASSERT(int64_t(block_nums.y) * block_dims.y <= std::numeric_limits<uint32_t>::max());
343            GGML_ASSERT(int64_t(block_nums.z) * block_dims.z <= std::numeric_limits<uint32_t>::max());
344
345            const uint3 ne3_fastdiv = init_fastdiv_values((uint32_t) ne3);
346            {
347                const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
348                ggml_cuda_kernel_launch(k_bin_bcast<bin_op, src0_t, src1_t, dst_t, type_for_index<const src1_t *, I>...>, launch_params,
349                    src0_dd, src1_dd, dst_dd, ne0, ne1, ne2, ne3_fastdiv, ne10, ne11, ne12, ne13,
350                  /*s0,*/ s1, s2,  s3,
351                    s00, s01, s02, s03,
352                    s10, s11, s12, s13, (const src1_t *) dst->src[I + 1]->data...);
353            }
354        }
355    }
356}
357
358template <typename T>
359static __global__ void k_repeat_back(
360    const T * __restrict__ src, T * __restrict__ dst, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
361    const size_t s00, const size_t s01, const size_t s02, const size_t s03,
362    const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3) {
363
364    const int64_t tid0  = int64_t(blockIdx.x)*blockDim.x + threadIdx.x;
365    const int64_t tid1  = int64_t(blockIdx.y)*blockDim.y + threadIdx.y;
366    const int64_t tid23 = int64_t(blockIdx.z)*blockDim.z + threadIdx.z;
367    const int64_t tid2  = tid23 % ne2;
368    const int64_t tid3  = tid23 / ne2;
369
370    if (tid0 >= ne0) {
371        return;
372    }
373
374    T sum = 0;
375    ggml_cuda_pdl_sync();
376    for (int64_t i3 = tid3; i3 < ne03; i3 += ne3) {
377        for (int64_t i2 = tid2; i2 < ne02; i2 += ne2) {
378            for (int64_t i1 = tid1; i1 < ne01; i1 += ne1) {
379                for (int64_t i0 = tid0; i0 < ne00; i0 += ne0) {
380                    sum += src[i3*s03 + i2*s02 + i1*s01 + i0*s00];
381                }
382            }
383        }
384    }
385    dst[tid3*ne2*ne1*ne0 + tid2*ne1*ne0 + tid1*ne0 + tid0] = sum;
386}
387
388template <float (*bin_op)(const float, const float), int n_fuse = 1>
389struct bin_bcast_cuda {
390    template<typename src0_t, typename src1_t, typename dst_t>
391    void operator()(const struct ggml_tensor * src0, const struct ggml_tensor * src1, struct ggml_tensor * dst,
392            const src0_t * src0_dd, const src1_t * src1_dd, dst_t * dst_dd,
393            cudaStream_t stream) {
394        launch_bin_bcast_pack<bin_op, src0_t, src1_t, dst_t>(
395            src0, src1, dst, src0_dd, src1_dd, dst_dd, stream, std::make_index_sequence<n_fuse>{});
396    }
397};
398
399template <typename T>
400static void repeat_back_cuda(
401    const T * src, T * dst, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
402    const size_t s00, const size_t s01, const size_t s02, const size_t s03,
403    const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
404
405    const dim3 block_dims(WARP_SIZE, 1, 1);
406    const dim3 block_nums((ne0 + WARP_SIZE - 1) / WARP_SIZE, ne1, ne2*ne3);
407    k_repeat_back<T><<<block_nums, block_dims, 0, stream>>>
408        (src, dst, ne00, ne01, ne02, ne03, s00, s01, s02, s03, ne0, ne1, ne2, ne3);
409}
410
411template<class op>
412static void ggml_cuda_op_bin_bcast(
413    const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
414    const void * src0_dd, const void * src1_dd, void * dst_dd, cudaStream_t stream) {
415
416    GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
417
418    if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
419        op()(src0, src1, dst, (const float *)src0_dd, (const float *)src1_dd, (float *)dst_dd, stream);
420    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
421        op()(src0, src1, dst, (const half *) src0_dd, (const half *)src1_dd, (half *) dst_dd, stream);
422    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
423        op()(src0, src1, dst, (const half *) src0_dd, (const float *)src1_dd, (half *) dst_dd, stream);
424    } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) {
425        op()(src0, src1, dst, (const half *) src0_dd, (const float *)src1_dd, (float *)dst_dd, stream);
426    } else {
427        fprintf(stderr, "%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__,
428            ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type));
429        GGML_ABORT("fatal error");
430    }
431}
432
433void ggml_cuda_op_repeat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
434    ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_repeat, 0>>(dst, dst->src[0], dst, nullptr, dst->src[0]->data, dst->data, ctx.stream());
435}
436
437void ggml_cuda_op_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
438    ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_add>>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream());
439}
440
441void ggml_cuda_op_sub(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
442    ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_sub>>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream());
443}
444
445void ggml_cuda_op_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
446    ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_mul>>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream());
447}
448
449void ggml_cuda_op_div(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
450    ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_div>>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream());
451}
452
453template <float (*op)(const float, const float), int n_fuse>
454static void ggml_cuda_op_fused_binbcast_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
455    cudaStream_t stream = ctx.stream();
456
457    const ggml_tensor * src0 = dst->src[0];
458    const ggml_tensor * src1 = dst->src[1];
459
460    if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
461        launch_bin_bcast_pack<op, float, float, float>(src0, src1, dst,
462            (const float *) src0->data, (const float *) src1->data, (float *) dst->data,
463            stream, std::make_index_sequence<n_fuse>{});
464    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
465        launch_bin_bcast_pack<op, half, half, half>(src0, src1, dst,
466            (const half *) src0->data, (const half *) src1->data, (half *) dst->data,
467            stream, std::make_index_sequence<n_fuse>{});
468    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
469        launch_bin_bcast_pack<op, half, float, half>(src0, src1, dst,
470            (const half *) src0->data, (const float *) src1->data, (half *) dst->data,
471            stream, std::make_index_sequence<n_fuse>{});
472    } else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32) {
473        launch_bin_bcast_pack<op, half, float, float>(src0, src1, dst,
474            (const half *) src0->data, (const float *) src1->data, (float *) dst->data,
475            stream, std::make_index_sequence<n_fuse>{});
476    } else {
477        fprintf(stderr,
478                "%s: unsupported types for fusion: dst: %s, src0: %s, src1: %s\n",
479                __func__, ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type));
480        GGML_ABORT("fatal error");
481    }
482}
483
484
485void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse) {
486    GGML_ASSERT(2 <= n_fuse && n_fuse <= 8);
487
488    switch (n_fuse) {
489        case 2:
490            ggml_cuda_op_fused_binbcast_impl<op_add, 2>(ctx, dst);
491            break;
492        case 3:
493            ggml_cuda_op_fused_binbcast_impl<op_add, 3>(ctx, dst);
494            break;
495        case 4:
496            ggml_cuda_op_fused_binbcast_impl<op_add, 4>(ctx, dst);
497            break;
498        case 5:
499            ggml_cuda_op_fused_binbcast_impl<op_add, 5>(ctx, dst);
500            break;
501        case 6:
502            ggml_cuda_op_fused_binbcast_impl<op_add, 6>(ctx, dst);
503            break;
504        case 7:
505            ggml_cuda_op_fused_binbcast_impl<op_add, 7>(ctx, dst);
506            break;
507        case 8:
508            ggml_cuda_op_fused_binbcast_impl<op_add, 8>(ctx, dst);
509            break;
510        default:
511            GGML_ASSERT(false && "Unsupported n_fuse value");
512    }
513}
514
515void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse) {
516    GGML_ASSERT(2 <= n_fuse && n_fuse <= 8);
517
518    switch (n_fuse) {
519        case 2:
520            ggml_cuda_op_fused_binbcast_impl<op_mul, 2>(ctx, dst);
521            break;
522        case 3:
523            ggml_cuda_op_fused_binbcast_impl<op_mul, 3>(ctx, dst);
524            break;
525        case 4:
526            ggml_cuda_op_fused_binbcast_impl<op_mul, 4>(ctx, dst);
527            break;
528        case 5:
529            ggml_cuda_op_fused_binbcast_impl<op_mul, 5>(ctx, dst);
530            break;
531        case 6:
532            ggml_cuda_op_fused_binbcast_impl<op_mul, 6>(ctx, dst);
533            break;
534        case 7:
535            ggml_cuda_op_fused_binbcast_impl<op_mul, 7>(ctx, dst);
536            break;
537        case 8:
538            ggml_cuda_op_fused_binbcast_impl<op_mul, 8>(ctx, dst);
539            break;
540        default:
541            GGML_ASSERT(false && "Unsupported n_fuse value");
542    }
543}
544
545void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
546    const ggml_tensor * src0 = dst->src[0];
547
548    GGML_ASSERT(src0->type == dst->type);
549    GGML_ASSERT(ggml_is_contiguous(dst));
550    GGML_ASSERT(ggml_can_repeat(dst, src0));
551
552    cudaStream_t stream = ctx.stream();
553
554    GGML_TENSOR_UNARY_OP_LOCALS;
555
556    GGML_ASSERT(ne2*ne3 <= (1 << 15));
557
558    const size_t ts = ggml_type_size(src0->type);
559    const size_t s00 = nb00 / ts;
560    const size_t s01 = nb01 / ts;
561    const size_t s02 = nb02 / ts;
562    const size_t s03 = nb03 / ts;
563
564    switch (dst->type) {
565        case GGML_TYPE_F32: {
566            const float * src0_d = (const float *) src0->data;
567            float       * dst_d  = (float       *) dst->data;
568            repeat_back_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s00, s01, s02, s03, ne0, ne1, ne2, ne3, stream);
569        } break;
570        default: {
571            GGML_ASSERT(false);
572        } break;
573    }
574}