char-fork/llama.cpp

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

Max Krasnyanskyhexagon: support for multi-device model split (aka row-split) (#28589)eafe15a5e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
22.0 KiB436 linesraw
1#pragma clang diagnostic ignored "-Wunused-variable"
2#pragma clang diagnostic ignored "-Wunused-function"
3#pragma clang diagnostic ignored "-Wunused-but-set-variable"
4
5#include <HAP_farf.h>
6#include <HAP_perf.h>
7#include <qurt_memory.h>
8
9#include <math.h>
10#include <string.h>
11
12#define GGML_COMMON_DECL_C
13#include "ggml-common.h"
14#include "htp-ctx.h"
15#include "htp-ops.h"
16#include "htp-ops.h"
17#include "hvx-utils.h"
18#include "htp-tensor.h"
19#include "htp-fence.h"
20
21struct htp_copy_context {
22    struct htp_ops_context * octx;
23
24    uint32_t          src0_type_size;
25    uint32_t          src0_block_size;
26
27    uint32_t          dst_type_size;
28    uint32_t          dst_block_size;
29
30    uint32_t          src0_blocks_per_row;
31    uint32_t          dst_blocks_per_row;
32
33    uint32_t          elem_start;
34    uint32_t          nelem;
35    uint32_t          elem_per_thread;
36
37    uint32_t          src0_nrows_per_thread;
38    uint32_t          row_start;
39    uint32_t          nrows;
40
41    struct fastdiv_values div_ne01;
42    struct fastdiv_values div_ne02_ne01;
43
44    struct fastdiv_values div_ne0;
45    struct fastdiv_values div_ne1_ne0;
46    struct fastdiv_values div_ne2_ne1_ne0;
47    struct fastdiv_values div_ne00;
48    struct fastdiv_values div_ne01_ne00;
49    struct fastdiv_values div_ne02_ne01_ne00;
50};
51
52#define cpy_preamble                              \
53    const struct htp_tensor *src0 = octx->src[0]; \
54    const struct htp_tensor *dst  = octx->dst;    \
55                                                  \
56    const uint32_t ne00 = src0->ne[0];            \
57    const uint32_t ne01 = src0->ne[1];            \
58    const uint32_t ne02 = src0->ne[2];            \
59    const uint32_t ne03 = src0->ne[3];            \
60                                                  \
61    const uint32_t nb00 = src0->nb[0];            \
62    const uint32_t nb01 = src0->nb[1];            \
63    const uint32_t nb02 = src0->nb[2];            \
64    const uint32_t nb03 = src0->nb[3];            \
65                                                  \
66    const uint32_t  ne0 = dst->ne[0];             \
67    const uint32_t  ne1 = dst->ne[1];             \
68    const uint32_t  ne2 = dst->ne[2];             \
69    const uint32_t  ne3 = dst->ne[3];             \
70                                                  \
71    const uint32_t  nb0 = dst->nb[0];             \
72    const uint32_t  nb1 = dst->nb[1];             \
73    const uint32_t  nb2 = dst->nb[2];             \
74    const uint32_t  nb3 = dst->nb[3];
75
76#define DEFINE_CPY_SAMESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                       \
77static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) {   \
78    struct htp_copy_context * ct = (struct htp_copy_context *) data;                           \
79    struct htp_ops_context * octx = ct->octx;                                                  \
80    cpy_preamble;                                                                              \
81    const uint32_t dr  = ct->src0_nrows_per_thread;                                            \
82    const uint32_t ir0 = ct->row_start + dr * ith;                                             \
83    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);                             \
84    if (ir0 >= ir1) return;                                                                    \
85    const bool contiguous = (nb01 == ne00 * ELEM_SIZE) && (nb1 == nb01) &&                     \
86                            (nb02 == ne01 * nb01)      && (nb2 == nb02) &&                     \
87                            (nb03 == ne02 * nb02)      && (nb3 == nb03);                       \
88    const uint32_t ne02_ne01 = ne02 * ne01;                                                    \
89    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);                                           \
90    uint32_t rem = ir0 - i03 * ne02_ne01;                                                      \
91    uint32_t i02 = fastdiv(rem, &ct->div_ne01);                                                \
92    uint32_t i01 = rem - i02 * ne01;                                                           \
93    uint8_t * dst_ptr  = (uint8_t *) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;               \
94    uint8_t * src0_ptr = (uint8_t *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;              \
95    if (contiguous) {                                                                          \
96        hvx_copy_uu(dst_ptr, src0_ptr, (ir1 - ir0) * ne00, ELEM_SIZE);                         \
97        return;                                                                                \
98    }                                                                                          \
99    for (uint32_t r = ir0; r < ir1; r++) {                                                     \
100        hex_l2fetch(src0_ptr, ne00 * ELEM_SIZE, nb01, 2);                                      \
101        hvx_copy_uu(dst_ptr, src0_ptr, ne00, ELEM_SIZE);                                       \
102        dst_ptr  += nb1;                                                                       \
103        src0_ptr += nb01;                                                                      \
104        if (++i01 == ne01) {                                                                   \
105            i01 = 0;                                                                           \
106            if (++i02 == ne02) {                                                               \
107                i02 = 0;                                                                       \
108                i03++;                                                                         \
109            }                                                                                  \
110            dst_ptr  = (uint8_t *) dst->data  + i02*nb2  + i03*nb3;                            \
111            src0_ptr = (uint8_t *) src0->data + i02*nb02 + i03*nb03;                           \
112        }                                                                                      \
113    }                                                                                          \
114}
115
116DEFINE_CPY_SAMESHAPE(f32,  float, 4)
117DEFINE_CPY_SAMESHAPE(f16, __fp16, 2)
118
119#define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE)                                               \
120static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void * data) {           \
121    struct htp_copy_context * ct = (struct htp_copy_context *) data;                                 \
122    struct htp_ops_context * octx = ct->octx;                                                        \
123    cpy_preamble;                                                                                    \
124    const uint32_t th_nelem = ct->elem_per_thread;                                                   \
125    const uint32_t th_start = ct->elem_start + ith * th_nelem;                                       \
126    const uint32_t th_end   = MIN(th_start + th_nelem, ct->elem_start + ct->nelem);                  \
127    if (th_start >= th_end) return;                                                                  \
128                                                                                                     \
129    const uint32_t ne01_ne00      = ne01 * ne00;                                                     \
130    const uint32_t ne02_ne01_ne00 = ne02 * ne01_ne00;                                                \
131    const uint32_t ne1_ne0        = ne1 * ne0;                                                       \
132    const uint32_t ne2_ne1_ne0    = ne2 * ne1_ne0;                                                   \
133                                                                                                     \
134    uint32_t e = th_start;                                                                           \
135    uint32_t i13 = fastdiv(e, &ct->div_ne2_ne1_ne0);                                                 \
136    uint32_t rem = e - i13 * ne2_ne1_ne0;                                                            \
137    uint32_t i12 = fastdiv(rem, &ct->div_ne1_ne0);                                                   \
138    uint32_t rem2 = rem - i12 * ne1_ne0;                                                             \
139    uint32_t i11 = fastdiv(rem2, &ct->div_ne0);                                                      \
140    uint32_t i10 = rem2 - i11 * ne0;                                                                 \
141                                                                                                     \
142    uint32_t i03 = fastdiv(e, &ct->div_ne02_ne01_ne00);                                              \
143    uint32_t rem_s = e - i03 * ne02_ne01_ne00;                                                       \
144    uint32_t i02 = fastdiv(rem_s, &ct->div_ne01_ne00);                                               \
145    uint32_t rem2_s = rem_s - i02 * ne01_ne00;                                                       \
146    uint32_t i01 = fastdiv(rem2_s, &ct->div_ne00);                                                   \
147    uint32_t i00 = rem2_s - i01 * ne00;                                                              \
148                                                                                                     \
149    char * dst_ptr        = (char *)       dst->data  + i10*nb0  + i11*nb1  + i12*nb2  + i13*nb3;    \
150    const char * src0_ptr = (const char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03;   \
151                                                                                                     \
152    for (; e < th_end; e++) {                                                                        \
153        *((ELEM_TYPE *) dst_ptr) = *((const ELEM_TYPE *) src0_ptr);                                  \
154                                                                                                     \
155        dst_ptr += nb0;                                                                              \
156        if (++i10 == ne0) {                                                                          \
157            i10 = 0;                                                                                 \
158            if (++i11 == ne1) {                                                                      \
159                i11 = 0;                                                                             \
160                if (++i12 == ne2) {                                                                  \
161                    i12 = 0;                                                                         \
162                    i13++;                                                                           \
163                }                                                                                    \
164            }                                                                                        \
165            dst_ptr = (char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3;                              \
166        }                                                                                            \
167                                                                                                     \
168        src0_ptr += nb00;                                                                            \
169        if (++i00 == ne00) {                                                                         \
170            i00 = 0;                                                                                 \
171            if (++i01 == ne01) {                                                                     \
172                i01 = 0;                                                                             \
173                if (++i02 == ne02) {                                                                 \
174                    i02 = 0;                                                                         \
175                    i03++;                                                                           \
176                }                                                                                    \
177            }                                                                                        \
178            src0_ptr = (const char *) src0->data + i01*nb01 + i02*nb02 + i03*nb03;                   \
179        }                                                                                            \
180    }                                                                                                \
181}
182
183DEFINE_CPY_RESHAPE(f32,  float, 4)
184DEFINE_CPY_RESHAPE(f16, __fp16, 2)
185
186static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
187    struct htp_copy_context * ct = (struct htp_copy_context *) data;
188    struct htp_ops_context * octx = ct->octx;
189    cpy_preamble;
190
191    const uint32_t dr  = ct->src0_nrows_per_thread;
192    const uint32_t ir0 = ct->row_start + dr * ith;
193    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
194    if (ir0 >= ir1) return;
195
196    const uint32_t ne02_ne01 = ne02 * ne01;
197    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
198    uint32_t rem = ir0 - i03 * ne02_ne01;
199    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
200    uint32_t i01 = rem - i02 * ne01;
201
202    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
203    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
204
205    for (uint32_t r = ir0; r < ir1; r++) {
206        hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
207        hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00);
208        dst_ptr  += nb1;
209        src0_ptr += nb01;
210        if (++i01 == ne01) {
211            i01 = 0;
212            if (++i02 == ne02) {
213                i02 = 0;
214                i03++;
215            }
216            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
217            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
218        }
219    }
220}
221
222static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, void * data) {
223    struct htp_copy_context * ct = (struct htp_copy_context *) data;
224    struct htp_ops_context * octx = ct->octx;
225    cpy_preamble;
226
227    const uint32_t dr  = ct->src0_nrows_per_thread;
228    const uint32_t ir0 = ct->row_start + dr * ith;
229    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
230    if (ir0 >= ir1) return;
231
232    const uint32_t ne02_ne01 = ne02 * ne01;
233    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
234    uint32_t rem = ir0 - i03 * ne02_ne01;
235    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
236    uint32_t i01 = rem - i02 * ne01;
237
238    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
239    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
240
241    for (uint32_t r = ir0; r < ir1; r++) {
242        hex_l2fetch(src0_ptr, ne00 * sizeof(__fp16), nb01, 2);
243        hvx_copy_f32_f16_uu(dst_ptr, src0_ptr, ne00);
244        dst_ptr  += nb1;
245        src0_ptr += nb01;
246        if (++i01 == ne01) {
247            i01 = 0;
248            if (++i02 == ne02) {
249                i02 = 0;
250                i03++;
251            }
252            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
253            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
254        }
255    }
256}
257
258static inline void cpy_dma_sametype_sameshape(
259    struct htp_ops_context * octx,
260    const struct htp_tensor * dst,
261    const struct htp_tensor * src0,
262    uint32_t elem_size,
263    uint32_t ne00, uint32_t ne01, uint32_t ne02, uint32_t ne03,
264    uint32_t nb01, uint32_t nb02, uint32_t nb03,
265    uint32_t  nb1, uint32_t  nb2, uint32_t nb3
266) {
267    const bool contiguous_outer =
268        (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) &&
269        (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2));
270
271    dma_queue * q = octx->ctx->dma[0];
272
273    if (contiguous_outer) {
274        if (!dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03)) {
275            dma_queue_flush(q);
276            dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03);
277        }
278        dma_queue_flush(q);
279        return;
280    }
281
282    for (uint32_t i03 = 0; i03 < ne03; i03++) {
283        for (uint32_t i02 = 0; i02 < ne02; i02++) {
284            uint8_t * dst_ptr  = (uint8_t *) dst->data  + i02 * nb2  + i03 * nb3;
285            uint8_t * src0_ptr = (uint8_t *) src0->data + i02 * nb02 + i03 * nb03;
286
287            if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) {
288                dma_queue_flush(q);
289                dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01);
290            }
291        }
292    }
293
294    dma_queue_flush(q);
295}
296
297static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
298    cpy_preamble;
299    *use_dma = false;
300
301    struct htp_copy_context ct;
302    ct.octx = octx;
303
304    switch (src0->type) {
305    case HTP_TYPE_F32: ct.src0_type_size = 4; ct.src0_block_size = 1; ct.src0_blocks_per_row = ne00 / 1; break;
306    case HTP_TYPE_F16: ct.src0_type_size = 2; ct.src0_block_size = 1; ct.src0_blocks_per_row = ne00 / 1; break;
307    default:
308        return HTP_STATUS_NO_SUPPORT;
309    }
310
311    switch (dst->type) {
312    case HTP_TYPE_F32: ct.dst_type_size = 4; ct.dst_block_size = 1; ct.dst_blocks_per_row = ne0 / 1; break;
313    case HTP_TYPE_F16: ct.dst_type_size = 2; ct.dst_block_size = 1; ct.dst_blocks_per_row = ne0 / 1; break;
314    default:
315        return HTP_STATUS_NO_SUPPORT;
316    }
317
318    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
319        return HTP_STATUS_OK;
320    }
321
322    const bool sametype   = (src0->type == dst->type);
323    const bool transposed = (nb00 > nb01) || (nb0 > nb1) ||
324                            (nb00 != ct.src0_type_size) || (nb0 != ct.dst_type_size) ||
325                            (nb01 < ne00 * ct.src0_type_size) || (nb1 < ne0 * ct.dst_type_size);
326    const bool sameshape  = !transposed && (ne00 == ne0 && ne01 == ne1 && ne02 == ne2 && ne03 == ne3);
327
328    const uint32_t n_threads = octx->n_threads;
329
330    const bool dst_is_contiguous = htp_tensor_is_contiguous(dst, ct.dst_type_size);
331
332    if (sameshape) {
333        const uint32_t total_rows = ne01 * ne02 * ne03;
334        const uint32_t row_size   = ne00 * ct.dst_type_size;
335
336        ct.div_ne01      = init_fastdiv_values(ne01);
337        ct.div_ne02_ne01 = init_fastdiv_values(ne02 * ne01);
338
339        uint32_t row_start = 0;
340        uint32_t nrows     = total_rows;
341
342        if (octx->ctx->mdev.count > 1) {
343            const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
344            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
345            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, can_split ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
346            row_start = range.start;
347            nrows     = range.count;
348        }
349
350        if (nrows == 0) {
351            return HTP_STATUS_OK;
352        }
353
354        ct.row_start = row_start;
355        ct.nrows     = nrows;
356        ct.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
357
358        if (sametype && octx->ctx->mdev.count <= 1) {
359            *use_dma = true;
360            cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3);
361        } else {
362            work_queue_func_t copy_fun = NULL;
363            if (sametype) {
364                copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_sameshape : cpy_thread_f16_sameshape;
365            } else if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) {
366                copy_fun = cpy_thread_f16_f32_sameshape;
367            } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) {
368                copy_fun = cpy_thread_f32_f16_sameshape;
369            } else {
370                return HTP_STATUS_NO_SUPPORT;
371            }
372            work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
373        }
374    } else if (sametype) {
375        const uint32_t total_elems = ne0 * ne1 * ne2 * ne3;
376        const uint32_t elems_per_line = (ct.dst_type_size == 4) ? 32 : 64;
377
378        ct.div_ne0            = init_fastdiv_values(ne0);
379        ct.div_ne1_ne0        = init_fastdiv_values(ne1 * ne0);
380        ct.div_ne2_ne1_ne0    = init_fastdiv_values(ne2 * ne1 * ne0);
381        ct.div_ne00           = init_fastdiv_values(ne00);
382        ct.div_ne01_ne00      = init_fastdiv_values(ne01 * ne00);
383        ct.div_ne02_ne01_ne00 = init_fastdiv_values(ne02 * ne01 * ne00);
384
385        uint32_t elem_start = 0;
386        uint32_t nelem      = total_elems;
387
388        if (octx->ctx->mdev.count > 1) {
389            const bool can_split = htp_tensor_mdev_data_aligned(dst) && dst_is_contiguous;
390            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_elems, can_split ? elems_per_line : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
391            elem_start = range.start;
392            nelem      = range.count;
393        }
394
395        if (nelem == 0) {
396            return HTP_STATUS_OK;
397        }
398
399        ct.elem_start      = elem_start;
400        ct.nelem           = nelem;
401        ct.elem_per_thread = fastdiv(nelem + n_threads - 1, &octx->n_threads_div);
402
403        work_queue_func_t copy_fun = (src0->type == HTP_TYPE_F32) ? cpy_thread_f32_reshape : cpy_thread_f16_reshape;
404        work_queue_run(octx->ctx->work_queue, copy_fun, &ct, n_threads);
405    } else {
406        return HTP_STATUS_NO_SUPPORT;
407    }
408
409    return HTP_STATUS_OK;
410}
411
412int op_cpy(struct htp_ops_context * octx) {
413    bool use_dma = false;
414    int status = exec_cpy(octx, &use_dma);
415
416    htp_ops_context_set_status(octx, status);
417
418    if (octx->op == HTP_OP_CPY_FENCE) {
419        if (!use_dma) {
420            htp_flush_dirty_ranges(octx->ctx);
421        }
422
423        htp_mdev_group_barrier(octx);
424
425        if (octx->ctx->mdev.idx == 0) {
426            const struct htp_tensor * sync = octx->src[1];
427            const uint32_t seq = (uint32_t) octx->op_params[0];
428            atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data;
429            htp_fence_write(sync_fence, seq, octx->status);
430
431            FARF(HIGH, "ggml-hex: sync-release : fence %p seq 0x%x status %d\n", sync_fence, seq, octx->status);
432        }
433    }
434
435    return octx->status;
436}